diff --git a/tests/pass_through_unit_tests/test_claude_code_marketplace.py b/tests/pass_through_unit_tests/test_claude_code_marketplace.py index 288b151eadf..27de520d8fe 100644 --- a/tests/pass_through_unit_tests/test_claude_code_marketplace.py +++ b/tests/pass_through_unit_tests/test_claude_code_marketplace.py @@ -219,18 +219,12 @@ async def test_register_plugin(mock_prisma_client): assert response["plugin"]["enabled"] is True # Verify the plugin was stored in the mock - stored_plugin = ( - await mock_prisma_client.db.litellm_claudecodeplugintable.find_unique( - where={"name": plugin_name} - ) - ) + stored_plugin = await mock_prisma_client.db.litellm_claudecodeplugintable.find_unique(where={"name": plugin_name}) assert stored_plugin is not None assert stored_plugin.name == plugin_name # Cleanup - delete the plugin - await mock_prisma_client.db.litellm_claudecodeplugintable.delete( - where={"name": plugin_name} - ) + await mock_prisma_client.db.litellm_claudecodeplugintable.delete(where={"name": plugin_name}) @pytest.mark.asyncio @@ -281,9 +275,7 @@ async def test_get_marketplace(mock_prisma_client): assert our_plugin["version"] == "2.0.0" # Cleanup - await mock_prisma_client.db.litellm_claudecodeplugintable.delete( - where={"name": plugin_name} - ) + await mock_prisma_client.db.litellm_claudecodeplugintable.delete(where={"name": plugin_name}) @pytest.mark.asyncio @@ -323,9 +315,7 @@ async def test_get_marketplace_no_key_unaffected_by_imported_disabled_skills( "name": imported_disabled_name, "version": "1.0.0", "description": "Imported but not yet enabled", - "manifest_json": json.dumps( - {"source": {"source": "github", "repo": "org/private-skill"}} - ), + "manifest_json": json.dumps({"source": {"source": "github", "repo": "org/private-skill"}}), "enabled": False, } ) @@ -344,12 +334,8 @@ async def test_get_marketplace_no_key_unaffected_by_imported_disabled_skills( assert imported_disabled_name not in names # Cleanup - await mock_prisma_client.db.litellm_claudecodeplugintable.delete( - where={"name": public_plugin_name} - ) - await mock_prisma_client.db.litellm_claudecodeplugintable.delete( - where={"name": imported_disabled_name} - ) + await mock_prisma_client.db.litellm_claudecodeplugintable.delete(where={"name": public_plugin_name}) + await mock_prisma_client.db.litellm_claudecodeplugintable.delete(where={"name": imported_disabled_name}) @pytest.mark.asyncio @@ -371,9 +357,7 @@ async def test_get_marketplace_with_key_unlocks_allowed_imported_skill( "name": imported_disabled_name, "version": "1.0.0", "description": "Imported, granted to one key only", - "manifest_json": json.dumps( - {"source": {"source": "github", "repo": "org/scoped-skill"}} - ), + "manifest_json": json.dumps({"source": {"source": "github", "repo": "org/scoped-skill"}}), "enabled": False, } ) @@ -396,9 +380,7 @@ async def test_get_marketplace_with_key_unlocks_allowed_imported_skill( assert imported_disabled_name in names_with_key # Cleanup - await mock_prisma_client.db.litellm_claudecodeplugintable.delete( - where={"name": imported_disabled_name} - ) + await mock_prisma_client.db.litellm_claudecodeplugintable.delete(where={"name": imported_disabled_name}) @pytest.mark.asyncio @@ -430,9 +412,7 @@ async def test_get_marketplace_hides_skill_from_disabled_marketplace_even_with_g "name": skill_name, "version": "1.0.0", "description": "Skill owned by a since-disabled marketplace", - "manifest_json": json.dumps( - {"source": {"source": "github", "repo": "org/untrusted-skill"}} - ), + "manifest_json": json.dumps({"source": {"source": "github", "repo": "org/untrusted-skill"}}), "enabled": True, # cascade-disable races aside, this asserts the belt-and-suspenders check too "marketplace_id": marketplace_id, } @@ -442,9 +422,7 @@ async def test_get_marketplace_hides_skill_from_disabled_marketplace_even_with_g user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-granted", user_id="granted-user", - object_permission=LiteLLM_ObjectPermissionTable( - object_permission_id="perm-2", allowed_skills=[skill_name] - ), + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="perm-2", allowed_skills=[skill_name]), ) response = await get_marketplace(user_api_key_dict=granted_key) body = json.loads(response.body.decode()) @@ -492,14 +470,9 @@ async def test_register_plugin_git_subdir(mock_prisma_client): assert response["action"] == "created" assert response["plugin"]["name"] == plugin_name assert response["plugin"]["source"]["source"] == "git-subdir" - assert ( - response["plugin"]["source"]["url"] - == "https://github.com/test-org/monorepo.git" - ) + assert response["plugin"]["source"]["url"] == "https://github.com/test-org/monorepo.git" assert response["plugin"]["source"]["path"] == "plugins/my-plugin" assert response["plugin"]["enabled"] is True # Cleanup - await mock_prisma_client.db.litellm_claudecodeplugintable.delete( - where={"name": plugin_name} - ) + await mock_prisma_client.db.litellm_claudecodeplugintable.delete(where={"name": plugin_name}) diff --git a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py index 5c6d305bfa9..fc395a68c1b 100644 --- a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py @@ -44,9 +44,7 @@ async def test_set_object_permission(): mock_created_permission = MagicMock() mock_created_permission.object_permission_id = "test_perm_id_123" - mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock( - return_value=mock_created_permission - ) + mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock(return_value=mock_created_permission) # Test data with object_permission data_json = { @@ -62,9 +60,7 @@ async def test_set_object_permission(): } # Call the function - result = await _set_object_permission( - data_json=data_json, prisma_client=mock_prisma_client - ) + result = await _set_object_permission(data_json=data_json, prisma_client=mock_prisma_client) # Verify object_permission_id was added to result assert result["object_permission_id"] == "test_perm_id_123" @@ -105,9 +101,7 @@ async def test_set_object_permission_persists_mcp_tool_search_enabled(): mock_prisma_client = MagicMock() mock_created_permission = MagicMock() mock_created_permission.object_permission_id = "perm_id" - mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock( - return_value=mock_created_permission - ) + mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock(return_value=mock_created_permission) data_json = { "object_permission": { @@ -118,11 +112,7 @@ async def test_set_object_permission_persists_mcp_tool_search_enabled(): await _set_object_permission(data_json=data_json, prisma_client=mock_prisma_client) - created_data = ( - mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs[ - "data" - ] - ) + created_data = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"] assert created_data["mcp_tool_search_enabled"] is True @@ -203,11 +193,7 @@ def _make_team_obj( mock_team = MagicMock() mock_team.team_id = team_id - if ( - mcp_servers is not None - or mcp_access_groups is not None - or mcp_tool_permissions is not None - ): + if mcp_servers is not None or mcp_access_groups is not None or mcp_tool_permissions is not None: mock_team.object_permission = MagicMock(spec=LiteLLM_ObjectPermissionTable) mock_team.object_permission.mcp_servers = mcp_servers or [] mock_team.object_permission.mcp_access_groups = mcp_access_groups or [] @@ -278,9 +264,7 @@ async def test_validate_no_object_permission(mock_access_groups, mock_allow_all) new_callable=AsyncMock, return_value=[], ) -async def test_validate_key_servers_within_team_scope( - mock_access_groups, mock_allow_all -): +async def test_validate_key_servers_within_team_scope(mock_access_groups, mock_allow_all): """Key requests servers that are in the team's scope — should pass.""" team_obj = _make_team_obj(mcp_servers=["server-1", "server-2", "server-3"]) await validate_key_mcp_servers_against_team( @@ -303,9 +287,7 @@ async def test_validate_key_servers_within_team_scope( new_callable=AsyncMock, return_value=[], ) -async def test_validate_key_servers_outside_team_scope_raises( - mock_access_groups, mock_allow_all -): +async def test_validate_key_servers_outside_team_scope_raises(mock_access_groups, mock_allow_all): """Key requests a server that exists but is NOT in the team's scope — should raise 403.""" team_obj = _make_team_obj(mcp_servers=["server-1"]) with pytest.raises(HTTPException) as exc_info: @@ -331,9 +313,7 @@ async def test_validate_key_servers_outside_team_scope_raises( new_callable=AsyncMock, return_value=[], ) -async def test_validate_allow_all_keys_servers_always_allowed( - mock_access_groups, mock_allow_all -): +async def test_validate_allow_all_keys_servers_always_allowed(mock_access_groups, mock_allow_all): """allow_all_keys servers should be accessible even if not in team scope.""" team_obj = _make_team_obj(mcp_servers=["server-1"]) await validate_key_mcp_servers_against_team( @@ -378,9 +358,7 @@ async def test_validate_no_team_only_allow_all_keys(mock_access_groups, mock_all new_callable=AsyncMock, return_value=[], ) -async def test_validate_no_team_non_global_server_raises( - mock_access_groups, mock_allow_all -): +async def test_validate_no_team_non_global_server_raises(mock_access_groups, mock_allow_all): """Key without a team requesting an existing non-global server — should raise 403.""" with pytest.raises(HTTPException) as exc_info: await validate_key_mcp_servers_against_team( @@ -405,9 +383,7 @@ async def test_validate_no_team_non_global_server_raises( new_callable=AsyncMock, return_value=[], ) -async def test_validate_no_team_proxy_admin_can_assign_private_server( - mock_access_groups, mock_allow_all -): +async def test_validate_no_team_proxy_admin_can_assign_private_server(mock_access_groups, mock_allow_all): """Proxy admin assigning a non-global server to a teamless key — should pass (LIT-3815).""" result = await validate_key_mcp_servers_against_team( object_permission={"mcp_servers": ["private-server"]}, @@ -431,9 +407,7 @@ async def test_validate_no_team_proxy_admin_can_assign_private_server( new_callable=AsyncMock, return_value=[], ) -async def test_validate_no_team_non_admin_private_server_still_raises( - mock_access_groups, mock_allow_all -): +async def test_validate_no_team_non_admin_private_server_still_raises(mock_access_groups, mock_allow_all): """The teamless override is gated on proxy admin — a non-admin still gets 403.""" with pytest.raises(HTTPException) as exc_info: await validate_key_mcp_servers_against_team( @@ -454,9 +428,7 @@ async def test_validate_no_team_non_admin_private_server_still_raises( new_callable=AsyncMock, return_value=[], ) -async def test_validate_no_team_proxy_admin_can_assign_access_group( - mock_access_groups, mock_allow_all -): +async def test_validate_no_team_proxy_admin_can_assign_access_group(mock_access_groups, mock_allow_all): """Proxy admin assigning an access group to a teamless key — should pass (LIT-3815).""" result = await validate_key_mcp_servers_against_team( object_permission={"mcp_access_groups": ["group-1"]}, @@ -480,9 +452,7 @@ async def test_validate_no_team_proxy_admin_can_assign_access_group( new_callable=AsyncMock, return_value=[], ) -async def test_validate_proxy_admin_still_bounded_by_team_scope( - mock_access_groups, mock_allow_all -): +async def test_validate_proxy_admin_still_bounded_by_team_scope(mock_access_groups, mock_allow_all): """The override is scoped to teamless keys — an admin assigning beyond a team's scope still raises.""" team_obj = _make_team_obj(mcp_servers=["server-1"]) with pytest.raises(HTTPException) as exc_info: @@ -509,9 +479,7 @@ async def test_validate_proxy_admin_still_bounded_by_team_scope( new_callable=AsyncMock, return_value=[], ) -async def test_validate_team_no_mcp_config_blocks_all( - mock_access_groups, mock_allow_all -): +async def test_validate_team_no_mcp_config_blocks_all(mock_access_groups, mock_allow_all): """Team with no object_permission — key can't use any non-global MCP servers.""" team_obj = _make_team_obj() # No object_permission with pytest.raises(HTTPException) as exc_info: @@ -536,9 +504,7 @@ async def test_validate_team_no_mcp_config_blocks_all( new_callable=AsyncMock, return_value=[], ) -async def test_validate_tool_permissions_validated_against_team( - mock_access_groups, mock_allow_all -): +async def test_validate_tool_permissions_validated_against_team(mock_access_groups, mock_allow_all): """Server IDs in mcp_tool_permissions should also be validated when they exist.""" team_obj = _make_team_obj(mcp_servers=["server-1"]) with pytest.raises(HTTPException) as exc_info: @@ -564,9 +530,7 @@ async def test_validate_tool_permissions_validated_against_team( new_callable=AsyncMock, return_value=[], ) -async def test_validate_stale_mcp_server_ids_are_silently_dropped( - mock_access_groups, mock_allow_all -): +async def test_validate_stale_mcp_server_ids_are_silently_dropped(mock_access_groups, mock_allow_all): """ Stale MCP server IDs (servers deleted and no longer in the registry) must not block a key save with a 403. They are silently stripped instead. @@ -596,9 +560,7 @@ async def test_validate_stale_mcp_server_ids_are_silently_dropped( new_callable=AsyncMock, return_value=[], ) -async def test_validate_stale_ids_in_mcp_tool_permissions_silently_dropped( - mock_access_groups, mock_allow_all -): +async def test_validate_stale_ids_in_mcp_tool_permissions_silently_dropped(mock_access_groups, mock_allow_all): """ Stale server IDs referenced only as keys in mcp_tool_permissions (not in mcp_servers) must also be silently stripped rather than raising a 403. @@ -626,9 +588,7 @@ async def test_validate_stale_ids_in_mcp_tool_permissions_silently_dropped( new_callable=AsyncMock, return_value=[], ) -async def test_validate_stale_mcp_server_ids_are_removed_from_object_permission( - mock_access_groups, mock_allow_all -): +async def test_validate_stale_mcp_server_ids_are_removed_from_object_permission(mock_access_groups, mock_allow_all): team_obj = _make_team_obj(mcp_servers=["s3", "s4"]) object_permission = {"mcp_servers": ["s1-stale", "s2-stale"]} await validate_key_mcp_servers_against_team( @@ -661,9 +621,7 @@ async def test_validate_stale_mcp_server_ids_are_removed_from_object_permission( new_callable=AsyncMock, return_value=[], ) -async def test_validate_mcp_server_alias_outside_team_scope_raises( - mock_access_groups, mock_allow_all -): +async def test_validate_mcp_server_alias_outside_team_scope_raises(mock_access_groups, mock_allow_all): team_obj = _make_team_obj(mcp_servers=["team-server"]) with pytest.raises(HTTPException) as exc_info: await validate_key_mcp_servers_against_team( @@ -696,9 +654,7 @@ async def test_validate_mcp_server_alias_outside_team_scope_raises( new_callable=AsyncMock, return_value=[], ) -async def test_validate_mcp_server_alias_is_normalized_before_save( - mock_access_groups, mock_allow_all -): +async def test_validate_mcp_server_alias_is_normalized_before_save(mock_access_groups, mock_allow_all): team_obj = _make_team_obj(mcp_servers=["allowed-server-id"]) object_permission = { "mcp_servers": ["allowed-alias"], @@ -736,9 +692,7 @@ async def test_validate_db_mcp_server_alias_outside_team_scope_raises_when_regis mock_db_server.server_id = "private-server-id" mock_db_server.alias = "private-alias" mock_db_server.server_name = "Private Server" - mock_prisma_client.db.litellm_mcpservertable.find_many = AsyncMock( - return_value=[mock_db_server] - ) + mock_prisma_client.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[mock_db_server]) team_obj = _make_team_obj(mcp_servers=[]) with pytest.raises(HTTPException) as exc_info: @@ -762,9 +716,7 @@ async def test_validate_db_mcp_server_alias_outside_team_scope_raises_when_regis new_callable=AsyncMock, return_value=[], ) -async def test_validate_access_groups_within_team_scope( - mock_access_groups, mock_allow_all -): +async def test_validate_access_groups_within_team_scope(mock_access_groups, mock_allow_all): """Key requests access groups that are in the team's scope — should pass.""" team_obj = _make_team_obj(mcp_access_groups=["group-a", "group-b"]) await validate_key_mcp_servers_against_team( @@ -783,9 +735,7 @@ async def test_validate_access_groups_within_team_scope( new_callable=AsyncMock, return_value=[], ) -async def test_validate_access_groups_outside_team_scope_raises( - mock_access_groups, mock_allow_all -): +async def test_validate_access_groups_outside_team_scope_raises(mock_access_groups, mock_allow_all): """Key requests access groups NOT in the team's scope — should raise 403.""" team_obj = _make_team_obj(mcp_access_groups=["group-a"]) with pytest.raises(HTTPException) as exc_info: @@ -807,9 +757,7 @@ async def test_validate_access_groups_outside_team_scope_raises( new_callable=AsyncMock, return_value=[], ) -async def test_validate_access_groups_no_team_raises( - mock_access_groups, mock_allow_all -): +async def test_validate_access_groups_no_team_raises(mock_access_groups, mock_allow_all): """Key without a team requesting access groups — should raise 403.""" with pytest.raises(HTTPException) as exc_info: await validate_key_mcp_servers_against_team( @@ -830,9 +778,7 @@ async def test_validate_access_groups_no_team_raises( new_callable=AsyncMock, return_value=["server-from-group"], ) -async def test_validate_team_access_groups_resolve_to_servers( - mock_access_groups, mock_allow_all -): +async def test_validate_team_access_groups_resolve_to_servers(mock_access_groups, mock_allow_all): """Team access groups should resolve to server IDs and be included in allowed set.""" team_obj = _make_team_obj(mcp_access_groups=["group-a"]) # Key requests a server that comes from the team's access group @@ -937,9 +883,7 @@ async def test_resolve_team_all_proxy_sentinel_resolves_dynamically(mock_access_ new_callable=AsyncMock, return_value=[], ) -async def test_validate_key_scoped_to_server_added_after_team_all_proxy( - mock_access_groups, mock_allow_all -): +async def test_validate_key_scoped_to_server_added_after_team_all_proxy(mock_access_groups, mock_allow_all): """The exact user scenario: a team scoped to the all-proxy sentinel, a server (srv-z) registered afterwards, and a key scoped to just srv-z. Because the team ceiling resolves to every registered server, the key passes validation @@ -968,9 +912,7 @@ async def test_validate_key_scoped_to_server_added_after_team_all_proxy( new_callable=AsyncMock, return_value=[], ) -async def test_validate_key_scoped_to_server_rejected_when_team_not_all_proxy( - mock_access_groups, mock_allow_all -): +async def test_validate_key_scoped_to_server_rejected_when_team_not_all_proxy(mock_access_groups, mock_allow_all): """Contrast with the sentinel case: a team scoped to a concrete server list (srv-x, not the sentinel) does NOT unlock srv-z for a key. It is the sentinel specifically, not a blanket allow, that widens the team ceiling.""" @@ -1111,9 +1053,7 @@ async def test_validate_search_tools_raises_when_not_subset(): new_callable=AsyncMock, return_value=[], ) -async def test_personal_non_admin_cannot_assign_mcp_toolsets( - mock_access_groups, mock_allow_all -): +async def test_personal_non_admin_cannot_assign_mcp_toolsets(mock_access_groups, mock_allow_all): with pytest.raises(HTTPException) as exc: await validate_key_mcp_servers_against_team( object_permission={"mcp_toolsets": ["ts-private"]}, @@ -1134,9 +1074,7 @@ async def test_personal_non_admin_cannot_assign_mcp_toolsets( new_callable=AsyncMock, return_value=[], ) -async def test_personal_admin_can_assign_mcp_toolsets( - mock_access_groups, mock_allow_all -): +async def test_personal_admin_can_assign_mcp_toolsets(mock_access_groups, mock_allow_all): await validate_key_mcp_servers_against_team( object_permission={"mcp_toolsets": ["ts-private"]}, team_obj=None,