Add retries to Run service requests (GoogleCloudPlatform#6087)

averikitsch · web-flow · commit 0a88585dea4e · 2021-06-09T21:22:06.000Z
## Description Fixes GoogleCloudPlatform#6086 GoogleCloudPlatform#6085 Note: It's a good idea to open an issue first for discussion. ## Checklist - [ ] I have followed [Sample Guidelines from AUTHORING_GUIDE.MD](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/AUTHORING_GUIDE.md) - [ ] README is updated to include [all relevant information](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/AUTHORING_GUIDE.md#readme-file) - [ ] **Tests** pass: `nox -s py-3.6` (see [Test Environment Setup](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/AUTHORING_GUIDE.md#test-environment-setup)) - [ ] **Lint** pass: `nox -s lint` (see [Test Environment Setup](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/AUTHORING_GUIDE.md#test-environment-setup)) - [ ] These samples need a new **API enabled** in testing projects to pass (let us know which ones) - [ ] These samples need a new/updated **env vars** in testing projects set to pass (let us know which ones) - [ ] Please **merge** this PR for me once it is approved. - [ ] This sample adds a new sample directory, and I updated the [CODEOWNERS file](https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/.github/CODEOWNERS) with the codeowners for this sample
diff --git a/run/idp-sql/test/e2e_test.py b/run/idp-sql/test/e2e_test.py
@@ -25,6 +25,8 @@
 from firebase_admin import auth  # noqa: F401
 import pytest
 import requests
+from requests.adapters import HTTPAdapter
+from requests.packages.urllib3.util.retry import Retry
 
 default_app = firebase_admin.initialize_app()
 
@@ -179,7 +181,24 @@ def test_end_to_end(jwt_token: str, deployed_service: str) -> None:
     token = jwt_token
     service_url = deployed_service
 
+    retry_strategy = Retry(
+        total=3,
+        status_forcelist=[400, 401, 403, 500, 502, 503, 504],
+        allowed_methods=["GET", "POST"],
+        backoff_factor=3
+    )
+
+    retry_strategy_500 = Retry(
+        total=3,
+        status_forcelist=[500, 502, 503, 504],
+        allowed_methods=["GET", "POST"],
+        backoff_factor=3
+    )
+
+    adapter = HTTPAdapter(max_retries=retry_strategy)
+
     client = requests.session()
+    client.mount("https://", adapter)
 
     # Can successfully make a request
     response = client.get(service_url)
@@ -197,6 +216,8 @@ def test_end_to_end(jwt_token: str, deployed_service: str) -> None:
     assert response.status_code == 200
     assert "🐶" in response.content.decode("UTF-8")
 
+    adapter = HTTPAdapter(max_retries=retry_strategy_500)
+    client.mount("https://", adapter)
     # Cannot make post with bad token
     response = client.post(
         service_url,
diff --git a/run/logging-manual/e2e_test.py b/run/logging-manual/e2e_test.py
@@ -20,13 +20,16 @@
 import os
 import subprocess
 import time
-from urllib import request
 import uuid
 
 from google.cloud import logging_v2
 
 import pytest
 
+import requests
+from requests.adapters import HTTPAdapter
+from requests.packages.urllib3.util.retry import Retry
+
 # Unique suffix to create distinct service names
 SUFFIX = uuid.uuid4().hex[:10]
 PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"]
@@ -148,18 +151,26 @@ def test_end_to_end(service_url_auth_token, deployed_service):
     service_url, auth_token = service_url_auth_token
 
     # Test that the service is responding
-    req = request.Request(
+    retry_strategy = Retry(
+        total=3,
+        status_forcelist=[400, 401, 403, 500, 502, 503, 504],
+        allowed_methods=["GET", "POST"],
+        backoff_factor=3
+    )
+    adapter = HTTPAdapter(max_retries=retry_strategy)
+
+    client = requests.session()
+    client.mount("https://", adapter)
+
+    response = client.get(
         service_url,
         headers={
             "Authorization": f"Bearer {auth_token}",
             "X-Cloud-Trace-Context": "foo/bar",
         },
     )
-    response = request.urlopen(req)
-    assert response.status == 200
-
-    body = response.read()
-    assert body.decode() == "Hello Logger!"
+    assert response.status_code == 200
+    assert "Hello Logger!" in response.content.decode("UTF-8")
 
     # Test that the logs are writing properly to stackdriver
     time.sleep(10)  # Slight delay writing to stackdriver
