(fix) Fix code coverage reporting

eng/pipelines/templates/steps/build-test.yml

@@ -141,6 +141,13 @@ steps:
     displayName: Report Coverage
     condition: and(succeeded(), ${{ parameters.RunCoverage }})
 
+  - task: PythonScript@0
+    displayName: Clean Up Test Isolates
+    condition: and(always(), ${{ parameters.RunCoverage }})
+    continueOnError: true
+    inputs:
+      scriptPath: eng/scripts/cleanup_isolate_dirs.py
+
   - ${{ if eq('true', parameters.UseFederatedAuth) }}:
     - task: AzurePowerShell@5
       displayName: Test Samples (AzurePowerShell@5)

eng/scripts/cleanup_isolate_dirs.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python
+
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import argparse
+import logging
+import os
+import shutil
+from typing import List
+
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+root_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+
+
+def find_isolate_dirs(repo_root: str) -> List[str]:
+    """Find azpysdk isolate directories under the repository-managed virtual environment root."""
+    isolate_root = os.path.join(os.path.abspath(repo_root), ".venv")
+    if not os.path.isdir(isolate_root):
+        return []
+
+    isolate_dirs = []
+    with os.scandir(isolate_root) as packages:
+        for package in packages:
+            if package.is_symlink() or not package.is_dir(follow_symlinks=False):
+                continue
+            with os.scandir(package.path) as environments:
+                isolate_dirs.extend(
+                    environment.path
+                    for environment in environments
+                    if environment.name.startswith(".venv_")
+                    and not environment.is_symlink()
+                    and environment.is_dir(follow_symlinks=False)
+                )
+    return sorted(isolate_dirs)
+
+
+def cleanup_isolate_dirs(repo_root: str) -> int:
+    """Remove azpysdk isolate directories and return the number that could not be removed."""
+    isolate_dirs = find_isolate_dirs(repo_root)
+    if not isolate_dirs:
+        logger.info("No azpysdk isolate directories found for cleanup.")
+        return 0
+
+    failures = 0
+    for isolate_dir in isolate_dirs:
+        try:
+            logger.info("Removing azpysdk isolate directory %s", isolate_dir)
+            shutil.rmtree(isolate_dir)
+        except OSError as exc:
+            failures += 1
+            logger.warning(
+                "Failed to remove isolate directory %s: %s", isolate_dir, exc
+            )
+
+    logger.info(
+        "Isolate directory cleanup complete: %d removed, %d failed.",
+        len(isolate_dirs) - failures,
+        failures,
+    )
+    return failures
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Remove azpysdk isolate directories after coverage reporting."
+    )
+    parser.add_argument(
+        "--repo-root",
+        default=root_dir,
+        help="Repository root containing the azpysdk .venv directory.",
+    )
+    args = parser.parse_args()
+    raise SystemExit(1 if cleanup_isolate_dirs(args.repo_root) else 0)

eng/scripts/dispatch_checks.py

@@ -9,7 +9,7 @@
 import subprocess
 import re
 from dataclasses import dataclass
-from typing import IO, List, Optional
+from typing import IO, List, Optional, Sequence
 
 from ci_tools.functions import discover_targeted_packages
 from ci_tools.variables import in_ci
@@ -55,6 +55,17 @@ class ProxyProcess:
     "latestdependency",
     "mindependency",
 }
+# Checks that actually emit package coverage data consumed by the post-dispatch
+# coverage report. Only these justify preserving isolate directories; other checks
+# (e.g. pylint, pyright, samples) produce no coverage, so their environments must be
+# cleaned up immediately even when coverage is not explicitly disabled.
+COVERAGE_PRODUCING_CHECKS = {
+    "whl",
+    "whl_no_aio",
+    "sdist",
+    "devtest",
+    "optional",
+}
 SHARED_RESTORE_ENV = "__shared_restore__"
 
 
@@ -73,6 +84,24 @@ def _cleanup_isolate_dirs() -> None:
     ISOLATE_DIRS_TO_CLEAN.clear()
 
 
+def _finalize_isolate_dirs(checks: Sequence[str], coverage_enabled: bool) -> None:
+    # Azure Pipelines generates the combined coverage report after dispatch completes.
+    # Only preserve isolate directories when the selected checks actually produce
+    # coverage; otherwise (e.g. pylint/pyright/samples runs) clean them up immediately
+    # so they do not accumulate for the remainder of the job.
+    produces_coverage = bool(set(checks) & COVERAGE_PRODUCING_CHECKS)
+    if coverage_enabled and produces_coverage and in_ci() == 1:
+        if ISOLATE_DIRS_TO_CLEAN:
+            logger.info(
+                "Preserving isolate directories for coverage report generation; "
+                "the CI cleanup step removes them afterward."
+            )
+        ISOLATE_DIRS_TO_CLEAN.clear()
+        return
+
+    _cleanup_isolate_dirs()
+
+
 def _normalize_newlines(text: str) -> str:
     return text.replace("\r\n", "\n").replace("\r", "\n")
 
