151 lines
6.4 KiB
Python
151 lines
6.4 KiB
Python
"""The climate cache backend switch (data/climate_store.py + climate.py dispatch).
|
|
|
|
Two layers:
|
|
|
|
* **Hermetic** (always run): with no THERMOGRAPH_DATABASE_URL the store falls back
|
|
to the parquet cache, so ``is_postgres()`` is False, ``recent_stamp`` reads the
|
|
file mtime, and the backend helpers round-trip through parquet. This is the path
|
|
CI exercises — no Postgres required.
|
|
* **Gated** (skipped unless ``THERMOGRAPH_TEST_DATABASE_URL`` points at a reachable
|
|
Postgres): round-trips history/recent frames through the real psycopg + COPY
|
|
bridge and asserts dtype fidelity, upsert keep-last, and the recent_stamp token
|
|
semantics.
|
|
"""
|
|
import datetime
|
|
import os
|
|
|
|
import polars as pl
|
|
import pytest
|
|
|
|
from data import climate
|
|
from data import climate_store
|
|
|
|
|
|
def _hist_frame(n=5, tmax0=70.0):
|
|
"""A raw daily frame with the persisted columns (no doy — the write path drops it)."""
|
|
d0 = datetime.date(1995, 1, 1)
|
|
cols = {"date": [d0 + datetime.timedelta(days=i) for i in range(n)]}
|
|
for c in ("tmax", "tmin", "precip", "wind", "gust", "humid", "fmax", "fmin", "feels", "sun"):
|
|
cols[c] = [float(tmax0 + i) for i in range(n)]
|
|
return pl.DataFrame(cols)
|
|
|
|
|
|
# --- hermetic (parquet fallback) --------------------------------------------
|
|
|
|
def test_selects_parquet_when_no_database_url():
|
|
# The test env sets no THERMOGRAPH_DATABASE_URL (see conftest), so the climate
|
|
# record is served from parquet, not Postgres.
|
|
assert climate_store.is_postgres() is False
|
|
|
|
|
|
def test_recent_stamp_reads_mtime_on_parquet(monkeypatch, tmp_path):
|
|
monkeypatch.setattr(climate, "CACHE_DIR", str(tmp_path))
|
|
cell_id = "12_-34"
|
|
assert climate.recent_stamp(cell_id) == 0 # absent -> 0
|
|
climate._write_recent_backed(cell_id, _hist_frame())
|
|
stamp = climate.recent_stamp(cell_id)
|
|
assert stamp == int(os.path.getmtime(climate._rf_cache_path(cell_id)))
|
|
assert stamp > 0
|
|
|
|
|
|
def test_history_backend_roundtrips_through_parquet(monkeypatch, tmp_path):
|
|
monkeypatch.setattr(climate, "CACHE_DIR", str(tmp_path))
|
|
cell_id = "5_6"
|
|
assert climate._read_history_backed(cell_id) is None # nothing cached yet
|
|
climate._write_history_backed(cell_id, _hist_frame(n=4))
|
|
hit = climate._read_history_backed(cell_id)
|
|
assert hit is not None
|
|
df, age_s = hit
|
|
assert df.height == 4 and "doy" not in df.columns
|
|
assert age_s >= 0.0
|
|
|
|
|
|
def test_recent_backend_roundtrips_through_parquet(monkeypatch, tmp_path):
|
|
monkeypatch.setattr(climate, "CACHE_DIR", str(tmp_path))
|
|
cell_id = "7_8"
|
|
assert climate._read_recent_backed(cell_id) is None
|
|
climate._write_recent_backed(cell_id, _hist_frame(n=3))
|
|
hit = climate._read_recent_backed(cell_id)
|
|
assert hit is not None and hit[0].height == 3
|
|
|
|
|
|
# --- gated Postgres integration ---------------------------------------------
|
|
|
|
@pytest.fixture
|
|
def pg_store():
|
|
"""Point climate_store at THERMOGRAPH_TEST_DATABASE_URL for one test, creating
|
|
the tables (plain — the bridge logic doesn't need the hypertable) and wiping
|
|
them. Restores the module's dialect state afterwards."""
|
|
url = os.environ.get("THERMOGRAPH_TEST_DATABASE_URL", "").strip()
|
|
if not url:
|
|
pytest.skip("set THERMOGRAPH_TEST_DATABASE_URL to run the Postgres bridge test")
|
|
import psycopg
|
|
|
|
saved = (climate_store.IS_POSTGRES, climate_store._PG_DSN, climate_store._pool_obj)
|
|
climate_store.IS_POSTGRES = True
|
|
climate_store._PG_DSN = climate_store._libpq_dsn(url)
|
|
climate_store._pool_obj = None # force a fresh pool bound to the test DSN
|
|
|
|
ddl_cols = ", ".join(f"{c} DOUBLE PRECISION" for c in climate_store.COLS if c != "date")
|
|
with psycopg.connect(climate_store._PG_DSN, autocommit=True) as conn:
|
|
for tbl in ("climate_history", "climate_recent"):
|
|
conn.execute(f"CREATE TABLE IF NOT EXISTS {tbl} "
|
|
f"(cell_id TEXT NOT NULL, date DATE NOT NULL, {ddl_cols}, "
|
|
f"PRIMARY KEY (cell_id, date))")
|
|
conn.execute("CREATE TABLE IF NOT EXISTS climate_sync "
|
|
"(cell_id TEXT PRIMARY KEY, history_synced_at DOUBLE PRECISION, "
|
|
"recent_synced_at DOUBLE PRECISION, cols_version SMALLINT NOT NULL DEFAULT 1)")
|
|
conn.execute("TRUNCATE climate_history, climate_recent, climate_sync")
|
|
|
|
yield climate_store
|
|
|
|
if climate_store._pool_obj is not None:
|
|
climate_store._pool_obj.close()
|
|
climate_store.IS_POSTGRES, climate_store._PG_DSN, climate_store._pool_obj = saved
|
|
|
|
|
|
def test_pg_history_roundtrip_and_dtypes(pg_store):
|
|
pg_store.write_history("1_2", _hist_frame(n=5, tmax0=70.0), 1000.0)
|
|
hit = pg_store.read_history("1_2")
|
|
assert hit is not None
|
|
df, synced = hit
|
|
assert synced == 1000.0
|
|
assert df.schema["date"] == pl.Date and df.schema["tmax"] == pl.Float64
|
|
assert "doy" not in df.columns
|
|
assert df.height == 5 and df["tmax"].to_list() == [70, 71, 72, 73, 74]
|
|
|
|
|
|
def test_pg_history_upsert_keeps_last(pg_store):
|
|
pg_store.write_history("1_2", _hist_frame(n=3, tmax0=10.0), 1000.0)
|
|
# Re-write day 2 with a new value + append a 4th day.
|
|
d0 = datetime.date(1995, 1, 1)
|
|
tail = pl.DataFrame({
|
|
"date": [d0 + datetime.timedelta(days=1), d0 + datetime.timedelta(days=3)],
|
|
**{c: [999.0, 40.0] for c in
|
|
("tmax", "tmin", "precip", "wind", "gust", "humid", "fmax", "fmin", "feels")},
|
|
})
|
|
pg_store.write_history("1_2", tail, 2000.0)
|
|
df, synced = pg_store.read_history("1_2")
|
|
assert synced == 2000.0 and df.height == 4
|
|
assert df.filter(pl.col("date") == d0 + datetime.timedelta(days=1))["tmax"].item() == 999.0
|
|
|
|
|
|
def test_pg_recent_stamp_advances_only_on_write(pg_store):
|
|
assert pg_store.recent_synced_at("1_2") == 0.0
|
|
pg_store.write_recent("1_2", _hist_frame(n=4), 5000.0)
|
|
assert pg_store.recent_synced_at("1_2") == 5000.0
|
|
# A history write must not disturb the recent stamp, and vice-versa.
|
|
pg_store.write_history("1_2", _hist_frame(n=2), 6000.0)
|
|
assert pg_store.recent_synced_at("1_2") == 5000.0
|
|
assert pg_store.history_synced_at("1_2") == 6000.0
|
|
|
|
|
|
def test_pg_cols_version_gate(pg_store):
|
|
pg_store.write_history("1_2", _hist_frame(n=2), 1000.0)
|
|
assert pg_store.read_history("1_2") is not None
|
|
pg_store.COLS_VERSION = pg_store.COLS_VERSION + 1
|
|
try:
|
|
assert pg_store.read_history("1_2") is None # stale schema -> refetch
|
|
assert pg_store.read_history_raw("1_2") is not None # raw ignores the gate
|
|
finally:
|
|
pg_store.COLS_VERSION -= 1
|