TJWaterServerBinary/tests/unit/test_scada_cleaning.py

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

import pandas as pd
import pytest
from fastapi import HTTPException

from app.api.v1.endpoints.timeseries import composite as composite_endpoint
from app.services import timeseries_analysis as composite_queries


class _FakeTimescaleConnection:
    @asynccontextmanager
    async def transaction(self):
        yield


def test_clean_scada_uses_current_project_metadata(monkeypatch):
    """Fengyang data must not be classified with the global tjwater metadata."""

    monkeypatch.setattr(
        composite_queries.ScadaInfoRepository,
        "get_scadas",
        AsyncMock(
            return_value=[{"device_id": "fengyang-pressure-1", "device_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()
    update_mock.return_value = 1
    monkeypatch.setattr(
        composite_queries.ScadaRepository,
        "update_scada_field_batch",
        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.TimeseriesAnalysisService.clean_scada_data(
            _FakeTimescaleConnection(),
            object(),
            ["fengyang-pressure-1"],
            datetime(2026, 6, 1, tzinfo=timezone.utc),
            datetime(2026, 6, 2, tzinfo=timezone.utc),
        )
    )

    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=[{"device_id": "other-device", "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.TimeseriesAnalysisService.clean_scada_data(
                _FakeTimescaleConnection(),
                _FakeTimescaleConnection(),
                ["fengyang-pressure-1"],
                datetime(2026, 6, 1, tzinfo=timezone.utc),
                datetime(2026, 6, 2, tzinfo=timezone.utc),
            )
        )

    query_mock.assert_not_awaited()


def test_clean_scada_rejects_zero_database_updates(monkeypatch):
    monkeypatch.setattr(
        composite_queries.ScadaInfoRepository,
        "get_scadas",
        AsyncMock(
            return_value=[{"device_id": "fengyang-pressure-1", "device_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_batch",
        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.TimeseriesAnalysisService.clean_scada_data(
                _FakeTimescaleConnection(),
                object(),
                ["fengyang-pressure-1"],
                datetime(2026, 6, 1, tzinfo=timezone.utc),
                datetime(2026, 6, 2, tzinfo=timezone.utc),
            )
        )

    update_mock.assert_not_awaited()


def test_clean_scada_propagates_write_failures(monkeypatch):
    monkeypatch.setattr(
        composite_queries.ScadaInfoRepository,
        "get_scadas",
        AsyncMock(
            return_value=[{"device_id": "fengyang-pressure-1", "device_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_batch",
        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.TimeseriesAnalysisService.clean_scada_data(
                _FakeTimescaleConnection(),
                object(),
                ["fengyang-pressure-1"],
                datetime(2026, 6, 1, tzinfo=timezone.utc),
                datetime(2026, 6, 2, tzinfo=timezone.utc),
            )
        )


def test_clean_scada_endpoint_returns_http_400_for_validation_error(monkeypatch):
    monkeypatch.setattr(
        composite_endpoint.TimeseriesAnalysisService,
        "clean_scada_data",
        AsyncMock(side_effect=ValueError("当前项目没有可清洗的 SCADA 设备")),
    )

    with pytest.raises(HTTPException) as exc_info:
        asyncio.run(
            composite_endpoint.clean_scada_data(
                device_ids="all",
                start_time=datetime(2026, 6, 1, tzinfo=timezone.utc),
                end_time=datetime(2026, 6, 2, tzinfo=timezone.utc),
                timescale_conn=object(),
                postgres_conn=object(),
            )
        )

    assert exc_info.value.status_code == 400
    assert exc_info.value.detail == "当前项目没有可清洗的 SCADA 设备"