85 lines
2.9 KiB
Python
85 lines
2.9 KiB
Python
import polars as pl
|
|
import pytest
|
|
|
|
from pipeline.utils.poi_counts import POI_GROUPS, count_pois_within_radius
|
|
|
|
|
|
@pytest.fixture
|
|
def pois():
|
|
"""POIs clustered around two locations: central London and 10km away."""
|
|
return pl.DataFrame({
|
|
"lat": [51.5074, 51.5075, 51.5080, 51.5076, 51.5073, 51.60],
|
|
"lng": [-0.1278, -0.1280, -0.1275, -0.1279, -0.1277, -0.20],
|
|
"category": [
|
|
"Restaurant",
|
|
"Fast Food",
|
|
"Supermarket",
|
|
"Park",
|
|
"Station",
|
|
"Restaurant", # too far from any property
|
|
],
|
|
})
|
|
|
|
|
|
@pytest.fixture
|
|
def properties():
|
|
"""Two properties at the same postcode near central London, one at a distant postcode."""
|
|
return pl.DataFrame({
|
|
"postcode": ["EC1A 1BB", "EC1A 1BB", "ZZ99 9ZZ"],
|
|
"lat": [51.5074, 51.5074, 55.0],
|
|
"lon": [-0.1278, -0.1278, -3.0],
|
|
})
|
|
|
|
|
|
def test_counts_pois_within_radius(properties, pois):
|
|
result = count_pois_within_radius(properties, pois, radius_km=2.0)
|
|
|
|
assert set(result.keys()) == {f"{g}_2km" for g in POI_GROUPS}
|
|
|
|
# Result Series must be aligned to properties (3 rows)
|
|
for col, series in result.items():
|
|
assert len(series) == 3, f"{col} has {len(series)} rows, expected 3"
|
|
|
|
# First two rows share a postcode near the central London cluster
|
|
assert result["restaurants_2km"][0] == 2 # Restaurant + Fast Food
|
|
assert result["groceries_2km"][0] == 1 # Supermarket
|
|
assert result["parks_2km"][0] == 1 # Park
|
|
assert result["public_transport_2km"][0] == 1 # Station
|
|
|
|
# Second row is the same postcode, so same counts
|
|
assert result["restaurants_2km"][1] == result["restaurants_2km"][0]
|
|
|
|
# Third row (ZZ99 9ZZ) is far from all POIs → zero counts
|
|
for group in POI_GROUPS:
|
|
assert result[f"{group}_2km"][2] == 0
|
|
|
|
|
|
def test_no_pois_returns_zeros(properties):
|
|
empty_pois = pl.DataFrame({
|
|
"lat": pl.Series([], dtype=pl.Float64),
|
|
"lng": pl.Series([], dtype=pl.Float64),
|
|
"category": pl.Series([], dtype=pl.String),
|
|
})
|
|
result = count_pois_within_radius(properties, empty_pois, radius_km=2.0)
|
|
|
|
for group in POI_GROUPS:
|
|
col = f"{group}_2km"
|
|
assert col in result
|
|
assert result[col].to_list() == [0, 0, 0]
|
|
|
|
|
|
def test_custom_radius(pois):
|
|
"""A tiny radius should exclude POIs that are even slightly away."""
|
|
properties = pl.DataFrame({
|
|
"postcode": ["EC1A 1BB"],
|
|
"lat": [51.5074],
|
|
"lon": [-0.1278],
|
|
})
|
|
|
|
# 0.01 km = 10m — only the POI at the exact same location should match
|
|
result = count_pois_within_radius(properties, pois, radius_km=0.01)
|
|
# The Restaurant at (51.5074, -0.1278) is at distance 0
|
|
assert result["restaurants_0km"][0] >= 1
|
|
# POIs >100m away should not be counted
|
|
total = sum(result[f"{g}_0km"][0] for g in POI_GROUPS)
|
|
assert total <= 2 # at most the co-located POIs
|