From fe571264c3d4e5cf3dc9dd4665aa8b16c99f86ee Mon Sep 17 00:00:00 2001
From: quettabit <27509167+quettabit@users.noreply.github.com>
Date: Thu, 10 Sep 2026 12:55:04 -0700
Subject: [PATCH] initial commit

---
 src/s2_sdk/_client.py | 36 +++++++++++++++++++++---------------
 tests/test_client.py  |  2 +-
 2 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/src/s2_sdk/_client.py b/src/s2_sdk/_client.py
index c2b7255..7150bbc 100644
--- a/src/s2_sdk/_client.py
+++ b/src/s2_sdk/_client.py
@@ -339,9 +339,9 @@ async def checkout(self, base_url: str) -> _Checkout:
 
         self._ensure_reaper()
 
-        result = self._try_checkout(base_url)
-        if result is not None:
-            return result
+        checkout = self._try_checkout(base_url)
+        if checkout is not None:
+            return checkout
 
         lock = self._host_locks.get(base_url)
         if lock is None:
@@ -351,9 +351,9 @@ async def checkout(self, base_url: str) -> _Checkout:
         async with lock:
             # Re-check after acquiring lock — another caller may have
             # created a connection while we waited.
-            result = self._try_checkout(base_url)
-            if result is not None:
-                return result
+            checkout = self._try_checkout(base_url)
+            if checkout is not None:
+                return checkout
 
             scheme, host, port = _origin(base_url)
             use_ssl = self._ssl_context if scheme == "https" else None
@@ -375,26 +375,32 @@ async def checkout(self, base_url: str) -> _Checkout:
                 )
             except asyncio.TimeoutError:
                 pass  # Proceed with h2 defaults
+
             if conn._recv_dead:
                 await conn.close()
                 raise ConnectError(
                     f"Connection to {host}:{port} closed before HTTP/2 SETTINGS"
                 )
 
-            pc = _PooledConnection(conn)
-            conns = self._hosts.get(base_url)
-            if conns is None:
-                conns = [pc]
-                self._hosts[base_url] = conns
-            else:
-                conns.append(pc)
             if conn._settings_received.is_set() and conn.max_concurrent_streams <= 0:
-                await pc.close()
-                conns.remove(pc)
+                await conn.close()
                 raise ProtocolError("Connection has no available stream capacity")
+
+            pc = await self._add_connection(base_url, conn)
             state = pc._conn.reserve_stream()
             return _Checkout(pc, state)
 
+    async def _add_connection(
+        self, base_url: str, conn: Connection
+    ) -> _PooledConnection:
+        if self._closed:
+            await conn.close()
+            raise S2ClientError("Pool is closed")
+
+        pc = _PooledConnection(conn)
+        self._hosts.setdefault(base_url, []).append(pc)
+        return pc
+
     def _try_checkout(self, base_url: str) -> _Checkout | None:
         conns = self._hosts.get(base_url)
         if conns is not None:
diff --git a/tests/test_client.py b/tests/test_client.py
index ae42779..7286f21 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -180,7 +180,7 @@ async def test_new_connection_with_zero_stream_capacity_raises(pool: ConnectionP
             await pool.checkout("https://example.com")
 
         conn.close.assert_awaited_once()
-        assert pool._hosts["https://example.com"] == []
+        assert "https://example.com" not in pool._hosts
     await pool.close()