diff options
Diffstat (limited to 'tests/test_batch.py')
| -rw-r--r-- | tests/test_batch.py | 534 |
1 files changed, 534 insertions, 0 deletions
diff --git a/tests/test_batch.py b/tests/test_batch.py new file mode 100644 index 0000000..2f0e834 --- /dev/null +++ b/tests/test_batch.py @@ -0,0 +1,534 @@ +import io +import shlex +import subprocess +import sys +import tempfile +import threading +import unittest +from contextlib import redirect_stderr, redirect_stdout +from pathlib import Path + +HERE = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(HERE)) +sys.path.insert(0, str(HERE / "tests")) + +import wf_res # noqa: E402 +from test_res_io import Box, hist_rec, seed_history # noqa: E402 + +OUT = "out/wf-batch-2026-10-01-1400.md" + + +class BatchBase(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.box = Box(Path(self._tmp.name)) + bin_ = self.box.tmp / "bin" + bin_.mkdir() + self.claude = bin_ / "claude" + self.claude.write_text("#!/bin/sh\n") + self.claude.chmod(0o755) + self.box.caller["PATH"] = f"{bin_}:/usr/bin" + + def tearDown(self): + self._tmp.cleanup() + + def batch(self, *argv): + out, err = io.StringIO(), io.StringIO() + with redirect_stdout(out), redirect_stderr(err): + code = wf_res.batch_main(list(argv), self.box.env) + return code, out.getvalue(), err.getvalue() + + def unit_argv(self): + return self.box.fake.ran("systemd-run")[0] + + +class Start(BatchBase): + def test_starts_claude_print_unit(self): + code, out, err = self.batch("2", "--lanes", "fast,slow") + self.assertEqual((code, err), (0, "")) + log = self.box.env.state / "logs" / "r-1.log" + self.assertEqual(out, f"r-1 started; log {log}; ETA ~20:00 (estimate: not killed when over)\n" + f"summary: {OUT} · status: wf batch --status\n") + argv = self.unit_argv() + i = argv.index(str(self.claude)) + self.assertEqual(argv[i - 1], "sh") + self.assertEqual(argv[i + 1:-1], ["-p", "--model", "opus", "--permission-mode", "auto", + "--permission-prompts", "none"]) + self.assertIn("--setenv=CLAUDE_CODE_PRINT_BG_WAIT_CEILING_MS=21600000", argv) + self.assertIn(f"MemoryMax={4 * 1024 ** 3}", argv) + e = self.box.ledger().get("r-1") + self.assertEqual((e.title, e.est_min, e.project), ("wf-batch", 360, "proj")) + + def test_prompt_from_template(self): + self.batch("2", "--lanes", "fast,slow") + prompt = self.unit_argv()[-1] + self.assertIn("for up to 2 tasks (lanes: fast, slow)", prompt) + self.assertIn(f"append one line per task to {OUT}", prompt) + self.assertIn("/projects/public/workflow/docs/orchestrator.md", prompt.replace(str(HERE), "/projects/public/workflow")) + self.assertIn("foreground", prompt) + self.assertNotIn("{", prompt) + + def test_dry_run_writes_no_summary(self): + self.batch("2", "--dry-run") + self.assertFalse((self.box.env.root() / OUT).exists()) + + def test_existing_summary_not_overwritten(self): + f = self.box.env.root() / OUT + f.parent.mkdir(exist_ok=True) + f.write_text("live\n") + self.batch("2") + self.assertTrue(f.read_text().startswith("live\n# wf-batch")) + + def test_size_lane(self): + code, out, err = self.batch("2", "--lanes", "fast", "--dry-run") + self.assertEqual((code, err), (0, "")) + self.assertIn("lanes: fast", out) + + def test_default_lanes(self): + self.batch("3") + self.assertIn("for up to 3 tasks (lanes: every lane with ready tasks, see wf lanes)", self.unit_argv()[-1]) + + def test_for_sets_ceiling_and_model(self): + self.batch("1", "--for", "2h", "--model", "sonnet", "--mem", "2G") + argv = self.unit_argv() + self.assertIn("--setenv=CLAUDE_CODE_PRINT_BG_WAIT_CEILING_MS=7200000", argv) + self.assertEqual(argv[argv.index("--model") + 1], "sonnet") + self.assertIn(f"MemoryMax={2 * 1024 ** 3}", argv) + + def test_no_claude(self): + self.box.caller["PATH"] = "/nonexistent" + code, out, err = self.batch("2") + self.assertEqual((code, out, err), (1, "", "wf: claude not found on PATH\n")) + self.assertEqual(self.box.fake.ran("systemd-run"), []) + + def test_bad_count(self): + code, _, err = self.batch("0") + self.assertEqual(code, 2) + self.assertIn("N must be ≥ 1", err) + + def test_summary_header_written(self): + code, out, err = self.batch("2", "--lanes", "fast") + self.assertEqual(code, 0) + text = (self.box.env.root() / OUT).read_text() + self.assertEqual(text, "# wf-batch proj 2026-10-01 14:00 (N=2, lanes: fast)\n") + + def test_dry_run(self): + code, out, err = self.batch("2", "--dry-run") + self.assertEqual((code, err), (0, "")) + self.assertIn("CLAUDE_CODE_PRINT_BG_WAIT_CEILING_MS=21600000 wf res run --mem 4G --for 6h --title wf-batch -- ", out) + self.assertIn(f"{self.claude} -p --model opus", out) + self.assertEqual(self.box.fake.ran("systemd-run"), []) + + def test_mem_scales_with_n(self): + for args, mem in ((("4",), "6G"), (("1",), "4G"), (("9",), "8G"), (("4", "--mem", "3G"), "3G")): + code, out, err = self.batch(*args, "--dry-run") + self.assertIn(f"wf res run --mem {mem} --for", out, args) + + def test_defaults_from_history(self): + # peaks 2,3,4 → p95 4 ×1.15 = 4.6 GB; durations 20,30,40 → p90 40 ×1.5 = 60 min; ceiling stays 6h + seed_history(self.box, [hist_rec(f"r-{i}", title="wf-batch", peak=p, minutes=m) + for i, (p, m) in enumerate(((2.0, 20.0), (3.0, 30.0), (4.0, 40.0)))]) + code, out, err = self.batch("4", "--dry-run") + self.assertEqual((code, err), (0, "")) + self.assertIn("CEILING_MS=21600000 wf res run --mem 4.6G --for 60m --title wf-batch -- ", out) + out = self.batch("4", "--mem", "3G", "--for", "2h", "--dry-run")[1] + self.assertIn("CEILING_MS=7200000 wf res run --mem 3G --for 2h --title wf-batch -- ", out) + + def test_busy_exit_3(self): + self.box.available_gb = 1 + code, out, _ = self.batch("2") + self.assertEqual(code, 3) + self.assertIn("busy", out) + + def test_force_past_unused_claim(self): + self.box.wf("run", "--mem", "10G", "--for", "40m", "--title", "big", "--", "x") + code, out, _ = self.batch("2") + self.assertEqual(code, 3) + code, out, err = self.batch("2", "--force") + self.assertEqual(code, 0, (out, err)) + self.assertIn("r-2 started", out) + + +class Status(BatchBase): + def test_none(self): + code, out, err = self.batch("--status") + self.assertEqual((code, out, err), (0, "no batch in this project\n", "")) + + def test_running_and_summary(self): + self.batch("2") + md = self.box.tmp / "proj" / OUT + with md.open("a") as f: + f.write("t-a sonnet done abc123\n") + (md.parent / "wf-batch-2026-09-30-0100.md").write_text("old\n") + code, out, _ = self.batch("--status") + self.assertEqual(code, 0) + lines = out.splitlines() + self.assertTrue(lines[0].startswith('r-1 proj "wf-batch" running'), lines[0]) + self.assertIn(f"log {self.box.env.state / 'logs' / 'r-1.log'}", lines[1]) + self.assertEqual(lines[2:], [f"{OUT}:", "# wf-batch proj 2026-10-01 14:00 (N=2, lanes: every lane with ready tasks, see wf lanes)", + "t-a sonnet done abc123"]) + + def test_orch_log_progress(self): + self.batch("2") + out_dir = self.box.tmp / "proj" / "out" + (out_dir / "wf-orch.log").write_text( + "2026-10-01 13:00 fast sonnet t-old done aaa\n2026-10-01 14:05 fast sonnet t-new done bbb\n") + _, out, _ = self.batch("--status") + self.assertIn("finished since batch start", out) + self.assertIn("t-new done bbb", out) + self.assertNotIn("t-old", out) + + def test_stop_creates_file_and_status_reports(self): + code, out, err = self.batch("--stop") + self.assertEqual((code, err), (0, "")) + f = self.box.tmp / "proj" / "out" / "wf-batch.stop" + self.assertTrue(f.exists()) + self.assertIn("out/wf-batch.stop", out) + code, out, _ = self.batch("--status") + self.assertIn("stop requested: out/wf-batch.stop", out) + + def test_stale_stop_cleared_on_start(self): + f = self.box.tmp / "proj" / "out" / "wf-batch.stop" + f.parent.mkdir(exist_ok=True) + f.touch() + code, out, err = self.batch("2") + self.assertEqual(code, 0) + self.assertIn("stale out/wf-batch.stop cleared", err) + + def test_stop_consumed_on_exit(self): + f = self.box.tmp / "proj" / "out" / "wf-batch.stop" + f.parent.mkdir(exist_ok=True) + orig = wf_res.cmd_run + + def fake(env, cfg, args): + f.touch() + return 0 + wf_res.cmd_run = fake + try: + code, _, _ = self.batch("2") + finally: + wf_res.cmd_run = orig + self.assertEqual(code, 0) + self.assertFalse(f.exists()) + + def test_prompt_has_stop_check(self): + code, out, _ = self.batch("2", "--dry-run") + self.assertIn("out/wf-batch.stop", out) + self.assertIn("stopped: stop file", out) + + def test_prompt_cloud_lane(self): + code, out, _ = self.batch("2", "--dry-run") + self.assertIn("wf orch pick cloud", out) + self.assertIn("never run `wf cloud pull`", out) + self.assertIn("sidecar pulls every 10 min", out) + self.assertIn(f"-- {self.claude} -p", out) # not a cloud project: plain claude -p job + + def test_prompt_stale_branch_once(self): + code, out, _ = self.batch("2", "--dry-run") + self.assertIn("pre-existing branch", out) + self.assertIn("ONCE", out) + + def test_prompt_checks_stop_before_every_spawn(self): + code, out, _ = self.batch("2", "--dry-run") + self.assertIn("Before EVERY spawn", out) + self.assertIn(f"test -e {self.box.tmp / 'proj' / 'out' / 'wf-batch.stop'}", out) + self.assertNotIn("Each round: first", out) + code, out, _ = self.batch("--stop") + self.assertIn("before its next spawn", out) + + def test_n_required_without_status(self): + code, _, err = self.batch() + self.assertEqual(code, 2) + + +PREP_TASKS = """\ +## Awaiting your decision + +## Pending + +- **t-ready** [P1] (<1h): Ready. + Done: works. + +- **t-low** [P3] (<1h): Low, no Done. + +- **t-hi** [P1] (1h): High, no Done. + +- **t-mid** [P2] (1h): Mid, no Done. + +## Needs human + +## Deferred +""" + + +class Prep(BatchBase): + def setUp(self): + super().setUp() + root = self.box.env.root() + (root / "workflow.toml").write_text('format = 1\ntasks = "TASKS.md"\narchive = "archive.md"\n') + (root / "TASKS.md").write_text(PREP_TASKS) + (root / "archive.md").write_text("# Archive\n") + + def test_prompt_lists_targets(self): + code, out, err = self.batch("2", "--prep") + self.assertEqual((code, err), (0, "")) + prompt = self.unit_argv()[-1] + self.assertIn("Tasks: t-hi t-mid.", prompt) + self.assertIn("never implement", prompt) + self.assertIn("wf set <id> --done", prompt) + self.assertIn("wf add -s awaiting", prompt) + self.assertIn(f"to {OUT}", prompt) + for ph in ("never invert recorded original behaviour", "wf add --parent <id> slices", + "concrete files/runs", "[opus: <why>]", "[Model: <m>]"): + self.assertIn(ph, prompt) + self.assertNotIn("{", prompt) + self.assertEqual((self.box.env.root() / OUT).read_text(), + "# wf-batch proj 2026-10-01 14:00 (prep K=2: t-hi t-mid)\n") + self.assertEqual(self.box.ledger().get("r-1").title, "wf-batch") + + def test_lanes_filter(self): + self.batch("5", "--prep", "--lanes", "fast") + self.assertIn("Tasks: t-low.", self.unit_argv()[-1]) + + def test_prep_alias_all(self): + out, err = io.StringIO(), io.StringIO() + with redirect_stdout(out), redirect_stderr(err): + code = wf_res.prep_main(["all"], self.box.env) + self.assertEqual((code, err.getvalue()), (0, "")) + self.assertIn("Tasks: t-hi t-mid", self.unit_argv()[-1]) + + def test_nothing_to_prep_starts_nothing(self): + code, out, err = self.batch("3", "--prep", "--lanes", "nolane") + self.assertEqual((code, out, err), (0, "prep: no pending task without Done (lanes: nolane); nothing started\n", "")) + self.assertEqual(self.box.fake.ran("systemd-run"), []) + self.assertFalse((self.box.env.root() / OUT).exists()) + + def test_prep_mem_default_small(self): + for args, mem in ((("3", "--prep"), "1G"), (("3", "--prep", "--mem", "2G"), "2G")): + code, out, err = self.batch(*args, "--dry-run") + self.assertIn(f"wf res run --mem {mem} --for", out, args) + + def test_dry_run(self): + code, out, err = self.batch("1", "--prep", "--dry-run") + self.assertEqual((code, err), (0, "")) + self.assertIn("--title wf-batch -- ", out) + self.assertIn("t-hi", out) + self.assertEqual(self.box.fake.ran("systemd-run"), []) + + def test_no_project(self): + (self.box.env.root() / "workflow.toml").unlink() + code, out, err = self.batch("1", "--prep") + self.assertEqual(code, 1) + self.assertTrue(err.startswith("wf: "), err) + + +class Dispatch(unittest.TestCase): + def test_wf_forwards_batch(self): + r = subprocess.run([sys.executable, str(HERE / "wf.py"), "batch", "-h"], capture_output=True, text=True) + self.assertEqual(r.returncode, 0) + self.assertIn("usage: wf batch", r.stdout) + + def test_listed_in_wf_help(self): + r = subprocess.run([sys.executable, str(HERE / "wf.py"), "-h"], capture_output=True, text=True) + self.assertIn("batch", r.stdout) + + +if __name__ == "__main__": + unittest.main() + + +FAKE_WF = """import sys, pathlib +d = pathlib.Path(sys.argv[0]).parent +with open(d / "calls.log", "a") as f: + f.write(" ".join(sys.argv[1:]) + "\\n") +if sys.argv[1:3] == ["cloud", "pull"] and (d / "clear").exists(): + for rec in (d / ".wf" / "cloud").glob("*.json"): + rec.unlink() +if sys.argv[1:3] == ["cloud", "pull"] and not (d / "pulled").exists(): + (d / "pulled").touch() + print("t-a: done (session_1, $0.40 usage): merged abc1234") + print("archive session_1 failed: x; archive by hand (wf cloud archive --ended)") + print("report: commit abc1234") + print("t-b: running (session_2, 12 min)") + print("t-c: handback (session_3, $0.10 usage): cloud handback: stuck") + print("t-d: pulled by another wf cloud pull, skipped") + sys.exit(1) +""" + + +class Sidecar(unittest.TestCase): + """wf batch on a cloud = true project: the job runs the orchestrator under the pull sidecar.""" + + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.root = Path(self._tmp.name) + (self.root / "out").mkdir() + self.summary = self.root / OUT + self.summary.write_text("# wf-batch\n") + self.fake = self.root / "fakewf.py" + self.fake.write_text(FAKE_WF) + + def tearDown(self): + self._tmp.cleanup() + + def run_sidecar(self, child_s, every=0.2, rc=7, **kw): + child = [sys.executable, "-c", f"import time, sys; time.sleep({child_s}); sys.exit({rc})"] + out = io.StringIO() + with redirect_stdout(out): + code = wf_res.sidecar(child, self.root, self.summary, every, [sys.executable, str(self.fake)], **kw) + return code, out.getvalue() + + def records(self, *ids): + d = self.root / ".wf" / "cloud" + d.mkdir(parents=True) + for i in ids: + (d / f"{i}.json").write_text("{}") + + def tail_line(self): + return [ln[6:] for ln in self.summary.read_text().splitlines()[1:] if "tail" in ln] + + def test_tail_ends_on_empty_records(self): + self.records("t-b") + (self.root / "clear").touch() + code, out = self.run_sidecar(0.05, every=0.1) + self.assertEqual(code, 7) + self.assertEqual(len([c for c in self.calls() if c.startswith("cloud pull")]), 1) + self.assertIn("orch post t-a cloud --result done --no-pick --commit abc1234", self.calls()) + self.assertEqual(self.tail_line(), ["cloud tail ended: no cloud records left; 1 pulls; left: none (sidecar)"]) + + def test_tail_ends_on_stop_file(self): + self.records("t-b", "t-e") + threading.Timer(0.5, (self.root / wf_res.STOP_FILE).touch).start() + code, _ = self.run_sidecar(0.05, every=0.2) + pulls = [c for c in self.calls() if c.startswith("cloud pull")] + self.assertGreaterEqual(len(pulls), 1) # pulled in the tail until the stop file + self.assertEqual(self.tail_line(), [f"cloud tail ended: stop file; {len(pulls)} pulls; left: t-b t-e (sidecar)"]) + self.assertEqual(code, 7) + + def test_tail_ends_on_cap(self): + self.records("t-b") + shrunk = [] + code, out = self.run_sidecar(0.05, every=0.1, cap=0.35, shrink=lambda: shrunk.append(1) or "shrunk") + self.assertEqual((code, shrunk), (7, [1])) # reservation shrunk once, entering the tail + self.assertIn("shrunk", out) + pulls = [c for c in self.calls() if c.startswith("cloud pull")] + self.assertGreaterEqual(len(pulls), 2) + self.assertTrue(all(c.startswith("cloud pull") or c.startswith("orch post") for c in self.calls())) # no picks + self.assertEqual(self.tail_line(), + [f"cloud tail ended: {0.35 / 3600:g}h cap; {len(pulls)} pulls; left: t-b (sidecar)"]) + + def test_no_tail_without_records(self): + shrunk = [] + code, _ = self.run_sidecar(0.05, every=0.1, shrink=lambda: shrunk.append(1) or "") + self.assertEqual((code, shrunk, self.calls(), self.tail_line()), (7, [], [], [])) + + def test_shrink_reservation(self): + b = BatchBase("run") + b.setUp() + try: + b.batch("4", "--mem", "6G") + msg = wf_res.shrink_reservation(b.box.env, "r-1") + e = b.box.ledger().get("r-1") + self.assertEqual((e.mem_gb, e.cpus, e.state), (0.2, 1, "running")) + self.assertIn("r-1 reservation 6", msg) + self.assertIn("not running", wf_res.shrink_reservation(b.box.env, "r-9")) + finally: + b.tearDown() + + def calls(self): + f = self.root / "calls.log" + return f.read_text().splitlines() if f.exists() else [] + + def test_pulls_and_posts_while_child_lives(self): + code, out = self.run_sidecar(1.0) + self.assertEqual(code, 7) # the orchestrator's exit code + calls = self.calls() + pulls = [c for c in calls if c.startswith("cloud pull")] + self.assertGreaterEqual(len(pulls), 2) # repeats every interval while the child lives + self.assertEqual(pulls[0], f"cloud pull --all --project {self.root}") + self.assertEqual([c for c in calls if c.startswith("orch post")], + ["orch post t-a cloud --result done --no-pick --commit abc1234", + "orch post t-c cloud --result handback --no-pick"]) + lines = self.summary.read_text().splitlines()[1:] + self.assertEqual([ln[6:] for ln in lines], ["cloud t-a done abc1234 (sidecar pull)", + "cloud t-c handback - (sidecar pull)"]) + self.assertIn("t-a: done", out) # pull output lands in the job log + + def test_stops_with_child(self): + code, _ = self.run_sidecar(0.1, every=5, rc=0) + self.assertEqual((code, self.calls()), (0, [])) + + def test_stop_file_ends_pulls(self): + (self.root / wf_res.STOP_FILE).touch() + code, out = self.run_sidecar(0.7) + self.assertEqual((code, self.calls()), (7, [])) + self.assertIn("sidecar: stop file, no more pulls", out) + + def test_cloud_project_job_wraps_claude(self): + b = BatchBase("run") + b.setUp() + try: + root = b.box.env.root() + (root / "workflow.toml").write_text('format = 1\ntasks = "TASKS.md"\narchive = "a.md"\ncloud = true\n') + code, out, _ = b.batch("2", "--dry-run") + self.assertEqual(code, 0) + argv = shlex.split(out.split(" -- ", 1)[1]) + self.assertEqual(argv[:9], [sys.executable, str(HERE / "wf_res.py"), "batch-sidecar", "--root", str(root), + "--summary", OUT, "--every", "600"]) + self.assertEqual(argv[9:12], ["--", str(b.claude), "-p"]) + finally: + b.tearDown() + + def test_main_entry(self): + r = subprocess.run([sys.executable, str(HERE / "wf_res.py"), "batch-sidecar", "--root", str(self.root), + "--summary", OUT, "--every", "5", "--wf", f"{sys.executable} {self.fake}", "--", + sys.executable, "-c", "import sys; sys.exit(3)"], capture_output=True, text=True) + self.assertEqual((r.returncode, r.stderr), (3, "")) + + +LOG_ROWS = "".join(f"2026-10-01T1{i}:00:00 slow opus t-{i} done abc {m}m00s\n" for i, m in enumerate((20, 25, 31))) + + +class Fit(BatchBase): + def log(self, text): + (self.box.env.root() / "out").mkdir(exist_ok=True) + (self.box.env.root() / "out" / "wf-orch.log").write_text(text) + + def test_left_shrinks_k_and_sets_for(self): + self.log(LOG_ROWS) # p90 of 20,25,31 = 31 min + code, out, err = self.batch("4", "--left", "1h40m", "--dry-run") + self.assertEqual((code, err), (0, "")) + self.assertIn("fit: 3 of 4 tasks (left 1h40m, task p90 31m from 3 runs)\n", out) + self.assertIn("CEILING_MS=6000000 wf res run --mem 4.5G --for 1h40m --title wf-batch -- ", out) + self.assertIn("for up to 3 tasks", out) + + def test_left_default_30m(self): + code, out, err = self.batch("4", "--left", "45m", "--dry-run") + self.assertIn("fit: 1 of 4 tasks (left 45m, task p90 30m default (0 runs))\n", out) + self.assertIn("--for 45m", out) + + def test_left_too_short_starts_nothing(self): + self.log(LOG_ROWS) + code, out, err = self.batch("4", "--left", "30m") + self.assertEqual((code, err), (0, "")) + self.assertEqual(out, "fit: 0 of 4 tasks (left 30m, task p90 31m from 3 runs); nothing started\n") + self.assertEqual(self.box.fake.ran("systemd-run"), []) + + def test_prompt_has_deadline_check(self): + code, out, _ = self.batch("2", "--for", "2h", "--dry-run") + self.assertIn("wf batch --time-left 2026-10-01T16:00+00:00", out) + self.assertIn("stopped: deadline", out) + + def test_time_left(self): + self.log(LOG_ROWS) + code, out, err = self.batch("--time-left", "2026-10-01T15:00+00:00") + self.assertEqual((code, out, err), (0, "time left 1h, task p90 31m (3 runs): spawn\n", "")) + code, out, err = self.batch("--time-left", "2026-10-01T14:30+00:00") + self.assertEqual((code, out), (0, "time left 30m, task p90 31m (3 runs): stop\n")) + code, out, err = self.batch("--time-left", "2026-10-01T13:00+00:00") + self.assertEqual((code, out), (0, "time left 0m, task p90 31m (3 runs): stop\n")) + + def test_time_left_bad(self): + code, out, err = self.batch("--time-left", "soon") + self.assertEqual(code, 1) + self.assertIn("wf: --time-left", err) |