test_work_graph_plan.py 18.7 KB
"""U-011 tests: c_work_graph_plan, mocked at the subprocess boundary (CLAUDE.md
Testing constraint). Real-call coverage is in src/test_work_graph_plan_e2e.py
(RUN_REAL_AGENT_E2E=1) -- that real call is what proved --json-schema doesn't
constrain output when --agent is set, which is why these mocks return
free-form ```json-fenced text in `result`, not `structured_output`.

Run: python3 -m unittest src.test_work_graph_plan -v
"""

import json
import subprocess
import unittest
from unittest.mock import patch

from src import schema, work_graph_execute, work_graph_orchestrate, work_graph_plan

_MANIFEST = {
    "nodes": [
        {"node_id": "n1", "artifact_type": "brd", "depends_on": []},
        {"node_id": "n2", "artifact_type": "prd", "depends_on": ["n1"]},
        {"node_id": "n3", "artifact_type": "sdd", "depends_on": ["n2"]},
        {"node_id": "n4", "artifact_type": "benchmark", "depends_on": ["n3"]},
    ],
    "edges": [
        {"from": "n1", "to": "n2", "edge_type": "needs_judgment"},
        {"from": "n2", "to": "n3", "edge_type": "needs_judgment"},
        {"from": "n3", "to": "n4", "edge_type": "needs_judgment"},
    ],
}


def _completed_process(returncode=0, stdout_obj=None, stderr=""):
    stdout = json.dumps(stdout_obj) if stdout_obj is not None else ""
    return subprocess.CompletedProcess(args=["claude"], returncode=returncode, stdout=stdout, stderr=stderr)


def _fenced_result(manifest, wrapper_key=None):
    body = {wrapper_key: manifest} if wrapper_key else manifest
    return "Here is my plan:\n\n```json\n" + json.dumps(body) + "\n```\n\nDone."


class PlanWorkGraphTestBase(unittest.TestCase):
    def setUp(self):
        # check_same_thread=False: GenericPathEndToEndViaRunGraphTest below
        # drives a real multi-node concurrent layer through run_graph()
        # (G), which calls into this conn from ThreadPoolExecutor worker
        # threads via _LockedConn -- see test_work_graph_orchestrate.py's
        # RunGraphTestBase for the same reasoning.
        self.conn = schema.connect(":memory:", check_same_thread=False)
        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', 'queued', 'planner')"
        )
        self.conn.commit()

    def tearDown(self):
        self.conn.close()


class ExtractGraphManifestTest(unittest.TestCase):
    def test_bare_json_result(self):
        self.assertEqual(work_graph_plan._extract_graph_manifest(json.dumps(_MANIFEST)), _MANIFEST)

    def test_fenced_json_block_with_prose_around_it(self):
        self.assertEqual(work_graph_plan._extract_graph_manifest(_fenced_result(_MANIFEST)), _MANIFEST)

    def test_fenced_json_wrapped_under_graph_manifest_key(self):
        text = _fenced_result(_MANIFEST, wrapper_key="graph_manifest")
        self.assertEqual(work_graph_plan._extract_graph_manifest(text), _MANIFEST)

    def test_no_json_anywhere_returns_none(self):
        self.assertIsNone(work_graph_plan._extract_graph_manifest("sorry, I can't do that"))


