41 lines
1.1 KiB
Python
41 lines
1.1 KiB
Python
import polars as pl
|
|
|
|
from pipeline.transform.poi_proximity import _build_poi_category_groups
|
|
|
|
|
|
def test_dynamic_poi_groups_include_requested_categories_only() -> None:
|
|
pois = pl.DataFrame(
|
|
{
|
|
"group": (
|
|
["Public Transport"] * 2
|
|
+ ["Leisure"] * 2
|
|
+ ["Groceries"] * 101
|
|
+ ["Groceries"] * 100
|
|
+ ["Education"] * 200
|
|
+ ["Health"] * 200
|
|
),
|
|
"category": (
|
|
["Rail station", "Bus stop"]
|
|
+ ["Café", "Restaurant"]
|
|
+ ["Tesco"] * 101
|
|
+ ["Waitrose"] * 100
|
|
+ ["School"] * 200
|
|
+ ["Pharmacy"] * 200
|
|
),
|
|
"lat": [51.5] * 605,
|
|
"lng": [-0.1] * 605,
|
|
}
|
|
)
|
|
|
|
groups, display_names = _build_poi_category_groups(pois)
|
|
|
|
assert set(display_names.values()) == {
|
|
"Bus stop",
|
|
"Café",
|
|
"Rail station",
|
|
"Restaurant",
|
|
"Tesco",
|
|
}
|
|
assert "poi_waitrose" not in groups
|
|
assert "poi_school" not in groups
|
|
assert "poi_pharmacy" not in groups
|