From 516b18fecaa49a39d08217d30c3f59faf149e871 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 26 Feb 2026 20:07:51 -0800 Subject: [PATCH] [Feature] Access group CRUD: Add bidirectional sync for teams/keys When creating, updating, or deleting access groups, automatically keep team and key access_group_ids in sync with the access group's assigned_team_ids and assigned_key_ids. Includes transaction-based DB updates, cache patching, and handles out-of-sync data by unioning assigned_* fields with hasSome queries. Adds 12 new tests covering sync behavior across all three CRUD operations. Co-Authored-By: Claude Sonnet 4.6 --- .../access_group_endpoints.py | 363 +++++++++++++----- .../test_access_group_endpoints.py | 284 ++++++++++++++ 2 files changed, 560 insertions(+), 87 deletions(-) diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 12aa748bbc..c6958240d9 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -1,4 +1,4 @@ -from typing import List +from typing import List, Set from fastapi import APIRouter, Depends, HTTPException, status @@ -94,6 +94,183 @@ async def _invalidate_cache_access_group(access_group_id: str) -> None: ) +# --------------------------------------------------------------------------- +# DB sync helpers (called inside a Prisma transaction) +# --------------------------------------------------------------------------- + + +async def _sync_add_access_group_to_teams( + tx, team_ids: List[str], access_group_id: str +) -> None: + """Add access_group_id to each team's access_group_ids (idempotent).""" + for team_id in team_ids: + team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id}) + if team is not None and access_group_id not in (team.access_group_ids or []): + await tx.litellm_teamtable.update( + where={"team_id": team_id}, + data={"access_group_ids": list(team.access_group_ids or []) + [access_group_id]}, + ) + + +async def _sync_remove_access_group_from_teams( + tx, team_ids: List[str], access_group_id: str +) -> None: + """Remove access_group_id from each team's access_group_ids (idempotent).""" + for team_id in team_ids: + team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id}) + if team is not None and access_group_id in (team.access_group_ids or []): + await tx.litellm_teamtable.update( + where={"team_id": team_id}, + data={"access_group_ids": [ag for ag in team.access_group_ids if ag != access_group_id]}, + ) + + +async def _sync_add_access_group_to_keys( + tx, key_tokens: List[str], access_group_id: str +) -> None: + """Add access_group_id to each key's access_group_ids (idempotent).""" + for token in key_tokens: + key = await tx.litellm_verificationtoken.find_unique(where={"token": token}) + if key is not None and access_group_id not in (key.access_group_ids or []): + await tx.litellm_verificationtoken.update( + where={"token": token}, + data={"access_group_ids": list(key.access_group_ids or []) + [access_group_id]}, + ) + + +async def _sync_remove_access_group_from_keys( + tx, key_tokens: List[str], access_group_id: str +) -> None: + """Remove access_group_id from each key's access_group_ids (idempotent).""" + for token in key_tokens: + key = await tx.litellm_verificationtoken.find_unique(where={"token": token}) + if key is not None and access_group_id in (key.access_group_ids or []): + await tx.litellm_verificationtoken.update( + where={"token": token}, + data={"access_group_ids": [ag for ag in key.access_group_ids if ag != access_group_id]}, + ) + + +# --------------------------------------------------------------------------- +# Cache patch helpers +# --------------------------------------------------------------------------- + + +async def _patch_team_caches_add_access_group( + team_ids: List[str], + access_group_id: str, + user_api_key_cache, + proxy_logging_obj, +) -> None: + """Patch cached team objects to include access_group_id.""" + for team_id in team_ids: + cached_team = await _get_team_object_from_cache( + key="team_id:{}".format(team_id), + proxy_logging_obj=proxy_logging_obj, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + ) + if cached_team is None: + continue + if cached_team.access_group_ids is None: + cached_team.access_group_ids = [access_group_id] + elif access_group_id not in cached_team.access_group_ids: + cached_team.access_group_ids = list(cached_team.access_group_ids) + [access_group_id] + else: + continue + await _cache_team_object( + team_id=team_id, + team_table=cached_team, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + +async def _patch_team_caches_remove_access_group( + team_ids: List[str], + access_group_id: str, + user_api_key_cache, + proxy_logging_obj, +) -> None: + """Patch cached team objects to remove access_group_id.""" + for team_id in team_ids: + cached_team = await _get_team_object_from_cache( + key="team_id:{}".format(team_id), + proxy_logging_obj=proxy_logging_obj, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + ) + if cached_team is not None and cached_team.access_group_ids: + cached_team.access_group_ids = [ + ag for ag in cached_team.access_group_ids if ag != access_group_id + ] + await _cache_team_object( + team_id=team_id, + team_table=cached_team, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + +async def _patch_key_caches_add_access_group( + key_tokens: List[str], + access_group_id: str, + user_api_key_cache, + proxy_logging_obj, +) -> None: + """Patch cached key objects to include access_group_id.""" + for token in key_tokens: + cached_key = await user_api_key_cache.async_get_cache(key=token) + if cached_key is None: + continue + if isinstance(cached_key, dict): + cached_key = UserAPIKeyAuth(**cached_key) + if not isinstance(cached_key, UserAPIKeyAuth): + continue + if cached_key.access_group_ids is None: + cached_key.access_group_ids = [access_group_id] + elif access_group_id not in cached_key.access_group_ids: + cached_key.access_group_ids = list(cached_key.access_group_ids) + [access_group_id] + else: + continue + await _cache_key_object( + hashed_token=token, + user_api_key_obj=cached_key, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + +async def _patch_key_caches_remove_access_group( + key_tokens: List[str], + access_group_id: str, + user_api_key_cache, + proxy_logging_obj, +) -> None: + """Patch cached key objects to remove access_group_id.""" + for token in key_tokens: + cached_key = await user_api_key_cache.async_get_cache(key=token) + if cached_key is None: + continue + if isinstance(cached_key, dict): + cached_key = UserAPIKeyAuth(**cached_key) + if isinstance(cached_key, UserAPIKeyAuth) and cached_key.access_group_ids: + cached_key.access_group_ids = [ + ag for ag in cached_key.access_group_ids if ag != access_group_id + ] + await _cache_key_object( + hashed_token=token, + user_api_key_obj=cached_key, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + +# --------------------------------------------------------------------------- +# CRUD endpoints +# --------------------------------------------------------------------------- + + @router.post( "/v1/access_group", response_model=AccessGroupResponse, @@ -106,32 +283,42 @@ async def create_access_group( _require_proxy_admin(user_api_key_dict) prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) - existing = await prisma_client.db.litellm_accessgrouptable.find_unique( - where={"access_group_name": data.access_group_name} - ) - if existing is not None: - raise HTTPException( - status_code=status.HTTP_409_CONFLICT, - detail=f"Access group '{data.access_group_name}' already exists", - ) - try: - record = await prisma_client.db.litellm_accessgrouptable.create( - data={ - "access_group_name": data.access_group_name, - "description": data.description, - "access_model_names": data.access_model_names or [], - "access_mcp_server_ids": data.access_mcp_server_ids or [], - "access_agent_ids": data.access_agent_ids or [], - "assigned_team_ids": data.assigned_team_ids or [], - "assigned_key_ids": data.assigned_key_ids or [], - "created_by": user_api_key_dict.user_id, - "updated_by": user_api_key_dict.user_id, - } - ) + async with prisma_client.db.tx() as tx: + existing = await tx.litellm_accessgrouptable.find_unique( + where={"access_group_name": data.access_group_name} + ) + if existing is not None: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Access group '{data.access_group_name}' already exists", + ) + + record = await tx.litellm_accessgrouptable.create( + data={ + "access_group_name": data.access_group_name, + "description": data.description, + "access_model_names": data.access_model_names or [], + "access_mcp_server_ids": data.access_mcp_server_ids or [], + "access_agent_ids": data.access_agent_ids or [], + "assigned_team_ids": data.assigned_team_ids or [], + "assigned_key_ids": data.assigned_key_ids or [], + "created_by": user_api_key_dict.user_id, + "updated_by": user_api_key_dict.user_id, + } + ) + + # Sync team and key tables to reference the new access group + await _sync_add_access_group_to_teams( + tx, data.assigned_team_ids or [], record.access_group_id + ) + await _sync_add_access_group_to_keys( + tx, data.assigned_key_ids or [], record.access_group_id + ) + except HTTPException: + raise except Exception as e: # Race condition: another request created the same name between find_unique and create. - # Prisma raises UniqueViolationError (P2002) or similar for unique constraint. if "unique constraint" in str(e).lower() or "P2002" in str(e): raise HTTPException( status_code=status.HTTP_409_CONFLICT, @@ -139,8 +326,15 @@ async def create_access_group( ) raise - # Cache the newly created access group for read-heavy access patterns + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + await _cache_access_group_record(record) + await _patch_team_caches_add_access_group( + data.assigned_team_ids or [], record.access_group_id, user_api_key_cache, proxy_logging_obj + ) + await _patch_key_caches_add_access_group( + data.assigned_key_ids or [], record.access_group_id, user_api_key_cache, proxy_logging_obj + ) return _record_to_response(record) @@ -204,15 +398,35 @@ async def update_access_group( detail=f"Access group '{access_group_id}' not found", ) + # Compute team/key assignment deltas before the transaction + update_fields = data.model_dump(exclude_unset=True) + old_team_ids: Set[str] = set(existing.assigned_team_ids or []) + old_key_ids: Set[str] = set(existing.assigned_key_ids or []) + new_team_ids: Set[str] = set(update_fields["assigned_team_ids"]) if "assigned_team_ids" in update_fields else old_team_ids + new_key_ids: Set[str] = set(update_fields["assigned_key_ids"]) if "assigned_key_ids" in update_fields else old_key_ids + + teams_to_add = list(new_team_ids - old_team_ids) + teams_to_remove = list(old_team_ids - new_team_ids) + keys_to_add = list(new_key_ids - old_key_ids) + keys_to_remove = list(old_key_ids - new_key_ids) + update_data: dict = {"updated_by": user_api_key_dict.user_id} - for field, value in data.model_dump(exclude_unset=True).items(): + for field, value in update_fields.items(): update_data[field] = value try: - record = await prisma_client.db.litellm_accessgrouptable.update( - where={"access_group_id": access_group_id}, - data=update_data, - ) + async with prisma_client.db.tx() as tx: + record = await tx.litellm_accessgrouptable.update( + where={"access_group_id": access_group_id}, + data=update_data, + ) + + await _sync_add_access_group_to_teams(tx, teams_to_add, access_group_id) + await _sync_remove_access_group_from_teams(tx, teams_to_remove, access_group_id) + await _sync_add_access_group_to_keys(tx, keys_to_add, access_group_id) + await _sync_remove_access_group_from_keys(tx, keys_to_remove, access_group_id) + except HTTPException: + raise except Exception as e: # Unique constraint violation (e.g. access_group_name already exists). if "unique constraint" in str(e).lower() or "P2002" in str(e): @@ -222,8 +436,13 @@ async def update_access_group( ) raise - # Write the updated record into cache (same key, overwrites stale entry) + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + await _cache_access_group_record(record) + await _patch_team_caches_add_access_group(teams_to_add, access_group_id, user_api_key_cache, proxy_logging_obj) + await _patch_team_caches_remove_access_group(teams_to_remove, access_group_id, user_api_key_cache, proxy_logging_obj) + await _patch_key_caches_add_access_group(keys_to_add, access_group_id, user_api_key_cache, proxy_logging_obj) + await _patch_key_caches_remove_access_group(keys_to_remove, access_group_id, user_api_key_cache, proxy_logging_obj) return _record_to_response(record) @@ -240,9 +459,8 @@ async def delete_access_group( prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) try: - # Track affected team IDs and key tokens for cache invalidation - affected_team_ids: list = [] - affected_key_tokens: list = [] + affected_team_ids: List[str] = [] + affected_key_tokens: List[str] = [] async with prisma_client.db.tx() as tx: existing = await tx.litellm_accessgrouptable.find_unique( @@ -254,73 +472,44 @@ async def delete_access_group( detail=f"Access group '{access_group_id}' not found", ) - # Remove access_group_id from teams and keys that reference it + # Union of: teams that have this access_group_id in their own access_group_ids + # AND teams listed in assigned_team_ids (handles out-of-sync data from before this sync was added) teams_with_group = await tx.litellm_teamtable.find_many( where={"access_group_ids": {"hasSome": [access_group_id]}} ) - for team in teams_with_group: - affected_team_ids.append(team.team_id) - updated_ids = [tid for tid in (team.access_group_ids or []) if tid != access_group_id] - await tx.litellm_teamtable.update( - where={"team_id": team.team_id}, - data={"access_group_ids": updated_ids}, - ) + all_affected_team_ids: Set[str] = ( + {team.team_id for team in teams_with_group} + | set(existing.assigned_team_ids or []) + ) + affected_team_ids = list(all_affected_team_ids) + # Union of: keys that have this access_group_id in their own access_group_ids + # AND keys listed in assigned_key_ids (handles out-of-sync data) keys_with_group = await tx.litellm_verificationtoken.find_many( where={"access_group_ids": {"hasSome": [access_group_id]}} ) - for key in keys_with_group: - affected_key_tokens.append(key.token) - updated_ids = [kid for kid in (key.access_group_ids or []) if kid != access_group_id] - await tx.litellm_verificationtoken.update( - where={"token": key.token}, - data={"access_group_ids": updated_ids}, - ) + all_affected_key_tokens: Set[str] = ( + {key.token for key in keys_with_group} + | set(existing.assigned_key_ids or []) + ) + affected_key_tokens = list(all_affected_key_tokens) + + await _sync_remove_access_group_from_teams(tx, affected_team_ids, access_group_id) + await _sync_remove_access_group_from_keys(tx, affected_key_tokens, access_group_id) await tx.litellm_accessgrouptable.delete( where={"access_group_id": access_group_id} ) - # Invalidate the deleted access group from cache - await _invalidate_cache_access_group(access_group_id) - - # Patch cached team and key objects to remove the deleted access_group_id - # instead of fully invalidating them (keeps cache warm, avoids DB re-fetch) from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache - for team_id in affected_team_ids: - cached_team = await _get_team_object_from_cache( - key="team_id:{}".format(team_id), - proxy_logging_obj=proxy_logging_obj, - user_api_key_cache=user_api_key_cache, - parent_otel_span=None, - ) - if cached_team is not None and cached_team.access_group_ids: - cached_team.access_group_ids = [ - ag_id for ag_id in cached_team.access_group_ids if ag_id != access_group_id - ] - await _cache_team_object( - team_id=team_id, - team_table=cached_team, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - for token in affected_key_tokens: - cached_key = await user_api_key_cache.async_get_cache(key=token) - if cached_key is not None: - if isinstance(cached_key, dict): - cached_key = UserAPIKeyAuth(**cached_key) - if isinstance(cached_key, UserAPIKeyAuth) and cached_key.access_group_ids: - cached_key.access_group_ids = [ - ag_id for ag_id in cached_key.access_group_ids if ag_id != access_group_id - ] - await _cache_key_object( - hashed_token=token, - user_api_key_obj=cached_key, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + await _invalidate_cache_access_group(access_group_id) + await _patch_team_caches_remove_access_group( + affected_team_ids, access_group_id, user_api_key_cache, proxy_logging_obj + ) + await _patch_key_caches_remove_access_group( + affected_key_tokens, access_group_id, user_api_key_cache, proxy_logging_obj + ) except HTTPException: raise diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py index 9b6e063176..a8842f7448 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py @@ -99,10 +99,12 @@ def client_and_mocks(monkeypatch): mock_team_table = MagicMock() mock_team_table.find_many = AsyncMock(return_value=[]) + mock_team_table.find_unique = AsyncMock(return_value=None) mock_team_table.update = AsyncMock(return_value=None) mock_key_table = MagicMock() mock_key_table.find_many = AsyncMock(return_value=[]) + mock_key_table.find_unique = AsyncMock(return_value=None) mock_key_table.update = AsyncMock(return_value=None) @asynccontextmanager @@ -570,11 +572,13 @@ def test_delete_access_group_cleans_up_teams_and_keys(client_and_mocks): team_with_group.team_id = "team-1" team_with_group.access_group_ids = ["ag-to-delete", "ag-other"] mock_team_table.find_many = AsyncMock(return_value=[team_with_group]) + mock_team_table.find_unique = AsyncMock(return_value=team_with_group) key_with_group = MagicMock() key_with_group.token = "key-token-1" key_with_group.access_group_ids = ["ag-to-delete"] mock_key_table.find_many = AsyncMock(return_value=[key_with_group]) + mock_key_table.find_unique = AsyncMock(return_value=key_with_group) resp = client.delete("/v1/access_group/ag-to-delete") assert resp.status_code == 204 @@ -669,11 +673,13 @@ def test_delete_access_group_patches_cached_team_and_key( team_with_group.team_id = "team-1" team_with_group.access_group_ids = ["ag-to-delete", "ag-keep"] mock_team_table.find_many = AsyncMock(return_value=[team_with_group]) + mock_team_table.find_unique = AsyncMock(return_value=team_with_group) key_with_group = MagicMock() key_with_group.token = "hashed-key-1" key_with_group.access_group_ids = ["ag-to-delete"] mock_key_table.find_many = AsyncMock(return_value=[key_with_group]) + mock_key_table.find_unique = AsyncMock(return_value=key_with_group) # Build cached team object (returned from proxy_logging dual cache) if team_cache_group_ids is not None: @@ -762,6 +768,7 @@ def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks): key_with_group.token = "hashed-key-dict" key_with_group.access_group_ids = ["ag-to-delete", "ag-other"] mock_key_table.find_many = AsyncMock(return_value=[key_with_group]) + mock_key_table.find_unique = AsyncMock(return_value=key_with_group) # No team in cache mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock( @@ -882,3 +889,280 @@ def test_record_to_access_group_table(): assert result.access_group_name == "unit-test-group" assert result.access_model_names == ["gpt-4", "claude-3"] assert result.access_agent_ids == ["agent-1"] + + +# --------------------------------------------------------------------------- +# Sync tests: CREATE +# --------------------------------------------------------------------------- + + +def test_create_access_group_syncs_assigned_teams(client_and_mocks): + """Create adds access_group_id to each assigned team's access_group_ids in DB.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + team_record = MagicMock() + team_record.team_id = "team-1" + team_record.access_group_ids = [] + mock_team_table.find_unique = AsyncMock(return_value=team_record) + + resp = client.post( + "/v1/access_group", + json={"access_group_name": "new-group", "assigned_team_ids": ["team-1"]}, + ) + assert resp.status_code == 201 + + mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-1"}) + mock_team_table.update.assert_awaited_once() + call_kwargs = mock_team_table.update.call_args.kwargs + assert call_kwargs["where"] == {"team_id": "team-1"} + # The newly created access group id ("ag-new") should be in the updated list + assert "ag-new" in call_kwargs["data"]["access_group_ids"] + + +def test_create_access_group_syncs_assigned_keys(client_and_mocks): + """Create adds access_group_id to each assigned key's access_group_ids in DB.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_key_table = mock_prisma.db.litellm_verificationtoken + + key_record = MagicMock() + key_record.token = "hashed-token-1" + key_record.access_group_ids = [] + mock_key_table.find_unique = AsyncMock(return_value=key_record) + + resp = client.post( + "/v1/access_group", + json={"access_group_name": "new-group", "assigned_key_ids": ["hashed-token-1"]}, + ) + assert resp.status_code == 201 + + mock_key_table.find_unique.assert_awaited_once_with(where={"token": "hashed-token-1"}) + mock_key_table.update.assert_awaited_once() + call_kwargs = mock_key_table.update.call_args.kwargs + assert call_kwargs["where"] == {"token": "hashed-token-1"} + assert "ag-new" in call_kwargs["data"]["access_group_ids"] + + +def test_create_access_group_skips_sync_for_nonexistent_team(client_and_mocks): + """Create skips updating a team that doesn't exist in DB.""" + client, mock_prisma, _, mock_cache, mock_proxy_logging = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + mock_team_table.find_unique = AsyncMock(return_value=None) + + resp = client.post( + "/v1/access_group", + json={"access_group_name": "new-group", "assigned_team_ids": ["nonexistent-team"]}, + ) + assert resp.status_code == 201 + mock_team_table.update.assert_not_awaited() + + +def test_create_access_group_idempotent_team_sync(client_and_mocks): + """Create skips updating a team that already has the access_group_id.""" + client, mock_prisma, _, mock_cache, mock_proxy_logging = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + team_record = MagicMock() + team_record.team_id = "team-1" + team_record.access_group_ids = ["ag-new"] # already synced + mock_team_table.find_unique = AsyncMock(return_value=team_record) + + resp = client.post( + "/v1/access_group", + json={"access_group_name": "new-group", "assigned_team_ids": ["team-1"]}, + ) + assert resp.status_code == 201 + mock_team_table.update.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# Sync tests: UPDATE +# --------------------------------------------------------------------------- + + +def test_update_access_group_syncs_added_teams(client_and_mocks): + """Update adds access_group_id to newly assigned teams.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + existing = _make_access_group_record( + access_group_id="ag-update", assigned_team_ids=["team-existing"] + ) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + team_record = MagicMock() + team_record.team_id = "team-new" + team_record.access_group_ids = [] + mock_team_table.find_unique = AsyncMock(return_value=team_record) + + resp = client.put( + "/v1/access_group/ag-update", + json={"assigned_team_ids": ["team-existing", "team-new"]}, + ) + assert resp.status_code == 200 + + mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-new"}) + mock_team_table.update.assert_awaited_once() + call_kwargs = mock_team_table.update.call_args.kwargs + assert call_kwargs["where"] == {"team_id": "team-new"} + assert "ag-update" in call_kwargs["data"]["access_group_ids"] + + +def test_update_access_group_syncs_removed_teams(client_and_mocks): + """Update removes access_group_id from de-assigned teams.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + existing = _make_access_group_record( + access_group_id="ag-update", assigned_team_ids=["team-keep", "team-remove"] + ) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + team_to_remove = MagicMock() + team_to_remove.team_id = "team-remove" + team_to_remove.access_group_ids = ["ag-update"] + mock_team_table.find_unique = AsyncMock(return_value=team_to_remove) + + resp = client.put( + "/v1/access_group/ag-update", + json={"assigned_team_ids": ["team-keep"]}, + ) + assert resp.status_code == 200 + + mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-remove"}) + mock_team_table.update.assert_awaited_once() + call_kwargs = mock_team_table.update.call_args.kwargs + assert call_kwargs["where"] == {"team_id": "team-remove"} + assert "ag-update" not in call_kwargs["data"]["access_group_ids"] + + +def test_update_access_group_no_team_sync_when_ids_not_in_payload(client_and_mocks): + """Update does not sync teams when assigned_team_ids is absent from the payload.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + existing = _make_access_group_record( + access_group_id="ag-update", assigned_team_ids=["team-1"] + ) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + resp = client.put("/v1/access_group/ag-update", json={"description": "new desc"}) + assert resp.status_code == 200 + + mock_team_table.find_unique.assert_not_awaited() + mock_team_table.update.assert_not_awaited() + + +def test_update_access_group_syncs_added_keys(client_and_mocks): + """Update adds access_group_id to newly assigned keys.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_key_table = mock_prisma.db.litellm_verificationtoken + + existing = _make_access_group_record( + access_group_id="ag-update", assigned_key_ids=["old-token"] + ) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + key_record = MagicMock() + key_record.token = "new-token" + key_record.access_group_ids = [] + mock_key_table.find_unique = AsyncMock(return_value=key_record) + + resp = client.put( + "/v1/access_group/ag-update", + json={"assigned_key_ids": ["old-token", "new-token"]}, + ) + assert resp.status_code == 200 + + mock_key_table.find_unique.assert_awaited_once_with(where={"token": "new-token"}) + mock_key_table.update.assert_awaited_once() + call_kwargs = mock_key_table.update.call_args.kwargs + assert call_kwargs["where"] == {"token": "new-token"} + assert "ag-update" in call_kwargs["data"]["access_group_ids"] + + +def test_update_access_group_syncs_removed_keys(client_and_mocks): + """Update removes access_group_id from de-assigned keys.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_key_table = mock_prisma.db.litellm_verificationtoken + + existing = _make_access_group_record( + access_group_id="ag-update", assigned_key_ids=["keep-token", "remove-token"] + ) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + key_to_remove = MagicMock() + key_to_remove.token = "remove-token" + key_to_remove.access_group_ids = ["ag-update"] + mock_key_table.find_unique = AsyncMock(return_value=key_to_remove) + + resp = client.put( + "/v1/access_group/ag-update", + json={"assigned_key_ids": ["keep-token"]}, + ) + assert resp.status_code == 200 + + mock_key_table.find_unique.assert_awaited_once_with(where={"token": "remove-token"}) + mock_key_table.update.assert_awaited_once() + call_kwargs = mock_key_table.update.call_args.kwargs + assert call_kwargs["where"] == {"token": "remove-token"} + assert "ag-update" not in call_kwargs["data"]["access_group_ids"] + + +# --------------------------------------------------------------------------- +# Sync tests: DELETE (out-of-sync data handling) +# --------------------------------------------------------------------------- + + +def test_delete_access_group_handles_out_of_sync_assigned_teams(client_and_mocks): + """Delete includes teams from assigned_team_ids even when not found by hasSome query.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + + # Access group has assigned_team_ids but the team's access_group_ids is not synced + existing = _make_access_group_record( + access_group_id="ag-to-delete", + assigned_team_ids=["team-out-of-sync"], + ) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + # hasSome query finds nothing (team's own access_group_ids is out of sync) + mock_team_table.find_many = AsyncMock(return_value=[]) + + out_of_sync_team = MagicMock() + out_of_sync_team.team_id = "team-out-of-sync" + out_of_sync_team.access_group_ids = [] # already clean, no update needed + mock_team_table.find_unique = AsyncMock(return_value=out_of_sync_team) + + resp = client.delete("/v1/access_group/ag-to-delete") + assert resp.status_code == 204 + + # find_unique is called for the out-of-sync team (included via union with assigned_team_ids) + mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-out-of-sync"}) + # No update needed since team's access_group_ids doesn't contain "ag-to-delete" + mock_team_table.update.assert_not_awaited() + + +def test_delete_access_group_handles_out_of_sync_assigned_keys(client_and_mocks): + """Delete includes keys from assigned_key_ids even when not found by hasSome query.""" + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks + mock_key_table = mock_prisma.db.litellm_verificationtoken + + existing = _make_access_group_record( + access_group_id="ag-to-delete", + assigned_key_ids=["token-out-of-sync"], + ) + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + mock_key_table.find_many = AsyncMock(return_value=[]) + + out_of_sync_key = MagicMock() + out_of_sync_key.token = "token-out-of-sync" + out_of_sync_key.access_group_ids = [] + mock_key_table.find_unique = AsyncMock(return_value=out_of_sync_key) + + resp = client.delete("/v1/access_group/ag-to-delete") + assert resp.status_code == 204 + + mock_key_table.find_unique.assert_awaited_once_with(where={"token": "token-out-of-sync"}) + mock_key_table.update.assert_not_awaited()