"""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, backfill_dam, parse_dam_range_records, 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", } ], }, ], } def _range_row(date, **overrides): """One api/dam row, shaped like the real 2024-09-24 Mae Ngat response.""" row = { "DAM_ID": "200103", "DATE_curr": date, "DAM_Name": "เขื่อนแม่งัดสมบูรณ์ชล", "DAM_Region": "เหนือ", "DAM_QMax": "323.00", "DAM_QStore": "265.00", "DAM_QUsage": "253.00", "DUL_Useless": "12.00", "DMD_ULevel_curr": "395.91", "DMD_Q_curr": "0.00", "DMD_Inflow_curr": "19.06", "DMD_Outflow_curr": "0.13", "VAL_DMD_Q_curr": "242.87", "DMD_QUse_curr": "254.87", "PERCENT_DMD_QUse_curr": "96.18", # Same calendar date one year earlier -> must be ignored, not stored "DMD_ULevel_prev": "391.52", "DMD_QUse_prev": "195.64", "PERCENT_DMD_QUse_prev": "73.83", "DMD_Inflow_prev": "1.64", "DMD_Outflow_prev": "3.79", } row.update(overrides) return row def _range_payload(dates=("2024-09-24",), rows=None): return { "dam_id": "200103", "dam_name": "เขื่อนแม่งัดสมบูรณ์ชล", "dam_region": "เหนือ", "dam_coordinates": {"lat": 19.16138, "lng": 99.04011}, "dam_data": rows if rows is not None else [_range_row(d) for d in dates], "max": {"max_DMD_QUse_curr": "254.87"}, "min": {"min_DMD_QUse_curr": "254.87"}, } 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 FakeRangeClient: """Serves every day in the requested window, minus `gaps` (the real endpoint omits scattered days rather than returning empty rows).""" def __init__(self, gaps=(), fail_chunks=()): self.gaps = {datetime.date.fromisoformat(d) for d in gaps} self.fail_chunks = set(fail_chunks) # (start, end) tuples that raise self.calls = [] def fetch_dam_range(self, dam_id, start, end): self.calls.append((dam_id, start, end)) if (start, end) in self.fail_chunks: raise ConnectionError("boom") days = [ start + datetime.timedelta(days=i) for i in range((end - start).days + 1) ] return parse_dam_range_records( _range_payload(rows=[_range_row(d.isoformat()) for d in days if d not in self.gaps]) ) 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({}) == [] def test_parse_dam_range_records(self): records = parse_dam_range_records(_range_payload(("2024-09-24",))) assert len(records) == 1 # the _prev columns are last year, not a row row = records[0] assert row["dam_id"] == MAE_NGAT_DAM_ID assert row["date"] == datetime.date(2024, 9, 24) assert row["region"] == "เหนือ" assert row["latitude"] == 19.16138 assert row["longitude"] == 99.04011 assert row["capacity_max_mcm"] == 323.0 assert row["capacity_normal_mcm"] == 265.0 # Agrees with api/dams for this date, at higher percent precision assert row["storage_mcm"] == 254.87 assert row["storage_pct"] == 96.18 assert row["inflow_mcm"] == 19.06 assert row["outflow_mcm"] == 0.13 assert row["level_msl"] == 395.91 # DMD_ULevel, absent from api/dams def test_parse_dam_range_records_produces_same_keys_as_dams(self): assert set(parse_dam_range_records(_range_payload())[0]) == set( parse_dam_records(_dams_payload())[0] ) def test_parse_dam_range_zero_level_is_missing_not_a_reading(self): payload = _range_payload(rows=[_range_row("2024-09-24", DMD_ULevel_curr="0.00")]) assert parse_dam_range_records(payload)[0]["level_msl"] is None def test_parse_dam_range_skips_broken_rows(self): payload = _range_payload( rows=[ _range_row("n/a"), # unparseable date _range_row("2024-09-24"), ] ) assert len(parse_dam_range_records(payload)) == 1 def test_parse_dam_range_empty_payload(self): assert parse_dam_range_records({}) == [] assert parse_dam_range_records(_range_payload(rows=[])) == [] 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() def test_daily_collector_does_not_blank_a_backfilled_level(self, store): date = "2026-08-13" assert store.save(parse_dam_range_records(_range_payload((date,)))) == 1 # api/dams has no level column at all; its rewrite must not clear one dams_row = [r for r in parse_dam_records(_dams_payload(date)) if r["dam_id"] == MAE_NGAT_DAM_ID] assert dams_row[0]["level_msl"] is None store.save(dams_row) from sqlalchemy import text with store.engine.begin() as conn: row = conn.execute( text( "SELECT level_msl, storage_mcm FROM rid_reservoir_daily " "WHERE dam_id = :d" ), {"d": MAE_NGAT_DAM_ID}, ).fetchone() assert float(row[0]) == 395.91 # kept assert float(row[1]) == 222.01 # published columns still overwritten def test_present_dates_per_dam(self, store): lo, hi = datetime.date(2026, 8, 1), datetime.date(2026, 8, 31) store.save(parse_dam_records(_dams_payload())) day = datetime.date(2026, 8, 13) assert store.present_dates(lo, hi, dam_id=MAE_NGAT_DAM_ID) == {day} # Another dam having the date must not mark this one done assert store.present_dates(lo, hi, dam_id="999999") == set() def test_null_metadata_does_not_wipe_known_dam_details(self, store): store.save(parse_dam_records(_dams_payload())) blank = parse_dam_records(_dams_payload("2026-08-14")) for record in blank: # a payload that omits metadata record.update({"name_th": None, "latitude": None, "region": None}) assert store.save(blank) == 2 from sqlalchemy import text with store.engine.begin() as conn: row = conn.execute( text( "SELECT name_th, latitude, region FROM rid_dams " "WHERE dam_id = :d" ), {"d": MAE_NGAT_DAM_ID}, ).fetchone() assert row[0] == "เขื่อนแม่งัดสมบูรณ์ชล" assert float(row[1]) == 19.16138 assert row[2] == "เหนือ" 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 class TestBackfillDam: @pytest.fixture def store(self, tmp_path): store = RidReservoirStore(f"sqlite:///{tmp_path}/range.db", "sqlite") assert store.connect() return store def test_whole_range_in_one_request(self, store): client = FakeRangeClient() saved = backfill_dam( store, start=datetime.date(2024, 9, 24), end=datetime.date(2024, 10, 6), client=client, throttle_seconds=0, ) assert saved == 13 assert client.calls == [ (MAE_NGAT_DAM_ID, datetime.date(2024, 9, 24), datetime.date(2024, 10, 6)) ] def test_chunks_long_ranges(self, store): client = FakeRangeClient() saved = backfill_dam( store, start=datetime.date(2024, 1, 1), end=datetime.date(2024, 1, 10), client=client, chunk_days=4, throttle_seconds=0, ) assert saved == 10 assert client.calls == [ (MAE_NGAT_DAM_ID, datetime.date(2024, 1, 1), datetime.date(2024, 1, 4)), (MAE_NGAT_DAM_ID, datetime.date(2024, 1, 5), datetime.date(2024, 1, 8)), (MAE_NGAT_DAM_ID, datetime.date(2024, 1, 9), datetime.date(2024, 1, 10)), ] def test_source_gaps_are_tolerated(self, store): client = FakeRangeClient(gaps=("2024-01-03",)) saved = backfill_dam( store, start=datetime.date(2024, 1, 1), end=datetime.date(2024, 1, 5), client=client, throttle_seconds=0, ) assert saved == 4 # the day the source never published stays absent assert store.present_dates( datetime.date(2024, 1, 1), datetime.date(2024, 1, 5) ) == { datetime.date(2024, 1, d) for d in (1, 2, 4, 5) } def test_stored_days_are_skipped_and_whole_chunks_cost_no_request(self, store): client = FakeRangeClient() backfill_dam( store, start=datetime.date(2024, 1, 1), end=datetime.date(2024, 1, 4), client=client, throttle_seconds=0, ) # Rerun over a wider window: the stored chunk is not re-requested and # only the missing days are written saved = backfill_dam( store, start=datetime.date(2024, 1, 1), end=datetime.date(2024, 1, 8), client=client, chunk_days=4, throttle_seconds=0, ) assert saved == 4 assert client.calls[1:] == [ (MAE_NGAT_DAM_ID, datetime.date(2024, 1, 5), datetime.date(2024, 1, 8)) ] def test_partly_stored_chunk_saves_only_missing_days(self, store): client = FakeRangeClient() backfill_dam( store, start=datetime.date(2024, 1, 3), end=datetime.date(2024, 1, 3), client=client, throttle_seconds=0, ) saved = backfill_dam( store, start=datetime.date(2024, 1, 1), end=datetime.date(2024, 1, 5), client=client, throttle_seconds=0, ) assert saved == 4 # day 3 already present, requested but not rewritten def test_refresh_rewrites_stored_days(self, store): client = FakeRangeClient() window = dict( start=datetime.date(2024, 1, 1), end=datetime.date(2024, 1, 3), client=client, throttle_seconds=0, ) assert backfill_dam(store, **window) == 3 assert backfill_dam(store, skip_present=False, **window) == 3 assert len(client.calls) == 2 def test_junk_values_are_nulled_not_fatal(self, store): client = FakeRangeClient() client.fetch_dam_range = lambda dam_id, start, end: parse_dam_range_records( _range_payload( rows=[ _range_row( "2019-01-05", DMD_QUse_curr="343292.00", PERCENT_DMD_QUse_curr="87798.47", DMD_Outflow_curr="1e12", # beyond NUMERIC(10,2) -> NULL ) ] ) ) assert ( backfill_dam( store, start=datetime.date(2019, 1, 5), end=datetime.date(2019, 1, 5), client=client, throttle_seconds=0, ) == 1 ) from sqlalchemy import text with store.engine.begin() as conn: row = conn.execute( text( "SELECT storage_pct, outflow_mcm FROM rid_reservoir_daily " "WHERE date = '2019-01-05'" ) ).fetchone() assert float(row[0]) == 87798.47 assert row[1] is None def test_stats_separate_an_empty_rerun_from_an_outage(self, store): window = dict( start=datetime.date(2024, 1, 1), end=datetime.date(2024, 1, 3), throttle_seconds=0, ) backfill_dam(store, client=FakeRangeClient(), **window) # Everything already stored: no request, no failure, still a success stats = {} assert backfill_dam(store, client=FakeRangeClient(), stats=stats, **window) == 0 assert stats == {"requests": 0, "failures": 0, "aborted": False} # A source gap keeps requesting, but still reports no failure gapped = FakeRangeClient(gaps=("2024-01-05",)) stats = {} assert ( backfill_dam( store, client=gapped, stats=stats, start=datetime.date(2024, 1, 1), end=datetime.date(2024, 1, 5), throttle_seconds=0, ) == 1 ) assert stats["requests"] == 1 and not stats["failures"] def test_stats_record_failures(self, store): chunk = (datetime.date(2024, 1, 1), datetime.date(2024, 1, 3)) client = FakeRangeClient(fail_chunks=[chunk]) stats = {} assert ( backfill_dam( store, client=client, stats=stats, start=chunk[0], end=chunk[1], throttle_seconds=0, ) == 0 ) assert stats["failures"] == 1 and not stats["aborted"] def test_aborts_after_consecutive_failures(self, store): start = datetime.date(2024, 1, 1) chunks = [ (start + datetime.timedelta(days=i), start + datetime.timedelta(days=i)) for i in range(30) ] client = FakeRangeClient(fail_chunks=chunks[1:]) stats = {} saved = backfill_dam( store, start=start, end=start + datetime.timedelta(days=29), client=client, chunk_days=1, throttle_seconds=0, stats=stats, ) assert saved == 1 # first chunk succeeded, then 5 failures -> abort assert len(client.calls) == 6 assert stats["aborted"] and stats["failures"] == 5 def test_aborts_when_db_saves_nothing(self, store, monkeypatch): monkeypatch.setattr(store, "save", lambda records: 0) # broken DB client = FakeRangeClient() backfill_dam( store, start=datetime.date(2024, 1, 1), end=datetime.date(2024, 3, 1), client=client, chunk_days=1, throttle_seconds=0, ) assert len(client.calls) == 5 # aborted, not one request per chunk