added tests

src/traffic/core/flight.py

@@ -2598,16 +2598,18 @@ def get_toff_runway(
         threshold_alt: float = 1500,
         min_groundspeed_kts: float = 30,
         min_vert_rate_ftmin: float = 257,
+        max_bearing_deg: float = 10,
         max_dist_nm: float = 5,
     ) -> Optional[str]:
         """
         Determines the taking-off runway of a flight based on its trajectory.
         Parameters:
-            flight (Flight): The flight data object containing flight information.
-            airport: The IATA code of the airport. Defaults to "LSZH".
+            flight (Flight): The flight data object containing flight data.
+            airport: The ICAO code of the airport. Defaults to None.
             threshold_alt (float): Altitude threshold above airport altitude.
             min_groundspeed_kts (float): Minimum groundspeed to consider.
             min_vert_rate_ftmin (float): Minimum vertical rate to consider.
+            maximum_bearing_deg (float): Maximum bearing difference to consider.
             max_dist_nm (float): Maximum distance from the airport to consider.
 
         Returns:
@@ -2629,16 +2631,22 @@ def get_toff_runway(
         runways_names = runways.name
 
         filtered_flight = self.distance(airport).query(
-            "distance < @max_dist_nm"
+            f"distance < {max_dist_nm}"
         )
         if filtered_flight is None or filtered_flight.data.empty:
             return None
-
-        query_str = (
-            f"geoaltitude < {alt_max} and "
-            f"vertical_rate > {min_vert_rate_ftmin} and "
-            f"groundspeed > {min_groundspeed_kts}"
-        )
+        if "geoaltitude" in filtered_flight.data.columns:
+            query_str = (
+                f"geoaltitude < {alt_max} and "
+                f"vertical_rate > {min_vert_rate_ftmin} and "
+                f"groundspeed > {min_groundspeed_kts}"
+            )
+        else:
+            query_str = (
+                f"altitude < {alt_max} and "
+                f"vertical_rate > {min_vert_rate_ftmin} and "
+                f"groundspeed > {min_groundspeed_kts}"
+            )
         filtered_flight = filtered_flight.query(query_str)
         if (
             filtered_flight is None
@@ -2668,15 +2676,15 @@ def get_toff_runway(
                 rounded_bearings % 360
             )  # Ensure bearings stay within 0-359
 
-            # Find the bearing closest to the median track using minimal angular difference
+            # Find the bearing closest to the median track
             bearing_diff = rounded_bearings.apply(
                 lambda x: minimal_angular_difference(x, median_track)
             )
 
             # Identify all runways where bearing diff is less than 10 deg
 
             candidate_runways = runways.loc[
-                bearing_diff[bearing_diff < 10].index
+                bearing_diff[bearing_diff < max_bearing_deg].index
             ]
             if candidate_runways.empty:
                 return None
@@ -2696,7 +2704,7 @@ def get_toff_runway(
                 if closest is None:
                     return None
                 lon_1, _, lon_2, _ = closest.bounds
-                # find the candidate_runways latitude and longitude that is closest to the closest bounds
+                # candidate_runways = lat\lon that is closest to the closest bounds
                 eps = 1 / 1000
                 closest_runway = candidate_runways.query(
                     f"(abs(longitude-{lon_1})<{eps}) or (abs(longitude-{lon_2})<{eps})"

tests/test_flight.py

@@ -8,10 +8,6 @@
 import httpx
 import pytest
 from cartes.osm import Overpass
-
-import numpy as np
-import pandas as pd
-from pandas.testing import assert_frame_equal
 from traffic.algorithms.douglas_peucker import douglas_peucker
 from traffic.core import Flight, Traffic
 from traffic.data import airports, eurofirs, navaids, runways
@@ -25,6 +21,10 @@
     zurich_airport,
 )
 
+import numpy as np
+import pandas as pd
+from pandas.testing import assert_frame_equal
+
 # This part only serves on travis when the downloaded file is corrupted
 # This shouldn't happen much as caching is now activated.
 skip_runways = False
@@ -378,6 +378,39 @@ def test_takeoff_runway() -> None:
         assert takeoff is None or aligned.max("ILS") == takeoff.max("runway")
 
 
+def test_get_toff_runway_specific_cases() -> None:
+    # Airport ICAO codes
+    lszh = "LSZH"
+    lfpg = "LFPG"
+
+    # Flight SWR5220 - Should take off from runway 28 at LSZH
+    flight_swr5220 = zurich_airport["SWR5220"]
+    runway_swr5220 = flight_swr5220.get_toff_runway(lszh)
+    assert (
+        runway_swr5220 == "28"
+    ), f"SWR5220 should take off from runway 28, got {runway_swr5220}"
+
+    # Flight VJT796 - Should return None since it's a landing at LSZH
+    flight_vjt796 = zurich_airport["VJT796"]
+    runway_vjt796 = flight_vjt796.get_toff_runway(lszh)
+    assert (
+        runway_vjt796 is None
+    ), f"VJT796 should return None, got {runway_vjt796}"
+
+    # Flight ACA879 - Should take off from runway 16 at LSZH
+    flight_aca879 = zurich_airport["ACA879"]
+    runway_aca879 = flight_aca879.get_toff_runway(lszh)
+    assert (
+        runway_aca879 == "16"
+    ), f"ACA879 should take off from runway 16, got {runway_aca879}"
+
+    # Flight SWR5220 tested at LFPG - Should return None
+    runway_swr5220_lfpg = flight_swr5220.get_toff_runway(lfpg)
+    assert (
+        runway_swr5220_lfpg is None
+    ), f"SWR5220 should return None for LFPG, got {runway_swr5220_lfpg}"
+
+
 @pytest.mark.skipif(True, reason="only for local debug")
 def test_takeoff_goaround() -> None:
     from traffic.data.datasets import landing_zurich_2019