SPipe/backend/tests/test_dem.py

"""DEM store: import (GeoTIFF, zipped ASC), bilinear sampling across CRSs, preview, TIN from points, HTTP."""
import io
import zipfile

import numpy as np
import pytest
import rasterio
from fastapi.testclient import TestClient
from pyproj import CRS, Transformer
from rasterio.transform import from_origin

from backend import dem as dem_store
from backend.batch import BatchError
from backend.main import app

EPSG = "EPSG:5186"
X0, Y0, CELL, W, H = 200_000.0, 550_000.0, 10.0, 60, 40  # top-left corner


def plane(x, y):
    return 100 + 0.01 * (x - X0) - 0.02 * (Y0 - y)


@pytest.fixture(autouse=True)
def store(tmp_path, monkeypatch):
    monkeypatch.setattr(dem_store, "DEM_DIR", tmp_path / "dem")


def make_tif(path):
    xs = X0 + (np.arange(W) + 0.5) * CELL
    ys = Y0 - (np.arange(H) + 0.5) * CELL
    z = plane(*np.meshgrid(xs, ys)).astype("float32")
    z[0, 0] = -9999
    with rasterio.open(path, "w", driver="GTiff", width=W, height=H, count=1, dtype="float32", crs=EPSG,
                       transform=from_origin(X0, Y0, CELL, CELL), nodata=-9999) as dst:
        dst.write(z, 1)


def test_import_sample_is_bilinear_and_handles_outside(tmp_path):
    src = tmp_path / "up.bin"
    make_tif(tmp_path / "a.tif")
    src.write_bytes((tmp_path / "a.tif").read_bytes())
    info = dem_store.import_dem(str(src), "a.tif")
    assert info["width"] == W and info["resolution"] == [CELL, CELL] and len(info["corners"]) == 4
    pts = [[X0 + 123.4, Y0 - 77.7], [X0 + 300, Y0 - 200], [X0 - 50, Y0 - 50]]
    vals = dem_store.sample(info["id"], pts, EPSG)
    assert vals[0] == pytest.approx(plane(*pts[0]), abs=1e-3)
    assert vals[1] == pytest.approx(plane(*pts[1]), abs=1e-3)
    assert vals[2] is None
    # Same point given in WGS84 lands on the same value.
    lon, lat = Transformer.from_crs(EPSG, "EPSG:4326", always_xy=True).transform(*pts[0])
    assert dem_store.sample(info["id"], [[lon, lat]], "EPSG:4326")[0] == pytest.approx(vals[0], abs=1e-3)
    assert dem_store.preview_png(info["id"])[:8] == b"\x89PNG\r\n\x1a\n"
    assert dem_store.dem_info(info["id"])["id"] == info["id"]


def test_zipped_asc_with_prj(tmp_path):
    xs = X0 + (np.arange(W) + 0.5) * CELL
    ys = Y0 - (np.arange(H) + 0.5) * CELL
    z = plane(*np.meshgrid(xs, ys))
    asc = f"ncols {W}\nnrows {H}\nxllcorner {X0}\nyllcorner {Y0 - H * CELL}\ncellsize {CELL}\nNODATA_value -9999\n"
    asc += "\n".join(" ".join(f"{v:.4f}" for v in row) for row in z)
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w") as zf:
        zf.writestr("dem/ground.asc", asc)
        zf.writestr("dem/ground.prj", CRS.from_user_input(EPSG).to_wkt("WKT1_ESRI"))
    up = tmp_path / "up.zip"
    up.write_bytes(buf.getvalue())
    info = dem_store.import_dem(str(up), "ground.zip")
    p = [X0 + 250, Y0 - 150]
    assert dem_store.sample(info["id"], [p], EPSG)[0] == pytest.approx(plane(*p), abs=1e-3)


def test_errors(tmp_path):
    with pytest.raises(BatchError):
        dem_store.dem_info("0" * 32)
    with pytest.raises(BatchError):
        dem_store.dem_info("../../etc")
    bad = tmp_path / "x.bin"
    bad.write_bytes(b"not a raster")
    with pytest.raises(BatchError):
        dem_store.import_dem(str(bad), "x.tif")
    with pytest.raises(BatchError):
        dem_store.import_dem(str(bad), "x.docx")


def test_tin_from_points_reproduces_a_plane():
    rng = np.random.default_rng(1)
    xy = np.column_stack([X0 + rng.uniform(0, 500, 300), Y0 - rng.uniform(0, 300, 300)])
    pts = [[x, y, plane(x, y)] for x, y in xy]
    pts += [[X0, Y0, plane(X0, Y0)], [X0 + 500, Y0, plane(X0 + 500, Y0)], [X0, Y0 - 300, plane(X0, Y0 - 300)], [X0 + 500, Y0 - 300, plane(X0 + 500, Y0 - 300)]]
    info = dem_store.from_points(pts, EPSG, 5.0)
    q = [X0 + 211, Y0 - 133]
    assert dem_store.sample(info["id"], [q], EPSG)[0] == pytest.approx(plane(*q), abs=0.05)
    with pytest.raises(BatchError):
        dem_store.from_points([[0, 0, 1], [1, 1, 2], [2, 2, 3]], EPSG, 1.0)  # collinear


def test_http_roundtrip(tmp_path):
    make_tif(tmp_path / "a.tif")
    with TestClient(app) as c:
        r = c.post("/api/dem", files={"file": ("a.tif", (tmp_path / "a.tif").read_bytes(), "image/tiff")})
        assert r.status_code == 200, r.text
        dem_id = r.json()["id"]
        assert c.get(f"/api/dem/{dem_id}/preview.png").headers["content-type"] == "image/png"
        s = c.post(f"/api/dem/{dem_id}/sample", json={"points": [[X0 + 100, Y0 - 100]], "crs": EPSG})
        assert s.json()["values"][0] == pytest.approx(plane(X0 + 100, Y0 - 100), abs=1e-3)
        assert c.get("/api/dem/" + "f" * 32).status_code == 400
        assert c.post("/api/dem", files={"file": ("a.txt", b"x", "text/plain")}).status_code == 400


def test_terrain_tile_terrarium(tmp_path):
    make_tif(tmp_path / "a.tif")
    up = tmp_path / "up.bin"
    up.write_bytes((tmp_path / "a.tif").read_bytes())
    info = dem_store.import_dem(str(up), "a.tif")
    import math
    lon, lat = Transformer.from_crs(EPSG, "EPSG:4326", always_xy=True).transform(X0 + 300, Y0 - 200)
    z = 15
    n = 2 ** z
    fx = (lon + 180) / 360 * n
    fy = (1 - math.asinh(math.tan(math.radians(lat))) / math.pi) / 2 * n
    xt, yt = int(fx), int(fy)
    png = dem_store.terrain_tile(info["id"], z, xt, yt)
    with rasterio.MemoryFile(png) as mem, mem.open() as ds:
        rgb = ds.read().astype(float)
    elev = rgb[0] * 256 + rgb[1] + rgb[2] / 256 - 32768
    px, py = int((fx - xt) * 256), int((fy - yt) * 256)
    assert elev[py, px] == pytest.approx(plane(X0 + 300, Y0 - 200), abs=0.2)
    with pytest.raises(BatchError):
        dem_store.terrain_tile(info["id"], 3, 99, 0)