litellm/tests/unit/test_router/test_router_endpoints.py
devin-ai-integration[bot] fa2c8984ba
test: move the unit half of 126 mixed legacy files into tests/unit (#45090)
* test: move the unit half of 126 mixed legacy files into tests/unit

* test: restore litellm globals that moved tests set

* test: finalize migration test cleanup

* test: restore original bodies of moved legacy tests

The move into tests/unit had rewritten 612 test bodies, and some of the rewrites dropped assertions. Each moved test now carries its original body from the legacy file, with only the imports, helpers, fake provider credentials and monkeypatched env it needs to run under tests/unit

test_timeout_streaming goes back to tests/local_testing because it needs the fake OpenAI endpoint server. The image payload fixture moves with its only user, and two tests that leaked global state (a registered model cost entry and queued logging tasks) are now isolated

* test: drop module imports shadowed by restored local imports

* test: assert on LiteLLM output in no-assertion moved tests and isolate leaks

Twenty no-assertion candidates get one assertion on the value LiteLLM returns, with the original lines unchanged. Four tests go back to their legacy files because they only check types or imports, write into the working directory, or cannot assert without a body change

Two moved tests leaked globals into later tests in the same worker, so monkeypatch fixtures now restore the retry-after header parser and the end user cost tracking flags

* test: drain queued logging tasks before the Phoenix span test

The moved Phoenix test counted spans from logging tasks that earlier tests had queued, so the drain fixture moves to tests/unit/conftest.py and both it and the Datadog batch test use it. test_factory_function goes back to its legacy file because its returned wrapper calls the real Assistants API and cannot be asserted on without a body change

---------

Co-authored-by: yuneng <yuneng@berri.ai>
2026-10-07 14:07:43 -07:00

1224 lines
37 KiB
Python

import asyncio
import os
from unittest.mock import AsyncMock, MagicMock, Mock, patch
import litellm
import pytest
from litellm import Router
@pytest.fixture
def model_list():
return [
{
"model_name": "gpt-5-mini",
"litellm_params": {
"model": "gpt-5-mini",
"api_key": os.getenv("OPENAI_API_KEY"),
},
},
{
"model_name": "gpt-5.5",
"litellm_params": {
"model": "gpt-5.5",
"api_key": os.getenv("OPENAI_API_KEY"),
},
},
{
"model_name": "gpt-image-1",
"litellm_params": {
"model": "gpt-image-1",
"api_key": os.getenv("OPENAI_API_KEY"),
},
},
{
"model_name": "cohere-rerank",
"litellm_params": {
"model": "cohere/rerank-english-v3.0",
"api_key": os.getenv("COHERE_API_KEY"),
},
},
{
"model_name": "claude-sonnet-4-5-20250929",
"litellm_params": {
"model": "gpt-5-mini",
"mock_response": "hi this is macintosh.",
},
},
]
@pytest.mark.asyncio
async def test_aspeech_fallbacks_on_deployment_failure():
router = Router(
model_list=[
{
"model_name": "tts-main",
"litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"},
},
{
"model_name": "tts-backup",
"litellm_params": {"model": "openai/tts-1-hd", "api_key": "fake-key"},
},
],
fallbacks=[{"tts-main": ["tts-backup"]}],
num_retries=0,
)
called_models = []
async def mock_aspeech(*args, **kwargs):
called_models.append(kwargs["model"])
if kwargs["model"] == "openai/tts-1":
raise litellm.InternalServerError(
message="deployment down",
llm_provider="openai",
model="tts-1",
)
return MagicMock()
with patch("litellm.aspeech", side_effect=mock_aspeech):
response = await router.aspeech(
model="tts-main",
input="the quick brown fox jumped over the lazy dogs",
voice="alloy",
)
assert response is not None
assert called_models == ["openai/tts-1", "openai/tts-1-hd"]
@pytest.mark.asyncio
async def test_aspeech_success_returns_response():
router = Router(
model_list=[
{
"model_name": "tts",
"litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"},
},
]
)
mock_response = MagicMock()
with patch("litellm.aspeech", return_value=mock_response) as mock_aspeech:
response = await router.aspeech(
model="tts",
input="the quick brown fox jumped over the lazy dogs",
voice="alloy",
)
assert response is mock_response
mock_aspeech.assert_called_once()
assert mock_aspeech.call_args.kwargs["model"] == "openai/tts-1"
@pytest.mark.asyncio
async def test_aspeech_sets_deployment_metadata():
router = Router(
model_list=[
{
"model_name": "tts",
"litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"},
},
]
)
mock_response = MagicMock()
with patch("litellm.aspeech", return_value=mock_response) as mock_aspeech:
response = await router._aspeech(
model="tts",
input="the quick brown fox jumped over the lazy dogs",
voice="alloy",
)
assert response is mock_response
metadata = mock_aspeech.call_args.kwargs["metadata"]
assert metadata["deployment"] == "openai/tts-1"
assert metadata["deployment_model_name"] == "tts"
assert metadata["model_info"]["id"] is not None
@pytest.mark.asyncio()
async def test_moderation_endpoint_with_api_base():
"""
Test that the moderation endpoint respects api_base configuration
"""
from unittest.mock import AsyncMock, MagicMock, patch
custom_api_base = "https://us.api.openai.com/v1"
router = Router(
model_list=[
{
"model_name": "openai/omni-moderation-latest",
"litellm_params": {
"model": "openai/omni-moderation-latest",
"api_base": custom_api_base,
"api_key": "test-key",
},
},
]
)
# Mock the OpenAI client to verify api_base is passed
with patch(
"litellm.main.openai_chat_completions.get_openai_client"
) as mock_get_client:
mock_client = AsyncMock()
mock_response = MagicMock()
mock_response.model_dump.return_value = {
"id": "modr-123",
"model": "omni-moderation-latest",
"results": [
{
"flagged": False,
"categories": {},
"category_scores": {},
"category_applied_input_types": {},
}
],
}
mock_client.moderations.create = AsyncMock(return_value=mock_response)
mock_get_client.return_value = mock_client
response = await router.amoderation(
model="openai/omni-moderation-latest", input="hello this is a test"
)
# Verify that get_openai_client was called with the custom api_base
mock_get_client.assert_called()
call_kwargs = mock_get_client.call_args.kwargs
assert (
call_kwargs.get("api_base") == custom_api_base
), f"Expected api_base to be {custom_api_base}, but got {call_kwargs.get('api_base')}"
print(f"✓ Moderation endpoint correctly uses api_base: {custom_api_base}")
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_aaaaatext_completion_endpoint(model_list, sync_mode):
router = Router(model_list=model_list)
if sync_mode:
response = router.text_completion(
model="gpt-5-mini",
prompt="Hello, how are you?",
mock_response="I'm fine, thank you!",
)
else:
response = await router.atext_completion(
model="gpt-5-mini",
prompt="Hello, how are you?",
mock_response="I'm fine, thank you!",
)
response_2 = await router._atext_completion(
model="gpt-5-mini",
prompt="Hello, how are you?",
mock_response="I'm fine, thank you!",
)
assert response_2.choices[0].text == "I'm fine, thank you!"
assert response.choices[0].text == "I'm fine, thank you!"
@pytest.mark.asyncio
async def test_router_with_empty_choices(model_list):
"""
https://github.com/BerriAI/litellm/issues/8306
"""
router = Router(model_list=model_list)
mock_response = litellm.ModelResponse(
choices=[],
usage=litellm.Usage(
prompt_tokens=10,
completion_tokens=10,
total_tokens=20,
),
model="gpt-5-mini",
object="chat.completion",
created=1723081200,
).model_dump()
response = await router.acompletion(
model="gpt-5-mini",
messages=[{"role": "user", "content": "Hello, how are you?"}],
mock_response=mock_response,
)
assert response is not None
@pytest.mark.parametrize("sync_mode", [True, False])
def test_generic_api_call_with_fallbacks_basic(sync_mode):
"""
Test both the sync and async versions of generic_api_call_with_fallbacks with a basic successful call
"""
# Create a mock function that will be passed to generic_api_call_with_fallbacks
if sync_mode:
from unittest.mock import Mock
mock_function = Mock()
mock_function.__name__ = "test_function"
else:
mock_function = AsyncMock()
mock_function.__name__ = "test_function"
# Create a mock response
mock_response = {
"id": "resp_123456",
"role": "assistant",
"content": "This is a test response",
"model": "test-model",
"usage": {"input_tokens": 10, "output_tokens": 20},
}
mock_function.return_value = mock_response
# Create a router with a test model
router = Router(
model_list=[
{
"model_name": "test-model-alias",
"litellm_params": {
"model": "anthropic/test-model",
"api_key": "fake-api-key",
},
}
]
)
# Call the appropriate generic_api_call_with_fallbacks method
if sync_mode:
response = router._generic_api_call_with_fallbacks(
model="test-model-alias",
original_function=mock_function,
messages=[{"role": "user", "content": "Hello"}],
max_tokens=100,
)
else:
response = asyncio.run(
router._ageneric_api_call_with_fallbacks(
model="test-model-alias",
original_function=mock_function,
messages=[{"role": "user", "content": "Hello"}],
max_tokens=100,
)
)
# Verify the mock function was called
mock_function.assert_called_once()
# Verify the response
assert response == mock_response
@pytest.mark.asyncio
async def test_aadapter_completion():
"""
Test the aadapter_completion method which uses async_function_with_fallbacks
"""
mock_response = {
"id": "adapter_resp_123",
"object": "adapter.completion",
"created": 1677858242,
"model": "test-model-with-adapter",
"choices": [
{
"text": "This is a test adapter response",
"index": 0,
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
}
with patch.object(Router, "_aadapter_completion", new_callable=AsyncMock) as mock_method:
mock_method.return_value = mock_response
router = Router(
model_list=[
{
"model_name": "test-adapter-model",
"litellm_params": {
"model": "anthropic/test-model",
"api_key": "fake-api-key",
},
}
]
)
router.async_function_with_fallbacks = AsyncMock(return_value=mock_response)
response = await router.aadapter_completion(
adapter_id="test-adapter-id",
model="test-adapter-model",
prompt="This is a test prompt",
max_tokens=100,
)
assert response == mock_response
router.async_function_with_fallbacks.assert_called_once()
call_kwargs = router.async_function_with_fallbacks.call_args.kwargs
assert call_kwargs["adapter_id"] == "test-adapter-id"
assert call_kwargs["model"] == "test-adapter-model"
assert call_kwargs["prompt"] == "This is a test prompt"
assert call_kwargs["max_tokens"] == 100
assert call_kwargs["original_function"] == router._aadapter_completion
assert "metadata" in call_kwargs
assert call_kwargs["metadata"]["model_group"] == "test-adapter-model"
@pytest.mark.asyncio
async def test__aadapter_completion():
"""
Test the _aadapter_completion method directly
"""
mock_response = {
"id": "adapter_resp_123",
"object": "adapter.completion",
"created": 1677858242,
"model": "test-model-with-adapter",
"choices": [
{
"text": "This is a test adapter response",
"index": 0,
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
}
with patch("litellm.aadapter_completion", new_callable=AsyncMock) as mock_adapter_completion:
mock_adapter_completion.return_value = mock_response
router = Router(
model_list=[
{
"model_name": "test-adapter-model",
"litellm_params": {
"model": "anthropic/test-model",
"api_key": "fake-api-key",
},
}
]
)
router.async_get_available_deployment = AsyncMock(
return_value={
"model_name": "test-adapter-model",
"litellm_params": {
"model": "test-model",
"api_key": "fake-api-key",
},
"model_info": {
"id": "test-unique-id",
},
}
)
router.async_routing_strategy_pre_call_checks = AsyncMock()
response = await router._aadapter_completion(
adapter_id="test-adapter-id",
model="test-adapter-model",
prompt="This is a test prompt",
max_tokens=100,
)
assert response == mock_response
mock_adapter_completion.assert_called_once()
call_kwargs = mock_adapter_completion.call_args.kwargs
assert call_kwargs["adapter_id"] == "test-adapter-id"
assert call_kwargs["model"] == "test-model"
assert call_kwargs["prompt"] == "This is a test prompt"
assert call_kwargs["max_tokens"] == 100
assert call_kwargs["api_key"] == "fake-api-key"
assert call_kwargs["caching"] == router.cache_responses
assert router.success_calls["test-model"] == 1
assert router.total_calls["test-model"] == 1
router.async_routing_strategy_pre_call_checks.assert_called_once()
def test_initialize_router_endpoints():
"""
Test that initialize_router_endpoints correctly sets up all router endpoints
"""
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "anthropic/test-model",
"api_key": "fake-api-key",
},
}
]
)
router.initialize_router_endpoints()
assert hasattr(router, "amoderation")
assert hasattr(router, "aanthropic_messages")
assert hasattr(router, "aresponses")
assert hasattr(router, "responses")
assert hasattr(router, "aget_responses")
assert hasattr(router, "adelete_responses")
assert callable(router.amoderation)
assert callable(router.aanthropic_messages)
assert callable(router.aresponses)
assert callable(router.responses)
assert callable(router.aget_responses)
assert callable(router.adelete_responses)
@pytest.mark.asyncio
async def test_init_responses_api_endpoints():
"""
A simpler test for _init_responses_api_endpoints that focuses on the basic functionality
"""
from litellm.responses.utils import ResponsesAPIRequestUtils
# Create a router with a basic model
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "openai/test-model",
"api_key": "fake-api-key",
},
}
]
)
# Just mock the _ageneric_api_call_with_fallbacks method
router._ageneric_api_call_with_fallbacks = AsyncMock()
# Add a mock implementation of _get_model_id_from_response_id to the Router instance
ResponsesAPIRequestUtils.get_model_id_from_response_id = MagicMock(
return_value=None
)
# Call without a response_id (no model extraction should happen)
await router._init_responses_api_endpoints(
original_function=AsyncMock(), thread_id="thread_xyz"
)
# Verify _ageneric_api_call_with_fallbacks was called but model wasn't changed
first_call_kwargs = router._ageneric_api_call_with_fallbacks.call_args.kwargs
assert "model" not in first_call_kwargs
assert first_call_kwargs["thread_id"] == "thread_xyz"
# Reset the mock
router._ageneric_api_call_with_fallbacks.reset_mock()
# Change the return value for the second call
ResponsesAPIRequestUtils.get_model_id_from_response_id.return_value = (
"claude-3-sonnet"
)
# Call with a response_id
await router._init_responses_api_endpoints(
original_function=AsyncMock(), response_id="resp_claude_123"
)
# Verify model was updated in the kwargs
second_call_kwargs = router._ageneric_api_call_with_fallbacks.call_args.kwargs
assert second_call_kwargs["model"] == "claude-3-sonnet"
assert second_call_kwargs["response_id"] == "resp_claude_123"
@pytest.mark.asyncio
async def test_init_vector_store_api_endpoints():
"""
Test that _init_vector_store_api_endpoints correctly passes custom_llm_provider to kwargs
"""
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "openai/test-model",
"api_key": "fake-api-key",
},
}
]
)
mock_original_function = AsyncMock(return_value={"status": "success"})
result = await router._init_vector_store_api_endpoints(
original_function=mock_original_function, vector_store_id="test-store"
)
mock_original_function.assert_called_once_with(vector_store_id="test-store")
assert result == {"status": "success"}
mock_original_function.reset_mock()
await router._init_vector_store_api_endpoints(
original_function=mock_original_function,
custom_llm_provider="openai",
vector_store_id="test-store",
)
mock_original_function.assert_called_once_with(vector_store_id="test-store", custom_llm_provider="openai")
def test_apply_default_settings():
"""
Test the apply_default_settings method.
This test verifies that apply_default_settings correctly initializes
default pre-call checks and doesn't modify existing router state.
"""
router = Router()
initial_optional_callbacks = router.optional_callbacks
result = router.apply_default_settings()
assert result is None
assert router.optional_callbacks == initial_optional_callbacks
router_with_callbacks = Router()
mock_callback = MagicMock()
router_with_callbacks.optional_callbacks = [mock_callback]
result = router_with_callbacks.apply_default_settings()
assert result is None
assert mock_callback in router_with_callbacks.optional_callbacks
with patch.object(Router, "apply_default_settings") as mock_apply:
Router()
mock_apply.assert_called_once()
router_test = Router()
with patch.object(router_test, "add_optional_pre_call_checks") as mock_add_checks:
router_test.apply_default_settings()
mock_add_checks.assert_called_once_with([])
def test_initialize_core_endpoints():
"""
Test that _initialize_core_endpoints correctly sets up all core router endpoints.
"""
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "anthropic/test-model",
"api_key": "fake-api-key",
},
}
]
)
router._initialize_core_endpoints()
core_endpoints = [
"amoderation",
"aanthropic_messages",
"agenerate_content",
"aadapter_generate_content",
"aresponses",
"afile_delete",
"afile_content",
"responses",
"aget_responses",
"acancel_responses",
"adelete_responses",
"alist_input_items",
"_arealtime",
"acreate_fine_tuning_job",
"acancel_fine_tuning_job",
"alist_fine_tuning_jobs",
"aretrieve_fine_tuning_job",
"afile_list",
"aimage_edit",
"allm_passthrough_route",
]
for endpoint in core_endpoints:
assert hasattr(router, endpoint)
assert callable(getattr(router, endpoint))
def test_initialize_specialized_endpoints():
"""
Test that _initialize_specialized_endpoints correctly sets up specialized endpoints.
"""
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "openai/test-model",
"api_key": "fake-api-key",
},
}
]
)
router._initialize_specialized_endpoints()
specialized_endpoints = [
"avector_store_search",
"avector_store_create",
"vector_store_search",
"vector_store_create",
"agenerate_content",
"generate_content",
"agenerate_content_stream",
"generate_content_stream",
"aocr",
"ocr",
"asearch",
"search",
"avideo_generation",
"video_generation",
"avideo_list",
"video_list",
"avideo_status",
"video_status",
"avideo_content",
"video_content",
"avideo_remix",
"video_remix",
"acreate_container",
"create_container",
"alist_containers",
"list_containers",
"aretrieve_container",
"retrieve_container",
"adelete_container",
"delete_container",
"acreate_skill",
"alist_skills",
"aget_skill",
"adelete_skill",
]
for endpoint in specialized_endpoints:
assert hasattr(router, endpoint)
assert callable(getattr(router, endpoint))
def test_initialize_vector_store_endpoints():
"""
Test that _initialize_vector_store_endpoints correctly sets up vector store endpoints.
"""
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "openai/test-model",
"api_key": "fake-api-key",
},
}
]
)
router._initialize_vector_store_endpoints()
vector_store_endpoints = [
"avector_store_search",
"avector_store_create",
"vector_store_search",
"vector_store_create",
]
for endpoint in vector_store_endpoints:
assert hasattr(router, endpoint)
assert callable(getattr(router, endpoint))
def test_initialize_vector_store_file_endpoints():
"""
Test that _initialize_vector_store_file_endpoints correctly sets up vector store file endpoints.
"""
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "openai/test-model",
"api_key": "fake-api-key",
},
}
]
)
router._initialize_vector_store_file_endpoints()
vector_store_file_endpoints = [
"avector_store_file_create",
"vector_store_file_create",
"avector_store_file_list",
"vector_store_file_list",
"avector_store_file_retrieve",
"vector_store_file_retrieve",
"avector_store_file_content",
"vector_store_file_content",
"avector_store_file_update",
"vector_store_file_update",
"avector_store_file_delete",
"vector_store_file_delete",
]
for endpoint in vector_store_file_endpoints:
assert hasattr(router, endpoint)
assert callable(getattr(router, endpoint))
def test_initialize_google_genai_endpoints():
"""
Test that _initialize_google_genai_endpoints correctly sets up Google GenAI endpoints.
"""
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "openai/test-model",
"api_key": "fake-api-key",
},
}
]
)
router._initialize_google_genai_endpoints()
google_genai_endpoints = [
"agenerate_content",
"generate_content",
"agenerate_content_stream",
"generate_content_stream",
]
for endpoint in google_genai_endpoints:
assert hasattr(router, endpoint)
assert callable(getattr(router, endpoint))
def test_initialize_ocr_search_endpoints():
"""
Test that _initialize_ocr_search_endpoints correctly sets up OCR and search endpoints.
"""
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "openai/test-model",
"api_key": "fake-api-key",
},
}
]
)
router._initialize_ocr_search_endpoints()
ocr_search_endpoints = [
"aocr",
"ocr",
"asearch",
"search",
]
for endpoint in ocr_search_endpoints:
assert hasattr(router, endpoint)
assert callable(getattr(router, endpoint))
def test_initialize_video_endpoints():
"""
Test that _initialize_video_endpoints correctly sets up video endpoints.
"""
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "openai/test-model",
"api_key": "fake-api-key",
},
}
]
)
router._initialize_video_endpoints()
video_endpoints = [
"avideo_generation",
"video_generation",
"avideo_list",
"video_list",
"avideo_status",
"video_status",
"avideo_content",
"video_content",
"avideo_remix",
"video_remix",
]
for endpoint in video_endpoints:
assert hasattr(router, endpoint)
assert callable(getattr(router, endpoint))
def test_initialize_container_endpoints():
"""
Test that _initialize_container_endpoints correctly sets up container endpoints.
"""
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "openai/test-model",
"api_key": "fake-api-key",
},
}
]
)
router._initialize_container_endpoints()
container_endpoints = [
"acreate_container",
"create_container",
"alist_containers",
"list_containers",
"aretrieve_container",
"retrieve_container",
"adelete_container",
"delete_container",
]
for endpoint in container_endpoints:
assert hasattr(router, endpoint)
assert callable(getattr(router, endpoint))
def test_initialize_skills_endpoints():
"""
Test that _initialize_skills_endpoints correctly sets up skills endpoints.
"""
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "anthropic/test-model",
"api_key": "fake-api-key",
},
}
]
)
router._initialize_skills_endpoints()
skills_endpoints = [
"acreate_skill",
"alist_skills",
"aget_skill",
"adelete_skill",
]
for endpoint in skills_endpoints:
assert hasattr(router, endpoint)
assert callable(getattr(router, endpoint))
@pytest.mark.asyncio
async def test_init_containers_api_endpoints():
"""
Test that _init_containers_api_endpoints calls the original function
directly when there is no managed container ID (no embedded model_id).
"""
router = Router(model_list=[])
mock_response = {"id": "cntr_test", "name": "Test Container"}
mock_original_function = AsyncMock(return_value=mock_response)
result = await router._init_containers_api_endpoints(
original_function=mock_original_function,
custom_llm_provider="openai",
name="Test Container",
)
mock_original_function.assert_called_once_with(custom_llm_provider="openai", name="Test Container")
assert result == mock_response
@pytest.mark.asyncio
async def test_init_containers_api_endpoints_managed_id_routes_via_generic_fallbacks():
"""
Managed ``cntr_`` IDs embed ``model_id``; router should decode and use
``_ageneric_api_call_with_fallbacks`` so deployment credentials apply.
"""
from litellm.responses.utils import ResponsesAPIRequestUtils
router = Router(
model_list=[
{
"model_name": "azure-router-model",
"litellm_params": {
"model": "azure/gpt-5.5",
"api_key": "fake-key",
"api_base": "https://westus.api.cognitive.microsoft.com",
},
}
]
)
router._ageneric_api_call_with_fallbacks = AsyncMock()
managed_id = ResponsesAPIRequestUtils.build_container_id(
custom_llm_provider="azure",
model_id="azure-router-model",
container_id="cfile_upstream_abc",
)
await router._init_containers_api_endpoints(
original_function=AsyncMock(),
custom_llm_provider="openai",
container_id=managed_id,
file_id="cfile_xyz",
)
router._ageneric_api_call_with_fallbacks.assert_called_once()
call_kw = router._ageneric_api_call_with_fallbacks.call_args.kwargs
assert call_kw["model"] == "azure-router-model"
assert call_kw["container_id"] == "cfile_upstream_abc"
assert call_kw["file_id"] == "cfile_xyz"
assert call_kw["custom_llm_provider"] == "azure"
@pytest.mark.asyncio
async def test_init_containers_api_endpoints_managed_id_without_model_id_unwraps():
"""
Managed ``cntr_`` IDs may be encoded with an empty ``model_id`` (e.g. when a
streaming response had no router metadata). The router must still unwrap the
managed ID before calling the upstream provider — otherwise the raw
``cntr_...`` token leaks downstream and the provider rejects it.
"""
from litellm.responses.utils import ResponsesAPIRequestUtils
router = Router(model_list=[])
mock_original_function = AsyncMock(return_value={"ok": True})
managed_id = ResponsesAPIRequestUtils.build_container_id(
custom_llm_provider="openai",
model_id=None,
container_id="cfile_upstream_abc",
)
await router._init_containers_api_endpoints(
original_function=mock_original_function,
custom_llm_provider="openai",
container_id=managed_id,
file_id="cfile_xyz",
)
mock_original_function.assert_called_once()
call_kw = mock_original_function.call_args.kwargs
assert call_kw["container_id"] == "cfile_upstream_abc"
assert call_kw["file_id"] == "cfile_xyz"
assert call_kw["custom_llm_provider"] == "openai"
@pytest.mark.asyncio
async def test_init_containers_api_endpoints_managed_id_without_model_id_applies_decoded_provider():
"""
A managed ``cntr_`` ID can encode a non-OpenAI provider (e.g. ``azure``) with
an empty ``model_id`` (streaming events without router ``model_info.id``).
The router must still apply the decoded provider so the request routes to
the correct upstream — not stay on the default ``openai``.
"""
from litellm.responses.utils import ResponsesAPIRequestUtils
router = Router(model_list=[])
mock_original_function = AsyncMock(return_value={"ok": True})
managed_id = ResponsesAPIRequestUtils.build_container_id(
custom_llm_provider="azure",
model_id=None,
container_id="cfile_upstream_abc",
)
await router._init_containers_api_endpoints(
original_function=mock_original_function,
custom_llm_provider="openai",
container_id=managed_id,
file_id="cfile_xyz",
)
mock_original_function.assert_called_once()
call_kw = mock_original_function.call_args.kwargs
assert call_kw["container_id"] == "cfile_upstream_abc"
assert call_kw["file_id"] == "cfile_xyz"
assert call_kw["custom_llm_provider"] == "azure"
@pytest.mark.asyncio
async def test_init_containers_api_endpoints_create_with_model_uses_deployment_credentials(monkeypatch):
"""
``POST /v1/containers`` carries no container ID, so a ``model`` in the request
body is the only way to pick a deployment. The upstream call must receive that
deployment's ``api_key``/``api_base`` instead of falling back to the global
``OPENAI_API_KEY`` (which may be unset on the proxy).
"""
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
router = Router(
model_list=[
{
"model_name": "gpt-5.4",
"litellm_params": {
"model": "openai/gpt-5.4",
"api_key": "sk-model-list-key",
"api_base": "https://custom.openai.example/v1",
},
}
]
)
mock_original_function = AsyncMock(return_value={"id": "cntr_test", "name": "Test Container"})
await router._init_containers_api_endpoints(
original_function=mock_original_function,
custom_llm_provider="openai",
name="Test Container",
model="gpt-5.4",
)
mock_original_function.assert_called_once()
call_kw = mock_original_function.call_args.kwargs
assert call_kw["api_key"] == "sk-model-list-key"
assert call_kw["api_base"] == "https://custom.openai.example/v1"
assert call_kw["model"] == "openai/gpt-5.4"
assert call_kw["name"] == "Test Container"
@pytest.mark.asyncio
async def test_init_containers_api_endpoints_create_without_model_calls_directly():
"""
Without ``model`` (or with ``model=None`` as the proxy forwards it), create/list
must keep calling the handler directly with global provider credentials.
"""
router = Router(model_list=[])
router._ageneric_api_call_with_fallbacks = AsyncMock()
mock_original_function = AsyncMock(return_value={"id": "cntr_test"})
await router._init_containers_api_endpoints(
original_function=mock_original_function,
custom_llm_provider="openai",
name="Test Container",
model=None,
)
router._ageneric_api_call_with_fallbacks.assert_not_called()
mock_original_function.assert_called_once_with(custom_llm_provider="openai", name="Test Container", model=None)
@pytest.mark.asyncio
async def test_init_containers_api_endpoints_create_with_unknown_model_passes_through(monkeypatch):
"""
A ``model`` that names no configured deployment must not turn into a 400. The call
falls through to the handler with the caller's model and no injected deployment
credentials, matching the behaviour before model-based routing existed.
"""
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
router = Router(
model_list=[
{
"model_name": "gpt-5.4",
"litellm_params": {"model": "openai/gpt-5.4", "api_key": "sk-model-list-key"},
}
]
)
mock_original_function = AsyncMock(return_value={"id": "cntr_test"})
await router._init_containers_api_endpoints(
original_function=mock_original_function,
custom_llm_provider="openai",
name="Test Container",
model="does-not-exist",
)
mock_original_function.assert_called_once()
call_kw = mock_original_function.call_args.kwargs
assert call_kw["model"] == "does-not-exist"
assert call_kw["name"] == "Test Container"
assert "api_key" not in call_kw
assert "api_base" not in call_kw
def test_router_model_group_encrypted_content_affinity_callback_registration():
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
)
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
EncryptedContentAffinityCheck,
)
model_group = "openai.gpt-5.1-codex"
model_group_affinity_config = {
model_group: ["encrypted_content_affinity"],
}
router = Router(
model_list=[
{
"model_name": model_group,
"litellm_params": {
"model": "openai/gpt-5.1-codex",
"api_key": "mock-api-key",
},
}
],
model_group_affinity_config=model_group_affinity_config,
num_retries=0,
)
try:
callbacks = router.optional_callbacks or []
encrypted_content_callbacks = [
cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)
]
deployment_callback = next(
cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)
)
assert len(encrypted_content_callbacks) == 1
assert encrypted_content_callbacks[0].enable_global_affinity is False
assert (
encrypted_content_callbacks[0].model_group_affinity_config
== model_group_affinity_config
)
assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index(
deployment_callback
)
router._add_encrypted_content_affinity_check(enable_global_affinity=True)
callbacks = router.optional_callbacks or []
encrypted_content_callbacks = [
cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)
]
assert len(encrypted_content_callbacks) == 1
assert encrypted_content_callbacks[0].enable_global_affinity is True
assert encrypted_content_callbacks[0].router is router
finally:
router.discard()