mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-21 08:26:34 +00:00
[Fix] SCIM - Add SCIM PATCH and PUT Ops for Users (#11863)
* fix SCIM memberships Patch * fixes for SCIM updates * fixes for SCIM * working provisioning for teams on SCIM * working user patch / PUT ops SCIM * fixes SCIM * test_scim_v2_endpoints.py * handle_existing_user_by_email * fixes for provisioning SCIMUser * fixes SCIM provisioning * test scim v2 * fixes for linting * fix _apply_patch_ops * fixes code QA check for team membership checks
This commit is contained in:
Ishaan Jaff
@@ -1,13 +0,0 @@
|
||||
from fastapi import HTTPException
|
||||
|
||||
|
||||
class ScimUserAlreadyExists(HTTPException):
|
||||
"""
|
||||
Exception raised when a user already exists in the database.
|
||||
"""
|
||||
def __init__(self, message: str, scim_type: str = "uniqueness"):
|
||||
super().__init__(status_code=409, detail=message)
|
||||
self.message = message
|
||||
self.scim_type = scim_type
|
||||
self.schemas = ["urn:ietf:params:scim:api:messages:2.0:Error"]
|
||||
|
||||
@@ -124,6 +124,8 @@ class ScimTransformations:
|
||||
# Get team members
|
||||
scim_members: List[SCIMMember] = []
|
||||
for member in team.members_with_roles or []:
|
||||
if isinstance(member, dict):
|
||||
member = Member(**member)
|
||||
scim_members.append(
|
||||
SCIMMember(
|
||||
value=ScimTransformations._get_scim_member_value(member),
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,20 +0,0 @@
|
||||
from fastapi import HTTPException
|
||||
|
||||
|
||||
async def _check_user_exists(prisma_client, user_name: str) -> bool:
|
||||
"""Check if user already exists by username"""
|
||||
if not user_name:
|
||||
return False
|
||||
|
||||
existing_user = await prisma_client.db.litellm_usertable.find_unique(
|
||||
where={"user_id": user_name}
|
||||
)
|
||||
return existing_user is not None
|
||||
|
||||
|
||||
def _extract_error_message(http_exception: HTTPException) -> str:
|
||||
"""Extract error message from HTTPException detail"""
|
||||
if isinstance(http_exception.detail, dict):
|
||||
return http_exception.detail.get("error", "User already exists")
|
||||
return str(http_exception.detail)
|
||||
|
||||
@@ -22,8 +22,8 @@ class SCIMResource(BaseModel):
|
||||
|
||||
|
||||
class SCIMUserName(BaseModel):
|
||||
familyName: str
|
||||
givenName: str
|
||||
familyName: Optional[str] = None
|
||||
givenName: Optional[str] = None
|
||||
formatted: Optional[str] = None
|
||||
middleName: Optional[str] = None
|
||||
honorificPrefix: Optional[str] = None
|
||||
@@ -43,8 +43,8 @@ class SCIMUserGroup(BaseModel):
|
||||
|
||||
|
||||
class SCIMUser(SCIMResource):
|
||||
userName: str
|
||||
name: SCIMUserName
|
||||
userName: Optional[str] = None
|
||||
name: Optional[SCIMUserName] = None
|
||||
displayName: Optional[str] = None
|
||||
active: bool = True
|
||||
emails: Optional[List[SCIMUserEmail]] = None
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._types import LiteLLM_UserTable
|
||||
from litellm.proxy.management_endpoints.scim.scim_v2 import patch_user
|
||||
from litellm.types.proxy.management_endpoints.scim_v2 import (
|
||||
SCIMPatchOp,
|
||||
SCIMPatchOperation,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_patch_user_updates_fields():
|
||||
mock_user = LiteLLM_UserTable(
|
||||
user_id="user-1",
|
||||
user_email="test@example.com",
|
||||
user_alias="Old",
|
||||
teams=[],
|
||||
metadata={},
|
||||
)
|
||||
|
||||
async def mock_update(*, where, data):
|
||||
if "user_alias" in data:
|
||||
mock_user.user_alias = data["user_alias"]
|
||||
if "metadata" in data:
|
||||
mock_user.metadata = data["metadata"]
|
||||
if "teams" in data:
|
||||
mock_user.teams = data["teams"]
|
||||
if "sso_user_id" in data:
|
||||
mock_user.sso_user_id = data["sso_user_id"]
|
||||
return mock_user
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_db = MagicMock()
|
||||
mock_client.db = mock_db
|
||||
mock_db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user)
|
||||
mock_db.litellm_usertable.update = AsyncMock(side_effect=mock_update)
|
||||
mock_db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
patch_ops = SCIMPatchOp(
|
||||
Operations=[
|
||||
SCIMPatchOperation(op="replace", path="displayName", value="New Name"),
|
||||
SCIMPatchOperation(op="replace", path="active", value="False"),
|
||||
]
|
||||
)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_client):
|
||||
result = await patch_user(user_id="user-1", patch_ops=patch_ops)
|
||||
|
||||
mock_db.litellm_usertable.update.assert_called_once()
|
||||
assert result.displayName == "New Name"
|
||||
assert mock_user.metadata.get("scim_active") is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_patch_user_manages_group_memberships():
|
||||
mock_user = LiteLLM_UserTable(
|
||||
user_id="user-2",
|
||||
user_email="test@example.com",
|
||||
user_alias="Old",
|
||||
teams=["old-team"],
|
||||
metadata={},
|
||||
)
|
||||
|
||||
async def mock_update(*, where, data):
|
||||
if "teams" in data:
|
||||
mock_user.teams = data["teams"]
|
||||
if "metadata" in data:
|
||||
mock_user.metadata = data["metadata"]
|
||||
return mock_user
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_db = MagicMock()
|
||||
mock_client.db = mock_db
|
||||
mock_db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user)
|
||||
mock_db.litellm_usertable.update = AsyncMock(side_effect=mock_update)
|
||||
|
||||
async def mock_add(data, user_api_key_dict):
|
||||
mock_user.teams.append(data.team_id)
|
||||
|
||||
async def mock_delete(data, user_api_key_dict):
|
||||
if data.team_id in mock_user.teams:
|
||||
mock_user.teams.remove(data.team_id)
|
||||
|
||||
patch_ops = SCIMPatchOp(
|
||||
Operations=[
|
||||
SCIMPatchOperation(op="add", path="groups", value=[{"value": "new-team"}]),
|
||||
SCIMPatchOperation(op="remove", path="groups", value=[{"value": "old-team"}]),
|
||||
]
|
||||
)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", mock_client), patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.team_member_add",
|
||||
AsyncMock(side_effect=mock_add),
|
||||
) as mock_add_fn, patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete",
|
||||
AsyncMock(side_effect=mock_delete),
|
||||
) as mock_del_fn:
|
||||
await patch_user(user_id="user-2", patch_ops=patch_ops)
|
||||
|
||||
assert mock_add_fn.called
|
||||
assert mock_del_fn.called
|
||||
assert mock_user.teams == ["new-team"]
|
||||
|
||||
@@ -1,12 +1,23 @@
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import NewUserRequest, ProxyException
|
||||
from litellm.proxy.management_endpoints.scim.scim_errors import ScimUserAlreadyExists
|
||||
from litellm.proxy.management_endpoints.scim.scim_v2 import create_user
|
||||
from litellm.proxy.management_endpoints.scim.scim_v2 import (
|
||||
UserProvisionerHelpers,
|
||||
_handle_team_membership_changes,
|
||||
create_user,
|
||||
patch_user,
|
||||
update_user,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.scim_v2 import (
|
||||
SCIMPatchOp,
|
||||
SCIMPatchOperation,
|
||||
SCIMUser,
|
||||
SCIMUserEmail,
|
||||
SCIMUserGroup,
|
||||
SCIMUserName,
|
||||
)
|
||||
|
||||
@@ -22,21 +33,432 @@ async def test_create_user_existing_user_conflict(mocker):
|
||||
emails=[SCIMUserEmail(value="existing@example.com")],
|
||||
)
|
||||
|
||||
mock_prisma = mocker.MagicMock()
|
||||
# Create a properly structured mock for the prisma client
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
mock_prisma_client.db = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value={"user_id": "existing-user"})
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
# Mock the _get_prisma_client_or_raise_exception to return our mock
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists",
|
||||
AsyncMock(return_value=True),
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
|
||||
AsyncMock(return_value=mock_prisma_client),
|
||||
)
|
||||
|
||||
mocked_new_user = mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.new_user",
|
||||
AsyncMock(),
|
||||
)
|
||||
|
||||
with pytest.raises(ScimUserAlreadyExists) as exc_info:
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await create_user(user=scim_user)
|
||||
|
||||
# Check that it's an HTTPException with status 409
|
||||
assert exc_info.value.status_code == 409
|
||||
assert "existing-user" in exc_info.value.message
|
||||
assert "existing-user" in str(exc_info.value.detail)
|
||||
mocked_new_user.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_existing_user_by_email_no_email(mocker):
|
||||
"""Should return None when new_user_request has no email"""
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
|
||||
new_user_request = NewUserRequest(
|
||||
user_id="test-user",
|
||||
user_email=None, # No email provided
|
||||
user_alias="Test User",
|
||||
teams=[],
|
||||
metadata={},
|
||||
auto_create_key=False,
|
||||
)
|
||||
|
||||
result = await UserProvisionerHelpers.handle_existing_user_by_email(
|
||||
prisma_client=mock_prisma_client,
|
||||
new_user_request=new_user_request
|
||||
)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_existing_user_by_email_no_existing_user(mocker):
|
||||
"""Should return None when no existing user is found with the email"""
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
mock_prisma_client.db = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
|
||||
|
||||
new_user_request = NewUserRequest(
|
||||
user_id="test-user",
|
||||
user_email="test@example.com",
|
||||
user_alias="Test User",
|
||||
teams=["team1"],
|
||||
metadata={"key": "value"},
|
||||
auto_create_key=False,
|
||||
)
|
||||
|
||||
result = await UserProvisionerHelpers.handle_existing_user_by_email(
|
||||
prisma_client=mock_prisma_client,
|
||||
new_user_request=new_user_request
|
||||
)
|
||||
|
||||
assert result is None
|
||||
mock_prisma_client.db.litellm_usertable.find_first.assert_called_once_with(
|
||||
where={"user_email": "test@example.com"}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_existing_user_by_email_existing_user_updated(mocker):
|
||||
"""Should update existing user and return SCIMUser when user with email exists"""
|
||||
# Mock existing user - create a proper mock object with attributes
|
||||
existing_user = mocker.MagicMock()
|
||||
existing_user.user_id = "old-user-id"
|
||||
existing_user.user_email = "test@example.com"
|
||||
existing_user.user_alias = "Old Name"
|
||||
existing_user.teams = ["old-team"]
|
||||
existing_user.metadata = {"old": "data"}
|
||||
|
||||
# Mock updated user
|
||||
updated_user = {
|
||||
"user_id": "new-user-id",
|
||||
"user_email": "test@example.com",
|
||||
"user_alias": "New Name",
|
||||
"teams": ["new-team"],
|
||||
"metadata": '{"new": "data"}'
|
||||
}
|
||||
|
||||
# Mock SCIM user to be returned
|
||||
mock_scim_user = SCIMUser(
|
||||
schemas=["urn:ietf:params:scim:schemas:core:2.0:User"],
|
||||
id="new-user-id",
|
||||
userName="new-user-id",
|
||||
name=SCIMUserName(familyName="Name", givenName="New"),
|
||||
emails=[SCIMUserEmail(value="test@example.com")],
|
||||
)
|
||||
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
mock_prisma_client.db = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user)
|
||||
mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user)
|
||||
|
||||
# Mock the transformation function
|
||||
mock_transform = mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user",
|
||||
AsyncMock(return_value=mock_scim_user)
|
||||
)
|
||||
|
||||
new_user_request = NewUserRequest(
|
||||
user_id="new-user-id",
|
||||
user_email="test@example.com",
|
||||
user_alias="New Name",
|
||||
teams=["new-team"],
|
||||
metadata={"new": "data"},
|
||||
auto_create_key=False,
|
||||
)
|
||||
|
||||
result = await UserProvisionerHelpers.handle_existing_user_by_email(
|
||||
prisma_client=mock_prisma_client,
|
||||
new_user_request=new_user_request
|
||||
)
|
||||
|
||||
# Verify the result
|
||||
assert result == mock_scim_user
|
||||
|
||||
# Verify database operations
|
||||
mock_prisma_client.db.litellm_usertable.find_first.assert_called_once_with(
|
||||
where={"user_email": "test@example.com"}
|
||||
)
|
||||
|
||||
mock_prisma_client.db.litellm_usertable.update.assert_called_once_with(
|
||||
where={"user_id": "old-user-id"},
|
||||
data={
|
||||
"user_id": "new-user-id",
|
||||
"user_email": "test@example.com",
|
||||
"user_alias": "New Name",
|
||||
"teams": ["new-team"],
|
||||
"metadata": '{"new": "data"}',
|
||||
},
|
||||
)
|
||||
|
||||
# Verify transformation was called
|
||||
mock_transform.assert_called_once_with(updated_user)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_team_membership_changes_no_changes(mocker):
|
||||
"""Should not call patch_team_membership when existing teams equal new teams"""
|
||||
mock_patch_team_membership = mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership",
|
||||
AsyncMock()
|
||||
)
|
||||
|
||||
# Same teams - no changes
|
||||
await _handle_team_membership_changes(
|
||||
user_id="test-user",
|
||||
existing_teams=["team1", "team2"],
|
||||
new_teams=["team1", "team2"]
|
||||
)
|
||||
|
||||
# Should not be called since no changes
|
||||
mock_patch_team_membership.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_team_membership_changes_add_teams(mocker):
|
||||
"""Should call patch_team_membership with teams to add"""
|
||||
mock_patch_team_membership = mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership",
|
||||
AsyncMock()
|
||||
)
|
||||
|
||||
# Adding teams
|
||||
await _handle_team_membership_changes(
|
||||
user_id="test-user",
|
||||
existing_teams=["team1"],
|
||||
new_teams=["team1", "team2", "team3"]
|
||||
)
|
||||
|
||||
mock_patch_team_membership.assert_called_once_with(
|
||||
user_id="test-user",
|
||||
teams_ids_to_add_user_to=["team2", "team3"], # Order might vary due to set operations
|
||||
teams_ids_to_remove_user_from=[]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_team_membership_changes_remove_teams(mocker):
|
||||
"""Should call patch_team_membership with teams to remove"""
|
||||
mock_patch_team_membership = mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership",
|
||||
AsyncMock()
|
||||
)
|
||||
|
||||
# Removing teams
|
||||
await _handle_team_membership_changes(
|
||||
user_id="test-user",
|
||||
existing_teams=["team1", "team2", "team3"],
|
||||
new_teams=["team1"]
|
||||
)
|
||||
|
||||
mock_patch_team_membership.assert_called_once_with(
|
||||
user_id="test-user",
|
||||
teams_ids_to_add_user_to=[],
|
||||
teams_ids_to_remove_user_from=["team2", "team3"] # Order might vary due to set operations
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_team_membership_changes_add_and_remove(mocker):
|
||||
"""Should call patch_team_membership with both teams to add and remove"""
|
||||
mock_patch_team_membership = mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership",
|
||||
AsyncMock()
|
||||
)
|
||||
|
||||
# Both adding and removing teams
|
||||
await _handle_team_membership_changes(
|
||||
user_id="test-user",
|
||||
existing_teams=["team1", "team2"],
|
||||
new_teams=["team2", "team3"]
|
||||
)
|
||||
|
||||
# team1 should be removed, team3 should be added, team2 stays
|
||||
mock_patch_team_membership.assert_called_once_with(
|
||||
user_id="test-user",
|
||||
teams_ids_to_add_user_to=["team3"],
|
||||
teams_ids_to_remove_user_from=["team1"]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_user_success(mocker):
|
||||
"""Should successfully update user with PUT request"""
|
||||
# Mock existing user
|
||||
existing_user = mocker.MagicMock()
|
||||
existing_user.teams = ["old-team"]
|
||||
|
||||
# Mock updated user
|
||||
updated_user = {
|
||||
"user_id": "test-user",
|
||||
"user_email": "updated@example.com",
|
||||
"user_alias": "Updated User",
|
||||
"teams": ["new-team"],
|
||||
"metadata": '{"scim_metadata": {"givenName": "Updated", "familyName": "User"}}'
|
||||
}
|
||||
|
||||
# Mock SCIM user for request
|
||||
scim_user = SCIMUser(
|
||||
schemas=["urn:ietf:params:scim:schemas:core:2.0:User"],
|
||||
userName="test-user",
|
||||
name=SCIMUserName(familyName="User", givenName="Updated"),
|
||||
emails=[SCIMUserEmail(value="updated@example.com")],
|
||||
groups=[SCIMUserGroup(value="new-team")]
|
||||
)
|
||||
|
||||
# Mock SCIM user for response
|
||||
response_scim_user = SCIMUser(
|
||||
schemas=["urn:ietf:params:scim:schemas:core:2.0:User"],
|
||||
id="test-user",
|
||||
userName="test-user",
|
||||
name=SCIMUserName(familyName="User", givenName="Updated"),
|
||||
emails=[SCIMUserEmail(value="updated@example.com")],
|
||||
)
|
||||
|
||||
# Mock prisma client
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
mock_prisma_client.db = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user)
|
||||
|
||||
# Mock dependencies
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
|
||||
AsyncMock(return_value=mock_prisma_client)
|
||||
)
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists",
|
||||
AsyncMock(return_value=existing_user)
|
||||
)
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes",
|
||||
AsyncMock()
|
||||
)
|
||||
mock_transform = mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user",
|
||||
AsyncMock(return_value=response_scim_user)
|
||||
)
|
||||
|
||||
# Call update_user
|
||||
result = await update_user(user_id="test-user", user=scim_user)
|
||||
|
||||
# Verify result
|
||||
assert result == response_scim_user
|
||||
|
||||
# Verify database update was called with correct data
|
||||
mock_prisma_client.db.litellm_usertable.update.assert_called_once()
|
||||
call_args = mock_prisma_client.db.litellm_usertable.update.call_args
|
||||
assert call_args[1]["where"] == {"user_id": "test-user"}
|
||||
assert call_args[1]["data"]["user_email"] == "updated@example.com"
|
||||
assert call_args[1]["data"]["teams"] == ["new-team"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_user_not_found(mocker):
|
||||
"""Should raise 404 when user doesn't exist"""
|
||||
scim_user = SCIMUser(
|
||||
schemas=["urn:ietf:params:scim:schemas:core:2.0:User"],
|
||||
userName="nonexistent-user",
|
||||
name=SCIMUserName(familyName="User", givenName="Test"),
|
||||
emails=[SCIMUserEmail(value="test@example.com")],
|
||||
)
|
||||
|
||||
# Mock dependencies to raise HTTPException for user not found
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
|
||||
AsyncMock(return_value=mocker.MagicMock())
|
||||
)
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists",
|
||||
AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "User not found"}))
|
||||
)
|
||||
|
||||
# Should raise ProxyException (which wraps the HTTPException)
|
||||
with pytest.raises(ProxyException):
|
||||
await update_user(user_id="nonexistent-user", user=scim_user)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_patch_user_success(mocker):
|
||||
"""Should successfully patch user with PATCH request"""
|
||||
# Mock existing user
|
||||
existing_user = mocker.MagicMock()
|
||||
existing_user.teams = ["team1"]
|
||||
existing_user.metadata = {}
|
||||
|
||||
# Mock updated user
|
||||
updated_user = {
|
||||
"user_id": "test-user",
|
||||
"user_alias": "Patched User",
|
||||
"teams": ["team1", "team2"],
|
||||
"metadata": '{"scim_metadata": {}}'
|
||||
}
|
||||
|
||||
# Mock patch operations
|
||||
patch_ops = SCIMPatchOp(
|
||||
schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"],
|
||||
Operations=[
|
||||
SCIMPatchOperation(op="replace", path="displayName", value="Patched User"),
|
||||
SCIMPatchOperation(op="add", path="groups", value=[{"value": "team2"}])
|
||||
]
|
||||
)
|
||||
|
||||
# Mock response SCIM user
|
||||
response_scim_user = SCIMUser(
|
||||
schemas=["urn:ietf:params:scim:schemas:core:2.0:User"],
|
||||
id="test-user",
|
||||
userName="test-user",
|
||||
name=SCIMUserName(familyName="User", givenName="Patched"),
|
||||
)
|
||||
|
||||
# Mock prisma client
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
mock_prisma_client.db = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user)
|
||||
|
||||
# Mock dependencies
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
|
||||
AsyncMock(return_value=mock_prisma_client)
|
||||
)
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists",
|
||||
AsyncMock(return_value=existing_user)
|
||||
)
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._handle_team_membership_changes",
|
||||
AsyncMock()
|
||||
)
|
||||
mock_transform = mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user",
|
||||
AsyncMock(return_value=response_scim_user)
|
||||
)
|
||||
|
||||
# Call patch_user
|
||||
result = await patch_user(user_id="test-user", patch_ops=patch_ops)
|
||||
|
||||
# Verify result
|
||||
assert result == response_scim_user
|
||||
|
||||
# Verify database update was called
|
||||
mock_prisma_client.db.litellm_usertable.update.assert_called_once()
|
||||
call_args = mock_prisma_client.db.litellm_usertable.update.call_args
|
||||
assert call_args[1]["where"] == {"user_id": "test-user"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_patch_user_not_found(mocker):
|
||||
"""Should raise 404 when user doesn't exist for patch"""
|
||||
patch_ops = SCIMPatchOp(
|
||||
schemas=["urn:ietf:params:scim:api:messages:2.0:PatchOp"],
|
||||
Operations=[
|
||||
SCIMPatchOperation(op="replace", path="displayName", value="New Name")
|
||||
]
|
||||
)
|
||||
|
||||
# Mock dependencies to raise HTTPException for user not found
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
|
||||
AsyncMock(return_value=mocker.MagicMock())
|
||||
)
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists",
|
||||
AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "User not found"}))
|
||||
)
|
||||
|
||||
# Should raise ProxyException (which wraps the HTTPException)
|
||||
with pytest.raises(ProxyException):
|
||||
await patch_user(user_id="nonexistent-user", patch_ops=patch_ops)
|
||||
|
||||
Reference in New Issue
Block a user