Skip to content

Commit ea74b63

Browse files
test(litellm): Replace mocks with httpx types in rate-limit test (#5975)
Replace mocks with `httpx` types to avoid test failures when library internals change.
1 parent 06ed1bc commit ea74b63

File tree

2 files changed

+40
-13
lines changed

2 files changed

+40
-13
lines changed

tests/conftest.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1127,6 +1127,28 @@ def inner(response_content, serialize_pydantic=False, request_headers=None):
11271127
return inner
11281128

11291129

1130+
@pytest.fixture
1131+
def get_rate_limit_model_response():
1132+
def inner(request_headers=None):
1133+
if request_headers is None:
1134+
request_headers = {}
1135+
1136+
model_request = HttpxRequest(
1137+
"POST",
1138+
"/responses",
1139+
headers=request_headers,
1140+
)
1141+
1142+
response = HttpxResponse(
1143+
429,
1144+
request=model_request,
1145+
)
1146+
1147+
return response
1148+
1149+
return inner
1150+
1151+
11301152
@pytest.fixture
11311153
def streaming_chat_completions_model_response():
11321154
return [

tests/integrations/litellm/test_litellm.py

Lines changed: 18 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -465,7 +465,9 @@ def test_embeddings_no_pii(
465465
assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span["data"]
466466

467467

468-
def test_exception_handling(sentry_init, capture_events):
468+
def test_exception_handling(
469+
reset_litellm_executor, sentry_init, capture_events, get_rate_limit_model_response
470+
):
469471
sentry_init(
470472
integrations=[LiteLLMIntegration()],
471473
traces_sample_rate=1.0,
@@ -474,19 +476,22 @@ def test_exception_handling(sentry_init, capture_events):
474476

475477
messages = [{"role": "user", "content": "Hello!"}]
476478

477-
with start_transaction(name="litellm test"):
478-
kwargs = {
479-
"model": "gpt-3.5-turbo",
480-
"messages": messages,
481-
}
479+
client = OpenAI(api_key="test-key")
482480

483-
_input_callback(kwargs)
484-
_failure_callback(
485-
kwargs,
486-
Exception("API rate limit reached"),
487-
datetime.now(),
488-
datetime.now(),
489-
)
481+
model_response = get_rate_limit_model_response()
482+
483+
with mock.patch.object(
484+
client.completions._client._client,
485+
"send",
486+
return_value=model_response,
487+
):
488+
with start_transaction(name="litellm test"):
489+
with pytest.raises(litellm.RateLimitError):
490+
litellm.completion(
491+
model="gpt-3.5-turbo",
492+
messages=messages,
493+
client=client,
494+
)
490495

491496
# Should have error event and transaction
492497
assert len(events) >= 1

0 commit comments

Comments
 (0)