grabowski
6eafb353b1
POST app.rid.go.th/reservoir/api/dams (open, archive >=2009) collected hourly into rid_dams + rid_reservoir_daily; backfill script fetches only missing days so reruns repair holes and are safe alongside the live collector. /api/stats counts the new table via an engine fallback that works when HII collection is disabled. Mae Ngat (DAM_ID 200103) hit 113% usable capacity with ~19 MCM/day inflow in the Oct 2024 flood — candidate features for the next retrain.
219 lines
7.7 KiB
Python
219 lines
7.7 KiB
Python
"""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_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
|