212 lines
9.8 KiB
Python
212 lines
9.8 KiB
Python
import contextlib
|
|
import io
|
|
import json
|
|
import os
|
|
import shutil
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
SKILL = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(SKILL))
|
|
|
|
from grm import cli # noqa: E402
|
|
|
|
FAKE = str(SKILL / "tests" / "fake_claude.py")
|
|
MODEL = "claude-opus-5-5"
|
|
|
|
|
|
class CliTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = Path(tempfile.mkdtemp())
|
|
self.addCleanup(shutil.rmtree, self.tmp, True)
|
|
repo = self.tmp / "repo"
|
|
repo.mkdir()
|
|
for a in (["init", "-q"], ["-c", "user.email=t@t", "-c", "user.name=t", "commit", "-q",
|
|
"--allow-empty", "-m", "x"]):
|
|
subprocess.run(["git", *a], cwd=repo, check=True, capture_output=True)
|
|
self.run_md = str(self.tmp / "run.md")
|
|
self.state = self.tmp / "run-gremlins"
|
|
self.log = self.tmp / "fake.jsonl"
|
|
self.brief = self.tmp / "brief.md"
|
|
self.brief.write_text("do the thing\n", encoding="utf-8")
|
|
env = mock.patch.dict(os.environ, {"FAKE_CLAUDE_LOG": str(self.log), "FAKE_CLAUDE_SCENARIO": "ok"})
|
|
env.start()
|
|
self.addCleanup(env.stop)
|
|
code, out = self.g("init", "--run", self.run_md, "--model", MODEL, "--cwd", str(repo), "--binary", sys.executable)
|
|
self.assertEqual(code, 0, out)
|
|
run = json.loads((self.state / "run.json").read_text(encoding="utf-8"))
|
|
run["argv"] = [sys.executable, FAKE] + run["argv"][1:]
|
|
(self.state / "run.json").write_text(json.dumps(run), encoding="utf-8")
|
|
self.stored = run["argv"]
|
|
|
|
def g(self, *args, scenario=None):
|
|
if scenario:
|
|
os.environ["FAKE_CLAUDE_SCENARIO"] = scenario
|
|
buf = io.StringIO()
|
|
with contextlib.redirect_stdout(buf):
|
|
code = cli.main(list(args))
|
|
os.environ["FAKE_CLAUDE_SCENARIO"] = "ok"
|
|
return code, json.loads(buf.getvalue().splitlines()[-1])
|
|
|
|
def train(self, scenario=None):
|
|
return self.g("train", "--run", self.run_md, "--base", "B1", "--brief", str(self.brief),
|
|
scenario=scenario)
|
|
|
|
def spawn(self, job, base="B1", scenario=None):
|
|
return self.g("spawn", "--run", self.run_md, "--base", base, "--job", job,
|
|
"--brief", str(self.brief), scenario=scenario)
|
|
|
|
def calls(self):
|
|
if not self.log.exists():
|
|
return []
|
|
return [json.loads(ln) for ln in self.log.read_text(encoding="utf-8").splitlines()]
|
|
|
|
def ledger(self):
|
|
return [json.loads(ln) for ln in (self.state / "ledger.jsonl").read_text(encoding="utf-8").splitlines()]
|
|
|
|
def base_rec(self):
|
|
p = self.state / "bases" / "B1.json"
|
|
return json.loads(p.read_text(encoding="utf-8")) if p.exists() else None
|
|
|
|
def test_train_spawn_happy(self):
|
|
"""B1, B2, B5: base keeps the --session-id; fork argv = stored argv + resume/fork; ledger."""
|
|
code, out = self.train()
|
|
self.assertEqual((code, out["status"]), (0, "ok"))
|
|
sid = self.calls()[0]["argv"][-1]
|
|
# the fake logs its sys.argv: the stored argv without the interpreter
|
|
self.assertEqual(self.calls()[0]["argv"], self.stored[1:] + ["--session-id", sid])
|
|
self.assertEqual(self.base_rec()["session_id"], sid)
|
|
code, out = self.spawn("1.1")
|
|
self.assertEqual((code, out["status"], out["cache"]), (0, "ok", "hit"))
|
|
argv = self.calls()[1]["argv"]
|
|
n = len(self.stored) - 1
|
|
self.assertEqual(argv[:n], self.stored[1:])
|
|
self.assertEqual(argv[n:n + 3], ["--resume", sid, "--fork-session"])
|
|
line = self.ledger()[-1]
|
|
self.assertEqual((line["model"], line["job"], line["cache"]), (MODEL, "1.1", "hit"))
|
|
for k in ("input", "output", "cache_creation", "cache_read"):
|
|
self.assertIsInstance(line["tokens"][k], int)
|
|
self.assertAlmostEqual(line["cost_delta_usd"],
|
|
line["cost_cumulative_usd"] - self.base_rec()["cost_cumulative_usd"], places=6)
|
|
self.assertAlmostEqual(line["cost_delta_usd"], 0.02, places=6)
|
|
|
|
def test_missing_binary(self):
|
|
"""B1: a missing binary is refused before launch; no job file."""
|
|
run = json.loads((self.state / "run.json").read_text(encoding="utf-8"))
|
|
run["argv"][0] = str(self.tmp / "nope.exe")
|
|
(self.state / "run.json").write_text(json.dumps(run), encoding="utf-8")
|
|
code, out = self.train()
|
|
self.assertEqual(code, 2)
|
|
self.assertIn("binary", out["error"])
|
|
self.assertEqual(list((self.state / "jobs").iterdir()), [])
|
|
self.assertEqual(self.calls(), [])
|
|
|
|
def test_train_failures(self):
|
|
"""B1: no result + exit 1, and an API error result: status error, no base file."""
|
|
for scenario, reason in (("noresult", "exit_code"), ("api_error", "usage_limit")):
|
|
code, out = self.train(scenario)
|
|
self.assertEqual((code, out["status"], out["reason"]), (1, "error", reason), scenario)
|
|
self.assertIsNone(self.base_rec())
|
|
|
|
def test_unknown_base(self):
|
|
"""B2: an unknown base is refused with no launch."""
|
|
code, out = self.spawn("1.1", base="NOPE")
|
|
self.assertEqual((code, out["error"]), (2, "unknown base"))
|
|
self.assertEqual(self.calls(), [])
|
|
self.assertFalse((self.state / "jobs" / "1.1.json").exists())
|
|
|
|
def test_cache_miss_is_reported(self):
|
|
"""B2: first-usage cache_read 0 is `cache: miss`, not an error."""
|
|
self.train()
|
|
code, out = self.spawn("1.1", scenario="miss")
|
|
self.assertEqual((code, out["status"], out["cache"]), (0, "ok", "miss"))
|
|
|
|
def test_parallel_spawns(self):
|
|
"""B3: two spawns in parallel threads give two job files and two whole ledger lines."""
|
|
self.train()
|
|
ts = [threading.Thread(target=cli.main, args=(["spawn", "--run", self.run_md, "--base", "B1",
|
|
"--job", j, "--brief", str(self.brief)],))
|
|
for j in ("p1", "p2")]
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
for t in ts:
|
|
t.start()
|
|
for t in ts:
|
|
t.join(60)
|
|
for j in ("p1", "p2"):
|
|
rec = json.loads((self.state / "jobs" / (j + ".json")).read_text(encoding="utf-8"))
|
|
self.assertEqual(rec["status"], "ok")
|
|
lines = self.ledger()
|
|
self.assertEqual(len(lines), 3)
|
|
self.assertEqual(sorted(ln["job"] for ln in lines[1:]), ["p1", "p2"])
|
|
|
|
def test_stop_signal(self):
|
|
"""B4: a stop signal to the helper ends the child, marks the job stopped, keeps raw."""
|
|
self.train()
|
|
env = dict(os.environ, FAKE_CLAUDE_SCENARIO="sleep")
|
|
kw = {"creationflags": subprocess.CREATE_NEW_PROCESS_GROUP} if os.name == "nt" else {}
|
|
h = subprocess.Popen([sys.executable, str(SKILL / "gremlins.py"), "spawn", "--run", self.run_md,
|
|
"--base", "B1", "--job", "s1", "--brief", str(self.brief)],
|
|
env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, **kw)
|
|
jpath, raw = self.state / "jobs" / "s1.json", self.state / "raw" / "s1.jsonl"
|
|
deadline = time.time() + 30
|
|
while time.time() < deadline and not (
|
|
jpath.exists() and json.loads(jpath.read_text(encoding="utf-8")).get("child_pid")
|
|
and raw.exists() and b'"assistant"' in raw.read_bytes()):
|
|
time.sleep(0.2)
|
|
pid = json.loads(jpath.read_text(encoding="utf-8"))["child_pid"]
|
|
os.kill(h.pid, signal.CTRL_BREAK_EVENT if os.name == "nt" else signal.SIGTERM)
|
|
out, _ = h.communicate(timeout=30)
|
|
self.assertEqual(h.returncode, 1)
|
|
self.assertEqual(json.loads(out.decode("utf-8").splitlines()[-1])["status"], "stopped")
|
|
self.assertEqual(json.loads(jpath.read_text(encoding="utf-8"))["status"], "stopped")
|
|
self.assertIn(b'"init"', raw.read_bytes())
|
|
if os.name == "nt":
|
|
r = subprocess.run(["tasklist", "/FI", "PID eq %d" % pid, "/NH"], capture_output=True)
|
|
self.assertNotIn(str(pid).encode(), r.stdout)
|
|
else:
|
|
with self.assertRaises(ProcessLookupError):
|
|
os.kill(pid, 0)
|
|
|
|
def test_usage_models_and_cost(self):
|
|
"""B5: two models both in models_used; no total_cost_usd gives a null cost delta."""
|
|
self.train("two_models")
|
|
line = self.ledger()[-1]
|
|
self.assertEqual(sorted(line["models_used"]), ["claude-haiku-5-5", MODEL])
|
|
self.assertEqual(line["model"], MODEL)
|
|
(self.state / "bases" / "B1.json").unlink()
|
|
self.train("no_cost")
|
|
self.assertIsNone(self.ledger()[-1]["cost_delta_usd"])
|
|
|
|
def test_status_cold_without_hit(self):
|
|
"""B6: a base with no hit in the ledger is cold (and warm right after training)."""
|
|
self.train()
|
|
code, out = self.g("status", "--run", self.run_md)
|
|
self.assertEqual((code, out["bases"]["B1"]["base_state"]), (0, "warm"))
|
|
(self.state / "ledger.jsonl").write_bytes(b"")
|
|
code, out = self.g("status", "--run", self.run_md)
|
|
self.assertEqual(out["bases"]["B1"]["base_state"], "cold")
|
|
|
|
def test_rebase(self):
|
|
"""B7: rebase resumes without a fork; the base takes the result's session id."""
|
|
self.train()
|
|
old = self.base_rec()["session_id"]
|
|
code, out = self.g("rebase", "--run", self.run_md, "--base", "B1", "--brief", str(self.brief))
|
|
self.assertEqual((code, out["status"]), (0, "ok"))
|
|
argv = self.calls()[-1]["argv"]
|
|
self.assertEqual(argv[-2:], ["--resume", old])
|
|
self.assertNotIn("--fork-session", argv)
|
|
new = self.base_rec()["session_id"]
|
|
self.assertTrue(new.startswith("resumed-"))
|
|
self.assertEqual(new, out["session_id"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|