98 lines
3 KiB
Python
98 lines
3 KiB
Python
import polars as pl
|
|
import pytest
|
|
|
|
from pipeline.utils.poi_counts import count_pois_per_postcode
|
|
|
|
POI_GROUPS = {
|
|
"restaurants": ["Restaurant", "Fast Food"],
|
|
"groceries": ["Supermarket"],
|
|
"parks": ["Park"],
|
|
"public_transport": ["Station"],
|
|
}
|
|
|
|
|
|
@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 postcodes():
|
|
"""Two postcodes: one near central London, one far away."""
|
|
return pl.DataFrame(
|
|
{
|
|
"postcode": ["EC1A 1BB", "ZZ99 9ZZ"],
|
|
"lat": [51.5074, 55.0],
|
|
"lon": [-0.1278, -3.0],
|
|
}
|
|
)
|
|
|
|
|
|
def test_counts_pois_within_radius(postcodes, pois):
|
|
result = count_pois_per_postcode(postcodes, pois, groups=POI_GROUPS, radius_km=2.0)
|
|
|
|
expected_cols = {f"{g}_2km" for g in POI_GROUPS}
|
|
assert expected_cols.issubset(set(result.columns))
|
|
|
|
# Result must be aligned to postcodes (2 rows)
|
|
assert len(result) == 2
|
|
|
|
ec1a = result.filter(pl.col("postcode") == "EC1A 1BB")
|
|
assert ec1a["restaurants_2km"][0] == 2 # Restaurant + Fast Food
|
|
assert ec1a["groceries_2km"][0] == 1 # Supermarket
|
|
assert ec1a["parks_2km"][0] == 1 # Park
|
|
assert ec1a["public_transport_2km"][0] == 1 # Station
|
|
|
|
# Far-away postcode should have zero counts
|
|
zz99 = result.filter(pl.col("postcode") == "ZZ99 9ZZ")
|
|
for group in POI_GROUPS:
|
|
assert zz99[f"{group}_2km"][0] == 0
|
|
|
|
|
|
def test_no_pois_returns_zeros(postcodes):
|
|
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_per_postcode(postcodes, empty_pois, groups=POI_GROUPS, radius_km=2.0)
|
|
|
|
for group in POI_GROUPS:
|
|
col = f"{group}_2km"
|
|
assert col in result.columns
|
|
assert result[col].to_list() == [0, 0]
|
|
|
|
|
|
def test_custom_radius(pois):
|
|
"""A tiny radius should exclude POIs that are even slightly away."""
|
|
postcodes = 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_per_postcode(postcodes, pois, groups=POI_GROUPS, 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
|