diff --git a/app/router/locations.py b/app/router/locations.py index d0e03c46..af4b1cfd 100644 --- a/app/router/locations.py +++ b/app/router/locations.py @@ -1,4 +1,4 @@ -from fastapi import Request +from fastapi import HTTPException, Request from ..enums.sources import Sources from ..models.location import LocationResponse as Location @@ -39,6 +39,8 @@ def get_locations( locations = [location for location in locations if str(getattr(location, key)).lower() == str(value)] except AttributeError: pass + if not locations: + raise HTTPException(404, detail=f"Source `{source}` does not have the desired location data.") # Return final serialized data. return { diff --git a/tests/test_routes.py b/tests/test_routes.py index 7c4b1f03..9e1c03ef 100644 --- a/tests/test_routes.py +++ b/tests/test_routes.py @@ -120,20 +120,23 @@ def tearDown(self): @pytest.mark.parametrize( - "query_params", + "query_params,expected_status", [ - {"source": "csbs"}, - {"source": "jhu"}, - {"timelines": True}, - {"timelines": "true"}, - {"timelines": 1}, - {"source": "jhu", "timelines": True}, + ({"source": "csbs"}, 200), + ({"source": "jhu"}, 200), + ({"timelines": True}, 200), + ({"timelines": "true"}, 200), + ({"timelines": 1}, 200), + ({"source": "jhu", "timelines": True}, 200), + ({"source": "csbs", "country_code": "US"}, 200), + ({"source": "jhu", "country_code": "US"}, 404), ], ) -def test_locations_status_code(api_client, query_params): +def test_locations_status_code(api_client, query_params, expected_status): response = api_client.get("/v2/locations", params=query_params) print(f"GET {response.url}\n{response}") - assert response.status_code == 200 + print(f"\tjson:\n{pf(response.json())[:1000]}\n\t...") + assert response.status_code == expected_status @pytest.mark.parametrize(