@@ -385,6 +414,7 @@ async def run_all_checks(
     wheel_dir,
     mark_arg: Optional[str],
     injected_packages: str,
+    disablecov: bool,
     dest_dir: Optional[str] = None,
     service: Optional[str] = None,
 ):
@@ -403,6 +433,8 @@ async def run_all_checks(
     :rtype: int
     """
     base_args = [sys.executable, "-m", "azpysdk.main"]
+    if disablecov:
+        base_args.append("--disablecov")
     tasks = []
     semaphore = asyncio.Semaphore(max_parallel)
     combos = [(p, c) for p in packages for c in checks]
@@ -739,6 +771,7 @@ def handler(signum, frame):
                 temp_wheel_dir,
                 args.mark_arg,
                 args.injected_packages,
+                args.disablecov,
                 args.dest_dir,
                 effective_service,
             )
@@ -747,5 +780,5 @@ def handler(signum, frame):
         logger.error("Aborted by user.")
         exit_code = 130
     finally:
-        _cleanup_isolate_dirs()
+        _finalize_isolate_dirs(checks, coverage_enabled=not args.disablecov)
     sys.exit(exit_code)

eng/scripts/run_coverage.py

@@ -36,10 +36,16 @@
         required=True,
     )
 
+    parser.add_argument(
+        "--coverage-file",
+        dest="coverage_file",
+        help="The coverage data file to read. Defaults to <target_package>/.coverage.",
+    )
+
     args = parser.parse_args()
     pkg_details = ParsedSetup.from_path(args.target_package)
 
-    possible_coverage_file = os.path.join(args.target_package, ".coverage")
+    possible_coverage_file = args.coverage_file or os.path.join(args.target_package, ".coverage")
 
     if os.path.exists(possible_coverage_file):
         total_coverage = get_total_coverage(possible_coverage_file, coveragerc_file, pkg_details.name, args.repo_root)
@@ -74,4 +80,3 @@
                     )
                     exit(1)
 
-

eng/tools/azure-sdk-tools/azpysdk/install_and_test.py

@@ -7,6 +7,7 @@
 from .Check import Check, DEPENDENCY_TOOLS_REQUIREMENTS, PACKAGING_REQUIREMENTS, TEST_TOOLS_REQUIREMENTS
 
 from ci_tools.functions import is_error_code_5_allowed, install_into_venv
+from ci_tools.parsing import ParsedSetup
 from ci_tools.scenario.generation import create_package_and_install
 from ci_tools.variables import discover_repo_root, set_envvar_defaults
 from ci_tools.logging import logger
@@ -74,28 +75,48 @@ def run(self, args: argparse.Namespace) -> int:
                 results.append(install_result)
                 continue
 
+            coverage_enabled = self.coverage_enabled and not getattr(args, "disablecov", False)
+            coverage_file = self.get_coverage_file(package_dir) if coverage_enabled else None
             pytest_args = self._build_pytest_args(package_dir, args)
-            pytest_result = self.run_pytest(executable, staging_directory, package_dir, package_name, pytest_args)
+            pytest_result = self.run_pytest(
+                executable,
+                staging_directory,
+                package_dir,
+                package_name,
+                pytest_args,
+                coverage_file,
+            )
             if pytest_result != 0:
                 results.append(pytest_result)
                 continue
 
-            if not self.coverage_enabled:
+            if not coverage_enabled:
                 continue
 
-            coverage_result = self.check_coverage(executable, package_dir, package_name)
+            coverage_result = self.check_coverage(executable, package_dir, package_name, coverage_file)
             if coverage_result != 0:
                 results.append(coverage_result)
 
         return max(results) if results else 0
 
-    def check_coverage(self, executable: str, package_dir: str, package_name: str) -> int:
+    def get_coverage_file(self, package_dir: str) -> str:
+        return os.path.join(package_dir, f".coverage.{self.display_name}")
+
+    def check_coverage(
+        self,
+        executable: str,
+        package_dir: str,
+        package_name: str,
+        coverage_file: str,
+    ) -> int:
         coverage_command = [
             os.path.join(REPO_ROOT, "eng/scripts/run_coverage.py"),
             "-t",
             package_dir,
             "-r",
             REPO_ROOT,
+            "--coverage-file",
+            coverage_file,
         ]
         coverage_result = self.run_venv_command(executable, coverage_command, cwd=package_dir)
         if coverage_result.returncode != 0:
@@ -108,12 +129,20 @@ def check_coverage(self, executable: str, package_dir: str, package_name: str) -
         return 0
 
     def run_pytest(
-        self, executable: str, staging_directory: str, package_dir: str, package_name: str, pytest_args: List[str]
+        self,
+        executable: str,
+        staging_directory: str,
+        package_dir: str,
+        package_name: str,
+        pytest_args: List[str],
+        coverage_file: Optional[str] = None,
     ) -> int:
         pytest_command = ["pytest", *pytest_args]
 
         environment = os.environ.copy()
         environment.update({"PYTHONPYCACHEPREFIX": staging_directory})
+        if coverage_file:
+            environment["COVERAGE_FILE"] = coverage_file
 
         logger.info(f"Running pytest for {package_name} with command: {pytest_command}")
         logger.debug(f"with environment vars: {environment}")
@@ -188,10 +217,15 @@ def _install_common_requirements(self, executable: str, package_dir: str) -> Non
             logger.warning(f"Test tools requirements file not found at {TEST_TOOLS_REQUIREMENTS}.")
 
     def _build_pytest_args(self, package_dir: str, args: argparse.Namespace) -> List[str]:
+        extra_args = list(self.additional_pytest_args)
+        if self.coverage_enabled and not getattr(args, "disablecov", False):
+            namespace = ParsedSetup.from_path(package_dir).namespace
+            extra_args.append(f"--cov={namespace}")
+
         return self._build_pytest_args_base(
             package_dir,
             args,
             ignore_globs=["**/.venv*", "**/.venv*/**"],
-            extra_args=self.additional_pytest_args,
+            extra_args=extra_args,
             test_target=package_dir,
         )

eng/tools/azure-sdk-tools/azpysdk/main.py

@@ -74,6 +74,12 @@ def build_parser() -> argparse.ArgumentParser:
             "Passed through to 'uv venv --python'. Requires --isolate and uv."
         ),
     )
+    parser.add_argument(
+        "--disablecov",
+        action="store_true",
+        default=False,
+        help="Disable code coverage collection for test checks.",
+    )
 
     # mutually exclusive logging options
     log_group = parser.add_mutually_exclusive_group()
@@ -118,6 +124,12 @@ def build_parser() -> argparse.ArgumentParser:
             "Passed through to 'uv venv --python'. Requires --isolate and uv."
         ),
     )
+    common.add_argument(
+        "--disablecov",
+        action="store_true",
+        default=argparse.SUPPRESS,
+        help="Disable code coverage collection for test checks.",
+    )
     common.add_argument(
         "--service",
         default=None,

eng/tools/azure-sdk-tools/tests/test_cleanup_isolate_dirs.py

@@ -0,0 +1,59 @@
+import os
+import sys
+from unittest.mock import patch
+
+REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", ".."))
+if REPO_ROOT not in sys.path:
+    sys.path.insert(0, REPO_ROOT)
+
+from eng.scripts.cleanup_isolate_dirs import cleanup_isolate_dirs, find_isolate_dirs
+
+
+def test_find_isolate_dirs_returns_only_azpysdk_environments(tmp_path):
+    isolate_root = tmp_path / ".venv"
+    whl_dir = isolate_root / "azure-core" / ".venv_whl"
+    versioned_dir = isolate_root / "azure-core" / ".venv_whl_py311"
+    unrelated_dir = isolate_root / "azure-core" / "shared"
+    root_environment = isolate_root / "repository-environment"
+    for directory in (whl_dir, versioned_dir, unrelated_dir, root_environment):
+        directory.mkdir(parents=True)
+
+    assert find_isolate_dirs(os.fspath(tmp_path)) == sorted([os.fspath(whl_dir), os.fspath(versioned_dir)])
+
+
+def test_cleanup_isolate_dirs_removes_isolates_and_preserves_other_directories(
+    tmp_path,
+):
+    isolate_dir = tmp_path / ".venv" / "azure-core" / ".venv_whl"
+    unrelated_dir = tmp_path / ".venv" / "azure-core" / "shared"
+    isolate_dir.mkdir(parents=True)
+    unrelated_dir.mkdir()
+
+    assert cleanup_isolate_dirs(os.fspath(tmp_path)) == 0
+    assert not isolate_dir.exists()
+    assert unrelated_dir.exists()
+
+
+def test_cleanup_isolate_dirs_succeeds_when_root_does_not_exist(tmp_path):
+    assert cleanup_isolate_dirs(os.fspath(tmp_path)) == 0
+
+
+def test_cleanup_isolate_dirs_reports_failures_and_continues(tmp_path):
+    first_dir = tmp_path / ".venv" / "azure-core" / ".venv_whl"
+    second_dir = tmp_path / ".venv" / "azure-core" / ".venv_sdist"
+    first_dir.mkdir(parents=True)
+    second_dir.mkdir()
+
+    def remove_with_failure(path):
+        if path == os.fspath(first_dir):
+            raise OSError("directory is in use")
+        os.rmdir(path)
+
+    with patch(
+        "eng.scripts.cleanup_isolate_dirs.shutil.rmtree",
+        side_effect=remove_with_failure,
+    ):
+        assert cleanup_isolate_dirs(os.fspath(tmp_path)) == 1
+
+    assert first_dir.exists()
+    assert not second_dir.exists()