stack/dev/scripts/nb_integration.py

"""Headless notebook integration test.

Runs `marimo export session` over a set of notebooks (executing them and
writing JSON session snapshots), parses each snapshot for cell errors, and
emits a consolidated JSON report. Exits 1 if any notebook outside the
expected-failures allowlist errored.

Why: the pytest suite only exercises src/ — nothing ran the notebooks
themselves, so kernel-level breakage (dead imports, schema drift, service
API changes) only surfaced when a human opened the notebook.

Config: infra/marimo/nb-tests.toml
    [ci_smoke]
    notebooks = ["sample.py"]          # data-independent, runs in CI
    [expected_failures]
    "acodb_explorer.py" = "reason"     # reported but non-fatal

Usage:
    uv run python dev/scripts/nb_integration.py --set ci-smoke
    uv run python dev/scripts/nb_integration.py --set all --file-issues
"""

from __future__ import annotations

import argparse
import html
import json
import re
import subprocess
import sys
import time
import tomllib
from pathlib import Path

# In the repo this file sits at dev/scripts/, two levels below the root.
# Deployed contexts (nb-watcher mounts it at /scripts, the nightly run
# docker-cps it to /tmp/nbtest) are shallower — fall back to "/" there;
# those contexts always override paths via --nb-dir / env anyway.
_parents = Path(__file__).resolve().parents
ROOT = _parents[2] if len(_parents) > 2 else Path("/")
NB_DIR = ROOT / "notebooks"
# Config lives in infra/marimo/ (notebooks/ holds only .py content — see
# tests/test_notebook_layout.py). Overridable via --config for contexts
# where the repo layout doesn't exist (nightly docker exec).
CONFIG_PATH = ROOT / "infra" / "marimo" / "nb-tests.toml"
REPORT_PATH = ROOT / "data" / "nb-integration-report.json"
TIMEOUT_S = 600


# ── snapshot parsing (pure) ──────────────────────────────────────


def _strip_html(text: str) -> str:
    return html.unescape(re.sub(r"<[^>]+>", "", text))


def parse_snapshot(path: Path) -> list[dict]:
    """Extract cell errors from a __marimo__/session/<nb>.py.json snapshot."""
    if not path.exists():
        return [
            {
                "cell": "",
                "ename": "snapshot-missing",
                "evalue": f"no session snapshot at {path.name}",
                "detail": "",
            }
        ]
    try:
        snap = json.loads(path.read_text())
    except (OSError, json.JSONDecodeError) as e:
        return [
            {
                "cell": "",
                "ename": "snapshot-unreadable",
                "evalue": str(e)[:200],
                "detail": "",
            }
        ]

    errors: list[dict] = []
    for cell in snap.get("cells", []):
        for out in cell.get("outputs", []):
            if out.get("type") != "error":
                continue
            # Cascade noise: marimo marks every descendant of a failed or
            # stopped cell with its own error output. The root cause is
            # always present as a non-ancestor error in the same snapshot,
            # so these only inflate the report (and the issue tracker).
            if out.get("ename") == "ancestor-stopped" or out.get(
                "evalue", ""
            ).startswith("An ancestor raised an exception"):
                continue
            detail = ""
            for con in cell.get("console", []):
                if con.get("name") == "stderr":
                    detail += _strip_html(con.get("text", "")) + "\n"
            errors.append(
                {
                    "cell": cell.get("id", ""),
                    "ename": out.get("ename", "error"),
                    "evalue": out.get("evalue", ""),
                    "detail": detail.strip(),
                }
            )
    return errors


def classify(notebook: str, errors: list[dict], expected: dict[str, str]) -> str:
    if not errors:
        return "pass"
    return "expected-fail" if notebook in expected else "fail"


def exit_code(results: dict[str, dict]) -> int:
    return 1 if any(r["status"] == "fail" for r in results.values()) else 0


# ── execution ────────────────────────────────────────────────────


def load_config(path: Path) -> dict:
    if not path.exists():
        return {"ci_smoke": {"notebooks": []}, "expected_failures": {}}
    return tomllib.loads(path.read_text())


