diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 7cd1107abce..a8588adfd44 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -55,13 +55,18 @@ async def test_model_creation_disabled_before_any_write( from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model prisma: Final = MagicMock() - prisma.writer_db.litellm_uisettings.find_unique = AsyncMock(return_value=SimpleNamespace( - ui_settings={"disable_model_add_for_internal_users": True}, - )) - prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable( - team_id="test_team", members_with_roles=[Member(user_id="test_user", role=team_role)], - team_member_permissions=["/auto_router/manage"], - )) + prisma.writer_db.litellm_uisettings.find_unique = AsyncMock( + return_value=SimpleNamespace( + ui_settings={"disable_model_add_for_internal_users": True}, + ) + ) + prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=LiteLLM_TeamTable( + team_id="test_team", + members_with_roles=[Member(user_id="test_user", role=team_role)], + team_member_permissions=["/auto_router/manage"], + ) + ) monkeypatch.setattr(proxy_server, "prisma_client", prisma) monkeypatch.setattr(proxy_server, "general_settings", {"disable_model_add_for_internal_users": False}) monkeypatch.setattr(proxy_server, "premium_user", True) @@ -73,7 +78,8 @@ async def test_model_creation_disabled_before_any_write( model_info={"team_id": "test_team"}, ), user_api_key_dict=UserAPIKeyAuth( - user_id="test_user", user_role=LitellmUserRoles.INTERNAL_USER, + user_id="test_user", + user_role=LitellmUserRoles.INTERNAL_USER, team_id="test_team" if team_key else None, ), ) @@ -371,9 +377,7 @@ class TestModelManagementAuthChecks: ModelManagementAuthChecks.can_user_attach_credential( litellm_params=litellm_params(litellm_credential_name=None), user_api_key_dict=self.team_admin_user, - existing_litellm_params=LiteLLM_Params( - model="test_model", litellm_credential_name="shared-credential" - ), + existing_litellm_params=LiteLLM_Params(model="test_model", litellm_credential_name="shared-credential"), null_detaches=True, ) @@ -386,9 +390,7 @@ class TestModelManagementAuthChecks: result = ModelManagementAuthChecks.can_user_attach_credential( litellm_params=litellm_params(litellm_credential_name=None), user_api_key_dict=self.admin_user, - existing_litellm_params=LiteLLM_Params( - model="test_model", litellm_credential_name="shared-credential" - ), + existing_litellm_params=LiteLLM_Params(model="test_model", litellm_credential_name="shared-credential"), null_detaches=True, ) @@ -412,9 +414,7 @@ class TestModelManagementAuthChecks: result = ModelManagementAuthChecks.can_user_attach_credential( litellm_params=litellm_params(litellm_credential_name=None), user_api_key_dict=self.team_admin_user, - existing_litellm_params=LiteLLM_Params( - model="test_model", litellm_credential_name="shared-credential" - ), + existing_litellm_params=LiteLLM_Params(model="test_model", litellm_credential_name="shared-credential"), ) assert result is True @@ -440,9 +440,15 @@ class TestModelManagementAuthChecks: mock_prisma = MagicMock() mock_prisma.writer_db.litellm_uisettings.find_unique = AsyncMock(return_value=None) with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.premium_user", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch( # test-quality-ok: prior auth check needs a live DB; only the credential check is under test "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None), @@ -477,10 +483,18 @@ class TestModelManagementAuthChecks: model_info={"id": model_id}, ) with ( - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.llm_router", MagicMock()), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.llm_router", MagicMock() + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.premium_user", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch( # test-quality-ok: stubs the DB row fetch; only the credential check is under test "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", new=AsyncMock(return_value=db_model), @@ -571,9 +585,15 @@ class TestModelManagementAuthChecks: mock_prisma = MagicMock() mock_prisma.writer_db.litellm_uisettings.find_unique = AsyncMock(return_value=None) with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.premium_user", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch( # test-quality-ok: prior auth check needs a live DB; only the session tag check is under test "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None), @@ -614,10 +634,18 @@ class TestModelManagementAuthChecks: model_info={"id": model_id, "team_id": "test_team"}, ) with ( - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.llm_router", MagicMock()), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.llm_router", MagicMock() + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.premium_user", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch( # test-quality-ok: stubs the DB row fetch; only the session tag check is under test "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", new=AsyncMock(return_value=db_model), @@ -666,10 +694,18 @@ class TestModelManagementAuthChecks: mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) mock_prisma.db.litellm_proxymodeltable.update = AsyncMock() with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.llm_router", None + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.premium_user", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch( # test-quality-ok: prior auth check needs a live DB; only the session tag check is under test "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None), @@ -1497,16 +1533,16 @@ class TestUpdateModel: ): await update_model( model_params=updateDeployment( - litellm_params=updateLiteLLMParams( - model="openai/gpt-4o-mini", litellm_credential_name=None - ), + litellm_params=updateLiteLLMParams(model="openai/gpt-4o-mini", litellm_credential_name=None), model_info=ModelInfo(id=model_id), ), user_api_key_dict=team_admin, ) mock_prisma.db.litellm_proxymodeltable.update.assert_awaited_once() - persisted = json.loads(mock_prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"]["litellm_params"]) + persisted = json.loads( + mock_prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"]["litellm_params"] + ) assert persisted["litellm_credential_name"] == "shared-credential" @@ -1970,7 +2006,9 @@ class TestTeamModelUpdate: "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.allow_team_model_action", AsyncMock(return_value=True), ), - patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: team models are premium-gated through a proxy global with no injection seam + patch( + "litellm.proxy.proxy_server.premium_user", True + ), # test-quality-ok: team models are premium-gated through a proxy global with no injection seam patch( # test-quality-ok: the team list write is the collaborator whose ordering is asserted "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add", side_effect=team_add, @@ -4361,9 +4399,7 @@ class TestModelInfoCostMapEchoFilter: for deployment_entry in (litellm.get_model_info(deployment_model),) for key_entry in (litellm.get_model_info(key_model),) for differing_fields in ( - frozenset( - k for k in deployment_entry if k in key_entry and deployment_entry[k] != key_entry[k] - ), + frozenset(k for k in deployment_entry if k in key_entry and deployment_entry[k] != key_entry[k]), ) if differing_fields ) @@ -4456,9 +4492,7 @@ class TestUpdateDBModelClearCacheControlInjectionPoints: ), model_info=ModelInfo(id="dep-cache-0"), ) - patch = updateDeployment( - litellm_params=updateLiteLLMParams(cache_control_injection_points=None) - ) + patch = updateDeployment(litellm_params=updateLiteLLMParams(cache_control_injection_points=None)) result = update_db_model(db_model=db_model, updated_patch=patch) @@ -4504,11 +4538,12 @@ class TestUpdateDBModelClearCredentialName: ), model_info=ModelInfo(id="dep-cred-1", team_id="team-keep", access_groups=["prod"]), ) - update_patch: Final = updateDeployment( - litellm_params=updateLiteLLMParams(litellm_credential_name=None) - ) + update_patch: Final = updateDeployment(litellm_params=updateLiteLLMParams(litellm_credential_name=None)) - with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + with patch( + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + side_effect=lambda value, **kwargs: value, + ): result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) params: Final = json.loads(result["litellm_params"]) @@ -4537,7 +4572,10 @@ class TestUpdateDBModelClearCredentialName: ) update_patch: Final = updateDeployment(litellm_params=updateLiteLLMParams(tpm=10)) - with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + with patch( + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + side_effect=lambda value, **kwargs: value, + ): result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) params: Final = json.loads(result["litellm_params"]) @@ -4552,11 +4590,12 @@ class TestUpdateDBModelClearCredentialName: litellm_params=LiteLLM_Params(model="openai/gpt-4o", api_base="https://api.openai.com/v1"), model_info=ModelInfo(id="dep-cred-1", team_id="team-keep", access_groups=["prod"]), ) - update_patch: Final = updateDeployment( - litellm_params=updateLiteLLMParams(litellm_credential_name=None) - ) + update_patch: Final = updateDeployment(litellm_params=updateLiteLLMParams(litellm_credential_name=None)) - with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + with patch( + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + side_effect=lambda value, **kwargs: value, + ): result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) params: Final = json.loads(result["litellm_params"]) @@ -4584,7 +4623,10 @@ class TestUpdateDBModelClearCredentialName: ) ) - with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + with patch( + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + side_effect=lambda value, **kwargs: value, + ): result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) params: Final = json.loads(result["litellm_params"]) @@ -4611,7 +4653,10 @@ class TestUpdateDBModelClearCredentialName: litellm_params=updateLiteLLMParams(litellm_credential_name="other-credential") ) - with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + with patch( + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + side_effect=lambda value, **kwargs: value, + ): result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) params: Final = json.loads(result["litellm_params"]) @@ -4690,9 +4735,7 @@ class TestPatchModelCredentialName: "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value, ), - patch( - "litellm.proxy.management_endpoints.model_management_endpoints.raise_if_reload_degraded_serving" - ), + patch("litellm.proxy.management_endpoints.model_management_endpoints.raise_if_reload_degraded_serving"), patch( "litellm.proxy.management_endpoints.model_management_endpoints.create_object_audit_log", new=AsyncMock(), @@ -4958,9 +5001,7 @@ class TestPatchModelBlockedAuthGate: existing_row.model_dump_json.return_value = "{}" mock_prisma = MagicMock() - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=existing_row - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) with ( patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), @@ -5001,12 +5042,8 @@ class TestPatchModelBlockedAuthGate: updated_row.model_dump_json.return_value = "{}" mock_prisma = MagicMock() - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=existing_row - ) - mock_prisma.db.litellm_proxymodeltable.update = AsyncMock( - return_value=updated_row - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row) with ( patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), @@ -5019,9 +5056,7 @@ class TestPatchModelBlockedAuthGate: ), 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) - ), + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), ), ): result = await patch_model( @@ -5056,25 +5091,29 @@ class TestPatchModelRowDeletedBeforeWrite: existing_row.model_dump_json.return_value = "{}" mock_prisma = MagicMock() - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=existing_row - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) mock_prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=None) with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy_server module global is the endpoint's only injection point - patch("litellm.proxy.proxy_server.llm_router", MagicMock(**{"get_model_ids.return_value": ["m1"]})), # test-quality-ok: proxy_server module global is the endpoint's only injection point - patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: proxy_server module global is the endpoint's only injection point - patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch( + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch( + "litellm.proxy.proxy_server.llm_router", MagicMock(**{"get_model_ids.return_value": ["m1"]}) + ), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch( + "litellm.proxy.proxy_server.premium_user", True + ), # test-quality-ok: proxy_server module global is the endpoint's only injection point patch( # test-quality-ok: stubs the auth gate so the test exercises the not-found branch under test "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None), ), patch( # test-quality-ok: stubs the cache write so the test observes only the DB result handling "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", - new=AsyncMock( - return_value=ReconcileOutcome(still_desired=None, live_after=None) - ), + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), ), ): with pytest.raises(ProxyException) as exc_info: @@ -5158,9 +5197,7 @@ class TestWriteSurfacesReloadDrop: ) with pytest.raises(ProxyException, match="m-gone"): - raise_if_reload_degraded_serving( - before=frozenset(), written_models=[("m-gone", None)], action="update" - ) + raise_if_reload_degraded_serving(before=frozenset(), written_models=[("m-gone", None)], action="update") with pytest.raises(ProxyException, match="m-collateral"): raise_if_reload_degraded_serving( @@ -5275,10 +5312,7 @@ class TestConcurrentModelWritesDoNotEvictEachOther: config = ProxyConfig() await asyncio.gather( - *[ - config.add_deployment(prisma_client=MagicMock(), proxy_logging_obj=MagicMock()) - for _ in range(5) - ] + *[config.add_deployment(prisma_client=MagicMock(), proxy_logging_obj=MagicMock()) for _ in range(5)] ) assert observed_max == 1 @@ -5507,9 +5541,7 @@ class TestDeleteEvictionsHoldTheReconcileLock: ) async def call() -> None: - await delete_team_models( - team_ids=["team-1"], prisma_client=prisma, llm_router=router - ) + await delete_team_models(team_ids=["team-1"], prisma_client=prisma, llm_router=router) await self._assert_evicts_under_lock(monkeypatch, call, model_id, config) router.delete_deployment.assert_called_once_with(id=model_id) @@ -6039,7 +6071,9 @@ class TestStrategyRouterWriteValidation: **_FORECAST_BASE, "classifier_type": "capability", "capability_classifier_config": { - "efficient_tier": "SIMPLE", "capable_tier": "REASONING", "base_threshold": 0.7, + "efficient_tier": "SIMPLE", + "capable_tier": "REASONING", + "base_threshold": 0.7, }, } _FUSE = { @@ -6047,8 +6081,10 @@ class TestStrategyRouterWriteValidation: "classifier_type": "llm_v2", "adaptive": False, "llm_v2_config": { - "efficient_profile": "Small solver", "capable_profile": "Large solver", - "harness": "One attempt", "max_quality_gap": 0.05, + "efficient_profile": "Small solver", + "capable_profile": "Large solver", + "harness": "One attempt", + "max_quality_gap": 0.05, }, } _CUSTOM_TIERS = { @@ -6099,17 +6135,17 @@ class TestStrategyRouterWriteValidation: ("no-config", _V2, _V2), ], ) - def test_effective_complexity_router_config( - self, incoming: object, existing: object, expected: object - ) -> None: + def test_effective_complexity_router_config(self, incoming: object, existing: object, expected: object) -> None: """A write is judged on the config it leaves on the row: the incoming one when it carries one, else the stored one.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( _effective_complexity_router_config, ) from litellm.types.router import updateLiteLLMParams - incoming_params = None if incoming is None else updateLiteLLMParams( - complexity_router_config=None if incoming == "no-config" else incoming + incoming_params = ( + None + if incoming is None + else updateLiteLLMParams(complexity_router_config=None if incoming == "no-config" else incoming) ) existing_params = None if existing is None else updateLiteLLMParams(complexity_router_config=existing) assert _effective_complexity_router_config(incoming_params, existing_params) == expected @@ -6118,32 +6154,207 @@ class TestStrategyRouterWriteValidation: @pytest.mark.parametrize( "limit,effective_params,db_models,config_config,model_id,expected", [ - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _CAPABILITY}, ["auto_router/complexity_router"], None, None, "refused"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _CAPABILITY}, [], _CAPABILITY, None, "refused"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _CAPABILITY}, [], _FUSE, None, "reserved"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _CAPABILITY}, [], None, "held-id", "reserved"), - (None, {"model": "auto_router/complexity_router", "complexity_router_config": _CAPABILITY}, ["auto_router/complexity_router"], _CAPABILITY, None, "plain"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _FUSE}, ["auto_router/complexity_router"], None, None, "refused"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _FUSE}, [], _FUSE, None, "refused"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _FUSE}, [], _CAPABILITY, None, "reserved"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _FUSE}, [], None, "held-id", "reserved"), - (None, {"model": "auto_router/complexity_router", "complexity_router_config": _FUSE}, ["auto_router/complexity_router"], _FUSE, None, "plain"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, ["auto_router/complexity_router"], None, None, "refused"), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _CAPABILITY}, + ["auto_router/complexity_router"], + None, + None, + "refused", + ), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _CAPABILITY}, + [], + _CAPABILITY, + None, + "refused", + ), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _CAPABILITY}, + [], + _FUSE, + None, + "reserved", + ), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _CAPABILITY}, + [], + None, + "held-id", + "reserved", + ), + ( + None, + {"model": "auto_router/complexity_router", "complexity_router_config": _CAPABILITY}, + ["auto_router/complexity_router"], + _CAPABILITY, + None, + "plain", + ), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _FUSE}, + ["auto_router/complexity_router"], + None, + None, + "refused", + ), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _FUSE}, + [], + _FUSE, + None, + "refused", + ), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _FUSE}, + [], + _CAPABILITY, + None, + "reserved", + ), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _FUSE}, + [], + None, + "held-id", + "reserved", + ), + ( + None, + {"model": "auto_router/complexity_router", "complexity_router_config": _FUSE}, + ["auto_router/complexity_router"], + _FUSE, + None, + "plain", + ), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, + ["auto_router/complexity_router"], + None, + None, + "refused", + ), (1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, [], _V2, None, "refused"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, [], None, None, "reserved"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, [], None, "held-id", "reserved"), - (2, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, ["auto_router/complexity_router"], None, None, "reserved"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_TIERS}, ["openai/gpt-4o"], None, None, "reserved"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_TIERS}, [], _CUSTOM_PROMPT, None, "refused"), - (1, {"model": "openai/gpt-4o", "complexity_router_config": _CUSTOM_TIERS}, ["auto_router/complexity_router"], None, None, "plain"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _V1}, ["auto_router/complexity_router"], _V2, None, "plain"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": None}, ["auto_router/complexity_router"], _V2, None, "plain"), - (None, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, ["auto_router/complexity_router"], _V2, None, "plain"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _TIER_LABELS_ONLY}, ["auto_router/complexity_router"], _CUSTOM_TIERS, None, "plain"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_PROMPT}, ["auto_router/complexity_router"], None, None, "refused"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_EXAMPLES}, [], _CUSTOM_TIERS, None, "refused"), - (1, {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_OPENING_PROMPT}, [], _CUSTOM_PROMPT, None, "refused"), - (None, {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_EXAMPLES}, ["auto_router/complexity_router"], _CUSTOM_TIERS, None, "plain"), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, + [], + None, + None, + "reserved", + ), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, + [], + None, + "held-id", + "reserved", + ), + ( + 2, + {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, + ["auto_router/complexity_router"], + None, + None, + "reserved", + ), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_TIERS}, + ["openai/gpt-4o"], + None, + None, + "reserved", + ), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_TIERS}, + [], + _CUSTOM_PROMPT, + None, + "refused", + ), + ( + 1, + {"model": "openai/gpt-4o", "complexity_router_config": _CUSTOM_TIERS}, + ["auto_router/complexity_router"], + None, + None, + "plain", + ), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _V1}, + ["auto_router/complexity_router"], + _V2, + None, + "plain", + ), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": None}, + ["auto_router/complexity_router"], + _V2, + None, + "plain", + ), + ( + None, + {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, + ["auto_router/complexity_router"], + _V2, + None, + "plain", + ), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _TIER_LABELS_ONLY}, + ["auto_router/complexity_router"], + _CUSTOM_TIERS, + None, + "plain", + ), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_PROMPT}, + ["auto_router/complexity_router"], + None, + None, + "refused", + ), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_EXAMPLES}, + [], + _CUSTOM_TIERS, + None, + "refused", + ), + ( + 1, + {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_OPENING_PROMPT}, + [], + _CUSTOM_PROMPT, + None, + "refused", + ), + ( + None, + {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_EXAMPLES}, + ["auto_router/complexity_router"], + _CUSTOM_TIERS, + None, + "plain", + ), ], ) async def test_auto_router_capability_slot_matrix( @@ -6172,10 +6383,16 @@ class TestStrategyRouterWriteValidation: capabilities: Final = gated_capabilities_of(effective_params) fake = self._FakeDb(db_models) - live_router = self._live_router_holding_one_capability(limit, config_config) if config_config is not None else None + live_router = ( + self._live_router_holding_one_capability(limit, config_config) if config_config is not None else None + ) with ( - patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: limit), # test-quality-ok: the guard reads the proxy license singleton with no injection seam - patch("litellm.proxy.proxy_server.llm_router", live_router), # test-quality-ok: the guard reads the proxy router global with no injection seam + patch( + "litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: limit + ), # test-quality-ok: the guard reads the proxy license singleton with no injection seam + patch( + "litellm.proxy.proxy_server.llm_router", live_router + ), # test-quality-ok: the guard reads the proxy router global with no injection seam patch( # test-quality-ok: the cross-pod publish is the side effect under test; redis is not configured here "litellm.proxy.management_endpoints.model_management_endpoints.publish_config_change", new=AsyncMock(), @@ -6191,7 +6408,9 @@ class TestStrategyRouterWriteValidation: assert capabilities[0].subject in str(exc_info.value.detail) assert "'auto_router' feature lifts the limit" in str(exc_info.value.detail) return - async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=model_id) as tables: + async with _auto_router_capability_slot( + fake, effective_params=effective_params, model_id=model_id + ) as tables: handle = tables if expected == "plain": await handle.create(data={}) @@ -6233,17 +6452,26 @@ class TestStrategyRouterWriteValidation: "local_heuristic": "heuristic_v2", **( {"heuristic_first_max_tier": "MEDIUM"} - if classifier == "heuristic_first" else {"hybrid_boundary_margin": 0.1} + if classifier == "heuristic_first" + else {"hybrid_boundary_margin": 0.1} ), }, } fake: Final = self._FakeDb([]) router: Final = self._live_router_holding_one_capability(1, other_config) if other_config is not None else None with ( - patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", return_value=1), # test-quality-ok: [TQ008] the guard reads the proxy license singleton - patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] the guard reads the proxy router singleton - patch("litellm.proxy.proxy_server.heuristic_v1_tuning_baselines", {}), # test-quality-ok: [TQ008] the guard reads the proxy baseline singleton - patch("litellm.proxy.management_endpoints.model_management_endpoints.publish_config_change", new=AsyncMock()), # test-quality-ok: [TQ008] Redis publication is outside admission under test + patch( + "litellm.proxy.proxy_server._license_check.auto_router_capability_limit", return_value=1 + ), # test-quality-ok: [TQ008] the guard reads the proxy license singleton + patch( + "litellm.proxy.proxy_server.llm_router", router + ), # test-quality-ok: [TQ008] the guard reads the proxy router singleton + patch( + "litellm.proxy.proxy_server.heuristic_v1_tuning_baselines", {} + ), # test-quality-ok: [TQ008] the guard reads the proxy baseline singleton + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.publish_config_change", new=AsyncMock() + ), # test-quality-ok: [TQ008] Redis publication is outside admission under test ): if denied_subject is not None: with pytest.raises(HTTPException) as failure: @@ -6314,7 +6542,10 @@ class TestStrategyRouterWriteValidation: "_MODELS_ONLY_B": self._MODELS_ONLY_B, } baselines = snapshot_tuning_baselines( - [self._db_router_row(row_id, configs["_TUNED_A" if row_id == "a" else "_TUNED_B"]) for row_id in baseline_rows] + [ + self._db_router_row(row_id, configs["_TUNED_A" if row_id == "a" else "_TUNED_B"]) + for row_id in baseline_rows + ] ) effective_params = { "model": "auto_router/complexity_router", @@ -6337,19 +6568,29 @@ class TestStrategyRouterWriteValidation: ], ) with ( - patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: limit), # test-quality-ok: the guard reads the proxy license singleton with no injection seam - patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: the guard reads the proxy router global with no injection seam - patch("litellm.proxy.proxy_server.heuristic_v1_tuning_baselines", baselines), # test-quality-ok: baselines are a startup-loaded proxy global with no injection seam + patch( + "litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: limit + ), # test-quality-ok: the guard reads the proxy license singleton with no injection seam + patch( + "litellm.proxy.proxy_server.llm_router", None + ), # test-quality-ok: the guard reads the proxy router global with no injection seam + patch( + "litellm.proxy.proxy_server.heuristic_v1_tuning_baselines", baselines + ), # test-quality-ok: baselines are a startup-loaded proxy global with no injection seam ): if expected == "refused": with pytest.raises(HTTPException) as exc_info: - async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=candidate_id): + async with _auto_router_capability_slot( + fake, effective_params=effective_params, model_id=candidate_id + ): pass assert exc_info.value.status_code == 403 assert "changed heuristic scoring rules" in str(exc_info.value.detail) assert "'auto_router' feature lifts the limit" in str(exc_info.value.detail) return - async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=candidate_id) as table: + async with _auto_router_capability_slot( + fake, effective_params=effective_params, model_id=candidate_id + ) as table: assert hasattr(table, "create") @pytest.mark.asyncio @@ -6423,12 +6664,24 @@ class TestStrategyRouterWriteValidation: ], ) with ( - patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.heuristic_v1_tuning_baselines", baselines), # test-quality-ok: baselines are a startup-loaded proxy global with no injection seam + patch( + "litellm.proxy.proxy_server.prisma_client", fake + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.llm_router", None + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.premium_user", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1 + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.heuristic_v1_tuning_baselines", baselines + ), # test-quality-ok: baselines are a startup-loaded proxy global with no injection seam patch( # test-quality-ok: prior auth check needs a live DB; only the tuning quota is under test "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None), @@ -6460,9 +6713,15 @@ class TestStrategyRouterWriteValidation: fake = self._FakeDb([]) with ( - patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: the guard reads the proxy license singleton with no injection seam - patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: the guard reads the proxy router global with no injection seam - patch("litellm.proxy.proxy_server.heuristic_v1_tuning_baselines", None), # test-quality-ok: baselines are a startup-loaded proxy global with no injection seam + patch( + "litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1 + ), # test-quality-ok: the guard reads the proxy license singleton with no injection seam + patch( + "litellm.proxy.proxy_server.llm_router", None + ), # test-quality-ok: the guard reads the proxy router global with no injection seam + patch( + "litellm.proxy.proxy_server.heuristic_v1_tuning_baselines", None + ), # test-quality-ok: baselines are a startup-loaded proxy global with no injection seam ): async with _auto_router_capability_slot( fake, @@ -6520,7 +6779,9 @@ class TestStrategyRouterWriteValidation: @pytest.mark.asyncio @pytest.mark.parametrize("config", [_V2, _CAPABILITY, _FUSE]) - async def test_add_new_model_refuses_a_second_gated_classifier_router_before_the_db_write(self, config: Mapping[str, object]) -> None: + async def test_add_new_model_refuses_a_second_gated_classifier_router_before_the_db_write( + self, config: Mapping[str, object] + ) -> None: from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.model_management_endpoints import ( add_new_model, @@ -6530,11 +6791,21 @@ class TestStrategyRouterWriteValidation: fake = self._FakeDb(["auto_router/complexity_router"]) with ( - patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.prisma_client", fake + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.llm_router", None + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.premium_user", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1 + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None), @@ -6548,7 +6819,9 @@ class TestStrategyRouterWriteValidation: await add_new_model( model_params=Deployment( model_name="second-v2", - litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=config), + litellm_params=LiteLLM_Params( + model="auto_router/complexity_router", complexity_router_config=config + ), ), user_api_key_dict=admin, ) @@ -6573,10 +6846,18 @@ class TestStrategyRouterWriteValidation: ) fake = self._FakeDb([]) with ( - patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy globals with no injection seam - patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: config-based lookup must be absent to drive the stored-row branch - patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reaches its DB-write branch only with this process setting - patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: authorization branch reads the proxy-wide premium flag + patch( + "litellm.proxy.proxy_server.prisma_client", fake + ), # test-quality-ok: endpoint reads proxy globals with no injection seam + patch( + "litellm.proxy.proxy_server.llm_router", None + ), # test-quality-ok: config-based lookup must be absent to drive the stored-row branch + patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), # test-quality-ok: endpoint reaches its DB-write branch only with this process setting + patch( + "litellm.proxy.proxy_server.premium_user", True + ), # test-quality-ok: authorization branch reads the proxy-wide premium flag patch( # test-quality-ok: inject stored regular row without a database "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", new=AsyncMock(return_value=regular), @@ -6620,10 +6901,18 @@ class TestStrategyRouterWriteValidation: existing_row.litellm_params = regular.litellm_params.model_dump() fake = self._FakeDb([], existing_row=existing_row) with ( - patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy globals with no injection seam - patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: config-based lookup must be absent to drive the stored-row branch - patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reaches its DB-write branch only with this process setting - patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: authorization branch reads the proxy-wide premium flag + patch( + "litellm.proxy.proxy_server.prisma_client", fake + ), # test-quality-ok: endpoint reads proxy globals with no injection seam + patch( + "litellm.proxy.proxy_server.llm_router", None + ), # test-quality-ok: config-based lookup must be absent to drive the stored-row branch + patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), # test-quality-ok: endpoint reaches its DB-write branch only with this process setting + patch( + "litellm.proxy.proxy_server.premium_user", True + ), # test-quality-ok: authorization branch reads the proxy-wide premium flag patch( # test-quality-ok: endpoint must reject before database authorization needs a live store "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None), @@ -6646,7 +6935,9 @@ class TestStrategyRouterWriteValidation: @pytest.mark.asyncio @pytest.mark.parametrize("config", [_V2, _CAPABILITY, _FUSE]) - async def test_patch_model_refuses_switching_another_router_to_gated_classifier(self, config: Mapping[str, object]) -> None: + async def test_patch_model_refuses_switching_another_router_to_gated_classifier( + self, config: Mapping[str, object] + ) -> None: """patch_model relays HTTPException as-is, so the license refusal reaches the client as a plain 403.""" from fastapi import HTTPException @@ -6660,11 +6951,21 @@ class TestStrategyRouterWriteValidation: fake = self._FakeDb(["auto_router/complexity_router"]) with ( - patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.prisma_client", fake + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.llm_router", None + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.premium_user", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1 + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch( # test-quality-ok: the write must be refused before this DB step runs "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", new=AsyncMock(return_value=self._db_complexity_router(model_id)), @@ -6690,7 +6991,9 @@ class TestStrategyRouterWriteValidation: @pytest.mark.asyncio @pytest.mark.parametrize("config", [_V2, _CAPABILITY, _FUSE]) - async def test_update_model_refuses_switching_another_router_to_gated_classifier(self, config: Mapping[str, object]) -> None: + async def test_update_model_refuses_switching_another_router_to_gated_classifier( + self, config: Mapping[str, object] + ) -> None: from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.model_management_endpoints import ( update_model, @@ -6713,11 +7016,21 @@ class TestStrategyRouterWriteValidation: fake = self._FakeDb(["auto_router/complexity_router"], existing_row=existing_row) with ( - patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam - patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.prisma_client", fake + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.llm_router", None + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server.premium_user", True + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( + "litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1 + ), # test-quality-ok: endpoint reads proxy server globals with no injection seam patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None), @@ -6807,13 +7120,17 @@ class TestAutoRouterClassifierDefaultPrompt: from litellm.router_strategy.complexity_router import ClassificationRubric, classification_system_prompt for preset in ClassificationRubric: - response = await get_auto_router_classifier_default_prompt(context_window_size=5, classification_rubric=preset) + response = await get_auto_router_classifier_default_prompt( + context_window_size=5, classification_rubric=preset + ) assert response.system_prompt == classification_system_prompt(5, classification_rubric=preset) agentic = await get_auto_router_classifier_default_prompt( context_window_size=5, classification_rubric=ClassificationRubric.AGENTIC ) - chat = await get_auto_router_classifier_default_prompt(context_window_size=5, classification_rubric=ClassificationRubric.CHAT) + chat = await get_auto_router_classifier_default_prompt( + context_window_size=5, classification_rubric=ClassificationRubric.CHAT + ) unset = await get_auto_router_classifier_default_prompt(context_window_size=5) assert "Calibration on engineering tasks" in agentic.system_prompt assert "Calibration on engineering tasks" not in chat.system_prompt @@ -7147,9 +7464,7 @@ class TestEnforceRpmTpmOnModelAdd: class TestBlockModelResponseSerialization: - @pytest.mark.parametrize( - ("route", "blocked"), [("/model/block", True), ("/model/unblock", False)] - ) + @pytest.mark.parametrize(("route", "blocked"), [("/model/block", True), ("/model/unblock", False)]) def test_block_routes_serialize_prisma_row_to_200(self, route, blocked): from datetime import datetime, timezone @@ -7180,13 +7495,19 @@ class TestBlockModelResponseSerialization: app.dependency_overrides[ps.user_api_key_auth] = lambda: admin try: with ( - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: proxy_server module global is the endpoint's only injection point - patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch( + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), # test-quality-ok: proxy_server module global is the endpoint's only injection point patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point "litellm.proxy.proxy_server.llm_router", MagicMock(**{"get_model_ids.return_value": ["m-block-1"]}), ), - patch("litellm.proxy.proxy_server.redis_usage_cache", None), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch( + "litellm.proxy.proxy_server.redis_usage_cache", None + ), # test-quality-ok: proxy_server module global is the endpoint's only injection point patch( # test-quality-ok: stubs the cache write so the test observes only response serialization "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), @@ -7515,16 +7836,36 @@ class TestTeamMemberAutoRouterWrites: @contextlib.contextmanager def _environment(self, database: MagicMock, row: LiteLLM_ProxyModelTable) -> Iterator[None]: with ( - patch("litellm.proxy.proxy_server.prisma_client", database), # test-quality-ok: [TQ008] endpoint storage singleton injection - patch("litellm.proxy.proxy_server.llm_router", self._catalog()), # test-quality-ok: [TQ008] inject real destination model catalog - patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] endpoint storage mode singleton - patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] inject licensed process state - patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", return_value=None), # test-quality-ok: [TQ008] inject unlimited license result - patch("litellm.proxy.management_endpoints.model_management_endpoints.publish_config_change", new=AsyncMock()), # test-quality-ok: [TQ008] pubsub I/O boundary - patch("litellm.proxy.management_endpoints.model_management_endpoints.create_object_audit_log", new=AsyncMock()), # test-quality-ok: [TQ008] audit database I/O boundary - patch("litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", new=AsyncMock(return_value=ReconcileOutcome( # test-quality-ok: [TQ008] model reload I/O boundary - still_desired=frozenset((row.model_id, "allowed-id")), live_after=frozenset((row.model_id, "allowed-id")) - ))), + patch( + "litellm.proxy.proxy_server.prisma_client", database + ), # test-quality-ok: [TQ008] endpoint storage singleton injection + patch( + "litellm.proxy.proxy_server.llm_router", self._catalog() + ), # test-quality-ok: [TQ008] inject real destination model catalog + patch( + "litellm.proxy.proxy_server.store_model_in_db", True + ), # test-quality-ok: [TQ008] endpoint storage mode singleton + patch( + "litellm.proxy.proxy_server.premium_user", True + ), # test-quality-ok: [TQ008] inject licensed process state + patch( + "litellm.proxy.proxy_server._license_check.auto_router_capability_limit", return_value=None + ), # test-quality-ok: [TQ008] inject unlimited license result + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.publish_config_change", new=AsyncMock() + ), # test-quality-ok: [TQ008] pubsub I/O boundary + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.create_object_audit_log", new=AsyncMock() + ), # test-quality-ok: [TQ008] audit database I/O boundary + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock( + return_value=ReconcileOutcome( # test-quality-ok: [TQ008] model reload I/O boundary + still_desired=frozenset((row.model_id, "allowed-id")), + live_after=frozenset((row.model_id, "allowed-id")), + ) + ), + ), ): yield @@ -7581,15 +7922,23 @@ class TestTeamMemberAutoRouterWrites: litellm_proxymodeltable=table, tx=MagicMock(return_value=context), ) - return MagicMock(db=db, writer_db=MagicMock(litellm_uisettings=MagicMock(find_unique=AsyncMock(return_value=None))), transaction=transaction) + return MagicMock( + db=db, + writer_db=MagicMock(litellm_uisettings=MagicMock(find_unique=AsyncMock(return_value=None))), + transaction=transaction, + ) @staticmethod def _catalog() -> Router: - return Router(model_list=[{ - "model_name": "allowed", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"}, - "model_info": {"id": "allowed-id"}, - }]) + return Router( + model_list=[ + { + "model_name": "allowed", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"}, + "model_info": {"id": "allowed-id"}, + } + ] + ) @staticmethod def _classifier_config(classifier: Mapping[str, object], legacy: bool) -> Mapping[str, object]: @@ -7603,7 +7952,12 @@ class TestTeamMemberAutoRouterWrites: @pytest.mark.parametrize("team_id", [None, "member-team"]) @pytest.mark.parametrize( "legacy,provider,model", - [(True, "typesafe", "jev-latest"), (False, "jev", "jev-latest"), (True, "laya", "english"), (False, "laya", "english")], + [ + (True, "typesafe", "jev-latest"), + (False, "jev", "jev-latest"), + (True, "laya", "english"), + (False, "laya", "english"), + ], ) async def test_classifier_create_stores_only_canonical_configuration( self, team_id: str | None, legacy: bool, provider: str, model: str @@ -7613,8 +7967,10 @@ class TestTeamMemberAutoRouterWrites: row: Final = self._row() database: Final = self._database(self._team(), row) classifier: Final = { - "provider": provider, "model": model, - "api_base": "https://decision.test", "api_key": "stored-secret", + "provider": provider, + "model": model, + "api_base": "https://decision.test", + "api_key": "stored-secret", } deployment: Final = Deployment( model_name="new-classifier-router", @@ -7627,10 +7983,17 @@ class TestTeamMemberAutoRouterWrites: actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) with ( self._environment(database, row), - patch("litellm.proxy.proxy_server.proxy_config.add_deployment", new=AsyncMock(return_value=ReconcileOutcome( # test-quality-ok: [TQ008] model reload I/O boundary - still_desired=frozenset((row.model_id,)), live_after=frozenset((row.model_id,)) - ))), - patch("litellm.proxy.management_endpoints.model_management_endpoints.append_team_models", new=AsyncMock()), # test-quality-ok: [TQ008] team allowlist persistence boundary + patch( + "litellm.proxy.proxy_server.proxy_config.add_deployment", + new=AsyncMock( + return_value=ReconcileOutcome( # test-quality-ok: [TQ008] model reload I/O boundary + still_desired=frozenset((row.model_id,)), live_after=frozenset((row.model_id,)) + ) + ), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.append_team_models", new=AsyncMock() + ), # test-quality-ok: [TQ008] team allowlist persistence boundary ): await add_new_model(deployment, actor) written: Final = database.db.litellm_proxymodeltable.create.await_args.kwargs["data"] @@ -7657,13 +8020,16 @@ class TestTeamMemberAutoRouterWrites: } actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) request: Final = updateDeployment( - litellm_params=updateLiteLLMParams(complexity_router_config=config), model_info=ModelInfo(id=row.model_id), + litellm_params=updateLiteLLMParams(complexity_router_config=config), + model_info=ModelInfo(id=row.model_id), ) operation: Final = ( add_new_model( Deployment( model_name="ambiguous-classifier-router", - litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=config), + litellm_params=LiteLLM_Params( + model="auto_router/complexity_router", complexity_router_config=config + ), model_info=ModelInfo(id=row.model_id), ), actor, @@ -7712,7 +8078,9 @@ class TestTeamMemberAutoRouterWrites: @pytest.mark.asyncio @pytest.mark.parametrize("endpoint", ["patch", "legacy"]) - @pytest.mark.parametrize("stored_legacy,supplied_legacy", [(True, True), (True, False), (False, True), (False, False)]) + @pytest.mark.parametrize( + "stored_legacy,supplied_legacy", [(True, True), (True, False), (False, True), (False, False)] + ) @pytest.mark.parametrize("change", ["save", "rotate", "move", "move-without-key", "reset", "heuristic"]) async def test_jev_dashboard_save_preserves_server_transport( self, endpoint: str, change: str, stored_legacy: bool, supplied_legacy: bool @@ -7776,61 +8144,106 @@ class TestTeamMemberAutoRouterWrites: @pytest.mark.asyncio @pytest.mark.parametrize("endpoint", ["patch", "legacy"]) - @pytest.mark.parametrize("stored_legacy,supplied_legacy", [(True, True), (True, False), (False, True), (False, False)]) + @pytest.mark.parametrize( + "stored_legacy,supplied_legacy", [(True, True), (True, False), (False, True), (False, False)] + ) @pytest.mark.parametrize( "stored_provider,stored_base,supplied,expected_transport", [ - ("bespoke", "https://decision.test", {"provider": "bespoke", "model": "nimble-latest"}, {"api_base": "https://decision.test", "api_key": "stored-secret"}), - ("bespoke", "https://decision.test", {"provider": "bespoke", "model": "nimble-latest", "api_base": "https://new.test"}, {}), + ( + "bespoke", + "https://decision.test", + {"provider": "bespoke", "model": "nimble-latest"}, + {"api_base": "https://decision.test", "api_key": "stored-secret"}, + ), + ( + "bespoke", + "https://decision.test", + {"provider": "bespoke", "model": "nimble-latest", "api_base": "https://new.test"}, + {}, + ), ("bespoke", "https://decision.test", {"provider": "laya", "model": "english"}, {}), ("laya", "https://decision.test", {"provider": "bespoke", "model": "nimble-latest"}, {}), - ("laya", "https://decision.test", {"provider": "laya", "model": "english"}, {"api_base": "https://decision.test", "api_key": "stored-secret"}), - ("laya", "https://decision.test", {"provider": "laya", "model": "english", "api_base": "https://decision.test"}, {"api_base": "https://decision.test", "api_key": "stored-secret"}), + ( + "laya", + "https://decision.test", + {"provider": "laya", "model": "english"}, + {"api_base": "https://decision.test", "api_key": "stored-secret"}, + ), + ( + "laya", + "https://decision.test", + {"provider": "laya", "model": "english", "api_base": "https://decision.test"}, + {"api_base": "https://decision.test", "api_key": "stored-secret"}, + ), ( "laya", "https://decision.test", {"provider": "laya", "model": "english", "api_key": None}, {"api_base": "https://decision.test"}, ), - ("laya", "https://decision.test", {"provider": "laya", "model": "english", "api_base": "https://new.test"}, {}), + ( + "laya", + "https://decision.test", + {"provider": "laya", "model": "english", "api_base": "https://new.test"}, + {}, + ), ("laya", "https://decision.test", {"provider": "laya", "model": "english", "api_base": None}, {}), ("laya", None, {"provider": "laya", "model": "english", "api_base": None}, {}), ("laya", "https://decision.test", {"provider": "typesafe", "model": "jev-latest"}, {}), ( - "laya", "https://decision.test", {"model": "english", "timeout_ms": 8100}, + "laya", + "https://decision.test", + {"model": "english", "timeout_ms": 8100}, {"provider": "laya", "api_base": "https://decision.test", "api_key": "stored-secret"}, ), ("typesafe", "https://decision.test", {"provider": "laya", "model": "english"}, {}), ( - "typesafe", "https://decision.test", {"provider": "jev", "model": "jev-latest"}, + "typesafe", + "https://decision.test", + {"provider": "jev", "model": "jev-latest"}, {"api_base": "https://decision.test", "api_key": "stored-secret"}, ), ( - "jev", "https://decision.test", {"provider": "typesafe", "model": "jev-latest"}, + "jev", + "https://decision.test", + {"provider": "typesafe", "model": "jev-latest"}, {"api_base": "https://decision.test", "api_key": "stored-secret"}, ), ], ) async def test_decision_provider_changes_cannot_reuse_a_stored_key( - self, endpoint: str, stored_provider: str, stored_base: str | None, - supplied: Mapping[str, object], expected_transport: Mapping[str, object], - stored_legacy: bool, supplied_legacy: bool, + self, + endpoint: str, + stored_provider: str, + stored_base: str | None, + supplied: Mapping[str, object], + expected_transport: Mapping[str, object], + stored_legacy: bool, + supplied_legacy: bool, ) -> None: original: Final = self._row() - row: Final = original.model_copy(update={"litellm_params": { - "model": "auto_router/complexity_router", - "complexity_router_config": self._classifier_config( - { - "provider": stored_provider, "model": {"laya": "english", "bespoke": "nimble-latest"}.get(stored_provider, "jev-latest"), - "api_base": stored_base, "api_key": "stored-secret", - }, - stored_legacy, - ), - }}) + row: Final = original.model_copy( + update={ + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": self._classifier_config( + { + "provider": stored_provider, + "model": {"laya": "english", "bespoke": "nimble-latest"}.get(stored_provider, "jev-latest"), + "api_base": stored_base, + "api_key": "stored-secret", + }, + stored_legacy, + ), + } + } + ) database: Final = self._database(self._team(), row) config: Final = self._classifier_config(supplied, supplied_legacy) request: Final = updateDeployment( - litellm_params=updateLiteLLMParams(complexity_router_config=config), model_info=ModelInfo(id=row.model_id), + litellm_params=updateLiteLLMParams(complexity_router_config=config), + model_info=ModelInfo(id=row.model_id), ) actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) with self._environment(database, row): @@ -7842,7 +8255,8 @@ class TestTeamMemberAutoRouterWrites: "classifier_type": "oss_classifier", "tiers": {"SIMPLE": "allowed"}, "opensource_classifier_config": { - **expected_transport, **supplied, + **expected_transport, + **supplied, "provider": "jev" if expected_provider == "typesafe" else expected_provider, }, } @@ -7852,9 +8266,13 @@ class TestTeamMemberAutoRouterWrites: @pytest.mark.parametrize( "string_params,reset_field,config_shape", [ - (False, None, "full"), (True, None, "full"), (False, "api_key", "full"), - (False, "api_base", "full"), (False, None, "omit-provider"), - (False, None, "omit-config"), (False, None, "null-config"), + (False, None, "full"), + (True, None, "full"), + (False, "api_key", "full"), + (False, "api_base", "full"), + (False, None, "omit-provider"), + (False, None, "omit-config"), + (False, None, "null-config"), ], ) async def test_member_save_protects_stored_classifier_connection( @@ -7862,14 +8280,18 @@ class TestTeamMemberAutoRouterWrites: ) -> None: original: Final = self._row() config: Final = { - "classifier_type": "jev", "tiers": {"SIMPLE": "allowed"}, + "classifier_type": "jev", + "tiers": {"SIMPLE": "allowed"}, "jev_classifier_config": {"provider": "laya", "model": "english"}, } secret_params: Final = { "model": "auto_router/complexity_router", "complexity_router_config": { - **config, "jev_classifier_config": { - **config["jev_classifier_config"], "api_key": "retained-laya-secret", "api_base": "https://laya.test", + **config, + "jev_classifier_config": { + **config["jev_classifier_config"], + "api_key": "retained-laya-secret", + "api_base": "https://laya.test", }, }, } @@ -7880,9 +8302,11 @@ class TestTeamMemberAutoRouterWrites: update={"litellm_params": json.dumps(secret_params) if string_params else secret_params} ) supplied_config: Final = { - **config, "jev_classifier_config": { + **config, + "jev_classifier_config": { **{ - key: value for key, value in config["jev_classifier_config"].items() + key: value + for key, value in config["jev_classifier_config"].items() if key != "provider" or config_shape != "omit-provider" }, **({reset_field: None} if reset_field is not None else {}), @@ -7900,14 +8324,19 @@ class TestTeamMemberAutoRouterWrites: model_info=ModelInfo(id=row.model_id, team_id="member-team"), ) actor: Final = UserAPIKeyAuth( - user_id="owner", user_role=LitellmUserRoles.INTERNAL_USER, models=["allowed", "laya/english"], config={"timeout": 60}, + user_id="owner", + user_role=LitellmUserRoles.INTERNAL_USER, + models=["allowed", "laya/english"], + config={"timeout": 60}, ) with self._environment(database, row): if reset_field is not None: expected_error: Final = HTTPException if endpoint == "patch" else ProxyException with pytest.raises(expected_error, match="Team members cannot change classifier connections") as denied: await ( - patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor) + patch_model(row.model_id, request, actor) + if endpoint == "patch" + else update_model(request, actor) ) assert ( denied.value.status_code if isinstance(denied.value, HTTPException) else int(denied.value.code) @@ -7915,7 +8344,9 @@ class TestTeamMemberAutoRouterWrites: database.transaction.litellm_proxymodeltable.update.assert_not_awaited() assert row.litellm_params == secret_params return - response: Final = await (patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor)) + response: Final = await ( + patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor) + ) written: Final = database.transaction.litellm_proxymodeltable.update.await_args.kwargs["data"] saved_config: Final = json.loads(written["litellm_params"])["complexity_router_config"] untouched: Final = config_shape in ("omit-config", "null-config") @@ -7925,17 +8356,27 @@ class TestTeamMemberAutoRouterWrites: if untouched: from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper - assert decrypt_value_helper( - json.loads(written["litellm_params"])["complexity_router_default_model"], - key="complexity_router_default_model", return_original_value=True, - ) == "allowed" + assert ( + decrypt_value_helper( + json.loads(written["litellm_params"])["complexity_router_default_model"], + key="complexity_router_default_model", + return_original_value=True, + ) + == "allowed" + ) response_payload: Final = jsonable_encoder(response) assert "retained-laya-secret" not in json.dumps(response_payload) - response_params: Final = json.loads(response_payload["litellm_params"]) if string_params else response_payload["litellm_params"] + response_params: Final = ( + json.loads(response_payload["litellm_params"]) if string_params else response_payload["litellm_params"] + ) assert response_params == { - **secret_params, "complexity_router_config": { - **config, "jev_classifier_config": { - **config["jev_classifier_config"], "api_key": "REDACTED", "api_base": "https://laya.test", + **secret_params, + "complexity_router_config": { + **config, + "jev_classifier_config": { + **config["jev_classifier_config"], + "api_key": "REDACTED", + "api_base": "https://laya.test", }, }, } @@ -7944,9 +8385,7 @@ class TestTeamMemberAutoRouterWrites: @pytest.mark.asyncio @pytest.mark.parametrize("endpoint", ["patch", "legacy"]) @pytest.mark.parametrize("access", ["owner", "peer", "limited-key"]) - async def test_both_update_entries_enforce_creator_and_stamp_member_scope( - self, endpoint: str, access: str - ) -> None: + async def test_both_update_entries_enforce_creator_and_stamp_member_scope(self, endpoint: str, access: str) -> None: from fastapi import HTTPException from litellm.proxy._types import ProxyException @@ -7955,15 +8394,21 @@ class TestTeamMemberAutoRouterWrites: row: Final = self._row() database: Final = self._database(self._team(), row) request: Final = updateDeployment( - litellm_params=updateLiteLLMParams(complexity_router_config={"tiers": {"SIMPLE": "allowed"}, "session_affinity": True}), + litellm_params=updateLiteLLMParams( + complexity_router_config={"tiers": {"SIMPLE": "allowed"}, "session_affinity": True} + ), model_info=ModelInfo(id=row.model_id, team_id="member-team"), ) actor: Final = UserAPIKeyAuth( - user_id="peer" if access == "peer" else "owner", user_role=LitellmUserRoles.INTERNAL_USER, - models=["personal-router"] if access == "limited-key" else ["allowed"], config={"timeout": 60}, + user_id="peer" if access == "peer" else "owner", + user_role=LitellmUserRoles.INTERNAL_USER, + models=["personal-router"] if access == "limited-key" else ["allowed"], + config={"timeout": 60}, ) with self._environment(database, row): - operation: Final = patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor) + operation: Final = ( + patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor) + ) if access != "owner": with pytest.raises((HTTPException, ProxyException)): await operation @@ -7985,35 +8430,52 @@ class TestTeamMemberAutoRouterWrites: from fastapi import HTTPException from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot - from litellm.proxy.management_helpers.auto_router_permissions import MemberAutoRouterWrite, validate_member_auto_router_config + from litellm.proxy.management_helpers.auto_router_permissions import ( + MemberAutoRouterWrite, + validate_member_auto_router_config, + ) row: Final = self._row() database: Final = self._database(self._team(), row) if changed_state == "revoked": database.transaction.litellm_teamtable.find_unique.return_value = self._team(enabled=False) elif changed_state == "moved": - database.transaction.litellm_proxymodeltable.find_unique.return_value = row.model_copy(update={"model_info": {"team_id": "other-team"}}) + database.transaction.litellm_proxymodeltable.find_unique.return_value = row.model_copy( + update={"model_info": {"team_id": "other-team"}} + ) elif changed_state == "creator": - database.transaction.litellm_proxymodeltable.find_unique.return_value = row.model_copy(update={"created_by": "peer"}) + database.transaction.litellm_proxymodeltable.find_unique.return_value = row.model_copy( + update={"created_by": "peer"} + ) elif changed_state == "collision": database.transaction.litellm_proxymodeltable.find_many.return_value = [row] config: Final = validate_member_auto_router_config({"tiers": {"SIMPLE": "allowed"}}) grant: Final = MemberAutoRouterWrite( actor=UserAPIKeyAuth(user_id="owner", user_role=LitellmUserRoles.INTERNAL_USER, models=["allowed"]), - team_id="member-team", model_id=None if changed_state in ("collision", "global-alias") else row.model_id, - public_name="personal-router", updated_at=None, config=config, default_model="allowed", + team_id="member-team", + model_id=None if changed_state in ("collision", "global-alias") else row.model_id, + public_name="personal-router", + updated_at=None, + config=config, + default_model="allowed", ) with ( self._environment(database, row), - patch("litellm.model_alias_map", {"personal-router": "allowed"} if changed_state == "global-alias" else {}), # test-quality-ok: [TQ008] inject alias namespace for collision behavior + patch( + "litellm.model_alias_map", {"personal-router": "allowed"} if changed_state == "global-alias" else {} + ), # test-quality-ok: [TQ008] inject alias namespace for collision behavior ): if changed_state != "allowed": with pytest.raises(HTTPException) as denied: - async with _auto_router_capability_slot(database, effective_params={}, model_id=grant.model_id, member_write=grant): + async with _auto_router_capability_slot( + database, effective_params={}, model_id=grant.model_id, member_write=grant + ): pytest.fail("An invalidated grant reached the database writer") assert denied.value.status_code == (409 if changed_state in ("collision", "global-alias") else 403) return - async with _auto_router_capability_slot(database, effective_params={}, model_id=grant.model_id, member_write=grant) as table: + async with _auto_router_capability_slot( + database, effective_params={}, model_id=grant.model_id, member_write=grant + ) as table: await table.update(where={"model_id": row.model_id}, data={"updated_by": "owner"}) assert database.transaction.query_raw.await_count == 2 database.transaction.litellm_proxymodeltable.update.assert_awaited_once() @@ -8027,20 +8489,32 @@ class TestTeamMemberAutoRouterWrites: row: Final = self._row() database: Final = self._database(self._team(enabled=access != "opt-out"), row) actor: Final = UserAPIKeyAuth( - user_id="owner", user_role=LitellmUserRoles.INTERNAL_USER, - models=["personal-router"] if access == "limited-key" else ["allowed"], config={"timeout": 60}, + user_id="owner", + user_role=LitellmUserRoles.INTERNAL_USER, + models=["personal-router"] if access == "limited-key" else ["allowed"], + config={"timeout": 60}, ) deployment: Final = Deployment( model_name="new-personal-router", - litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config={"tiers": {"SIMPLE": "allowed"}}), + litellm_params=LiteLLM_Params( + model="auto_router/complexity_router", complexity_router_config={"tiers": {"SIMPLE": "allowed"}} + ), model_info=ModelInfo(id=row.model_id, team_id="member-team"), ) with ( self._environment(database, row), - patch("litellm.proxy.proxy_server.proxy_config.add_deployment", new=AsyncMock(return_value=ReconcileOutcome( # test-quality-ok: [TQ008] model reload I/O boundary - still_desired=frozenset((row.model_id, "allowed-id")), live_after=frozenset((row.model_id, "allowed-id")) - ))), - patch("litellm.proxy.management_endpoints.model_management_endpoints.append_team_models", new=AsyncMock()) as appended, # test-quality-ok: [TQ008] persistence boundary; the appended scope is asserted + patch( + "litellm.proxy.proxy_server.proxy_config.add_deployment", + new=AsyncMock( + return_value=ReconcileOutcome( # test-quality-ok: [TQ008] model reload I/O boundary + still_desired=frozenset((row.model_id, "allowed-id")), + live_after=frozenset((row.model_id, "allowed-id")), + ) + ), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.append_team_models", new=AsyncMock() + ) as appended, # test-quality-ok: [TQ008] persistence boundary; the appended scope is asserted ): if access != "allowed": with pytest.raises(ProxyException) as denied: @@ -8072,10 +8546,18 @@ class TestModelManagementActorEdges: 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 + 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) @@ -8089,9 +8571,7 @@ class TestModelManagementActorEdges: 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 - ) + 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", @@ -8099,10 +8579,18 @@ class TestModelManagementActorEdges: 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 + 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) @@ -8124,10 +8612,18 @@ class TestModelManagementActorEdges: 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 + 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) @@ -8158,10 +8654,18 @@ class TestModelManagementActorEdges: 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( + "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, @@ -8207,10 +8711,18 @@ class TestModelManagementActorEdges: 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( + "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, @@ -8250,10 +8762,18 @@ class TestModelManagementActorEdges: ) 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 + 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( @@ -8285,10 +8805,18 @@ class TestModelManagementActorEdges: prisma: Final = MagicMock() prisma.writer_db.litellm_uisettings.find_unique = AsyncMock(return_value=None) 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 + 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( @@ -8329,10 +8857,18 @@ class TestModelManagementActorEdges: 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( + "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, @@ -8374,10 +8910,18 @@ class TestModelManagementActorEdges: ) 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( + "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( @@ -8871,11 +9415,7 @@ class TestNonAdminCannotPersistWifFieldsOnModel: "api_key": "stored-api-key", "client_id": "stored-client-id", "client_secret": "stored-client-secret", - **( - {"token_exchange_endpoint": "https://identity.example.com/token"} - if server_owned_oauth - else {} - ), + **({"token_exchange_endpoint": "https://identity.example.com/token"} if server_owned_oauth else {}), } existing_row.model_dump.return_value = { "model_name": "copilot", @@ -8934,9 +9474,7 @@ class TestNonAdminCannotPersistWifFieldsOnModel: with pytest.raises(ProxyException) as exc_info: await patch_model( model_id="oauth-client-patch-1", - patch_data=updateDeployment( - litellm_params=updateLiteLLMParams.model_validate({field: value}) - ), + patch_data=updateDeployment(litellm_params=updateLiteLLMParams.model_validate({field: value})), user_api_key_dict=team_admin, ) @@ -9755,9 +10293,7 @@ class TestNonAdminCannotSetPerUserOauthOnModel: mock_prisma.writer_db.litellm_uisettings.find_unique = AsyncMock(return_value=None) per_user_credential_row = MagicMock() per_user_credential_row.credential_values = {"github_copilot_auth_type": "per_user_oauth"} - mock_prisma.db.litellm_credentialstable.find_unique = AsyncMock( - return_value=per_user_credential_row - ) + mock_prisma.db.litellm_credentialstable.find_unique = AsyncMock(return_value=per_user_credential_row) with ( patch( # test-quality-ok: the proxy wiring under test is what this patches @@ -9798,9 +10334,7 @@ class TestNonAdminCannotSetPerUserOauthOnModel: mock_prisma.db.litellm_proxymodeltable.create = AsyncMock(return_value=created_row) per_user_credential_row = MagicMock() per_user_credential_row.credential_values = {"github_copilot_auth_type": "per_user_oauth"} - mock_prisma.db.litellm_credentialstable.find_unique = AsyncMock( - return_value=per_user_credential_row - ) + mock_prisma.db.litellm_credentialstable.find_unique = AsyncMock(return_value=per_user_credential_row) with ( patch( # test-quality-ok: the proxy wiring under test is what this patches diff --git a/tests/unit/proxy/management_endpoints/test_r2_9_api_base_team_gate.py b/tests/unit/proxy/management_endpoints/test_r2_9_api_base_team_gate.py index 8e114b75091..524677ce145 100644 --- a/tests/unit/proxy/management_endpoints/test_r2_9_api_base_team_gate.py +++ b/tests/unit/proxy/management_endpoints/test_r2_9_api_base_team_gate.py @@ -44,9 +44,7 @@ async def test_team_admin_bring_your_own_key_move_allowed(): from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth stored = GenericLiteLLMParams(api_key="server-owned-key") - incoming = GenericLiteLLMParams( - api_base="https://my-team-proxy.example", api_key="my-team-key" - ) + incoming = GenericLiteLLMParams(api_base="https://my-team-proxy.example", api_key="my-team-key") model_params = MagicMock() model_params.litellm_params = stored @@ -80,9 +78,7 @@ async def test_team_admin_wif_association_still_guarded(): model_params.litellm_params = stored with ( - patch.object( - mme, "effective_server_owned_wif_fields", AsyncMock(return_value=("oidc_token_url",)) - ), + patch.object(mme, "effective_server_owned_wif_fields", AsyncMock(return_value=("oidc_token_url",))), ): try: await ModelManagementAuthChecks._reject_non_admin_wif_write( @@ -96,8 +92,6 @@ async def test_team_admin_wif_association_still_guarded(): except Exception as guard_error: # On this stack both the static-key harvest gate and the WIF # rule can fire for the same shape; accept either interdiction. - assert ("api_key stays server-owned" in str(guard_error)) or ( - "workload identity" in str(guard_error) - ) + assert ("api_key stays server-owned" in str(guard_error)) or ("workload identity" in str(guard_error)) return raise AssertionError("expected a server-owned-key interdiction to fire")