Run the tests in the release run's six shards on pull requests

.github/workflows/pr.yml

@@ -31,6 +31,8 @@ jobs:
       - test_cloud_run_container_builds
       - test_env_vars
       - test
+      - check-test-shards
+      - upload-test-coverage
     runs-on: ubuntu-latest
     permissions:
       contents: read
@@ -249,7 +251,11 @@ jobs:
           HUGGING_FACE_TOKEN: ${{ secrets.PE_UK_PRIVATE_HF_READ_TOKEN }}
           POLICYENGINE_DB_PASSWORD: ${{ secrets.POLICYENGINE_DB_PASSWORD }}
   test:
-    name: Test
+    # make test on six runners, one shard each, split by file as push.yml's
+    # release-tests job splits it (scripts/release_test_shards.py), so a
+    # change that breaks the split fails here, before it merges.
+    # check-test-shards checks that each test ran in exactly one shard.
+    name: Test (shard ${{ matrix.shard_index }} of 6)
     runs-on: ubuntu-latest
     needs:
       - ensure-policyengine-bundle-supported-by-simulation-api
@@ -261,6 +267,11 @@ jobs:
       # uv installs uv.lock on the Cloud Run image's Python, not the 3.14
       # that .python-version selects.
       UV_PYTHON: "3.12"
+    strategy:
+      # Every shard runs to the end, so one run shows every failure.
+      fail-fast: false
+      matrix:
+        shard_index: [1, 2, 3, 4, 5, 6]
     steps:
       - name: Checkout repo
         uses: actions/checkout@v4
@@ -287,15 +298,80 @@ jobs:
           project_id: policyengine-api
       - name: Install locked dependencies
         run: make install
+      # PYTEST_ADDOPTS loads the shard plugin as run_release_tests.sh does.
       - name: Test the API
         run: make test
         env:
+          PYTEST_ADDOPTS: -p tests.release_test_shard_plugin
+          RELEASE_TEST_SHARD_INDEX: ${{ matrix.shard_index }}
+          RELEASE_TEST_SHARD_COUNT: "6"
+          RELEASE_TEST_SHARD_REPORT_DIRECTORY: ${{ runner.temp }}/release-test-shard
           POLICYENGINE_DB_PASSWORD: ${{ secrets.POLICYENGINE_DB_PASSWORD }}
           POLICYENGINE_GITHUB_MICRODATA_AUTH_TOKEN: ${{ secrets.POLICYENGINE_GITHUB_MICRODATA_AUTH_TOKEN }}
           OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
           HUGGING_FACE_TOKEN: ${{ secrets.PE_UK_PRIVATE_HF_READ_TOKEN }}
+      - name: Upload the tests this shard collected and ran
+        uses: actions/upload-artifact@v4
+        with:
+          name: release-test-shard-${{ matrix.shard_index }}
+          path: ${{ runner.temp }}/release-test-shard
+          if-no-files-found: error
+          retention-days: 7
+      # upload-test-coverage sends every shard's report to Codecov at once.
+      - name: Keep this shard's coverage report
+        uses: actions/upload-artifact@v4
+        with:
+          name: test-coverage-shard-${{ matrix.shard_index }}
+          path: coverage.xml
+          if-no-files-found: error
+          retention-days: 7
+
+  upload-test-coverage:
+    # Sends the six shards' coverage reports to Codecov in one upload, so
+    # Codecov gets one upload for the test suite, as it did before the suite
+    # was split. With an upload from each shard, Codecov updated its checks
+    # after each upload and left codecov/patch in progress on #4145.
+    name: Upload the test coverage to Codecov
+    runs-on: ubuntu-latest
+    needs: test
+    permissions:
+      contents: read
+    steps:
+      - name: Checkout repo
+        uses: actions/checkout@v4
+      - name: Download every shard's coverage report
+        uses: actions/download-artifact@v4
+        with:
+          pattern: test-coverage-shard-*
+          path: ${{ runner.temp }}/test-coverage
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v5
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
           slug: PolicyEngine/policyengine-api
+          directory: ${{ runner.temp }}/test-coverage
+
+  check-test-shards:
+    # Fails unless every shard collected the same tests and each test ran in
+    # exactly one shard (scripts/release_test_shards.py check-coverage). It
+    # needs every shard, so it is skipped when one fails, and CI complete
+    # fails with it.
+    name: Check that every test ran in exactly one shard
+    runs-on: ubuntu-latest
+    needs: test
+    permissions:
+      contents: read
+    steps:
+      - name: Checkout repo
+        uses: actions/checkout@v4
+      - name: Setup Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - name: Download every shard's report
+        uses: actions/download-artifact@v4
+        with:
+          pattern: release-test-shard-*
+          path: ${{ runner.temp }}/release-test-shards
+      - name: Check that every release test ran in exactly one shard
+        run: python -m scripts.release_test_shards check-coverage --shard-count 6 "${{ runner.temp }}/release-test-shards"

