From e07036ad5dffb6af6f65389e37d8a99076bea1d0 Mon Sep 17 00:00:00 2001
From: Ed_ <edwardgz@gmail.com>
Date: Mon, 8 Jun 2026 00:46:12 -0400
Subject: [PATCH] feat(batcher): implement Batch dataclass and plan() function

---
 tests/batcher.py | 146 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 146 insertions(+)
 create mode 100644 tests/batcher.py

diff --git a/tests/batcher.py b/tests/batcher.py
new file mode 100644
index 00000000..c26d3d71
--- /dev/null
+++ b/tests/batcher.py
@@ -0,0 +1,146 @@
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from categorizer import CategoryRecord, FixtureClass, Speed
+
+@dataclass(frozen=True)
+class Batch:
+ tier: str
+ label: str
+ files: list[Path]
+ pytest_args: list[str]
+ estimated_seconds: float
+ skip_reason: str | None = None
+
+_TIER_ORDER: tuple[str, ...] = ("0", "1", "2", "3", "H", "P")
+
+_SPEED_SECONDS: dict[str, float] = {
+ "fast": 0.5,
+ "medium": 3.0,
+ "slow": 15.0,
+ "very_slow": 60.0,
+}
+
+def _est(r: CategoryRecord) -> float:
+ return _SPEED_SECONDS.get(r.speed.value, 3.0)
+
+def _env_set(name: str) -> bool:
+ return bool(os.environ.get(name))
+
+def _batches_for_unit(records: list[CategoryRecord], xdist: bool) -> list[Batch]:
+ by_group: dict[str, list[CategoryRecord]] = {}
+ for r in records:
+  by_group.setdefault(r.batch_group or "core", []).append(r)
+ batches: list[Batch] = []
+ for group in sorted(by_group):
+  files = [Path("tests") / r.filename for r in by_group[group]]
+  args: list[str] = ["--maxfail=10"]
+  if xdist:
+   args = ["-n", "auto"] + args
+  batches.append(Batch(
+   tier="1",
+   label=f"tier-1-unit-{group}",
+   files=files,
+   pytest_args=args,
+   estimated_seconds=sum(_est(r) for r in by_group[group]),
+  ))
+ return batches
+
+def _batches_for_mock_app(records: list[CategoryRecord]) -> list[Batch]:
+ by_group: dict[str, list[CategoryRecord]] = {}
+ for r in records:
+  by_group.setdefault(r.batch_group or "core", []).append(r)
+ batches: list[Batch] = []
+ for group in sorted(by_group):
+  files = [Path("tests") / r.filename for r in by_group[group]]
+  batches.append(Batch(
+   tier="2",
+   label=f"tier-2-mock_app-{group}",
+   files=files,
+   pytest_args=["--maxfail=5"],
+   estimated_seconds=sum(_est(r) for r in by_group[group]),
+  ))
+ return batches
+
+def _batches_for_live_gui(records: list[CategoryRecord]) -> list[Batch]:
+ if not records:
+  return []
+ files = [Path("tests") / r.filename for r in records]
+ return [Batch(
+  tier="3",
+  label="tier-3-live_gui",
+  files=files,
+  pytest_args=["--maxfail=1"],
+  estimated_seconds=sum(_est(r) for r in records),
+ )]
+
+def _batches_for_headless(records: list[CategoryRecord]) -> list[Batch]:
+ if not records:
+  return []
+ files = [Path("tests") / r.filename for r in records]
+ return [Batch(
+  tier="H",
+  label="tier-H-headless",
+  files=files,
+  pytest_args=["--maxfail=5"],
+  estimated_seconds=sum(_est(r) for r in records),
+ )]
+
+def _batches_for_performance(records: list[CategoryRecord]) -> list[Batch]:
+ if not records:
+  return []
+ files = [Path("tests") / r.filename for r in records]
+ return [Batch(
+  tier="P",
+  label="tier-P-performance",
+  files=files,
+  pytest_args=["--maxfail=1"],
+  estimated_seconds=sum(_est(r) for r in records),
+ )]
+
+def _batches_for_opt_in(records: list[CategoryRecord], include_opt_in: bool) -> list[Batch]:
+ batches: list[Batch] = []
+ for r in records:
+  files = [Path("tests") / r.filename]
+  skip_reason: str | None = None
+  if not include_opt_in:
+   skip_reason = "--include-opt-in not set"
+  elif r.filename.startswith("test_clean_install") and not _env_set("RUN_CLEAN_INSTALL_TEST"):
+   skip_reason = "RUN_CLEAN_INSTALL_TEST not set"
+  elif r.filename.startswith("test_docker_build") and not _env_set("RUN_DOCKER_TEST"):
+   skip_reason = "RUN_DOCKER_TEST not set"
+  batches.append(Batch(
+   tier="0",
+   label=f"tier-0-opt_in-{r.filename.removeprefix('test_').removesuffix('.py')}",
+   files=files,
+   pytest_args=["--maxfail=1"],
+   estimated_seconds=_est(r),
+   skip_reason=skip_reason,
+  ))
+ return batches
+
+def plan(
+ records: list[CategoryRecord],
+ *,
+ tiers: set[str] = set(_TIER_ORDER),
+ include_opt_in: bool = False,
+ xdist: bool = True,
+) -> list[Batch]:
+ by_fc: dict[FixtureClass, list[CategoryRecord]] = {fc: [] for fc in FixtureClass}
+ for r in records:
+  by_fc[r.fixture_class].append(r)
+ out: list[Batch] = []
+ if "0" in tiers:
+  out.extend(_batches_for_opt_in(by_fc[FixtureClass.OPT_IN], include_opt_in))
+ if "1" in tiers:
+  out.extend(_batches_for_unit(by_fc[FixtureClass.UNIT], xdist))
+ if "2" in tiers:
+  out.extend(_batches_for_mock_app(by_fc[FixtureClass.MOCK_APP]))
+ if "3" in tiers:
+  out.extend(_batches_for_live_gui(by_fc[FixtureClass.LIVE_GUI]))
+ if "H" in tiers:
+  out.extend(_batches_for_headless(by_fc[FixtureClass.HEADLESS]))
+ if "P" in tiers:
+  out.extend(_batches_for_performance(by_fc[FixtureClass.PERFORMANCE]))
+ out.sort(key=lambda b: (_TIER_ORDER.index(b.tier), b.label))
+ return out