fix all tests

Author: balogh.adam@icloud.com

src/opengradient/client/tee_connection.py

@@ -142,7 +142,7 @@ def _connect(self) -> ActiveTEE:
         """Resolve TEE from registry and create a secure HTTP client."""
         tee = self._resolve_tee()
 
-        ssl_ctx = build_ssl_context_from_der(tee.tls_cert_der) 
+        ssl_ctx = build_ssl_context_from_der(tee.tls_cert_der)
         return ActiveTEE(
             endpoint=tee.endpoint,
             http_client=x402HttpxClient(self._x402_client, verify=ssl_ctx),

tests/client_test.py

@@ -22,9 +22,12 @@
 @pytest.fixture
 def mock_tee_registry():
     """Mock the TEE registry so LLM.__init__ doesn't need a live registry."""
-    with patch("src.opengradient.client.llm.TEERegistry") as mock_tee_registry, patch(
-        "src.opengradient.client.tee_connection.build_ssl_context_from_der",
-        return_value=MagicMock(),
+    with (
+        patch("src.opengradient.client.llm.TEERegistry") as mock_tee_registry,
+        patch(
+            "src.opengradient.client.tee_connection.build_ssl_context_from_der",
+            return_value=MagicMock(),
+        ),
     ):
         mock_tee = MagicMock()
         mock_tee.endpoint = "https://test.tee.server"

tests/tee_connection_test.py

@@ -42,12 +42,15 @@ def _mock_x402_client():
 def _make_registry_connection(*, registry=None, http_factory=None):
     """Build a RegistryTEEConnection with patched externals."""
     factory = http_factory or FakeHTTPClient
-    with patch(
-        "src.opengradient.client.tee_connection.x402HttpxClient",
-        side_effect=factory,
-    ), patch(
-        "src.opengradient.client.tee_connection.build_ssl_context_from_der",
-        return_value=MagicMock(spec=ssl.SSLContext),
+    with (
+        patch(
+            "src.opengradient.client.tee_connection.x402HttpxClient",
+            side_effect=factory,
+        ),
+        patch(
+            "src.opengradient.client.tee_connection.build_ssl_context_from_der",
+            return_value=MagicMock(spec=ssl.SSLContext),
+        ),
     ):
         return RegistryTEEConnection(
             x402_client=_mock_x402_client(),
@@ -182,12 +185,15 @@ def make_client(*args, **kwargs):
 
         mock_reg = _mock_registry_with_tee()
 
-        with patch(
-            "src.opengradient.client.tee_connection.x402HttpxClient",
-            side_effect=make_client,
-        ), patch(
-            "src.opengradient.client.tee_connection.build_ssl_context_from_der",
-            return_value=MagicMock(spec=ssl.SSLContext),
+        with (
+            patch(
+                "src.opengradient.client.tee_connection.x402HttpxClient",
+                side_effect=make_client,
+            ),
+            patch(
+                "src.opengradient.client.tee_connection.build_ssl_context_from_der",
+                return_value=MagicMock(spec=ssl.SSLContext),
+            ),
         ):
             conn = RegistryTEEConnection(
                 x402_client=_mock_x402_client(),
@@ -199,7 +205,6 @@ def make_client(*args, **kwargs):
         assert conn.get().http_client is not old_client
         assert len(clients_created) == 2
 
-
     async def test_reconnect_swallows_close_failure(self):
         mock_reg = _mock_registry_with_tee()
         conn = _make_registry_connection(registry=mock_reg)
@@ -224,7 +229,13 @@ def slow_connect(self):
         mock_reg = _mock_registry_with_tee()
         conn = _make_registry_connection(registry=mock_reg)
 
-        with patch.object(RegistryTEEConnection, "_connect", slow_connect):
+        with patch.object(RegistryTEEConnection, "_connect", slow_connect), patch(
+            "src.opengradient.client.tee_connection.build_ssl_context_from_der",
+            return_value=MagicMock(spec=ssl.SSLContext),
+        ), patch(
+            "src.opengradient.client.tee_connection.x402HttpxClient",
+            side_effect=FakeHTTPClient,
+        ):
             await asyncio.gather(conn.reconnect(), conn.reconnect())
 
         assert call_order == ["start", "end", "start", "end"]