changelog.d/pull-request-test-shards.changed.md

@@ -0,0 +1 @@
+Pull requests run the test suite in the release run's six shards, check that each test ran in exactly one, and send Codecov the six shards' coverage in one upload.

docs/engineering/skills/testing.md

@@ -188,10 +188,18 @@ the same tests and each test ran in exactly one shard. The plugin reads its
 settings once, when pytest configures it, and then removes them and its `-p`
 option from `PYTEST_ADDOPTS`, so a pytest process that a test starts collects
 its own tests without the split; a shard writes its report once, and fails if
-one is already there. When a test file
-becomes much slower or faster, measure it again with `--durations=0` and
-update `MEASURED_FILE_SECONDS`; an out-of-date value only makes the shards
-finish less evenly.
+one is already there. When a test file becomes much slower or faster, measure
+it again with `--durations=0` and update `MEASURED_FILE_SECONDS`; an
+out-of-date value only makes the shards finish less evenly.
+
+Every pull request runs the same six shards in `pr.yml`'s `test` job, which
+runs `make test` with the plugin in `PYTEST_ADDOPTS`, as
+`run_release_tests.sh` does, and its `check-test-shards` job runs the same
+check, so a change that breaks the split fails before it merges. Each shard
+keeps its `coverage.xml` as an artifact, and the `upload-test-coverage` job
+sends all six to Codecov in one upload: with an upload from each shard, Codecov
+updated its checks after every upload and could leave one in progress. "CI
+complete" needs all three jobs.
 
 The release run skips lint, the quality guards, both Alembic checks and the
 cross-database integration tests when its merge run passed them on the same

docs/generated/rules.md

@@ -338,7 +338,7 @@ flowchart LR
 | `api.quality.ruff-format` | All Python files are formatted with ruff format. | mechanical | enforced | yes | ruff format | `pr.yml:lint` | — |
 | `api.quality.ruff-lint` | All Python code passes ruff check with every ruff rule selected (ALL), with no ratchet, except the ignores in pyproject.toml: COM812 and ISC001, which conflict with ruff format, and D203 and D213, which conflict with the kept D211 and D212. | mechanical | enforced | yes | ruff check (ALL) | `pr.yml:lint` | — |
 | `api.quality.runtime-evaluated-models-resolve` | Every package model is one ruff treats as runtime-evaluated (a base resolves, through the import in its own module, to a name in runtime-evaluated-base-classes, or is a same-module class that does), and every Pydantic model and SQLAlchemy mapper resolves its annotations once the package is imported. | mechanical | enforced | yes | pytest `tests/unit/test_runtime_evaluated_annotations.py` | `pr.yml:test` | **unproven** |
-| `api.quality.tests-pass` | The unit, contract and to-refactor test suites pass. | mechanical | enforced | yes | CI job `test` | `pr.yml:test` | — |
+| `api.quality.tests-pass` | The unit, contract and to-refactor test suites pass, each test in exactly one of the shards they run in. | mechanical | enforced | yes | CI job `test`, `check-test-shards` | `pr.yml:test`, `pr.yml:check-test-shards` | — |
 | `api.quality.tests-ruff-exemptions` | Test files may use pytest's bare assert, need no docstrings, annotations or named constants for magic numbers, and may reach into private members, live outside packages and import inside a test (S101, D1, ANN, PLR2004, SLF001, INP001, PLC0415). Their fakes may name methods after the interface they replace, such as a gRPC stub's Export (N802). They may also use fake credentials, naive datetimes, subprocess calls to repository scripts, SQL built from strings and prints (S105, S106, DTZ001, DTZ011, S603, S607, S608, T201), may import the lazily loaded packages at module level (TID253), and may leave a stand-in's arguments unused (ARG001, ARG002, ARG005). Every other ruff rule applies to tests. | mechanical | enforced | yes | ruff per-file ignore of ANN, ARG001, ARG002, ARG005, D1, DTZ001, DTZ011, INP001, N802, PLC0415, PLR2004, S101, S105, S106, S603, S607, S608, SLF001, T201, TID253 for `**/tests/**` | `pr.yml:lint` | — |
 | `api.quality.unittest-ruff-exemptions` | The shared rules tool's tests use unittest, because the tool needs only the standard library, so ruff's pytest-style rules (PT) don't apply to them. | mechanical | enforced | yes | ruff per-file ignore of PT for `scripts/rules/tests/**` | `pr.yml:lint` | — |
 
