test_work_graph_plan.py 6.47 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_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):
        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', '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 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)


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