tests

Author: dixitaniket

tests/llm_test.py

@@ -5,9 +5,10 @@
 """
 
 import json
+import ssl
 from contextlib import asynccontextmanager
 from typing import List
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
 import pytest
@@ -31,6 +32,8 @@ def __init__(self, *_args, **_kwargs):
         self._response_body: bytes = b"{}"
         self._post_calls: List[dict] = []
         self._stream_response = None
+        self._error_on_next: BaseException | None = None
+        self._stream_error_on_next: BaseException | None = None
 
     def set_response(self, status_code: int, body: dict) -> None:
         self._response_status = status_code
@@ -43,8 +46,19 @@ def set_stream_response(self, status_code: int, chunks: List[bytes]) -> None:
     def post_calls(self) -> List[dict]:
         return self._post_calls
 
+    def fail_next_post(self, exc: BaseException) -> None:
+        """Make the next post() call raise *exc*, then revert to normal."""
+        self._error_on_next = exc
+
+    def fail_next_stream(self, exc: BaseException) -> None:
+        """Make the next stream() call raise *exc*, then revert to normal."""
+        self._stream_error_on_next = exc
+
     async def post(self, url: str, *, json=None, headers=None, timeout=None) -> "_FakeResponse":
         self._post_calls.append({"url": url, "json": json, "headers": headers, "timeout": timeout})
+        if self._error_on_next is not None:
+            exc, self._error_on_next = self._error_on_next, None
+            raise exc
         resp = _FakeResponse(self._response_status, self._response_body)
         if self._response_status >= 400:
             resp.raise_for_status = MagicMock(side_effect=httpx.HTTPStatusError("error", request=MagicMock(), response=MagicMock()))
@@ -53,6 +67,9 @@ async def post(self, url: str, *, json=None, headers=None, timeout=None) -> "_Fa
     @asynccontextmanager
     async def stream(self, method: str, url: str, *, json=None, headers=None, timeout=None):
         self._post_calls.append({"method": method, "url": url, "json": json, "headers": headers, "timeout": timeout})
+        if self._stream_error_on_next is not None:
+            exc, self._stream_error_on_next = self._stream_error_on_next, None
+            raise exc
         yield self._stream_response
 
     async def aclose(self):
@@ -535,3 +552,248 @@ def test_registry_success(self):
             assert cert == b"cert-bytes"
             assert tee_id == "tee-42"
             assert pay_addr == "0xPay"
+
+
+# ── SSL cause detection tests ────────────────────────────────────────
+
+
+class TestHasSSLCause:
+    def test_direct_ssl_error(self):
+        assert LLM._has_ssl_cause(ssl.SSLError("cert expired")) is True
+
+    def test_wrapped_ssl_error_via_cause(self):
+        root = ssl.SSLError("handshake failed")
+        wrapper = RuntimeError("connection error")
+        wrapper.__cause__ = root
+        assert LLM._has_ssl_cause(wrapper) is True
+
+    def test_wrapped_ssl_error_via_context(self):
+        root = ssl.SSLError("cert verify failed")
+        wrapper = OSError("transport error")
+        wrapper.__context__ = root
+        assert LLM._has_ssl_cause(wrapper) is True
+
+    def test_deeply_nested_ssl(self):
+        root = ssl.SSLError("deep")
+        mid = OSError("mid")
+        mid.__cause__ = root
+        top = RuntimeError("top")
+        top.__cause__ = mid
+        assert LLM._has_ssl_cause(top) is True
+
+    def test_non_ssl_error(self):
+        assert LLM._has_ssl_cause(ValueError("not ssl")) is False
+
+    def test_non_ssl_chain(self):
+        root = TimeoutError("timed out")
+        wrapper = RuntimeError("oops")
+        wrapper.__cause__ = root
+        assert LLM._has_ssl_cause(wrapper) is False
+
+    def test_cycle_detection(self):
+        """Self-referencing chain should not loop forever."""
+        exc = RuntimeError("cycle")
+        exc.__cause__ = exc
+        assert LLM._has_ssl_cause(exc) is False
+
+
+# ── SSL retry tests (non-streaming) ──────────────────────────────────
+
+
+def _make_ssl_error(msg: str = "certificate verify failed") -> Exception:
+    """Create a RuntimeError wrapping an SSLError, mimicking httpx behaviour."""
+    root = ssl.SSLError(msg)
+    wrapper = RuntimeError(f"connection failed: {msg}")
+    wrapper.__cause__ = root
+    return wrapper
+
+
+@pytest.mark.asyncio
+class TestSSLRetryCompletion:
+    async def test_retries_on_ssl_and_succeeds(self, fake_http):
+        """First call hits SSL error → refresh → second call succeeds."""
+        fake_http.set_response(200, {"completion": "retried ok", "tee_signature": "s", "tee_timestamp": "t"})
+        fake_http.fail_next_post(_make_ssl_error())
+        llm = _make_llm()
+
+        result = await llm.completion(model=TEE_LLM.GPT_5, prompt="Hi")
+
+        assert result.completion_output == "retried ok"
+        # Two post calls: the failed one + the retry
+        assert len(fake_http.post_calls) == 2
+
+    async def test_non_ssl_error_not_retried(self, fake_http):
+        """A non-SSL error should propagate immediately, no retry."""
+        fake_http.fail_next_post(ConnectionError("DNS failed"))
+        llm = _make_llm()
+
+        with pytest.raises(RuntimeError, match="TEE LLM completion failed"):
+            await llm.completion(model=TEE_LLM.GPT_5, prompt="Hi")
+        assert len(fake_http.post_calls) == 1
+
+    async def test_second_ssl_failure_propagates(self, fake_http):
+        """If the retry also hits SSL, the error should propagate."""
+        fake_http.set_response(200, {"completion": "ok"})
+
+        call_count = 0
+
+        async def always_ssl(*args, **kwargs):
+            nonlocal call_count
+            call_count += 1
+            raise _make_ssl_error()
+
+        fake_http.post = always_ssl
+        llm = _make_llm()
+
+        # The second SSL error bubbles out of _call_with_ssl_retry as a
+        # RuntimeError (our wrapper), then caught by completion()'s outer
+        # handler which re-wraps it.
+        with pytest.raises(RuntimeError):
+            await llm.completion(model=TEE_LLM.GPT_5, prompt="Hi")
+        assert call_count == 2  # original + retry
+
+
+@pytest.mark.asyncio
+class TestSSLRetryChat:
+    async def test_retries_on_ssl_and_succeeds(self, fake_http):
+        fake_http.set_response(
+            200,
+            {"choices": [{"message": {"role": "assistant", "content": "retry ok"}, "finish_reason": "stop"}]},
+        )
+        fake_http.fail_next_post(_make_ssl_error())
+        llm = _make_llm()
+
+        result = await llm.chat(model=TEE_LLM.GPT_5, messages=[{"role": "user", "content": "Hi"}])
+
+        assert result.chat_output["content"] == "retry ok"
+        assert len(fake_http.post_calls) == 2
+
+    async def test_non_ssl_error_not_retried(self, fake_http):
+        fake_http.fail_next_post(TimeoutError("timed out"))
+        llm = _make_llm()
+
+        with pytest.raises(RuntimeError, match="TEE LLM chat failed"):
+            await llm.chat(model=TEE_LLM.GPT_5, messages=[{"role": "user", "content": "Hi"}])
+        assert len(fake_http.post_calls) == 1
+
+
+# ── SSL retry tests (streaming) ──────────────────────────────────────
+
+
+@pytest.mark.asyncio
+class TestSSLRetryStreaming:
+    async def test_retries_stream_on_ssl_before_chunks(self, fake_http):
+        """SSL failure during stream setup (no chunks yielded) → retry succeeds."""
+        fake_http.set_stream_response(
+            200,
+            [
+                b'data: {"model":"m","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":"stop"}]}\n\n',
+                b"data: [DONE]\n\n",
+            ],
+        )
+        fake_http.fail_next_stream(_make_ssl_error())
+        llm = _make_llm()
+
+        gen = await llm.chat(
+            model=TEE_LLM.GPT_5,
+            messages=[{"role": "user", "content": "Hi"}],
+            stream=True,
+        )
+        chunks = [c async for c in gen]
+
+        assert len(chunks) == 1
+        assert chunks[0].choices[0].delta.content == "ok"
+        # Two stream attempts: failed + retry
+        assert len(fake_http.post_calls) == 2
+
+    async def test_no_retry_after_chunks_yielded(self, fake_http):
+        """SSL failure AFTER chunks were yielded must raise, not retry."""
+
+        class _FailMidStream:
+            """Yields one chunk then raises SSL."""
+
+            def __init__(self):
+                self.status_code = 200
+
+            async def aiter_raw(self):
+                yield b'data: {"model":"m","choices":[{"index":0,"delta":{"content":"partial"},"finish_reason":null}]}\n\n'
+                raise _make_ssl_error()
+
+            async def aread(self) -> bytes:
+                return b""
+
+        fake_http._stream_response = _FailMidStream()
+        llm = _make_llm()
+
+        gen = await llm.chat(
+            model=TEE_LLM.GPT_5,
+            messages=[{"role": "user", "content": "Hi"}],
+            stream=True,
+        )
+
+        with pytest.raises(RuntimeError):
+            _ = [c async for c in gen]
+
+        # Only one stream call — no retry after partial output
+        assert len(fake_http.post_calls) == 1
+
+    async def test_non_ssl_stream_error_not_retried(self, fake_http):
+        fake_http.fail_next_stream(ConnectionError("reset"))
+        llm = _make_llm()
+
+        gen = await llm.chat(
+            model=TEE_LLM.GPT_5,
+            messages=[{"role": "user", "content": "Hi"}],
+            stream=True,
+        )
+
+        with pytest.raises(ConnectionError):
+            _ = [c async for c in gen]
+        assert len(fake_http.post_calls) == 1
+
+
+# ── _refresh_tee_and_reset tests ─────────────────────────────────────
+
+
+@pytest.mark.asyncio
+class TestRefreshTeeAndReset:
+    async def test_replaces_http_client(self):
+        """After refresh, the http client should be a new instance."""
+        clients_created = []
+
+        def make_client(*args, **kwargs):
+            c = FakeHTTPClient()
+            clients_created.append(c)
+            return c
+
+        with (
+            patch(_PATCHES["x402_httpx"], side_effect=make_client),
+            patch(_PATCHES["x402_client"]),
+            patch(_PATCHES["signer"]),
+            patch(_PATCHES["register_exact"]),
+            patch(_PATCHES["register_upto"]),
+        ):
+            llm = _make_llm()
+            old_client = llm._http_client
+
+            await llm._refresh_tee_and_reset()
+
+            assert llm._http_client is not old_client
+            assert len(clients_created) == 2  # init + refresh
+
+    async def test_closes_old_client(self, fake_http):
+        llm = _make_llm()
+        old_client = llm._http_client
+        old_client.aclose = AsyncMock()
+
+        await llm._refresh_tee_and_reset()
+
+        old_client.aclose.assert_awaited_once()
+
+    async def test_close_failure_is_swallowed(self, fake_http):
+        llm = _make_llm()
+        old_client = llm._http_client
+        old_client.aclose = AsyncMock(side_effect=OSError("already closed"))
+
+        # Should not raise
+        await llm._refresh_tee_and_reset()