@@ -370,7 +370,8 @@ flowchart LR
   e_pytest_tests_unit_test_runtime_evaluated_annotations_py["pytest tests/unit/test_runtime_evaluated_annotations.py"]
   j_CI_pr_yml_test["CI pr.yml:test"]
   r_api_quality_tests_pass("api.quality.tests-pass (enforced)")
-  e_CI_job_test["CI job test"]
+  e_CI_job_test__check_test_shards["CI job test, check-test-shards"]
+  j_CI_pr_yml_check_test_shards["CI pr.yml:check-test-shards"]
   r_api_quality_tests_ruff_exemptions("api.quality.tests-ruff-exemptions (enforced)")
   e_ruff_per_file_ignore_of_ANN__ARG001__ARG002__ARG005__D1__DTZ001__DTZ011__INP001__N802__PLC0415__PLR2004__S101__S105__S106__S603__S607__S608__SLF001__T201__TID253_for____tests___["ruff per-file ignore of ANN, ARG001, ARG002, ARG005, D1, DTZ001, DTZ011, INP001, N802, PLC0415, PLR2004, S101, S105, S106, S603, S607, S608, SLF001, T201, TID253 for **/tests/**"]
   r_api_quality_unittest_ruff_exemptions("api.quality.unittest-ruff-exemptions (enforced)")
@@ -394,8 +395,9 @@ flowchart LR
   e_ruff_check__ALL_ --> j_CI_pr_yml_lint
   r_api_quality_runtime_evaluated_models_resolve --> e_pytest_tests_unit_test_runtime_evaluated_annotations_py
   e_pytest_tests_unit_test_runtime_evaluated_annotations_py --> j_CI_pr_yml_test
-  r_api_quality_tests_pass --> e_CI_job_test
-  e_CI_job_test --> j_CI_pr_yml_test
+  r_api_quality_tests_pass --> e_CI_job_test__check_test_shards
+  e_CI_job_test__check_test_shards --> j_CI_pr_yml_test
+  e_CI_job_test__check_test_shards --> j_CI_pr_yml_check_test_shards
   r_api_quality_tests_ruff_exemptions --> e_ruff_per_file_ignore_of_ANN__ARG001__ARG002__ARG005__D1__DTZ001__DTZ011__INP001__N802__PLC0415__PLR2004__S101__S105__S106__S603__S607__S608__SLF001__T201__TID253_for____tests___
   e_ruff_per_file_ignore_of_ANN__ARG001__ARG002__ARG005__D1__DTZ001__DTZ011__INP001__N802__PLC0415__PLR2004__S101__S105__S106__S603__S607__S608__SLF001__T201__TID253_for____tests___ --> j_CI_pr_yml_lint
   r_api_quality_unittest_ruff_exemptions --> e_ruff_per_file_ignore_of_PT_for_scripts_rules_tests___

rules/registry.toml

@@ -57,6 +57,9 @@ title = "Process and repository hygiene"
 
 [discovery]
 workflows = [".github/workflows/pr.yml"]
