mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-15 04:32:22 +00:00
140 lines
4.8 KiB
Python
140 lines
4.8 KiB
Python
import asyncio
|
|
import os
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
|
from litellm.proxy.proxy_server import app, initialize
|
|
|
|
|
|
def _mock_user_api_key_auth():
|
|
"""Bypass auth for tests so /v1/audio/speech doesn't require a real key."""
|
|
return MagicMock()
|
|
|
|
|
|
def _make_mock_tts_response():
|
|
"""Mock response that simulates HttpxBinaryResponseContent with aiter_bytes."""
|
|
|
|
async def _chunks():
|
|
yield b"\xff\xfb"
|
|
|
|
def _aiter_bytes(chunk_size=8192):
|
|
async def _wrapper():
|
|
return _chunks()
|
|
|
|
return _wrapper()
|
|
|
|
inner = MagicMock()
|
|
inner.aiter_bytes = _aiter_bytes
|
|
inner._hidden_params = {}
|
|
|
|
async def _resolver():
|
|
return inner
|
|
|
|
return _resolver()
|
|
|
|
|
|
@pytest.fixture
|
|
def client_no_auth():
|
|
from litellm.proxy.proxy_server import cleanup_router_config_variables
|
|
|
|
cleanup_router_config_variables()
|
|
filepath = os.path.dirname(os.path.abspath(__file__))
|
|
config_fp = os.path.join(filepath, "test_configs", "test_config_no_auth.yaml")
|
|
asyncio.run(initialize(config=config_fp, debug=True))
|
|
return TestClient(app)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.retry(retries=0)
|
|
async def test_audio_speech_success_does_not_call_post_call_success_hook(
|
|
client_no_auth,
|
|
):
|
|
"""TTS success path must NOT call post_call_success_hook.
|
|
|
|
TTS returns a streaming binary response (HttpxBinaryResponseContent) which
|
|
is not in LLMResponseTypes. Prometheus metrics for successful requests are
|
|
tracked at the litellm level via async_log_success_event, not here.
|
|
"""
|
|
mock_success_hook = AsyncMock()
|
|
mock_failure_hook = AsyncMock()
|
|
mock_pre_call = AsyncMock(side_effect=lambda *, data, **kw: data)
|
|
mock_update_status = AsyncMock()
|
|
|
|
mock_logging = MagicMock()
|
|
mock_logging.post_call_success_hook = mock_success_hook
|
|
mock_logging.post_call_failure_hook = mock_failure_hook
|
|
mock_logging.pre_call_hook = mock_pre_call
|
|
mock_logging.update_request_status = mock_update_status
|
|
|
|
async def _mock_route_request(*, data, route_type, llm_router, user_model):
|
|
assert route_type == "aspeech"
|
|
return _make_mock_tts_response()
|
|
|
|
original_overrides = app.dependency_overrides.copy()
|
|
app.dependency_overrides[user_api_key_auth] = _mock_user_api_key_auth
|
|
try:
|
|
with (
|
|
patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_logging),
|
|
patch(
|
|
"litellm.proxy.proxy_server.route_request",
|
|
side_effect=_mock_route_request,
|
|
),
|
|
):
|
|
response = client_no_auth.post(
|
|
"/v1/audio/speech",
|
|
json={"model": "tts-1", "input": "hello"},
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
finally:
|
|
app.dependency_overrides = original_overrides
|
|
|
|
assert response.status_code == 200
|
|
mock_success_hook.assert_not_called()
|
|
mock_failure_hook.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.retry(retries=0)
|
|
async def test_audio_speech_failure_calls_post_call_failure_hook(client_no_auth):
|
|
"""TTS failure path must call proxy_logging_obj.post_call_failure_hook (Prometheus failed requests)."""
|
|
mock_success_hook = AsyncMock()
|
|
mock_failure_hook = AsyncMock()
|
|
mock_pre_call = AsyncMock(side_effect=lambda *, data, **kw: data)
|
|
|
|
mock_logging = MagicMock()
|
|
mock_logging.post_call_success_hook = mock_success_hook
|
|
mock_logging.post_call_failure_hook = mock_failure_hook
|
|
mock_logging.pre_call_hook = mock_pre_call
|
|
|
|
async def _mock_route_request_raise(*, data, route_type, llm_router, user_model):
|
|
raise ValueError("mock rate limit")
|
|
|
|
original_overrides = app.dependency_overrides.copy()
|
|
app.dependency_overrides[user_api_key_auth] = _mock_user_api_key_auth
|
|
# Don't re-raise server exceptions so we get the 500 response instead of ValueError
|
|
client = TestClient(app, raise_server_exceptions=False)
|
|
try:
|
|
with (
|
|
patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_logging),
|
|
patch(
|
|
"litellm.proxy.proxy_server.route_request",
|
|
side_effect=_mock_route_request_raise,
|
|
),
|
|
):
|
|
response = client.post(
|
|
"/v1/audio/speech",
|
|
json={"model": "tts-1", "input": "hello"},
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
finally:
|
|
app.dependency_overrides = original_overrides
|
|
|
|
assert response.status_code == 500
|
|
mock_failure_hook.assert_awaited_once()
|
|
mock_success_hook.assert_not_called()
|
|
call_kw = mock_failure_hook.call_args.kwargs
|
|
assert "user_api_key_dict" in call_kw and "original_exception" in call_kw
|