fix: add 52x range to retry transient failures and tests

Clayton Walker · Clayton Walker · commit 49286751eff1 · 2022-03-14T11:43:41.000-06:00
diff --git a/gitlab/client.py b/gitlab/client.py
@@ -694,9 +694,9 @@ def http_request(
                 )
             except requests.ConnectionError:
                 if retry_transient_errors and (
-                        max_retries == -1 or cur_retries < max_retries
+                    max_retries == -1 or cur_retries < max_retries
                 ):
-                    wait_time = 2 ** cur_retries * 0.1
+                    wait_time = 2**cur_retries * 0.1
                     cur_retries += 1
                     time.sleep(wait_time)
                     continue
@@ -712,7 +712,11 @@ def http_request(
                 "retry_transient_errors", self.retry_transient_errors
             )
             if (429 == result.status_code and obey_rate_limit) or (
-                result.status_code in [500, 502, 503, 504] and retry_transient_errors
+                (
+                    result.status_code in [500, 502, 503, 504]
+                    or 520 <= result.status_code <= 530
+                )
+                and retry_transient_errors
             ):
                 # Response headers documentation:
                 # https://docs.gitlab.com/ee/user/admin_area/settings/user_and_ip_rate_limits.html#response-headers
diff --git a/tests/unit/test_gitlab_http_methods.py b/tests/unit/test_gitlab_http_methods.py
@@ -51,7 +51,7 @@ def test_http_request_404(gl):
 
 
 @responses.activate
-@pytest.mark.parametrize("status_code", [500, 502, 503, 504])
+@pytest.mark.parametrize("status_code", [500, 502, 503, 504] + list(range(520, 531)))
 def test_http_request_with_only_failures(gl, status_code):
     url = "http://localhost/api/v4/projects"
     responses.add(
@@ -97,6 +97,37 @@ def request_callback(request):
     assert len(responses.calls) == calls_before_success
 
 
+@responses.activate
+def test_http_request_with_retry_on_method_for_transient_network_failures(gl):
+    call_count = 0
+    calls_before_success = 3
+
+    url = "http://localhost/api/v4/projects"
+
+    def request_callback(request):
+        nonlocal call_count
+        call_count += 1
+        status_code = 200
+        headers = {}
+        body = "[]"
+
+        if call_count >= calls_before_success:
+            return (status_code, headers, body)
+        raise requests.ConnectionError("Connection aborted.")
+
+    responses.add_callback(
+        method=responses.GET,
+        url=url,
+        callback=request_callback,
+        content_type="application/json",
+    )
+
+    http_r = gl.http_request("get", "/projects", retry_transient_errors=True)
+
+    assert http_r.status_code == 200
+    assert len(responses.calls) == calls_before_success
+
+
 @responses.activate
 def test_http_request_with_retry_on_class_for_transient_failures(gl_retry):
     call_count = 0
@@ -126,6 +157,37 @@ def request_callback(request: requests.models.PreparedRequest):
     assert len(responses.calls) == calls_before_success
 
 
+@responses.activate
+def test_http_request_with_retry_on_class_for_transient_network_failures(gl_retry):
+    call_count = 0
+    calls_before_success = 3
+
+    url = "http://localhost/api/v4/projects"
+
+    def request_callback(request: requests.models.PreparedRequest):
+        nonlocal call_count
+        call_count += 1
+        status_code = 200
+        headers = {}
+        body = "[]"
+
+        if call_count >= calls_before_success:
+            return (status_code, headers, body)
+        raise requests.ConnectionError("Connection aborted.")
+
+    responses.add_callback(
+        method=responses.GET,
+        url=url,
+        callback=request_callback,
+        content_type="application/json",
+    )
+
+    http_r = gl_retry.http_request("get", "/projects", retry_transient_errors=True)
+
+    assert http_r.status_code == 200
+    assert len(responses.calls) == calls_before_success
+
+
 @responses.activate
 def test_http_request_with_retry_on_class_and_method_for_transient_failures(gl_retry):
     call_count = 0
@@ -155,6 +217,39 @@ def request_callback(request):
     assert len(responses.calls) == 1
 
 
+@responses.activate
+def test_http_request_with_retry_on_class_and_method_for_transient_network_failures(
+    gl_retry,
+):
+    call_count = 0
+    calls_before_success = 3
+
+    url = "http://localhost/api/v4/projects"
+
+    def request_callback(request):
+        nonlocal call_count
+        call_count += 1
+        status_code = 200
+        headers = {}
+        body = "[]"
+
+        if call_count >= calls_before_success:
+            return (status_code, headers, body)
+        raise requests.ConnectionError("Connection aborted.")
+
+    responses.add_callback(
+        method=responses.GET,
+        url=url,
+        callback=request_callback,
+        content_type="application/json",
+    )
+
+    with pytest.raises(requests.ConnectionError):
+        gl_retry.http_request("get", "/projects", retry_transient_errors=False)
+
+    assert len(responses.calls) == 1
+
+
 def create_redirect_response(
     *, response: requests.models.Response, http_method: str, api_path: str
 ) -> requests.models.Response:
