eng-lab/labs/el01_hover/check.py

"""Run the EL01 lab tests one lab step at a time.

From the eng-lab folder, with the venv's Python:
    python labs/el01_hover/check.py          every step (1-9)
    python labs/el01_hover/check.py 3        steps 1-3 only; use this while you work

A step "passes" only when all of its tests pass. When one fails, read the test's
name and body first (tests/test_*.py shows the exact inputs and the arithmetic
it expected), re-read that lab step in EL01, and climb the step's hint ladder
one rung at a time.
"""
import argparse
import sys
import unittest
from pathlib import Path

HERE = Path(__file__).resolve().parent
STEP_NAMES = {
    1: "weight", 2: "Newton's second law", 3: "thrust-to-weight", 4: "energy and flight time",
    5: "mass budget", 6: "the arm (CAD)", 7: "requirements", 8: "phone accelerometer", 9: "tinker log",
}
MODULE_STEPS = {
    "test_hover": {1, 2, 3}, "test_energy": {4}, "test_massbudget": {5}, "test_arm": {6},
    "test_requirements": {7}, "test_kinematics": {8}, "test_tinker": {9},
}


def step_of(test):
    """Step number from the test's class name (Step3ThrustToWeight -> 3)."""
    test = getattr(test, "test_case", test)            # unwrap a failing subTest
    name = type(test).__name__
    if name == "_ErrorHolder":                         # setUpClass failed: "setUpClass (test_arm.Step6Arm)"
        name = test.description.rsplit(".", 1)[-1].rstrip(")")
    if name.startswith("Step") and name[4:5].isdigit():
        return int(name[4])
    return None


def flatten(suite):
    for item in suite:
        if isinstance(item, unittest.TestSuite):
            yield from flatten(item)
        else:
            yield item


def phone_status():
    """Report (not test) whether the physical test's prediction was written before its data."""
    root = HERE.parents[1]
    pred = HERE / "predictions.toml"
    data_dir = root / "data" / "tests" / "el01_phone"
    data = sorted(data_dir.glob("*.csv")) if data_dir.exists() else []
    if not data:
        return "no phone data yet (data/tests/el01_phone/*.csv)"
    if not pred.exists():
        return "phone data found, but no predictions.toml: write your prediction BEFORE measuring"
    first = min(p.stat().st_mtime for p in data)
    return ("prediction written before the data: yes" if pred.stat().st_mtime < first
            else "prediction file is newer than the data: was it edited after measuring?")


def main():
    for stream in (sys.stdout, sys.stderr):
        try:
            stream.reconfigure(encoding="utf-8", errors="replace")
        except (AttributeError, ValueError):
            pass
    parser = argparse.ArgumentParser(description="Run the EL01 lab tests by step.")
    parser.add_argument("step", nargs="?", type=int, default=9, help="run steps 1..STEP (default 9)")
    args = parser.parse_args()
    if not 1 <= args.step <= 9:
        parser.error("step must be between 1 and 9")

    wanted = set(range(1, args.step + 1))
    sys.path.insert(0, str(HERE / "tests"))
    loader = unittest.TestLoader()
    suite = unittest.TestSuite()
    for module, steps in MODULE_STEPS.items():
        if not steps & wanted:
            continue
        for test in flatten(loader.loadTestsFromName(module)):
            step = step_of(test)
            if step is None or step in wanted:          # None: an import error, always show it
                suite.addTest(test)

    print(f"EL01 lab: running steps 1-{args.step}\n")
    result = unittest.TextTestRunner(verbosity=2, stream=sys.stdout).run(suite)

    broken = {}
    for test, trace in result.failures + result.errors:
        step = step_of(test) or "import"
        broken.setdefault(step, []).append((test, "NotImplementedError" in trace))

    print("\nSummary")
    for step in sorted(wanted):
        if step in broken:
            tests = broken[step]
            if all(not_written for _, not_written in tests):
                status = "not written yet"
            else:
                first = next(t for t, nw in tests if not nw)
                first = getattr(first, "test_case", first)
                label = (first.description if type(first).__name__ == "_ErrorHolder"
                         else f"{type(first).__name__}.{first._testMethodName}")
                status = f"FAIL  first failing test: {label}"
        else:
            status = "pass"
        print(f"  step {step} ({STEP_NAMES[step]}): {status}")
    if "import" in broken:
        print("  A test file could not import your code. Read the traceback above first.")
    if 8 in wanted:
        print(f"  phone test: {phone_status()}")
    print(f"\n{result.testsRun} tests run, {len(result.failures)} failures, "
          f"{len(result.errors)} errors, {len(result.skipped)} skipped")
    sys.exit(0 if result.wasSuccessful() else 1)


if __name__ == "__main__":
    main()