diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 87ff4a66e08..0c85e9ba4f2 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -893,6 +893,7 @@ class KeyRequestBase(GenerateRequestBase): Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm router_settings: Optional[UpdateRouterConfig] = None + access_group_ids: Optional[List[str]] = None class LiteLLMKeyType(str, enum.Enum): @@ -1502,6 +1503,7 @@ class TeamBase(LiteLLMPydanticObjectBase): models: list = [] blocked: bool = False router_settings: Optional[dict] = None + access_group_ids: Optional[List[str]] = None class NewTeamRequest(TeamBase): @@ -1589,6 +1591,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): model_tpm_limit: Optional[Dict[str, int]] = None allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None router_settings: Optional[dict] = None + access_group_ids: Optional[List[str]] = None class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase): @@ -2177,6 +2180,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): updated_by: Optional[str] = None object_permission_id: Optional[str] = None object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + access_group_ids: Optional[List[str]] = None rotation_count: Optional[int] = 0 # Number of times key has been rotated auto_rotate: Optional[bool] = False # Whether this key should be auto-rotated rotation_interval: Optional[str] = None # How often to rotate (e.g., "30d", "90d") diff --git a/litellm/proxy/management_endpoints/access_groups_endpoints.py b/litellm/proxy/management_endpoints/access_groups_endpoints.py index d03d5a3588c..ba432dff84b 100644 --- a/litellm/proxy/management_endpoints/access_groups_endpoints.py +++ b/litellm/proxy/management_endpoints/access_groups_endpoints.py @@ -54,7 +54,7 @@ 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_unifiedaccessgroup.find_unique( + existing = await prisma_client.db.litellm_accessgrouptable.find_unique( where={"access_group_name": data.access_group_name} ) if existing is not None: @@ -63,7 +63,7 @@ async def create_access_group( detail=f"Access group '{data.access_group_name}' already exists", ) - record = await prisma_client.db.litellm_unifiedaccessgroup.create( + record = await prisma_client.db.litellm_accessgrouptable.create( data={ "access_group_name": data.access_group_name, "description": data.description, @@ -89,7 +89,7 @@ async def list_access_groups( _require_proxy_admin(user_api_key_dict) prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) - records = await prisma_client.db.litellm_unifiedaccessgroup.find_many( + records = await prisma_client.db.litellm_accessgrouptable.find_many( order={"created_at": "desc"} ) return [_record_to_response(r) for r in records] @@ -106,7 +106,7 @@ async def get_access_group( _require_proxy_admin(user_api_key_dict) prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value) - record = await prisma_client.db.litellm_unifiedaccessgroup.find_unique( + record = await prisma_client.db.litellm_accessgrouptable.find_unique( where={"access_group_id": access_group_id} ) if record is None: @@ -129,7 +129,7 @@ async def update_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_unifiedaccessgroup.find_unique( + existing = await prisma_client.db.litellm_accessgrouptable.find_unique( where={"access_group_id": access_group_id} ) if existing is None: @@ -142,7 +142,7 @@ async def update_access_group( for field, value in data.model_dump(exclude_unset=True).items(): update_data[field] = value - record = await prisma_client.db.litellm_unifiedaccessgroup.update( + record = await prisma_client.db.litellm_accessgrouptable.update( where={"access_group_id": access_group_id}, data=update_data, ) @@ -160,7 +160,7 @@ 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_unifiedaccessgroup.find_unique( + existing = await prisma_client.db.litellm_accessgrouptable.find_unique( where={"access_group_id": access_group_id} ) if existing is None: @@ -169,6 +169,6 @@ async def delete_access_group( detail=f"Access group '{access_group_id}' not found", ) - await prisma_client.db.litellm_unifiedaccessgroup.delete( + await prisma_client.db.litellm_accessgrouptable.delete( where={"access_group_id": access_group_id} ) 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 new file mode 100644 index 00000000000..a0f81559855 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py @@ -0,0 +1,381 @@ +""" +Tests for access group management endpoints. +""" + +import os +import sys +import types +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi.testclient import TestClient + +import litellm.proxy.proxy_server as ps +from litellm.proxy.proxy_server import app +from litellm.proxy._types import ( + CommonProxyErrors, + LitellmUserRoles, + UserAPIKeyAuth, +) + +sys.path.insert(0, os.path.abspath("../../../")) + + +def _make_access_group_record( + access_group_id: str = "ag-123", + access_group_name: str = "test-group", + description: str | None = "Test description", + access_model_ids: list | None = None, + access_mcp_server_ids: list | None = None, + access_agent_ids: list | None = None, + assigned_team_ids: list | None = None, + assigned_key_ids: list | None = None, + created_by: str | None = "admin-user", + updated_by: str | None = "admin-user", +): + record = MagicMock() + record.access_group_id = access_group_id + record.access_group_name = access_group_name + record.description = description + record.access_model_ids = access_model_ids or [] + record.access_mcp_server_ids = access_mcp_server_ids or [] + 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_by = created_by + record.updated_at = datetime.now() + record.updated_by = updated_by + return record + + +@pytest.fixture +def client_and_mocks(monkeypatch): + """Setup mock prisma and admin auth for access group endpoints.""" + mock_access_group_table = MagicMock() + mock_prisma = MagicMock() + + def _create_side_effect(*, data): + return _make_access_group_record( + access_group_id="ag-new", + access_group_name=data.get("access_group_name", "new"), + description=data.get("description"), + access_model_ids=data.get("access_model_ids", []), + access_mcp_server_ids=data.get("access_mcp_server_ids", []), + access_agent_ids=data.get("access_agent_ids", []), + assigned_team_ids=data.get("assigned_team_ids", []), + assigned_key_ids=data.get("assigned_key_ids", []), + created_by=data.get("created_by"), + updated_by=data.get("updated_by"), + ) + + mock_access_group_table.create = AsyncMock(side_effect=_create_side_effect) + mock_access_group_table.find_unique = AsyncMock(return_value=None) + mock_access_group_table.find_many = AsyncMock(return_value=[]) + mock_access_group_table.update = AsyncMock(side_effect=lambda *, where, data: _make_access_group_record( + access_group_id=where.get("access_group_id", "ag-123"), + access_group_name=data.get("access_group_name", "updated"), + description=data.get("description"), + access_model_ids=data.get("access_model_ids", []), + access_mcp_server_ids=data.get("access_mcp_server_ids", []), + access_agent_ids=data.get("access_agent_ids", []), + assigned_team_ids=data.get("assigned_team_ids", []), + assigned_key_ids=data.get("assigned_key_ids", []), + updated_by=data.get("updated_by"), + )) + mock_access_group_table.delete = AsyncMock(return_value=None) + + mock_prisma.db = types.SimpleNamespace( + litellm_accessgrouptable=mock_access_group_table, + ) + + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + + admin_user = UserAPIKeyAuth( + user_id="admin_user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: admin_user + + client = TestClient(app) + + yield client, mock_prisma, mock_access_group_table + + app.dependency_overrides.clear() + monkeypatch.setattr(ps, "prisma_client", ps.prisma_client) + + +# --------------------------------------------------------------------------- +# CREATE +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "payload", + [ + {"access_group_name": "group-a"}, + { + "access_group_name": "group-b", + "description": "Group B description", + "access_model_ids": ["model-1"], + "access_mcp_server_ids": ["mcp-1"], + "assigned_team_ids": ["team-1"], + }, + ], +) +def test_create_access_group_success(client_and_mocks, payload): + """Create access group with various payloads returns 201.""" + client, _, mock_table = client_and_mocks + + resp = client.post("/v1/access_group", json=payload) + assert resp.status_code == 201 + body = resp.json() + assert body["access_group_name"] == payload["access_group_name"] + assert body.get("access_group_id") is not None + mock_table.create.assert_awaited_once() + + +def test_create_access_group_duplicate_name_conflict(client_and_mocks): + """Create with duplicate name returns 409.""" + client, _, mock_table = client_and_mocks + + existing = _make_access_group_record(access_group_name="existing-group") + mock_table.find_unique = AsyncMock(return_value=existing) + + resp = client.post("/v1/access_group", json={"access_group_name": "existing-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.""" + client, _, _ = client_and_mocks + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="regular_user", + user_role=user_role, + ) + + resp = client.post("/v1/access_group", json={"access_group_name": "forbidden"}) + assert resp.status_code == 403 + assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value + + +# --------------------------------------------------------------------------- +# LIST +# --------------------------------------------------------------------------- + + +def test_list_access_groups_success_empty(client_and_mocks): + """List access groups returns empty list when none exist.""" + client, _, mock_table = client_and_mocks + + resp = client.get("/v1/access_group") + assert resp.status_code == 200 + assert resp.json() == [] + mock_table.find_many.assert_awaited_once() + + +def test_list_access_groups_success_with_items(client_and_mocks): + """List access groups returns items when they exist.""" + client, _, mock_table = client_and_mocks + + records = [ + _make_access_group_record(access_group_id="ag-1", access_group_name="group-1"), + _make_access_group_record(access_group_id="ag-2", access_group_name="group-2"), + ] + mock_table.find_many = AsyncMock(return_value=records) + + resp = client.get("/v1/access_group") + assert resp.status_code == 200 + body = resp.json() + assert len(body) == 2 + assert body[0]["access_group_name"] == "group-1" + assert body[1]["access_group_name"] == "group-2" + + +@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.""" + client, _, _ = client_and_mocks + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="regular_user", + user_role=user_role, + ) + + resp = client.get("/v1/access_group") + assert resp.status_code == 403 + assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value + + +# --------------------------------------------------------------------------- +# GET +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("access_group_id", ["ag-123", "ag-other-id"]) +def test_get_access_group_success(client_and_mocks, access_group_id): + """Get access group by id returns record when found.""" + client, _, mock_table = client_and_mocks + + record = _make_access_group_record(access_group_id=access_group_id) + mock_table.find_unique = AsyncMock(return_value=record) + + resp = client.get(f"/v1/access_group/{access_group_id}") + assert resp.status_code == 200 + assert resp.json()["access_group_id"] == access_group_id + + +def test_get_access_group_not_found(client_and_mocks): + """Get access group returns 404 when not found.""" + client, _, mock_table = client_and_mocks + + mock_table.find_unique = AsyncMock(return_value=None) + + resp = client.get("/v1/access_group/nonexistent-id") + assert resp.status_code == 404 + assert "not found" in resp.json()["detail"] + + +@pytest.mark.parametrize("user_role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY]) +def test_get_access_group_forbidden_non_admin(client_and_mocks, user_role): + """Non-admin users cannot get access group.""" + client, _, _ = client_and_mocks + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="regular_user", + user_role=user_role, + ) + + resp = client.get("/v1/access_group/ag-123") + assert resp.status_code == 403 + assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value + + +# --------------------------------------------------------------------------- +# UPDATE +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "update_payload", + [ + {"description": "Updated description"}, + {"access_model_ids": ["model-1", "model-2"]}, + {"assigned_team_ids": [], "assigned_key_ids": ["key-1"]}, + ], +) +def test_update_access_group_success(client_and_mocks, update_payload): + """Update access group with various payloads returns 200.""" + client, _, mock_table = client_and_mocks + + existing = _make_access_group_record(access_group_id="ag-update") + mock_table.find_unique = AsyncMock(return_value=existing) + + resp = client.put("/v1/access_group/ag-update", json=update_payload) + assert resp.status_code == 200 + mock_table.update.assert_awaited_once() + + +def test_update_access_group_not_found(client_and_mocks): + """Update access group returns 404 when not found.""" + client, _, mock_table = client_and_mocks + + mock_table.find_unique = AsyncMock(return_value=None) + + resp = client.put( + "/v1/access_group/nonexistent-id", + json={"description": "Updated"}, + ) + assert resp.status_code == 404 + assert "not found" in resp.json()["detail"] + mock_table.update.assert_not_awaited() + + +@pytest.mark.parametrize("user_role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY]) +def test_update_access_group_forbidden_non_admin(client_and_mocks, user_role): + """Non-admin users cannot update access groups.""" + client, _, _ = client_and_mocks + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="regular_user", + user_role=user_role, + ) + + resp = client.put("/v1/access_group/ag-123", json={"description": "Updated"}) + assert resp.status_code == 403 + assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value + + +# --------------------------------------------------------------------------- +# DELETE +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("access_group_id", ["ag-123", "ag-delete-me"]) +def test_delete_access_group_success(client_and_mocks, access_group_id): + """Delete access group returns 204 when found.""" + client, _, mock_table = client_and_mocks + + existing = _make_access_group_record(access_group_id=access_group_id) + mock_table.find_unique = AsyncMock(return_value=existing) + + resp = client.delete(f"/v1/access_group/{access_group_id}") + assert resp.status_code == 204 + mock_table.delete.assert_awaited_once() + + +def test_delete_access_group_not_found(client_and_mocks): + """Delete access group returns 404 when not found.""" + client, _, mock_table = client_and_mocks + + mock_table.find_unique = AsyncMock(return_value=None) + + resp = client.delete("/v1/access_group/nonexistent-id") + assert resp.status_code == 404 + assert "not found" in resp.json()["detail"] + mock_table.delete.assert_not_awaited() + + +@pytest.mark.parametrize("user_role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY]) +def test_delete_access_group_forbidden_non_admin(client_and_mocks, user_role): + """Non-admin users cannot delete access groups.""" + client, _, _ = client_and_mocks + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="regular_user", + user_role=user_role, + ) + + resp = client.delete("/v1/access_group/ag-123") + assert resp.status_code == 403 + assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value + + +# --------------------------------------------------------------------------- +# DB NOT CONNECTED +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "method,url,factory", + [ + ("post", "/v1/access_group", lambda: {"json": {"access_group_name": "test"}}), + ("get", "/v1/access_group", lambda: {}), + ("get", "/v1/access_group/ag-123", lambda: {}), + ("put", "/v1/access_group/ag-123", lambda: {"json": {"description": "x"}}), + ("delete", "/v1/access_group/ag-123", lambda: {}), + ], +) +def test_access_group_endpoints_db_not_connected(client_and_mocks, monkeypatch, method, url, factory): + """All endpoints return 500 when DB is not connected.""" + client, _, _ = client_and_mocks + + monkeypatch.setattr(ps, "prisma_client", None) + + resp = getattr(client, method)(url, **factory()) + assert resp.status_code == 500 + assert resp.json()["detail"]["error"] == CommonProxyErrors.db_not_connected_error.value