mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
fix(scim): retire directory resources with local deletions
This commit is contained in:
Joshua Valluru
parent
d5d2c5bea1
commit
f9eaae589e
5 changed files with 201 additions and 121 deletions
|
|
@ -1,13 +1,13 @@
|
|||
import re
|
||||
from collections import deque
|
||||
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping
|
||||
from collections.abc import AsyncGenerator, Mapping
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
from functools import reduce, wraps
|
||||
from functools import reduce
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Concatenate, Final, Literal, ParamSpec, TypeVar
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
|
@ -214,23 +214,6 @@ def group_members_after_patch(group: SCIMGroup, patch: SCIMPatchOp) -> SCIMGroup
|
|||
return reduce(_patch_group_operation, patch.Operations, group)
|
||||
|
||||
|
||||
Parameters = ParamSpec("Parameters")
|
||||
Result = TypeVar("Result")
|
||||
|
||||
|
||||
def serialized_source(
|
||||
operation: Callable[Concatenate["AgentProvisioningService", Parameters], Awaitable[Result]],
|
||||
) -> Callable[Concatenate["AgentProvisioningService", Parameters], Awaitable[Result]]:
|
||||
@wraps(operation)
|
||||
async def execute(
|
||||
service: "AgentProvisioningService", *args: Parameters.args, **kwargs: Parameters.kwargs
|
||||
) -> Result:
|
||||
async with service.source_transaction():
|
||||
return await operation(service, *args, **kwargs)
|
||||
|
||||
return execute
|
||||
|
||||
|
||||
def _identity_patch_children(value: object) -> tuple[object, ...] | Literal[True]:
|
||||
if isinstance(value, dict):
|
||||
fields: Final = TypeAdapter(dict[str, object]).validate_python(value)
|
||||
|
|
@ -500,27 +483,24 @@ class AgentProvisioningService:
|
|||
await SourceHumanProvisioner.finish_update(result)
|
||||
return result.document
|
||||
|
||||
@serialized_source
|
||||
async def delete(self, kind: Literal["Users", "Groups"], resource_id: str) -> None:
|
||||
from litellm.proxy.management_endpoints.scim import scim_v2
|
||||
|
||||
where: Final[LiteLLM_SCIMResourceWhereUniqueInput] = {"id": resource_id}
|
||||
async with self.client.tx() as tx:
|
||||
admin_group: Final = await scim_v2.provisioning_group_admin_role() if kind == "Groups" else None
|
||||
async with self.source_transaction() as tx:
|
||||
row: Final = await tx.litellm_scimresource.find_unique(where=where)
|
||||
if row is None or row.source_id != self.source.source_id or row.kind != kind:
|
||||
raise HTTPException(404, "SCIM resource not found in this provisioning source")
|
||||
if row.deleted:
|
||||
return
|
||||
if row.local_id is not None:
|
||||
try:
|
||||
if kind == "Groups":
|
||||
await scim_v2.delete_group(group_id=row.local_id)
|
||||
elif user_document(row).agent_user is None:
|
||||
await scim_v2.delete_user(user_id=row.local_id)
|
||||
except HTTPException as exc:
|
||||
if exc.status_code != 404:
|
||||
raise
|
||||
async with self.client.tx() as tx:
|
||||
user_deletion: Final = (
|
||||
await scim_v2.write_scim_user_deletion(tx, self.client, row.local_id)
|
||||
if row.local_id is not None and kind == "Users" and user_document(row).agent_user is None
|
||||
else None
|
||||
)
|
||||
if row.local_id is not None and kind == "Groups":
|
||||
await scim_v2.write_scim_group_deletion(tx, row.local_id, admin_group)
|
||||
if kind == "Users":
|
||||
memberships: Final[LiteLLM_SCIMResourceWhereInput] = {
|
||||
"source_id": self.source.source_id,
|
||||
|
|
@ -532,6 +512,7 @@ class AgentProvisioningService:
|
|||
await remove_group_member(tx, group, row.id)
|
||||
retired: Final[LiteLLM_SCIMResourceUpdateInput] = {"active": False, "deleted": True, "member_ids": []}
|
||||
await tx.litellm_scimresource.update(where=where, data=retired)
|
||||
await scim_v2.finish_scim_user_deletion(user_deletion)
|
||||
|
||||
async def create_group(self, group: SCIMGroup) -> SCIMGroup:
|
||||
from litellm.proxy.management_endpoints.scim import scim_v2
|
||||
|
|
|
|||
|
|
@ -1082,7 +1082,7 @@ async def write_user_keys_blocked(
|
|||
return tuple(key.token for key in affected_keys)
|
||||
|
||||
|
||||
async def _delete_rows_referencing_user(prisma_client: PrismaClient, *, user_id: str) -> None:
|
||||
async def _delete_rows_referencing_user(prisma_client: "PrismaClient | _GroupWriteDatabase", *, user_id: str) -> None:
|
||||
"""Drop rows whose foreign keys reference ``LiteLLM_UserTable.user_id``.
|
||||
|
||||
Required before deleting the user row itself, otherwise Postgres rejects
|
||||
|
|
@ -1887,6 +1887,74 @@ async def update_user(
|
|||
raise handle_exception_on_proxy(e)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SCIMUserDeletion:
|
||||
user_id: str
|
||||
tokens: tuple[str, ...]
|
||||
removals: tuple[TeamMemberRemoval, ...]
|
||||
|
||||
|
||||
async def write_scim_user_deletion(tx: "Prisma", client: PrismaClient, user_id: str) -> SCIMUserDeletion | None:
|
||||
database: Final = _GroupWriteDatabase(tx)
|
||||
user: Final = await UserRepository(database).find_by_id(user_id, id_field="user_id")
|
||||
if user is None:
|
||||
return None
|
||||
teams: Final = tuple(
|
||||
[await TeamRepository(database).find_by_id(team_id, id_field="team_id") for team_id in sorted(user.teams or ())]
|
||||
)
|
||||
for team in teams:
|
||||
if team is not None and user_id in (team.members or ()):
|
||||
await _table(TeamRepository(database)).update(
|
||||
where={"team_id": team.team_id},
|
||||
data={"members": [member for member in team.members if member != user_id]},
|
||||
)
|
||||
removals: Final = tuple(
|
||||
[
|
||||
await delete_team_member_in_transaction(
|
||||
tx=tx,
|
||||
data=TeamMemberDeleteRequest(team_id=team.team_id, user_id=user_id),
|
||||
existing_team_row=team,
|
||||
prisma_client=client,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
for team in teams
|
||||
if team is not None and any(member.user_id == user_id for member in team.members_with_roles or ())
|
||||
]
|
||||
)
|
||||
tokens: Final = await write_user_keys_blocked(
|
||||
_table(VerificationTokenRepository(database)), user_id=user_id, blocked=True
|
||||
)
|
||||
await _delete_rows_referencing_user(database, user_id=user_id)
|
||||
await _table(UserRepository(database)).delete(where={"user_id": user_id})
|
||||
return SCIMUserDeletion(user_id=user_id, tokens=tokens, removals=removals)
|
||||
|
||||
|
||||
async def finish_scim_user_deletion(result: SCIMUserDeletion | None) -> None:
|
||||
if result is None:
|
||||
return
|
||||
for removal in result.removals:
|
||||
await finish_team_member_removal(removal, UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN))
|
||||
await finish_provisioned_user_update(result.user_id, result.tokens)
|
||||
|
||||
|
||||
async def write_scim_group_deletion(tx: "Prisma", group_id: str, admin_group: str | None) -> tuple[str, ...]:
|
||||
database: Final = _GroupWriteDatabase(tx)
|
||||
await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, group_id)
|
||||
team: Final = await TeamRepository(database).find_by_id(group_id, id_field="team_id")
|
||||
if team is None:
|
||||
return ()
|
||||
member_ids: Final = tuple(await _get_team_member_user_ids_from_team(team))
|
||||
for member_id in member_ids:
|
||||
user: Final = await UserRepository(database).find_by_id(member_id, id_field="user_id")
|
||||
if user is not None and group_id in (user.teams or ()):
|
||||
await _table(UserRepository(database)).update(
|
||||
where={"user_id": member_id}, data={"teams": [value for value in user.teams if value != group_id]}
|
||||
)
|
||||
await _table(TeamRepository(database)).delete(where={"team_id": group_id})
|
||||
await write_scim_member_roles(database, member_ids, admin_group)
|
||||
return member_ids
|
||||
|
||||
|
||||
@scim_router.delete(
|
||||
"/Users/{user_id}",
|
||||
status_code=204,
|
||||
|
|
@ -1901,43 +1969,10 @@ async def delete_user(
|
|||
verbose_proxy_logger.debug("SCIM DELETE USER request for user_id=%s", user_id)
|
||||
try:
|
||||
prisma_client: Final = await _get_prisma_client_or_raise_exception()
|
||||
existing_user: Final = await _check_user_exists(user_id)
|
||||
|
||||
# Get teams user belongs to
|
||||
found_teams: Final = tuple(
|
||||
[
|
||||
await _table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id})
|
||||
for team_id in existing_user.teams or []
|
||||
]
|
||||
)
|
||||
teams: Final = tuple(team for team in found_teams if team)
|
||||
|
||||
# Remove user from all teams
|
||||
for team in teams:
|
||||
current_members: Sequence[str] = team.members or []
|
||||
if user_id in current_members:
|
||||
new_members = [m for m in current_members if m != user_id]
|
||||
await _table(TeamRepository(prisma_client)).update(
|
||||
where={"team_id": team.team_id}, data={"members": new_members}
|
||||
)
|
||||
|
||||
team_row = LiteLLM_TeamTable.model_validate(team.model_dump())
|
||||
if any(member.user_id == user_id for member in team_row.members_with_roles or []):
|
||||
await team_member_delete(
|
||||
data=TeamMemberDeleteRequest(team_id=team_row.team_id, user_id=user_id),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
await _set_user_keys_blocked(user_id=user_id, blocked=True)
|
||||
|
||||
await _delete_rows_referencing_user(prisma_client, user_id=user_id)
|
||||
|
||||
# Delete user
|
||||
await _table(UserRepository(prisma_client)).delete(where={"user_id": user_id})
|
||||
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
await evict_and_broadcast(cache_keys=(user_id,), user_api_key_cache=user_api_key_cache)
|
||||
await _check_user_exists(user_id)
|
||||
async with prisma_client.tx() as tx:
|
||||
result: Final = await write_scim_user_deletion(tx, prisma_client, user_id)
|
||||
await finish_scim_user_deletion(result)
|
||||
|
||||
return Response(status_code=204)
|
||||
except Exception as e:
|
||||
|
|
@ -2874,25 +2909,10 @@ async def delete_group(
|
|||
verbose_proxy_logger.debug("SCIM DELETE GROUP request for group_id=%s", group_id)
|
||||
try:
|
||||
prisma_client: Final = await _get_prisma_client_or_raise_exception()
|
||||
existing_team: Final = await _check_team_exists(group_id)
|
||||
|
||||
member_ids: Final = await _get_team_member_user_ids_from_team(existing_team)
|
||||
|
||||
# For each member, remove this team from their teams list
|
||||
for member_id in member_ids:
|
||||
user = await _table(UserRepository(prisma_client)).find_unique(where={"user_id": member_id})
|
||||
if user:
|
||||
current_teams = user.teams or []
|
||||
if group_id in current_teams:
|
||||
new_teams = [t for t in current_teams if t != group_id]
|
||||
await _table(UserRepository(prisma_client)).update(
|
||||
where={"user_id": member_id}, data={"teams": new_teams}
|
||||
)
|
||||
|
||||
await _recompute_scim_member_roles(prisma_client, member_ids)
|
||||
|
||||
# Delete team
|
||||
await _table(TeamRepository(prisma_client)).delete(where={"team_id": group_id})
|
||||
await _check_team_exists(group_id)
|
||||
admin_group: Final = await _get_scim_admin_group()
|
||||
async with prisma_client.tx() as tx:
|
||||
await write_scim_group_deletion(tx, group_id, admin_group)
|
||||
|
||||
return Response(status_code=204)
|
||||
|
||||
|
|
|
|||
|
|
@ -383,32 +383,26 @@ async def test_concurrent_native_profile_update_returns_conflict() -> None:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kind", ["Users", "Groups"])
|
||||
async def test_scoped_delete_propagates_human_and_team_deprovisioning(
|
||||
kind: str,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy.management_endpoints.scim import scim_v2
|
||||
async def test_scoped_delete_propagates_human_and_team_deprovisioning(kind: str) -> None:
|
||||
from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable
|
||||
|
||||
service, tx, native = provisioning_fixture()
|
||||
document: Final = SCIMUser(schemas=[], userName="human@example.com").model_dump(mode="json")
|
||||
row: Final = native.model_copy(update={"kind": kind, "document": document})
|
||||
row: Final = native.model_copy(update={"kind": kind, "document": SCIMUser(schemas=[], userName="human").model_dump()})
|
||||
tx.litellm_scimresource.find_unique.return_value = row
|
||||
delete_user: Final = AsyncMock()
|
||||
delete_group: Final = AsyncMock()
|
||||
monkeypatch.setattr(scim_v2, "delete_user", delete_user)
|
||||
monkeypatch.setattr(scim_v2, "delete_group", delete_group)
|
||||
tx.query_raw = AsyncMock(return_value=[])
|
||||
tx.litellm_usertable.find_unique = AsyncMock(return_value=LiteLLM_UserTable(user_id=row.local_id, teams=[]))
|
||||
tx.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable(team_id=row.local_id, members_with_roles=[]))
|
||||
tx.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
tx.litellm_usertable.delete = AsyncMock()
|
||||
tx.litellm_teamtable.delete = AsyncMock()
|
||||
tx.litellm_invitationlink.delete_many = AsyncMock()
|
||||
tx.litellm_organizationmembership.delete_many = AsyncMock()
|
||||
tx.litellm_teammembership.delete_many = AsyncMock()
|
||||
await service.delete(kind, row.id)
|
||||
if kind == "Users":
|
||||
delete_user.assert_awaited_once_with(user_id=row.local_id)
|
||||
delete_group.assert_not_awaited()
|
||||
else:
|
||||
delete_group.assert_awaited_once_with(group_id=row.local_id)
|
||||
delete_user.assert_not_awaited()
|
||||
deletion: Final = tx.litellm_usertable.delete if kind == "Users" else tx.litellm_teamtable.delete
|
||||
deletion.assert_awaited_once_with(where={"user_id" if kind == "Users" else "team_id": row.local_id})
|
||||
tx.litellm_scimresource.update.assert_awaited_once_with(
|
||||
where={"id": row.id},
|
||||
data={"active": False, "deleted": True, "member_ids": []},
|
||||
where={"id": row.id}, data={"active": False, "deleted": True, "member_ids": []},
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -423,7 +417,7 @@ async def test_failed_human_deletion_remains_retryable(monkeypatch: pytest.Monke
|
|||
service, tx, native = provisioning_fixture()
|
||||
row: Final = native.model_copy(update={"document": SCIMUser(schemas=[], userName="human@example.com").model_dump()})
|
||||
tx.litellm_scimresource.find_unique.return_value = row
|
||||
monkeypatch.setattr(scim_v2, "delete_user", AsyncMock(side_effect=HTTPException(503, "unavailable")))
|
||||
tx.litellm_usertable.find_unique = AsyncMock(side_effect=HTTPException(503, "unavailable"))
|
||||
with pytest.raises(HTTPException) as failure:
|
||||
await service.delete("Users", row.id)
|
||||
assert failure.value.status_code == 503
|
||||
|
|
@ -1244,3 +1238,78 @@ async def test_group_addition_reuses_member_writes_and_defers_existing_audit_and
|
|||
assert any(call.kwargs.get("cache_keys") == (human.local_id,) for call in invalidate.await_args_list)
|
||||
audit.assert_called_once()
|
||||
assert audit.call_args.kwargs["existing_user_ids"] == frozenset([human.local_id])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kind", ["Users", "Groups"])
|
||||
async def test_delete_retry_retires_directory_when_local_record_is_already_missing(kind, monkeypatch):
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.management_endpoints.scim import scim_v2
|
||||
|
||||
service, tx, native = provisioning_fixture()
|
||||
row = native.model_copy(update={"kind": kind, "document": SCIMUser(schemas=[], userName="human").model_dump()})
|
||||
tx.litellm_scimresource.find_unique.return_value = row
|
||||
tx.litellm_usertable.find_unique = AsyncMock(return_value=None)
|
||||
tx.litellm_teamtable.find_unique = AsyncMock(return_value=None)
|
||||
missing = ProxyException(message="not found", type="not_found", param=None, code=404)
|
||||
monkeypatch.setattr(scim_v2, "delete_user", AsyncMock(side_effect=missing))
|
||||
monkeypatch.setattr(scim_v2, "delete_group", AsyncMock(side_effect=missing))
|
||||
await service.delete(kind, row.id)
|
||||
assert tx.litellm_scimresource.update.call_args.kwargs["data"] == {
|
||||
"active": False, "deleted": True, "member_ids": [],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("kind", ["Users", "Groups"])
|
||||
async def test_directory_retirement_failure_rolls_back_local_deletion(kind, monkeypatch):
|
||||
from contextlib import asynccontextmanager
|
||||
from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable
|
||||
from litellm.proxy.management_endpoints.scim import scim_v2
|
||||
|
||||
service, tx, native = provisioning_fixture()
|
||||
row = native.model_copy(update={"kind": kind, "document": SCIMUser(schemas=[], userName="human").model_dump()})
|
||||
state = {"local_exists": True, "directory_active": True}
|
||||
transaction_writes = []
|
||||
|
||||
@asynccontextmanager
|
||||
async def transaction(**kwargs):
|
||||
before = dict(state)
|
||||
writes = set()
|
||||
transaction_writes.append(writes)
|
||||
try:
|
||||
yield tx
|
||||
except BaseException:
|
||||
state.update({key: before[key] for key in writes})
|
||||
raise
|
||||
finally:
|
||||
transaction_writes.pop()
|
||||
|
||||
async def delete_local(**kwargs):
|
||||
if "where" in kwargs:
|
||||
transaction_writes[-1].add("local_exists")
|
||||
state["local_exists"] = False
|
||||
|
||||
async def retire(**kwargs):
|
||||
transaction_writes[-1].add("directory_active")
|
||||
state["directory_active"] = False
|
||||
raise RuntimeError("directory write failed")
|
||||
|
||||
service.client.tx = transaction
|
||||
tx.query_raw = AsyncMock(return_value=[])
|
||||
tx.litellm_scimresource.find_unique.return_value = row
|
||||
tx.litellm_scimresource.update = AsyncMock(side_effect=retire)
|
||||
tx.litellm_usertable.find_unique = AsyncMock(return_value=LiteLLM_UserTable(user_id=row.local_id, teams=[]))
|
||||
tx.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable(team_id=row.local_id, members_with_roles=[], max_budget=99))
|
||||
tx.litellm_usertable.delete = AsyncMock(side_effect=delete_local)
|
||||
tx.litellm_teamtable.delete = AsyncMock(side_effect=delete_local)
|
||||
tx.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
tx.litellm_invitationlink.delete_many = AsyncMock()
|
||||
tx.litellm_organizationmembership.delete_many = AsyncMock()
|
||||
tx.litellm_teammembership.delete_many = AsyncMock()
|
||||
monkeypatch.setattr(scim_v2, "delete_user", delete_local)
|
||||
monkeypatch.setattr(scim_v2, "delete_group", delete_local)
|
||||
with pytest.raises(RuntimeError, match="directory write failed"):
|
||||
await service.delete(kind, row.id)
|
||||
assert state == {"local_exists": True, "directory_active": True}
|
||||
assert tx.litellm_teamtable.find_unique.return_value.max_budget == 99
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ def _build_prisma_with_keys(user_keys, mock_user=None, updated_user=None):
|
|||
mock_client = MagicMock()
|
||||
mock_db = MagicMock()
|
||||
mock_client.db = mock_db
|
||||
mock_client.tx.return_value.__aenter__.return_value = mock_db
|
||||
if mock_user is not None:
|
||||
mock_db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user)
|
||||
if updated_user is not None:
|
||||
|
|
|
|||
|
|
@ -3032,17 +3032,19 @@ async def test_delete_group_recomputes_roles_for_members(mocker):
|
|||
deleting the admin group demotes everyone who was only admin through it."""
|
||||
from litellm.proxy._types import Member
|
||||
|
||||
existing_team = mocker.MagicMock()
|
||||
existing_team = LiteLLM_TeamTable(team_id="test-team-123")
|
||||
existing_team.members_with_roles = [
|
||||
Member(user_id="user1", role="user"),
|
||||
Member(user_id="user2", role="user"),
|
||||
]
|
||||
|
||||
member = mocker.MagicMock()
|
||||
member = LiteLLM_UserTable(user_id="user1")
|
||||
member.teams = ["test-team-123"]
|
||||
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
mock_prisma_client.db = mocker.MagicMock()
|
||||
mock_prisma_client.tx.return_value.__aenter__.return_value = mock_prisma_client.db
|
||||
mock_prisma_client.db.query_raw = AsyncMock(return_value=[])
|
||||
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team)
|
||||
mock_prisma_client.db.litellm_teamtable.delete = AsyncMock()
|
||||
|
|
@ -3056,7 +3058,7 @@ async def test_delete_group_recomputes_roles_for_members(mocker):
|
|||
AsyncMock(return_value=mock_prisma_client),
|
||||
)
|
||||
recompute_mock = mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._recompute_scim_member_roles",
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.write_scim_member_roles",
|
||||
AsyncMock(),
|
||||
)
|
||||
|
||||
|
|
@ -3516,6 +3518,8 @@ async def test_apply_group_patch_updates_does_not_write_legacy_members(mocker):
|
|||
def _mock_prisma_for_delete_user(mocker, team):
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
mock_prisma_client.db = mocker.MagicMock()
|
||||
mock_prisma_client.tx.return_value.__aenter__.return_value = mock_prisma_client.db
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
|
||||
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team)
|
||||
mock_prisma_client.db.litellm_teamtable.update = AsyncMock()
|
||||
|
|
@ -3525,6 +3529,8 @@ def _mock_prisma_for_delete_user(mocker, team):
|
|||
|
||||
|
||||
def _patch_delete_user_dependencies(mocker, mock_prisma_client, existing_user):
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=existing_user)
|
||||
mocker.patch("litellm.proxy.management_endpoints.scim.scim_v2.finish_scim_user_deletion", AsyncMock())
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
|
||||
AsyncMock(return_value=mock_prisma_client),
|
||||
|
|
@ -3551,7 +3557,7 @@ async def test_delete_user_prunes_members_with_roles(mocker):
|
|||
to the now-deleted user."""
|
||||
user_id = "scim-del-user"
|
||||
|
||||
existing_user = mocker.MagicMock()
|
||||
existing_user = LiteLLM_UserTable(user_id=user_id)
|
||||
existing_user.teams = ["team-1"]
|
||||
|
||||
team = LiteLLM_TeamTable(
|
||||
|
|
@ -3563,7 +3569,7 @@ async def test_delete_user_prunes_members_with_roles(mocker):
|
|||
mock_prisma_client = _mock_prisma_for_delete_user(mocker, team)
|
||||
_patch_delete_user_dependencies(mocker, mock_prisma_client, existing_user)
|
||||
team_member_delete_mock = mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete",
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.delete_team_member_in_transaction",
|
||||
AsyncMock(),
|
||||
)
|
||||
|
||||
|
|
@ -3585,7 +3591,7 @@ async def test_delete_user_surfaces_prune_failure_and_keeps_user(mocker):
|
|||
so the IdP retries)."""
|
||||
user_id = "scim-del-user"
|
||||
|
||||
existing_user = mocker.MagicMock()
|
||||
existing_user = LiteLLM_UserTable(user_id=user_id)
|
||||
existing_user.teams = ["team-1"]
|
||||
|
||||
team = LiteLLM_TeamTable(
|
||||
|
|
@ -3597,7 +3603,7 @@ async def test_delete_user_surfaces_prune_failure_and_keeps_user(mocker):
|
|||
mock_prisma_client = _mock_prisma_for_delete_user(mocker, team)
|
||||
_patch_delete_user_dependencies(mocker, mock_prisma_client, existing_user)
|
||||
mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete",
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.delete_team_member_in_transaction",
|
||||
AsyncMock(side_effect=Exception("database connection lost")),
|
||||
)
|
||||
|
||||
|
|
@ -3614,7 +3620,7 @@ async def test_delete_user_skips_teams_where_not_a_member(mocker):
|
|||
user, so a stale legacy membership can't block the delete."""
|
||||
user_id = "scim-del-user"
|
||||
|
||||
existing_user = mocker.MagicMock()
|
||||
existing_user = LiteLLM_UserTable(user_id=user_id)
|
||||
existing_user.teams = ["team-1"]
|
||||
|
||||
team = LiteLLM_TeamTable(
|
||||
|
|
@ -3626,7 +3632,7 @@ async def test_delete_user_skips_teams_where_not_a_member(mocker):
|
|||
mock_prisma_client = _mock_prisma_for_delete_user(mocker, team)
|
||||
_patch_delete_user_dependencies(mocker, mock_prisma_client, existing_user)
|
||||
team_member_delete_mock = mocker.patch(
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete",
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.delete_team_member_in_transaction",
|
||||
AsyncMock(),
|
||||
)
|
||||
|
||||
|
|
@ -6033,6 +6039,7 @@ def _shadow_tenant_prisma(
|
|||
|
||||
prisma_client = mocker.MagicMock()
|
||||
prisma_client.db = mocker.MagicMock()
|
||||
prisma_client.tx.return_value.__aenter__.return_value = prisma_client.db
|
||||
prisma_client.db.litellm_usertable = mocker.MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=find_unique)
|
||||
prisma_client.db.litellm_usertable.find_many = AsyncMock(side_effect=identity_rows)
|
||||
|
|
@ -6076,10 +6083,12 @@ async def test_merge_placeholder_hands_the_group_to_the_shadowed_account(mocker,
|
|||
)
|
||||
team_member_delete_mock = (
|
||||
mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the endpoint
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", AsyncMock()
|
||||
"litellm.proxy.management_endpoints.scim.scim_v2.delete_team_member_in_transaction", AsyncMock()
|
||||
)
|
||||
)
|
||||
|
||||
mocker.patch("litellm.proxy.management_endpoints.scim.scim_v2.finish_scim_user_deletion", AsyncMock())
|
||||
|
||||
with pytest.raises(HTTPException) as before:
|
||||
await _push_shadow_member(shadowed_tenant)
|
||||
assert before.value.status_code == 400
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue