mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
fix(team): schedule membership audit writes after commit and lock the roster on role updates
The member add, delete and role-change audit rows were awaited on the request path, so a slow audit sink held the response, and the roster was serialized before checking whether audit logging is enabled at all. Membership audit work is now scheduled after the transaction commits and skipped outright when auditing is off. member_update read the roster outside the team advisory lock and wrote it back, so a concurrent add or delete could be lost. It now takes the lock, rereads the roster, and builds the before and after snapshots from that read. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng
parent
c4272d894f
commit
181406e05f
3 changed files with 310 additions and 89 deletions
|
|
@ -3053,6 +3053,35 @@ async def _add_team_members_to_team(
|
|||
return updated_team, updated_users, updated_team_memberships
|
||||
|
||||
|
||||
async def _update_team_member_role(
|
||||
prisma_client: PrismaClient,
|
||||
team_id: str,
|
||||
user_id: str,
|
||||
role: Literal["admin", "user"],
|
||||
user_email: str | None,
|
||||
) -> tuple[tuple[Member, ...], tuple[Member, ...]]:
|
||||
"""Rewrite one member's role from the roster read under the team lock; returns (before, after)."""
|
||||
async with prisma_client.tx() as tx:
|
||||
await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id)
|
||||
|
||||
locked_members: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, team_id)
|
||||
if locked_members is None:
|
||||
raise HTTPException(status_code=404, detail={"error": f"Team id={team_id} does not exist in db"})
|
||||
|
||||
before: Final = tuple(locked_members)
|
||||
after: Final = tuple(
|
||||
Member(user_id=member.user_id, role=role, user_email=user_email or member.user_email)
|
||||
if member.user_id == user_id
|
||||
else member
|
||||
for member in before
|
||||
)
|
||||
await _team_tx_db(tx).update(
|
||||
where={"team_id": team_id},
|
||||
data={"members_with_roles": json.dumps([m.model_dump() for m in after])},
|
||||
)
|
||||
return before, after
|
||||
|
||||
|
||||
def _emit_team_members_metric(team: LiteLLM_TeamTable) -> None:
|
||||
"""Update the Prometheus team members gauge after a membership change.
|
||||
|
||||
|
|
@ -3164,7 +3193,7 @@ def _members_audit_value(team_alias: str | None, members: Sequence[Member]) -> s
|
|||
)
|
||||
|
||||
|
||||
async def _create_team_membership_audit_log(
|
||||
def _schedule_team_membership_audit_log(
|
||||
team_id: str,
|
||||
team_alias: str | None,
|
||||
before_members: Sequence[Member],
|
||||
|
|
@ -3172,21 +3201,29 @@ async def _create_team_membership_audit_log(
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_proxy_admin_name: str,
|
||||
) -> None:
|
||||
from litellm.proxy.management_helpers.audit_logs import create_object_audit_log
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_object_audit_log,
|
||||
is_audit_logging_enabled,
|
||||
)
|
||||
|
||||
await create_object_audit_log(
|
||||
object_id=team_id,
|
||||
action="updated",
|
||||
litellm_changed_by=None,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
table_name=LitellmTableNames.TEAM_TABLE_NAME,
|
||||
before_value=_members_audit_value(team_alias, before_members),
|
||||
after_value=_members_audit_value(team_alias, after_members),
|
||||
if not is_audit_logging_enabled() or tuple(before_members) == tuple(after_members):
|
||||
return
|
||||
|
||||
asyncio.create_task(
|
||||
create_object_audit_log(
|
||||
object_id=team_id,
|
||||
action="updated",
|
||||
litellm_changed_by=None,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
table_name=LitellmTableNames.TEAM_TABLE_NAME,
|
||||
before_value=_members_audit_value(team_alias, before_members),
|
||||
after_value=_members_audit_value(team_alias, after_members),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _create_team_member_add_audit_logs(
|
||||
def _schedule_team_member_add_audit_logs(
|
||||
team_id: str,
|
||||
team_alias: str | None,
|
||||
updated_users: Sequence[LiteLLM_UserTable],
|
||||
|
|
@ -3196,29 +3233,32 @@ async def _create_team_member_add_audit_logs(
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_proxy_admin_name: str,
|
||||
) -> None:
|
||||
"""Record the membership change, and any user row it created, in the audit log.
|
||||
|
||||
The entries are written concurrently so a request adding many members does
|
||||
not pay for them one after another.
|
||||
"""
|
||||
from litellm.proxy.management_helpers.audit_logs import create_object_audit_log
|
||||
|
||||
created_user_entries: Final = tuple(
|
||||
create_object_audit_log(
|
||||
object_id=user.user_id,
|
||||
action="created",
|
||||
litellm_changed_by=None,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
table_name=LitellmTableNames.USER_TABLE_NAME,
|
||||
before_value=None,
|
||||
after_value=safe_dumps(user.model_dump(exclude_none=True)),
|
||||
)
|
||||
for user in updated_users
|
||||
if user.user_id is not None and user.user_id not in existing_user_ids
|
||||
"""Record the membership change, and any user row it created, in the audit log."""
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
create_object_audit_log,
|
||||
is_audit_logging_enabled,
|
||||
)
|
||||
|
||||
membership_entry: Final = _create_team_membership_audit_log(
|
||||
if not is_audit_logging_enabled():
|
||||
return
|
||||
|
||||
for user in updated_users:
|
||||
if user.user_id is None or user.user_id in existing_user_ids:
|
||||
continue
|
||||
asyncio.create_task(
|
||||
create_object_audit_log(
|
||||
object_id=user.user_id,
|
||||
action="created",
|
||||
litellm_changed_by=None,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
table_name=LitellmTableNames.USER_TABLE_NAME,
|
||||
before_value=None,
|
||||
after_value=safe_dumps(user.model_dump(exclude_none=True)),
|
||||
)
|
||||
)
|
||||
|
||||
_schedule_team_membership_audit_log(
|
||||
team_id=team_id,
|
||||
team_alias=team_alias,
|
||||
before_members=before_members,
|
||||
|
|
@ -3227,8 +3267,6 @@ async def _create_team_member_add_audit_logs(
|
|||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
)
|
||||
|
||||
await asyncio.gather(*created_user_entries, membership_entry)
|
||||
|
||||
|
||||
async def _validate_and_populate_member_user_info(
|
||||
member: Member,
|
||||
|
|
@ -3462,7 +3500,7 @@ async def team_member_add(
|
|||
|
||||
_emit_team_members_metric(complete_team_data)
|
||||
|
||||
await _create_team_member_add_audit_logs(
|
||||
_schedule_team_member_add_audit_logs(
|
||||
team_id=data.team_id,
|
||||
team_alias=complete_team_data.team_alias,
|
||||
updated_users=updated_users,
|
||||
|
|
@ -3540,15 +3578,14 @@ async def team_member_delete(
|
|||
data=data, user_api_key_dict=user_api_key_dict
|
||||
)
|
||||
|
||||
if before_members != after_members:
|
||||
await _create_team_membership_audit_log(
|
||||
team_id=existing_team_row.team_id,
|
||||
team_alias=existing_team_row.team_alias,
|
||||
before_members=before_members,
|
||||
after_members=after_members,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
)
|
||||
_schedule_team_membership_audit_log(
|
||||
team_id=existing_team_row.team_id,
|
||||
team_alias=existing_team_row.team_alias,
|
||||
before_members=before_members,
|
||||
after_members=after_members,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
)
|
||||
|
||||
return existing_team_row
|
||||
|
||||
|
|
@ -3858,39 +3895,22 @@ async def team_member_update(
|
|||
|
||||
### update team member role
|
||||
if data.role is not None:
|
||||
members_before_role_update: Final = tuple(
|
||||
Member(user_id=member.user_id, user_email=member.user_email, role=member.role)
|
||||
for member in team_table.members_with_roles
|
||||
members_before_role_update, team_members = await _update_team_member_role(
|
||||
prisma_client=prisma_client,
|
||||
team_id=data.team_id,
|
||||
user_id=received_user_id,
|
||||
role=data.role,
|
||||
user_email=data.user_email,
|
||||
)
|
||||
team_members: Final[list[Member]] = []
|
||||
for member in members_before_role_update:
|
||||
if member.user_id == received_user_id:
|
||||
team_members.append(
|
||||
Member(
|
||||
user_id=member.user_id,
|
||||
role=data.role,
|
||||
user_email=data.user_email or member.user_email,
|
||||
)
|
||||
)
|
||||
else:
|
||||
team_members.append(member)
|
||||
|
||||
team_table.members_with_roles = team_members
|
||||
|
||||
_db_team_members: Final[list[dict]] = [m.model_dump() for m in team_members]
|
||||
await _team_db(prisma_client).update(
|
||||
where={"team_id": data.team_id},
|
||||
data={"members_with_roles": json.dumps(_db_team_members)},
|
||||
team_table.members_with_roles = list(team_members)
|
||||
_schedule_team_membership_audit_log(
|
||||
team_id=data.team_id,
|
||||
team_alias=team_table.team_alias,
|
||||
before_members=members_before_role_update,
|
||||
after_members=team_members,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
)
|
||||
if members_before_role_update != tuple(team_members):
|
||||
await _create_team_membership_audit_log(
|
||||
team_id=data.team_id,
|
||||
team_alias=team_table.team_alias,
|
||||
before_members=members_before_role_update,
|
||||
after_members=team_members,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
)
|
||||
|
||||
return TeamMemberUpdateResponse(
|
||||
team_id=data.team_id,
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import json
|
|||
from contextlib import asynccontextmanager, contextmanager
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from collections.abc import Sequence
|
||||
from typing import Final, Optional, cast
|
||||
from unittest.mock import AsyncMock, MagicMock, PropertyMock, call, patch
|
||||
|
||||
|
|
@ -148,11 +149,11 @@ def _wire_member_add_tx(prisma_client):
|
|||
|
||||
|
||||
def _wire_member_delete_tx(prisma_client):
|
||||
"""/team/member_delete's four cleanups, plus the advisory-lock re-read that now guards
|
||||
them, run inside one transaction, so a mocked client has to hand back its own table
|
||||
mocks (and a `query_raw` that answers the locked re-read from the same team row the
|
||||
test already configured on `find_unique`) out of `tx()` for the existing per-table
|
||||
assertions to keep seeing the calls."""
|
||||
"""/team/member_delete's four cleanups and /team/member_update's role rewrite, plus the
|
||||
advisory-lock re-read that guards them, run inside one transaction, so a mocked client
|
||||
has to hand back its own table mocks (and a `query_raw` that answers the locked re-read
|
||||
from the same team row the test already configured on `find_unique`) out of `tx()` for
|
||||
the existing per-table assertions to keep seeing the calls."""
|
||||
|
||||
async def _query_raw(sql, team_id):
|
||||
if sql != TEAM_ADVISORY_LOCK_SQL:
|
||||
|
|
@ -168,10 +169,12 @@ def _wire_member_delete_tx(prisma_client):
|
|||
return getattr(prisma_client.db, table_name)
|
||||
|
||||
tx = _Tx()
|
||||
tx.query_raw = AsyncMock(side_effect=_query_raw)
|
||||
tx_cm = MagicMock()
|
||||
tx_cm.__aenter__ = AsyncMock(return_value=tx)
|
||||
tx_cm.__aexit__ = AsyncMock(return_value=None)
|
||||
prisma_client.tx = MagicMock(return_value=tx_cm)
|
||||
return tx
|
||||
|
||||
|
||||
def _wire_team_delete_tx(prisma_client):
|
||||
|
|
@ -13311,8 +13314,7 @@ async def test_team_member_add_audits_a_user_created_from_a_list_payload(monkeyp
|
|||
side_effect=fake_add_team_members_to_team,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.team_endpoints._create_team_member_add_audit_logs",
|
||||
new_callable=AsyncMock,
|
||||
"litellm.proxy.management_endpoints.team_endpoints._schedule_team_member_add_audit_logs",
|
||||
) as mock_audit,
|
||||
):
|
||||
await team_member_add(
|
||||
|
|
@ -13517,12 +13519,10 @@ async def test_team_member_update_role_change_emits_a_roster_audit_event(monkeyp
|
|||
}
|
||||
|
||||
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
|
||||
mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=team_row)
|
||||
mock_prisma_client.db.litellm_teamtable.update = AsyncMock(side_effect=_roster_writer(team_row))
|
||||
mock_prisma_client.db.litellm_auditlog.create = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mock_tx = AsyncMock()
|
||||
mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx)
|
||||
mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None)
|
||||
_wire_member_delete_tx(mock_prisma_client)
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
|
||||
|
|
@ -13563,6 +13563,197 @@ async def test_team_member_update_role_change_emits_a_roster_audit_event(monkeyp
|
|||
)
|
||||
|
||||
|
||||
def _roster_writer(team_row: LiteLLM_TeamTable):
|
||||
"""An `update` side effect that lands `members_with_roles` on the team row later reads see."""
|
||||
|
||||
async def _update(where, data):
|
||||
team_row.members_with_roles = [Member(**m) for m in json.loads(data["members_with_roles"])]
|
||||
return team_row
|
||||
|
||||
return _update
|
||||
|
||||
|
||||
def _member_update_patches(team_snapshot: LiteLLM_TeamTable):
|
||||
return (
|
||||
patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
|
||||
"litellm.proxy.management_endpoints.team_endpoints.team_info",
|
||||
AsyncMock(
|
||||
return_value={
|
||||
"team_info": TeamInfoResponseObjectTeamTable(**team_snapshot.model_dump()),
|
||||
"team_memberships": [
|
||||
LiteLLM_TeamMembership(user_id="bob", team_id=team_snapshot.team_id, budget_id=None)
|
||||
],
|
||||
}
|
||||
),
|
||||
),
|
||||
patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
|
||||
"litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership",
|
||||
AsyncMock(),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_update_role_change_rewrites_the_roster_it_read_under_the_lock(monkeypatch):
|
||||
"""Regression: a member added between /team/member_update's permission checks and its write
|
||||
was dropped, because the new roster was built from the pre-check snapshot."""
|
||||
audit_logger = _wire_audit_log_callback(monkeypatch)
|
||||
|
||||
stale_snapshot = LiteLLM_TeamTable(
|
||||
team_id="team-race",
|
||||
team_alias="race",
|
||||
metadata={},
|
||||
members_with_roles=[Member(user_id="alice", role="admin"), Member(user_id="bob", role="user")],
|
||||
)
|
||||
team_row = LiteLLM_TeamTable(
|
||||
**{
|
||||
**stale_snapshot.model_dump(),
|
||||
"members_with_roles": [*stale_snapshot.members_with_roles, Member(user_id="carol", role="user")],
|
||||
}
|
||||
)
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
|
||||
mock_prisma_client.db.litellm_teamtable.update = AsyncMock(side_effect=_roster_writer(team_row))
|
||||
mock_prisma_client.db.litellm_auditlog.create = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
tx = _wire_member_delete_tx(mock_prisma_client)
|
||||
|
||||
team_info_patch, upsert_patch = _member_update_patches(stale_snapshot)
|
||||
with team_info_patch, upsert_patch:
|
||||
response = await team_member_update(
|
||||
data=TeamMemberUpdateRequest(team_id="team-race", user_id="bob", role="admin"),
|
||||
http_request=MagicMock(),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
|
||||
),
|
||||
)
|
||||
await _settle_audit_log_tasks()
|
||||
|
||||
assert {m.user_id: m.role for m in team_row.members_with_roles} == {
|
||||
"alice": "admin",
|
||||
"bob": "admin",
|
||||
"carol": "user",
|
||||
}
|
||||
assert response.team_id == "team-race" and response.user_id == "bob"
|
||||
assert tx.query_raw.await_args_list[0].args == (TEAM_ADVISORY_LOCK_SQL, "team-race"), (
|
||||
"the roster must be read only after the team advisory lock is held"
|
||||
)
|
||||
|
||||
[event] = _team_roster_events(audit_logger, "updated")
|
||||
assert _roster_user_roles(event["before_value"]) == {"alice": "admin", "bob": "user", "carol": "user"}
|
||||
assert _roster_user_roles(event["updated_values"]) == {"alice": "admin", "bob": "admin", "carol": "user"}
|
||||
assert _roster_team_alias(event["updated_values"]) == "race"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_update_role_change_404s_when_the_team_is_gone_under_the_lock(monkeypatch):
|
||||
_wire_audit_log_callback(monkeypatch)
|
||||
snapshot = LiteLLM_TeamTable(
|
||||
team_id="team-gone-race",
|
||||
team_alias="gone-race",
|
||||
metadata={},
|
||||
members_with_roles=[Member(user_id="bob", role="user")],
|
||||
)
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(side_effect=[snapshot, None])
|
||||
mock_prisma_client.db.litellm_teamtable.update = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
_wire_member_delete_tx(mock_prisma_client)
|
||||
|
||||
team_info_patch, upsert_patch = _member_update_patches(snapshot)
|
||||
with team_info_patch, upsert_patch, pytest.raises(HTTPException) as exc_info:
|
||||
await team_member_update(
|
||||
data=TeamMemberUpdateRequest(team_id="team-gone-race", user_id="bob", role="admin"),
|
||||
http_request=MagicMock(),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
|
||||
),
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
mock_prisma_client.db.litellm_teamtable.update.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_delete_response_does_not_wait_for_the_audit_insert(
|
||||
monkeypatch, mock_db_client, mock_admin_auth
|
||||
):
|
||||
"""Regression: the roster audit row was awaited on the request path, so a slow audit table
|
||||
held every /team/member_delete response."""
|
||||
from litellm.proxy._types import TeamMemberDeleteRequest
|
||||
|
||||
audit_logger = _wire_audit_log_callback(monkeypatch)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
|
||||
team_row = LiteLLM_TeamTable(
|
||||
team_id="team-slow-audit",
|
||||
team_alias="slow-audit",
|
||||
metadata={},
|
||||
members_with_roles=[Member(user_id="alice", role="admin"), Member(user_id="bob", role="user")],
|
||||
)
|
||||
mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
|
||||
mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_row)
|
||||
mock_db_client.db.litellm_usertable.find_many = AsyncMock(return_value=[])
|
||||
mock_db_client.db.litellm_teammembership.delete_many = AsyncMock(return_value=MagicMock())
|
||||
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock(return_value=MagicMock())
|
||||
_wire_member_delete_tx(mock_db_client)
|
||||
|
||||
audit_table_answers = asyncio.Event()
|
||||
audit_rows = []
|
||||
|
||||
async def _blocked_create(data):
|
||||
await audit_table_answers.wait()
|
||||
audit_rows.append(data)
|
||||
|
||||
mock_db_client.db.litellm_auditlog.create = AsyncMock(side_effect=_blocked_create)
|
||||
|
||||
await asyncio.wait_for(
|
||||
team_member_delete(
|
||||
data=TeamMemberDeleteRequest(team_id="team-slow-audit", user_id="bob"),
|
||||
user_api_key_dict=mock_admin_auth,
|
||||
),
|
||||
timeout=1,
|
||||
)
|
||||
assert audit_rows == [], "the response returned while the audit table was still blocked"
|
||||
|
||||
audit_table_answers.set()
|
||||
await _settle_audit_log_tasks()
|
||||
|
||||
assert [row["object_id"] for row in audit_rows] == ["team-slow-audit"]
|
||||
[event] = _team_roster_events(audit_logger, "updated")
|
||||
assert _roster_user_roles(event["before_value"]) == {"alice": "admin", "bob": "user"}
|
||||
assert _roster_user_roles(event["updated_values"]) == {"alice": "admin"}
|
||||
|
||||
|
||||
class _UntouchableRoster(Sequence[Member]):
|
||||
"""A roster that fails the test the moment anything reads it."""
|
||||
|
||||
def __getitem__(self, index):
|
||||
raise AssertionError("the roster was read while audit logging is off")
|
||||
|
||||
def __len__(self) -> int:
|
||||
raise AssertionError("the roster was read while audit logging is off")
|
||||
|
||||
|
||||
def test_membership_audit_scheduling_skips_the_roster_entirely_when_audit_logging_is_off(monkeypatch):
|
||||
"""Regression: the before/after rosters were serialized on every membership change, even
|
||||
when audit logs are not stored."""
|
||||
from litellm.proxy.management_endpoints.team_endpoints import _schedule_team_membership_audit_log
|
||||
|
||||
monkeypatch.setattr("litellm.store_audit_logs", False)
|
||||
|
||||
_schedule_team_membership_audit_log(
|
||||
team_id="team-quiet",
|
||||
team_alias="quiet",
|
||||
before_members=_UntouchableRoster(),
|
||||
after_members=_UntouchableRoster(),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user"),
|
||||
litellm_proxy_admin_name="admin",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_team_emits_only_the_deleted_audit_event(monkeypatch):
|
||||
from litellm.proxy._types import DeleteTeamRequest
|
||||
|
|
|
|||
|
|
@ -14,7 +14,10 @@ from litellm.proxy._types import (
|
|||
TeamMemberUpdateRequest,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.team_endpoints import team_member_update
|
||||
from litellm.proxy.management_endpoints.team_endpoints import (
|
||||
TEAM_ADVISORY_LOCK_SQL,
|
||||
team_member_update,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -65,13 +68,20 @@ def happy_path_upsert(monkeypatch):
|
|||
prisma_client.db.litellm_teamtable.update = AsyncMock()
|
||||
|
||||
class _FakeTx:
|
||||
litellm_teamtable = prisma_client.db.litellm_teamtable
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
return False
|
||||
|
||||
prisma_client.db.tx = MagicMock(return_value=_FakeTx())
|
||||
async def query_raw(self, sql, team_id):
|
||||
if sql == TEAM_ADVISORY_LOCK_SQL:
|
||||
return []
|
||||
return [{"members_with_roles": team_row.model_dump()["members_with_roles"]}]
|
||||
|
||||
prisma_client.tx = MagicMock(return_value=_FakeTx())
|
||||
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
|
||||
monkeypatch.setattr(proxy_server, "premium_user", False)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue