From e9815f060ac042422f5db184964eaa41e882ed79 Mon Sep 17 00:00:00 2001
From: Nick Sweeting <git@sweeting.me>
Date: Sat, 25 Jul 2026 18:49:53 -0700
Subject: [PATCH] Shard ArchiveBox tests to reduce CI queue time

---
 .github/scripts/discover_test_matrix.py | 48 +++++++++++++++++++++----
 .github/workflows/test-parallel.yml     | 13 ++++---
 2 files changed, 49 insertions(+), 12 deletions(-)

diff --git a/.github/scripts/discover_test_matrix.py b/.github/scripts/discover_test_matrix.py
index 2c3c3494..3cb6723b 100755
--- a/.github/scripts/discover_test_matrix.py
+++ b/.github/scripts/discover_test_matrix.py
@@ -1,9 +1,28 @@
 #!/usr/bin/env python3
-"""Discover every ArchiveBox test file for one CI matrix."""
+"""Discover every ArchiveBox test file and pack it into one CI matrix."""
 
 import json
 from pathlib import Path
 
+TARGET_SHARDS = 18
+
+SLOW_TEST_WEIGHTS = {
+    "test_cli_archiveresult.py": 200,
+    "test_opencode_agent.py": 180,
+    "test_cli_piping.py": 135,
+    "test_api_v1_crawls_crawl_crawl_id.py": 130,
+    "test_binary_service.py": 80,
+    "test_api_v1_workflow_core_token_auth_side_effects.py": 55,
+    "test_api_v1_cli_schedule.py": 50,
+    "test_api_v1_cli_remove.py": 40,
+    "test_api_v1_workflow_frozen_crawl_config_sources.py": 35,
+    "test_cli_help.py": 35,
+}
+
+
+def test_weight(path: Path) -> int:
+    return SLOW_TEST_WEIGHTS.get(path.name, 25)
+
 
 def main() -> None:
     root = Path.cwd().resolve()
@@ -12,21 +31,36 @@ def main() -> None:
     if not archivebox_tests:
         raise SystemExit("No ArchiveBox tests discovered")
 
+    shard_count = min(TARGET_SHARDS, len(archivebox_tests))
+    shards: list[list[Path]] = [[] for _ in range(shard_count)]
+    shard_weights = [0 for _ in range(shard_count)]
+
+    for path in sorted(archivebox_tests, key=lambda item: (-test_weight(item), item.as_posix())):
+        shard_index = min(range(shard_count), key=lambda index: (shard_weights[index], index))
+        shards[shard_index].append(path)
+        shard_weights[shard_index] += test_weight(path)
+
     matrix: list[dict[str, object]] = []
-    for path in archivebox_tests:
+    for index, shard in enumerate(shards, start=1):
+        shard_paths = sorted(path.relative_to(root).as_posix() for path in shard)
+        paths_arg = " ".join(shard_paths)
         matrix.append(
             {
-                "name": f"main/{path.stem.removeprefix('test_')}",
-                "path": path.relative_to(root).as_posix(),
-                "extra": "ldap" if path.name == "test_auth_ldap.py" else "",
+                "name": f"main/shard-{index:02d}",
+                "paths": shard_paths,
+                "paths_arg": paths_arg,
+                "extra": "ldap" if any(path.endswith("/test_auth_ldap.py") for path in shard_paths) else "",
+                "count": len(shard_paths),
             },
         )
 
-    discovered_paths = [str(entry["path"]) for entry in matrix]
+    discovered_paths = [path for entry in matrix for path in entry["paths"]]
     if len(discovered_paths) != len(set(discovered_paths)):
         raise SystemExit("Tests were not discovered exactly once")
+    if sorted(discovered_paths) != [path.relative_to(root).as_posix() for path in archivebox_tests]:
+        raise SystemExit("Discovered test shard coverage does not match test files")
 
-    print(f"Discovered {len(matrix)} test files exactly once")
+    print(f"Discovered {len(discovered_paths)} test files exactly once across {len(matrix)} balanced shards")
     print(json.dumps(matrix, separators=(",", ":")))
 
 
diff --git a/.github/workflows/test-parallel.yml b/.github/workflows/test-parallel.yml
index abb27f7a..f3421aba 100644
--- a/.github/workflows/test-parallel.yml
+++ b/.github/workflows/test-parallel.yml
@@ -64,7 +64,7 @@ jobs:
           test "${#output[@]}" -eq 2
           echo "tests=${output[1]}" >> "$GITHUB_OUTPUT"
           echo "${output[0]}"
-          "$JQ_BINARY" -e 'length > 0 and length <= 256' <<< "${output[1]}"
+          "$JQ_BINARY" -e 'length > 0 and length <= 32 and ([.[].paths[]] | length == (unique | length))' <<< "${output[1]}"
 
   tests:
     name: ${{ matrix.test.name }}
@@ -124,7 +124,7 @@ jobs:
           uv sync --locked --dev "${extra_args[@]}"
 
       - name: Install PostgreSQL server binaries
-        if: contains(matrix.test.path, 'test_postgres_backend')
+        if: contains(matrix.test.paths_arg, 'test_postgres_backend')
         run: |
           set -Eeuo pipefail
           if ! ls /usr/lib/postgresql/*/bin/initdb >/dev/null 2>&1 && ! command -v initdb >/dev/null 2>&1; then
@@ -133,10 +133,13 @@ jobs:
 
       - name: Run ${{ matrix.test.name }}
         env:
-          TEST_PATH: ${{ matrix.test.path }}
+          TEST_PATHS: ${{ matrix.test.paths_arg }}
         run: |
           set -Eeuo pipefail
-          test -f "$TEST_PATH"
+          read -r -a test_paths <<< "$TEST_PATHS"
+          for test_path in "${test_paths[@]}"; do
+            test -f "$test_path"
+          done
           mkdir -p "$DATA_DIR"
-          uv run --no-sync --no-sources pytest -vs "$TEST_PATH" \
+          uv run --no-sync --no-sources pytest -vs "${test_paths[@]}" \
             --basetemp="$RUNNER_TEMP/pytest"