thermograph/backend/tests/data/test_store.py

"""data/store.py: hermetic SQLite tests (always run, via the tmp_store fixture in
conftest.py) plus a gated Postgres round-trip through the bounded connection pool
(skipped unless THERMOGRAPH_TEST_DATABASE_URL points at a reachable Postgres)."""
import json
import sqlite3
import time

import pytest


def test_payload_round_trip(tmp_store):
    payload = {"cell": "1_2", "days": [1, 2, 3]}
    body = tmp_store.put_payload("grade", "1_2", "k", "tok", payload)
    assert json.loads(body) == payload
    assert tmp_store.get_payload("grade", "1_2", "k", "tok") == body
    assert tmp_store.get_json("grade", "1_2", "k", "tok") == payload


def test_token_mismatch_is_a_miss(tmp_store):
    tmp_store.put_payload("grade", "1_2", "k", "tok-a", {"v": 1})
    assert tmp_store.get_payload("grade", "1_2", "k", "tok-b") is None
    # A rewrite under the new token replaces the row (same primary key).
    tmp_store.put_payload("grade", "1_2", "k", "tok-b", {"v": 2})
    assert tmp_store.get_json("grade", "1_2", "k", "tok-b") == {"v": 2}
    assert tmp_store.get_payload("grade", "1_2", "k", "tok-a") is None


def test_cache_false_serves_without_persisting(tmp_store):
    body = tmp_store.put_payload("grade", "1_2", "k", "tok", {"v": 1}, cache=False)
    assert json.loads(body) == {"v": 1}
    assert tmp_store.get_payload("grade", "1_2", "k", "tok") is None


def test_put_payload_rejects_nan(tmp_store):
    import pytest
    with pytest.raises(ValueError):
        tmp_store.put_payload("grade", "1_2", "k", "tok", {"v": float("nan")})


def test_revgeo_round_trip(tmp_store):
    key = tmp_store.revgeo_key(47.60623, -122.33305)
    assert key == "v2:47.606,-122.333"   # v2: label format now carries country
    assert tmp_store.get_revgeo(key) == (False, None)
    tmp_store.put_revgeo(key, "Belltown, Seattle, Washington, United States")
    assert tmp_store.get_revgeo(key) == (True, "Belltown, Seattle, Washington, United States")


def test_revgeo_cached_miss_retries_after_ttl(tmp_store):
    key = tmp_store.revgeo_key(1.0, 2.0)
    tmp_store.put_revgeo(key, None)
    assert tmp_store.get_revgeo(key) == (True, None)   # fresh miss: don't re-ask yet
    # Age the row past the TTL directly in SQLite.
    conn = sqlite3.connect(tmp_store.DB_PATH)
    conn.execute("UPDATE revgeo SET updated_at=? WHERE key=?",
                 (time.time() - tmp_store.REVGEO_MISS_TTL - 1, key))
    conn.commit()
    conn.close()
    assert tmp_store.get_revgeo(key) == (False, None)  # stale miss: retryable


def test_store_is_a_pure_accelerator_when_db_unavailable(tmp_store, monkeypatch):
    # Point the store somewhere unwritable: every helper degrades, none raises.
    monkeypatch.setattr(tmp_store, "DB_PATH", "/proc/nope/store.sqlite")
    import threading
    monkeypatch.setattr(tmp_store, "_local", threading.local())
    assert tmp_store.get_payload("grade", "c", "k", "t") is None
    body = tmp_store.put_payload("grade", "c", "k", "t", {"v": 1})
    assert json.loads(body) == {"v": 1}   # still serves the encoded payload
    assert tmp_store.get_revgeo("x") == (False, None)
    tmp_store.put_revgeo("x", "label")   # no-op, no exception


# --- gated Postgres integration (the bounded pool, _get_pg_pool) -----------

