Model/tests/repositories/geospatial/test_geospatial_repository.py

"""GeospatialRepo resolves a Property's coordinates from the OS Open-UPRN data.

A reference-data lookup, not a Fetcher (ADR-0011): no live OS API call. The
adapter reads the partitioned Open-UPRN parquet via an injected reader, so the
test exercises the partition lookup + filter against real fixture parquets with
no network.
"""

from __future__ import annotations

from collections.abc import Callable
from pathlib import Path

import pandas as pd

from typing import Optional

from domain.geospatial.coordinates import Coordinates
from domain.geospatial.planning_restrictions import PlanningRestrictions
from domain.geospatial.spatial_reference import SpatialReference
from repositories.geospatial.geospatial_s3_repository import GeospatialS3Repository


def _reader(base: Path) -> Callable[[str], pd.DataFrame]:
    def read(key: str) -> pd.DataFrame:
        return pd.read_parquet(base / key)

    return read


def _write_open_uprn(base: Path) -> None:
    spatial = base / "spatial"
    spatial.mkdir(parents=True, exist_ok=True)
    pd.DataFrame(
        {"lower": [0], "upper": [100000], "filenames": ["0_100000.parquet"]}
    ).to_parquet(spatial / "filename_meta.parquet")
    pd.DataFrame(
        {
            "UPRN": [12345, 12346],
            "LATITUDE": [51.5074, 51.6000],
            "LONGITUDE": [-0.1278, -0.2000],
            # Planning flags co-located with the coordinates in the partition
            # (legacy column names — confirm exact names in the S3 deep-dive).
            "conservation_status": [True, False],
            "is_listed_building": [False, True],
            "is_heritage_building": [False, False],
        }
    ).to_parquet(spatial / "0_100000.parquet")


def test_coordinates_for_returns_lon_lat(tmp_path: Path) -> None:
    # Arrange
    _write_open_uprn(tmp_path)
    repo = GeospatialS3Repository(_reader(tmp_path))

    # Act
    coords = repo.coordinates_for(12345)

    # Assert
    assert coords == Coordinates(longitude=-0.1278, latitude=51.5074)


def test_coordinates_for_returns_none_when_uprn_absent(tmp_path: Path) -> None:
    # Arrange
    _write_open_uprn(tmp_path)
    repo = GeospatialS3Repository(_reader(tmp_path))

    # Act / Assert — uprn inside the partition range but not present in the data
    assert repo.coordinates_for(99999) is None


def test_coordinates_for_returns_none_when_no_partition_covers_uprn(
    tmp_path: Path,
) -> None:
    # Arrange
    _write_open_uprn(tmp_path)
    repo = GeospatialS3Repository(_reader(tmp_path))

    # Act / Assert — uprn beyond every partition's range
    assert repo.coordinates_for(500000) is None


def test_planning_restrictions_for_reads_the_co_located_flags(tmp_path: Path) -> None:
    # Arrange — same partition, planning flags alongside the coordinates.
    _write_open_uprn(tmp_path)
    repo = GeospatialS3Repository(_reader(tmp_path))

    # Act
    restrictions = repo.planning_restrictions_for(12345)

    # Assert — the three flags come back as the Property's PlanningRestrictions.
    assert restrictions == PlanningRestrictions(
        in_conservation_area=True, is_listed=False, is_heritage=False
    )


def test_planning_restrictions_for_returns_none_when_uprn_absent(
    tmp_path: Path,
) -> None:
    # Arrange
    _write_open_uprn(tmp_path)
    repo = GeospatialS3Repository(_reader(tmp_path))

    # Act / Assert
    assert repo.planning_restrictions_for(99999) is None


def test_spatial_for_returns_coordinates_and_restrictions_together(
    tmp_path: Path,
) -> None:
    # Arrange — one partition row carries the coordinates and the planning flags.
    _write_open_uprn(tmp_path)
    repo = GeospatialS3Repository(_reader(tmp_path))

    # Act — a single reference lookup yields both, so Ingestion reads the row once.
    reference: Optional[SpatialReference] = repo.spatial_for(12346)

    # Assert
    assert reference == SpatialReference(
        coordinates=Coordinates(longitude=-0.2000, latitude=51.6000),
        restrictions=PlanningRestrictions(
            in_conservation_area=False, is_listed=True, is_heritage=False
        ),
    )


def test_spatial_for_returns_none_when_uprn_absent(tmp_path: Path) -> None:
    # Arrange
    _write_open_uprn(tmp_path)
    repo = GeospatialS3Repository(_reader(tmp_path))

    # Act / Assert
    assert repo.spatial_for(99999) is None


def _write_two_partition_open_uprn(base: Path) -> None:
    """Two UPRN-range partitions, so the batch lookup must span both."""
    spatial = base / "spatial"
    spatial.mkdir(parents=True, exist_ok=True)
    pd.DataFrame(
        {
            "lower": [0, 100001],
            "upper": [100000, 200000],
            "filenames": ["0_100000.parquet", "100001_200000.parquet"],
        }
    ).to_parquet(spatial / "filename_meta.parquet")
    pd.DataFrame(
        {"UPRN": [10, 11], "LATITUDE": [51.0, 51.1], "LONGITUDE": [-1.0, -1.1]}
    ).to_parquet(spatial / "0_100000.parquet")
    pd.DataFrame(
        {"UPRN": [150000], "LATITUDE": [52.0], "LONGITUDE": [-2.0]}
    ).to_parquet(spatial / "100001_200000.parquet")


def test_coordinates_for_uprns_resolves_a_batch_across_partitions(
    tmp_path: Path,
) -> None:
    # Arrange — UPRNs spanning two partitions, plus one absent and one off-scale.
    _write_two_partition_open_uprn(tmp_path)
    repo = GeospatialS3Repository(_reader(tmp_path))

    # Act
    resolved = repo.coordinates_for_uprns([10, 11, 150000, 99999, 500000])

    # Assert — present UPRNs resolved; absent (99999) and uncovered (500000) omitted.
    assert resolved == {
        10: Coordinates(longitude=-1.0, latitude=51.0),
        11: Coordinates(longitude=-1.1, latitude=51.1),
        150000: Coordinates(longitude=-2.0, latitude=52.0),
    }


def test_coordinates_for_uprns_reads_each_partition_once(tmp_path: Path) -> None:
    # Arrange — count reads so co-located UPRNs don't re-read their partition.
    _write_two_partition_open_uprn(tmp_path)
    reads: list[str] = []

    def counting_reader(key: str) -> pd.DataFrame:
        reads.append(key)
        return pd.read_parquet(tmp_path / key)

    repo = GeospatialS3Repository(counting_reader)

    # Act — two UPRNs share partition 0; one is in partition 1.
    repo.coordinates_for_uprns([10, 11, 150000])

    # Assert — the meta once + each of the two partitions once (3 reads, not 4).
    assert reads.count("spatial/0_100000.parquet") == 1
    assert reads.count("spatial/100001_200000.parquet") == 1
    assert reads.count("spatial/filename_meta.parquet") == 1