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 --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 slices", "concrete files/runs", "[opus: ]", "[Model: ]"): 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)