@pytest.fixture
def pg_store():
    """Point store.py at THERMOGRAPH_TEST_DATABASE_URL for one test -- exercises
    the real bounded ConnectionPool (borrow-per-operation, configure=_configure_pg
    running the schema DDL) rather than the SQLite fallback.

    store.py picks its dialect-specific SQL text (``?`` vs ``%s`` placeholders,
    INSERT OR REPLACE vs ON CONFLICT) once, at import time, from IS_POSTGRES --
    correct in production (THERMOGRAPH_DATABASE_URL is set before the module is
    ever imported) but not something flipping IS_POSTGRES back on afterwards
    retroactively fixes, so this fixture swaps the SQL constants too, not just
    IS_POSTGRES/_PG_DSN. Restores all of it afterwards."""
    import os

    from data import store

    url = os.environ.get("THERMOGRAPH_TEST_DATABASE_URL", "").strip()
    if not url:
        pytest.skip("set THERMOGRAPH_TEST_DATABASE_URL to run the Postgres pool test")

    saved = (store.IS_POSTGRES, store._PG_DSN, store._pg_pool,
             store._SQL_GET_PAYLOAD, store._SQL_PUT_PAYLOAD,
             store._SQL_GET_REVGEO, store._SQL_PUT_REVGEO)
    store.IS_POSTGRES = True
    store._PG_DSN = store._libpq_dsn(url)
    store._pg_pool = None   # force a fresh pool bound to the test DSN
    store._SQL_GET_PAYLOAD = "SELECT token, payload FROM derived WHERE kind=%s AND cell_id=%s AND key=%s"
    store._SQL_PUT_PAYLOAD = (
        "INSERT INTO derived (kind, cell_id, key, token, payload, updated_at) "
        "VALUES (%s,%s,%s,%s,%s,%s) "
        "ON CONFLICT (kind, cell_id, key) DO UPDATE SET "
        "token=EXCLUDED.token, payload=EXCLUDED.payload, updated_at=EXCLUDED.updated_at"
    )
    store._SQL_GET_REVGEO = "SELECT label, updated_at FROM revgeo WHERE key=%s"
    store._SQL_PUT_REVGEO = (
        "INSERT INTO revgeo (key, label, updated_at) VALUES (%s,%s,%s) "
        "ON CONFLICT (key) DO UPDATE SET label=EXCLUDED.label, updated_at=EXCLUDED.updated_at"
    )

    import psycopg
    with psycopg.connect(store._PG_DSN, autocommit=True) as conn:
        conn.execute("DROP TABLE IF EXISTS derived")
        conn.execute("DROP TABLE IF EXISTS revgeo")

    yield store

    if store._pg_pool is not None:
        store._pg_pool.close()
    (store.IS_POSTGRES, store._PG_DSN, store._pg_pool,
     store._SQL_GET_PAYLOAD, store._SQL_PUT_PAYLOAD,
     store._SQL_GET_REVGEO, store._SQL_PUT_REVGEO) = saved


def test_pg_payload_round_trip_through_the_pool(pg_store):
    """Two consecutive calls borrow (and return) separate pooled connections --
    unlike the old thread-local model there is no single connection object tying
    the write and the read together, so this also exercises that the schema
    (applied via the pool's configure= callback) is visible across borrows."""
    body = pg_store.put_payload("grade", "1_2", "k", "tok", {"v": 1})
    assert json.loads(body) == {"v": 1}
    assert pg_store.get_payload("grade", "1_2", "k", "tok") == body
    assert pg_store.get_payload("grade", "1_2", "k", "tok-mismatch") is None


def test_pg_revgeo_round_trip_through_the_pool(pg_store):
    key = pg_store.revgeo_key(47.6, -122.3)
    assert pg_store.get_revgeo(key) == (False, None)
    pg_store.put_revgeo(key, "Somewhere, WA")
    assert pg_store.get_revgeo(key) == (True, "Somewhere, WA")


def test_pg_pool_survives_many_sequential_borrows(pg_store):
    """A crude stand-in for the finding this pool exists to fix: many operations
    in a row must keep working off a small bounded pool (max_size=8, well below
    this loop) instead of piling up unclosed connections."""
    for i in range(25):
        pg_store.put_payload("grade", "cell", f"k{i}", "tok", {"i": i})
    for i in range(25):
        assert pg_store.get_json("grade", "cell", f"k{i}", "tok") == {"i": i}