thermograph/backend/tests/data/test_climate_store.py

"""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


def test_history_max_date_reads_parquet_tail_without_full_load(monkeypatch, tmp_path):
    monkeypatch.setattr(climate, "CACHE_DIR", str(tmp_path))
    cell_id = "9_10"
    assert climate.history_max_date(cell_id) is None            # nothing cached -> None
    climate._write_history_backed(cell_id, _hist_frame(n=5))     # dates 1995-01-01..05
    assert climate.history_max_date(cell_id) == "1995-01-05"
    # A tail top-up that appends a newer day advances the max date.
    climate._write_history_backed(cell_id, _hist_frame(n=7))
    assert climate.history_max_date(cell_id) == "1995-01-07"


# --- 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_history_max_date(pg_store):
    assert pg_store.history_max_date("1_2") is None            # no rows -> None
    pg_store.write_history("1_2", _hist_frame(n=5), 1000.0)    # dates 1995-01-01..05
    assert pg_store.history_max_date("1_2") == "1995-01-05"
    # Another cell's rows must not leak into this cell's max.
    pg_store.write_history("3_4", _hist_frame(n=9), 1000.0)
    assert pg_store.history_max_date("1_2") == "1995-01-05"


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