diff --git a/marta/api.py b/marta/api.py index fda7e35..df8739e 100644 --- a/marta/api.py +++ b/marta/api.py @@ -1,21 +1,29 @@ +import json +from functools import wraps +from os import getenv from typing import List, Union -import requests import requests_cache -import json -from os import getenv -from functools import wraps from .exceptions import APIKeyError, InvalidDirectionError from .vehicles import Bus, Train _CACHE_EXPIRE = int(getenv('MARTA_CACHE_EXPIRE', 30)) -_BASE_URL = 'http://developer.itsmarta.com' +_BASE_URL = 'https://developer.itsmarta.com' _TRAIN_PATH = '/RealtimeTrain/RestServiceNextTrain/GetRealtimeArrivals' _BUS_PATH = '/BRDRestService/RestBusRealTimeService/GetAllBus' _BUS_ROUTE_PATH = '/BRDRestService/RestBusRealTimeService/GetBusByRoute/' - -requests_cache.install_cache('marta_api_cache', backend='sqlite', expire_after=_CACHE_EXPIRE) +_REQUEST_TIMEOUT = (3.05, 10) +_MAX_RESPONSE_BYTES = 5 * 1024 * 1024 +_JSON_CONTENT_TYPES = {'application/json', 'text/json'} + +# Keep caching private to this library and in memory. Importing marta must not +# patch requests globally or persist request URLs containing credentials. +CACHE = requests_cache.CachedSession( + backend='memory', + expire_after=_CACHE_EXPIRE, + ignored_parameters=['apikey'], +) def require_api_key(func): """ @@ -31,41 +39,84 @@ def with_key(self, *args, **kwargs): return with_key +def get_bus_direction(user_direction) -> str: + direction = user_direction.lower() + if direction.startswith('n'): + return 'Northbound' + elif direction.startswith('s'): + return 'Southbound' + elif direction.startswith('e'): + return 'Eastbound' + elif direction.startswith('w'): + return 'Westbound' + else: + return None + +def get_train_direction(user_direction) -> str: + train_direction = user_direction.lower() + if train_direction.startswith('n'): + return 'N' + elif train_direction.startswith('s'): + return 'S' + elif train_direction.startswith('e'): + return 'E' + elif train_direction.startswith('w'): + return 'W' + else: + return None + + def _convert_direction(user_direction: str, vehicle_type: str = 'bus') -> Union[str, None]: if not user_direction: return None if vehicle_type == 'bus': - if user_direction.lower().startswith('n'): - return 'Northbound' - elif user_direction.lower().startswith('s'): - return 'Southbound' - elif user_direction.lower().startswith('e'): - return 'Eastbound' - elif user_direction.lower().startswith('w'): - return 'Westbound' + bus_direction = get_bus_direction(user_direction) + if bus_direction is not None: + return bus_direction else: raise InvalidDirectionError(direction_provided=user_direction) elif vehicle_type == 'train': - if user_direction.lower().startswith('n'): - return 'N' - elif user_direction.lower().startswith('s'): - return 'S' - elif user_direction.lower().startswith('e'): - return 'E' - elif user_direction.lower().startswith('w'): - return 'W' + train_direction = get_train_direction(user_direction) + if train_direction is not None: + return train_direction else: raise InvalidDirectionError(direction_provided=user_direction) else: return user_direction -def _get_data(endpoint: str, api_key: str) -> dict: - url = f'{_BASE_URL}{endpoint}?apikey={api_key}' - response = requests.get(url) - if response.status_code == 401 or response.status_code == 403: - raise APIKeyError(f'Your API key seems to be invalid. Try visiting {url}.') - return json.loads(response.text) +def _get_data(endpoint: str, api_key: str) -> List[dict]: + url = f'{_BASE_URL}{endpoint}' + response = CACHE.get( + url, + params={'apikey': api_key}, + timeout=_REQUEST_TIMEOUT, + stream=True, + ) + if response.status_code in (401, 403): + raise APIKeyError('Your MARTA API key was rejected.') + response.raise_for_status() + + content_type = response.headers.get('Content-Type', '').split(';', 1)[0].lower() + if content_type and content_type not in _JSON_CONTENT_TYPES: + raise ValueError(f'Unexpected MARTA API content type: {content_type}') + + content_length = response.headers.get('Content-Length') + if content_length and int(content_length) > _MAX_RESPONSE_BYTES: + raise ValueError('MARTA API response exceeds the maximum allowed size') + + chunks = [] + received = 0 + for chunk in response.iter_content(chunk_size=64 * 1024): + received += len(chunk) + if received > _MAX_RESPONSE_BYTES: + raise ValueError('MARTA API response exceeds the maximum allowed size') + chunks.append(chunk) + + data = json.loads(b''.join(chunks).decode(response.encoding or 'utf-8')) + if not isinstance(data, list) or any(not isinstance(item, dict) for item in data): + raise ValueError('MARTA API response must be a list of objects') + return data def _filter_response(response: dict, filters: dict) -> List[dict]: valid_items = [] diff --git a/marta/exceptions.py b/marta/exceptions.py index dab3668..65a08de 100644 --- a/marta/exceptions.py +++ b/marta/exceptions.py @@ -4,7 +4,7 @@ def __init__(self, message: str = None): if not message: message = 'API Key is missing. Please set MARTA_API_KEY or use api_key kwarg.' - super(Exception, self).__init__(message) + super().__init__(message) class InvalidDirectionError(Exception): """Exception thrown for an invalid bus/train direction""" @@ -12,4 +12,4 @@ def __init__(self, direction_provided: str, message: str = None): if not message: message = f'{direction_provided} is an invalid direction.' - super(Exception, self).__init__(message) \ No newline at end of file + super().__init__(message) diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..c7b23ec --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +pythonpath = . +testpaths = tests diff --git a/requirements.txt b/requirements.txt index 6493c2e..9a6e556 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,2 @@ -requests -requests-cache +requests>=2.32.4,<3 +requests-cache>=1.2,<2 diff --git a/setup.py b/setup.py index fac285d..97db2bf 100644 --- a/setup.py +++ b/setup.py @@ -3,11 +3,11 @@ setup( name='marta', description='Python library for accessing MARTA real-time API', - url='http://www.itsmarta.com/app-developer-resources.aspx', + url='https://www.itsmarta.com/app-developer-resources.aspx', packages=['marta'], install_requires=[ 'pytest-runner', - 'requests', - 'requests-cache' + 'requests>=2.32.4,<3', + 'requests-cache>=1.2,<2' ], ) diff --git a/tests/test_api.py b/tests/test_api.py index f35e453..bd27ff7 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1,82 +1,176 @@ +import json + import pytest -import requests_mock +import requests -import marta -from marta.api import get_buses, get_trains +from marta.api import MARTA, _get_data from marta.exceptions import APIKeyError from marta.vehicles import Bus, Train -def test_get_trains(train_response): - with requests_mock.mock() as m: - m.get(requests_mock.ANY, text=train_response) - trains = get_trains() - for t in trains: - assert isinstance(t, Train) +class FakeResponse: + def __init__(self, payload=None, status_code=200, headers=None, chunks=None): + self.payload = payload + self.status_code = status_code + self.headers = ( + {'Content-Type': 'application/json'} if headers is None else headers + ) + self.encoding = 'utf-8' + self._chunks = chunks + + def raise_for_status(self): + if self.status_code >= 400: + raise requests.HTTPError() + + def iter_content(self, chunk_size): + if self._chunks is not None: + yield from self._chunks + return + yield json.dumps(self.payload).encode() + + +@pytest.fixture +def mock_response(monkeypatch): + calls = [] + + def install(response): + def get(url, **kwargs): + calls.append((url, kwargs)) + return response + + monkeypatch.setattr('marta.api.CACHE.get', get) + return calls + + return install + + +def test_get_trains(mock_response, train_response): + calls = mock_response(FakeResponse(json.loads(train_response))) + + trains = MARTA(api_key='secret').get_trains() + + assert trains + assert all(isinstance(train, Train) for train in trains) + url, kwargs = calls[0] + assert url == ( + 'https://developer.itsmarta.com/' + 'RealtimeTrain/RestServiceNextTrain/GetRealtimeArrivals' + ) + assert 'secret' not in url + assert kwargs == { + 'params': {'apikey': 'secret'}, + 'timeout': (3.05, 10), + 'stream': True, + } + + +@pytest.mark.parametrize( + ('filters', 'attribute', 'expected'), + [ + ({'line': 'BLUE'}, 'line', 'BLUE'), + ({'station': 'INDIAN CREEK STATION'}, 'station', 'INDIAN CREEK STATION'), + ], +) +def test_get_trains_filters_locally( + mock_response, train_response, filters, attribute, expected +): + mock_response(FakeResponse(json.loads(train_response))) + + trains = MARTA(api_key='secret').get_trains(**filters) + + assert trains + assert all(getattr(train, attribute) == expected for train in trains) + + +def test_get_trains_combines_filters(mock_response, train_response): + mock_response(FakeResponse(json.loads(train_response))) + + trains = MARTA(api_key='secret').get_trains( + line='BLUE', station='INDIAN CREEK STATION' + ) + + assert len(trains) == 1 + assert trains[0].line == 'BLUE' + assert trains[0].station == 'INDIAN CREEK STATION' + + +def test_get_buses(mock_response, bus_all_response): + calls = mock_response(FakeResponse(json.loads(bus_all_response))) + + buses = MARTA(api_key='secret').get_buses() + + assert buses + assert all(isinstance(bus, Bus) for bus in buses) + assert calls[0][0] == ( + 'https://developer.itsmarta.com/' + 'BRDRestService/RestBusRealTimeService/GetAllBus' + ) + + +def test_get_buses_by_route(mock_response, bus_route_response): + mock_response(FakeResponse(json.loads(bus_route_response))) + + buses = MARTA(api_key='secret').get_buses(route=1) + + assert buses + assert all(bus.route == 1 for bus in buses) + + +def test_missing_api_key_is_deterministic(monkeypatch): + monkeypatch.delenv('MARTA_API_KEY', raising=False) + + with pytest.raises(APIKeyError, match='API Key is missing'): + MARTA().get_buses() -def test_get_trains_by_line(train_response): - line = "BLUE" - with requests_mock.mock() as m: - m.get(requests_mock.ANY, text=train_response) - trains = get_trains(line=line) +@pytest.mark.parametrize('status_code', [401, 403]) +def test_rejected_key_is_not_disclosed(mock_response, status_code): + mock_response(FakeResponse(status_code=status_code)) - assert len(trains) > 0 - for t in trains: - assert isinstance(t, Train) - assert t.line == line + with pytest.raises(APIKeyError) as error: + _get_data('/endpoint', 'super-secret') + assert 'super-secret' not in str(error.value) + assert str(error.value) == 'Your MARTA API key was rejected.' -def test_get_trains_by_station(train_response): - station = "INDIAN CREEK STATION" - with requests_mock.mock() as m: - m.get(requests_mock.ANY, text=train_response) - trains = get_trains(station=station) - assert len(trains) > 0 - for t in trains: - assert isinstance(t, Train) - assert t.station == station +def test_raises_for_other_http_errors(mock_response): + mock_response(FakeResponse(status_code=500)) + with pytest.raises(requests.HTTPError): + _get_data('/endpoint', 'secret') -def test_get_trains_by_line_and_station(train_response): - station = "INDIAN CREEK STATION" - line = "BLUE" - with requests_mock.mock() as m: - m.get(requests_mock.ANY, text=train_response) - trains = get_trains(line=line, station=station) - assert len(trains) > 0 - for t in trains: - assert isinstance(t, Train) - assert t.station == station - assert t.line == line +def test_rejects_non_json_response(mock_response): + mock_response(FakeResponse(headers={'Content-Type': 'text/html'})) + with pytest.raises(ValueError, match='content type'): + _get_data('/endpoint', 'secret') -def test_get_buses(bus_all_response): - with requests_mock.mock() as m: - m.get(requests_mock.ANY, text=bus_all_response) - buses = get_buses() - assert len(buses) > 0 +def test_rejects_oversized_content_length(mock_response): + mock_response( + FakeResponse( + headers={ + 'Content-Type': 'application/json', + 'Content-Length': str(5 * 1024 * 1024 + 1), + } + ) + ) - for b in buses: - assert isinstance(b, Bus) + with pytest.raises(ValueError, match='maximum allowed size'): + _get_data('/endpoint', 'secret') -def test_get_buses_by_route(bus_route_response): - with requests_mock.mock() as m: - m.get(requests_mock.ANY, text=bus_route_response) - buses = get_buses(route=1) +def test_rejects_oversized_stream(mock_response): + mock_response(FakeResponse(chunks=[b'x' * (5 * 1024 * 1024 + 1)])) - assert len(buses) > 0 + with pytest.raises(ValueError, match='maximum allowed size'): + _get_data('/endpoint', 'secret') - for b in buses: - assert isinstance(b, Bus) - assert b.route == 1 +def test_rejects_unexpected_json_shape(mock_response): + mock_response(FakeResponse({'not': 'a list'})) -def test_missing_api_key(): - marta.api._API_KEY = None - with pytest.raises(APIKeyError): - buses = get_buses() + with pytest.raises(ValueError, match='list of objects'): + _get_data('/endpoint', 'secret')