diff --git a/py_nextbus/client.py b/py_nextbus/client.py index 5b57b49..b2fa56e 100644 --- a/py_nextbus/client.py +++ b/py_nextbus/client.py @@ -89,10 +89,9 @@ def route_details( def predictions_for_stop( self, stop_id: str | int, - route_id: str, + route_id: str | None = None, direction_id: str | None = None, agency_id: str | None = None, - unfiltered: bool = False, ) -> list[dict[str, Any]]: agency_id = agency_id or self.agency_id if not agency_id: @@ -100,17 +99,23 @@ def predictions_for_stop( params: dict[str, Any] = {"coincident": True} if direction_id: + if not route_id: + raise NextBusValidationError("Direction ID provided without route ID") params["direction"] = direction_id + route_component = "" + if route_id: + route_component = f"routes/{route_id}/" + result = self._get( - f"agencies/{agency_id}/routes/{route_id}/stops/{stop_id}/predictions", + f"agencies/{agency_id}/{route_component}stops/{stop_id}/predictions", params, ) predictions = cast(list[dict[str, Any]], result) - # If unfiltered, return all predictions as the API returned them - if unfiltered: + # If route not provided, return all predictions as the API returned them + if not route_id: return predictions # HACK: Filter predictions based on stop and route because the API seems to ignore the route @@ -124,7 +129,7 @@ def predictions_for_stop( ] # HACK: Filter predictions based on direction in case the API returns extra predictions - if direction_id is not None: + if direction_id: for prediction_result in predictions: prediction_result["values"] = [ prediction diff --git a/setup.py b/setup.py index d78d44c..d870627 100644 --- a/setup.py +++ b/setup.py @@ -6,7 +6,7 @@ setup( name="py_nextbusnext", - version="2.0.4", + version="2.0.5", author="ViViDboarder", description="Minimalistic Python client for the NextBus public API for real-time transit " "arrival data", diff --git a/tests/client_test.py b/tests/client_test.py index 9f88625..16e6562 100644 --- a/tests/client_test.py +++ b/tests/client_test.py @@ -3,7 +3,8 @@ import unittest.mock from py_nextbus.client import NextBusClient -from tests.helpers.mock_responses import MOCK_PREDICTIONS_RESPONSE +from tests.helpers.mock_responses import MOCK_PREDICTIONS_RESPONSE_NO_ROUTE +from tests.helpers.mock_responses import MOCK_PREDICTIONS_RESPONSE_WITH_ROUTE from tests.helpers.mock_responses import TEST_AGENCY_ID from tests.helpers.mock_responses import TEST_DIRECTION_ID from tests.helpers.mock_responses import TEST_ROUTE_ID @@ -16,8 +17,26 @@ def setUp(self): self.client = NextBusClient() @unittest.mock.patch("py_nextbus.client.NextBusClient._get") - def test_predictions_for_stop(self, mock_get): - mock_get.return_value = MOCK_PREDICTIONS_RESPONSE + def test_predictions_for_stop_no_route(self, mock_get): + mock_get.return_value = MOCK_PREDICTIONS_RESPONSE_NO_ROUTE + + result = self.client.predictions_for_stop( + TEST_STOP_ID, + agency_id=TEST_AGENCY_ID + ) + + self.assertEqual({r["stop"]["id"] for r in result}, {TEST_STOP_ID}) + self.assertEqual(len(result), 3) # Results include all routes + + mock_get.assert_called_once() + mock_get.assert_called_with( + f"agencies/{TEST_AGENCY_ID}/stops/{TEST_STOP_ID}/predictions", + {"coincident": True}, + ) + + @unittest.mock.patch("py_nextbus.client.NextBusClient._get") + def test_predictions_for_stop_with_route(self, mock_get): + mock_get.return_value = MOCK_PREDICTIONS_RESPONSE_WITH_ROUTE result = self.client.predictions_for_stop( TEST_STOP_ID, diff --git a/tests/helpers/mock_responses.py b/tests/helpers/mock_responses.py index de78f7c..986869b 100644 --- a/tests/helpers/mock_responses.py +++ b/tests/helpers/mock_responses.py @@ -105,7 +105,143 @@ "timestamp": "2024-06-23T03:06:58Z", } -MOCK_PREDICTIONS_RESPONSE = [ +MOCK_PREDICTIONS_RESPONSE_NO_ROUTE = [ + { + "serverTimestamp": 1724038210798, + "nxbs2RedirectUrl": "", + "route": { + "id": "LOWL", + "title": "Lowl Owl Taraval", + "description": "10pm-5am nightly", + "color": "666666", + "textColor": "ffffff", + "hidden": False + }, + "stop": { + "id": "5184", + "lat": 37.8071299, + "lon": -122.41732, + "name": "Jones St & Beach St", + "code": "15184", + "hidden": False, + "showDestinationSelector": True, + "route": "LOWL" + }, + "values": [] + }, + { + "serverTimestamp": 1724038210798, + "nxbs2RedirectUrl": "", + "route": { + "id": "FBUS", + "title": "Fbus Market & Wharves", + "description": "", + "color": "b49a36", + "textColor": "000000", + "hidden": False + }, + "stop": { + "id": "5184", + "lat": 37.8071299, + "lon": -122.41732, + "name": "Jones St & Beach St", + "code": "15184", + "hidden": False, + "showDestinationSelector": True, + "route": "FBUS" + }, + "values": [] + }, + { + "serverTimestamp": 1724038210798, + "nxbs2RedirectUrl": "", + "route": { + "id": "F", + "title": "F Market & Wharves", + "description": "7am-10pm daily", + "color": "b49a36", + "textColor": "000000", + "hidden": False + }, + "stop": { + "id": "5184", + "lat": 37.8071299, + "lon": -122.41732, + "name": "Jones St & Beach St", + "code": "15184", + "hidden": False, + "showDestinationSelector": True, + "route": "F" + }, + "values": [ + { + "timestamp": 1724038309178, + "minutes": 1, + "affectedByLayover": True, + "isDeparture": True, + "occupancyStatus": -1, + "occupancyDescription": "Unknown", + "vehiclesInConsist": 1, + "linkedVehicleIds": "1078", + "vehicleId": "1078", + "vehicleType": "Historic Street Car_VC1", + "direction": { + "id": "F_0_var1", + "name": "Castro + Market", + "destinationName": "Castro + Market" + }, + "tripId": "11593249_M13", + "delay": 0, + "predUsingNavigationTm": False, + "departure": True + }, + { + "timestamp": 1724039160000, + "minutes": 15, + "affectedByLayover": True, + "isDeparture": True, + "occupancyStatus": -1, + "occupancyDescription": "Unknown", + "vehiclesInConsist": 1, + "linkedVehicleIds": "1080", + "vehicleId": "1080", + "vehicleType": "Historic Street Car_VC1", + "direction": { + "id": "F_0_var0", + "name": "Castro", + "destinationName": "Castro" + }, + "tripId": "11593252_M13", + "delay": 0, + "predUsingNavigationTm": False, + "departure": True + }, + { + "timestamp": 1724041320000, + "minutes": 51, + "affectedByLayover": True, + "isDeparture": True, + "occupancyStatus": -1, + "occupancyDescription": "Unknown", + "vehiclesInConsist": 1, + "linkedVehicleIds": "1056", + "vehicleId": "1056", + "vehicleType": "Historic Street Car_VC1", + "direction": { + "id": "F_0_var1", + "name": "Castro + Market", + "destinationName": "Castro + Market" + }, + "tripId": "11593256_M13", + "delay": 0, + "predUsingNavigationTm": False, + "departure": True + } + ] + } +] + +MOCK_PREDICTIONS_RESPONSE_WITH_ROUTE = [ { "serverTimestamp": 1720034290432, "nxbs2RedirectUrl": "",