256 lines
7.2 KiB
Python
256 lines
7.2 KiB
Python
import unittest
|
|
from types import SimpleNamespace
|
|
|
|
from api_retry_utils import (
|
|
APIRequestError,
|
|
api_error_from_response,
|
|
call_with_retry,
|
|
is_retryable_api_error,
|
|
normalize_max_attempts,
|
|
)
|
|
|
|
|
|
class KeyErrorAttributeResponse:
|
|
"""Mimic DashScope responses that raise KeyError for missing fields."""
|
|
|
|
def __init__(self, **values):
|
|
self._values = values
|
|
|
|
def __getattr__(self, name):
|
|
try:
|
|
return self._values[name]
|
|
except KeyError:
|
|
raise KeyError(name) from None
|
|
|
|
|
|
class APIErrorTests(unittest.TestCase):
|
|
def test_response_error_is_sanitized(self):
|
|
response = SimpleNamespace(
|
|
status_code=429,
|
|
code="Throttled",
|
|
message="secret response body",
|
|
request_id="request-123",
|
|
headers={"Retry-After": "4"},
|
|
)
|
|
|
|
error = api_error_from_response(response, "test request")
|
|
|
|
self.assertEqual(error.status_code, 429)
|
|
self.assertEqual(error.retry_after, 4.0)
|
|
self.assertEqual(error.service_message, "secret response body")
|
|
self.assertIn("code=Throttled", str(error))
|
|
self.assertIn("request_id=request-123", str(error))
|
|
self.assertNotIn("secret response body", str(error))
|
|
|
|
def test_dashscope_style_missing_fields_are_safe(self):
|
|
response = KeyErrorAttributeResponse(
|
|
status_code=429,
|
|
code="Throttled",
|
|
message="secret response body",
|
|
request_id="request-123",
|
|
)
|
|
|
|
error = api_error_from_response(response, "test request")
|
|
|
|
self.assertEqual(error.status_code, 429)
|
|
self.assertIsNone(error.retry_after)
|
|
self.assertTrue(is_retryable_api_error(error))
|
|
self.assertNotIn("secret response body", str(error))
|
|
|
|
def test_none_retry_after_falls_back_to_header(self):
|
|
response = SimpleNamespace(
|
|
status_code=429,
|
|
retry_after=None,
|
|
headers={"Retry-After": "4"},
|
|
)
|
|
|
|
error = api_error_from_response(response, "test request")
|
|
|
|
self.assertEqual(error.retry_after, 4.0)
|
|
|
|
def test_invalid_retry_after_falls_back_to_case_insensitive_header(self):
|
|
response = SimpleNamespace(
|
|
status_code=503,
|
|
retry_after="not-a-delay",
|
|
headers={"retry-after": "5"},
|
|
)
|
|
|
|
error = api_error_from_response(response, "test request")
|
|
|
|
self.assertEqual(error.retry_after, 5.0)
|
|
|
|
def test_valid_retry_after_attribute_takes_precedence(self):
|
|
response = SimpleNamespace(
|
|
status_code=503,
|
|
retry_after="2",
|
|
headers={"Retry-After": "4"},
|
|
)
|
|
|
|
error = api_error_from_response(response, "test request")
|
|
|
|
self.assertEqual(error.retry_after, 2.0)
|
|
|
|
def test_mapping_response_is_supported(self):
|
|
response = {
|
|
"status_code": 500,
|
|
"code": "InternalError",
|
|
"headers": {"Retry-After": "3"},
|
|
}
|
|
|
|
error = api_error_from_response(response, "test request")
|
|
|
|
self.assertEqual(error.status_code, 500)
|
|
self.assertEqual(error.code, "InternalError")
|
|
self.assertEqual(error.retry_after, 3.0)
|
|
|
|
def test_attempt_count_is_capped(self):
|
|
messages = []
|
|
|
|
self.assertEqual(normalize_max_attempts(10000, messages.append), 10)
|
|
self.assertIn("capping at 10", messages[0])
|
|
|
|
|
|
class RetryTests(unittest.TestCase):
|
|
def _run(self, operation, max_attempts=3, **kwargs):
|
|
sleeps = []
|
|
messages = []
|
|
result = call_with_retry(
|
|
operation,
|
|
max_attempts=max_attempts,
|
|
operation_name="test request",
|
|
sleep_fn=sleeps.append,
|
|
random_fn=lambda: 0.0,
|
|
interrupt_check=kwargs.pop("interrupt_check", lambda: None),
|
|
interrupt_error_check=lambda error: False,
|
|
logger=messages.append,
|
|
**kwargs,
|
|
)
|
|
return result, sleeps, messages
|
|
|
|
def test_retries_transient_status_then_succeeds(self):
|
|
calls = 0
|
|
|
|
def operation():
|
|
nonlocal calls
|
|
calls += 1
|
|
if calls == 1:
|
|
raise APIRequestError("test", status_code=429)
|
|
return "ok"
|
|
|
|
result, sleeps, messages = self._run(operation)
|
|
|
|
self.assertEqual(result, "ok")
|
|
self.assertEqual(calls, 2)
|
|
self.assertAlmostEqual(sum(sleeps), 1.0)
|
|
self.assertEqual(len(messages), 1)
|
|
|
|
def test_retries_timeout_then_succeeds(self):
|
|
calls = 0
|
|
|
|
def operation():
|
|
nonlocal calls
|
|
calls += 1
|
|
if calls == 1:
|
|
raise TimeoutError("timed out")
|
|
return "ok"
|
|
|
|
result, _, _ = self._run(operation)
|
|
|
|
self.assertEqual(result, "ok")
|
|
self.assertEqual(calls, 2)
|
|
|
|
def test_honors_retry_after(self):
|
|
calls = 0
|
|
|
|
def operation():
|
|
nonlocal calls
|
|
calls += 1
|
|
if calls == 1:
|
|
raise APIRequestError(
|
|
"test",
|
|
status_code=503,
|
|
retry_after=4.0,
|
|
)
|
|
return "ok"
|
|
|
|
_, sleeps, _ = self._run(operation)
|
|
|
|
self.assertAlmostEqual(sum(sleeps), 4.0)
|
|
|
|
def test_does_not_retry_non_transient_status(self):
|
|
calls = 0
|
|
|
|
def operation():
|
|
nonlocal calls
|
|
calls += 1
|
|
raise APIRequestError("test", status_code=401)
|
|
|
|
with self.assertRaises(APIRequestError):
|
|
self._run(operation)
|
|
|
|
self.assertEqual(calls, 1)
|
|
|
|
def test_does_not_retry_response_parsing_error(self):
|
|
calls = 0
|
|
|
|
def operation():
|
|
nonlocal calls
|
|
calls += 1
|
|
raise ValueError("invalid response")
|
|
|
|
with self.assertRaises(ValueError):
|
|
self._run(operation)
|
|
|
|
self.assertEqual(calls, 1)
|
|
|
|
def test_stops_after_last_attempt_without_final_sleep(self):
|
|
calls = 0
|
|
|
|
def operation():
|
|
nonlocal calls
|
|
calls += 1
|
|
raise APIRequestError("test", status_code=500)
|
|
|
|
sleeps = []
|
|
with self.assertRaises(APIRequestError):
|
|
call_with_retry(
|
|
operation,
|
|
max_attempts=3,
|
|
operation_name="test request",
|
|
sleep_fn=sleeps.append,
|
|
random_fn=lambda: 0.0,
|
|
interrupt_check=lambda: None,
|
|
interrupt_error_check=lambda error: False,
|
|
logger=lambda message: None,
|
|
)
|
|
|
|
self.assertEqual(calls, 3)
|
|
self.assertAlmostEqual(sum(sleeps), 3.0)
|
|
|
|
def test_interrupt_during_backoff_stops_retrying(self):
|
|
calls = 0
|
|
checks = 0
|
|
|
|
class Interrupted(Exception):
|
|
pass
|
|
|
|
def operation():
|
|
nonlocal calls
|
|
calls += 1
|
|
raise APIRequestError("test", status_code=503)
|
|
|
|
def interrupt_check():
|
|
nonlocal checks
|
|
checks += 1
|
|
if checks == 2:
|
|
raise Interrupted()
|
|
|
|
with self.assertRaises(Interrupted):
|
|
self._run(operation, interrupt_check=interrupt_check)
|
|
|
|
self.assertEqual(calls, 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|