Refactor/batch transfer in graph and compile

tests/compile/test_batch_transfer_protocol.py

@@ -0,0 +1,296 @@
+"""
+批量转运编译器测试
+
+覆盖：单物料退化、刚好一批、多批次、空操作、AGV 配置发现、children dict 状态。
+"""
+
+import pytest
+import networkx as nx
+
+from unilabos.compile.batch_transfer_protocol import generate_batch_transfer_protocol
+from unilabos.compile.agv_transfer_protocol import generate_agv_transfer_protocol
+from unilabos.compile._agv_utils import find_agv_config, get_agv_capacity, split_batches
+
+
+# ============ 构建测试用设备图 ============
+
+def _make_graph(capacity_x=2, capacity_y=1, capacity_z=1):
+    """构建包含 AGV 节点的测试设备图"""
+    G = nx.DiGraph()
+
+    # AGV 节点
+    G.add_node("AGV", **{
+        "type": "device",
+        "class_": "agv_transport_station",
+        "config": {
+            "protocol_type": ["AGVTransferProtocol", "BatchTransferProtocol"],
+            "device_roles": {
+                "navigator": "zhixing_agv",
+                "arm": "zhixing_ur_arm"
+            },
+            "route_table": {
+                "StationA->StationB": {
+                    "nav_command": '{"target": "LM1"}',
+                    "arm_pick": '{"task_name": "pick.urp"}',
+                    "arm_place": '{"task_name": "place.urp"}'
+                },
+                "AGV->StationA": {
+                    "nav_command": '{"target": "LM1"}',
+                    "arm_pick": '{"task_name": "pick.urp"}',
+                    "arm_place": '{"task_name": "place.urp"}'
+                },
+                "StationA->StationA": {
+                    "nav_command": '{"target": "LM1"}',
+                    "arm_pick": '{"task_name": "pick.urp"}',
+                    "arm_place": '{"task_name": "place.urp"}'
+                },
+            }
+        }
+    })
+
+    # AGV 子设备
+    G.add_node("zhixing_agv", type="device", class_="zhixing_agv")
+    G.add_node("zhixing_ur_arm", type="device", class_="zhixing_ur_arm")
+    G.add_edge("AGV", "zhixing_agv")
+    G.add_edge("AGV", "zhixing_ur_arm")
+
+    # AGV Warehouse 子资源
+    G.add_node("agv_platform", **{
+        "type": "warehouse",
+        "config": {
+            "name": "agv_platform",
+            "num_items_x": capacity_x,
+            "num_items_y": capacity_y,
+            "num_items_z": capacity_z,
+        }
+    })
+    G.add_edge("AGV", "agv_platform")
+
+    # 来源/目标工站
+    G.add_node("StationA", type="device", class_="workstation")
+    G.add_node("StationB", type="device", class_="workstation")
+
+    return G
+
+
+def _make_repos(items_count=2):
+    """构建测试用的 from_repo 和 to_repo dict"""
+    children = {}
+    for i in range(items_count):
+        pos = f"A{i + 1:02d}"
+        children[pos] = {
+            "id": f"resource_{i + 1}",
+            "name": f"R{i + 1}",
+            "parent": "StationA",
+            "type": "resource",
+        }
+
+    from_repo = {
+        "StationA": {
+            "id": "StationA",
+            "name": "StationA",
+            "children": children,
+        }
+    }
+    to_repo = {
+        "StationB": {
+            "id": "StationB",
+            "name": "StationB",
+            "children": {},
+        }
+    }
+    return from_repo, to_repo
+
+
+def _make_items(count=2):
+    """构建 transfer_resources / from_positions / to_positions"""
+    resources = [
+        {
+            "id": f"resource_{i + 1}",
+            "name": f"R{i + 1}",
+            "sample_id": f"uuid-{i + 1}",
+            "parent": "StationA",
+            "type": "resource",
+        }
+        for i in range(count)
+    ]
+    from_positions = [f"A{i + 1:02d}" for i in range(count)]
+    to_positions = [f"A{i + 1:02d}" for i in range(count)]
+    return resources, from_positions, to_positions
+
+
+# ============ _agv_utils 测试 ============
+
+class TestAGVUtils:
+    def test_find_agv_config(self):
+        G = _make_graph()
+        cfg = find_agv_config(G)
+        assert cfg["agv_id"] == "AGV"
+        assert cfg["device_roles"]["navigator"] == "zhixing_agv"
+        assert cfg["device_roles"]["arm"] == "zhixing_ur_arm"
+        assert "StationA->StationB" in cfg["route_table"]
+
+    def test_find_agv_config_by_id(self):
+        G = _make_graph()
+        cfg = find_agv_config(G, agv_id="AGV")
+        assert cfg["agv_id"] == "AGV"
+
+    def test_find_agv_config_not_found(self):
+        G = nx.DiGraph()
+        G.add_node("SomeDevice", type="device", class_="pump")
+        with pytest.raises(ValueError, match="未找到 AGV"):
+            find_agv_config(G)
+
+    def test_get_agv_capacity(self):
+        G = _make_graph(capacity_x=2, capacity_y=1, capacity_z=1)
+        assert get_agv_capacity(G, "AGV") == 2
+
+    def test_get_agv_capacity_multi_layer(self):
+        G = _make_graph(capacity_x=1, capacity_y=2, capacity_z=3)
+        assert get_agv_capacity(G, "AGV") == 6
+
+    def test_split_batches_exact(self):
+        assert split_batches([1, 2], 2) == [[1, 2]]
+
+    def test_split_batches_overflow(self):
+        assert split_batches([1, 2, 3], 2) == [[1, 2], [3]]
+
+    def test_split_batches_single(self):
+        assert split_batches([1], 4) == [[1]]
+
+    def test_split_batches_zero_capacity(self):
+        with pytest.raises(ValueError):
+            split_batches([1], 0)
+
+
+# ============ 批量转运编译器测试 ============
+
+class TestBatchTransferProtocol:
+    def test_empty_items(self):
+        """空物料列表返回空 steps"""
+        G = _make_graph()
+        from_repo, to_repo = _make_repos(0)
+        steps = generate_batch_transfer_protocol(G, from_repo, to_repo, [], [], [])
+        assert steps == []
+
+    def test_single_item(self):
+        """单物料转运（BatchTransfer 退化为单物料）"""
+        G = _make_graph(capacity_x=2)
+        from_repo, to_repo = _make_repos(1)
+        resources, from_pos, to_pos = _make_items(1)
+        steps = generate_batch_transfer_protocol(G, from_repo, to_repo, resources, from_pos, to_pos)
+
+        # 应该有: nav到来源 + 1个pick + nav到目标 + 1个place = 4 steps
+        assert len(steps) == 4
+        assert steps[0]["action_name"] == "send_nav_task"
+        assert steps[1]["action_name"] == "move_pos_task"
+        assert steps[1]["_transfer_meta"]["phase"] == "pick"
+        assert steps[2]["action_name"] == "send_nav_task"
+        assert steps[3]["action_name"] == "move_pos_task"
+        assert steps[3]["_transfer_meta"]["phase"] == "place"
+
+    def test_exact_capacity(self):
+        """物料数 = AGV 容量，刚好一批"""
+        G = _make_graph(capacity_x=2)
+        from_repo, to_repo = _make_repos(2)
+        resources, from_pos, to_pos = _make_items(2)
+        steps = generate_batch_transfer_protocol(G, from_repo, to_repo, resources, from_pos, to_pos)
+
+        # nav + 2 pick + nav + 2 place = 6 steps
+        assert len(steps) == 6
+        pick_steps = [s for s in steps if s.get("_transfer_meta", {}).get("phase") == "pick"]
+        place_steps = [s for s in steps if s.get("_transfer_meta", {}).get("phase") == "place"]
+        assert len(pick_steps) == 2
+        assert len(place_steps) == 2
+
+    def test_multi_batch(self):
+        """物料数 > AGV 容量，自动分批"""
+        G = _make_graph(capacity_x=2)
+        from_repo, to_repo = _make_repos(3)
+        resources, from_pos, to_pos = _make_items(3)
+        steps = generate_batch_transfer_protocol(G, from_repo, to_repo, resources, from_pos, to_pos)
+
+        # 批次1: nav + 2 pick + nav + 2 place + nav(返回) = 7
+        # 批次2: nav + 1 pick + nav + 1 place = 4
+        # 总计 11 steps
+        assert len(steps) == 11
+
+        nav_steps = [s for s in steps if s["action_name"] == "send_nav_task"]
+        # 批次1: 2 nav(去来源+去目标) + 1 nav(返回) + 批次2: 2 nav = 5 nav
+        assert len(nav_steps) == 5
+
+    def test_children_dict_updated(self):
+        """compile 阶段三方 children dict 状态正确"""
+        G = _make_graph(capacity_x=2)
+        from_repo, to_repo = _make_repos(2)
+        resources, from_pos, to_pos = _make_items(2)
+
+        assert "A01" in from_repo["StationA"]["children"]
+        assert "A02" in from_repo["StationA"]["children"]
+        assert len(to_repo["StationB"]["children"]) == 0
+
+        generate_batch_transfer_protocol(G, from_repo, to_repo, resources, from_pos, to_pos)
+
+        # compile 后 from_repo 的 children 应该被 pop 掉
+        assert "A01" not in from_repo["StationA"]["children"]
+        assert "A02" not in from_repo["StationA"]["children"]
+        # to_repo 应该有新物料
+        assert "A01" in to_repo["StationB"]["children"]
+        assert "A02" in to_repo["StationB"]["children"]
+        assert to_repo["StationB"]["children"]["A01"]["id"] == "resource_1"
+
+    def test_device_ids_from_config(self):
+        """设备 ID 全部从配置读取，不硬编码"""
+        G = _make_graph()
+        from_repo, to_repo = _make_repos(1)
+        resources, from_pos, to_pos = _make_items(1)
+        steps = generate_batch_transfer_protocol(G, from_repo, to_repo, resources, from_pos, to_pos)
+
+        device_ids = {s["device_id"] for s in steps}
+        assert "zhixing_agv" in device_ids
+        assert "zhixing_ur_arm" in device_ids
+
+    def test_route_not_found(self):
+        """路由表中无对应路线时报错"""
+        G = _make_graph()
+        from_repo = {"Unknown": {"id": "Unknown", "children": {"A01": {"id": "R1", "parent": "Unknown"}}}}
+        to_repo = {"Other": {"id": "Other", "children": {}}}
+        resources = [{"id": "R1", "name": "R1"}]
+        with pytest.raises(KeyError, match="路由表"):
+            generate_batch_transfer_protocol(G, from_repo, to_repo, resources, ["A01"], ["B01"])
+
+    def test_length_mismatch(self):
+        """三个数组长度不一致时报错"""
+        G = _make_graph()
+        from_repo, to_repo = _make_repos(2)
+        resources = [{"id": "R1"}]
+        with pytest.raises(ValueError, match="长度不一致"):
+            generate_batch_transfer_protocol(G, from_repo, to_repo, resources, ["A01", "A02"], ["B01"])
+
+
+# ============ 改造后的 AGV 单物料编译器测试 ============
+
+class TestAGVTransferProtocol:
+    def test_single_transfer_from_config(self):
+        """改造后的单物料编译器从 G 读取配置"""
+        G = _make_graph()
+        from_repo = {"StationA": {"id": "StationA", "children": {"A01": {"id": "R1", "parent": "StationA"}}}}
+        to_repo = {"StationB": {"id": "StationB", "children": {}}}
+        steps = generate_agv_transfer_protocol(G, from_repo, "A01", to_repo, "B01")
+
+        assert len(steps) == 2
+        assert steps[0]["device_id"] == "zhixing_agv"
+        assert steps[0]["action_name"] == "send_nav_task"
+        assert steps[1]["device_id"] == "zhixing_ur_arm"
+        assert steps[1]["action_name"] == "move_pos_task"
+
+    def test_children_updated(self):
+        """单物料编译后 children dict 正确更新"""
+        G = _make_graph()
+        from_repo = {"StationA": {"id": "StationA", "children": {"A01": {"id": "R1", "parent": "StationA"}}}}
+        to_repo = {"StationB": {"id": "StationB", "children": {}}}
+        generate_agv_transfer_protocol(G, from_repo, "A01", to_repo, "B01")
+
+        assert "A01" not in from_repo["StationA"]["children"]
+        assert "B01" in to_repo["StationB"]["children"]
+        assert to_repo["StationB"]["children"]["B01"]["parent"] == "StationB"