class NormalizeGraphManifestTest(unittest.TestCase):
    """U-031 regression: a real Planner call (2026-08-01) came back using
    `id` instead of the documented `node_id`, and never populated any
    node's own `depends_on` -- dependencies existed only as `edges`. This
    is a real, observed deviation, not a hypothetical one; the fixture
    below is the actual shape captured from that call (fields trimmed to
    what normalization touches)."""

    def test_already_well_formed_manifest_untouched(self):
        manifest = {
            "nodes": [
                {"node_id": "brd", "artifact_type": "brd", "depends_on": []},
                {"node_id": "prd", "artifact_type": "prd", "depends_on": ["brd"]},
            ],
            "edges": [{"from": "brd", "to": "prd", "edge_type": "needs_judgment"}],
        }
        result = work_graph_plan._normalize_graph_manifest(manifest)
        self.assertEqual(result["nodes"][1]["depends_on"], ["brd"])

    def test_id_key_normalized_to_node_id(self):
        manifest = {
            "nodes": [{"id": "brd", "artifact_type": "brd"}, {"id": "prd", "artifact_type": "prd"}],
            "edges": [{"from": "brd", "to": "prd", "edge_type": "needs_judgment"}],
        }
        result = work_graph_plan._normalize_graph_manifest(manifest)
        self.assertEqual([n["node_id"] for n in result["nodes"]], ["brd", "prd"])

    def test_depends_on_derived_from_edges_when_absent(self):
        manifest = {
            "nodes": [{"id": "brd"}, {"id": "prd"}, {"id": "sdd"}, {"id": "benchmark"}],
            "edges": [
                {"from": "brd", "to": "prd", "edge_type": "needs_judgment"},
                {"from": "prd", "to": "sdd", "edge_type": "needs_judgment"},
                {"from": "sdd", "to": "benchmark", "edge_type": "needs_judgment"},
            ],
        }
        result = work_graph_plan._normalize_graph_manifest(manifest)
        by_id = {n["node_id"]: n for n in result["nodes"]}
        self.assertEqual(by_id["brd"]["depends_on"], [])
        self.assertEqual(by_id["prd"]["depends_on"], ["brd"])
        self.assertEqual(by_id["sdd"]["depends_on"], ["prd"])
        self.assertEqual(by_id["benchmark"]["depends_on"], ["sdd"])

    def test_does_not_duplicate_dependency_already_present(self):
        manifest = {
            "nodes": [{"node_id": "brd", "depends_on": []}, {"node_id": "prd", "depends_on": ["brd"]}],
            "edges": [{"from": "brd", "to": "prd", "edge_type": "needs_judgment"}],
        }
        result = work_graph_plan._normalize_graph_manifest(manifest)
        self.assertEqual(result["nodes"][1]["depends_on"], ["brd"])

    def test_edge_referencing_unknown_node_id_ignored(self):
        manifest = {"nodes": [{"id": "brd"}], "edges": [{"from": "brd", "to": "no-such-node"}]}
        result = work_graph_plan._normalize_graph_manifest(manifest)
        self.assertEqual(result["nodes"][0]["depends_on"], [])

    def test_canonical_doc_chain_enforced_even_with_zero_dependency_info(self):
        # 2026-08-03 regression: a real call omitted depends_on AND edges
        # entirely -- worse than test_depends_on_derived_from_edges_when_absent,
        # which at least had edges to derive from. With nothing to
        # normalize, topological_order() scheduled 'benchmark' before 'sdd'
        # (alphabetically first among zero-in-degree nodes), and
        # e2e_pipeline's build_prompt() crashed formatting a template that
        # needed sdd's not-yet-produced output.
        manifest = {
            "nodes": [
                {"node_id": "brd", "artifact_type": "brd"},
                {"node_id": "prd", "artifact_type": "prd"},
                {"node_id": "sdd", "artifact_type": "sdd"},
                {"node_id": "benchmark", "artifact_type": "benchmark"},
                {"node_id": "test_dataset", "artifact_type": "test_dataset"},
                {"node_id": "test_script", "artifact_type": "test_script"},
            ],
            "edges": [],
        }
        result = work_graph_plan._normalize_graph_manifest(manifest)
        by_type = {n["artifact_type"]: n for n in result["nodes"]}
        self.assertEqual(by_type["brd"]["depends_on"], [])
        self.assertEqual(by_type["prd"]["depends_on"], ["brd"])
        self.assertEqual(by_type["sdd"]["depends_on"], ["prd"])
        self.assertEqual(by_type["benchmark"]["depends_on"], ["sdd"])
        self.assertEqual(by_type["test_dataset"]["depends_on"], ["benchmark"])
        # test_script needs both -- its own prompt references
        # {benchmark_content} directly, not just what survives into
        # test_dataset_content (see work_graph_plan._normalize_graph_manifest).
        self.assertEqual(by_type["test_script"]["depends_on"], ["test_dataset", "benchmark"])

    def test_canonical_chain_enforcement_uses_real_node_ids_not_artifact_types(self):
        # node_id and artifact_type can differ (e.g. real LLM output naming
        # nodes "n1"/"n2"/... while still setting artifact_type correctly) --
        # the enforced dependency must reference the real node_id.
        manifest = {
            "nodes": [
                {"node_id": "step-1", "artifact_type": "brd"},
                {"node_id": "step-2", "artifact_type": "prd"},
                {"node_id": "step-3", "artifact_type": "sdd"},
                {"node_id": "step-4", "artifact_type": "benchmark"},
                {"node_id": "step-5", "artifact_type": "test_dataset"},
                {"node_id": "step-6", "artifact_type": "test_script"},
            ],
            "edges": [],
        }
        result = work_graph_plan._normalize_graph_manifest(manifest)
        by_type = {n["artifact_type"]: n for n in result["nodes"]}
        self.assertEqual(by_type["prd"]["depends_on"], ["step-1"])
        self.assertEqual(by_type["sdd"]["depends_on"], ["step-2"])
        self.assertEqual(by_type["benchmark"]["depends_on"], ["step-3"])
        self.assertEqual(by_type["test_dataset"]["depends_on"], ["step-4"])
        self.assertEqual(by_type["test_script"]["depends_on"], ["step-5", "step-4"])

    def test_partial_canonical_set_does_not_trigger_enforcement(self):
        # only brd/prd/sdd present, no benchmark -- not a recognized
        # canonical 4-node doc chain, must not guess at a shape.
        manifest = {
            "nodes": [
                {"node_id": "brd", "artifact_type": "brd"},
                {"node_id": "prd", "artifact_type": "prd"},
                {"node_id": "sdd", "artifact_type": "sdd"},
            ],
            "edges": [],
        }
        result = work_graph_plan._normalize_graph_manifest(manifest)
        for node in result["nodes"]:
            self.assertEqual(node["depends_on"], [])


