Northern-Thailand-Ping-River-Monitor/tests/test_hii_collector.py

"""Tests for the HII/ThaiWater api-v3 collector (parsing + persistence)."""

import datetime

import pytest

from src.hii_collector import (
    HiiStore,
    parse_rain_records,
    parse_waterlevel_records,
    rid_code_from_oldcode,
)


def _rain_payload():
    return {
        "result": "OK",
        "data": [
            {
                "id": 306091240,
                "rain_1h": 0,
                "rain_24h": "49.6",
                "rainfall_datetime": "2026-08-11 13:00",
                "agency": {"agency_shortname": {"en": "HII", "th": "สสน."}},
                "basin": {"basin_code": 6, "basin_name": {"en": "Ping Basin"}},
                "station": {
                    "id": 418,
                    "sub_basin_id": "0604",
                    "tele_station_lat": 19.12207,
                    "tele_station_long": 98.94447,
                    "tele_station_name": {"en": "Chiang Mai 5", "th": "แม่แตง"},
                    "tele_station_oldcode": "CHM005",
                },
            },
            # Wrong basin -> filtered out
            {
                "id": 2,
                "rain_24h": 10,
                "rainfall_datetime": "2026-08-11 13:00",
                "basin": {"basin_code": 7},
                "station": {"id": 99},
            },
            # No timestamp -> skipped
            {
                "id": 3,
                "rain_24h": 5,
                "rainfall_datetime": None,
                "basin": {"basin_code": 6},
                "station": {"id": 100},
            },
        ],
    }


def _waterlevel_payload():
    return {
        "waterlevel_data": {
            "result": "OK",
            "data": [
                {
                    "id": 1286124160,
                    "waterlevel_datetime": "2026-08-11 13:00",
                    "waterlevel_m": None,
                    "waterlevel_msl": "303.27",
                    "discharge": "335.00",
                    "flow_rate": None,
                    "storage_percent": "81.21",
                    "situation_level": 4,
                    "diff_wl_bank": "0.93",
                    "river_name": "แม่น้ำปิง",
                    "agency": {"agency_shortname": {"en": "RID"}},
                    "basin": {"basin_code": 6},
                    "station": {
                        "id": 3226,
                        "tele_station_lat": 18.786961,
                        "tele_station_long": 99.005089,
                        "tele_station_name": {"th": "สะพานนวรัฐ"},
                        "tele_station_oldcode": "P.1",
                        "offset": 300.5,
                        "ground_level": 299.25,
                        "min_bank": 304.2,
                        "critical_level_msl": 304.2,
                        "critical_level_m": 3.7,
                        "qmax": 425,
                        "is_key_station": True,
                    },
                },
                {
                    "id": 2,
                    "waterlevel_datetime": "2026-08-11 13:00",
                    "waterlevel_msl": "200.0",
                    "basin": {"basin_code": 6},
                    "station": {"id": 4000, "tele_station_oldcode": "ridhydro_P.67"},
                },
                # Wrong basin -> filtered out
                {
                    "id": 3,
                    "waterlevel_datetime": "2026-08-11 13:00",
                    "basin": {"basin_code": 10},
                    "station": {"id": 5000},
                },
            ],
        }
    }


class TestRidCodeNormalization:
    def test_plain_code(self):
        assert rid_code_from_oldcode("P.1") == "P.1"

    def test_ridhydro_prefix(self):
        assert rid_code_from_oldcode("ridhydro_P.67") == "P.67"

    def test_letter_suffix(self):
        assert rid_code_from_oldcode("ridhydro_P.4A") == "P.4A"

    def test_non_rid_codes(self):
        assert rid_code_from_oldcode("CHM005") is None
        assert rid_code_from_oldcode("ridtele_TUP.14") is None
        assert rid_code_from_oldcode(None) is None


class TestParseRain:
    def test_filters_and_parses(self):
        records = parse_rain_records(_rain_payload())
        assert len(records) == 1
        r = records[0]
        assert r["station_id"] == 418
        assert r["oldcode"] == "CHM005"
        assert r["name_en"] == "Chiang Mai 5"
        assert r["rain_1h"] == 0.0
        assert r["rain_24h"] == 49.6
        assert r["timestamp"] == datetime.datetime(2026, 8, 11, 13, 0)
        assert r["agency"] == "HII"

    def test_empty_payload(self):
        assert parse_rain_records({}) == []


class TestParseWaterlevel:
    def test_filters_and_parses(self):
        records = parse_waterlevel_records(_waterlevel_payload())
        assert len(records) == 2
        p1 = records[0]
        assert p1["station_id"] == 3226
        assert p1["rid_code"] == "P.1"
        assert p1["wl_msl"] == 303.27
        assert p1["discharge"] == 335.0
        assert p1["flow_rate"] is None
        assert p1["storage_percent"] == 81.21
        assert p1["situation_level"] == 4
        assert p1["offset_msl"] == 300.5
        assert p1["is_key_station"] is True
        # MSL minus station offset recovers the familiar gauge level
        assert p1["wl_msl"] - p1["offset_msl"] == pytest.approx(2.77)
        assert records[1]["rid_code"] == "P.67"

    def test_empty_payload(self):
        assert parse_waterlevel_records({}) == []


class TestHiiStore:
    @pytest.fixture
    def store(self, tmp_path):
        store = HiiStore(f"sqlite:///{tmp_path}/hii_test.db", "sqlite")
        assert store.connect()
        return store

    def test_rejects_non_sql_backend(self):
        with pytest.raises(ValueError):
            HiiStore("http://localhost:8428", "victoriametrics")

    def test_rain_roundtrip_and_upsert(self, store):
        records = parse_rain_records(_rain_payload())
        assert store.save_rain(records) == 1
        # Same snapshot again -> upsert, still one row
        assert store.save_rain(records) == 1

        from sqlalchemy import text

        with store.engine.connect() as conn:
            rows = conn.execute(text("SELECT COUNT(*) FROM hii_rainfall")).scalar()
            stations = conn.execute(
                text("SELECT oldcode FROM hii_rain_stations")
            ).fetchall()
        assert rows == 1
        assert stations == [("CHM005",)]

    def test_waterlevel_roundtrip(self, store):
        records = parse_waterlevel_records(_waterlevel_payload())
        assert store.save_waterlevel(records) == 2

        from sqlalchemy import text

        with store.engine.connect() as conn:
            row = conn.execute(
                text(
                    "SELECT s.rid_code, m.wl_msl, m.situation_level "
                    "FROM hii_waterlevel m JOIN hii_wl_stations s ON s.id = m.station_id "
                    "WHERE s.oldcode = 'P.1'"
                )
            ).fetchone()
        assert row is not None
        assert row[0] == "P.1"
        assert float(row[1]) == 303.27
        assert row[2] == 4

    def test_save_empty(self, store):
        assert store.save_rain([]) == 0