+ignore_jobs = [
+  { job = "upload-test-coverage", reason = "sends the test shards' coverage reports to Codecov in one upload; it checks nothing" },
+]
 guard_test_globs = [
   "tests/**/test_*boundar*.py",
   "tests/**/test_import_side_effects.py",
@@ -1061,7 +1064,7 @@ invocation = "mypy"
 [[rule]]
 id = "api.quality.tests-pass"
 category = "quality"
-statement = "The unit, contract and to-refactor test suites pass."
+statement = "The unit, contract and to-refactor test suites pass, each test in exactly one of the shards they run in."
 rationale = "docs/engineering/skills/testing.md"
 scope = ["tests/**/*.py"]
 tier = "mechanical"
@@ -1073,6 +1076,10 @@ kind = "workflow-job"
 workflow = ".github/workflows/pr.yml"
 job = "test"
 invocation = "make test"
+[[rule.ci]]
+workflow = ".github/workflows/pr.yml"
+job = "check-test-shards"
+invocation = "scripts.release_test_shards check-coverage"
 
 [[rule]]
 id = "api.quality.environment-variables-load"

tests/unit/cloud_run_deploy_support.py

@@ -310,6 +310,11 @@ def push_workflow_jobs() -> dict[str, dict]:
     return yaml.safe_load(push_workflow())["jobs"]
 
 
+def pr_workflow_jobs() -> dict[str, dict]:
+    """Return pr.yml's jobs, parsed, keyed by job id."""
+    return yaml.safe_load(pr_workflow())["jobs"]
+
+
 def job_needs(job: dict) -> list[str]:
     """Return the ids of the jobs ``job`` needs, as a list."""
     needs = job.get("needs") or []

tests/unit/test_release_test_shards.py

@@ -3,6 +3,7 @@
 push.yml's release-tests job runs make test on six runners, one shard each
 (scripts/release_test_shards.py, tests/release_test_shard_plugin.py), and
 check-release-test-shards fails unless each collected test ran in exactly one.
+pr.yml's test and check-test-shards jobs do the same on every pull request.
 """
 
 from __future__ import annotations
@@ -20,6 +21,7 @@
     REPO,
     job_needs,
     job_step,
+    pr_workflow_jobs,
     push_workflow_jobs,
 )
 
@@ -551,3 +553,91 @@ def test_run_release_tests_loads_the_shard_plugin():
         "RELEASE_TEST_SHARD_REPORT_DIRECTORY",
     ):
         assert f"${{{variable}:?" in script
+
+
+# --- pr.yml ---------------------------------------------------------------------
+
+
+def test_pull_requests_run_the_tests_in_the_release_runs_shards():
+    # A pull request runs the split the release run does, so a change that
+    # breaks the split fails before it merges. Nothing ran the split before
+    # #4134 merged, so release run 41063091599 was the first to find #4139.
+    pull_request_tests = pr_workflow_jobs()["test"]
+    release_tests = push_workflow_jobs()["release-tests"]
+    pull_request_run_step = job_step(pull_request_tests, "Test the API")
+    release_run_step = job_step(release_tests, "Run release tests")
+    upload_step_name = "Upload the tests this shard collected and ran"
+    release_script = (REPO / ".github/scripts/run_release_tests.sh").read_text(
+        encoding="utf-8"
+    )
+
+    assert pull_request_tests["strategy"] == release_tests["strategy"]
+    # The job runs make test itself, which the rules registry's pr.yml rules
+    # look for, and loads the shard plugin as run_release_tests.sh does.
+    assert pull_request_run_step["run"] == "make test"
+    pytest_options = pull_request_run_step["env"]["PYTEST_ADDOPTS"]
+    assert f'PYTEST_ADDOPTS="{pytest_options}"' in release_script
+    assert {
+        name: value
+        for name, value in pull_request_run_step["env"].items()
+        if name.startswith("RELEASE_TEST_SHARD_")
+    } == {
+        name: value
+        for name, value in release_run_step["env"].items()
+        if name.startswith("RELEASE_TEST_SHARD_")
+    }
+    assert job_step(pull_request_tests, upload_step_name) == job_step(
+        release_tests, upload_step_name
+    )
+
+
+def test_ci_complete_needs_every_pull_request_test_to_have_run_once():
+    jobs = pr_workflow_jobs()
+    shard_check = jobs["check-test-shards"]
+
+    assert job_needs(shard_check) == ["test"]
+    # No condition of its own: when a shard fails, or on a stacked pull request
+    # whose shards were skipped, this job is skipped and CI complete fails.
+    assert "if" not in shard_check
+    assert (
+        shard_check["steps"]
+        == push_workflow_jobs()["check-release-test-shards"]["steps"]
+    )
+    assert "check-test-shards" in job_needs(jobs["ci-complete"])
+
+
+def test_codecov_gets_one_upload_for_every_pull_request_shard():
+    # With an upload from each shard, Codecov updated its checks after each of
+    # the seven uploads #4145's run sent and left codecov/patch in progress.
+    # The shards hand their reports to one job that uploads them together, so
+    # Codecov gets one upload for the test suite, as it did before the split.
+    jobs = pr_workflow_jobs()
+    shard = jobs["test"]
+    coverage_upload = jobs["upload-test-coverage"]
+    report_step = job_step(shard, "Keep this shard's coverage report")
+    download_step = job_step(coverage_upload, "Download every shard's coverage report")
+    codecov_steps = [
+        step
+        for step in coverage_upload["steps"]
+        if step.get("uses", "").startswith("codecov/codecov-action@")
+    ]
+
+    assert not any(
+        step.get("uses", "").startswith("codecov/codecov-action@")
+        for step in shard["steps"]
+    )
+    assert report_step["uses"] == "actions/upload-artifact@v4"
+    assert report_step["with"] == {
+        "name": "test-coverage-shard-${{ matrix.shard_index }}",
+        "path": "coverage.xml",
+        "if-no-files-found": "error",
+        "retention-days": 7,
+    }
+    assert job_needs(coverage_upload) == ["test"]
+    # No condition of its own, like check-test-shards.
+    assert "if" not in coverage_upload
+    assert download_step["uses"] == "actions/download-artifact@v4"
+    assert download_step["with"]["pattern"] == "test-coverage-shard-*"
+    assert len(codecov_steps) == 1
+    assert codecov_steps[0]["with"]["directory"] == download_step["with"]["path"]
+    assert "upload-test-coverage" in job_needs(jobs["ci-complete"])