emi
370a4ffdc1
All checks were successful
secrets-guard / encrypted (push) Successful in 10s
Deploy backend to LAN dev server / build (push) Successful in 1m13s
Build + push backend image (Forgejo registry) / build-push (push) Successful in 1m39s
Deploy backend to LAN dev server / deploy (push) Successful in 34s
100 lines
3.8 KiB
Python
100 lines
3.8 KiB
Python
"""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
|