TJWaterServerBinary/tests/unit/test_scada_cleaning.py

import asyncio
from datetime import datetime, timezone
from unittest.mock import AsyncMock

import pandas as pd
import pytest

from app.api.v1.endpoints import project_data
from app.infra.db.timescaledb import composite_queries


PROJECT_SCADA = {
    "id": "fengyang-pressure-1",
    "type": "pressure",
    "associated_element_id": "J1",
    "api_query_id": "query-1",
    "transmission_mode": "realtime",
    "transmission_frequency": None,
    "reliability": 1.0,
    "x": 117.1,
    "y": 32.9,
}
START_TIME = datetime(2026, 6, 1, tzinfo=timezone.utc)
END_TIME = datetime(2026, 6, 2, tzinfo=timezone.utc)


def _cleaning_args():
    return (
        object(),
        object(),
        ["fengyang-pressure-1"],
        datetime(2026, 6, 1, tzinfo=timezone.utc),
        datetime(2026, 6, 2, tzinfo=timezone.utc),
    )


def test_clean_scada_uses_current_project_metadata(monkeypatch):
    monkeypatch.setattr(
        composite_queries.ScadaInfoRepository,
        "get_scadas",
        AsyncMock(
            return_value=[{"id": "fengyang-pressure-1", "type": "pressure"}]
        ),
    )
    monkeypatch.setattr(
        composite_queries.ScadaRepository,
        "get_scada_field_by_id_time_range",
        AsyncMock(
            return_value={
                "fengyang-pressure-1": [
                    {"time": "2026-06-01T00:00:00+08:00", "value": 26.5}
                ]
            }
        ),
    )
    update_mock = AsyncMock()
    monkeypatch.setattr(
        composite_queries.ScadaRepository,
        "update_scada_field",
        update_mock,
    )
    monkeypatch.setattr(
        composite_queries,
        "clean_pressure_data_df_km",
        lambda frame: pd.DataFrame(
            {
                "time": frame["time"],
                "fengyang-pressure-1": frame["fengyang-pressure-1"],
            }
        ),
    )

    result = asyncio.run(
        composite_queries.CompositeQueries.clean_scada_data(*_cleaning_args())
    )

    assert result == "success"
    update_mock.assert_awaited_once()


def test_clean_scada_rejects_devices_missing_from_project_metadata(monkeypatch):
    monkeypatch.setattr(
        composite_queries.ScadaInfoRepository,
        "get_scadas",
        AsyncMock(return_value=[{"id": "other-device", "type": "pressure"}]),
    )
    query_mock = AsyncMock()
    monkeypatch.setattr(
        composite_queries.ScadaRepository,
        "get_scada_field_by_id_time_range",
        query_mock,
    )

    with pytest.raises(ValueError, match="缺少元数据"):
        asyncio.run(
            composite_queries.CompositeQueries.clean_scada_data(*_cleaning_args())
        )

    query_mock.assert_not_awaited()


def test_clean_scada_rejects_zero_database_updates(monkeypatch):
    monkeypatch.setattr(
        composite_queries.ScadaInfoRepository,
        "get_scadas",
        AsyncMock(
            return_value=[{"id": "fengyang-pressure-1", "type": "pressure"}]
        ),
    )
    monkeypatch.setattr(
        composite_queries.ScadaRepository,
        "get_scada_field_by_id_time_range",
        AsyncMock(
            return_value={
                "fengyang-pressure-1": [
                    {"time": "2026-06-01T00:00:00+08:00", "value": 26.5}
                ]
            }
        ),
    )
    update_mock = AsyncMock()
    monkeypatch.setattr(
        composite_queries.ScadaRepository,
        "update_scada_field",
        update_mock,
    )
    monkeypatch.setattr(
        composite_queries,
        "clean_pressure_data_df_km",
        lambda _frame: pd.DataFrame(
            {"time": [], "fengyang-pressure-1": []}
        ),
    )

    with pytest.raises(ValueError, match="未产生任何数据库更新"):
        asyncio.run(
            composite_queries.CompositeQueries.clean_scada_data(*_cleaning_args())
        )

    update_mock.assert_not_awaited()


