"""U-031 tests: work_graph_orchestrate.py -- the graph_manifest -> per-node execution bridge closed during the REQ-01~14 regression audit. Real sqlite, run_fn mocked at the claude subprocess boundary (execute_work_graph is tested for real elsewhere, in test_work_graph_execute.py / the 2026-08-01 spike -- this module's own job is orchestration, not re-proving that). Run: python3 -m unittest src.test_work_graph_orchestrate -v """ import unittest from unittest.mock import patch from src import schema, work_graph_execute, work_graph_orchestrate as wgo class TopologicalOrderTest(unittest.TestCase): def test_linear_chain_via_depends_on(self): manifest = { "nodes": [ {"node_id": "brd", "depends_on": []}, {"node_id": "prd", "depends_on": ["brd"]}, {"node_id": "sdd", "depends_on": ["prd"]}, ], "edges": [], } self.assertEqual([n["node_id"] for n in wgo.topological_order(manifest)], ["brd", "prd", "sdd"]) def test_linear_chain_via_edges_only(self): # real LLM output might populate edges but not each node's own # depends_on -- must still order correctly manifest = { "nodes": [{"node_id": "a"}, {"node_id": "b"}, {"node_id": "c"}], "edges": [{"from": "a", "to": "b"}, {"from": "b", "to": "c"}], } self.assertEqual([n["node_id"] for n in wgo.topological_order(manifest)], ["a", "b", "c"]) def test_diamond_fan_out_fan_in(self): manifest = { "nodes": [ {"node_id": "n1", "depends_on": []}, {"node_id": "n2", "depends_on": ["n1"]}, {"node_id": "n3", "depends_on": ["n1"]}, {"node_id": "n4", "depends_on": ["n2", "n3"]}, ], "edges": [], } order = [n["node_id"] for n in wgo.topological_order(manifest)] self.assertEqual(order[0], "n1") self.assertEqual(order[-1], "n4") self.assertEqual(set(order[1:3]), {"n2", "n3"}) def test_cycle_raises(self): manifest = { "nodes": [{"node_id": "x", "depends_on": ["y"]}, {"node_id": "y", "depends_on": ["x"]}], "edges": [], } with self.assertRaises(ValueError): wgo.topological_order(manifest) def test_independent_nodes_ordered_deterministically(self): manifest = {"nodes": [{"node_id": "b"}, {"node_id": "a"}], "edges": []} self.assertEqual([n["node_id"] for n in wgo.topological_order(manifest)], ["a", "b"]) class RunGraphTestBase(unittest.TestCase): def setUp(self): self.conn = schema.connect(":memory:") schema.create_schema(self.conn) self.conn.execute( "INSERT INTO project (project_id, name, target_folder_path, has_git_repo, status) " "VALUES ('proj-1', 'P', '/tmp/proj-1', 1, 'active')" ) self.conn.execute( "INSERT INTO task (task_id, project_id, goal_text, task_type, status, assigned_org_node) " "VALUES ('task-1', 'proj-1', 'goal', 'goal', 'coding', 'planner')" ) self.conn.execute( "INSERT INTO work_graph_run (run_id, task_id, org_node, graph_manifest, status) " "VALUES ('run-1', 'task-1', 'planner', '{}', 'running')" ) self.conn.commit() def tearDown(self): self.conn.close() _OUT_SCHEMA = {"type": "object", "required": ["content"]} _IN_SCHEMA = {"type": "object", "required": ["node_id"]} class RunGraphHappyPathTest(RunGraphTestBase): def test_runs_nodes_in_order_and_returns_outputs(self): manifest = { "nodes": [ {"node_id": "brd", "artifact_type": "brd", "depends_on": []}, {"node_id": "prd", "artifact_type": "prd", "depends_on": ["brd"]}, ], "edges": [], } calls = [] def build_prompt(node, outputs): calls.append((node["node_id"], dict(outputs))) return "prompt for " + node["node_id"] with patch.object(work_graph_execute, "write_workflow_script"), \ patch.object(work_graph_execute, "execute_work_graph") as mock_execute: mock_execute.side_effect = [{"content": "brd text"}, {"content": "prd text"}] outputs = wgo.run_graph( self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest, build_prompt, _IN_SCHEMA, _OUT_SCHEMA, ) self.assertEqual(outputs, {"brd": {"content": "brd text"}, "prd": {"content": "prd text"}}) self.assertEqual(calls[0], ("brd", {})) self.assertEqual(calls[1], ("prd", {"brd": {"content": "brd text"}})) def test_creates_work_node_execution_rows_and_placeholder_skill(self): manifest = {"nodes": [{"node_id": "brd", "depends_on": []}], "edges": []} with patch.object(work_graph_execute, "write_workflow_script"), \ patch.object(work_graph_execute, "execute_work_graph", return_value={"content": "x"}): wgo.run_graph( self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest, lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA, ) row = self.conn.execute( "SELECT status, subagent_type, skill_ref FROM work_node_execution WHERE run_id = 'run-1' AND node_id = 'brd'" ).fetchone() self.assertEqual(row, ("done", "planner", "planner-doc-writing")) skill_row = self.conn.execute("SELECT owner_org_node FROM skill WHERE skill_id = 'planner-doc-writing'").fetchone() self.assertEqual(skill_row[0], "planner") def test_resumes_without_rerunning_done_nodes(self): manifest = { "nodes": [ {"node_id": "brd", "depends_on": []}, {"node_id": "prd", "depends_on": ["brd"]}, ], "edges": [], } calls = [] def build_prompt(node, outputs): return "p-" + node["node_id"] with patch.object(work_graph_execute, "write_workflow_script"), \ patch.object(work_graph_execute, "execute_work_graph") as mock_execute: mock_execute.side_effect = lambda cwd, name, args, schema, timeout=600: calls.append(name) or {"content": "c-" + name} wgo.run_graph( self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest, build_prompt, _IN_SCHEMA, _OUT_SCHEMA, ) self.assertEqual(len(calls), 2) # simulate a restart: run_graph called again for the same run_id with patch.object(work_graph_execute, "write_workflow_script"), \ patch.object(work_graph_execute, "execute_work_graph") as mock_execute2: mock_execute2.side_effect = AssertionError("should not be called -- both nodes are already done") outputs = wgo.run_graph( self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest, build_prompt, _IN_SCHEMA, _OUT_SCHEMA, ) self.assertEqual(set(outputs.keys()), {"brd", "prd"}) class RunGraphUsageLimitTest(RunGraphTestBase): def test_usage_limit_reached_translates_to_quota_wait(self): manifest = {"nodes": [{"node_id": "brd", "depends_on": []}], "edges": []} attempts = {"n": 0} def fake_execute(cwd, name, args, output_schema, timeout=600): attempts["n"] += 1 if attempts["n"] == 1: raise work_graph_execute.UsageLimitReachedError("hit your session limit", "session") return {"content": "recovered"} sleeps = [] with patch.object(work_graph_execute, "write_workflow_script"), \ patch.object(work_graph_execute, "execute_work_graph", side_effect=fake_execute): outputs = wgo.run_graph( self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest, lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA, sleep=lambda s: sleeps.append(s), now=lambda: 0.0, ) self.assertEqual(outputs["brd"], {"content": "recovered"}) self.assertEqual(sleeps, [5 * 60]) class RunGraphClassifierBlockedTest(RunGraphTestBase): def test_classifier_blocked_marks_node_blocked_and_propagates(self): manifest = {"nodes": [{"node_id": "brd", "depends_on": []}], "edges": []} with patch.object(work_graph_execute, "write_workflow_script"), \ patch.object( work_graph_execute, "execute_work_graph", side_effect=work_graph_execute.ClassifierBlockedError("blocked by classifier"), ): with self.assertRaises(work_graph_execute.ClassifierBlockedError): wgo.run_graph( self.conn, "run-1", "planner", "planner-doc-writing", "/tmp/proj-1", manifest, lambda node, outputs: "p", _IN_SCHEMA, _OUT_SCHEMA, ) node_exec_id = self.conn.execute( "SELECT node_exec_id FROM work_node_execution WHERE run_id = 'run-1' AND node_id = 'brd'" ).fetchone()[0] status = self.conn.execute( "SELECT status FROM work_node_execution WHERE node_exec_id = ?", (node_exec_id,) ).fetchone()[0] self.assertEqual(status, "blocked") if __name__ == "__main__": unittest.main()