fix black formatting; add router-resolved path test

Black reformats the long ternary in _update_litellm_params_for_health_check
call to a multi-line form. Also add a second regression test covering the
deployment-resolved-by-id path where server-side model_info carries
health_check_supports_max_tokens: False.
This commit is contained in:
Sebin P Johnson 2026-05-30 09:46:23 +05:30
parent 4e286a4b8b
commit 2e5e78c35d
2 changed files with 77 additions and 1 deletions

View file

@ -1849,7 +1849,11 @@ async def test_model_connection(
)
# Include health_check_params if provided
litellm_params = _update_litellm_params_for_health_check(
model_info=resolved_model_info if resolved_model_info is not None else (model_info or {}),
model_info=(
resolved_model_info
if resolved_model_info is not None
else (model_info or {})
),
litellm_params=litellm_params,
)
mode = mode or litellm_params.pop("mode", None)

View file

@ -761,6 +761,78 @@ async def test_test_model_connection_respects_health_check_supports_max_tokens()
)
@pytest.mark.asyncio
async def test_test_model_connection_uses_server_model_info_when_deployment_resolved():
"""
When the router resolves a deployment by id, the server-side model_info
must be forwarded to the health-check helper, not the caller-supplied one.
This covers the router-resolved path for the same bug as
test_test_model_connection_respects_health_check_supports_max_tokens.
"""
from litellm.types.router import Deployment, LiteLLM_Params
mock_request = MagicMock()
mock_user_api_key_dict = MagicMock()
mock_user_api_key_dict.user_id = "test-user"
mock_user_api_key_dict.token = "test-token"
mock_prisma_client = MagicMock()
mock_can_user_make_model_call = AsyncMock()
mock_ahealth_check = AsyncMock(return_value={"status": "healthy"})
mock_run_with_timeout = AsyncMock(return_value={"status": "healthy"})
server_deployment = Deployment(
model_name="my-model",
litellm_params=LiteLLM_Params(
model="azure/gpt-4o",
api_key="server-key",
api_base="https://server.invalid/v1",
),
model_info={"id": "dep-1", "health_check_supports_max_tokens": False},
)
mock_router = MagicMock()
mock_router.get_deployment.return_value = server_deployment
def mock_reject_os_environ(params):
return None
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
patch("litellm.proxy.proxy_server.llm_router", mock_router),
patch("litellm.proxy.proxy_server.premium_user", False),
patch(
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
mock_can_user_make_model_call,
),
patch(
"litellm.proxy.health_endpoints._health_endpoints.litellm.ahealth_check",
mock_ahealth_check,
),
patch(
"litellm.proxy.health_endpoints._health_endpoints.run_with_timeout",
mock_run_with_timeout,
),
patch(
"litellm.proxy.health_endpoints._health_endpoints._reject_os_environ_references",
mock_reject_os_environ,
),
):
await health_test_model_connection(
request=mock_request,
mode="chat",
litellm_params={"model": "azure/gpt-4o"},
model_info={"id": "dep-1"},
user_api_key_dict=mock_user_api_key_dict,
)
model_params = mock_ahealth_check.call_args.kwargs.get("model_params", {})
assert "max_tokens" not in model_params, (
"max_tokens must not be sent when server model_info has "
"health_check_supports_max_tokens: False; "
f"got model_params={model_params}"
)
@pytest.mark.asyncio
async def test_health_services_endpoint_datadog_llm_observability():
"""