RadioPropagationApi/api/tests/test_viewshed_and_canopy.py

from __future__ import annotations

import sys
from pathlib import Path
from types import SimpleNamespace

import numpy as np
import rasterio
from rasterio.transform import from_origin

from app.core.geo import GeoPoint, PathPoint
from app.core.viewshed import _find_dem_tiles, _radius_bbox
from app.services import terrain as terrain_service

sys.path.append(str(Path(__file__).resolve().parents[2] / "scripts"))
import bootstrap_canopy  # noqa: E402


def _write_dem(path: Path, west: float) -> None:
    data = np.ones((10, 10), dtype="float32")
    _write_raster(path, west, data)


def _write_raster(path: Path, west: float, data: np.ndarray) -> None:
    transform = from_origin(west, 60.0, 0.01, 0.01)
    with rasterio.open(
        path,
        "w",
        driver="GTiff",
        height=data.shape[0],
        width=data.shape[1],
        count=1,
        dtype=data.dtype,
        crs="EPSG:4326",
        transform=transform,
        nodata=-9999.0,
    ) as dataset:
        dataset.write(data, 1)


def _settings(dem_path: Path, canopy_path: Path) -> SimpleNamespace:
    return SimpleNamespace(dem_path=dem_path, canopy_path=canopy_path)


def test_viewshed_dem_search_uses_radius_bbox(tmp_path: Path) -> None:
    _write_dem(tmp_path / "west.tif", 30.0)
    _write_dem(tmp_path / "east.tif", 30.1)
    bbox = _radius_bbox(GeoPoint(lat=59.95, lon=30.095), 2_000)

    tiles = _find_dem_tiles(tmp_path, bbox)

    assert {tile.name for tile in tiles} == {"west.tif", "east.tif"}


def test_surface_profile_uses_canopy_heights(monkeypatch, tmp_path: Path) -> None:
    dem_dir = tmp_path / "dem"
    canopy_dir = tmp_path / "canopy"
    dem_dir.mkdir()
    canopy_dir.mkdir()
    _write_raster(dem_dir / "dem.tif", 30.0, np.full((10, 10), 2.0, dtype="float32"))
    _write_raster(
        canopy_dir / "canopy.tif",
        30.0,
        np.full((10, 10), 12.0, dtype="float32"),
    )
    monkeypatch.setattr(terrain_service, "get_settings", lambda: _settings(dem_dir, canopy_dir))
    points = [
        PathPoint(lat=59.95, lon=30.05, distance_m=0),
        PathPoint(lat=59.94, lon=30.06, distance_m=1_000),
    ]

    profile = terrain_service.surface_profile_from_points(
        points,
        include_buildings=False,
        include_canopy=True,
    )

    assert [sample.ground_m for sample in profile.samples] == [2.0, 2.0]
    assert [sample.canopy_m for sample in profile.samples] == [12.0, 12.0]
    assert [sample.surface_m for sample in profile.samples] == [14.0, 14.0]


def test_surface_profile_ignores_canopy_when_disabled(monkeypatch, tmp_path: Path) -> None:
    dem_dir = tmp_path / "dem"
    canopy_dir = tmp_path / "canopy"
    dem_dir.mkdir()
    canopy_dir.mkdir()
    _write_raster(dem_dir / "dem.tif", 30.0, np.full((10, 10), 2.0, dtype="float32"))
    _write_raster(
        canopy_dir / "canopy.tif",
        30.0,
        np.full((10, 10), 12.0, dtype="float32"),
    )
    monkeypatch.setattr(terrain_service, "get_settings", lambda: _settings(dem_dir, canopy_dir))
    points = [PathPoint(lat=59.95, lon=30.05, distance_m=0)]

    profile = terrain_service.surface_profile_from_points(
        points,
        include_buildings=False,
        include_canopy=False,
    )

    assert profile.samples[0].canopy_m == 0.0
    assert profile.samples[0].surface_m == 2.0


def test_surface_profile_treats_missing_canopy_as_optional(
    monkeypatch,
    tmp_path: Path,
) -> None:
    dem_dir = tmp_path / "dem"
    canopy_dir = tmp_path / "canopy"
    dem_dir.mkdir()
    _write_raster(dem_dir / "dem.tif", 30.0, np.full((10, 10), 2.0, dtype="float32"))
    monkeypatch.setattr(terrain_service, "get_settings", lambda: _settings(dem_dir, canopy_dir))
    points = [PathPoint(lat=59.95, lon=30.05, distance_m=0)]

    profile = terrain_service.surface_profile_from_points(
        points,
        include_buildings=False,
        include_canopy=True,
    )

    assert profile.samples[0].canopy_m == 0.0
    assert profile.samples[0].surface_m == 2.0


def test_canopy_selects_intersecting_tiles_from_geojson_index() -> None:
    index = {
        "type": "FeatureCollection",
        "features": [
            {
                "type": "Feature",
                "bbox": [30.0, 59.0, 31.0, 60.0],
                "properties": {"filename": "tile_a.tif"},
            },
            {
                "type": "Feature",
                "bbox": [40.0, 59.0, 41.0, 60.0],
                "properties": {"filename": "tile_b.tif"},
            },
        ],
    }

    urls = bootstrap_canopy.select_tile_urls(index, (30.2, 59.2, 30.4, 59.4))

    assert len(urls) == 1
    assert urls[0].endswith("/tile_a.tif")


def test_canopy_parse_socks5_proxy_url() -> None:
    proxy = bootstrap_canopy.parse_socks5_proxy(
        "socks5h://proxy_user:secret@194.33.35.46:38599"
    )
    assert proxy.scheme == "socks5h"
    assert proxy.host == "194.33.35.46"
    assert proxy.port == 38599
    assert proxy.username == "proxy_user"
    assert proxy.password == "secret"
    assert proxy.endpoint == "socks5h://194.33.35.46:38599"


def test_canopy_response_total_size_from_content_range() -> None:
    response = type(
        "Response",
        (),
        {"headers": {"Content-Range": "bytes 1048576-2097151/257131787"}},
    )()

    assert bootstrap_canopy._response_total_size(response, 1_048_576) == 257_131_787


def test_canopy_urls_file_ignores_comments(tmp_path: Path) -> None:
    urls_file = tmp_path / "urls.txt"
    urls_file.write_text(
        "# comment\n"
        "https://example.com/a.tif\n"
        "\n"
        "https://example.com/b.tif\n",
        encoding="utf-8",
    )

    assert bootstrap_canopy.urls_from_file(urls_file) == [
        "https://example.com/a.tif",
        "https://example.com/b.tif",
    ]