diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 89e31c22308..00ebb9d9aaa 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -134,11 +134,11 @@ async def _find_users_by_email( async def _create_user_row( - prisma_client: PrismaClient, tx: MemberWriteTx | None, user_data: Mapping[str, object] + prisma_client: PrismaClient, tx: MemberWriteTx | None, user_data: dict[str, object] ) -> _PrismaUserRecord | None: if tx is not None: - return await tx.litellm_usertable.create(data=jsonify_object(dict(user_data))) - return await prisma_client.insert_data(data=dict(user_data), table_name="user") + return await tx.litellm_usertable.create(data=jsonify_object(user_data)) + return await prisma_client.insert_data(data=user_data, table_name="user") def get_new_internal_user_defaults(user_id: str, user_email: str | None = None) -> dict[str, object]: diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index f01c1f62d6a..ff070de6017 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -4,7 +4,7 @@ from contextlib import asynccontextmanager from datetime import datetime, timezone from types import SimpleNamespace from typing import Optional, cast -from unittest.mock import AsyncMock, MagicMock, call, patch +from unittest.mock import AsyncMock, MagicMock, PropertyMock, call, patch import pytest from fastapi import HTTPException @@ -1954,11 +1954,28 @@ async def test_add_team_members_runs_member_writes_on_the_lock_holding_transacti _add_team_members_to_team, ) + added_user = MagicMock() + added_user.user_id = "bob" + added_user.model_dump.return_value = {"user_id": "bob", "teams": ["team-pool"]} + created_budget = MagicMock() + created_budget.budget_id = "budget-pool" + membership = MagicMock() + membership.model_dump.return_value = { + "team_id": "team-pool", + "user_id": "bob", + "budget_id": "budget-pool", + "litellm_budget_table": None, + } + tx = MagicMock() tx.query_raw = AsyncMock(return_value=[{"members_with_roles": []}]) tx.litellm_teamtable.update = AsyncMock( return_value=LiteLLM_TeamTable(team_id="team-pool", members_with_roles=[]) ) + tx.litellm_usertable.upsert = AsyncMock(return_value=added_user) + tx.litellm_usertable.update_many = AsyncMock() + tx.litellm_budgettable.create = AsyncMock(return_value=created_budget) + tx.litellm_teammembership.create = AsyncMock(return_value=membership) tx_cm = MagicMock() tx_cm.__aenter__ = AsyncMock(return_value=tx) @@ -1966,27 +1983,25 @@ async def test_add_team_members_runs_member_writes_on_the_lock_holding_transacti prisma_client = MagicMock() prisma_client.tx = MagicMock(return_value=tx_cm) - - process_team_members = AsyncMock(return_value=([], [])) - with patch( - "litellm.proxy.management_endpoints.team_endpoints._process_team_members", - new=process_team_members, - ): - await _add_team_members_to_team( - data=TeamMemberAddRequest( - team_id="team-pool", - member=Member(user_id="bob", role="user"), - ), - complete_team_data=LiteLLM_TeamTable(team_id="team-pool", members_with_roles=[]), - prisma_client=cast(object, prisma_client), - user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), - litellm_proxy_admin_name="admin", - ) - - assert process_team_members.call_args.kwargs["tx"] is tx, ( - "member writes must reuse the lock holder's connection, not check out another one" + type(prisma_client).db = PropertyMock( + side_effect=AssertionError("member writes must not reach for a second pooled connection") ) + _, updated_users, updated_team_memberships = await _add_team_members_to_team( + data=TeamMemberAddRequest( + team_id="team-pool", + member=Member(user_id="bob", role="user"), + max_budget_in_team=50.0, + ), + complete_team_data=LiteLLM_TeamTable(team_id="team-pool", members_with_roles=[]), + prisma_client=cast(object, prisma_client), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + litellm_proxy_admin_name="admin", + ) + + assert [user.user_id for user in updated_users] == ["bob"] + assert [tm.budget_id for tm in updated_team_memberships] == ["budget-pool"] + @pytest.mark.asyncio async def test_add_team_members_writes_nothing_when_the_team_is_deleted_mid_request():