mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-28 01:32:17 +00:00
test(model_management): cover actor edges and wildcard models
This commit is contained in:
Yuneng Jiang
parent
db37977307
commit
db17841b3d
2 changed files with 356 additions and 3 deletions
|
|
@ -1224,11 +1224,11 @@ class TestUpdateModel:
|
|||
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
|
||||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
patch(
|
||||
patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
side_effect=lambda value: value,
|
||||
),
|
||||
patch(
|
||||
patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
||||
new=AsyncMock(
|
||||
return_value=ReconcileOutcome(still_desired=None, live_after=None)
|
||||
|
|
@ -4021,7 +4021,7 @@ class TestPatchModelBlockedAuthGate:
|
|||
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
|
||||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
patch(
|
||||
patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
||||
new=AsyncMock(
|
||||
return_value=ReconcileOutcome(still_desired=None, live_after=None)
|
||||
|
|
@ -6631,3 +6631,337 @@ class TestTeamMemberAutoRouterWrites:
|
|||
assert json.loads(written["model_info"])["member_auto_router"] is True
|
||||
assert appended.await_args.kwargs["data"].models == ["new-personal-router"]
|
||||
assert appended.await_args.kwargs["data"].team_id == "member-team"
|
||||
|
||||
|
||||
class TestModelManagementActorEdges:
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_model_rejects_non_team_internal_user(self):
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
|
||||
|
||||
actor: Final = UserAPIKeyAuth(user_id="internal-user", user_role=LitellmUserRoles.INTERNAL_USER)
|
||||
prisma: Final = MagicMock()
|
||||
deployment: Final = Deployment(
|
||||
model_name="internal-model",
|
||||
litellm_params=LiteLLM_Params(model="openai/test-model"),
|
||||
model_info=ModelInfo(id="internal-model-id"),
|
||||
)
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await add_new_model(model_params=deployment, user_api_key_dict=actor)
|
||||
|
||||
assert str(exc_info.value.code) == "403"
|
||||
assert "permission" in str(exc_info.value).lower()
|
||||
prisma.db.litellm_proxymodeltable.create.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_model_rejects_proxy_admin_viewer(self):
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
|
||||
|
||||
actor: Final = UserAPIKeyAuth(
|
||||
user_id="view-only-user", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
|
||||
)
|
||||
prisma: Final = MagicMock()
|
||||
deployment: Final = Deployment(
|
||||
model_name="view-only-model",
|
||||
litellm_params=LiteLLM_Params(model="openai/test-model"),
|
||||
model_info=ModelInfo(id="view-only-model-id"),
|
||||
)
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await add_new_model(model_params=deployment, user_api_key_dict=actor)
|
||||
|
||||
assert str(exc_info.value.code) == "403"
|
||||
assert "view-only" in str(exc_info.value).lower()
|
||||
prisma.db.litellm_proxymodeltable.create.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_model_requires_database_storage(self):
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
|
||||
|
||||
actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
prisma: Final = MagicMock()
|
||||
deployment: Final = Deployment(
|
||||
model_name="database-disabled-model",
|
||||
litellm_params=LiteLLM_Params(model="openai/test-model"),
|
||||
model_info=ModelInfo(id="database-disabled-model-id"),
|
||||
)
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", False), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await add_new_model(model_params=deployment, user_api_key_dict=actor)
|
||||
|
||||
assert str(exc_info.value.code) == "500"
|
||||
assert "STORE_MODEL_IN_DB" in str(exc_info.value)
|
||||
prisma.db.litellm_proxymodeltable.create.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_legacy_model_update_persists_changed_field(self):
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import update_model
|
||||
|
||||
model_id: Final = "legacy-update-model-id"
|
||||
existing_row: Final = MagicMock()
|
||||
existing_row.litellm_params = {"model": "openai/test-model", "timeout": 30}
|
||||
existing_row.model_dump.return_value = {
|
||||
"model_name": "legacy-update-model",
|
||||
"litellm_params": existing_row.litellm_params,
|
||||
"model_info": {"id": model_id},
|
||||
}
|
||||
existing_row.model_dump_json.return_value = "{}"
|
||||
updated_row: Final = MagicMock()
|
||||
updated_row.model_dump_json.return_value = "{}"
|
||||
prisma: Final = MagicMock()
|
||||
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row)
|
||||
prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row)
|
||||
router: Final = MagicMock()
|
||||
router.get_model_ids.return_value = [model_id]
|
||||
actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
|
||||
patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
side_effect=lambda value: value,
|
||||
),
|
||||
patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
||||
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
|
||||
),
|
||||
):
|
||||
await update_model(
|
||||
model_params=updateDeployment(
|
||||
litellm_params=updateLiteLLMParams(timeout=42),
|
||||
model_info=ModelInfo(id=model_id),
|
||||
),
|
||||
user_api_key_dict=actor,
|
||||
)
|
||||
|
||||
written: Final = json.loads(
|
||||
prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"]["litellm_params"]
|
||||
)
|
||||
assert written["timeout"] == 42
|
||||
assert written["model"] == "openai/test-model"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_legacy_model_update_explicit_null_preserves_existing_field(self):
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import update_model
|
||||
|
||||
model_id: Final = "legacy-null-model-id"
|
||||
existing_row: Final = MagicMock()
|
||||
existing_row.litellm_params = {"model": "openai/test-model", "timeout": 30}
|
||||
existing_row.model_dump.return_value = {
|
||||
"model_name": "legacy-null-model",
|
||||
"litellm_params": existing_row.litellm_params,
|
||||
"model_info": {"id": model_id},
|
||||
}
|
||||
existing_row.model_dump_json.return_value = "{}"
|
||||
updated_row: Final = MagicMock()
|
||||
updated_row.model_dump_json.return_value = "{}"
|
||||
prisma: Final = MagicMock()
|
||||
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row)
|
||||
prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row)
|
||||
router: Final = MagicMock()
|
||||
router.get_model_ids.return_value = [model_id]
|
||||
actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
|
||||
patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
side_effect=lambda value: value,
|
||||
),
|
||||
patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
||||
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
|
||||
),
|
||||
):
|
||||
await update_model(
|
||||
model_params=updateDeployment(
|
||||
litellm_params=updateLiteLLMParams(timeout=None),
|
||||
model_info=ModelInfo(id=model_id),
|
||||
),
|
||||
user_api_key_dict=actor,
|
||||
)
|
||||
|
||||
written: Final = json.loads(
|
||||
prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"]["litellm_params"]
|
||||
)
|
||||
assert written["timeout"] == 30
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_patch_model_rejects_config_file_model(self):
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import patch_model
|
||||
|
||||
model_id: Final = "config-model-id"
|
||||
prisma: Final = MagicMock()
|
||||
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=None)
|
||||
prisma.db.litellm_proxymodeltable.update = AsyncMock()
|
||||
router: Final = MagicMock()
|
||||
router.get_deployment.return_value = Deployment(
|
||||
model_name="config-model",
|
||||
litellm_params=LiteLLM_Params(model="openai/test-model"),
|
||||
model_info=ModelInfo(id=model_id),
|
||||
)
|
||||
actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] patch endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] patch endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] patch endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] patch endpoint reads proxy-server state through its only test seam
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await patch_model(
|
||||
model_id=model_id,
|
||||
patch_data=updateDeployment(
|
||||
litellm_params=updateLiteLLMParams(timeout=42),
|
||||
model_info=ModelInfo(id=model_id),
|
||||
),
|
||||
user_api_key_dict=actor,
|
||||
)
|
||||
|
||||
assert str(exc_info.value.code) == "400"
|
||||
assert "Cannot edit config-based model" in str(exc_info.value)
|
||||
prisma.db.litellm_proxymodeltable.update.assert_not_awaited()
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _client_for(self, actor: UserAPIKeyAuth) -> Iterator[TestClient]:
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
app.dependency_overrides[proxy_server.user_api_key_auth] = lambda: actor
|
||||
try:
|
||||
yield TestClient(app)
|
||||
finally:
|
||||
app.dependency_overrides.pop(proxy_server.user_api_key_auth, None)
|
||||
|
||||
def test_post_model_new_binds_to_actor_guard(self):
|
||||
actor: Final = UserAPIKeyAuth(user_id="internal-user", user_role=LitellmUserRoles.INTERNAL_USER)
|
||||
prisma: Final = MagicMock()
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
|
||||
self._client_for(actor) as client,
|
||||
):
|
||||
response: Final = client.post(
|
||||
"/model/new",
|
||||
json={
|
||||
"model_name": "internal-model",
|
||||
"litellm_params": {"model": "openai/test-model"},
|
||||
"model_info": {"id": "internal-model-id"},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 403
|
||||
assert "permission" in response.text.lower()
|
||||
prisma.db.litellm_proxymodeltable.create.assert_not_called()
|
||||
|
||||
def test_post_legacy_model_update_binds_to_persistence(self):
|
||||
model_id: Final = "legacy-route-model-id"
|
||||
existing_row: Final = LiteLLM_ProxyModelTable(
|
||||
model_id=model_id,
|
||||
model_name="legacy-route-model",
|
||||
litellm_params={"model": "openai/test-model", "timeout": 30},
|
||||
model_info={"id": model_id},
|
||||
created_by="admin",
|
||||
updated_by="admin",
|
||||
)
|
||||
updated_row: Final = LiteLLM_ProxyModelTable(
|
||||
model_id=model_id,
|
||||
model_name="legacy-route-model",
|
||||
litellm_params={"model": "openai/test-model", "timeout": 42},
|
||||
model_info={"id": model_id},
|
||||
created_by="admin",
|
||||
updated_by="admin",
|
||||
)
|
||||
prisma: Final = MagicMock()
|
||||
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row)
|
||||
prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row)
|
||||
router: Final = MagicMock()
|
||||
router.get_model_ids.return_value = [model_id]
|
||||
actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
|
||||
patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
side_effect=lambda value: value,
|
||||
),
|
||||
patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
||||
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
|
||||
),
|
||||
patch( # test-quality-ok: [TQ008] audit logging is outside the persistence contract
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.create_object_audit_log",
|
||||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
self._client_for(actor) as client,
|
||||
):
|
||||
response: Final = client.post(
|
||||
"/model/update",
|
||||
json={
|
||||
"litellm_params": {"timeout": 42},
|
||||
"model_info": {"id": model_id},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
written: Final = json.loads(
|
||||
prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"]["litellm_params"]
|
||||
)
|
||||
assert written["timeout"] == 42
|
||||
|
||||
def test_patch_config_model_binds_to_patch_route(self):
|
||||
model_id: Final = "config-route-model-id"
|
||||
prisma: Final = MagicMock()
|
||||
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=None)
|
||||
prisma.db.litellm_proxymodeltable.update = AsyncMock()
|
||||
router: Final = MagicMock()
|
||||
router.get_deployment.return_value = Deployment(
|
||||
model_name="config-route-model",
|
||||
litellm_params=LiteLLM_Params(model="openai/test-model"),
|
||||
model_info=ModelInfo(id=model_id),
|
||||
)
|
||||
actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
|
||||
self._client_for(actor) as client,
|
||||
):
|
||||
response: Final = client.patch(
|
||||
f"/model/{model_id}/update",
|
||||
json={
|
||||
"litellm_params": {"timeout": 42},
|
||||
"model_info": {"id": model_id},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "Cannot edit config-based model" in response.text
|
||||
prisma.db.litellm_proxymodeltable.update.assert_not_awaited()
|
||||
|
|
|
|||
|
|
@ -101,4 +101,23 @@ describe("prepareModelAddRequest", () => {
|
|||
expect(deployment.litellmParamsObj.litellm_credential_name).toBe("from-json");
|
||||
expect(deployment.litellmParamsObj.timeout).toBe(5);
|
||||
});
|
||||
|
||||
it.each([
|
||||
["OpenAI", "openai/*"],
|
||||
["Azure_AI_Studio", "azure_ai/*"],
|
||||
["Petals", "petals/*"],
|
||||
])("composes wildcard names for the all-model selection", async (custom_llm_provider, wildcardModel) => {
|
||||
const formValues = {
|
||||
model_mappings: [],
|
||||
model: "all-wildcard",
|
||||
custom_llm_provider,
|
||||
};
|
||||
|
||||
const deployments = await prepareModelAddRequest({ ...formValues }, "token", null);
|
||||
|
||||
expect(deployments).toHaveLength(1);
|
||||
const [deployment] = deployments!;
|
||||
expect(deployment.modelName).toBe(wildcardModel);
|
||||
expect(deployment.litellmParamsObj.model).toBe(wildcardModel);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue