"""Tests for the RID large-dam daily collector (parsing + persistence).""" import datetime import pytest from src.rid_reservoir import ( MAE_NGAT_DAM_ID, RidReservoirCollector, RidReservoirStore, backfill, parse_dam_records, ) def _dams_payload(date="2026-08-13"): return { "date_th": date, "regions": [ { "region_name": "เหนือ", "dams": [ { "DAM_ID": "200103", "DAM_Name": "เขื่อนแม่งัดสมบูรณ์ชล", "DAM_Lat": 19.16138, "DAM_Lon": 99.04011, "DMD_Date": date, "DAM_QMax": "323.00", "DAM_QStore": "265.00", "DMD_QUse": "222.01", "PERCENT_DMD_QUse": "84.00", "DMD_Inflow": "3.44", "DMD_Outflow": "3.57", "DMD_Q": " - ", # placeholder -> None }, # No DAM_ID -> skipped {"DAM_Name": "broken", "DMD_Date": date}, # Bad date -> skipped {"DAM_ID": "200199", "DMD_Date": "n/a"}, ], }, { "region_name": "กลาง", "dams": [ { "DAM_ID": "200301", "DAM_Name": "เขื่อนป่าสักชลสิทธิ์", "DMD_Date": date, "DMD_QUse": "500.10", "DMD_Q": "255.57", } ], }, ], } class FakeClient: def __init__(self, payload=None, fail_dates=()): self.payload = payload or _dams_payload() self.fail_dates = set(fail_dates) self.calls = [] def fetch_day(self, date=None): self.calls.append(date) if date in self.fail_dates: raise ConnectionError("boom") if date is not None: return parse_dam_records(_dams_payload(date.isoformat())) return parse_dam_records(self.payload) class TestParsing: def test_parse_dam_records(self): records = parse_dam_records(_dams_payload()) assert len(records) == 2 # broken rows skipped ngat = next(r for r in records if r["dam_id"] == MAE_NGAT_DAM_ID) assert ngat["region"] == "เหนือ" assert ngat["date"] == datetime.date(2026, 8, 13) assert ngat["storage_mcm"] == 222.01 assert ngat["storage_pct"] == 84.0 assert ngat["inflow_mcm"] == 3.44 assert ngat["outflow_mcm"] == 3.57 assert ngat["level_msl"] is None # ' - ' placeholder assert ngat["capacity_normal_mcm"] == 265.0 def test_parse_empty_payload(self): assert parse_dam_records({}) == [] class TestStore: @pytest.fixture def store(self, tmp_path): store = RidReservoirStore(f"sqlite:///{tmp_path}/rsv_test.db", "sqlite") assert store.connect() return store def test_rejects_non_sql(self): with pytest.raises(ValueError): RidReservoirStore("http://localhost:8428", "victoriametrics") def test_roundtrip_and_upsert(self, store): records = parse_dam_records(_dams_payload()) assert store.save(records) == 2 # Second save of the same day updates in place, no duplicates records[0]["storage_mcm"] = 230.00 assert store.save(records) == 2 from sqlalchemy import text with store.engine.begin() as conn: rows = conn.execute( text( "SELECT dam_id, storage_mcm FROM rid_reservoir_daily " "ORDER BY dam_id" ) ).fetchall() dams = conn.execute(text("SELECT COUNT(*) FROM rid_dams")).scalar() assert len(rows) == 2 assert dams == 2 assert float(rows[0][1]) == 230.00 def test_save_empty(self, store): assert store.save([]) == 0 def test_junk_source_values_are_nulled_not_fatal(self, store): # Real junk from 2019-01-05: dam 100602 reported 87798% storage, # which overflowed NUMERIC(6,2) and discarded the whole 33-dam batch. payload = _dams_payload() payload["regions"][0]["dams"].append( { "DAM_ID": "100602", "DAM_Name": "junk", "DMD_Date": "2026-08-13", "DMD_QUse": "343292.00", "PERCENT_DMD_QUse": "87798.00", "DMD_Inflow": "0.59", "DMD_Outflow": "1e12", # beyond NUMERIC(10,2) -> NULL } ) assert store.save(parse_dam_records(payload)) == 3 from sqlalchemy import text with store.engine.begin() as conn: row = conn.execute( text( "SELECT storage_mcm, storage_pct, outflow_mcm " "FROM rid_reservoir_daily WHERE dam_id = '100602'" ) ).fetchone() assert float(row[0]) == 343292.00 # fits NUMERIC(10,2), kept raw assert float(row[1]) == 87798.00 # fits widened NUMERIC(8,2) assert row[2] is None # beyond capacity -> NULL, batch survives def test_present_dates(self, store): lo, hi = datetime.date(2026, 8, 1), datetime.date(2026, 8, 31) assert store.present_dates(lo, hi) == set() store.save(parse_dam_records(_dams_payload())) assert store.present_dates(lo, hi) == {datetime.date(2026, 8, 13)} # Outside the window -> excluded assert store.present_dates(lo, datetime.date(2026, 8, 12)) == set() class TestCollectorAndBackfill: def test_run_cycle_today_and_yesterday(self, tmp_path): collector = RidReservoirCollector( {"type": "sqlite", "connection_string": f"sqlite:///{tmp_path}/c.db"}, client=FakeClient(), ) saved = collector.run_cycle() assert saved == 4 # 2 dams x (today + yesterday) assert collector.client.calls[0] is None # today via empty date def test_run_cycle_survives_fetch_failure(self, tmp_path): yesterday = datetime.date.today() - datetime.timedelta(days=1) collector = RidReservoirCollector( {"type": "sqlite", "connection_string": f"sqlite:///{tmp_path}/f.db"}, client=FakeClient(fail_dates={yesterday}), ) assert collector.run_cycle() == 2 # today only def test_backfill_range(self, tmp_path): store = RidReservoirStore(f"sqlite:///{tmp_path}/b.db", "sqlite") assert store.connect() client = FakeClient() saved = backfill( store, datetime.date(2024, 9, 24), datetime.date(2024, 9, 26), client=client, throttle_seconds=0, ) assert saved == 6 # 3 days x 2 dams assert client.calls == [ datetime.date(2024, 9, 24), datetime.date(2024, 9, 25), datetime.date(2024, 9, 26), ] def test_backfill_skips_present_days_so_reruns_repair_holes(self, tmp_path): store = RidReservoirStore(f"sqlite:///{tmp_path}/skip.db", "sqlite") assert store.connect() # Day 25 already stored (e.g. by the live collector) store.save(parse_dam_records(_dams_payload("2024-09-25"))) client = FakeClient() saved = backfill( store, datetime.date(2024, 9, 24), datetime.date(2024, 9, 26), client=client, throttle_seconds=0, ) assert client.calls == [ datetime.date(2024, 9, 24), datetime.date(2024, 9, 26), ] assert saved == 4 def test_backfill_aborts_when_db_saves_nothing(self, tmp_path, monkeypatch): store = RidReservoirStore(f"sqlite:///{tmp_path}/deaddb.db", "sqlite") assert store.connect() monkeypatch.setattr(store, "save", lambda records: 0) # broken DB client = FakeClient() backfill( store, datetime.date(2024, 1, 1), datetime.date(2024, 3, 1), client=client, throttle_seconds=0, ) assert len(client.calls) == 5 # aborted, not one request per day def test_backfill_aborts_after_consecutive_failures(self, tmp_path): store = RidReservoirStore(f"sqlite:///{tmp_path}/a.db", "sqlite") assert store.connect() start = datetime.date(2024, 1, 1) fail_dates = {start + datetime.timedelta(days=i) for i in range(2, 30)} client = FakeClient(fail_dates=fail_dates) saved = backfill( store, start, datetime.date(2024, 3, 1), client=client, throttle_seconds=0, ) assert saved == 4 # first 2 days succeeded, then 5 failures -> abort assert len(client.calls) == 7