diff --git a/geocode/wikidata.py b/geocode/wikidata.py index 441d460..ec87711 100644 --- a/geocode/wikidata.py +++ b/geocode/wikidata.py @@ -106,6 +106,7 @@ def qid_to_commons_category(qid: str, check_p910: bool = True) -> str | None: Row = dict[str, dict[str, typing.Any]] +@backoff.on_exception(backoff.expo, QueryError, max_tries=5) def wdqs(query: str) -> list[Row]: """Pass query to the Wikidata Query Service.""" r = requests.post( diff --git a/tests/test_wikidata_api.py b/tests/test_wikidata_api.py index 275d05b..9877e2a 100644 --- a/tests/test_wikidata_api.py +++ b/tests/test_wikidata_api.py @@ -2,7 +2,7 @@ import pytest import pytest_mock import requests import responses -from geocode.wikidata import APIResponseError, api_call +from geocode.wikidata import APIResponseError, QueryError, api_call, wdqs max_tries = 5 @@ -51,3 +51,22 @@ def test_api_call_retries_on_connection_error( api_call({"action": "wbgetentities", "ids": "Q42"}) assert mocked_sleep.call_count == max_tries - 1 + + +def test_wdqs_retry(mocker: pytest_mock.plugin.MockerFixture) -> None: + """Test retry for WDQS API calls.""" + # Patch 'time.sleep' to instantly return, effectively skipping the sleep + mocked_sleep = mocker.patch("time.sleep", return_value=None) + + responses.add( + responses.POST, + "https://query.wikidata.org/bigdata/namespace/wdq/sparql", + body="bad request", + status=400, + ) + + with pytest.raises(QueryError): + wdqs("test query") + + max_tries = 5 + assert mocked_sleep.call_count == max_tries - 1