tests/test_play.py

"""Smoke tests for play.py.

Drives play.main() against the smallest CPU variant (neural_computer8_small,
1 KB memory) and the smallest ALU-only variant (neural_alu8). Each demo
prints a deterministic block on stdout; we capture and grep for the success
markers. Catches the kind of drift where play.py outpaces a refactor of the
gate naming convention (cmp8bit, mod5, etc.) without anyone running it.
"""
from __future__ import annotations

import io
import os
import sys
import unittest
from contextlib import redirect_stdout

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
for _p in (REPO_ROOT, os.path.join(REPO_ROOT, "tools")):  # play.py lives in tools/
    if _p not in sys.path:
        sys.path.insert(0, _p)

import play  # noqa: E402


def _run(model_path: str, skip_cpu: bool) -> str:
    argv_save = sys.argv
    sys.argv = ["play.py", "--model", model_path]
    if skip_cpu:
        sys.argv.append("--skip-cpu")
    buf = io.StringIO()
    try:
        with redirect_stdout(buf):
            rc = play.main()
    finally:
        sys.argv = argv_save
    if rc not in (0, None):
        raise AssertionError(f"play.main() returned {rc}")
    return buf.getvalue()


class TestPlay(unittest.TestCase):

    def test_alu_variant(self) -> None:
        path = os.path.join(REPO_ROOT, "variants", "neural_alu8.safetensors")
        out = _run(path, skip_cpu=True)
        # Booleans
        for line in ["and      00->0 01->0 10->0 11->1",
                     "or       00->0 01->1 10->1 11->1",
                     "xor      00->0 01->1 10->1 11->0"]:
            self.assertIn(line, out, f"missing boolean line: {line!r}")
        # ALU ADD/SUB
        self.assertIn("127 + 128 = 255 (carry=0)  expected 255  [OK]", out)
        self.assertIn("127 - 128 = 255 (no_borrow=0)  expected 255  [OK]", out)
        # CMP via cmp8bit cascade
        self.assertIn("50 vs  30 -> GT=1 LT=0 EQ=0", out)
        self.assertIn("77 vs  77 -> GT=0 LT=0 EQ=1", out)
        self.assertIn("128 vs 127 -> GT=1 LT=0 EQ=0", out)
        # MUL
        self.assertIn("12 *  11 = 132  expected 132  [OK]", out)
        # mod-5 (52 multiples of 5 in [0, 255])
        self.assertIn("52 hits", out)
        self.assertNotIn("FAIL", out)

    def test_cpu_variant(self) -> None:
        path = os.path.join(REPO_ROOT, "variants", "neural_computer8_small.safetensors")
        out = _run(path, skip_cpu=False)
        self.assertIn("M[0x0083] = 15  (expected 15)  [OK]", out)
        self.assertIn("R0=15", out)
        self.assertNotIn("FAIL", out)


if __name__ == "__main__":
    unittest.main()