Add retry for HTTP Code 500, 502 and 504 (#233)

* Add retry for HTTP Code 500, 502 and 504 Add unit tests Fix #232 * Add the same for sparql requests
LeMyst · Sep 25, 2021 · 0c77d51 · 0c77d51
1 parent 4764ce2
commit 0c77d51
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 10 deletions.
diff --git a/test/test_wbi_helpers.py b/test/test_wbi_helpers.py
@@ -7,13 +7,27 @@
 
 
 def test_connection():
+    data = {'format': 'json', 'action': 'wbgetentities', 'ids': 'Q42'}
+
+    mediawiki_api_call_helper(data=data, max_retries=2, retry_after=1, allow_anonymous=True)
+
     with unittest.TestCase().assertRaises(MWApiError):
-        mediawiki_api_call_helper(data={'format': 'json', 'action': 'wbgetentities', 'ids': 'Q42'}, mediawiki_api_url="https://www.wikidataaaaaaa.org", max_retries=3,
-                                  retry_after=1, allow_anonymous=True)
-    with unittest.TestCase().assertRaises(requests.HTTPError):
-        mediawiki_api_call_helper(data=None, mediawiki_api_url="https://httpbin.org/status/400", max_retries=3, retry_after=1, allow_anonymous=True)
+        mediawiki_api_call_helper(data=data, mediawiki_api_url="https://www.wikidataaaaaaa.org", max_retries=2, retry_after=1, allow_anonymous=True)
 
-    mediawiki_api_call_helper(data={'format': 'json', 'action': 'wbgetentities', 'ids': 'Q42'}, max_retries=3, retry_after=1, allow_anonymous=True)
+    with unittest.TestCase().assertRaises(MWApiError):
+        mediawiki_api_call_helper(data=data, mediawiki_api_url="https://httpbin.org/status/500", max_retries=2, retry_after=1, allow_anonymous=True)
+
+    with unittest.TestCase().assertRaises(MWApiError):
+        mediawiki_api_call_helper(data=data, mediawiki_api_url="https://httpbin.org/status/502", max_retries=2, retry_after=1, allow_anonymous=True)
+
+    with unittest.TestCase().assertRaises(MWApiError):
+        mediawiki_api_call_helper(data=data, mediawiki_api_url="https://httpbin.org/status/503", max_retries=2, retry_after=1, allow_anonymous=True)
+
+    with unittest.TestCase().assertRaises(MWApiError):
+        mediawiki_api_call_helper(data=data, mediawiki_api_url="https://httpbin.org/status/504", max_retries=2, retry_after=1, allow_anonymous=True)
+
+    with unittest.TestCase().assertRaises(requests.HTTPError):
+        mediawiki_api_call_helper(data=data, mediawiki_api_url="https://httpbin.org/status/400", max_retries=2, retry_after=1, allow_anonymous=True)
 
 
 def test_user_agent(capfd):

diff --git a/wikibaseintegrator/wbi_helpers.py b/wikibaseintegrator/wbi_helpers.py
@@ -21,7 +21,7 @@ class BColors:
     UNDERLINE = '\033[4m'
 
 
-def mediawiki_api_call(method, mediawiki_api_url=None, session=None, max_retries=1000, retry_after=60, **kwargs):
+def mediawiki_api_call(method, mediawiki_api_url=None, session=None, max_retries=100, retry_after=60, **kwargs):
     """
     :param method: 'GET' or 'POST'
     :param mediawiki_api_url:
@@ -53,8 +53,8 @@ def mediawiki_api_call(method, mediawiki_api_url=None, session=None, max_retries
             print(f"Connection error: {e}. Sleeping for {retry_after} seconds.")
             sleep(retry_after)
             continue
-        if response.status_code == 503:  # pragma: no cover
-            print(f"service unavailable. sleeping for {retry_after} seconds")
+        if response.status_code in (500, 502, 503, 504):
+            print(f"Service unavailable (HTTP Code {response.status_code}). Sleeping for {retry_after} seconds.")
             sleep(retry_after)
             continue
 
@@ -205,8 +205,8 @@ def execute_sparql_query(query, prefix=None, endpoint=None, user_agent=None, max
             print(f"Connection error: {e}. Sleeping for {retry_after} seconds.")
             sleep(retry_after)
             continue
-        if response.status_code == 503:
-            print(f"Service unavailable (503). Sleeping for {retry_after} seconds")
+        if response.status_code in (500, 502, 503, 504):
+            print(f"Service unavailable (HTTP Code {response.status_code}). Sleeping for {retry_after} seconds.")
             sleep(retry_after)
             continue
         if response.status_code == 429: