Livrare LOT 1 - Didi

ai_platform/modules/llm-inference/tests/test_retry.py

@@ -0,0 +1,180 @@
+"""Tests for retry functionality."""
+
+import asyncio
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from llm_inference.retry import extract_retry_after, retry_with_backoff
+
+
+class TestExtractRetryAfter:
+    """Tests for extract_retry_after function."""
+
+    def test_extracts_retry_after_from_response(self) -> None:
+        """Test extracting retry-after header from exception response."""
+        exc = MagicMock()
+        exc.response.headers = {"retry-after": "30"}
+
+        result = extract_retry_after(exc)
+
+        assert result == 30.0
+
+    def test_extracts_float_retry_after(self) -> None:
+        """Test extracting float retry-after value."""
+        exc = MagicMock()
+        exc.response.headers = {"retry-after": "45.5"}
+
+        result = extract_retry_after(exc)
+
+        assert result == 45.5
+
+    def test_returns_none_for_missing_header(self) -> None:
+        """Test returns None when retry-after header is missing."""
+        exc = MagicMock()
+        exc.response.headers = {}
+
+        result = extract_retry_after(exc)
+
+        assert result is None
+
+    def test_returns_none_for_no_response(self) -> None:
+        """Test returns None when exception has no response."""
+        exc = MagicMock()
+        exc.response = None
+
+        result = extract_retry_after(exc)
+
+        assert result is None
+
+    def test_returns_none_for_invalid_value(self) -> None:
+        """Test returns None when retry-after value is invalid."""
+        exc = MagicMock()
+        exc.response.headers = {"retry-after": "not-a-number"}
+
+        result = extract_retry_after(exc)
+
+        assert result is None
+
+
+class TestRetryWithBackoff:
+    """Tests for retry_with_backoff function."""
+
+    @pytest.mark.asyncio
+    async def test_successful_call_no_retry(self) -> None:
+        """Test that successful call returns immediately without retry."""
+        call_count = 0
+
+        async def success_func() -> str:
+            nonlocal call_count
+            call_count += 1
+            return "success"
+
+        result = await retry_with_backoff(
+            success_func,
+            max_retries=3,
+            min_wait=0.1,
+            max_wait=1.0,
+            backend="test",
+        )
+
+        assert result == "success"
+        assert call_count == 1
+
+    @pytest.mark.asyncio
+    async def test_retry_after_capped_at_max_wait(self) -> None:
+        """Test that retry-after header is capped at max_wait."""
+        import httpx
+
+        call_count = 0
+        sleep_times: list[float] = []
+
+        async def failing_func() -> str:
+            nonlocal call_count
+            call_count += 1
+            if call_count < 3:
+                # Create exception with very large retry-after
+                exc = httpx.TimeoutException("timeout")
+                exc.response = MagicMock()  # type: ignore
+                exc.response.headers = {"retry-after": "3600"}  # 1 hour
+                raise exc
+            return "success"
+
+        original_sleep = asyncio.sleep
+
+        async def mock_sleep(duration: float) -> None:
+            sleep_times.append(duration)
+            await original_sleep(0.001)  # Actually sleep very briefly
+
+        with patch("asyncio.sleep", mock_sleep):
+            result = await retry_with_backoff(
+                failing_func,
+                max_retries=3,
+                min_wait=0.1,
+                max_wait=5.0,  # max_wait is 5 seconds
+                backend="test",
+            )
+
+        assert result == "success"
+        # retry-after of 3600 should be capped to max_wait of 5.0
+        assert all(t <= 5.0 for t in sleep_times)
+
+    @pytest.mark.asyncio
+    async def test_exponential_backoff_with_jitter(self) -> None:
+        """Test that backoff uses exponential increase with jitter."""
+        import httpx
+
+        call_count = 0
+        sleep_times: list[float] = []
+
+        async def failing_func() -> str:
+            nonlocal call_count
+            call_count += 1
+            if call_count < 4:
+                raise httpx.TimeoutException("timeout")
+            return "success"
+
+        original_sleep = asyncio.sleep
+
+        async def mock_sleep(duration: float) -> None:
+            sleep_times.append(duration)
+            await original_sleep(0.001)
+
+        with patch("asyncio.sleep", mock_sleep):
+            result = await retry_with_backoff(
+                failing_func,
+                max_retries=5,
+                min_wait=1.0,
+                max_wait=60.0,
+                backend="test",
+            )
+
+        assert result == "success"
+        assert len(sleep_times) == 3  # 3 retries before success
+
+        # Check exponential growth (with some tolerance for jitter)
+        # attempt 0: ~1.0, attempt 1: ~2.0, attempt 2: ~4.0
+        assert 1.0 <= sleep_times[0] <= 1.25  # base + up to 25% jitter
+        assert 2.0 <= sleep_times[1] <= 2.5
+        assert 4.0 <= sleep_times[2] <= 5.0
+
+    @pytest.mark.asyncio
+    async def test_max_retries_exceeded_raises(self) -> None:
+        """Test that exceeding max retries raises translated exception."""
+        import httpx
+
+        from llm_inference.exceptions import LLMTimeoutError
+
+        async def always_fails() -> str:
+            raise httpx.TimeoutException("timeout")
+
+        with pytest.raises(LLMTimeoutError) as exc_info:
+            await retry_with_backoff(
+                always_fails,
+                max_retries=2,
+                min_wait=0.001,
+                max_wait=0.01,
+                backend="test",
+            )
+
+        assert exc_info.value.backend == "test"