def test_clean_scada_propagates_write_failures(monkeypatch):
    monkeypatch.setattr(
        composite_queries.ScadaInfoRepository,
        "get_scadas",
        AsyncMock(
            return_value=[{"id": "fengyang-pressure-1", "type": "pressure"}]
        ),
    )
    monkeypatch.setattr(
        composite_queries.ScadaRepository,
        "get_scada_field_by_id_time_range",
        AsyncMock(
            return_value={
                "fengyang-pressure-1": [
                    {"time": "2026-06-01T00:00:00+08:00", "value": 26.5}
                ]
            }
        ),
    )
    monkeypatch.setattr(
        composite_queries.ScadaRepository,
        "update_scada_field",
        AsyncMock(side_effect=RuntimeError("database write failed")),
    )
    monkeypatch.setattr(
        composite_queries,
        "clean_pressure_data_df_km",
        lambda frame: frame,
    )

    with pytest.raises(RuntimeError, match="database write failed"):
        asyncio.run(
            composite_queries.CompositeQueries.clean_scada_data(*_cleaning_args())
        )


def _patch_project_scada_metadata(monkeypatch):
    monkeypatch.setattr(
        composite_queries.ScadaInfoRepository,
        "get_scadas",
        AsyncMock(return_value=[PROJECT_SCADA.copy()]),
    )


def test_realtime_scada_simulation_uses_current_project_metadata(monkeypatch):
    _patch_project_scada_metadata(monkeypatch)
    query_mock = AsyncMock(return_value=[{"time": START_TIME, "value": 26.5}])
    monkeypatch.setattr(
        composite_queries.RealtimeRepository,
        "get_node_field_by_time_range",
        query_mock,
        raising=False,
    )

    result = asyncio.run(
        composite_queries.CompositeQueries.get_scada_associated_realtime_simulation_data(
            object(), object(), [PROJECT_SCADA["id"]], START_TIME, END_TIME
        )
    )

    assert result[PROJECT_SCADA["id"]][0]["scada_id"] == PROJECT_SCADA["id"]
    assert query_mock.await_args.args[1:] == (
        START_TIME,
        END_TIME,
        "J1",
        "pressure",
    )


def test_scheme_scada_simulation_uses_current_project_metadata(monkeypatch):
    _patch_project_scada_metadata(monkeypatch)
    query_mock = AsyncMock(return_value=[{"time": START_TIME, "value": 26.5}])
    monkeypatch.setattr(
        composite_queries.SchemeRepository,
        "get_node_field_by_scheme_and_time_range",
        query_mock,
        raising=False,
    )

    result = asyncio.run(
        composite_queries.CompositeQueries.get_scada_associated_scheme_simulation_data(
            object(),
            object(),
            [PROJECT_SCADA["id"]],
            START_TIME,
            END_TIME,
            "baseline",
            "scheme-1",
        )
    )

    assert result[PROJECT_SCADA["id"]][0]["scada_id"] == PROJECT_SCADA["id"]
    assert query_mock.await_args.args[5:] == ("J1", "pressure")


def test_element_scada_query_uses_current_project_metadata(monkeypatch):
    _patch_project_scada_metadata(monkeypatch)
    query_mock = AsyncMock(
        return_value={PROJECT_SCADA["id"]: [{"time": START_TIME, "value": 26.5}]}
    )
    monkeypatch.setattr(
        composite_queries.ScadaRepository,
        "get_scada_field_by_id_time_range",
        query_mock,
    )

    result = asyncio.run(
        composite_queries.CompositeQueries.get_element_associated_scada_data(
            object(), object(), "J1", START_TIME, END_TIME
        )
    )

    assert result == {"J1": [{"time": START_TIME, "value": 26.5}]}


def test_scada_info_endpoint_uses_current_project_connection(monkeypatch):
    monkeypatch.setattr(
        project_data.ScadaInfoRepository,
        "get_scadas",
        AsyncMock(return_value=[PROJECT_SCADA.copy()]),
    )

    result = asyncio.run(project_data.get_scada_info_with_connection(object()))

    assert result == {"success": True, "data": [PROJECT_SCADA], "count": 1}