scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-03 02:22:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80

test(key_generate): cover mismatched team org IDOR path on generate

Browse source 
Add test_generate_key_foreign_org_with_mismatched_team_still_enforces_membership
for the case where a team is present but request organization_id differs from
team_table.organization_id. Enterprise inheritance is no-op'd in the test so
the guard is exercised directly; membership validation must still run.

Addresses Greptile review on #29310.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Milan 2026-05-30 01:54:25 +03:00
parent 701fa44c47
commit a634d831c0
No known key found for this signature in database
1 changed files with 82 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

82
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -3193,6 +3193,88 @@ async def test_generate_key_foreign_org_without_team_still_enforces_membership()
assert mock_validate_org.call_args.kwargs["organization_id"] == foreign_org_id
@pytest.mark.asyncio
async def test_generate_key_foreign_org_with_mismatched_team_still_enforces_membership():
"""VERIA-55: when a team is present but its organization_id differs from the
organization_id on the key request, the org-membership check must still run."""
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.management_endpoints.key_management_endpoints import (
_common_key_generation_helper,
)
team_org_id = "other-org"
foreign_org_id = "someone-elses-org"
mock_team_table = MagicMock()
mock_team_table.organization_id = team_org_id
mock_team_table.metadata = None
mock_validate_org = AsyncMock()
mock_generate_key = AsyncMock(
return_value={
"key": "sk-test-key",
"expires": None,
"user_id": "alice",
"team_id": "team-1",
}
)
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.premium_user", True),
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.validate_key_mcp_servers_against_team",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.validate_key_search_tools_against_team",
new_callable=AsyncMock,
),
patch(
"litellm_enterprise.proxy.management_endpoints.key_management_endpoints.apply_enterprise_key_management_params",
side_effect=lambda data, team_table: data,
),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints._validate_caller_can_assign_key_org",
mock_validate_org,
),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.get_org_object",
new_callable=AsyncMock,
return_value=MagicMock(litellm_budget_table=None),
),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints._check_org_key_limits",
new_callable=AsyncMock,
),
patch(
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
mock_generate_key,
),
):
await _common_key_generation_helper(
data=GenerateKeyRequest(
user_id="alice",
team_id="team-1",
organization_id=foreign_org_id,
),
user_api_key_dict=UserAPIKeyAuth(
user_id="alice",
user_role=LitellmUserRoles.INTERNAL_USER.value,
),
litellm_changed_by=None,
team_table=mock_team_table,
)
mock_validate_org.assert_awaited_once()
assert mock_validate_org.call_args.kwargs["organization_id"] == foreign_org_id
# ============================================
# Organization Key Limit Tests
# ============================================