mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
fix(test): replace flaky test_vertex_ai_gemini_audio_ogg with mocked version
Previously made a real Vertex AI call with a Wikimedia URL that intermittently failed with URL_REJECTED-REJECTED_FC_TIMEOUT. Now mocks HTTPHandler.post and VertexBase._ensure_access_token so the test verifies the translation (OGG -> file_data with audio/ogg mime_type) without any real network calls. Runs in ~0.36s instead of ~60s.
This commit is contained in:
parent
53f3dfa989
commit
7bf79f1e20
1 changed files with 69 additions and 22 deletions
|
|
@ -3891,30 +3891,77 @@ def test_vertex_ai_gemini_2_5_pro_streaming():
|
|||
|
||||
|
||||
def test_vertex_ai_gemini_audio_ogg():
|
||||
load_vertex_ai_credentials()
|
||||
litellm._turn_on_debug()
|
||||
response = completion(
|
||||
model="vertex_ai/gemini-2.0-flash",
|
||||
messages=[
|
||||
"""
|
||||
Test that OGG audio files are correctly formatted as file_data with audio/ogg mime type
|
||||
in the request sent to Vertex AI. Uses mocked HTTP and auth to avoid flaky external
|
||||
URL fetches and credential requirements.
|
||||
"""
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = {"Content-Type": "application/json"}
|
||||
mock_response.json.return_value = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": [
|
||||
{"text": "generate a transcript of the speech.", "type": "text"}
|
||||
],
|
||||
"role": "user",
|
||||
},
|
||||
{
|
||||
"content": [
|
||||
{
|
||||
"file": {
|
||||
"file_id": "https://upload.wikimedia.org/wikipedia/commons/5/5f/En-us-public.ogg"
|
||||
},
|
||||
"type": "file",
|
||||
}
|
||||
],
|
||||
"role": "user",
|
||||
},
|
||||
"content": {
|
||||
"parts": [{"text": "public domain audio file"}],
|
||||
"role": "model",
|
||||
},
|
||||
"finishReason": "STOP",
|
||||
}
|
||||
],
|
||||
)
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 10,
|
||||
"candidatesTokenCount": 5,
|
||||
"totalTokenCount": 15,
|
||||
},
|
||||
}
|
||||
|
||||
client = HTTPHandler()
|
||||
httpx_mock = MagicMock(return_value=mock_response)
|
||||
|
||||
with patch.object(client, "post", new=httpx_mock), patch.object(
|
||||
VertexBase, "_ensure_access_token", return_value=("fake-token", "fake-project")
|
||||
):
|
||||
response = completion(
|
||||
model="vertex_ai/gemini-2.0-flash",
|
||||
messages=[
|
||||
{
|
||||
"content": [
|
||||
{"text": "generate a transcript of the speech.", "type": "text"}
|
||||
],
|
||||
"role": "user",
|
||||
},
|
||||
{
|
||||
"content": [
|
||||
{
|
||||
"file": {
|
||||
"file_id": "https://upload.wikimedia.org/wikipedia/commons/5/5f/En-us-public.ogg"
|
||||
},
|
||||
"type": "file",
|
||||
}
|
||||
],
|
||||
"role": "user",
|
||||
},
|
||||
],
|
||||
client=client,
|
||||
)
|
||||
|
||||
httpx_mock.assert_called_once()
|
||||
request_body = httpx_mock.call_args.kwargs["json"]
|
||||
# Verify OGG file is sent as file_data with correct mime type
|
||||
file_data_parts = [
|
||||
part
|
||||
for content in request_body["contents"]
|
||||
for part in content["parts"]
|
||||
if "file_data" in part
|
||||
]
|
||||
assert len(file_data_parts) == 1, f"Expected 1 file_data part, got: {file_data_parts}"
|
||||
file_data = file_data_parts[0]["file_data"]
|
||||
assert file_data["mime_type"] == "audio/ogg", f"Expected audio/ogg, got: {file_data['mime_type']}"
|
||||
assert "En-us-public.ogg" in file_data["file_uri"], f"Unexpected file_uri: {file_data['file_uri']}"
|
||||
print(response)
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue