lgtm

pipeline/download/test_naptan.py

@@ -2,6 +2,7 @@ import polars as pl
 import pytest
 
 from pipeline.download.naptan import (
+    DLR_STATION_CATEGORY,
     TRAM_METRO_CATEGORY,
     TUBE_STATION_CATEGORY,
     canonical_station_name,
@@ -131,6 +132,48 @@ def test_deduplicate_naptan_splits_london_underground_from_tram_metro():
     assert tram["id"][0] == "9400ZZMAWST"
 
 
+def test_deduplicate_naptan_splits_dlr_from_tram_metro():
+    # DLR stations arrive as the tram/metro family (MET station node + TMU
+    # entrance). The Beckton group carries a 9400ZZDL station node, so the
+    # merged POI is reclassified to "DLR station" even though its entrance
+    # carries a non-ZZDL ATCO code; the Metrolink group stays "Tram & Metro
+    # stop".
+    df = pl.DataFrame(
+        {
+            "id": [
+                "9400ZZDLBEC",
+                "490000254S",
+                "9400ZZMAWST",
+            ],
+            "name": [
+                "Beckton DLR Station",
+                "Beckton",
+                "Weaste (Manchester Metrolink)",
+            ],
+            "category": [TRAM_METRO_CATEGORY] * 3,
+            "lat": [51.5148, 51.5148, 53.4826],
+            "lng": [0.0613, 0.0614, -2.3087],
+            "locality": [None, None, None],
+            "entrance": [False, True, False],
+            "is_lu": [False, False, False],
+            "is_dlr": [True, False, False],
+        }
+    )
+
+    result = deduplicate_naptan(df).sort("category")
+
+    assert len(result) == 2
+    assert result["category"].to_list() == [
+        DLR_STATION_CATEGORY,
+        TRAM_METRO_CATEGORY,
+    ]
+    dlr = result.filter(pl.col("category") == DLR_STATION_CATEGORY)
+    # The station node (not the entrance) represents the merged POI.
+    assert dlr["id"][0] == "9400ZZDLBEC"
+    tram = result.filter(pl.col("category") == TRAM_METRO_CATEGORY)
+    assert tram["id"][0] == "9400ZZMAWST"
+
+
 def test_deduplicate_naptan_merges_bus_station_bays_and_entrances():
     # BCS bays and a BCE entrance of one bus station collapse to a single POI
     # represented by a non-entrance node; a different bus station in another