1 files changed, 250 insertions, 0 deletions
diff --git a/tests/run.py b/tests/run.py
new file mode 100755
index 0000000..ca0a319
--- /dev/null
+++ b/tests/run.py
@@ -0,0 +1,250 @@
+#!/usr/bin/env python3
+"""The front door for the test suite.
+
+There are two kinds of test file here and the difference matters a lot:
+
+  * **pure** — stdlib (sometimes + textual), no IDA, no worker, no binary.
+    Runs anywhere in about a second. ``tests/run.py --fast`` is exactly this
+    set, and it's what you run between edits.
+  * **IDA** — spawns a real idalib worker on a real target and drives the
+    Textual pilot against it. Minutes, needs a licensed IDA, and gets CPU
+    starved on a loaded box (a SIGKILLed worker looks like a hang, not a bug —
+    check ``uptime`` before believing a regression).
+
+Rather than keep that classification in a table here, where it would rot the
+first time someone adds a test, **each test file declares it**::
+
+    NEEDS_IDA = True   # or False
+
+``run.py`` reads that marker with :mod:`ast` (it never imports the file — these
+modules run their suite at import time). A test file with no marker is a hard
+error, so a new test can't quietly join the fast set and start needing IDA.
+
+Usage::
+
+    python3 tests/run.py              # everything
+    python3 tests/run.py --fast       # only the no-IDA files (seconds)
+    python3 tests/run.py --list       # what would run, and why
+    python3 tests/run.py trace graph  # only files matching these substrings
+    python3 tests/run.py --fast -x    # stop at the first failing file
+
+Exit code is 0 only if every file selected ran and passed.
+"""
+from __future__ import annotations
+
+import argparse
+import ast
+import os
+import re
+import subprocess
+import sys
+import time
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+TESTS = os.path.join(ROOT, "tests")
+
+#: The IDA-capable interpreter. The pilot tests need textual AND idapro in one
+#: python; the worker python is auto-detected separately by WorkerClient.
+DEFAULT_PY = os.path.expanduser("~/ida-venv/bin/python")
+
+#: Both shapes the suites print: "N passed, M failed" and "N checks, M failed".
+_TALLY = re.compile(r"^\s*(\d+)\s+(?:passed|checks),\s*(\d+)\s+failed\s*$", re.M)
+_SKIP = re.compile(r"^\s*skip(?:ped|ping)?\b[:.]", re.M | re.I)
+
+
+class Marker(Exception):
+    """A test file didn't declare NEEDS_IDA."""
+
+
+def needs_ida(path: str) -> bool:
+    """Read the module-level ``NEEDS_IDA`` without importing the module.
+
+    Importing is not an option: every one of these files runs its whole suite
+    under ``if __name__ == "__main__"`` but also defines its checks at module
+    scope, and some spawn workers on import of a helper. AST it is.
+    """
+    with open(path, encoding="utf-8") as fh:
+        tree = ast.parse(fh.read(), filename=path)
+    for node in tree.body:
+        if not isinstance(node, ast.Assign):
+            continue
+        for target in node.targets:
+            if isinstance(target, ast.Name) and target.id == "NEEDS_IDA":
+                value = ast.literal_eval(node.value)
+                if not isinstance(value, bool):
+                    raise Marker(f"{os.path.basename(path)}: "
+                                 f"NEEDS_IDA must be a bool, got {value!r}")
+                return value
+    raise Marker(
+        f"{os.path.basename(path)}: no NEEDS_IDA marker.\n"
+        f"  Add `NEEDS_IDA = True` (spawns a worker / drives the pilot) or\n"
+        f"  `NEEDS_IDA = False` (pure: stdlib, no IDA, runs anywhere) at module\n"
+        f"  scope, so tests/run.py --fast knows whether it can run you.")
+
+
+def discover() -> list[tuple[str, bool]]:
+    """Every ``tests/test_*.py``, with its IDA requirement. Sorted: fast first,
+    so ``run.py`` fails on a cheap file before burning minutes on a pilot."""
+    out = []
+    problems = []
+    for name in sorted(os.listdir(TESTS)):
+        if not (name.startswith("test_") and name.endswith(".py")):
+            continue
+        path = os.path.join(TESTS, name)
+        try:
+            out.append((path, needs_ida(path)))
+        except Marker as exc:
+            problems.append(str(exc))
+    if problems:
+        raise Marker("\n".join(problems))
+    out.sort(key=lambda p: (p[1], p[0]))
+    return out
+
+
+def tally(output: str) -> tuple[int, int] | None:
+    """The last ``N passed, M failed`` line a suite printed.
+
+    The last one, not the first: test_graph prints a per-section tally on the
+    way through and the total at the end.
+    """
+    found = _TALLY.findall(output)
+    if not found:
+        return None
+    passed, failed = found[-1]
+    return int(passed), int(failed)
+
+
+def run_one(path: str, python: str, extra: list[str], echo: bool) -> dict:
+    """Run one test file as a subprocess and summarise it."""
+    name = os.path.basename(path)[len("test_"):-len(".py")]
+    started = time.time()
+    proc = subprocess.run([python, path, *extra], cwd=ROOT,
+                          capture_output=not echo, text=True)
+    took = time.time() - started
+    out = "" if echo else (proc.stdout or "") + (proc.stderr or "")
+    counts = tally(out)
+    skipped = bool(_SKIP.search(out)) and (counts is None or counts == (0, 0))
+    return {
+        "name": name, "path": path, "code": proc.returncode, "took": took,
+        "passed": counts[0] if counts else 0,
+        "failed": counts[1] if counts else 0,
+        "counted": counts is not None,
+        "skipped": skipped, "output": out,
+    }
+
+
+def main(argv: list[str]) -> int:
+    ap = argparse.ArgumentParser(
+        prog="tests/run.py", description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("only", nargs="*", metavar="SUBSTR",
+                    help="only run test files whose name contains one of these")
+    ap.add_argument("--fast", action="store_true",
+                    help="skip every file that needs IDA (seconds, runs anywhere)")
+    ap.add_argument("--ida-only", action="store_true",
+                    help="only the files that need IDA")
+    ap.add_argument("--list", action="store_true",
+                    help="show what would run, and whether it needs IDA")
+    ap.add_argument("-x", "--exitfirst", action="store_true",
+                    help="stop after the first failing file")
+    ap.add_argument("-v", "--verbose", action="store_true",
+                    help="stream each suite's output instead of capturing it")
+    ap.add_argument("--python", default=os.environ.get("IDATUI_PYTHON", DEFAULT_PY),
+                    help=f"interpreter for the IDA suites (default {DEFAULT_PY})")
+    args, extra = ap.parse_known_args(argv)
+
+    try:
+        files = discover()
+    except Marker as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 2
+
+    selected = []
+    for path, ida in files:
+        if args.fast and ida:
+            continue
+        if args.ida_only and not ida:
+            continue
+        if args.only and not any(s in os.path.basename(path) for s in args.only):
+            continue
+        selected.append((path, ida))
+
+    if not selected:
+        print("nothing selected", file=sys.stderr)
+        return 2
+
+    if args.list:
+        for path, ida in selected:
+            print(f"{'ida ' if ida else 'pure'}  {os.path.basename(path)}")
+        return 0
+
+    # A pure file runs under whatever python invoked us (it needs nothing);
+    # an IDA file needs the interpreter that has textual + idapro.
+    pure_py = sys.executable
+    if any(ida for _, ida in selected) and not os.path.exists(args.python):
+        print(f"error: {args.python} not found — the IDA suites need an "
+              f"interpreter with textual + idapro.\n"
+              f"       Pass --python, set $IDATUI_PYTHON, or use --fast.",
+              file=sys.stderr)
+        return 2
+
+    results = []
+    started = time.time()
+    for path, ida in selected:
+        label = os.path.basename(path)
+        print(f"\033[1m>> {label}\033[0m{' (ida)' if ida else ''}", flush=True)
+        res = run_one(path, args.python if ida else pure_py, extra, args.verbose)
+        results.append(res)
+        bad = res["code"] != 0 or res["failed"]
+        if bad and not args.verbose:
+            print(res["output"].rstrip())
+        elif res["skipped"]:
+            print("   skipped")
+        else:
+            print(f"   {res['passed']} passed  ({res['took']:.1f}s)")
+        if bad and args.exitfirst:
+            print("\nstopping at the first failure (-x)", file=sys.stderr)
+            break
+
+    total = time.time() - started
+    print("\n" + "=" * 62)
+    width = max(len(r["name"]) for r in results)
+    passed = failed = 0
+    for r in results:
+        passed += r["passed"]
+        failed += r["failed"]
+        if r["skipped"]:
+            state = "\033[33mSKIP\033[0m"
+        elif r["code"] != 0 or r["failed"]:
+            state = "\033[31mFAIL\033[0m"
+        elif not r["counted"]:
+            state = "\033[33m ?  \033[0m"   # exit 0 but printed no tally
+        else:
+            state = "\033[32m ok \033[0m"
+        detail = f"{r['passed']:4d} passed"
+        if r["failed"]:
+            detail += f", \033[31m{r['failed']} failed\033[0m"
+        if r["code"] != 0 and not r["failed"]:
+            detail += f", \033[31mexit {r['code']}\033[0m"
+        print(f" {state}  {r['name']:<{width}}  {detail}   {r['took']:6.1f}s")
+    print("=" * 62)
+
+    ran, total_files = len(results), len(selected)
+    skipped = sum(1 for r in results if r["skipped"])
+    hurt = [r["name"] for r in results if r["code"] != 0 or r["failed"]]
+    summary = f"{passed} passed"
+    if failed:
+        summary += f", {failed} failed"
+    if skipped:
+        summary += f", {skipped} file(s) skipped"
+    if ran != total_files:
+        summary += f", {total_files - ran} file(s) not reached"
+    print(f"{summary}   [{ran} file(s), {total:.1f}s]")
+    if hurt:
+        print(f"\033[31mfailing files: {', '.join(hurt)}\033[0m")
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main(sys.argv[1:]))