cad-agent/tests/test_planning.py

"""`dwg_review.phases.planning` 单元测试。"""

from dwg_review.phases.planning import _sort_frames, plan
from tests.fakes import FakeMcpClient, frame, make_ctx, ok_text


# ---------------------------------------------------------------------------
#  阅读顺序排序
# ---------------------------------------------------------------------------
def test_sort_frames_top_to_bottom_then_left_to_right():
    # bbox = [min_x, min_y, max_x, max_y]
    top_left = frame([0, 100, 10, 200])
    top_right = frame([50, 100, 60, 200])
    bottom_left = frame([0, 0, 10, 50])
    bottom_right = frame([50, 0, 60, 50])

    ordered = _sort_frames([bottom_right, top_left, bottom_left, top_right])
    assert ordered == [top_left, top_right, bottom_left, bottom_right]


def test_sort_frames_same_row_by_x():
    a = frame([90, 0, 100, 10])
    b = frame([10, 0, 20, 10])
    assert _sort_frames([a, b]) == [b, a]


def test_sort_frames_handles_missing_bbox():
    with_bbox = frame([0, 0, 10, 10])
    without = {"kind": "x"}  # 缺 bbox -> 视作全 0
    ordered = _sort_frames([without, with_bbox])
    assert with_bbox in ordered and without in ordered


# ---------------------------------------------------------------------------
#  plan()
# ---------------------------------------------------------------------------
def build_plan_client(frames):
    return FakeMcpClient(
        responses={
            "get_drawing_info": ok_text({"units": "mm"}),
            "get_layers": ok_text({"layers": [{"name": "0"}]}),
            "get_layouts": ok_text({"layouts": [{"name": "Model"}]}),
            "detect_drawing_frames": ok_text({"items": frames}),
        }
    )


def test_plan_returns_all_sections():
    frames = [frame([0, 0, 10, 10]), frame([0, 20, 10, 30])]
    ctx = make_ctx(client=build_plan_client(frames))
    result = plan(ctx)
    assert result["drawing_info"] == {"units": "mm"}
    assert result["layers"] == {"layers": [{"name": "0"}]}
    assert result["layouts"] == {"layouts": [{"name": "Model"}]}
    assert len(result["frames"]) == 2


def test_plan_sorts_frames():
    lower = frame([0, 0, 10, 10])
    upper = frame([0, 50, 10, 60])
    ctx = make_ctx(client=build_plan_client([lower, upper]))
    assert plan(ctx)["frames"] == [upper, lower]


def test_plan_respects_max_frames():
    frames = [frame([0, i * 10, 10, i * 10 + 5]) for i in range(5)]
    ctx = make_ctx(client=build_plan_client(frames), max_frames=2)
    assert len(plan(ctx)["frames"]) == 2


def test_plan_calls_with_expected_arguments():
    ctx = make_ctx(client=build_plan_client([]))
    plan(ctx)
    client = ctx.client
    assert client.arguments_of("get_layers") == {"count_entities": True}
    assert client.arguments_of("detect_drawing_frames") == {"min_confidence": 0.3}


def test_plan_tolerates_non_dict_frames_data():
    client = FakeMcpClient(
        responses={
            "get_drawing_info": ok_text({}),
            "get_layers": ok_text({}),
            "get_layouts": ok_text({}),
            "detect_drawing_frames": ok_text("oops"),  # 非 dict
        }
    )
    assert plan(make_ctx(client=client))["frames"] == []


def test_plan_tolerates_missing_items_key():
    client = FakeMcpClient(responses={"detect_drawing_frames": ok_text({"other": 1})})
    assert plan(make_ctx(client=client))["frames"] == []


# ---------------------------------------------------------------------------
#  文件源分发（PDF / 图片）
# ---------------------------------------------------------------------------
def test_plan_dispatches_file_source_to_sources(tmp_path, monkeypatch):
    """ctx.source 为文件源时走 `sources.plan_file`，**完全不碰 CAD 工具**。"""
    from PIL import Image

    from dwg_review import sources

    path = tmp_path / "tiny.png"
    Image.new("RGB", (60, 40), (255, 255, 255)).save(path)

    ctx = make_ctx(connect_cad=False, source="image", source_path=str(path))
    try:
        result = plan(ctx)
    finally:
        sources.close_reader(ctx)

    assert result["drawing_info"]["unit"] == "px"
    assert len(result["frames"]) == 1
    assert result["frames"][0]["source_type"] == "image"
    # 离线：没有任何 MCP 调用（ctx.tools 为 None 本身就是证据）
    assert ctx.tools is None


def test_plan_file_source_does_not_sort_frames(tmp_path, monkeypatch):
    """文件源**保持原始页序**（阅读顺序由文档自身决定，不能按 bbox 重排）。"""
    from dwg_review import sources

    ctx = make_ctx(connect_cad=False, source="pdf", source_path="unused.pdf", max_frames=3)

    class _Reader:
        page_count = 3

        def frames(self):
            # 故意让 bbox 顺序与页序相反（后页在上方）
            return [
                {"bbox": [0.0, 0.0, 10.0, 10.0], "page_index": 0, "source_type": "pdf"},
                {"bbox": [0.0, 50.0, 10.0, 60.0], "page_index": 1, "source_type": "pdf"},
                {"bbox": [0.0, 90.0, 10.0, 100.0], "page_index": 2, "source_type": "pdf"},
            ]

        def close(self):
            pass

    ctx.source_reader = _Reader()
    result = sources.plan_file(ctx)
    assert [f["page_index"] for f in result["frames"]] == [0, 1, 2]