224 lines
7.0 KiB
Python
224 lines
7.0 KiB
Python
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
import requests
|
|
from celery.exceptions import Retry
|
|
from locations import utils as locations_utils
|
|
from locations.models import GeoLocation
|
|
from scrobbles.tasks import reverse_geocode_geolocation
|
|
|
|
LAT = 44.3652
|
|
LON = -70.5057
|
|
|
|
ADDRESS = {
|
|
"road": "Main St",
|
|
"city": "Springfield",
|
|
"state": "ME",
|
|
"postcode": "01101",
|
|
"country": "USA",
|
|
}
|
|
|
|
real_throttle = locations_utils._geocode_throttle
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _no_real_geocode_requests():
|
|
with patch("locations.utils._geocode_throttle"):
|
|
yield
|
|
|
|
|
|
def _mock_response(status_code=200, json_data=None, headers=None):
|
|
mock_resp = MagicMock()
|
|
mock_resp.status_code = status_code
|
|
mock_resp.headers = headers or {}
|
|
if json_data is not None:
|
|
mock_resp.json.return_value = json_data
|
|
if status_code >= 400:
|
|
mock_resp.raise_for_status.side_effect = requests.HTTPError(
|
|
f"{status_code} Client Error", response=mock_resp
|
|
)
|
|
return mock_resp
|
|
|
|
|
|
# --- reverse_geocode util ---
|
|
|
|
|
|
def test_reverse_geocode_success():
|
|
with patch(
|
|
"locations.utils.requests.get",
|
|
return_value=_mock_response(json_data={"address": ADDRESS}),
|
|
):
|
|
result = locations_utils.reverse_geocode(LAT, LON)
|
|
|
|
assert result == {
|
|
"street": "Main St",
|
|
"city": "Springfield",
|
|
"state_province": "ME",
|
|
"postal_code": "01101",
|
|
"country": "USA",
|
|
}
|
|
|
|
|
|
def test_reverse_geocode_429_raises_rate_limited():
|
|
mock_resp = _mock_response(status_code=429, headers={"Retry-After": "5"})
|
|
with patch("locations.utils.requests.get", return_value=mock_resp):
|
|
with pytest.raises(locations_utils.NominatimRateLimited) as excinfo:
|
|
locations_utils.reverse_geocode(LAT, LON)
|
|
|
|
assert excinfo.value.retry_after == 5
|
|
|
|
|
|
def test_reverse_geocode_429_without_retry_after():
|
|
mock_resp = _mock_response(status_code=429, headers={})
|
|
with patch("locations.utils.requests.get", return_value=mock_resp):
|
|
with pytest.raises(locations_utils.NominatimRateLimited) as excinfo:
|
|
locations_utils.reverse_geocode(LAT, LON)
|
|
|
|
assert excinfo.value.retry_after is None
|
|
|
|
|
|
def test_reverse_geocode_other_http_error_returns_none():
|
|
mock_resp = _mock_response(status_code=500)
|
|
with patch("locations.utils.requests.get", return_value=mock_resp):
|
|
assert locations_utils.reverse_geocode(LAT, LON) is None
|
|
|
|
|
|
def test_reverse_geocode_connection_error_returns_none():
|
|
with patch(
|
|
"locations.utils.requests.get",
|
|
side_effect=requests.ConnectionError("connection refused"),
|
|
):
|
|
assert locations_utils.reverse_geocode(LAT, LON) is None
|
|
|
|
|
|
# --- throttle ---
|
|
|
|
|
|
class _FakeRedis:
|
|
def __init__(self, count):
|
|
self.count = count
|
|
|
|
def incr(self, key):
|
|
self.count += 1
|
|
return self.count
|
|
|
|
def expire(self, key, ttl):
|
|
return True
|
|
|
|
|
|
@patch("locations.utils._geocode_throttle", real_throttle)
|
|
@patch("locations.utils._get_redis_client")
|
|
def test_geocode_throttle_raises_when_budget_exhausted(mock_client, settings):
|
|
settings.REDIS_URL = "redis://localhost:6379/0"
|
|
settings.GEOLOC_REVERSE_GEOCODE_RATE = 1
|
|
mock_client.return_value = _FakeRedis(count=1)
|
|
|
|
with pytest.raises(locations_utils.NominatimRateLimited) as excinfo:
|
|
locations_utils._geocode_throttle()
|
|
|
|
assert excinfo.value.retry_after == 1
|
|
|
|
|
|
@patch("locations.utils._geocode_throttle", real_throttle)
|
|
@patch("locations.utils._get_redis_client")
|
|
def test_geocode_throttle_passes_when_under_budget(mock_client, settings):
|
|
settings.REDIS_URL = "redis://localhost:6379/0"
|
|
settings.GEOLOC_REVERSE_GEOCODE_RATE = 1
|
|
mock_client.return_value = _FakeRedis(count=0)
|
|
|
|
assert locations_utils._geocode_throttle() is None
|
|
|
|
|
|
@patch("locations.utils._geocode_throttle", real_throttle)
|
|
def test_geocode_throttle_disabled_without_redis_url(settings):
|
|
settings.REDIS_URL = ""
|
|
settings.GEOLOC_REVERSE_GEOCODE_RATE = 1
|
|
|
|
assert locations_utils._geocode_throttle() is None
|
|
|
|
|
|
# --- celery task ---
|
|
|
|
|
|
@pytest.mark.django_db
|
|
def test_task_success_sets_postal_code():
|
|
location = GeoLocation.objects.create(lat=LAT, lon=LON)
|
|
with patch(
|
|
"locations.utils.requests.get",
|
|
return_value=_mock_response(json_data={"address": ADDRESS}),
|
|
):
|
|
reverse_geocode_geolocation.run(location.id)
|
|
|
|
location.refresh_from_db()
|
|
assert location.postal_code == "01101"
|
|
assert location.city == "Springfield"
|
|
|
|
|
|
@pytest.mark.django_db
|
|
def test_task_retries_with_retry_after(settings):
|
|
settings.CELERY_TASK_ALWAYS_EAGER = False
|
|
location = GeoLocation.objects.create(lat=LAT, lon=LON)
|
|
mock_resp = _mock_response(status_code=429, headers={"Retry-After": "30"})
|
|
|
|
mock_retry = MagicMock(side_effect=Retry())
|
|
fake_self = SimpleNamespace(request=SimpleNamespace(retries=0), retry=mock_retry)
|
|
task = reverse_geocode_geolocation
|
|
with patch("locations.utils.requests.get", return_value=mock_resp):
|
|
with pytest.raises(Retry):
|
|
task.run.__func__(fake_self, location.id)
|
|
|
|
mock_retry.assert_called_once()
|
|
countdown = mock_retry.call_args.kwargs["countdown"]
|
|
assert countdown == 30
|
|
location.refresh_from_db()
|
|
assert location.postal_code is None
|
|
|
|
|
|
@pytest.mark.django_db
|
|
def test_task_retry_backoff_without_retry_after(settings):
|
|
settings.CELERY_TASK_ALWAYS_EAGER = False
|
|
location = GeoLocation.objects.create(lat=LAT, lon=LON)
|
|
mock_resp = _mock_response(status_code=429, headers={})
|
|
|
|
mock_retry = MagicMock(side_effect=Retry())
|
|
fake_self = SimpleNamespace(request=SimpleNamespace(retries=2), retry=mock_retry)
|
|
task = reverse_geocode_geolocation
|
|
with patch("locations.utils.requests.get", return_value=mock_resp):
|
|
with pytest.raises(Retry):
|
|
task.run.__func__(fake_self, location.id)
|
|
|
|
countdown = mock_retry.call_args.kwargs["countdown"]
|
|
assert countdown >= 240
|
|
assert countdown <= 250
|
|
|
|
|
|
@pytest.mark.django_db
|
|
def test_task_gives_up_after_max_retries(settings):
|
|
settings.CELERY_TASK_ALWAYS_EAGER = False
|
|
settings.GEOLOC_REVERSE_GEOCODE_MAX_RETRIES = 3
|
|
location = GeoLocation.objects.create(lat=LAT, lon=LON)
|
|
mock_resp = _mock_response(status_code=429, headers={})
|
|
|
|
mock_retry = MagicMock(side_effect=Retry())
|
|
fake_self = SimpleNamespace(request=SimpleNamespace(retries=3), retry=mock_retry)
|
|
task = reverse_geocode_geolocation
|
|
with patch("locations.utils.requests.get", return_value=mock_resp):
|
|
task.run.__func__(fake_self, location.id)
|
|
|
|
mock_retry.assert_not_called()
|
|
location.refresh_from_db()
|
|
assert location.postal_code is None
|
|
|
|
|
|
@pytest.mark.django_db
|
|
def test_task_skips_when_postal_code_set():
|
|
location = GeoLocation.objects.create(
|
|
lat=LAT, lon=LON, postal_code="01101", city="Springfield"
|
|
)
|
|
with patch(
|
|
"locations.utils.requests.get",
|
|
side_effect=AssertionError("should not hit Nominatim"),
|
|
):
|
|
reverse_geocode_geolocation.run(location.id)
|