cad-Integration/CADDesigner-Code-main/test/test_codegen_event_bridge.py

import os
import sys
from datetime import datetime, timezone
from typing import Any, cast

import pytest


PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)
WORKSPACE_ROOT = os.path.join(PROJECT_ROOT, "workspace")


from SimpleLLMFunc.hooks.events import (
    ReActEventType,
    ToolCallEndEvent,
    ToolCallErrorEvent,
    ToolCallStartEvent,
)
from SimpleLLMFunc.hooks.stream import EventYield, ResponseYield

import tools.code_tools as code_tools_module
from tools.code_tools import cad_code_generator


class _FakeEmitter:
    def __init__(self):
        self.events: list[tuple[str, dict]] = []

    async def emit(self, event_name: str, data):
        self.events.append((event_name, data))


@pytest.mark.asyncio
async def test_cad_code_generator_bridges_nested_specialist_events(monkeypatch):
    monkeypatch.chdir(WORKSPACE_ROOT)

    async def fake_specialist(**kwargs):
        assert "Current working directory:" in kwargs["message"]
        assert "Skill root: use the preferred skill root below." in kwargs["message"]
        assert "validation_command: uv run python part/model.py" in kwargs["message"]
        assert "references/docs/api/README.md" in kwargs["message"]
        yield EventYield(
            event=ToolCallStartEvent(
                event_type=ReActEventType.TOOL_CALL_START,
                timestamp=datetime.now(timezone.utc),
                trace_id="trace-1",
                func_name="cad_code_generator_specialist",
                iteration=0,
                tool_name="read_file",
                tool_call_id="nested-1",
                arguments={"file_path": "part/model.py"},
                tool_call=cast(
                    Any,
                    {
                        "id": "nested-1",
                        "type": "function",
                        "function": {"name": "read_file", "arguments": "{}"},
                    },
                ),
            )
        )
        yield EventYield(
            event=ToolCallEndEvent(
                event_type=ReActEventType.TOOL_CALL_END,
                timestamp=datetime.now(timezone.utc),
                trace_id="trace-1",
                func_name="cad_code_generator_specialist",
                iteration=0,
                tool_name="read_file",
                tool_call_id="nested-1",
                arguments={"file_path": "part/model.py"},
                result="print('old')",
                execution_time=0.05,
                success=True,
            )
        )
        yield EventYield(
            event=ToolCallErrorEvent(
                event_type=ReActEventType.TOOL_CALL_ERROR,
                timestamp=datetime.now(timezone.utc),
                trace_id="trace-1",
                func_name="cad_code_generator_specialist",
                iteration=0,
                tool_name="execute_command",
                tool_call_id="nested-2",
                arguments={"command": "python model.py"},
                error=RuntimeError("boom"),
                error_message="boom",
                error_type="RuntimeError",
                execution_time=0.12,
            )
        )
        yield ResponseYield(response="STATUS: SUCCESS\nSUMMARY: ok", messages=[])

    monkeypatch.setattr(
        code_tools_module,
        "cad_code_generator_specialist",
        fake_specialist,
    )
    monkeypatch.setattr(
        code_tools_module, "_read_latest_code", lambda path: "print('ok')\n"
    )
    emitter = _FakeEmitter()
    result = await cad_code_generator(
        task="Create a cube as a new file. This is a create-new-file task.",
        target_file_path="part/model.py",
        event_emitter=emitter,
    )

    event_names = [event_name for event_name, _ in emitter.events]
    assert event_names == [
        "subagent_status",
        "subagent_tool_start",
        "subagent_tool_end",
        "subagent_tool_error",
        "subagent_response",
        "subagent_status",
    ]
    first_status_payload = emitter.events[0][1]
    assert first_status_payload["subagent_label"] == "CAD Code Specialist"
    assert first_status_payload["validation_command"].startswith(
        "uv run python part/model.py"
    )
    assert "ls part/*.stl" in first_status_payload["validation_command"]
    assert (
        "(ls part/*.step || ls part/*.stp)"
        in first_status_payload["validation_command"]
    )
    assert "STATUS: SUCCESS" in result
    assert "Latest code" in result
    response_payload = emitter.events[4][1]
    assert response_payload["delta_text"] == "STATUS: SUCCESS\nSUMMARY: ok"


@pytest.mark.asyncio
async def test_cad_code_generator_retries_when_first_attempt_does_not_produce_code(
    monkeypatch,
):
    monkeypatch.chdir(WORKSPACE_ROOT)
    calls = {"count": 0}

    async def fake_specialist(**kwargs):
        calls["count"] += 1
        if calls["count"] == 1:
            assert kwargs["history"] == []
            assert "Current working directory:" in kwargs["message"]
            assert (
                "validation_command: uv run python part/model.py" in kwargs["message"]
            )
            assert "references/docs/api/README.md" in kwargs["message"]
            yield ResponseYield(
                response=(
                    "STATUS: SUCCESS\n"
                    "SUMMARY: planned but no file write\n"
                    "```python\n"
                    "print('draft only')\n"
                    "```"
                ),
                messages=[],
            )
            return

        assert "must write the required Python script directly" in kwargs["message"]
        assert "part/model.py" in kwargs["message"]
        assert kwargs["history"][0]["role"] == "user"
        assert "target_file: part/model.py" in kwargs["history"][0]["content"]
        assert kwargs["history"][1]["role"] == "assistant"
        assert "planned but no file write" in kwargs["history"][1]["content"]
        yield ResponseYield(
            response="STATUS: SUCCESS\nCODE_WRITTEN: YES\nSUMMARY: wrote code",
            messages=[],
        )

    monkeypatch.setattr(
        code_tools_module,
        "cad_code_generator_specialist",
        fake_specialist,
    )

    code_reads = {"count": 0}

    def fake_read_latest_code(path):
        code_reads["count"] += 1
        if code_reads["count"] == 1:
            return None
        return "print('ok')\n"

    monkeypatch.setattr(code_tools_module, "_read_latest_code", fake_read_latest_code)

    result = await cad_code_generator(
        task="Create a cube as a new file. This is a create-new-file task.",
        target_file_path="part/model.py",
        event_emitter=None,
    )

    assert calls["count"] == 2
    assert "CODE_WRITTEN: YES" in result
    assert "Latest code" in result