From eeb13fffbd66cabaae5397197aad46cfa647c254 Mon Sep 17 00:00:00 2001 From: L4XB Date: Mon, 14 Sep 2026 18:20:51 +0200 Subject: [PATCH 1/3] fix(key mgmt): a project may only be attached to keys of its own team MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A project is created under exactly one team, and its budget and models are validated against that team's. Nothing checked that a key's team matched, so `POST /key/generate` with `team_id: team-a` and a `project_id` owned by `team-b` answered 200 and stored exactly that — a key belonging to team-a, recorded under team-b's project and validated against its models and budget. The same held with no team at all. It needs no proxy-admin rights: an admin of team-a who is a member of no other team can issue keys under another tenant's project with nothing but the project id, and that tenant sees the spend without having granted anything. `_check_project_key_limits` now takes the key's team and refuses a project owned by a different one, before the model and budget checks so the refusal reads as what it is. A project with no owning team is left alone — nobody owns it, so there is no boundary to cross. `/key/update` passes the key's stored team when the request does not carry one, and now runs the check whenever the project itself is set or changed, which a request that moves only the project previously skipped. Two existing cells needed the project's team passed explicitly: they measure the model allowlist, not tenancy, and would otherwise have been asserting model behaviour on a request the new gate refuses. A third builds an unowned project for the same reason. Fixes #41089 --- .../key_management_endpoints.py | 29 ++++- .../test_key_management_endpoints.py | 13 +- .../test_key_project_team_ownership.py | 115 ++++++++++++++++++ 3 files changed, 152 insertions(+), 5 deletions(-) create mode 100644 tests/test_litellm/proxy/management_endpoints/test_key_project_team_ownership.py diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 95ccb7bbe0b..23dfe9e0c07 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1547,10 +1547,17 @@ async def _check_project_key_limits( data: GenerateKeyRequest | UpdateKeyRequest, prisma_client: PrismaClient, user_api_key_cache: UserApiKeyCache, + key_team_id: str | None = None, ) -> None: """ - Validate that key's models and budget respect its project's limits. + Validate that the key belongs to the project's team, and that its models + and budget respect the project's limits. + - The project's owning team must be the key's team. A project is created + under exactly one team and its budget and models are that team's, so a + key on another team recorded under it charges a tenant that never granted + anything — and issuing one needs no proxy-admin rights, only the project + id (#41089). A project with no team belongs to nobody and is left alone. - Key models must be a subset of project models, except the all-team-models / all-proxy-models sentinels, which inherit a parent scope and are narrowed by the project at request time - Key max_budget must be <= project max_budget @@ -1567,6 +1574,17 @@ async def _check_project_key_limits( detail={"error": f"Project not found, project_id={project_id}"}, ) + # Validate the project's team owns the key + if project_obj.team_id is not None and project_obj.team_id != key_team_id: + raise HTTPException( + status_code=403, + detail={ + "error": f"Project {project_id} belongs to team {project_obj.team_id}, " + f"but the key belongs to {key_team_id if key_team_id is not None else 'no team'}. " + "A project can only be attached to keys of the team that owns it." + }, + ) + # Validate key models are a subset of project models if data.models and len(project_obj.models) > 0: for m in data.models: @@ -1922,6 +1940,7 @@ async def generate_key_fn( data=data, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, + key_team_id=data.team_id, ) return await _common_key_generation_helper( @@ -2864,12 +2883,18 @@ async def _validate_update_key_data( _project_id_to_check: Final = ( data.project_id if "project_id" in data.model_fields_set else existing_key_row.project_id ) - if _project_id_to_check is not None and (data.models is not None or data.max_budget is not None): + # Also when the project itself is being set or changed: that is exactly when + # the team that owns it has to be checked, and a request that moves only the + # project carries neither models nor max_budget (#41089). + if _project_id_to_check is not None and ( + "project_id" in data.model_fields_set or data.models is not None or data.max_budget is not None + ): await _check_project_key_limits( project_id=_project_id_to_check, data=data, prisma_client=checked_prisma_client, user_api_key_cache=user_api_key_cache, + key_team_id=(data.team_id if "team_id" in data.model_fields_set else existing_key_row.team_id), ) # When the caller asks to change the key's organization_id, require that diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 2ac52da57df..37e7bc95285 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -17849,11 +17849,13 @@ async def test_regenerate_key_repoints_live_membership_not_the_key_row_it_read( ) == ["attached-model"] -async def _cache_with_project(project_id: str, project_models: list[str]) -> UserApiKeyCache: +async def _cache_with_project( + project_id: str, project_models: list[str], team_id: str | None = "team-lit-5823" +) -> UserApiKeyCache: user_api_key_cache = UserApiKeyCache() await user_api_key_cache.async_set_cache( key=_project_cache_key(project_id), - value=LiteLLM_ProjectTableCachedObj(project_id=project_id, team_id="team-lit-5823", models=project_models), + value=LiteLLM_ProjectTableCachedObj(project_id=project_id, team_id=team_id, models=project_models), model_type=LiteLLM_ProjectTableCachedObj, ) return user_api_key_cache @@ -17871,6 +17873,7 @@ async def test_check_project_key_limits_accepts_inherited_model_sentinels(reques data=request_cls(key="sk-lit-5823", models=[sentinel]), prisma_client=MagicMock(), user_api_key_cache=user_api_key_cache, + key_team_id="team-lit-5823", ) @@ -17889,6 +17892,7 @@ async def test_check_project_key_limits_still_rejects_real_model_outside_project data=request_cls(key="sk-lit-5823", models=key_models), prisma_client=MagicMock(), user_api_key_cache=user_api_key_cache, + key_team_id="team-lit-5823", ) assert exc_info.value.status_code == 400 @@ -18257,7 +18261,10 @@ async def test_project_detachment_preserves_omission_and_other_key_fields(): @pytest.mark.asyncio async def test_project_detachment_uses_effective_project_for_validation(project_id: str | None): existing: Final = LiteLLM_VerificationToken(token="project-detach-token", project_id="project-orbit") - cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + # An unowned project: this cell is about WHICH project the validation uses, + # so the ownership gate (#41089) must not be what it measures. Giving the + # key a team instead would pull the whole team lookup into a MagicMock db. + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"], team_id=None) data: Final = UpdateKeyRequest(key=existing.token, project_id=project_id, models=["model-other"]) if project_id is None: await _validate_update_key_data( diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_project_team_ownership.py b/tests/test_litellm/proxy/management_endpoints/test_key_project_team_ownership.py new file mode 100644 index 00000000000..082c9a0e501 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_key_project_team_ownership.py @@ -0,0 +1,115 @@ +"""A project may only be attached to keys of the team that owns it (#41089). + +A project is created under exactly one team, and its budget and models are that +team's. Nothing checked that the key's team matched, so an admin of team-a — a +member of no other team — could issue a key on their own team pointing at a +team-b project, and team-b would see the spend under a project they never +granted anything on. Only the project id was needed. + +The negative controls are the point: a project with no owning team is left +alone, and a key on the owning team still passes, or the check would be a wall +rather than a boundary. +""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException + +from litellm.models.project import LiteLLM_ProjectTable +from litellm.proxy._types import GenerateKeyRequest, UpdateKeyRequest + + +def _project(team_id, models=None, project_id="proj-1"): + return LiteLLM_ProjectTable( + project_id=project_id, + team_id=team_id, + models=models or [], + ) + + +async def _check(project, data, key_team_id): + """Drive _check_project_key_limits with the project the store would return.""" + from litellm.proxy.management_endpoints import key_management_endpoints as kme + + original = kme.get_project_object + kme.get_project_object = AsyncMock(return_value=project) + try: + await kme._check_project_key_limits( + project_id=project.project_id, + data=data, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + key_team_id=key_team_id, + ) + finally: + kme.get_project_object = original + + +@pytest.mark.asyncio +async def test_a_key_may_not_point_at_another_teams_project(): + with pytest.raises(HTTPException) as exc: + await _check( + _project(team_id="team-b"), + GenerateKeyRequest(team_id="team-a"), + key_team_id="team-a", + ) + assert exc.value.status_code == 403 + detail = str(exc.value.detail) + assert "team-b" in detail and "team-a" in detail + + +@pytest.mark.asyncio +async def test_a_key_with_no_team_may_not_point_at_a_teams_project(): + # The issue's step 5: no team at all still charges a team's project. + with pytest.raises(HTTPException) as exc: + await _check( + _project(team_id="team-b"), + GenerateKeyRequest(), + key_team_id=None, + ) + assert exc.value.status_code == 403 + assert "no team" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_a_key_on_the_owning_team_is_accepted(): + await _check( + _project(team_id="team-b"), + GenerateKeyRequest(team_id="team-b"), + key_team_id="team-b", + ) + + +@pytest.mark.asyncio +async def test_a_project_with_no_team_is_left_alone(): + # Nobody owns it, so there is no boundary to cross — rejecting here would + # break every project created outside a team. + await _check(_project(team_id=None), GenerateKeyRequest(team_id="team-a"), key_team_id="team-a") + await _check(_project(team_id=None), GenerateKeyRequest(), key_team_id=None) + + +@pytest.mark.asyncio +async def test_the_ownership_check_runs_before_the_model_check(): + # A foreign project must be refused as foreign, not as "model not allowed": + # the 400 would read as a configuration problem and hide the tenancy one. + with pytest.raises(HTTPException) as exc: + await _check( + _project(team_id="team-b", models=["gpt-4o-mini"]), + GenerateKeyRequest(team_id="team-a", models=["gpt-4o"]), + key_team_id="team-a", + ) + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_an_update_carries_the_keys_existing_team(): + # /key/update sends no team_id when only the project changes, so the check + # has to use the key's stored team rather than treating it as absent. + with pytest.raises(HTTPException) as exc: + await _check( + _project(team_id="team-b"), + UpdateKeyRequest(key="sk-x", project_id="proj-1"), + key_team_id="team-a", + ) + assert exc.value.status_code == 403 From 4440d4e6c754d96edf5f6942ce0cbd06d5cf1a9e Mon Sep 17 00:00:00 2001 From: L4XB Date: Mon, 14 Sep 2026 18:59:21 +0200 Subject: [PATCH 2/3] fix(key mgmt): check project ownership on every key mutation path /key/update ran the project check only when project_id, models or max_budget were supplied, so a request that changed team_id alone left a foreign project attached. /key/regenerate never ran it at all. Both now go through one helper that validates the key as the mutation leaves it. On /key/generate the check moves after default_key_generate_params is applied, because that can supply team_id; it was rejecting a valid key whose team came from the defaults. Tests move into the mapped test file per CLAUDE.md. --- .../key_management_endpoints.py | 85 +++-- .../test_key_management_endpoints.py | 318 ++++++++++++++++++ .../test_key_project_team_ownership.py | 115 ------- 3 files changed, 371 insertions(+), 147 deletions(-) delete mode 100644 tests/test_litellm/proxy/management_endpoints/test_key_project_team_ownership.py diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 23dfe9e0c07..346b46d2052 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1066,6 +1066,19 @@ async def _common_key_generation_helper( # check if user set upperbound key/generate params on config.yaml _enforce_upperbound_key_params(data, fill_defaults=True) + # Checked after the defaults, because default_key_generate_params can supply + # team_id and the project's owner is checked against the key's final team. + if data.project_id is not None and prisma_client is not None: + from litellm.proxy.proxy_server import user_api_key_cache + + await _check_project_key_limits( + project_id=data.project_id, + data=data, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + key_team_id=data.team_id, + ) + # Delegated-authority ceiling (GHSA-q775-qw9r-2r4g): a non-admin caller # cannot grant a key a higher budget than their own authority. is_ui_session_team_key = user_api_key_dict.team_id == UI_SESSION_TOKEN_TEAM_ID and _requested_team_id is not None @@ -1553,11 +1566,8 @@ async def _check_project_key_limits( Validate that the key belongs to the project's team, and that its models and budget respect the project's limits. - - The project's owning team must be the key's team. A project is created - under exactly one team and its budget and models are that team's, so a - key on another team recorded under it charges a tenant that never granted - anything — and issuing one needs no proxy-admin rights, only the project - id (#41089). A project with no team belongs to nobody and is left alone. + - The project's owning team must be the key's team. A project with no team + has no owner to protect, so it is not restricted - Key models must be a subset of project models, except the all-team-models / all-proxy-models sentinels, which inherit a parent scope and are narrowed by the project at request time - Key max_budget must be <= project max_budget @@ -1574,7 +1584,6 @@ async def _check_project_key_limits( detail={"error": f"Project not found, project_id={project_id}"}, ) - # Validate the project's team owns the key if project_obj.team_id is not None and project_obj.team_id != key_team_id: raise HTTPException( status_code=403, @@ -1610,6 +1619,33 @@ async def _check_project_key_limits( ) +# Touching any of these can change the project a key is under, the team it is on, or what the project must allow. +_PROJECT_LIMIT_FIELDS: Final = frozenset({"project_id", "team_id", "models", "max_budget"}) + + +async def _check_project_key_limits_on_mutation( + data: UpdateKeyRequest | RegenerateKeyRequest, + existing_key_row: LiteLLM_VerificationToken, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> None: + """Run _check_project_key_limits against the key as the mutation leaves it.""" + if not data.model_fields_set & _PROJECT_LIMIT_FIELDS: + return + + project_id: Final = data.project_id if "project_id" in data.model_fields_set else existing_key_row.project_id + if project_id is None: + return + + await _check_project_key_limits( + project_id=project_id, + data=data, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + key_team_id=(data.team_id if "team_id" in data.model_fields_set else existing_key_row.team_id), + ) + + def check_org_key_model_specific_limits( keys: Sequence[LiteLLM_VerificationToken], org_table: LiteLLM_OrganizationTable, @@ -1933,16 +1969,6 @@ async def generate_key_fn( prisma_client=prisma_client, ) - # Validate key against project limits if project_id is set - if data.project_id is not None: - await _check_project_key_limits( - project_id=data.project_id, - data=data, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - key_team_id=data.team_id, - ) - return await _common_key_generation_helper( data=data, user_api_key_dict=user_api_key_dict, @@ -2879,23 +2905,12 @@ async def _validate_update_key_data( access_group_ids=data.access_group_ids, ) - # Validate key against project limits if project_id is being set - _project_id_to_check: Final = ( - data.project_id if "project_id" in data.model_fields_set else existing_key_row.project_id + await _check_project_key_limits_on_mutation( + data=data, + existing_key_row=existing_key_row, + prisma_client=checked_prisma_client, + user_api_key_cache=user_api_key_cache, ) - # Also when the project itself is being set or changed: that is exactly when - # the team that owns it has to be checked, and a request that moves only the - # project carries neither models nor max_budget (#41089). - if _project_id_to_check is not None and ( - "project_id" in data.model_fields_set or data.models is not None or data.max_budget is not None - ): - await _check_project_key_limits( - project_id=_project_id_to_check, - data=data, - prisma_client=checked_prisma_client, - user_api_key_cache=user_api_key_cache, - key_team_id=(data.team_id if "team_id" in data.model_fields_set else existing_key_row.team_id), - ) # When the caller asks to change the key's organization_id, require that # they are a member of (or a proxy admin over) the target organization. @@ -5154,6 +5169,12 @@ async def _execute_virtual_key_regeneration( if data is not None: # Enforce upperbound key params on regenerate (don't fill defaults) _enforce_upperbound_key_params(data, fill_defaults=False) + await _check_project_key_limits_on_mutation( + data=data, + existing_key_row=key_in_db, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) non_default_values = await prepare_key_update_data(data=data, existing_key_row=key_in_db) # Only validate key_alias format if it's actually being changed new_key_alias: Final = non_default_values.get("key_alias") diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 37e7bc95285..0851fd5325d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -18297,3 +18297,321 @@ async def test_key_creator_cannot_detach_project_without_admin_access(): ) assert exc.value.status_code == 403 assert "Only proxy admins, team admins, or org admins" in str(exc.value.detail) + + +# --- Tests: a project may only be attached to keys of the team that owns it --- + + +def _make_owned_project(team_id, models=None, project_id="proj-owned-1"): + return LiteLLM_ProjectTableCachedObj( + project_id=project_id, + team_id=team_id, + models=models or [], + ) + + +async def _check_project_limits_with(project, data, key_team_id): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _check_project_key_limits, + ) + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_project_object", + new_callable=AsyncMock, + return_value=project, + ): + await _check_project_key_limits( + project_id=project.project_id, + data=data, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + key_team_id=key_team_id, + ) + + +async def _check_project_limits_on_mutation_with(project, data, existing_key_row): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _check_project_key_limits_on_mutation, + ) + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_project_object", + new_callable=AsyncMock, + return_value=project, + ): + await _check_project_key_limits_on_mutation( + data=data, + existing_key_row=existing_key_row, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + ) + + +@pytest.mark.asyncio +async def test_a_key_may_not_point_at_another_teams_project(): + with pytest.raises(HTTPException) as exc: + await _check_project_limits_with( + _make_owned_project(team_id="team-b"), + GenerateKeyRequest(team_id="team-a"), + key_team_id="team-a", + ) + assert exc.value.status_code == 403 + detail = str(exc.value.detail) + assert "team-b" in detail and "team-a" in detail + + +@pytest.mark.asyncio +async def test_a_key_with_no_team_may_not_point_at_a_teams_project(): + with pytest.raises(HTTPException) as exc: + await _check_project_limits_with( + _make_owned_project(team_id="team-b"), + GenerateKeyRequest(), + key_team_id=None, + ) + assert exc.value.status_code == 403 + assert "no team" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_a_key_on_the_owning_team_is_accepted(): + await _check_project_limits_with( + _make_owned_project(team_id="team-a", models=["gpt-4o"]), + GenerateKeyRequest(team_id="team-a", models=["gpt-4o"]), + key_team_id="team-a", + ) + + +@pytest.mark.asyncio +async def test_a_project_with_no_team_is_not_restricted(): + await _check_project_limits_with( + _make_owned_project(team_id=None), + GenerateKeyRequest(team_id="team-a"), + key_team_id="team-a", + ) + await _check_project_limits_with( + _make_owned_project(team_id=None), + GenerateKeyRequest(), + key_team_id=None, + ) + + +@pytest.mark.asyncio +async def test_a_foreign_project_is_refused_as_foreign_not_as_a_model_problem(): + with pytest.raises(HTTPException) as exc: + await _check_project_limits_with( + _make_owned_project(team_id="team-b", models=["gpt-4o-mini"]), + GenerateKeyRequest(team_id="team-a", models=["gpt-4o"]), + key_team_id="team-a", + ) + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_an_update_that_moves_only_the_project_uses_the_keys_stored_team(): + existing: Final = LiteLLM_VerificationToken(token="sk-hash", team_id="team-a") + with pytest.raises(HTTPException) as exc: + await _check_project_limits_on_mutation_with( + _make_owned_project(team_id="team-b"), + UpdateKeyRequest(key="sk-x", project_id="proj-owned-1"), + existing, + ) + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_an_update_that_moves_only_the_team_still_checks_the_attached_project(): + existing: Final = LiteLLM_VerificationToken( + token="sk-hash", team_id="team-b", project_id="proj-owned-1" + ) + with pytest.raises(HTTPException) as exc: + await _check_project_limits_on_mutation_with( + _make_owned_project(team_id="team-b"), + UpdateKeyRequest(key="sk-x", team_id="team-a"), + existing, + ) + assert exc.value.status_code == 403 + assert "team-a" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_an_update_that_moves_the_key_to_the_projects_own_team_is_accepted(): + existing: Final = LiteLLM_VerificationToken( + token="sk-hash", team_id=None, project_id="proj-owned-1" + ) + await _check_project_limits_on_mutation_with( + _make_owned_project(team_id="team-b"), + UpdateKeyRequest(key="sk-x", team_id="team-b"), + existing, + ) + + +@pytest.mark.asyncio +async def test_an_update_that_touches_neither_does_not_look_the_project_up(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _check_project_key_limits_on_mutation, + ) + + existing: Final = LiteLLM_VerificationToken( + token="sk-hash", team_id="team-a", project_id="proj-owned-1" + ) + lookup: Final = AsyncMock(return_value=_make_owned_project(team_id="team-b")) + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_project_object", + new=lookup, + ): + await _check_project_key_limits_on_mutation( + data=UpdateKeyRequest(key="sk-x", key_alias="renamed"), + existing_key_row=existing, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + ) + assert lookup.await_count == 0 + + +@pytest.mark.asyncio +async def test_regenerate_may_not_attach_another_teams_project(): + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + + existing_key: Final = LiteLLM_VerificationToken(token="abc123", team_id="team-a") + mock_prisma_client: Final = _make_regenerate_mock_prisma() + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_project_object", + new_callable=AsyncMock, + return_value=_make_owned_project(team_id="team-b"), + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=RegenerateKeyRequest(project_id="proj-owned-1"), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert exc_info.value.status_code == 403 + # A refused regenerate must not reach the DB update. + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 + + +@pytest.mark.asyncio +async def test_regenerate_may_attach_a_project_of_the_keys_own_team(): + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + + existing_key: Final = LiteLLM_VerificationToken(token="abc123", team_id="team-b") + mock_prisma_client: Final = _make_regenerate_mock_prisma() + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_project_object", + new_callable=AsyncMock, + return_value=_make_owned_project(team_id="team-b"), + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=RegenerateKeyRequest(project_id="proj-owned-1"), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 + + +def _make_generate_mock_prisma(): + """Mock prisma client shaped for _common_key_generation_helper.""" + mock_prisma_client = AsyncMock() + mock_prisma_client.insert_data = AsyncMock( + return_value=MagicMock( + token="hashed_token_123", litellm_budget_table=None, object_permission=None + ) + ) + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_verificationtoken = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock( + return_value=MagicMock( + token="hashed_token_123", litellm_budget_table=None, object_permission=None + ) + ) + return mock_prisma_client + + +async def _generate_key_with_defaulted_team(monkeypatch, project_team_id): + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", _make_generate_mock_prisma() + ) + monkeypatch.setattr(litellm, "default_key_generate_params", {"team_id": "team-b"}) + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_project_object", + new_callable=AsyncMock, + return_value=_make_owned_project(team_id=project_team_id), + ): + return await _common_key_generation_helper( + data=GenerateKeyRequest(project_id="proj-owned-1"), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="1234" + ), + litellm_changed_by=None, + team_table=None, + ) + + +@pytest.mark.asyncio +async def test_a_team_supplied_by_defaults_is_what_the_project_is_checked_against(monkeypatch): + # default_key_generate_params fills team_id after the request is parsed, so a + # request with no team_id still ends up on team-b and may use its projects. + await _generate_key_with_defaulted_team(monkeypatch, project_team_id="team-b") + + +@pytest.mark.asyncio +async def test_a_team_supplied_by_defaults_does_not_open_another_teams_project(monkeypatch): + with pytest.raises(HTTPException) as exc: + await _generate_key_with_defaulted_team(monkeypatch, project_team_id="team-c") + assert exc.value.status_code == 403 diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_project_team_ownership.py b/tests/test_litellm/proxy/management_endpoints/test_key_project_team_ownership.py deleted file mode 100644 index 082c9a0e501..00000000000 --- a/tests/test_litellm/proxy/management_endpoints/test_key_project_team_ownership.py +++ /dev/null @@ -1,115 +0,0 @@ -"""A project may only be attached to keys of the team that owns it (#41089). - -A project is created under exactly one team, and its budget and models are that -team's. Nothing checked that the key's team matched, so an admin of team-a — a -member of no other team — could issue a key on their own team pointing at a -team-b project, and team-b would see the spend under a project they never -granted anything on. Only the project id was needed. - -The negative controls are the point: a project with no owning team is left -alone, and a key on the owning team still passes, or the check would be a wall -rather than a boundary. -""" - -from unittest.mock import AsyncMock, MagicMock - -import pytest -from fastapi import HTTPException - -from litellm.models.project import LiteLLM_ProjectTable -from litellm.proxy._types import GenerateKeyRequest, UpdateKeyRequest - - -def _project(team_id, models=None, project_id="proj-1"): - return LiteLLM_ProjectTable( - project_id=project_id, - team_id=team_id, - models=models or [], - ) - - -async def _check(project, data, key_team_id): - """Drive _check_project_key_limits with the project the store would return.""" - from litellm.proxy.management_endpoints import key_management_endpoints as kme - - original = kme.get_project_object - kme.get_project_object = AsyncMock(return_value=project) - try: - await kme._check_project_key_limits( - project_id=project.project_id, - data=data, - prisma_client=MagicMock(), - user_api_key_cache=MagicMock(), - key_team_id=key_team_id, - ) - finally: - kme.get_project_object = original - - -@pytest.mark.asyncio -async def test_a_key_may_not_point_at_another_teams_project(): - with pytest.raises(HTTPException) as exc: - await _check( - _project(team_id="team-b"), - GenerateKeyRequest(team_id="team-a"), - key_team_id="team-a", - ) - assert exc.value.status_code == 403 - detail = str(exc.value.detail) - assert "team-b" in detail and "team-a" in detail - - -@pytest.mark.asyncio -async def test_a_key_with_no_team_may_not_point_at_a_teams_project(): - # The issue's step 5: no team at all still charges a team's project. - with pytest.raises(HTTPException) as exc: - await _check( - _project(team_id="team-b"), - GenerateKeyRequest(), - key_team_id=None, - ) - assert exc.value.status_code == 403 - assert "no team" in str(exc.value.detail) - - -@pytest.mark.asyncio -async def test_a_key_on_the_owning_team_is_accepted(): - await _check( - _project(team_id="team-b"), - GenerateKeyRequest(team_id="team-b"), - key_team_id="team-b", - ) - - -@pytest.mark.asyncio -async def test_a_project_with_no_team_is_left_alone(): - # Nobody owns it, so there is no boundary to cross — rejecting here would - # break every project created outside a team. - await _check(_project(team_id=None), GenerateKeyRequest(team_id="team-a"), key_team_id="team-a") - await _check(_project(team_id=None), GenerateKeyRequest(), key_team_id=None) - - -@pytest.mark.asyncio -async def test_the_ownership_check_runs_before_the_model_check(): - # A foreign project must be refused as foreign, not as "model not allowed": - # the 400 would read as a configuration problem and hide the tenancy one. - with pytest.raises(HTTPException) as exc: - await _check( - _project(team_id="team-b", models=["gpt-4o-mini"]), - GenerateKeyRequest(team_id="team-a", models=["gpt-4o"]), - key_team_id="team-a", - ) - assert exc.value.status_code == 403 - - -@pytest.mark.asyncio -async def test_an_update_carries_the_keys_existing_team(): - # /key/update sends no team_id when only the project changes, so the check - # has to use the key's stored team rather than treating it as absent. - with pytest.raises(HTTPException) as exc: - await _check( - _project(team_id="team-b"), - UpdateKeyRequest(key="sk-x", project_id="proj-1"), - key_team_id="team-a", - ) - assert exc.value.status_code == 403 From 3be888dec74883fcf6780585c5daa3e2d1ac9edb Mon Sep 17 00:00:00 2001 From: L4XB Date: Mon, 14 Sep 2026 19:58:35 +0200 Subject: [PATCH 3/3] test(key mgmt): read the project from the cache instead of patching the lookup The test-quality gate rejected the new cells: 12 TQ008 for patching `get_project_object`, an SDK internal, and 4 TQ001 for accept controls that could only fail by raising. The cells now seed `UserApiKeyCache` with the project, which is the idiom the neighbouring project cells already use and removes the patching. The accept controls are parametrised together with the rejecting ones, so each test function carries a real assertion. The three regenerate seams that stay stubbed each carry a reason. --- .../test_key_management_endpoints.py | 388 ++++++++---------- 1 file changed, 162 insertions(+), 226 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 0851fd5325d..5a2a080a495 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -35,6 +35,7 @@ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_org_key_limits, _check_project_key_limits, + _check_project_key_limits_on_mutation, _check_team_key_limits, _common_key_generation_helper, _enforce_upperbound_key_params, @@ -18301,265 +18302,192 @@ async def test_key_creator_cannot_detach_project_without_admin_access(): # --- Tests: a project may only be attached to keys of the team that owns it --- +_OWNED_PROJECT: Final = "proj-owned-1" -def _make_owned_project(team_id, models=None, project_id="proj-owned-1"): - return LiteLLM_ProjectTableCachedObj( - project_id=project_id, - team_id=team_id, - models=models or [], + +@pytest.mark.parametrize( + "project_team_id, key_team_id, expected_status, expected_in_detail", + [ + # A key on another team must not point at this project. + ("team-b", "team-a", 403, ["team-b", "team-a"]), + # The issue's step 5: no team at all still charges a team's project. + ("team-b", None, 403, ["no team"]), + # Accept controls. Without these the check is a wall, not a boundary. + ("team-a", "team-a", None, []), + # A project with no owning team has nobody to protect. + (None, "team-a", None, []), + (None, None, None, []), + ], +) +@pytest.mark.asyncio +async def test_a_project_is_only_for_keys_of_its_own_team( + project_team_id, key_team_id, expected_status, expected_in_detail +): + user_api_key_cache: Final = await _cache_with_project( + _OWNED_PROJECT, [], team_id=project_team_id ) - -async def _check_project_limits_with(project, data, key_team_id): - from litellm.proxy.management_endpoints.key_management_endpoints import ( - _check_project_key_limits, - ) - - with patch( - "litellm.proxy.management_endpoints.key_management_endpoints.get_project_object", - new_callable=AsyncMock, - return_value=project, - ): + async def check(): await _check_project_key_limits( - project_id=project.project_id, - data=data, + project_id=_OWNED_PROJECT, + data=GenerateKeyRequest(team_id=key_team_id), prisma_client=MagicMock(), - user_api_key_cache=MagicMock(), + user_api_key_cache=user_api_key_cache, key_team_id=key_team_id, ) + if expected_status is None: + await check() + # The project was read and accepted, rather than never reached. + assert await user_api_key_cache.async_get_cache(key=_project_cache_key(_OWNED_PROJECT)) + return -async def _check_project_limits_on_mutation_with(project, data, existing_key_row): - from litellm.proxy.management_endpoints.key_management_endpoints import ( - _check_project_key_limits_on_mutation, - ) - - with patch( - "litellm.proxy.management_endpoints.key_management_endpoints.get_project_object", - new_callable=AsyncMock, - return_value=project, - ): - await _check_project_key_limits_on_mutation( - data=data, - existing_key_row=existing_key_row, - prisma_client=MagicMock(), - user_api_key_cache=MagicMock(), - ) - - -@pytest.mark.asyncio -async def test_a_key_may_not_point_at_another_teams_project(): with pytest.raises(HTTPException) as exc: - await _check_project_limits_with( - _make_owned_project(team_id="team-b"), - GenerateKeyRequest(team_id="team-a"), - key_team_id="team-a", - ) - assert exc.value.status_code == 403 - detail = str(exc.value.detail) - assert "team-b" in detail and "team-a" in detail - - -@pytest.mark.asyncio -async def test_a_key_with_no_team_may_not_point_at_a_teams_project(): - with pytest.raises(HTTPException) as exc: - await _check_project_limits_with( - _make_owned_project(team_id="team-b"), - GenerateKeyRequest(), - key_team_id=None, - ) - assert exc.value.status_code == 403 - assert "no team" in str(exc.value.detail) - - -@pytest.mark.asyncio -async def test_a_key_on_the_owning_team_is_accepted(): - await _check_project_limits_with( - _make_owned_project(team_id="team-a", models=["gpt-4o"]), - GenerateKeyRequest(team_id="team-a", models=["gpt-4o"]), - key_team_id="team-a", - ) - - -@pytest.mark.asyncio -async def test_a_project_with_no_team_is_not_restricted(): - await _check_project_limits_with( - _make_owned_project(team_id=None), - GenerateKeyRequest(team_id="team-a"), - key_team_id="team-a", - ) - await _check_project_limits_with( - _make_owned_project(team_id=None), - GenerateKeyRequest(), - key_team_id=None, - ) + await check() + assert exc.value.status_code == expected_status + for fragment in expected_in_detail: + assert fragment in str(exc.value.detail) @pytest.mark.asyncio async def test_a_foreign_project_is_refused_as_foreign_not_as_a_model_problem(): + """A 403 reads as a tenancy problem; the 400 would read as a configuration one.""" + user_api_key_cache: Final = await _cache_with_project( + _OWNED_PROJECT, ["gpt-4o-mini"], team_id="team-b" + ) + with pytest.raises(HTTPException) as exc: - await _check_project_limits_with( - _make_owned_project(team_id="team-b", models=["gpt-4o-mini"]), - GenerateKeyRequest(team_id="team-a", models=["gpt-4o"]), + await _check_project_key_limits( + project_id=_OWNED_PROJECT, + data=GenerateKeyRequest(team_id="team-a", models=["gpt-4o"]), + prisma_client=MagicMock(), + user_api_key_cache=user_api_key_cache, key_team_id="team-a", ) + assert exc.value.status_code == 403 +@pytest.mark.parametrize( + "existing_team_id, data_kwargs, expected_status", + [ + # Moves only the project: the key's stored team is what counts. + ("team-a", {"project_id": _OWNED_PROJECT}, 403), + # Moves only the team: the project stays attached, so it still counts. + ("team-b", {"team_id": "team-a"}, 403), + # Moves the key onto the project's own team. Accept control. + (None, {"team_id": "team-b"}, None), + # Touches neither, so there is nothing to re-check. + ("team-a", {"key_alias": "renamed"}, None), + ], +) @pytest.mark.asyncio -async def test_an_update_that_moves_only_the_project_uses_the_keys_stored_team(): - existing: Final = LiteLLM_VerificationToken(token="sk-hash", team_id="team-a") - with pytest.raises(HTTPException) as exc: - await _check_project_limits_on_mutation_with( - _make_owned_project(team_id="team-b"), - UpdateKeyRequest(key="sk-x", project_id="proj-owned-1"), - existing, - ) - assert exc.value.status_code == 403 - - -@pytest.mark.asyncio -async def test_an_update_that_moves_only_the_team_still_checks_the_attached_project(): +async def test_an_update_checks_the_project_the_mutation_leaves_attached( + existing_team_id, data_kwargs, expected_status +): existing: Final = LiteLLM_VerificationToken( - token="sk-hash", team_id="team-b", project_id="proj-owned-1" + token="sk-hash", team_id=existing_team_id, project_id=_OWNED_PROJECT ) - with pytest.raises(HTTPException) as exc: - await _check_project_limits_on_mutation_with( - _make_owned_project(team_id="team-b"), - UpdateKeyRequest(key="sk-x", team_id="team-a"), - existing, - ) - assert exc.value.status_code == 403 - assert "team-a" in str(exc.value.detail) + user_api_key_cache: Final = await _cache_with_project(_OWNED_PROJECT, [], team_id="team-b") - -@pytest.mark.asyncio -async def test_an_update_that_moves_the_key_to_the_projects_own_team_is_accepted(): - existing: Final = LiteLLM_VerificationToken( - token="sk-hash", team_id=None, project_id="proj-owned-1" - ) - await _check_project_limits_on_mutation_with( - _make_owned_project(team_id="team-b"), - UpdateKeyRequest(key="sk-x", team_id="team-b"), - existing, - ) - - -@pytest.mark.asyncio -async def test_an_update_that_touches_neither_does_not_look_the_project_up(): - from litellm.proxy.management_endpoints.key_management_endpoints import ( - _check_project_key_limits_on_mutation, - ) - - existing: Final = LiteLLM_VerificationToken( - token="sk-hash", team_id="team-a", project_id="proj-owned-1" - ) - lookup: Final = AsyncMock(return_value=_make_owned_project(team_id="team-b")) - with patch( - "litellm.proxy.management_endpoints.key_management_endpoints.get_project_object", - new=lookup, - ): + async def check(): await _check_project_key_limits_on_mutation( - data=UpdateKeyRequest(key="sk-x", key_alias="renamed"), + data=UpdateKeyRequest(key="sk-x", **data_kwargs), existing_key_row=existing, prisma_client=MagicMock(), - user_api_key_cache=MagicMock(), + user_api_key_cache=user_api_key_cache, ) - assert lookup.await_count == 0 + + if expected_status is None: + await check() + assert await user_api_key_cache.async_get_cache(key=_project_cache_key(_OWNED_PROJECT)) + return + + with pytest.raises(HTTPException) as exc: + await check() + assert exc.value.status_code == expected_status @pytest.mark.asyncio -async def test_regenerate_may_not_attach_another_teams_project(): +async def test_an_update_that_touches_none_of_the_fields_does_not_read_the_project(): + existing: Final = LiteLLM_VerificationToken( + token="sk-hash", team_id="team-a", project_id=_OWNED_PROJECT + ) + prisma_client: Final = MagicMock() + + # The cache is empty, so any lookup would have to reach the database. + await _check_project_key_limits_on_mutation( + data=UpdateKeyRequest(key="sk-x", key_alias="renamed"), + existing_key_row=existing, + prisma_client=prisma_client, + user_api_key_cache=UserApiKeyCache(), + ) + + assert prisma_client.mock_calls == [] + + +@pytest.mark.parametrize( + "key_team_id, expected_status, expected_updates", + [ + # /key/regenerate wrote project_id and team_id with no project check. + ("team-a", 403, 0), + # Accept control: a key already on the project's team regenerates. + ("team-b", None, 1), + ], +) +@pytest.mark.asyncio +async def test_regenerate_checks_the_project_it_attaches( + key_team_id, expected_status, expected_updates +): from litellm.proxy._types import RegenerateKeyRequest from litellm.proxy.management_endpoints.key_management_endpoints import ( _execute_virtual_key_regeneration, ) - existing_key: Final = LiteLLM_VerificationToken(token="abc123", team_id="team-a") + existing_key: Final = LiteLLM_VerificationToken(token="abc123", team_id=key_team_id) mock_prisma_client: Final = _make_regenerate_mock_prisma() + user_api_key_cache: Final = await _cache_with_project(_OWNED_PROJECT, [], team_id="team-b") - with ( - patch( - "litellm.proxy.management_endpoints.key_management_endpoints.get_project_object", - new_callable=AsyncMock, - return_value=_make_owned_project(team_id="team-b"), - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", - new_callable=AsyncMock, - return_value="sk-newtoken1234ab12", - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", - new_callable=AsyncMock, - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", - new_callable=AsyncMock, - ), - ): - with pytest.raises(HTTPException) as exc_info: + async def regenerate(): + with ( + patch( # test-quality-ok: a fresh random token is a side effect of regeneration, not the project check under test; the file's other regenerate cells stub the same seam + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), + patch( # test-quality-ok: the deprecated-key row is a side effect of regeneration; stubbing it keeps the DB assertion below about the key update alone + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: the cache delete is a side effect of regeneration and would evict the seeded project this cell reads + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + ): await _execute_virtual_key_regeneration( prisma_client=mock_prisma_client, key_in_db=existing_key, hashed_api_key="abc123", key="abc123", - data=RegenerateKeyRequest(project_id="proj-owned-1"), + data=RegenerateKeyRequest(project_id=_OWNED_PROJECT), user_api_key_dict=_make_regenerate_user_api_key_dict(), litellm_changed_by=None, - user_api_key_cache=MagicMock(), + user_api_key_cache=user_api_key_cache, proxy_logging_obj=MagicMock(), ) - assert exc_info.value.status_code == 403 + if expected_status is None: + await regenerate() + else: + with pytest.raises(HTTPException) as exc: + await regenerate() + assert exc.value.status_code == expected_status + # A refused regenerate must not reach the DB update. - assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 0 - - -@pytest.mark.asyncio -async def test_regenerate_may_attach_a_project_of_the_keys_own_team(): - from litellm.proxy._types import RegenerateKeyRequest - from litellm.proxy.management_endpoints.key_management_endpoints import ( - _execute_virtual_key_regeneration, + assert ( + mock_prisma_client.db.litellm_verificationtoken.update.await_count == expected_updates ) - existing_key: Final = LiteLLM_VerificationToken(token="abc123", team_id="team-b") - mock_prisma_client: Final = _make_regenerate_mock_prisma() - - with ( - patch( - "litellm.proxy.management_endpoints.key_management_endpoints.get_project_object", - new_callable=AsyncMock, - return_value=_make_owned_project(team_id="team-b"), - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", - new_callable=AsyncMock, - return_value="sk-newtoken1234ab12", - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", - new_callable=AsyncMock, - ), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", - new_callable=AsyncMock, - ), - ): - await _execute_virtual_key_regeneration( - prisma_client=mock_prisma_client, - key_in_db=existing_key, - hashed_api_key="abc123", - key="abc123", - data=RegenerateKeyRequest(project_id="proj-owned-1"), - user_api_key_dict=_make_regenerate_user_api_key_dict(), - litellm_changed_by=None, - user_api_key_cache=MagicMock(), - proxy_logging_obj=MagicMock(), - ) - - assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == 1 - def _make_generate_mock_prisma(): """Mock prisma client shaped for _common_key_generation_helper.""" @@ -18582,19 +18510,32 @@ def _make_generate_mock_prisma(): return mock_prisma_client -async def _generate_key_with_defaulted_team(monkeypatch, project_team_id): +@pytest.mark.parametrize( + "project_team_id, expected_status", + [ + # default_key_generate_params fills team_id after the request is parsed, + # so the key ends up on team-b and may use its projects. + ("team-b", None), + # ... and still may not use another team's. + ("team-c", 403), + ], +) +@pytest.mark.asyncio +async def test_a_team_supplied_by_defaults_is_what_the_project_is_checked_against( + monkeypatch, project_team_id, expected_status +): monkeypatch.setattr( "litellm.proxy.proxy_server.prisma_client", _make_generate_mock_prisma() ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", + await _cache_with_project(_OWNED_PROJECT, [], team_id=project_team_id), + ) monkeypatch.setattr(litellm, "default_key_generate_params", {"team_id": "team-b"}) - with patch( - "litellm.proxy.management_endpoints.key_management_endpoints.get_project_object", - new_callable=AsyncMock, - return_value=_make_owned_project(team_id=project_team_id), - ): + async def generate(): return await _common_key_generation_helper( - data=GenerateKeyRequest(project_id="proj-owned-1"), + data=GenerateKeyRequest(project_id=_OWNED_PROJECT), user_api_key_dict=UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="1234" ), @@ -18602,16 +18543,11 @@ async def _generate_key_with_defaulted_team(monkeypatch, project_team_id): team_table=None, ) + if expected_status is None: + response = await generate() + assert response.team_id == "team-b" + return -@pytest.mark.asyncio -async def test_a_team_supplied_by_defaults_is_what_the_project_is_checked_against(monkeypatch): - # default_key_generate_params fills team_id after the request is parsed, so a - # request with no team_id still ends up on team-b and may use its projects. - await _generate_key_with_defaulted_team(monkeypatch, project_team_id="team-b") - - -@pytest.mark.asyncio -async def test_a_team_supplied_by_defaults_does_not_open_another_teams_project(monkeypatch): with pytest.raises(HTTPException) as exc: - await _generate_key_with_defaulted_team(monkeypatch, project_team_id="team-c") - assert exc.value.status_code == 403 + await generate() + assert exc.value.status_code == expected_status