180 lines
5.4 KiB
Python
180 lines
5.4 KiB
Python
"""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"
|