diff --git a/tests/test_wikidata_api.py b/tests/test_wikidata_api.py index ae14cbc..275d05b 100644 --- a/tests/test_wikidata_api.py +++ b/tests/test_wikidata_api.py @@ -1,8 +1,11 @@ import pytest import pytest_mock +import requests import responses from geocode.wikidata import APIResponseError, api_call +max_tries = 5 + @responses.activate def test_api_call_retries_on_failure(mocker: pytest_mock.plugin.MockerFixture) -> None: @@ -20,9 +23,9 @@ def test_api_call_retries_on_failure(mocker: pytest_mock.plugin.MockerFixture) - ) with pytest.raises(APIResponseError): api_call({"action": "wbgetentities", "ids": "Q42"}) - assert len(responses.calls) == 5 # Assuming max_tries is 5 + assert len(responses.calls) == max_tries - assert mocked_sleep.call_count == 4 + assert mocked_sleep.call_count == max_tries - 1 mock_send_mail.assert_called() @@ -31,3 +34,20 @@ def test_api_call_retries_on_failure(mocker: pytest_mock.plugin.MockerFixture) - "Geocode error", "Error making Wikidata API call\n\nbad request", ) + + +def test_api_call_retries_on_connection_error( + mocker: pytest_mock.plugin.MockerFixture, +) -> None: + """Test retry for API calls on connection error.""" + # Patch 'time.sleep' to instantly return, effectively skipping the sleep + mocked_sleep = mocker.patch("time.sleep", return_value=None) + + # Patch 'requests.get' to raise a ConnectionError + mocker.patch("requests.get", side_effect=requests.exceptions.ConnectionError) + mocker.patch("geocode.mail.send_to_admin") + + with pytest.raises(requests.exceptions.ConnectionError): + api_call({"action": "wbgetentities", "ids": "Q42"}) + + assert mocked_sleep.call_count == max_tries - 1