"""Tests for the HII/ThaiWater api-v3 collector (parsing + persistence).""" import datetime import pytest from src.hii_backfill import chunk_date_range, parse_graph_rows, select_stations 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 def test_history_upsert_preserves_snapshot_columns(self, store): # A live snapshot row exists with extra columns populated store.save_waterlevel(parse_waterlevel_records(_waterlevel_payload())) # Backfill collides on the same (station, timestamp) with new values rows = [ { "timestamp": datetime.datetime(2026, 8, 11, 13, 0), "wl_msl": 303.30, "discharge": 340.0, }, { "timestamp": datetime.datetime(2019, 8, 1, 1, 0), "wl_msl": 301.71, "discharge": 11.7, }, ] assert store.save_waterlevel_history(3226, rows) == 2 from sqlalchemy import text with store.engine.connect() as conn: collided = conn.execute( text( "SELECT wl_msl, discharge, storage_percent, situation_level " "FROM hii_waterlevel WHERE station_id = 3226 " "AND timestamp = '2026-08-11 13:00:00'" ) ).fetchone() historical = conn.execute( text( "SELECT wl_msl FROM hii_waterlevel WHERE station_id = 3226 " "AND timestamp = '2019-08-01 01:00:00'" ) ).fetchone() # wl_msl/discharge updated, snapshot-only columns untouched assert float(collided[0]) == 303.30 assert float(collided[1]) == 340.0 assert float(collided[2]) == 81.21 assert collided[3] == 4 assert float(historical[0]) == 301.71 class TestParseGraphRows: def test_parses_and_skips_empty(self): payload = { "data": { "graph_data": [ {"datetime": "2024-10-05 12:00", "value": 305.8, "discharge": 656}, {"datetime": "2024-10-05 13:00", "value": None, "discharge": None}, {"datetime": None, "value": 300.0, "discharge": 1}, ] } } rows = parse_graph_rows(payload) assert rows == [ { "timestamp": datetime.datetime(2024, 10, 5, 12, 0), "wl_msl": 305.8, "discharge": 656.0, } ] def test_empty_payload(self): assert parse_graph_rows({}) == [] class TestHiiApiEndpoints: """Call the endpoint coroutines directly (the venv's httpx/starlette combination is incompatible with TestClient).""" @pytest.fixture def web_api(self, tmp_path, monkeypatch): from src import web_api from src.hii_collector import HiiCollector collector = HiiCollector( {"type": "sqlite", "connection_string": f"sqlite:///{tmp_path}/api.db"} ) now = datetime.datetime.now().replace(microsecond=0) rain = parse_rain_records(_rain_payload()) wl = parse_waterlevel_records(_waterlevel_payload()) for record in rain + wl: record["timestamp"] = now assert collector.store.save_rain(rain) == 1 assert collector.store.save_waterlevel(wl) == 2 monkeypatch.setitem(web_api.app_state, "hii_collector", collector) return web_api def test_rainfall_latest(self, web_api): import asyncio rows = asyncio.run(web_api.get_hii_rainfall_latest(hours=26)) assert len(rows) == 1 assert rows[0]["oldcode"] == "CHM005" assert rows[0]["rain_24h"] == 49.6 assert rows[0]["latitude"] == pytest.approx(19.12207) def test_waterlevel_latest(self, web_api): import asyncio rows = asyncio.run(web_api.get_hii_waterlevel_latest(hours=26)) assert len(rows) == 2 p1 = next(r for r in rows if r["oldcode"] == "P.1") assert p1["rid_code"] == "P.1" assert p1["wl_msl"] == 303.27 assert p1["offset_msl"] == 300.5 assert p1["situation_level"] == 4 def test_empty_when_collector_disabled(self, monkeypatch): import asyncio from src import web_api monkeypatch.setitem(web_api.app_state, "hii_collector", None) assert asyncio.run(web_api.get_hii_rainfall_latest(hours=26)) == [] assert asyncio.run(web_api.get_hii_waterlevel_latest(hours=26)) == [] class TestBackfillHelpers: def test_chunk_date_range(self): chunks = chunk_date_range( datetime.date(2024, 1, 1), datetime.date(2024, 3, 1), 31 ) assert chunks[0] == (datetime.date(2024, 1, 1), datetime.date(2024, 1, 31)) assert chunks[-1][1] == datetime.date(2024, 3, 1) # Contiguous, no overlap for (_, prev_end), (next_start, _) in zip(chunks, chunks[1:]): assert next_start == prev_end + datetime.timedelta(days=1) def test_chunk_single_day(self): d = datetime.date(2024, 1, 1) assert chunk_date_range(d, d, 365) == [(d, d)] def test_select_default_keeps_rid_and_key_stations(self): records = [ {"station_id": 1, "rid_code": "P.1", "is_key_station": True}, {"station_id": 2, "rid_code": None, "is_key_station": False}, {"station_id": 3, "rid_code": None, "is_key_station": True}, ] assert [r["station_id"] for r in select_stations(records)] == [1, 3] def test_select_by_code_matches_rid_code_and_oldcode(self): records = [ {"station_id": 1, "rid_code": "P.1", "oldcode": "ridhydro_P.1"}, {"station_id": 2, "rid_code": None, "oldcode": "CHM004"}, {"station_id": 3, "rid_code": "P.67", "oldcode": "P.67"}, ] selected = select_stations(records, codes=["p.1", "chm004"]) assert [r["station_id"] for r in selected] == [1, 2] def test_select_all(self): records = [{"station_id": 1}, {"station_id": 2}] assert select_stations(records, all_stations=True) == records