import pytest

from app.algorithms.dma_leakage_estimation.topology_partitioning import (
    build_dma_partitions,
)


def test_partition_assigns_nodes_to_nearest_connected_sensor():
    coordinates = {
        "A": {"x": 0.0, "y": 0.0},
        "B": {"x": 1.0, "y": 0.0},
        "C": {"x": 2.0, "y": 0.0},
        "D": {"x": 3.0, "y": 0.0},
    }
    links = ["P1:pipe:A:B", "P2:pipe:B:C", "P3:pipe:C:D"]

    area_map, areas = build_dma_partitions(
        ["A", "D"], coordinates, links, dma_count=2
    )

    assert area_map == {"A": "1", "B": "1", "C": "2", "D": "2"}
    assert [area["node_count"] for area in areas] == [2, 2]


def test_partition_rejects_more_areas_than_sensors():
    with pytest.raises(ValueError, match="传感器数量"):
        build_dma_partitions(
            ["A"],
            {"A": {"x": 0.0, "y": 0.0}, "B": {"x": 1.0, "y": 0.0}},
            ["P1:pipe:A:B"],
            dma_count=2,
        )


def test_partition_preserves_requested_area_count_for_coincident_sensors():
    coordinates = {
        "A": {"x": 0.0, "y": 0.0},
        "B": {"x": 0.0, "y": 0.0},
        "C": {"x": 0.0, "y": 0.0},
    }

    area_map, areas = build_dma_partitions(
        ["A", "B", "C"],
        coordinates,
        ["P1:pipe:A:B", "P2:pipe:B:C"],
        dma_count=2,
    )

    assert set(area_map.values()) == {"1", "2"}
    assert [area["area_id"] for area in areas] == ["1", "2"]