diff --git a/src/opengradient/client/tee_connection.py b/src/opengradient/client/tee_connection.py index 2807b88..86c18f5 100644 --- a/src/opengradient/client/tee_connection.py +++ b/src/opengradient/client/tee_connection.py @@ -153,8 +153,10 @@ def _connect(self) -> ActiveTEE: async def reconnect(self) -> None: """Connect to a new TEE from the registry and rebuild the HTTP client.""" async with self._refresh_lock: + old_client = self._active.http_client + self._active = self._connect() try: - self._active = self._connect() + await old_client.aclose() except Exception: logger.debug("Failed to close previous HTTP client during TEE refresh.", exc_info=True) @@ -198,3 +200,4 @@ async def close(self) -> None: self._refresh_task.cancel() self._refresh_task = None await self._active.http_client.aclose() + diff --git a/tests/tee_connection_test.py b/tests/tee_connection_test.py index 20d02f7..d3172fc 100644 --- a/tests/tee_connection_test.py +++ b/tests/tee_connection_test.py @@ -210,11 +210,17 @@ async def test_reconnect_swallows_close_failure(self): conn = _make_registry_connection(registry=mock_reg) conn.get().http_client.aclose = AsyncMock(side_effect=OSError("already closed")) - with patch( - "opengradient.client.tee_connection.x402HttpxClient", - side_effect=FakeHTTPClient, + with ( + patch( + "opengradient.client.tee_connection.x402HttpxClient", + side_effect=FakeHTTPClient, + ), + patch( + "opengradient.client.tee_connection.build_ssl_context_from_der", + return_value=MagicMock(spec=ssl.SSLContext), + ), ): - await conn.reconnect() # should not raise + await conn.reconnect() # should not raise even when old client aclose errors async def test_reconnect_is_serialized(self): call_order = []