thermograph/backend/tests/web/test_lake_app.py

"""The lake service over a local-dir lake fixture: health, the point hot path
(+ its disk cache), and the SELECT-only SQL surface."""
import datetime
import io

import polars as pl
import pytest
from fastapi.testclient import TestClient


@pytest.fixture
def lake(tmp_path, monkeypatch):
    """A tiny two-point lake (points, hive partitions, manifest) on disk."""
    from data import era5lake

    points = [(154, 1439), (154, 0)]           # around London
    d0 = datetime.date(1940, 1, 1)
    dates = [d0 + datetime.timedelta(days=k) for k in range(12000)]
    n = len(dates)
    rows = []
    for i, j in points:
        df = pl.DataFrame({
            "date": dates,
            "tmax": [70.0] * n, "tmin": [50.0] * n, "precip": [0.1] * n,
            "wind": [8.0] * n, "gust": [14.0] * n, "humid": [55.0] * n,
            "fmax": [71.0] * n, "fmin": [48.0] * n,
        })
        p = tmp_path / era5lake.point_key(i, j)
        p.parent.mkdir(parents=True, exist_ok=True)
        df.write_parquet(p)
        ti, tj = era5lake.tile_of(i, j)
        part = df.with_columns(pl.lit(i, dtype=pl.Int32).alias("lat_idx"),
                               pl.lit(j, dtype=pl.Int32).alias("lon_idx"))
        hp = tmp_path / era5lake.daily_part_key(ti, tj, 1940, 1)
        hp.parent.mkdir(parents=True, exist_ok=True)
        part.filter((pl.col("date").dt.year() == 1940)
                    & (pl.col("date").dt.month() == 1)).write_parquet(hp)
        lat, lon = era5lake.to_coords(i, j)
        rows.append({"lat_idx": i, "lon_idx": j, "lat": lat, "lon": lon,
                     "tile": f"{ti}_{tj}", "rows": n,
                     "date_min": "1940-01-01", "date_max": str(dates[-1])})
    (tmp_path / "era5").mkdir(exist_ok=True)
    pl.DataFrame(rows).write_parquet(tmp_path / era5lake.MANIFEST_KEY)

    monkeypatch.setenv("THERMOGRAPH_LAKE_LOCAL_DIR", str(tmp_path))
    monkeypatch.setenv("THERMOGRAPH_LAKE_CACHE", str(tmp_path / "cache"))
    import lake_app
    monkeypatch.setattr(lake_app, "CACHE_DIR", str(tmp_path / "cache"))
    monkeypatch.setattr(lake_app, "_manifest", None)
    return TestClient(lake_app.app)


def test_healthz_reports_points(lake):
    body = lake.get("/healthz").json()
    assert body["status"] == "ok"
    assert body["source"] == "local"
    assert body["points"] == 2


def test_history_serves_parquet_and_caches(lake, tmp_path):
    r = lake.get("/history", params={"lat": 51.5074, "lon": -0.1278})
    assert r.status_code == 200
    assert r.headers["x-lake-point"] == "154,1439"
    df = pl.read_parquet(io.BytesIO(r.content))
    assert df.height == 12000
    assert (tmp_path / "cache" / "p154_1439.parquet").exists()
    # Second hit comes off the cache file and is byte-identical.
    assert lake.get("/history",
                    params={"lat": 51.5074, "lon": -0.1278}).content == r.content


def test_history_404_for_point_outside_lake(lake):
    r = lake.get("/history", params={"lat": -45.0, "lon": 170.0})
    assert r.status_code == 404


def test_query_selects_with_partition_predicates(lake):
    r = lake.post("/query", json={"sql":
        "SELECT year, month, count(*) AS n FROM era5_daily "
        "WHERE year = 1940 AND month = 1 GROUP BY year, month"})
    assert r.status_code == 200
    body = r.json()
    assert body["columns"] == ["year", "month", "n"]
    assert body["rows"][0][2] == 2 * 31          # two points x January 1940


def test_query_manifest_view(lake):
    r = lake.post("/query", json={"sql": "SELECT count(*) FROM manifest"})
    assert r.status_code == 200
    assert r.json()["rows"][0][0] == 2


@pytest.mark.parametrize("sql", [
    "DROP TABLE era5_daily",
    "SELECT 1; SELECT 2",
    "INSERT INTO era5_daily VALUES (1)",
    "COPY era5_daily TO 'x'",
])
def test_query_rejects_non_select(lake, sql):
    assert lake.post("/query", json={"sql": sql}).status_code == 400