thermograph/backend/tests/data/test_era5lake.py

"""The ERA5 lake: grid math, layout keys, and the read/slice client (no network
— the local-dir source stands in for the bucket)."""
import datetime

import polars as pl
import pytest

from data import era5lake


def _daily_frame(start="1940-01-01", days=32000):
    d0 = datetime.date.fromisoformat(start)
    dates = [d0 + datetime.timedelta(days=k) for k in range(days)]
    n = len(dates)
    return 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,
    })


def test_grid_roundtrip_and_bounds():
    # Nearest-index snap, poles clamped, longitude wrapped to [0, 360).
    assert era5lake.to_idx(90.0, 0.0) == (0, 0)
    assert era5lake.to_idx(-90.0, 0.0) == (720, 0)
    assert era5lake.to_idx(51.5074, -0.1278) == (154, 1439)
    i, j = era5lake.to_idx(35.6762, 139.6503)
    lat, lon = era5lake.to_coords(i, j)
    assert abs(lat - 35.6762) <= era5lake.STEP / 2 + 1e-9
    assert abs(lon - 139.6503) <= era5lake.STEP / 2 + 1e-9


def test_tile_and_keys_align_with_hive_layout():
    assert era5lake.tile_of(154, 1439) == (12, 119)
    assert era5lake.point_key(154, 1439) == "era5/points/lat=154/lon=1439.parquet"
    assert era5lake.daily_part_key(12, 119, 1994, 3) == \
        "era5/daily/tile=12_119/year=1994/month=03/part.parquet"


def test_s3_config_requires_both_keys(monkeypatch):
    monkeypatch.delenv("THERMOGRAPH_LAKE_S3_ACCESS_KEY", raising=False)
    monkeypatch.delenv("THERMOGRAPH_LAKE_S3_SECRET_KEY", raising=False)
    assert era5lake.s3_config() is None
    monkeypatch.setenv("THERMOGRAPH_LAKE_S3_ACCESS_KEY", "k")
    assert era5lake.s3_config() is None
    monkeypatch.setenv("THERMOGRAPH_LAKE_S3_SECRET_KEY", "s")
    cfg = era5lake.s3_config()
    assert cfg["bucket"] == "era5-thermograph"
    opts = era5lake.storage_options(cfg)
    assert opts["aws_virtual_hosted_style_request"] == "false"


def _write_point(tmp_path, cell, df):
    i, j = era5lake.to_idx(cell["center_lat"], cell["center_lon"])
    path = tmp_path / era5lake.point_key(i, j)
    path.parent.mkdir(parents=True, exist_ok=True)
    df.write_parquet(path)


CELL = {"center_lat": 51.5, "center_lon": -0.125}


def test_fetch_slices_to_start_and_keeps_schema(monkeypatch, tmp_path):
    monkeypatch.setenv("THERMOGRAPH_LAKE_LOCAL_DIR", str(tmp_path))
    _write_point(tmp_path, CELL, _daily_frame())
    df = era5lake.fetch_history_lake(CELL, start="1980-01-01")
    assert df["date"].min() == datetime.date(1980, 1, 1)
    assert set(df.columns) == {"date", "tmax", "tmin", "precip", "wind",
                               "gust", "humid", "fmax", "fmin"}
    # Without a start, the whole 1940+ record comes back.
    assert era5lake.fetch_history_lake(CELL).height == 32000


def test_fetch_rejects_short_and_missing(monkeypatch, tmp_path):
    monkeypatch.setenv("THERMOGRAPH_LAKE_LOCAL_DIR", str(tmp_path))
    with pytest.raises(era5lake.LakeUnavailable):
        era5lake.fetch_history_lake(CELL)          # nothing written
    _write_point(tmp_path, CELL, _daily_frame(days=500))
    with pytest.raises(era5lake.LakeUnavailable):  # implausibly short record
        era5lake.fetch_history_lake(CELL)


def test_fetch_unconfigured_raises(monkeypatch):
    monkeypatch.delenv("THERMOGRAPH_LAKE_LOCAL_DIR", raising=False)
    monkeypatch.delenv("THERMOGRAPH_LAKE_URL", raising=False)
    monkeypatch.delenv("THERMOGRAPH_LAKE_S3_ACCESS_KEY", raising=False)
    monkeypatch.delenv("THERMOGRAPH_LAKE_S3_SECRET_KEY", raising=False)
    with pytest.raises(era5lake.LakeUnavailable):
        era5lake.fetch_history_lake(CELL)