def select_notebooks(which: str, cfg: dict) -> list[str]:
    if which == "ci-smoke":
        return list(cfg.get("ci_smoke", {}).get("notebooks", []))
    return sorted(
        p.name for p in NB_DIR.glob("*.py") if not p.name.startswith((".", "__"))
    )


def run_notebook(name: str) -> tuple[int, str]:
    """Execute one notebook via marimo export session; snapshot lands in
    notebooks/__marimo__/session/<name>.json regardless of exit code."""
    try:
        proc = subprocess.run(
            [
                "marimo",
                "export",
                "session",
                name,
                "--no-sandbox",
                "--force-overwrite",
                "--continue-on-error",
            ],
            cwd=NB_DIR,
            capture_output=True,
            text=True,
            timeout=TIMEOUT_S,
        )
        return proc.returncode, (proc.stderr or "")[-1500:]
    except subprocess.TimeoutExpired:
        return -1, f"timed out after {TIMEOUT_S}s"
    except FileNotFoundError:
        return -2, "marimo CLI not found on PATH"


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--set", dest="which", choices=["ci-smoke", "all"], default="ci-smoke"
    )
    parser.add_argument(
        "--file-issues",
        action="store_true",
        help="report failures via nb_issue_filer (dedup) "
        "and sweep resolved ones, instead of only exiting 1",
    )
    parser.add_argument("--report", default=str(REPORT_PATH))
    parser.add_argument("--nb-dir", help="notebooks directory (default: repo)")
    parser.add_argument("--config", default=str(CONFIG_PATH), help="nb-tests.toml path")
    args = parser.parse_args()

    if args.nb_dir:
        global NB_DIR
        NB_DIR = Path(args.nb_dir).resolve()

    cfg = load_config(Path(args.config))
    expected = {k: str(v) for k, v in cfg.get("expected_failures", {}).items()}
    notebooks = select_notebooks(args.which, cfg)
    if not notebooks:
        print(f"no notebooks selected for --set {args.which}")
        return 0

    results: dict[str, dict] = {}
    for name in notebooks:
        print(f"→ {name}", flush=True)
        rc, stderr_tail = run_notebook(name)
        if rc == -1:
            errors = [
                {"cell": "", "ename": "timeout", "evalue": stderr_tail, "detail": ""}
            ]
        elif rc == -2:
            print(stderr_tail, file=sys.stderr)
            return 2
        else:
            errors = parse_snapshot(NB_DIR / "__marimo__" / "session" / f"{name}.json")
        status = classify(name, errors, expected)
        results[name] = {"status": status, "errors": errors}
        tag = {"pass": "ok", "fail": "FAIL", "expected-fail": "expected-fail"}[status]
        print(
            f"  {tag}"
            + (
                f" — {errors[0]['ename']}: {errors[0]['evalue'][:120]}"
                if errors
                else ""
            )
        )

    report = {
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "set": args.which,
        "marimo_set_size": len(notebooks),
        "results": results,
    }
    report_path = Path(args.report)
    report_path.parent.mkdir(parents=True, exist_ok=True)
    report_path.write_text(json.dumps(report, indent=2) + "\n")
    print(f"report → {report_path}")

    counts: dict[str, int] = {}
    for r in results.values():
        counts[r["status"]] = counts.get(r["status"], 0) + 1
    print("summary:", ", ".join(f"{k}={v}" for k, v in sorted(counts.items())))

    if args.file_issues:
        sys.path.insert(0, str(Path(__file__).parent))
        import nb_issue_filer  # noqa: PLC0415 — sibling script, lazy by design

        findings = [
            {"notebook": nb, **e}
            for nb, r in results.items()
            if r["status"] in ("fail", "expected-fail")
            for e in r["errors"]
        ]
        source = f"nb-integration/{args.which}"
        if findings:
            nb_issue_filer.cmd_report(findings, source)
        else:
            print("no findings — nothing to file")
        # No sweep here: closing is the nb-watcher's job (24h staleness).
        # Sweeping with only this run's signatures would close issues for
        # errors this run can't reproduce (e.g. server-log-only errors the
        # watcher filed).
        return 0  # issue filing is the failure channel for scheduled runs

    return exit_code(results)


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