diff --git a/src/translink/crud.py b/src/translink/crud.py index 4886a72..b032502 100644 --- a/src/translink/crud.py +++ b/src/translink/crud.py @@ -25,6 +25,7 @@ STATIC_URL = "https://gtfs-static.translink.ca/gtfs/google_transit.zip" REALTIME_CACHE_ID = 1 REALTIME_CACHE_TTL_SECONDS = 90 +ARRIVED_GRACE_SECONDS = 60 REALTIME_CACHE_LOCK_ID = 2026062601 STATIC_CACHE_ID = 1 STATIC_CACHE_VERSION = 1 @@ -452,6 +453,8 @@ def _response_from_static_row(row: Any, delay: int = 0, status: BusStatus = BusS return [_response_from_static_row(row) for row in next_departures] # FeedMessage is generated at runtime, so the type checker can't find this function + now_ts = int(datetime.now(tz=TZ_INFO).timestamp()) + # Map all the realtime data to each bus's status realtime_map: dict[str, tuple[int, BusStatus]] = {} for entity in trip_feed.entity: @@ -473,8 +476,7 @@ def _response_from_static_row(row: Any, delay: int = 0, status: BusStatus = BusS if stop is None: continue - first_stop = min(trip_update.stop_time_update, key=lambda s: s.stop_sequence) - if first_stop.stop_id == stop_id: + if stop.departure.time <= now_ts + ARRIVED_GRACE_SECONDS: status = BusStatus.Arrived elif stop.departure.delay > 0: status = BusStatus.Delayed diff --git a/tests/unit/test_translink.py b/tests/unit/test_translink.py index 5dd2b72..32995cd 100644 --- a/tests/unit/test_translink.py +++ b/tests/unit/test_translink.py @@ -12,6 +12,7 @@ from constants import TZ_INFO from translink.crud import ( + ARRIVED_GRACE_SECONDS, BUS_DATA, STATIC_CACHE_UNAVAILABLE_MESSAGE, STATIC_CACHE_VERSION, @@ -655,6 +656,53 @@ async def test__get_departure_statuses_uses_timestamps_when_realtime_unavailable ] +@pytest.mark.parametrize( + ("departure_offset", "delay", "expected_status"), + [ + (-1, 0, BusStatus.Arrived), + (ARRIVED_GRACE_SECONDS - 1, 0, BusStatus.Arrived), + (ARRIVED_GRACE_SECONDS + 300, 120, BusStatus.Delayed), + (ARRIVED_GRACE_SECONDS + 300, 0, BusStatus.OnTime), + ], +) +async def test__get_departure_statuses_uses_realtime_departure_time_for_status( + departure_offset: int, + delay: int, + expected_status: BusStatus, +): + now = datetime.now(tz=TZ_INFO) + midnight = now.replace(hour=0, minute=0, second=0, microsecond=0) + feed = gtfs_realtime_pb2.FeedMessage() # pyright: ignore[reportAttributeAccessIssue] + feed.ParseFromString( + make_feed_bytes( + trip_id="trip_143", + route_id="6656", + direction_id=0, + stop_id="2836", + departure_unix=int(now.timestamp()) + departure_offset, + delay=delay, + ) + ) + static_schedule = [ + { + "trip_id": "trip_143", + "route_id": "6656", + "bus_number": "143", + "departure_time": "23:00:00", + "departure_seconds": int((now - midnight).total_seconds()) + 300, + } + ] + + with ( + patch("translink.crud.get_static_schedule", return_value=(now.date(), static_schedule)), + patch("translink.crud.get_or_fetch_realtime_feed", return_value=feed), + ): + result = await get_departure_statuses(mock_db_session(), AsyncMock(spec=AsyncClient)) + + assert len(result) == 1 + assert result[0].status == expected_status + + # --------------------------------------------------------------------------- # REST API endpoint tests # ---------------------------------------------------------------------------