fix: Add session retries to constructor to avoid starting new connect…

…ions
Allda · Jan 18, 2024 · 61901a7 · 61901a7
1 parent 071ba10
commit 61901a7
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 2 deletions.
diff --git a/coregio/registry_api.py b/coregio/registry_api.py
@@ -63,7 +63,12 @@ def __init__(
         self._original_url = url
         self.docker_cfg = docker_cfg
 
-        self.session = session or requests.Session()
+        if session:
+            self.session = session
+        else:
+            self.session = requests.Session()
+            utils.add_session_retries(self.session)
+
         self.proxy = proxy
 
         self.auth_header = None
@@ -196,7 +201,6 @@ def _get_session(self, auth_class: Callable[[Any], Any]) -> requests.Session:
                 auth_token, proxy=self.proxy
             )
         self.session.auth = self._auth_session_cache[auth_class]
-        utils.add_session_retries(self.session)
 
         return self.session
 

diff --git a/tests/test_registry_api.py b/tests/test_registry_api.py
@@ -7,6 +7,21 @@
 from coregio.registry_api import ContainerRegistry
 
 
+def test_init_new_session():
+    registry_api = ContainerRegistry(url="fake_url")
+    adapter = registry_api.session.adapters["https://"]
+
+    assert adapter.max_retries.total == 10
+
+
+def test_init_pass_session():
+    session = requests.Session()
+    registry_api = ContainerRegistry(url="fake_url", session=session)
+    adapter = registry_api.session.adapters["https://"]
+
+    assert adapter.max_retries.total == 0
+
+
 @pytest.mark.parametrize(
     ["url", "expected_url"],
     [