OBLITERATUS/scripts/select_pr_tests.py

#!/usr/bin/env python3
"""Select deterministic PR tests from the exact diff and source risk map."""

from __future__ import annotations

import argparse
import fnmatch
import json
import subprocess
from pathlib import Path
from typing import Any


DEFAULT_POLICY = Path("ci/pr-test-policy.json")
DEFAULT_RISK_MAP = Path("ci/test-risk-map.json")


def _load_object(path: Path, label: str) -> dict[str, Any]:
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise ValueError(f"cannot read {label}: {exc}") from exc
    if not isinstance(value, dict):
        raise ValueError(f"{label} root must be an object")
    return value


def _risk_map_at_ref(base_ref: str, risk_map: Path) -> dict[str, Any] | None:
    """Read the base risk map when it exists so deleted sources remain owned."""

    try:
        content = subprocess.run(
            ["git", "show", f"{base_ref}:{risk_map.as_posix()}"],
            check=True,
            capture_output=True,
            text=True,
        ).stdout
    except subprocess.CalledProcessError:
        return None
    try:
        value = json.loads(content)
    except json.JSONDecodeError as exc:
        raise ValueError(f"base test risk map is invalid JSON: {exc}") from exc
    if not isinstance(value, dict):
        raise ValueError("base test risk map root must be an object")
    return value


def changed_paths(base_ref: str) -> list[str]:
    """Return exact paths changed from base_ref to the checked-out head."""

    try:
        subprocess.run(
            ["git", "rev-parse", "--verify", f"{base_ref}^{{commit}}"],
            check=True,
            capture_output=True,
            text=True,
        )
        raw = subprocess.run(
            ["git", "diff", "--name-only", "-z", base_ref, "HEAD"],
            check=True,
            capture_output=True,
        ).stdout
    except subprocess.CalledProcessError as exc:
        raise ValueError(f"cannot calculate changed paths from {base_ref!r}") from exc
    return sorted(path.decode("utf-8") for path in raw.split(b"\0") if path)


def _surface_tests(risk_map: dict[str, Any]) -> dict[str, set[str]]:
    result: dict[str, set[str]] = {}
    surfaces = risk_map.get("contract_surfaces")
    if not isinstance(surfaces, list):
        raise ValueError("test risk map requires contract_surfaces")
    for surface in surfaces:
        if not isinstance(surface, dict):
            raise ValueError("test risk map contract surface must be an object")
        paths = surface.get("paths")
        tests = surface.get("required_tests")
        if not isinstance(paths, list) or not isinstance(tests, list):
            raise ValueError("test risk map contract surface requires paths and required_tests")
        selected = {test for test in tests if isinstance(test, str)}
        for path in paths:
            if not isinstance(path, str):
                raise ValueError("test risk map source paths must be strings")
            result.setdefault(path, set()).update(selected)
    return result


def _production_source(path: str) -> bool:
    return path == "app.py" or (path.startswith("obliteratus/") and path.endswith(".py"))


def select_tests(
    paths: list[str],
    *,
    policy: dict[str, Any],
    risk_maps: list[dict[str, Any]],
    project_root: Path,
) -> list[str]:
    """Return stable existing CPU tests required by the changed paths."""

    selected = set(policy.get("always_tests", []))
    infrastructure_tests = set(policy.get("infrastructure_tests", []))
    infrastructure_patterns = policy.get("infrastructure_paths", [])
    excluded_prefixes = tuple(policy.get("excluded_test_prefixes", []))
    mappings = [_surface_tests(risk_map) for risk_map in risk_maps]

    for path in paths:
        if path.startswith("tests/") and path.endswith(".py"):
            selected.add(path)
        if any(fnmatch.fnmatchcase(path, pattern) for pattern in infrastructure_patterns):
            selected.update(infrastructure_tests)
        if _production_source(path):
            mapped = set().union(*(mapping.get(path, set()) for mapping in mappings))
            if not mapped:
                raise ValueError(
                    f"changed production source is unmapped in head and base risk maps: {path}",
                )
            selected.update(mapped)

    usable = [
        test
        for test in selected
        if isinstance(test, str)
        and test.startswith("tests/")
        and not test.startswith(excluded_prefixes)
        and (project_root / test).is_file()
    ]
    if not usable:
        raise ValueError("PR test policy selected no runnable CPU tests")
    return sorted(usable)


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-ref", required=True, help="exact base commit or ref")
    parser.add_argument("--policy", type=Path, default=DEFAULT_POLICY)
    parser.add_argument("--risk-map", type=Path, default=DEFAULT_RISK_MAP)
    return parser


def main() -> int:
    args = _parser().parse_args()
    try:
        policy = _load_object(args.policy, "PR test policy")
        head_risk_map = _load_object(args.risk_map, "test risk map")
        base_risk_map = _risk_map_at_ref(args.base_ref, args.risk_map)
        tests = select_tests(
            changed_paths(args.base_ref),
            policy=policy,
            risk_maps=[head_risk_map, *([base_risk_map] if base_risk_map else [])],
            project_root=Path.cwd(),
        )
    except ValueError as exc:
        print(f"PR test selection failed: {exc}")
        return 1
    for test in tests:
        print(test)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())