More

pipeline/transform/test_transform_poi.py

@@ -0,0 +1,59 @@
+import polars as pl
+
+from pipeline.transform.transform_poi import transform_grocery_retail_points
+
+
+def test_transform_grocery_retail_points_outputs_chain_categories():
+    raw = pl.DataFrame(
+        {
+            "id": [101, 102, 103],
+            "retailer": ["Waitrose", "Sainsburys", "The Co-operative Group"],
+            "fascia": ["Waitrose", "Sainsbury's Local", "Co-op Food"],
+            "store_name": ["Waitrose Test", "Sainsbury''s Test", "Co-op Test"],
+            "long_wgs": [-0.141, -0.142, -0.143],
+            "lat_wgs": [51.515, 51.516, 51.517],
+        }
+    )
+
+    pois = transform_grocery_retail_points(raw)
+
+    assert pois.select("id", "name", "category", "group", "emoji").to_dicts() == [
+        {
+            "id": "glx-101",
+            "name": "Waitrose Test",
+            "category": "Waitrose",
+            "group": "Groceries",
+            "emoji": "🛒",
+        },
+        {
+            "id": "glx-102",
+            "name": "Sainsbury's Test",
+            "category": "Sainsbury's",
+            "group": "Groceries",
+            "emoji": "🛒",
+        },
+        {
+            "id": "glx-103",
+            "name": "Co-op Test",
+            "category": "Co-op",
+            "group": "Groceries",
+            "emoji": "🛒",
+        },
+    ]
+
+
+def test_transform_grocery_retail_points_drops_invalid_rows():
+    raw = pl.DataFrame(
+        {
+            "id": [101, 102],
+            "retailer": ["Waitrose", ""],
+            "fascia": ["Waitrose", "Tesco"],
+            "store_name": ["Waitrose Test", "Tesco Test"],
+            "long_wgs": [-0.141, -0.142],
+            "lat_wgs": [51.515, 51.516],
+        }
+    )
+
+    pois = transform_grocery_retail_points(raw)
+
+    assert pois["category"].to_list() == ["Waitrose"]