From ea8c89ea3f059ddf57f742d4f74ca3d9d2d83d64 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 12 Feb 2026 14:05:51 -0800 Subject: [PATCH] rename file and add tests --- ...endpoints.py => access_group_endpoints.py} | 105 +++++++--- litellm/proxy/proxy_server.py | 2 +- .../test_access_group_endpoints.py | 190 +++++++++++++++++- 3 files changed, 270 insertions(+), 27 deletions(-) rename litellm/proxy/management_endpoints/{access_groups_endpoints.py => access_group_endpoints.py} (60%) diff --git a/litellm/proxy/management_endpoints/access_groups_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py similarity index 60% rename from litellm/proxy/management_endpoints/access_groups_endpoints.py rename to litellm/proxy/management_endpoints/access_group_endpoints.py index 6ce401a5ee..33b81e8565 100644 --- a/litellm/proxy/management_endpoints/access_groups_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -2,8 +2,10 @@ from typing import List from fastapi import APIRouter, Depends, HTTPException, status +from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.utils import get_prisma_client_or_throw from litellm.types.access_group import ( AccessGroupCreateRequest, @@ -62,19 +64,29 @@ async def create_access_group( detail=f"Access group '{data.access_group_name}' already exists", ) - record = await prisma_client.db.litellm_accessgrouptable.create( - data={ - "access_group_name": data.access_group_name, - "description": data.description, - "access_model_ids": data.access_model_ids 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, - } - ) + try: + record = await prisma_client.db.litellm_accessgrouptable.create( + data={ + "access_group_name": data.access_group_name, + "description": data.description, + "access_model_ids": data.access_model_ids 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, + } + ) + 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, + detail=f"Access group '{data.access_group_name}' already exists", + ) + raise return _record_to_response(record) @@ -159,18 +171,63 @@ async def delete_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_id": access_group_id} - ) - if existing is None: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail=f"Access group '{access_group_id}' not found", - ) + try: + async with prisma_client.db.tx() as tx: + existing = await tx.litellm_accessgrouptable.find_unique( + where={"access_group_id": access_group_id} + ) + if existing is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Access group '{access_group_id}' not found", + ) - await prisma_client.db.litellm_accessgrouptable.delete( - where={"access_group_id": access_group_id} - ) + # Remove access_group_id from teams and keys that reference it + teams_with_group = await tx.litellm_teamtable.find_many( + where={"access_group_ids": {"hasSome": [access_group_id]}} + ) + for team in teams_with_group: + 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}, + ) + + keys_with_group = await tx.litellm_verificationtoken.find_many( + where={"access_group_ids": {"hasSome": [access_group_id]}} + ) + for key in keys_with_group: + 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}, + ) + + await tx.litellm_accessgrouptable.delete( + where={"access_group_id": access_group_id} + ) + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception( + "delete_access_group failed: access_group_id=%s error=%s", + access_group_id, + e, + ) + if PrismaDBExceptionHandler.is_database_connection_error(e): + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=CommonProxyErrors.db_not_connected_error.value, + ) + if "P2025" in str(e) or ("record" in str(e).lower() and "not found" in str(e).lower()): + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Access group '{access_group_id}' not found", + ) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Failed to delete access group. Please try again.", + ) # Alias routes for /v1/unified_access_group diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 436e7eaa4e..2ec45cf139 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -393,7 +393,7 @@ from litellm.proxy.management_endpoints.tag_management_endpoints import ( from litellm.proxy.management_endpoints.team_callback_endpoints import ( router as team_callback_router, ) -from litellm.proxy.management_endpoints.access_groups_endpoints import ( +from litellm.proxy.management_endpoints.access_group_endpoints import ( router as access_group_router, ) from litellm.proxy.management_endpoints.team_endpoints import router as team_router 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 ddcc78077b..5f204918d0 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 @@ -5,11 +5,13 @@ Tests for access group management endpoints. import os import sys import types +from contextlib import asynccontextmanager from datetime import datetime from unittest.mock import AsyncMock, MagicMock import pytest from fastapi.testclient import TestClient +from prisma.errors import PrismaError import litellm.proxy.proxy_server as ps from litellm.proxy.proxy_server import app @@ -33,6 +35,7 @@ def _make_access_group_record( assigned_key_ids: list | None = None, created_by: str | None = "admin-user", updated_by: str | None = "admin-user", + created_at: datetime | None = None, ): record = MagicMock() record.access_group_id = access_group_id @@ -43,7 +46,7 @@ def _make_access_group_record( record.access_agent_ids = access_agent_ids or [] record.assigned_team_ids = assigned_team_ids or [] record.assigned_key_ids = assigned_key_ids or [] - record.created_at = datetime.now() + record.created_at = created_at or datetime.now() record.created_by = created_by record.updated_at = datetime.now() record.updated_by = updated_by @@ -86,9 +89,30 @@ def client_and_mocks(monkeypatch): )) mock_access_group_table.delete = AsyncMock(return_value=None) - mock_prisma.db = types.SimpleNamespace( + mock_team_table = MagicMock() + mock_team_table.find_many = AsyncMock(return_value=[]) + mock_team_table.update = AsyncMock(return_value=None) + + mock_key_table = MagicMock() + mock_key_table.find_many = AsyncMock(return_value=[]) + mock_key_table.update = AsyncMock(return_value=None) + + @asynccontextmanager + async def mock_tx(): + tx = types.SimpleNamespace( + litellm_accessgrouptable=mock_access_group_table, + litellm_teamtable=mock_team_table, + litellm_verificationtoken=mock_key_table, + ) + yield tx + + mock_db = types.SimpleNamespace( litellm_accessgrouptable=mock_access_group_table, + litellm_teamtable=mock_team_table, + litellm_verificationtoken=mock_key_table, + tx=mock_tx, ) + mock_prisma.db = mock_db monkeypatch.setattr(ps, "prisma_client", mock_prisma) @@ -153,6 +177,26 @@ def test_create_access_group_duplicate_name_conflict(client_and_mocks): assert "already exists" in resp.json()["detail"] +@pytest.mark.parametrize( + "error_message", + [ + "Unique constraint failed on the fields: (`access_group_name`)", + "P2002: Unique constraint failed", + "unique constraint violation", + ], +) +def test_create_access_group_race_condition_returns_409(client_and_mocks, error_message): + """Create race condition: Prisma unique constraint surfaces as 409, not 500.""" + client, _, mock_table = client_and_mocks + + mock_table.find_unique = AsyncMock(return_value=None) + mock_table.create = AsyncMock(side_effect=Exception(error_message)) + + resp = client.post("/v1/access_group", json={"access_group_name": "race-group"}) + assert resp.status_code == 409 + assert "already exists" in resp.json()["detail"] + + @pytest.mark.parametrize("user_role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY]) def test_create_access_group_forbidden_non_admin(client_and_mocks, user_role): """Non-admin users cannot create access groups.""" @@ -168,6 +212,27 @@ def test_create_access_group_forbidden_non_admin(client_and_mocks, user_role): assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value +def test_create_access_group_validation_missing_name(client_and_mocks): + """Create with missing access_group_name returns 422.""" + client, _, _ = client_and_mocks + + resp = client.post("/v1/access_group", json={}) + assert resp.status_code == 422 + + +def test_create_access_group_500_on_non_constraint_prisma_error(client_and_mocks): + """Create with non-unique-constraint Prisma error returns 500.""" + client, _, mock_table = client_and_mocks + + mock_table.find_unique = AsyncMock(return_value=None) + mock_table.create = AsyncMock(side_effect=Exception("Some other database error")) + + # Use raise_server_exceptions=False so unhandled exceptions become 500 responses + test_client = TestClient(app, raise_server_exceptions=False) + resp = test_client.post("/v1/access_group", json={"access_group_name": "test-group"}) + assert resp.status_code == 500 + + # --------------------------------------------------------------------------- # LIST # --------------------------------------------------------------------------- @@ -203,6 +268,37 @@ def test_list_access_groups_success_with_items(client_and_mocks, base_path): assert body[1]["access_group_name"] == "group-2" +@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS) +def test_list_access_groups_ordered_by_created_at_desc(client_and_mocks, base_path): + """List access groups calls find_many with created_at desc order.""" + client, _, mock_table = client_and_mocks + + older = datetime(2025, 1, 1, 12, 0, 0) + newer = datetime(2025, 1, 2, 12, 0, 0) + records = [ + _make_access_group_record( + access_group_id="ag-newer", + access_group_name="newer-group", + created_at=newer, + ), + _make_access_group_record( + access_group_id="ag-older", + access_group_name="older-group", + created_at=older, + ), + ] + mock_table.find_many = AsyncMock(return_value=records) + + resp = client.get(base_path) + assert resp.status_code == 200 + body = resp.json() + assert len(body) == 2 + # Mock returns newest first (simulating Prisma order desc) + assert body[0]["access_group_name"] == "newer-group" + assert body[1]["access_group_name"] == "older-group" + mock_table.find_many.assert_awaited_once_with(order={"created_at": "desc"}) + + @pytest.mark.parametrize("user_role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY]) def test_list_access_groups_forbidden_non_admin(client_and_mocks, user_role): """Non-admin users cannot list access groups.""" @@ -319,6 +415,22 @@ def test_update_access_group_forbidden_non_admin(client_and_mocks, user_role): assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value +def test_update_access_group_empty_body(client_and_mocks): + """Update with empty body succeeds; only updated_by is set.""" + client, _, mock_table = client_and_mocks + + existing = _make_access_group_record(access_group_id="ag-update", access_group_name="unchanged") + mock_table.find_unique = AsyncMock(return_value=existing) + + resp = client.put("/v1/access_group/ag-update", json={}) + assert resp.status_code == 200 + mock_table.update.assert_awaited_once() + call_kwargs = mock_table.update.call_args.kwargs + assert call_kwargs["where"] == {"access_group_id": "ag-update"} + assert "updated_by" in call_kwargs["data"] + assert call_kwargs["data"]["updated_by"] == "admin_user" + + # --------------------------------------------------------------------------- # DELETE # --------------------------------------------------------------------------- @@ -365,6 +477,80 @@ def test_delete_access_group_forbidden_non_admin(client_and_mocks, user_role): assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value +def test_delete_access_group_cleans_up_teams_and_keys(client_and_mocks): + """Delete removes access_group_id from teams and keys before deleting the group.""" + client, mock_prisma, mock_access_group_table = client_and_mocks + mock_team_table = mock_prisma.db.litellm_teamtable + mock_key_table = mock_prisma.db.litellm_verificationtoken + + existing = _make_access_group_record(access_group_id="ag-to-delete") + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + team_with_group = MagicMock() + 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]) + + 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]) + + resp = client.delete("/v1/access_group/ag-to-delete") + assert resp.status_code == 204 + + mock_team_table.update.assert_awaited_once_with( + where={"team_id": "team-1"}, + data={"access_group_ids": ["ag-other"]}, + ) + mock_key_table.update.assert_awaited_once_with( + where={"token": "key-token-1"}, + data={"access_group_ids": []}, + ) + mock_access_group_table.delete.assert_awaited_once_with( + where={"access_group_id": "ag-to-delete"} + ) + + +def test_delete_access_group_503_on_db_connection_error(client_and_mocks): + """Delete returns 503 when DB connection error occurs during transaction.""" + client, _, mock_table = client_and_mocks + + existing = _make_access_group_record(access_group_id="ag-to-delete") + mock_table.find_unique = AsyncMock(return_value=existing) + mock_table.delete = AsyncMock(side_effect=PrismaError()) + + resp = client.delete("/v1/access_group/ag-to-delete") + assert resp.status_code == 503 + assert resp.json()["detail"] == CommonProxyErrors.db_not_connected_error.value + + +def test_delete_access_group_404_on_p2025_or_record_not_found(client_and_mocks): + """Delete returns 404 when Prisma raises P2025 or record-not-found error.""" + client, _, mock_table = client_and_mocks + + existing = _make_access_group_record(access_group_id="ag-to-delete") + mock_table.find_unique = AsyncMock(return_value=existing) + mock_table.delete = AsyncMock(side_effect=Exception("P2025: Record to delete does not exist")) + + resp = client.delete("/v1/access_group/ag-to-delete") + assert resp.status_code == 404 + assert "not found" in resp.json()["detail"] + + +def test_delete_access_group_500_on_generic_exception(client_and_mocks): + """Delete returns 500 when generic exception occurs during transaction.""" + client, _, mock_table = client_and_mocks + + existing = _make_access_group_record(access_group_id="ag-to-delete") + mock_table.find_unique = AsyncMock(return_value=existing) + mock_table.delete = AsyncMock(side_effect=RuntimeError("Unexpected error")) + + resp = client.delete("/v1/access_group/ag-to-delete") + assert resp.status_code == 500 + assert "Failed to delete access group" in resp.json()["detail"] + + # --------------------------------------------------------------------------- # DB NOT CONNECTED # ---------------------------------------------------------------------------