class HappyPathTest(PlanWorkGraphTestBase):
    def test_returns_run_id_and_graph_manifest(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(
                stdout_obj={"is_error": False, "result": _fenced_result(_MANIFEST)}
            )
            result = work_graph_plan.plan_work_graph(
                self.conn, "planner", "task-1", "实现贪吃蛇游戏", "/tmp/proj-1"
            )
        self.assertEqual(result["graph_manifest"], _MANIFEST)
        self.assertTrue(result["run_id"])

    def test_creates_work_graph_run_row_in_planning_status(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(
                stdout_obj={"is_error": False, "result": _fenced_result(_MANIFEST)}
            )
            result = work_graph_plan.plan_work_graph(
                self.conn, "planner", "task-1", "实现贪吃蛇游戏", "/tmp/proj-1"
            )
        row = self.conn.execute(
            "SELECT task_id, org_node, status, graph_manifest FROM work_graph_run WHERE run_id = ?",
            (result["run_id"],),
        ).fetchone()
        self.assertEqual(row[0], "task-1")
        self.assertEqual(row[1], "planner")
        self.assertEqual(row[2], "planning")
        self.assertEqual(json.loads(row[3]), _MANIFEST)

    def test_invokes_claude_with_expected_args(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(
                stdout_obj={"is_error": False, "result": _fenced_result(_MANIFEST)}
            )
            work_graph_plan.plan_work_graph(self.conn, "planner", "task-1", "goal", "/tmp/proj-1")

        args, kwargs = run.call_args
        argv = args[0]
        self.assertEqual(argv[:3], ["claude", "--agent", "planner"])
        self.assertIn("--permission-mode", argv)
        self.assertIn("auto", argv)
        self.assertNotIn("--bare", argv)
        # --json-schema deliberately dropped -- confirmed for real it's a no-op
        # when combined with --agent (see module docstring)
        self.assertNotIn("--json-schema", argv)
        self.assertEqual(kwargs["cwd"], "/tmp/proj-1")
        self.assertEqual(kwargs["stdin"], subprocess.DEVNULL)


class FailureTest(PlanWorkGraphTestBase):
    def test_nonzero_exit_raises_planning_failed(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(returncode=1, stdout_obj=None, stderr="boom")
            with self.assertRaises(work_graph_plan.PlanningFailedError):
                work_graph_plan.plan_work_graph(self.conn, "planner", "task-1", "goal", "/tmp/proj-1")

    def test_is_error_true_raises_planning_failed(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(stdout_obj={"is_error": True})
            with self.assertRaises(work_graph_plan.PlanningFailedError):
                work_graph_plan.plan_work_graph(self.conn, "planner", "task-1", "goal", "/tmp/proj-1")

    def test_unparseable_result_raises_planning_failed(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(stdout_obj={"is_error": False, "result": "no json here"})
            with self.assertRaises(work_graph_plan.PlanningFailedError):
                work_graph_plan.plan_work_graph(self.conn, "planner", "task-1", "goal", "/tmp/proj-1")
        self.assertEqual(self.conn.execute("SELECT COUNT(*) FROM work_graph_run").fetchone()[0], 0)


# E (2026-08-06): cross-Org-Agent dynamic planning path -- org_node != 'planner'
# gets _GENERIC_PROMPT_TEMPLATE instead of the fixed doc-chain prompt, and each
# node carries its own task_prompt instead of a fixed artifact_type.

_CODER_THREE_MODULE_MANIFEST = {
    "nodes": [
        {"node_id": "test-module-a", "depends_on": [], "task_prompt": "Write unit tests for module_a.py."},
        {"node_id": "test-module-b", "depends_on": [], "task_prompt": "Write unit tests for module_b.py."},
        {"node_id": "test-module-c", "depends_on": [], "task_prompt": "Write unit tests for module_c.py."},
    ],
    "edges": [],
}


class GenericPlanningPromptTest(PlanWorkGraphTestBase):
    def test_planner_still_gets_fixed_doc_chain_template(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(
                stdout_obj={"is_error": False, "result": _fenced_result(_MANIFEST)}
            )
            work_graph_plan.plan_work_graph(self.conn, "planner", "task-1", "goal", "/tmp/proj-1")
        prompt = run.call_args.args[0][run.call_args.args[0].index("-p") + 1]
        self.assertIn("BRD, PRD, SDD, Benchmark", prompt)
        self.assertNotIn("design-dynamic-work-graph", prompt)

    def test_coder_gets_generic_template_referencing_the_skill(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(
                stdout_obj={"is_error": False, "result": _fenced_result(_CODER_THREE_MODULE_MANIFEST)}
            )
            work_graph_plan.plan_work_graph(self.conn, "coder", "task-1", "goal", "/tmp/proj-1")
        prompt = run.call_args.args[0][run.call_args.args[0].index("-p") + 1]
        self.assertIn("design-dynamic-work-graph", prompt)
        self.assertNotIn("BRD, PRD, SDD, Benchmark", prompt)
        self.assertIn("task_prompt", prompt)


class CoderThreeIndependentModulesScenarioTest(PlanWorkGraphTestBase):
    """E acceptance criterion 1: a Coder scenario that naturally splits into
    2-3 independent sub-tasks produces a graph_manifest via the new generic
    path (not the old single trigger_org_node call), correctly marking those
    nodes as mutually independent."""

    def test_three_independent_modules_produce_correctly_shaped_manifest(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(
                stdout_obj={"is_error": False, "result": _fenced_result(_CODER_THREE_MODULE_MANIFEST)}
            )
            result = work_graph_plan.plan_work_graph(self.conn, "coder", "task-1", "goal", "/tmp/proj-1")

        manifest = result["graph_manifest"]
        self.assertEqual(len(manifest["nodes"]), 3)
        for node in manifest["nodes"]:
            self.assertEqual(node["depends_on"], [])
            self.assertIn("task_prompt", node)

        # E's whole point ties directly into G: independence is expressed
        # structurally (empty depends_on), so topological_layers() must put
        # all three in ONE concurrent layer, not three serial ones.
        layers = work_graph_orchestrate.topological_layers(manifest)
        self.assertEqual(len(layers), 1)
        self.assertEqual({n["node_id"] for n in layers[0]}, {"test-module-a", "test-module-b", "test-module-c"})


class GenericNodePromptBuilderTest(unittest.TestCase):
    def test_returns_the_node_own_task_prompt(self):
        node = {"node_id": "test-module-a", "task_prompt": "Write unit tests for module_a.py."}
        self.assertEqual(
            work_graph_plan.generic_node_prompt_builder(node, {}), "Write unit tests for module_a.py."
        )


class GenericPathEndToEndViaRunGraphTest(PlanWorkGraphTestBase):
    """E acceptance criterion (FR-5): the generated graph_manifest, driven by
    generic_node_prompt_builder, actually runs to completion through
    run_graph() -- not just a standalone parsing/shape check."""

    def test_three_independent_nodes_execute_and_return_outputs(self):
        with patch("src.work_graph_plan.subprocess.run") as run:
            run.return_value = _completed_process(
                stdout_obj={"is_error": False, "result": _fenced_result(_CODER_THREE_MODULE_MANIFEST)}
            )
            plan_result = work_graph_plan.plan_work_graph(self.conn, "coder", "task-1", "goal", "/tmp/proj-1")

        # plan_work_graph() already inserted the work_graph_run row
        # (status='planning') -- run_graph() doesn't require any particular
        # status on it, so nothing more to set up here.
        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: {"content": "done: " + name}
            outputs = work_graph_orchestrate.run_graph(
                self.conn, plan_result["run_id"], "coder", "coder-implementation", "/tmp/proj-1",
                plan_result["graph_manifest"], work_graph_plan.generic_node_prompt_builder,
                {"type": "object", "required": ["node_id"]}, {"type": "object", "required": ["content"]},
            )

        self.assertEqual(set(outputs.keys()), {"test-module-a", "test-module-b", "test-module-c"})
        self.assertEqual(mock_execute.call_count, 3)


if __name__ == "__main__":
    unittest.main()