mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-07 00:24:03 +00:00
[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 <noreply@anthropic.com>
This commit is contained in:
yuneng-jiang
co-authored by
Claude Sonnet 4.6
Claude Sonnet 4.6
parent
8d1c75c48a
commit
516b18feca
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user