TJWaterServerBinary/tests/unit/test_sensor_placement_service.py

from datetime import datetime, timezone
from unittest.mock import MagicMock
from uuid import uuid4

import pytest
from openpyxl import load_workbook

from app.infra.db.postgresql import sensor_placement as sensor_placement_repository
from app.services import sensor_placement


def _mock_project_cursor(monkeypatch):
    cursor = MagicMock()
    connection = MagicMock()
    connection.cursor.return_value.__enter__.return_value = cursor
    connection_context = MagicMock()
    connection_context.__enter__.return_value = connection
    monkeypatch.setattr(
        sensor_placement_repository,
        "project_connection",
        lambda _network: connection_context,
    )
    return cursor


def _run() -> dict:
    return {
        "run_id": uuid4(),
        "name": "北区测压点",
        "sensor_count": 2,
        "min_diameter": 300,
        "created_by": "alice",
        "created_at": datetime(2026, 7, 30, tzinfo=timezone.utc),
        "status": "completed",
        "sensor_locations": ["J1", "J2"],
    }


def _point() -> dict:
    return {
        "node_id": "J1",
        "max_pipe_diameter": 400.0,
        "project_x": 3038.94,
        "project_y": -34446.59,
        "map_x": 13525191.530279,
        "map_y": 3622984.760237,
        "longitude": 121.498863,
        "latitude": 30.924784,
        "elevation": 4.5,
    }


def test_build_workbook_uses_analysis_run_metadata(monkeypatch):
    monkeypatch.setattr(sensor_placement, "_sensor_points", lambda *_: [_point()])
    output = sensor_placement.build_sensor_placement_workbook(
        network="tjwater",
        scheme=_run(),
        sensor_location=["J1"],
        adjustment_status={"J1": "replaced"},
    )
    workbook = load_workbook(output)

    assert workbook.sheetnames == ["方案信息", "监测点清单"]
    assert workbook["监测点清单"]["J2"].value == "替换"
    assert workbook["方案信息"]["B8"].value == "未保存草稿"


def test_candidate_transforms_map_coordinates(monkeypatch):
    row = _point().copy()
    row.pop("longitude")
    row.pop("latitude")
    monkeypatch.setattr(
        sensor_placement.sensor_placement_repository,
        "get_sensor_placement_nodes",
        lambda network, node_ids: [row],
    )

    point = sensor_placement.get_sensor_placement_candidate("tjwater", "J1")

    assert point["project_x"] == 3038.94
    assert point["longitude"] == pytest.approx(121.498863, abs=1e-6)


def test_update_validates_nodes_before_write(monkeypatch):
    monkeypatch.setattr(
        sensor_placement.sensor_placement_repository,
        "get_sensor_placement_nodes",
        lambda network, node_ids: [],
    )

    with pytest.raises(sensor_placement.SensorPlacementValidationError, match="missing"):
        sensor_placement.update_sensor_placement_run(
            "tjwater",
            uuid4(),
            expected_sensor_locations=["J1"],
            sensor_locations=["missing"],
        )


def test_list_sensor_placements_batches_node_lookup(monkeypatch):
    first = _run()
    second = {**_run(), "run_id": uuid4(), "sensor_locations": ["J2", "J3"]}
    lookup_calls = []
    monkeypatch.setattr(
        sensor_placement.sensor_placement_repository,
        "get_all_sensor_placements",
        lambda _network: [first, second],
    )
    monkeypatch.setattr(
        sensor_placement.sensor_placement_repository,
        "get_sensor_placement_nodes",
        lambda _network, node_ids: lookup_calls.append(node_ids)
        or [
            {**_point(), "node_id": node_id}
            for node_id in node_ids
        ],
    )

    runs = sensor_placement.list_sensor_placement_runs("tjwater_v2")

    assert lookup_calls == [["J1", "J2", "J3"]]
    assert [len(run["sensor_points"]) for run in runs] == [2, 2]


def test_sensor_nodes_query_uses_new_network_and_gis_schemas(monkeypatch):
    cursor = _mock_project_cursor(monkeypatch)
    cursor.fetchall.return_value = []

    sensor_placement_repository.get_sensor_placement_nodes("tjwater", ["J1"])

    query = cursor.execute.call_args.args[0]
    assert "network.pipes" in query
    assert "network.links" in query
    assert "gis.node_geometries" in query
    assert "ST_Transform(g.geom, 3857)" in query
    assert "l.start_node_id = ANY(%s)" in query
    assert "l.end_node_id = ANY(%s)" in query
    assert "CROSS JOIN LATERAL" not in query
    assert cursor.execute.call_args.args[1] == (["J1"], ["J1"], ["J1"])


def test_create_sensor_placement_writes_run_and_result_atomically(monkeypatch):
    cursor = _mock_project_cursor(monkeypatch)
    created_at = datetime(2026, 8, 24, tzinfo=timezone.utc)
    cursor.fetchone.return_value = {
        "run_id": uuid4(),
        "name": "北区测压点",
        "created_by": "alice",
        "created_at": created_at,
        "status": "completed",
    }

    created = sensor_placement_repository.create_sensor_placement(
        "tjwater",
        run_name="北区测压点",
        min_diameter=300,
        created_by="alice",
        sensor_locations=["J1", "J2"],
    )

    assert created["sensor_count"] == 2
    statements = [call.args[0] for call in cursor.execute.call_args_list]
    assert "INSERT INTO analysis.runs" in statements[0]
    assert "INSERT INTO analysis.results" in statements[1]