From 0c77d514a26443f5442de545f3718a50c9e0eabd Mon Sep 17 00:00:00 2001
From: Myst <1592048+LeMyst@users.noreply.github.com>
Date: Sun, 26 Sep 2021 01:23:01 +0200
Subject: [PATCH] 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
---
 test/test_wbi_helpers.py          | 24 +++++++++++++++++++-----
 wikibaseintegrator/wbi_helpers.py | 10 +++++-----
 2 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/test/test_wbi_helpers.py b/test/test_wbi_helpers.py
index d361164b..6a804314 100644
--- 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
index ac0c990d..9b3e1a89 100644
--- 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: