From 04ac3407597611cf616bb22d81b71e81bd45668c Mon Sep 17 00:00:00 2001
From: xangcastle <xangcastle@gmail.com>
Date: Thu, 11 Jun 2026 06:36:56 -0600
Subject: [PATCH] fix: support recursive coverage and split postprocessing in
 pytest_main

---
 examples/pytest/BUILD.bazel                   |  9 ++++
 examples/pytest/run_coverage_dir_and_check.sh | 50 +++++++++++++++++++
 py/private/pytest.py.tmpl                     | 27 ++++++----
 3 files changed, 76 insertions(+), 10 deletions(-)
 create mode 100755 examples/pytest/run_coverage_dir_and_check.sh

diff --git a/examples/pytest/BUILD.bazel b/examples/pytest/BUILD.bazel
index f8965d868..f94e75512 100644
--- a/examples/pytest/BUILD.bazel
+++ b/examples/pytest/BUILD.bazel
@@ -170,3 +170,12 @@ sh_test(
         ":coverage_setup_test",
     ],
 )
+
+sh_test(
+    name = "coverage_dir_test",
+    srcs = ["run_coverage_dir_and_check.sh"],
+    data = [
+        ":coverage_manifest",
+        ":coverage_setup_test",
+    ],
+)
diff --git a/examples/pytest/run_coverage_dir_and_check.sh b/examples/pytest/run_coverage_dir_and_check.sh
new file mode 100755
index 000000000..e76d0c589
--- /dev/null
+++ b/examples/pytest/run_coverage_dir_and_check.sh
@@ -0,0 +1,50 @@
+#!/usr/bin/env bash
+# Verifies that when COVERAGE_DIR is set, pytest_main.py writes coverage
+# data to $COVERAGE_DIR/python_coverage.dat instead of COVERAGE_OUTPUT_FILE.
+# This exercises the --experimental_split_coverage_postprocessing code path.
+
+set -euo pipefail
+
+LAUNCHER="$TEST_SRCDIR/_main/examples/pytest/coverage_setup_test"
+MANIFEST="$TEST_SRCDIR/_main/examples/pytest/coverage_manifest.txt"
+
+[[ -x "$LAUNCHER" ]] || { echo "launcher not found or not executable: $LAUNCHER" >&2; exit 1; }
+[[ -f "$MANIFEST" ]] || { echo "manifest not found: $MANIFEST" >&2; exit 1; }
+
+WORK="$(mktemp -d)"
+trap 'rm -rf "$WORK"' EXIT
+
+COVERAGE_DIR="$WORK/cov_dir"
+mkdir -p "$COVERAGE_DIR"
+
+# COVERAGE_OUTPUT_FILE is set to a different path to prove the output does NOT
+# go there when COVERAGE_DIR is present.
+DECOY_OUTPUT="$WORK/decoy.lcov"
+
+COVERAGE_MANIFEST="$MANIFEST" \
+  COVERAGE_DIR="$COVERAGE_DIR" \
+  COVERAGE_OUTPUT_FILE="$DECOY_OUTPUT" \
+  "$LAUNCHER"
+
+EXPECTED="$COVERAGE_DIR/python_coverage.dat"
+[[ -s "$EXPECTED" ]] || {
+  echo "Expected $EXPECTED to exist and be non-empty." >&2
+  echo "Contents of COVERAGE_DIR:" >&2
+  ls -la "$COVERAGE_DIR" >&2
+  exit 1
+}
+
+# The decoy file must NOT have been written.
+[[ ! -f "$DECOY_OUTPUT" ]] || {
+  echo "Coverage was incorrectly written to COVERAGE_OUTPUT_FILE ($DECOY_OUTPUT) instead of COVERAGE_DIR." >&2
+  exit 1
+}
+
+# Basic shape check: SF: record must point to foo.py.
+grep -qE '^SF:.*examples/pytest/foo\.py$' "$EXPECTED" || {
+  echo "Expected SF: record for examples/pytest/foo.py not found." >&2
+  cat "$EXPECTED" >&2
+  exit 1
+}
+
+echo "OK: coverage data written to \$COVERAGE_DIR/python_coverage.dat as expected."
diff --git a/py/private/pytest.py.tmpl b/py/private/pytest.py.tmpl
index 60503bcd9..53805c276 100644
--- a/py/private/pytest.py.tmpl
+++ b/py/private/pytest.py.tmpl
@@ -35,13 +35,18 @@ if "COVERAGE_MANIFEST" in os.environ:
     try:
         import coverage
         # The lines are files that matched the --instrumentation_filter flag
-        with open(os.getenv("COVERAGE_MANIFEST"), "r") as mf:
-            manifest_entries = mf.read().splitlines()
-            cov = coverage.Coverage(include = manifest_entries)
-            # coveragepy incorrectly converts our entries by following symlinks
-            # record a mapping of their conversion so we can undo it later in reporting the coverage
-            coveragepy_absfile_mapping = {coverage.files.abs_file(mfe): mfe for mfe in manifest_entries}
-        cov.start()
+
+        existing_cov = coverage.Coverage.current()
+        if existing_cov is not None:
+            cov = existing_cov
+        else:
+            with open(os.getenv("COVERAGE_MANIFEST"), "r") as mf:
+                manifest_entries = mf.read().splitlines()
+                cov = coverage.Coverage(include = manifest_entries)
+                # coveragepy incorrectly converts our entries by following symlinks
+                # record a mapping of their conversion so we can undo it later in reporting the coverage
+                coveragepy_absfile_mapping = {coverage.files.abs_file(mfe): mfe for mfe in manifest_entries}
+            cov.start()
     except ModuleNotFoundError as e:
         print("WARNING: python coverage setup failed. Do you need to include the 'coverage' package as a dependency of py_pytest_main?", e)
         pass
@@ -108,12 +113,14 @@ if __name__ == "__main__":
         # https://bazel.build/configure/coverage
         coverage_output_file = os.getenv("COVERAGE_OUTPUT_FILE")
 
-        unfixed_dat = coverage_output_file + ".tmp"
+        coverage_dir = os.environ.get("COVERAGE_DIR")
+        target_dat = os.path.join(coverage_dir, "python_coverage.dat") if coverage_dir else coverage_output_file
+        unfixed_dat = target_dat + ".tmp"
         cov.lcov_report(outfile = unfixed_dat)
         cov.save()
-        
+
         with open(unfixed_dat, "r") as unfixed:
-          with open(coverage_output_file, "w") as output_file:
+          with open(target_dat, "w") as output_file:
             for line in unfixed:
               # Workaround https://github.com/nedbat/coveragepy/issues/963
               # by mapping SF: records to un-do the symlink-following