diff --git a/tests/litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/litellm/proxy/management_endpoints/test_team_endpoints.py index d3f3c11593..ec47196ae0 100644 --- a/tests/litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/litellm/proxy/management_endpoints/test_team_endpoints.py @@ -59,40 +59,47 @@ async def test_get_team_permissions_list_success(mock_db_client, mock_admin_auth Test successful retrieval of team member permissions. """ test_team_id = "test-team-123" + permissions = ["/key/generate", "/key/update"] mock_team_data = { "team_id": test_team_id, "team_alias": "Test Team", - "team_member_permissions": ["/key/generate", "/key/update"], + "team_member_permissions": permissions, "spend": 0.0, } mock_team_row = MagicMock() mock_team_row.model_dump.return_value = mock_team_data - mock_db_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_team_row - ) - # Override the dependency for this test - app.dependency_overrides[user_api_key_auth] = lambda: mock_admin_auth + # Set attributes directly on the mock object + mock_team_row.team_id = test_team_id + mock_team_row.team_alias = "Test Team" + mock_team_row.team_member_permissions = permissions + mock_team_row.spend = 0.0 - response = client.get(f"/team/permissions_list?team_id={test_team_id}") + # Mock the get_team_object function used in the endpoint + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_row, + ): + # Override the dependency for this test + app.dependency_overrides[user_api_key_auth] = lambda: mock_admin_auth - assert response.status_code == 200 - response_data = response.json() - assert response_data["team_id"] == test_team_id - assert ( - response_data["team_member_permissions"] - == mock_team_data["team_member_permissions"] - ) - assert ( - response_data["all_available_permissions"] - == TeamMemberPermissionChecks.get_all_available_team_member_permissions() - ) - mock_db_client.db.litellm_teamtable.find_unique.assert_awaited_once_with( - where={"team_id": test_team_id} - ) + response = client.get(f"/team/permissions_list?team_id={test_team_id}") - # Clean up dependency override - app.dependency_overrides = {} + assert response.status_code == 200 + response_data = response.json() + assert response_data["team_id"] == test_team_id + assert ( + response_data["team_member_permissions"] + == mock_team_data["team_member_permissions"] + ) + assert ( + response_data["all_available_permissions"] + == TeamMemberPermissionChecks.get_all_available_team_member_permissions() + ) + + # Clean up dependency override + app.dependency_overrides = {} # Test for /team/permissions_update endpoint (POST) @@ -102,15 +109,17 @@ async def test_update_team_permissions_success(mock_db_client, mock_admin_auth): Test successful update of team member permissions by an admin. """ test_team_id = "test-team-456" + update_permissions = ["/key/generate", "/key/update"] update_payload = { "team_id": test_team_id, - "team_member_permissions": ["/key/generate", "/key/update"], + "team_member_permissions": update_permissions, } + existing_permissions = ["/key/list"] mock_existing_team_data = { "team_id": test_team_id, "team_alias": "Existing Team", - "team_member_permissions": ["/key/list"], + "team_member_permissions": existing_permissions, "spend": 0.0, "models": [], } @@ -121,41 +130,50 @@ async def test_update_team_permissions_success(mock_db_client, mock_admin_auth): mock_existing_team_row = MagicMock(spec=LiteLLM_TeamTable) mock_existing_team_row.model_dump.return_value = mock_existing_team_data - # Set attributes directly if model_dump isn't enough for LiteLLM_TeamTable usage - for key, value in mock_existing_team_data.items(): - setattr(mock_existing_team_row, key, value) + + # Set attributes directly on the existing team mock + mock_existing_team_row.team_id = test_team_id + mock_existing_team_row.team_alias = "Existing Team" + mock_existing_team_row.team_member_permissions = existing_permissions + mock_existing_team_row.spend = 0.0 + mock_existing_team_row.models = [] mock_updated_team_row = MagicMock(spec=LiteLLM_TeamTable) mock_updated_team_row.model_dump.return_value = mock_updated_team_data - # Set attributes directly if model_dump isn't enough for LiteLLM_TeamTable usage - for key, value in mock_updated_team_data.items(): - setattr(mock_updated_team_row, key, value) - mock_db_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=mock_existing_team_row - ) - mock_db_client.db.litellm_teamtable.update = AsyncMock( - return_value=mock_updated_team_row - ) + # Set attributes directly on the updated team mock + mock_updated_team_row.team_id = test_team_id + mock_updated_team_row.team_alias = "Existing Team" + mock_updated_team_row.team_member_permissions = update_permissions + mock_updated_team_row.spend = 0.0 + mock_updated_team_row.models = [] - # Override the dependency for this test - app.dependency_overrides[user_api_key_auth] = lambda: mock_admin_auth + # Mock the get_team_object function used in the endpoint + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + new_callable=AsyncMock, + return_value=mock_existing_team_row, + ): + # Mock the database update function + mock_db_client.db.litellm_teamtable.update = AsyncMock( + return_value=mock_updated_team_row + ) - response = client.post("/team/permissions_update", json=update_payload) + # Override the dependency for this test + app.dependency_overrides[user_api_key_auth] = lambda: mock_admin_auth - assert response.status_code == 200 - response_data = response.json() + response = client.post("/team/permissions_update", json=update_payload) - # Use model_dump for comparison if the endpoint returns the Prisma model directly - assert response_data == mock_updated_team_row.model_dump() + assert response.status_code == 200 + response_data = response.json() - mock_db_client.db.litellm_teamtable.find_unique.assert_awaited_once_with( - where={"team_id": test_team_id} - ) - mock_db_client.db.litellm_teamtable.update.assert_awaited_once_with( - where={"team_id": test_team_id}, - data={"team_member_permissions": update_payload["team_member_permissions"]}, - ) + # Use model_dump for comparison if the endpoint returns the Prisma model directly + assert response_data == mock_updated_team_row.model_dump() - # Clean up dependency override - app.dependency_overrides = {} + mock_db_client.db.litellm_teamtable.update.assert_awaited_once_with( + where={"team_id": test_team_id}, + data={"team_member_permissions": update_payload["team_member_permissions"]}, + ) + + # Clean up dependency override + app.dependency_overrides = {}