TJWaterServerBinary/tests/unit/test_wndb_network_views.py

from app.native.wndb.gis import network_views, region_geometry
from app.native.wndb.model import elements


def test_network_node_coords_use_one_unified_view_query(monkeypatch) -> None:
    calls = []

    def fake_read_all(name, statement, params=None):
        calls.append((name, statement, params))
        return [
            {"id": "J-1", "x": 10.5, "y": 20.5, "node_type": "junction"},
            {"id": "R-1", "x": 30.0, "y": 40.0, "node_type": "reservoir"},
        ]

    monkeypatch.setattr(network_views, "read_all", fake_read_all)

    assert network_views.get_network_node_coords("project_a") == {
        "J-1": {"x": 10.5, "y": 20.5, "type": "junction"},
        "R-1": {"x": 30.0, "y": 40.0, "type": "reservoir"},
    }
    assert len(calls) == 1
    assert "FROM gis.network_nodes" in calls[0][1]


def test_network_link_nodes_use_one_unified_view_query(monkeypatch) -> None:
    calls = []

    def fake_read_all(name, statement, params=None):
        calls.append((name, statement, params))
        return [
            {
                "id": "P-1",
                "link_type": "pipe",
                "start_node_id": "J-1",
                "end_node_id": "J-2",
            },
            {
                "id": "PU-1",
                "link_type": "pump",
                "start_node_id": "R-1",
                "end_node_id": "J-1",
            },
        ]

    monkeypatch.setattr(network_views, "read_all", fake_read_all)

    assert network_views.get_network_link_nodes("project_a") == [
        "P-1:pipe:J-1:J-2",
        "PU-1:pump:R-1:J-1",
    ]
    assert len(calls) == 1
    assert "FROM gis.network_links" in calls[0][1]


def test_isolation_snapshot_reads_authoritative_links_statuses_and_sources_once(monkeypatch):
    calls = []

    def fake_read_all(name, statement, params=None):
        calls.append((name, statement, params))
        return [
            {"id": "P", "element_type": "pipe", "start_node_id": "A", "end_node_id": "B", "status": None},
            {"id": "P-closed", "element_type": "pipe", "start_node_id": "A", "end_node_id": "R", "status": "CLOSED"},
            {"id": "V", "element_type": "valve", "start_node_id": "B", "end_node_id": "R", "status": "CLOSED"},
            {"id": "PU", "element_type": "pump", "start_node_id": "T", "end_node_id": "A", "status": "CLOSED"},
            {"id": "R", "element_type": "reservoir", "start_node_id": None, "end_node_id": None, "status": None},
            {"id": "T", "element_type": "tank", "start_node_id": None, "end_node_id": None, "status": None},
        ]

    monkeypatch.setattr(network_views, "read_all", fake_read_all)

    assert network_views.get_valve_isolation_topology("project_a") == (
        ["P:pipe:A:B", "P-closed:pipe:A:R", "V:valve:B:R", "PU:pump:T:A"],
        ["R", "T"], ["P-closed", "V", "PU"],
    )
    assert len(calls) == 1
    name, statement, _params = calls[0]
    assert name == "project_a"
    assert "FROM network.links" in statement
    assert "network.link_initial_settings" in statement
    assert "COALESCE(s.status::text, p.status::text)" in statement
    assert "FROM network.nodes" in statement
    assert "UNION ALL" in statement


def test_topology_rows_are_loaded_in_two_batch_queries(monkeypatch) -> None:
    calls = []
    responses = [
        [{"id": "J-1", "x": 1.0, "y": 2.0, "node_type": "junction"}],
        [
            {
                "id": "P-1",
                "start_node_id": "J-1",
                "end_node_id": "J-2",
                "length": 12.5,
            }
        ],
    ]

    def fake_read_all(name, statement, params=None):
        calls.append((name, statement, params))
        return responses[len(calls) - 1]

    monkeypatch.setattr(network_views, "read_all", fake_read_all)

    nodes, links = network_views.get_topology_rows("project_a", ["J-1", "J-2"])

    assert nodes == responses[0]
    assert links == responses[1]
    assert len(calls) == 2
    assert "FROM network.nodes" in calls[0][1]
    assert "FROM network.links" in calls[1][1]
    assert "LEFT JOIN network.pipes" in calls[1][1]


def test_topology_builds_adjacency_from_batch_rows(monkeypatch) -> None:
    monkeypatch.setattr(
        region_geometry,
        "get_topology_rows",
        lambda _name, _node_ids: (
            [
                {"id": "J-1", "x": 1.0, "y": 2.0, "node_type": "junction"},
                {"id": "R-1", "x": 3.0, "y": 4.0, "node_type": "reservoir"},
            ],
            [
                {
                    "id": "P-1",
                    "start_node_id": "J-1",
                    "end_node_id": "R-1",
                    "length": 20.0,
                }
            ],
        ),
    )

    topology = region_geometry.Topology("project_a", ["J-1", "R-1"])

    assert topology.max_x_node() == "R-1"
    assert topology.nodes()["J-1"] == {
        "x": 1.0,
        "y": 2.0,
        "type": "junction",
        "links": ["P-1"],
    }
    assert topology.links()["P-1"] == {
        "node1": "J-1",
        "node2": "R-1",
        "length": 20.0,
    }


def test_junction_demands_are_mapped_from_authoritative_table(monkeypatch) -> None:
    monkeypatch.setattr(
        network_views,
        "read_all",
        lambda _name, _statement, _params: [
            {
                "junction_id": "J-1",
                "sequence_no": 0,
                "base_demand": 3.5,
                "pattern_id": "PAT-1",
                "category": None,
            }
        ],
    )

    assert network_views.get_junction_demands("project_a", ["J-1"]) == {
        "J-1": [
            {"demand": 3.5, "pattern": "PAT-1", "category": None}
        ]
    }


def test_all_node_links_scan_the_unified_view_once(monkeypatch) -> None:
    calls = []

    def fake_read_all_typed(name, statement, params):
        calls.append((name, statement, params))
        return [
            {
                "id": "P-1",
                "start_node_id": "J-1",
                "end_node_id": "J-2",
            }
        ]

    monkeypatch.setattr(elements, "read_all_typed", fake_read_all_typed)

    assert elements.get_all_node_links("project_a") == {
        "J-1": ["P-1"],
        "J-2": ["P-1"],
    }
    assert len(calls) == 1
    assert "FROM gis.network_links" in calls[0][1]