From 8cd19605bdea7fbc626e981f7dbbc89898160268 Mon Sep 17 00:00:00 2001
From: Haibo Lin <hlin@redhat.com>
Date: Oct 31 2022 04:29:44 +0000
Subject: Retry failed cts requests


JIRA: RHELCMP-10033
Signed-off-by: Haibo Lin <hlin@redhat.com>

---

diff --git a/pungi/compose.py b/pungi/compose.py
index 88ef7f6..ab62bbf 100644
--- a/pungi/compose.py
+++ b/pungi/compose.py
@@ -28,6 +28,8 @@ import socket
 
 import kobo.log
 import kobo.tback
+import requests
+from requests.exceptions import RequestException
 from productmd.composeinfo import ComposeInfo
 from productmd.images import Images
 from dogpile.cache import make_region
@@ -42,6 +44,7 @@ from pungi.util import (
     get_arch_variant_data,
     get_format_substs,
     get_variant_data,
+    retry,
     translate_path_raw,
 )
 from pungi.metadata import compose_to_composeinfo
@@ -54,6 +57,14 @@ except ImportError:
     SUPPORTED_MILESTONES = ["RC", "Update", "SecurityFix"]
 
 
+@retry(wait_on=RequestException)
+def retry_request(method, url, data=None, auth=None):
+    request_method = getattr(requests, method)
+    rv = request_method(url, json=data, auth=auth)
+    rv.raise_for_status()
+    return rv
+
+
 def get_compose_info(
     conf,
     compose_type="production",
@@ -86,10 +97,6 @@ def get_compose_info(
 
     cts_url = conf.get("cts_url", None)
     if cts_url:
-        # Import requests and requests-kerberos here so it is not needed
-        # if running without Compose Tracking Service.
-        import requests
-
         # Requests-kerberos cannot accept custom keytab, we need to use
         # environment variable for this. But we need to change environment
         # only temporarily just for this single requests.post.
@@ -113,8 +120,7 @@ def get_compose_info(
                 "parent_compose_ids": parent_compose_ids,
                 "respin_of": respin_of,
             }
-            rv = requests.post(url, json=data, auth=authentication)
-            rv.raise_for_status()
+            rv = retry_request("post", url, data=data, auth=authentication)
         finally:
             if cts_keytab:
                 shutil.rmtree(os.environ["KRB5CCNAME"].split(":", 1)[1])
@@ -156,8 +162,6 @@ def write_compose_info(compose_dir, ci):
 
 
 def update_compose_url(compose_id, compose_dir, conf):
-    import requests
-
     authentication = get_authentication(conf)
     cts_url = conf.get("cts_url", None)
     if cts_url:
@@ -168,7 +172,7 @@ def update_compose_url(compose_id, compose_dir, conf):
             "action": "set_url",
             "compose_url": compose_url,
         }
-        return requests.patch(url, json=data, auth=authentication)
+        return retry_request("patch", url, data=data, auth=authentication)
 
 
 def get_compose_dir(
diff --git a/tests/test_compose.py b/tests/test_compose.py
index 9494090..4474852 100644
--- a/tests/test_compose.py
+++ b/tests/test_compose.py
@@ -13,7 +13,9 @@ import tempfile
 import shutil
 import json
 
-from pungi.compose import Compose
+from requests.exceptions import HTTPError
+
+from pungi.compose import Compose, retry_request
 
 
 class ConfigWrapper(dict):
@@ -608,8 +610,9 @@ class ComposeTestCase(unittest.TestCase):
         ci_json = json.loads(ci.dumps())
         self.assertEqual(ci_json, self.ci_json)
 
+    @mock.patch("pungi.compose.requests")
     @mock.patch("time.strftime", new=lambda fmt, time: "20200526")
-    def test_get_compose_info_cts(self):
+    def test_get_compose_info_cts(self, mocked_requests):
         conf = ConfigWrapper(
             release_name="Test",
             release_version="1.0",
@@ -626,7 +629,6 @@ class ComposeTestCase(unittest.TestCase):
         ci_copy["header"]["version"] = "1.2"
         mocked_response = mock.MagicMock()
         mocked_response.text = json.dumps(self.ci_json)
-        mocked_requests = mock.MagicMock()
         mocked_requests.post.return_value = mocked_response
 
         mocked_requests_kerberos = mock.MagicMock()
@@ -637,7 +639,6 @@ class ComposeTestCase(unittest.TestCase):
         # `import`.
         with mock.patch.dict(
             "sys.modules",
-            requests=mocked_requests,
             requests_kerberos=mocked_requests_kerberos,
         ):
             ci = Compose.get_compose_info(conf, respin_of="Fedora-Rawhide-20200517.n.1")
@@ -807,3 +808,22 @@ class TracebackTest(unittest.TestCase):
     def test_with_detail(self):
         self.compose.traceback("extra-info")
         self.assertTraceback("traceback-extra-info")
+
+
+class RetryRequestTest(unittest.TestCase):
+    @mock.patch("pungi.compose.requests")
+    def test_retry_timeout(self, mocked_requests):
+        mocked_requests.post.side_effect = [
+            HTTPError("Gateway Timeout", response=mock.Mock(status_code=504)),
+            mock.Mock(status_code=200),
+        ]
+        url = "http://locahost/api/1/composes/"
+        rv = retry_request("post", url)
+        self.assertEqual(
+            mocked_requests.mock_calls,
+            [
+                mock.call.post(url, json=None, auth=None),
+                mock.call.post(url, json=None, auth=None),
+            ],
+        )
+        self.assertEqual(rv.status_code, 200)