mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
9f4e3c6009
Every model-write endpoint returned 200 off the DB write alone; a model the reload dropped (ignore_invalid_deployments, or a wholesale reload failure) stayed invisible on every channel at once, which is how the registry-leak defect went undiagnosed for three weeks. ProxyConfig.add_deployment and clear_cache now return whether the reload pass completed, and each write endpoint verifies the rows it wrote are live in this pod's router afterwards, distinguishing a deliberately environment-inactive model via the same predicate the Router's own gate uses. The access-group writers return the mutated id set instead of discarding it
247 lines
8.4 KiB
Python
247 lines
8.4 KiB
Python
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm.proxy._types import (
|
|
BlockModelRequest,
|
|
LitellmUserRoles,
|
|
ProxyException,
|
|
UserAPIKeyAuth,
|
|
)
|
|
from litellm.types.router import RouterRateLimitError
|
|
|
|
|
|
def _setup_model_block_mocks(monkeypatch, *, updated_blocked: bool):
|
|
model_id = "model-123"
|
|
|
|
existing_row = MagicMock()
|
|
existing_row.model_dump.return_value = {
|
|
"model_name": "gpt-4o",
|
|
"litellm_params": {"model": "openai/gpt-4o"},
|
|
"model_info": {"id": model_id},
|
|
}
|
|
|
|
updated_row = MagicMock()
|
|
updated_row.model_id = model_id
|
|
updated_row.blocked = updated_blocked
|
|
|
|
model_table = MagicMock()
|
|
model_table.find_unique = AsyncMock(return_value=existing_row)
|
|
model_table.update = AsyncMock(return_value=updated_row)
|
|
|
|
mock_prisma_client = MagicMock()
|
|
mock_prisma_client.db.litellm_proxymodeltable = model_table
|
|
|
|
mock_router = MagicMock()
|
|
mock_router.get_model_ids.return_value = [model_id]
|
|
|
|
mock_clear_cache = AsyncMock(return_value=True)
|
|
mock_audit_log = AsyncMock(return_value=None)
|
|
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
|
mock_clear_cache,
|
|
)
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.management_endpoints.model_management_endpoints.create_object_audit_log",
|
|
mock_audit_log,
|
|
)
|
|
|
|
return model_id, model_table, updated_row, mock_clear_cache, mock_audit_log
|
|
|
|
|
|
def _proxy_admin() -> UserAPIKeyAuth:
|
|
return UserAPIKeyAuth(
|
|
user_id="admin",
|
|
user_role=LitellmUserRoles.PROXY_ADMIN,
|
|
api_key="sk-admin",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_model_block_endpoint_sets_blocked_true(monkeypatch):
|
|
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
|
block_model,
|
|
)
|
|
|
|
model_id, model_table, updated_row, mock_clear_cache, mock_audit_log = (
|
|
_setup_model_block_mocks(monkeypatch, updated_blocked=True)
|
|
)
|
|
|
|
result = await block_model(
|
|
data=BlockModelRequest(model_id=model_id),
|
|
http_request=MagicMock(),
|
|
user_api_key_dict=_proxy_admin(),
|
|
litellm_changed_by="operator@example.com",
|
|
)
|
|
|
|
assert result == updated_row
|
|
model_table.update.assert_awaited_once()
|
|
update_kwargs = model_table.update.await_args.kwargs
|
|
assert update_kwargs["where"] == {"model_id": model_id}
|
|
assert update_kwargs["data"]["blocked"] is True
|
|
assert update_kwargs["data"]["updated_by"] == "admin"
|
|
assert "updated_at" in update_kwargs["data"]
|
|
mock_clear_cache.assert_awaited_once_with()
|
|
assert mock_audit_log.call_args.kwargs["action"] == "blocked"
|
|
assert (
|
|
mock_audit_log.call_args.kwargs["litellm_changed_by"] == "operator@example.com"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_model_unblock_endpoint_sets_blocked_false(monkeypatch):
|
|
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
|
unblock_model,
|
|
)
|
|
|
|
model_id, model_table, updated_row, mock_clear_cache, mock_audit_log = (
|
|
_setup_model_block_mocks(monkeypatch, updated_blocked=False)
|
|
)
|
|
|
|
result = await unblock_model(
|
|
data=BlockModelRequest(model_id=model_id),
|
|
http_request=MagicMock(),
|
|
user_api_key_dict=_proxy_admin(),
|
|
litellm_changed_by=None,
|
|
)
|
|
|
|
assert result == updated_row
|
|
model_table.update.assert_awaited_once()
|
|
assert model_table.update.await_args.kwargs["data"]["blocked"] is False
|
|
mock_clear_cache.assert_awaited_once_with()
|
|
assert mock_audit_log.call_args.kwargs["action"] == "unblocked"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_model_block_endpoint_requires_proxy_admin(monkeypatch):
|
|
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
|
block_model,
|
|
)
|
|
|
|
model_id, model_table, _, _, _ = _setup_model_block_mocks(
|
|
monkeypatch, updated_blocked=True
|
|
)
|
|
non_admin = UserAPIKeyAuth(
|
|
user_id="internal-user",
|
|
user_role=LitellmUserRoles.INTERNAL_USER,
|
|
api_key="sk-user",
|
|
)
|
|
|
|
with pytest.raises(ProxyException) as exc_info:
|
|
await block_model(
|
|
data=BlockModelRequest(model_id=model_id),
|
|
http_request=MagicMock(),
|
|
user_api_key_dict=non_admin,
|
|
litellm_changed_by=None,
|
|
)
|
|
|
|
assert exc_info.value.code == "403"
|
|
assert "Only proxy admins" in exc_info.value.message
|
|
model_table.update.assert_not_awaited()
|
|
|
|
|
|
def test_router_returns_no_healthy_deployment_when_model_is_fully_blocked():
|
|
router = litellm.Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-4o",
|
|
"litellm_params": {"model": "openai/gpt-4o-0"},
|
|
"model_info": {"id": "dep-0", "blocked": True},
|
|
},
|
|
{
|
|
"model_name": "gpt-4o",
|
|
"litellm_params": {"model": "openai/gpt-4o-1"},
|
|
"model_info": {"id": "dep-1", "blocked": True},
|
|
},
|
|
]
|
|
)
|
|
|
|
with pytest.raises(RouterRateLimitError) as exc_info:
|
|
router.get_available_deployment(model="gpt-4o", request_kwargs={})
|
|
|
|
assert "No deployments available for selected model" in str(exc_info.value)
|
|
assert "Passed model=gpt-4o" in str(exc_info.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_route_request_returns_403_when_model_is_fully_blocked(monkeypatch):
|
|
from litellm.proxy.route_llm_request import route_request
|
|
|
|
router = litellm.Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-4o",
|
|
"litellm_params": {"model": "openai/gpt-4o"},
|
|
"model_info": {"id": "dep-0", "blocked": True},
|
|
}
|
|
]
|
|
)
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.route_llm_request.add_shared_session_to_data",
|
|
AsyncMock(return_value=None),
|
|
)
|
|
|
|
with pytest.raises(litellm.PermissionDeniedError) as exc_info:
|
|
await route_request(
|
|
data={"model": "gpt-4o"},
|
|
llm_router=router,
|
|
user_model=None,
|
|
route_type="acreate_eval",
|
|
)
|
|
|
|
assert exc_info.value.status_code == 403
|
|
assert "Model is blocked" in exc_info.value.message
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_model_block_surfaces_wholesale_reload_failure(monkeypatch):
|
|
"""The write endpoints owe the caller an error when the pod failed to reload at all;
|
|
the DB row is saved but this pod is not serving the change."""
|
|
from litellm.proxy._types import ProxyException
|
|
from litellm.proxy.management_endpoints.model_management_endpoints import block_model
|
|
|
|
model_id, model_table, updated_row, mock_clear_cache, mock_audit_log = _setup_model_block_mocks(
|
|
monkeypatch, updated_blocked=True
|
|
)
|
|
wiped_router = MagicMock()
|
|
wiped_router.get_model_ids.side_effect = [[model_id], []]
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", wiped_router)
|
|
|
|
with pytest.raises(ProxyException, match=model_id):
|
|
await block_model(
|
|
data=BlockModelRequest(model_id=model_id),
|
|
http_request=MagicMock(),
|
|
user_api_key_dict=_proxy_admin(),
|
|
litellm_changed_by="operator@example.com",
|
|
)
|
|
|
|
assert mock_audit_log.call_args.kwargs["object_id"] == model_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_model_block_surfaces_model_dropped_by_reload(monkeypatch):
|
|
"""A reload that completes but drops the written model (ignore_invalid_deployments
|
|
swallowed its re-add) must not produce an unqualified success."""
|
|
from litellm.proxy._types import ProxyException
|
|
from litellm.proxy.management_endpoints.model_management_endpoints import block_model
|
|
|
|
model_id, model_table, updated_row, mock_clear_cache, _ = _setup_model_block_mocks(
|
|
monkeypatch, updated_blocked=True
|
|
)
|
|
dropped_router = MagicMock()
|
|
dropped_router.get_model_ids.return_value = []
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", dropped_router)
|
|
|
|
with pytest.raises(ProxyException, match=model_id):
|
|
await block_model(
|
|
data=BlockModelRequest(model_id=model_id),
|
|
http_request=MagicMock(),
|
|
user_api_key_dict=_proxy_admin(),
|
|
litellm_changed_by="operator@example.com",
|
|
)
|