diff options
Diffstat (limited to 'tests/test_model.py')
| -rw-r--r-- | tests/test_model.py | 247 |
1 files changed, 247 insertions, 0 deletions
diff --git a/tests/test_model.py b/tests/test_model.py new file mode 100644 index 0000000..b2e0878 --- /dev/null +++ b/tests/test_model.py @@ -0,0 +1,247 @@ +import json +import os +import subprocess +import sys +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from wflib import tasks as T +from wflib import lanes as L +from test_cli import Cli +from test_check import Base, CLEAN +from wflib import check as K +from wflib import config as C + +LANES = """\ +# Tasks — demo + +## Awaiting your decision + +## Pending + +- **t-one** [P1] (1h): One. + - Steps: a + Model: sonnet + - After: [[t-done]] + +- **t-two** [P2] (1h): Two. + - Model: haiku + +- **t-three** [P2] (<1h): Three. + +## Needs human + +## Deferred +""" + + +class ModelLineTest(unittest.TestCase): + def setUp(self): + self.doc = T.parse(LANES) + + def test_parsed_from_body_default_opus(self): + self.assertEqual([i.model for i in self.doc.section("pending").items], ["sonnet", "haiku", "opus"]) + + def test_extra_words_ignored(self): + item = T.parse_block("- **t-x** [P1] (1h): X.\n - Model: sonnet ok\n") + self.assertEqual(item.model, "sonnet") + + def test_set_replaces_in_place(self): + T.set_fields(self.doc, "t-one", model="opus") + self.assertEqual(self.doc.item("t-one").body, [" - Steps: a", " Model: opus", " - After: [[t-done]]"]) + + def test_set_adds_before_after_and_ref(self): + doc = T.parse(CLEAN) + T.set_fields(doc, "t-one", model="haiku") + self.assertEqual(doc.item("t-one").body, + [" Model: haiku", " - After: [[t-done]]", " Ref: DESIGN.md#terrain, docs/specs/terrain.md"]) + + def test_set_empty_removes(self): + T.set_fields(self.doc, "t-two", model="") + self.assertEqual(self.doc.item("t-two").body, []) + self.assertEqual(self.doc.item("t-two").model, "opus") + + def test_set_bad_value(self): + with self.assertRaisesRegex(T.TaskError, r"model 'gpt' \(want haiku, sonnet, opus\)"): + T.set_fields(self.doc, "t-one", model="gpt") + + def test_note_goes_before_model_line(self): + T.add_note(self.doc, "t-two", "why") + self.assertEqual(self.doc.item("t-two").body, [" - why", " - Model: haiku"]) + + +def doc(items): + return T.parse("## Awaiting your decision\n\n## Pending\n\n" + items + "\n## Needs human\n\n## Deferred\n") + + +class PickOrderTest(unittest.TestCase): + def pick(self, items, model="opus", archived=()): + item, _ = L.pick(doc(items), set(archived), L.DEFAULT_LANES, "1h", None, model) + return item.id if item else None + + def test_priority_inherited_from_waiting_task(self): + self.assertEqual(self.pick("- **t-c** [P1] (1h): C.\n\n- **t-b** [P3] (1h): B.\n\n" + "- **t-a** [P0] (1h): A.\n - After: [[t-b]]\n"), "t-b") + + def test_inherited_transitively(self): + self.assertEqual(self.pick("- **t-c** [P1] (1h): C.\n\n- **t-x** [P3] (1h): X.\n\n" + "- **t-b** [P3] (1h): B.\n - After: [[t-x]]\n\n" + "- **t-a** [P0] (1h): A.\n - After: [[t-b]]\n"), "t-x") + + def test_unblocking_only_low_work_does_not_jump(self): + self.assertEqual(self.pick("- **t-c** [P1] (1h): C.\n\n- **t-b** [P3] (1h): B.\n\n" + "- **t-a** [P3] (1h): A.\n - After: [[t-b]]\n"), "t-c") + + def test_other_lane_waiting_first(self): + self.assertEqual(self.pick("- **t-f** [P2] (1h): F.\n\n- **t-k** [P2] (1h): K.\n\n" + "- **t-l** [P2] (1h): L.\n - After: [[t-k]]\n\n- **t-g** [P2] (1h): G.\n\n" + "- **t-h** [P2] (<1h): H.\n - After: [[t-g]]\n"), "t-g") + + def test_more_waiting_first(self): + self.assertEqual(self.pick("- **t-f** [P2] (1h): F.\n - After: [[t-z]]\n\n- **t-g** [P2] (1h): G.\n\n" + "- **t-i** [P2] (1h): I.\n - After: [[t-g]]\n\n" + "- **t-j** [P2] (1h): J.\n - After: [[t-i]]\n\n" + "- **t-z** [P2] (1h): Z.\n", archived=()), "t-g") + + def test_parent_waits_on_its_slices(self): + self.assertEqual(self.pick("- **t-q** [P1] (1h): Q.\n\n- **t-p** [P0] (5h): P.\n - Slices: [[t-p-1]]\n\n" + "- **t-p-1** [P2] (1h): P1.\n"), "t-p-1") + + def test_lane_filter(self): + items = ("- **t-a** [P0] (1h): A.\n\n- **t-b** [P1] (1h): B.\n Model: sonnet\n\n" + "- **t-c** [P2] (1h): C.\n Model: haiku\n") + self.assertEqual([self.pick(items, m) for m in ("opus", "sonnet", "haiku", None)], + ["t-a", "t-b", "t-c", "t-a"]) + self.assertIsNone(self.pick("- **t-a** [P0] (1h): A.\n", "haiku")) + + def test_skipped_are_ranked_before_pick(self): + d = doc("- **t-a** [P0] (1h) (blocked: [[a-k]]): A.\n\n- **t-s** [P0] (1h): S.\n Model: sonnet\n\n" + "- **t-b** [P1] (1h): B.\n\n- **t-c** [P2] (1h): C.\n - After: [[t-b]]\n") + item, skipped = L.pick(d, set(), L.DEFAULT_LANES, "1h", "slow", "opus") + self.assertEqual((item.id, [(i.id, why) for i, why in skipped]), ("t-s", [("t-a", "blocked: a-k")])) + item, skipped = L.pick(d, set(), L.DEFAULT_LANES, "1h", "slow", "haiku") + self.assertEqual((item, skipped), (None, [])) + + +class ModelCheckTest(Base): + def test_values(self): + self.pending("- **t-a** [P1] (1h): A.\n Model: gpt\n\n" + "- **t-b** [P1] (1h): B.\n - Model: sonnet ok\n\n" + "- **t-c** [P1] (1h): C.\n Model: haiku\n Model: opus\n\n" + "- **t-d** [P1] (1h): D.\n") + errors, warnings = K.check(C.load(self.root)) + self.assertEqual([(p.id, p.message) for p in errors], + [("t-a", "Model 'gpt' (want haiku, sonnet, opus)"), ("t-c", "two Model lines")]) + self.assertEqual([(p.id, p.message) for p in warnings], + [("t-b", "Model line 'sonnet ok': write 'Model: sonnet' (wf set --model)")]) + + +class ModelCliTest(Cli): + tasks_text = LANES + + def test_list_column_and_filter(self): + self.assertEqual(self.ok("list"), + "t-one P1 1h - sonnet slow One\n" + "t-two P2 1h - haiku slow Two\n" + "t-three P2 <1h - opus fast Three\n" + "pending 3 · human 0 · awaiting 0 · deferred 0\n") + self.assertEqual(self.ok("list", "--model", "opus"), + "t-three P2 <1h - opus fast Three\n" + "pending 3 · human 0 · awaiting 0 · deferred 0\n") + + def test_add_and_set(self): + self.assertEqual(self.ok("add", "Four. Goal.", "-p", "3", "-e", "1h", "--model", "haiku", "--ref", "DESIGN.md"), + "- **t-four** [P3] (1h): Four. Goal.\n") + self.assertEqual(self.item("t-four"), "- **t-four** [P3] (1h): Four. Goal.\n Model: haiku\n Ref: DESIGN.md\n") + self.ok("set", "t-four", "--model", "sonnet") + self.assertEqual(self.item("t-four"), "- **t-four** [P3] (1h): Four. Goal.\n Model: sonnet\n Ref: DESIGN.md\n") + self.ok("set", "t-four", "--model", "") + self.assertEqual(self.item("t-four"), "- **t-four** [P3] (1h): Four. Goal.\n Ref: DESIGN.md\n") + + def test_next_model_ceiling(self): + # --as M takes tasks with Model ≤ M; no --lane = all lanes, priority order + self.assertEqual(self.ok("next", "--as", "haiku", "--brief"), "- **t-two** [P2] (1h): Two.\n - Model: haiku\n") + self.assertEqual(self.ok("next", "--as", "opus", "--brief").splitlines()[0], "- **t-one** [P1] (1h): One.") + self.assertEqual(self.ok("next", "--as", "sonnet", "--brief").splitlines()[0], "- **t-one** [P1] (1h): One.") + self.assertEqual(self.ok("next", "--lane", "fast", "--as", "opus", "--brief"), "- **t-three** [P2] (<1h): Three.\n") + self.ok("set", "t-two", "--model", "opus") + self.assertEqual(self.fails("next", "--as", "haiku", "--brief").splitlines()[-1], + "wf: nothing pickable for all lanes (haiku) in Pending") + + def test_next_without_as_is_haiku_with_hint(self): + out = self.ok("next", "--brief") + self.assertEqual(out, "no --as: treated as haiku; pass --as haiku|sonnet|opus\n" + "- **t-two** [P2] (1h): Two.\n - Model: haiku\n") + + def session_env(self, sock): + sock.write_text("") + return {"CLAUDE_CODE_MESSAGING_SOCKET": str(sock), "CLAUDE_PID": str(os.getpid()), + "CLAUDE_CODE_SESSION_ID": "sess-1"} + + def dead_pid(self): + p = subprocess.Popen(["true"]) + p.wait() + return p.pid + + def test_next_registers_session(self): + env = self.session_env(self.root / "me.sock") + self.ok("next", "--as", "opus", "--brief", env=env) + reg = json.loads((self.root / ".wf" / "sessions" / "all.json").read_text()) + self.assertEqual({k: reg[k] for k in ("lane", "model", "socket", "pid", "session")}, + {"lane": "all", "model": "opus", "socket": str(self.root / "me.sock"), "pid": os.getpid(), "session": "sess-1"}) + self.assertEqual((self.root / ".wf" / ".gitignore").read_text(), "*\n") + + def test_no_register_without_as_or_env(self): + none = {"CLAUDE_CODE_MESSAGING_SOCKET": "", "CLAUDE_PID": "", "CLAUDE_CODE_SESSION_ID": ""} + self.ok("next", "--as", "opus", env=none) + self.ok("next", env=self.session_env(self.root / "me.sock")) + self.assertFalse((self.root / ".wf").exists()) + + def register(self, lane, model, sock, pid): + d = self.root / ".wf" / "sessions" + d.mkdir(parents=True, exist_ok=True) + (d / f"{lane}.json").write_text(json.dumps({"lane": lane, "model": model, "socket": str(sock), "pid": pid, + "session": "s", "at": "2026-10-04T10:00"})) + + def test_next_shows_lanes_and_waiting(self): + self.ok("set", "t-three", "--after", "t-one") + self.ok("set", "t-two", "--model", "opus") + sock = self.root / "son.sock" + sock.write_text("") + self.register("slow", "sonnet", sock, os.getpid()) + self.register("fast", "haiku", self.root / "gone.sock", self.dead_pid()) + none = {"CLAUDE_CODE_MESSAGING_SOCKET": ""} + code, out, err = self.wf("next", "--lane", "fast", "--as", "opus", env=none) # fast: no fallback lane + self.assertEqual((code, err), (1, "wf: nothing pickable for fast (opus) in Pending\n")) + self.assertIn("===== Lanes =====\n" + "fast (you): 0 pickable · 1 waiting\n" + f"slow: 2 pickable · 0 waiting · session uds:{sock} (alive, sonnet)\n\n" + "===== Waiting on other lanes =====\n" + f'- t-three waits on t-one (slow lane) → message uds:{sock}: "t-one blocks my t-three, please take it"\n', out) + self.assertEqual(self.ok("lanes", env=none), + "fast: 0 pickable · 1 waiting · no session\n" + f"slow: 0 pickable · 0 waiting · 2 not runner-ready (no Done) · session uds:{sock} (alive, sonnet)\n") + + def test_next_single_lane_prints_no_lanes_block(self): + (self.root / "workflow.toml").write_text(self.toml + '[lanes.one]\nefforts = ["<1h", "1h"]\n') + self.assertNotIn("Lanes", self.ok("next", "--as", "opus", env={"CLAUDE_CODE_MESSAGING_SOCKET": ""})) + + def test_done_notifies_other_lanes(self): + self.ok("add", "Four.", "-p", "3", "-e", "<1h", "--model", "haiku", "--after", "t-three") + self.ok("add", "Five.", "-p", "3", "-e", "1h", "--after", "t-three") + self.ok("add", "Six.", "-p", "3", "-e", "1h", "--model", "sonnet", "--after", "t-three") + sock = self.root / "h.sock" + sock.write_text("") + self.register("slow", "opus", sock, os.getpid()) + out = self.ok("done", "t-three", "-m", "ok") + self.assertIn(f"notify slow uds:{sock}: now pickable t-five, t-six\n", out) + self.assertNotIn("t-four", out) # same lane as t-three (fast) + + def test_bad_model_is_usage_error(self): + self.fails("add", "Four.", "-p", "3", "-e", "1h", "--model", "gpt", code=2) + self.fails("set", "t-one", "--model", "gpt", code=2) + + +if __name__ == "__main__": + unittest.main() |