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