scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-29 01:42:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

test: alias replica_db on remaining prisma test doubles

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yuneng 2026-09-24 09:33:38 +00:00
parent 74d9fe2ba0
commit e1f9382378
16 changed files with 191 additions and 50 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
tests/proxy_unit_tests/test_user_api_key_auth.py
View file

@ -208,6 +208,8 @@ class _NoMembershipRowPrisma:
async def find_unique(where: dict[str, dict[str, str]], include: dict[str, bool]) -> None:
return None
replica_db = db
@pytest.mark.parametrize("key_ownership", ["user_key", "team_key"])
@pytest.mark.asyncio

73
tests/test_litellm/containers/test_container_proxy_ownership.py
View file

@ -42,9 +42,8 @@ def _container(container_id: str) -> ContainerObject:
async def test_should_record_container_owner_with_original_provider_id(monkeypatch):
table = AsyncMock()
table.find_unique.return_value = None
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
tables = SimpleNamespace(litellm_managedobjecttable=table)
prisma_client = SimpleNamespace(db=tables, replica_db=tables)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
@ -73,9 +72,8 @@ async def test_should_not_mutate_dict_container_response_when_recording_owner(
):
table = AsyncMock()
table.find_unique.return_value = None
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
tables = SimpleNamespace(litellm_managedobjecttable=table)
prisma_client = SimpleNamespace(db=tables, replica_db=tables)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
@ -101,9 +99,8 @@ async def test_should_not_mutate_dict_container_response_when_recording_owner(
async def test_should_record_team_owner_for_keys_without_user_id(monkeypatch):
table = AsyncMock()
table.find_unique.return_value = None
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
tables = SimpleNamespace(litellm_managedobjecttable=table)
prisma_client = SimpleNamespace(db=tables, replica_db=tables)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
@ -126,9 +123,8 @@ async def test_should_record_team_owner_for_keys_without_user_id(monkeypatch):
async def test_should_record_token_owner_for_keys_without_user_team_or_org(monkeypatch):
table = AsyncMock()
table.find_unique.return_value = None
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
tables = SimpleNamespace(litellm_managedobjecttable=table)
prisma_client = SimpleNamespace(db=tables, replica_db=tables)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
@ -174,9 +170,8 @@ async def test_should_reject_record_for_identityless_proxy_auth(monkeypatch):
@pytest.mark.asyncio
async def test_should_skip_owner_record_when_provider_response_has_no_id(monkeypatch):
table = AsyncMock()
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
tables = SimpleNamespace(litellm_managedobjecttable=table)
prisma_client = SimpleNamespace(db=tables, replica_db=tables)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
@ -221,9 +216,8 @@ async def test_should_not_reassign_existing_container_to_different_owner(monkeyp
file_purpose=ownership.CONTAINER_OBJECT_PURPOSE,
created_by="user-2",
)
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
tables = SimpleNamespace(litellm_managedobjecttable=table)
prisma_client = SimpleNamespace(db=tables, replica_db=tables)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
@ -245,10 +239,11 @@ async def test_should_not_reassign_existing_container_to_different_owner(monkeyp
def _owned_containers_in_db(monkeypatch, *model_object_ids: str) -> AsyncMock:
table = AsyncMock()
table.find_many.return_value = [SimpleNamespace(model_object_id=object_id) for object_id in model_object_ids]
tables = SimpleNamespace(litellm_managedobjecttable=table)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
AsyncMock(return_value=SimpleNamespace(db=SimpleNamespace(litellm_managedobjecttable=table))),
AsyncMock(return_value=SimpleNamespace(db=tables, replica_db=tables)),
)
return table
@ -815,9 +810,8 @@ async def test_get_container_owner_uses_cache_after_first_db_hit(monkeypatch):
created_by="user-1", file_purpose=ownership.CONTAINER_OBJECT_PURPOSE
)
table.find_first.return_value = fake_row
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
tables = SimpleNamespace(litellm_managedobjecttable=table)
prisma_client = SimpleNamespace(db=tables, replica_db=tables)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
@ -840,9 +834,8 @@ async def test_get_container_owner_caches_negative_lookups(monkeypatch):
"""`None` (untracked) must also be cached so repeated misses don't query."""
table = AsyncMock()
table.find_first.return_value = None
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
tables = SimpleNamespace(litellm_managedobjecttable=table)
prisma_client = SimpleNamespace(db=tables, replica_db=tables)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
@ -862,9 +855,8 @@ async def test_allowed_container_ids_uses_cache_after_first_db_hit(monkeypatch):
table.find_many.return_value = [
SimpleNamespace(model_object_id="container:openai:cntr_a"),
]
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
tables = SimpleNamespace(litellm_managedobjecttable=table)
prisma_client = SimpleNamespace(db=tables, replica_db=tables)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
@ -894,9 +886,8 @@ async def test_record_container_owner_invalidates_caller_list_cache(monkeypatch)
table.find_many.return_value = [
SimpleNamespace(model_object_id="container:openai:cntr_old"),
]
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
tables = SimpleNamespace(litellm_managedobjecttable=table)
prisma_client = SimpleNamespace(db=tables, replica_db=tables)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
@ -934,9 +925,8 @@ async def test_admin_with_identity_records_container_ownership(monkeypatch):
untracked."""
table = AsyncMock()
table.find_unique.return_value = None
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
tables = SimpleNamespace(litellm_managedobjecttable=table)
prisma_client = SimpleNamespace(db=tables, replica_db=tables)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
@ -964,9 +954,8 @@ async def test_should_record_containers_from_responses_output_for_service_accoun
):
table = AsyncMock()
table.find_unique.return_value = None
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
tables = SimpleNamespace(litellm_managedobjecttable=table)
prisma_client = SimpleNamespace(db=tables, replica_db=tables)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
@ -1019,9 +1008,8 @@ async def test_service_account_can_access_container_after_responses_tracking(
)
table = AsyncMock()
table.find_unique.return_value = None
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
tables = SimpleNamespace(litellm_managedobjecttable=table)
prisma_client = SimpleNamespace(db=tables, replica_db=tables)
monkeypatch.setattr(
ownership,
"_get_prisma_client",
@ -1089,9 +1077,8 @@ async def test_should_record_container_ownership_after_streaming_responses_finis
table = AsyncMock()
table.find_unique.return_value = None
prisma_client = SimpleNamespace(
db=SimpleNamespace(litellm_managedobjecttable=table)
)
tables = SimpleNamespace(litellm_managedobjecttable=table)
prisma_client = SimpleNamespace(db=tables, replica_db=tables)
monkeypatch.setattr(
ownership,
"_get_prisma_client",

13
tests/test_litellm/enterprise/proxy/test_managed_files_hook.py
View file

@ -183,6 +183,7 @@ def _make_managed_files_instance():
mock_cache = MagicMock()
mock_prisma = MagicMock()
mock_prisma.replica_db = mock_prisma.db
instance = _PROXY_LiteLLMManagedFiles(
internal_usage_cache=mock_cache,
@ -1122,7 +1123,7 @@ def _managed_files_with_deletion_row(unified_file_id, provider_file_id, file_obj
)
return _PROXY_LiteLLMManagedFiles(
internal_usage_cache=DualCache(),
prisma_client=MagicMock(db=MagicMock(litellm_managedfiletable=table)),
prisma_client=MagicMock(db=MagicMock(litellm_managedfiletable=table), replica_db=MagicMock(litellm_managedfiletable=table)),
), table
@ -1302,6 +1303,7 @@ def _make_real_managed_files_instance():
mock_cache.async_set_cache = AsyncMock()
mock_prisma = MagicMock()
mock_prisma.replica_db = mock_prisma.db
mock_prisma.db.litellm_managedfiletable.upsert = AsyncMock()
mock_prisma.db.litellm_managedfiletable.create = AsyncMock(
side_effect=AssertionError(
@ -1329,6 +1331,7 @@ def _make_object_store_instance():
mock_cache.async_set_cache = AsyncMock()
mock_prisma = MagicMock()
mock_prisma.replica_db = mock_prisma.db
mock_prisma.db.litellm_managedobjecttable.upsert = AsyncMock()
mock_prisma.db.litellm_managedobjecttable.update_many = AsyncMock()
@ -1891,7 +1894,8 @@ async def test_afile_delete_storage_backed_row_deletes_stored_content_not_provid
managed_files = _PROXY_LiteLLMManagedFiles(
internal_usage_cache=DualCache(),
prisma_client=MagicMock(
db=MagicMock(litellm_managedfiletable=file_table, litellm_managedfilecontenttable=content_table)
db=MagicMock(litellm_managedfiletable=file_table, litellm_managedfilecontenttable=content_table),
replica_db=MagicMock(litellm_managedfiletable=file_table, litellm_managedfilecontenttable=content_table),
),
)
router = MagicMock(
@ -1935,7 +1939,8 @@ async def test_afile_content_storage_backed_row_returns_stored_bytes_not_provide
managed_files = _PROXY_LiteLLMManagedFiles(
internal_usage_cache=DualCache(),
prisma_client=MagicMock(
db=MagicMock(litellm_managedfiletable=file_table, litellm_managedfilecontenttable=content_table)
db=MagicMock(litellm_managedfiletable=file_table, litellm_managedfilecontenttable=content_table),
replica_db=MagicMock(litellm_managedfiletable=file_table, litellm_managedfilecontenttable=content_table),
),
)
router = MagicMock(
@ -1992,7 +1997,7 @@ async def test_store_unified_file_id_caches_the_storage_location_the_db_row_gets
file_table = MagicMock(upsert=AsyncMock(), find_first=AsyncMock(side_effect=AssertionError("cache miss")))
managed_files = _PROXY_LiteLLMManagedFiles(
internal_usage_cache=DualCache(),
prisma_client=MagicMock(db=MagicMock(litellm_managedfiletable=file_table)),
prisma_client=MagicMock(db=MagicMock(litellm_managedfiletable=file_table), replica_db=MagicMock(litellm_managedfiletable=file_table)),
)
stored = _make_file_object("file-kept").model_copy(update={"purpose": "batch"})
stored._hidden_params = {"storage_backend": "litellm_db", "storage_url": "litellm_db://content-row-1"}

8
tests/test_litellm/integrations/test_prometheus_user_team_metrics.py
View file

@ -211,6 +211,7 @@ class TestPrometheusUserTeamCountMetrics:
return 2 if where is not None else 10
mock_prisma = MagicMock()
mock_prisma.replica_db = mock_prisma.db
mock_prisma.db.litellm_usertable.count = _user_count
mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=4)
@ -840,6 +841,7 @@ async def test_set_org_budget_metrics_after_api_request(prometheus_logger):
org_mock.litellm_budget_table = budget_mock
mock_prisma = MagicMock()
mock_prisma.replica_db = mock_prisma.db
mock_proxy_server = MagicMock()
mock_proxy_server.prisma_client = mock_prisma
mock_proxy_server.user_api_key_cache = MagicMock()
@ -906,6 +908,7 @@ async def test_initialize_org_budget_metrics(prometheus_logger):
org_mock.litellm_budget_table = budget_mock
mock_prisma = MagicMock()
mock_prisma.replica_db = mock_prisma.db
mock_prisma.db.litellm_organizationtable.find_many = AsyncMock(
return_value=[org_mock]
)
@ -1186,6 +1189,7 @@ async def test_initialize_customer_budget_metrics_emits_gauges_for_budgeted_cust
]
find_many = AsyncMock(return_value=rows)
mock_prisma = MagicMock()
mock_prisma.replica_db = mock_prisma.db
mock_prisma.db.litellm_endusertable.find_many = find_many
mock_prisma.db.litellm_endusertable.count = AsyncMock(return_value=len(rows))
mock_proxy_server = MagicMock()
@ -1220,6 +1224,7 @@ async def test_initialize_customer_budget_metrics_skips_when_end_user_tracking_o
find_many = AsyncMock(return_value=[_mock_customer_row("cust-a", 100.0, 500.0, None)])
mock_prisma = MagicMock()
mock_prisma.replica_db = mock_prisma.db
mock_prisma.db.litellm_endusertable.find_many = find_many
mock_prisma.db.litellm_endusertable.count = AsyncMock(return_value=1)
mock_proxy_server = MagicMock()
@ -1237,6 +1242,7 @@ async def test_initialize_remaining_budget_metrics_includes_customers(prometheus
import sys
mock_prisma = MagicMock()
mock_prisma.replica_db = mock_prisma.db
mock_prisma.db.litellm_endusertable.find_many = AsyncMock(
return_value=[_mock_customer_row("cust-startup", 5.0, 25.0, None)]
)
@ -1262,6 +1268,7 @@ async def test_initialize_customer_budget_metrics_counts_once_across_pages(prome
find_many = AsyncMock(side_effect=pages)
count = AsyncMock(return_value=101)
mock_prisma = MagicMock()
mock_prisma.replica_db = mock_prisma.db
mock_prisma.db.litellm_endusertable.find_many = find_many
mock_prisma.db.litellm_endusertable.count = count
mock_proxy_server = MagicMock()
@ -1294,6 +1301,7 @@ async def test_initialize_customer_budget_metrics_applies_default_budget_to_unbu
find_many = AsyncMock(return_value=[explicit_row, default_row])
find_unique = AsyncMock(return_value=default_budget)
mock_prisma = MagicMock()
mock_prisma.replica_db = mock_prisma.db
mock_prisma.db.litellm_endusertable.find_many = find_many
mock_prisma.db.litellm_endusertable.count = AsyncMock(return_value=2)
mock_prisma.db.litellm_budgettable.find_unique = find_unique

1
tests/test_litellm/integrations/test_responses_background_cost.py
View file

@ -281,6 +281,7 @@ class TestCheckResponsesCost:
"""Create a mock Prisma client"""
client = MagicMock()
client.db = MagicMock()
client.replica_db = client.db
client.db.litellm_managedobjecttable = MagicMock()
return client

3
tests/test_litellm/llms/base_llm/files/test_litellm_db_storage_backend.py
View file

@ -14,7 +14,8 @@ from litellm.llms.base_llm.files.litellm_db_storage_backend import (
def _backend_with_table():
table = MagicMock(create=AsyncMock(), find_unique=AsyncMock(), delete=AsyncMock())
prisma_client = MagicMock(db=MagicMock(litellm_managedfilecontenttable=table))
tables = MagicMock(litellm_managedfilecontenttable=table)
prisma_client = MagicMock(db=tables, replica_db=tables)
return LiteLLMDbStorageBackend(prisma_client), table

3
tests/test_litellm/llms/base_llm/files/test_storage_backend_factory.py
View file

@ -14,7 +14,8 @@ from litellm.llms.base_llm.files.storage_backend_factory import get_storage_back
@pytest.mark.asyncio
async def test_litellm_db_backend_stores_through_the_given_prisma_client():
table = MagicMock(create=AsyncMock(return_value=SimpleNamespace(id="row-1")))
prisma_client = MagicMock(db=MagicMock(litellm_managedfilecontenttable=table))
tables = MagicMock(litellm_managedfilecontenttable=table)
prisma_client = MagicMock(db=tables, replica_db=tables)
backend = get_storage_backend(LITELLM_DB_STORAGE_BACKEND_NAME, prisma_client=prisma_client)

3
tests/test_litellm/llms/base_llm/test_base_managed_resource.py
View file

@ -40,9 +40,11 @@ def _make_resource(records: List = None) -> _StubResource:
cache.async_get_cache = AsyncMock(return_value=None)
prisma = MagicMock()
prisma.replica_db = prisma.db
table = MagicMock()
table.find_many = AsyncMock(return_value=records or [])
prisma.db = MagicMock()
prisma.replica_db = prisma.db
setattr(prisma.db, "litellm_test_resource_table", table)
return _StubResource(internal_usage_cache=cache, prisma_client=prisma)
@ -122,6 +124,7 @@ async def test_can_access_uses_team_id_for_service_account(caller_team_id, expec
}
)
prisma = MagicMock()
prisma.replica_db = prisma.db
resource = _StubResource(internal_usage_cache=cache, prisma_client=prisma)
caller = (

7
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -58,6 +58,7 @@ async def test_vendor_credential_state_agrees_with_egress_token_resolution(monke
read = AsyncMock(return_value=credential)
refresh = AsyncMock(return_value=_stored_grant(access_token="fresh-token", expires_in_seconds=3600))
prisma = MagicMock()
prisma.replica_db = prisma.db
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
monkeypatch.setattr(mcp_db, "get_user_oauth_credential", read)
monkeypatch.setattr(mcp_db, "refresh_user_oauth_token", refresh)
@ -7847,6 +7848,7 @@ async def test_load_active_user_by_id_reads_the_row_from_the_database_not_the_ca
key="fresh-jwt-user", value=LiteLLM_UserTable(user_id="fresh-jwt-user", teams=[]), model_type=LiteLLM_UserTable
)
prisma = MagicMock()
prisma.replica_db = prisma.db
prisma.db.litellm_usertable.find_unique = AsyncMock(
return_value=LiteLLM_UserTable(user_id="fresh-jwt-user", teams=["team-a"])
)
@ -7881,6 +7883,7 @@ async def test_load_active_user_by_id_serves_a_cached_row_without_a_database_rea
model_type=LiteLLM_UserTable,
)
prisma = MagicMock()
prisma.replica_db = prisma.db
prisma.db.litellm_usertable.find_unique = AsyncMock(
return_value=LiteLLM_UserTable(user_id="cached-jwt-user", teams=[])
)
@ -9341,6 +9344,7 @@ async def test_reload_servers_from_database_hydrates_dcr_clients():
)
prisma = MagicMock()
prisma.replica_db = prisma.db
prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[])
hydrate_spy = AsyncMock()
@ -11780,6 +11784,7 @@ def jwt_oauth_identity(monkeypatch: pytest.MonkeyPatch) -> tuple["JWTHandler", "
monkeypatch.setattr(proxy_server, "premium_user", True)
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
prisma: Final = MagicMock()
prisma.replica_db = prisma.db
prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=None)
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
return handler, signing_key
@ -11869,6 +11874,7 @@ async def test_oauth_exchange_stores_token_for_validated_jwt_user(
return httpx.Response(200, json={"access_token": "upstream-token", "token_type": "Bearer"})
database: Final = MagicMock()
database.replica_db = database.db
users: Final = database.db.litellm_usertable
users.find_unique = AsyncMock(return_value=None)
users.find_first = AsyncMock(return_value=None)
@ -12149,6 +12155,7 @@ async def test_oauth_jwt_resolves_canonical_owner_without_cached_identity(
organization_memberships=[],
)
database: Final = MagicMock()
database.replica_db = database.db
table: Final = database.db.litellm_usertable
table.find_unique = AsyncMock(side_effect=[None, owner if identity == "sso" else None])
table.find_first = AsyncMock(return_value=owner)

10
tests/test_litellm/proxy/auth/test_handle_jwt.py
View file

@ -726,6 +726,7 @@ async def test_sync_user_role_and_teams():
)
prisma = AsyncMock()
prisma.replica_db = prisma.db
prisma.db.litellm_usertable.update = AsyncMock()
with patch(
@ -769,6 +770,7 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_role_change():
)
prisma = AsyncMock()
prisma.replica_db = prisma.db
prisma.db.litellm_usertable.update = AsyncMock()
await JWTAuthManager.sync_user_role_and_teams(
@ -812,6 +814,7 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_team_change():
)
prisma = AsyncMock()
prisma.replica_db = prisma.db
prisma.db.litellm_usertable.update = AsyncMock()
with patch(
@ -859,6 +862,7 @@ async def test_sync_user_role_and_teams_no_cache_write_when_nothing_changes():
)
prisma = AsyncMock()
prisma.replica_db = prisma.db
await JWTAuthManager.sync_user_role_and_teams(
jwt_handler, token, user, prisma, user_api_key_cache=mock_cache
@ -3371,6 +3375,8 @@ class _UnreachableMembershipPrisma:
async def find_unique(where: dict[str, dict[str, str]], include: dict[str, bool]) -> None:
raise httpx.ConnectError("All connection attempts failed")
replica_db = db
@pytest.mark.asyncio
async def test_auth_builder_single_team_fallback_membership_outage_raises_instead_of_dropping_the_team():
@ -5639,6 +5645,7 @@ async def test_sync_user_role_and_teams_no_claim_team_preservation(
teams=["team_a", "team_b"],
)
prisma = AsyncMock()
prisma.replica_db = prisma.db
with patch(
"litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership",
@ -7343,6 +7350,7 @@ async def test_jwt_identity_and_authorization_keep_provisioning_in_admission(
if existing_user:
cache.set_cache(user_id, user)
database: Final = MagicMock()
database.replica_db = database.db
users: Final = database.db.litellm_usertable
users.find_unique = AsyncMock(return_value=None)
users.find_first = AsyncMock(return_value=None)
@ -7654,6 +7662,7 @@ async def test_admin_jwt_team_header_only_provisions_during_admission(monkeypatc
handler.litellm_jwtauth.team_id_upsert = True
handler.litellm_jwtauth.admin_allowed_routes = ["openai_routes"]
database = MagicMock()
database.replica_db = database.db
database.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
create_team = AsyncMock(return_value=LiteLLM_TeamTable(team_id="new-team").model_dump())
monkeypatch.setattr(team_endpoints, "new_team", create_team)
@ -7692,6 +7701,7 @@ async def test_scope_admin_admission_resolves_existing_user_without_provisioning
if existing_user and warm_cache:
cache.set_cache(user_id, user)
database: Final = MagicMock()
database.replica_db = database.db
users: Final = database.db.litellm_usertable
users.find_unique = AsyncMock(return_value=user if existing_user else None)
users.find_first = AsyncMock(return_value=None)

21
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -142,6 +142,7 @@ def _tool_call_response(*names: str) -> object:
def _tool_usage_prisma() -> MagicMock:
prisma = MagicMock()
prisma.replica_db = prisma.db
prisma.tool_usage_transactions = []
prisma._tool_usage_transactions_lock = asyncio.Lock()
prisma.spend_log_transactions = []
@ -316,6 +317,7 @@ class _RecordingDb:
class _RecordingPrisma:
def __init__(self, execute_raw: Callable[[], int] | None = None) -> None:
self.db = _RecordingDb(execute_raw=execute_raw)
self.replica_db = self.db
def _row_values(statement: Statement, column: str) -> list[object]:
@ -903,6 +905,7 @@ async def test_commit_spend_updates_to_db_increments_agent_spend():
mock_prisma_client = MagicMock()
mock_prisma_client.db = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction)
mock_proxy_logging = MagicMock()
@ -943,6 +946,7 @@ def _team_member_flush_fixtures() -> tuple[AsyncMock, MagicMock]:
mock_prisma_client = MagicMock()
mock_prisma_client.db = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction)
return mock_transaction, mock_prisma_client
@ -1064,6 +1068,7 @@ async def test_org_spend_increments_organization_membership_row_for_the_calling_
mock_batcher: Final = MagicMock()
mock_prisma_client: Final = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(mock_batcher))
proxy_logging: Final = MagicMock()
proxy_logging.call_details = {}
@ -1098,6 +1103,7 @@ async def test_org_spend_without_user_id_leaves_organization_membership_untouche
mock_batcher: Final = MagicMock()
mock_prisma_client: Final = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(mock_batcher))
proxy_logging: Final = MagicMock()
proxy_logging.call_details = {}
@ -1126,6 +1132,7 @@ async def test_org_spend_keeps_member_attribution_when_ids_contain_the_key_delim
mock_batcher: Final = MagicMock()
mock_prisma_client: Final = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(mock_batcher))
proxy_logging: Final = MagicMock()
proxy_logging.call_details = {}
@ -1196,6 +1203,7 @@ async def test_project_spend_is_persisted_to_project_table_and_project_cache_is_
mock_batcher: Final = MagicMock()
mock_prisma_client: Final = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(mock_batcher))
user_api_key_cache: Final = MagicMock()
user_api_key_cache.async_delete_cache = AsyncMock()
@ -1842,6 +1850,7 @@ async def test_commit_key_spend_updates_includes_last_active():
mock_prisma_client = MagicMock()
mock_prisma_client.db = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction)
# Also mock the other table batchers to avoid errors
@ -1920,6 +1929,7 @@ async def test_commit_spend_updates_to_db_increments_key_total_spend_alongside_s
mock_prisma_client = MagicMock()
mock_prisma_client.db = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction)
db_spend_update_transactions = {
@ -2136,6 +2146,7 @@ async def test_commit_spend_updates_uses_pipeline():
db_writer.pod_lock_manager = mock_pod_lock_manager
mock_prisma_client = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_proxy_logging = MagicMock()
await db_writer._commit_spend_updates_to_db_with_redis(
@ -2481,6 +2492,7 @@ async def test_commit_spend_updates_iterates_in_sorted_order(
mock_prisma_client = MagicMock()
mock_prisma_client.db = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction)
mock_proxy_logging = MagicMock()
@ -2866,6 +2878,7 @@ class _WindowSpendFakeDB:
class _WindowSpendFakePrisma:
def __init__(self, db):
self.db = db
self.replica_db = self.db
def _window_spend_upserts(db):
@ -3282,6 +3295,7 @@ async def test_failed_per_entity_increment_from_redis_restores_only_what_may_sti
return False
mock_prisma_client = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.tx = MagicMock(return_value=_Transaction())
proxy_logging_obj = MagicMock()
proxy_logging_obj.failure_handler = AsyncMock()
@ -3425,6 +3439,7 @@ async def test_commit_spend_updates_to_db_does_not_stamp_key_settings_updated_at
mock_prisma_client = MagicMock()
mock_prisma_client.db = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction)
token = "hashed-token-abc"
@ -3459,6 +3474,7 @@ async def test_commit_spend_updates_to_db_does_not_stamp_key_settings_updated_at
async def test_commit_spend_updates_to_db_reports_each_completed_table():
db_writer = DBSpendUpdateWriter()
mock_prisma_client = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(MagicMock()))
proxy_logging_obj = MagicMock()
proxy_logging_obj.call_details = {}
@ -3512,6 +3528,7 @@ async def test_commit_spend_updates_to_db_reports_table_committed_before_cache_i
):
db_writer = DBSpendUpdateWriter()
mock_prisma_client = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(MagicMock()))
user_api_key_cache = MagicMock()
user_api_key_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis down"))
@ -3704,6 +3721,7 @@ async def test_commit_spend_updates_retries_deadlock_then_commits(monkeypatch):
mock_batcher = MagicMock()
mock_prisma_client = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.tx = MagicMock(side_effect=[_failing_tx(_deadlock_error()), _good_tx(mock_batcher)])
proxy_logging = MagicMock()
@ -3731,6 +3749,7 @@ async def test_commit_spend_updates_raises_after_exhausting_deadlock_retries(mon
monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", AsyncMock(return_value=None))
mock_prisma_client = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.tx = MagicMock(side_effect=lambda *a, **k: _failing_tx(_deadlock_error()))
proxy_logging = MagicMock()
@ -3760,6 +3779,7 @@ async def test_commit_spend_updates_does_not_retry_non_deadlock_data_error(monke
data={"user_facing_error": {"error_code": "P2002", "meta": {"table": "LiteLLM_VerificationToken"}}}
)
mock_prisma_client = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.tx = MagicMock(side_effect=lambda *a, **k: _failing_tx(non_deadlock))
proxy_logging = MagicMock()
@ -3826,6 +3846,7 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey
mock_batcher = MagicMock()
mock_prisma_client = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.tx = MagicMock(side_effect=[_failing_tx(_deadlock_error()), _good_tx(mock_batcher)])
proxy_logging = MagicMock()

67
tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py
View file

@ -84,6 +84,7 @@ async def test_create_user_existing_user_conflict(mocker):
# Create a properly structured mock for the prisma client
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value={"user_id": "existing-user"})
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
@ -121,6 +122,7 @@ async def test_create_user_defaults_to_viewer(mocker, monkeypatch):
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
@ -172,6 +174,7 @@ async def test_create_user_ingests_enterprise_extension(mocker, monkeypatch):
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
@ -229,6 +232,7 @@ async def test_create_user_ingests_entitlements_and_roles(mocker, monkeypatch):
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
@ -279,6 +283,7 @@ async def test_create_user_uses_default_internal_user_params_role(mocker, monkey
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
@ -313,6 +318,7 @@ async def test_create_user_uses_default_internal_user_params_role(mocker, monkey
def _mock_scim_create_user_deps(mocker: MockerFixture, scim_user: SCIMUser) -> AsyncMock:
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
@ -458,6 +464,7 @@ async def test_scim_create_user_respects_default_role_set_via_ui(mocker, monkeyp
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
@ -508,6 +515,7 @@ async def test_scim_collection_endpoints_clamp_requested_page_size(
"""SCIM list endpoints accept zero and cap larger client page requests."""
mock_prisma_client = MagicMock()
mock_prisma_client.db = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
table = MagicMock()
table.find_many = AsyncMock(return_value=[])
table.count = AsyncMock(return_value=0)
@ -548,6 +556,7 @@ async def test_get_users_filters_username_by_exposed_scim_username_for_okta(mock
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[user])
mock_prisma_client.db.litellm_usertable.count = AsyncMock(return_value=1)
@ -607,6 +616,7 @@ async def test_get_users_filters_email_value_by_user_email(mocker):
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=[user])
mock_prisma_client.db.litellm_usertable.count = AsyncMock(return_value=1)
@ -649,6 +659,7 @@ async def test_get_users_filters_email_value_by_user_email(mocker):
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()
mock_prisma_client.replica_db = mock_prisma_client.db
new_user_request = NewUserRequest(
user_id="test-user",
@ -671,6 +682,7 @@ 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.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
@ -724,6 +736,7 @@ async def test_handle_existing_user_by_email_existing_user_updated(mocker):
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
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)
@ -791,6 +804,7 @@ async def test_handle_existing_user_by_email_roster_changes_use_existing_user_id
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
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={})
@ -845,6 +859,7 @@ async def test_handle_existing_user_by_email_syncs_roster_and_dedups_teams(mocke
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
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={})
@ -901,6 +916,7 @@ async def test_handle_existing_user_by_email_without_teams_preserves_memberships
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
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={})
@ -962,6 +978,7 @@ async def test_handle_existing_user_by_email_roster_add_failure_blocks_teams_wri
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
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={})
@ -1003,6 +1020,7 @@ async def test_handle_existing_user_by_email_roster_add_already_member_is_noop(m
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
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={})
@ -1055,6 +1073,7 @@ async def test_handle_existing_user_by_email_roster_remove_failure_blocks_teams_
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
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={})
@ -1101,6 +1120,7 @@ async def test_handle_existing_user_by_email_roster_remove_already_absent_is_noo
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
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={})
@ -1269,6 +1289,7 @@ async def test_update_user_success(mocker):
# Mock prisma client
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user)
@ -1326,6 +1347,7 @@ async def test_update_user_put_with_valueless_entitlements_deactivates_user(scim
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user)
@ -1405,6 +1427,7 @@ async def test_update_user_without_groups_preserves_memberships_and_role(mocker,
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={"user_id": "okta-user"})
@ -1495,6 +1518,7 @@ async def test_patch_user_success(mocker):
# Mock prisma client
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user)
@ -1611,6 +1635,7 @@ async def test_update_group_metadata_serialization_issue(mocker):
# Create a properly structured mock for the prisma client
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
@ -1800,6 +1825,7 @@ async def test_update_group_e2e(mocker):
# Mock prisma client
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
@ -1965,6 +1991,7 @@ async def test_create_group_with_nonexistent_users_rejects(mocker, monkeypatch):
# Mock prisma client
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
@ -2040,6 +2067,7 @@ async def test_update_group_with_nonexistent_users_rejects(mocker, monkeypatch):
# Mock prisma client
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
@ -2120,6 +2148,7 @@ async def test_create_group_with_nonexistent_users_creates_when_flag_true(mocker
# Mock prisma client
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
@ -2212,6 +2241,7 @@ async def test_extract_group_member_ids_with_flag_true_creates_users(mocker, mon
# Mock prisma client
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
@ -2281,6 +2311,7 @@ async def test_extract_group_member_ids_with_flag_false_rejects(mocker, monkeypa
# Mock prisma client
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
@ -2340,6 +2371,7 @@ async def test_process_group_patch_operations_with_flag_true_creates_users(mocke
# Mock prisma client
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
# Mock user lookup - new-user-1 doesn't exist
@ -2397,6 +2429,7 @@ async def test_process_group_patch_operations_with_flag_false_rejects(mocker, mo
# Mock prisma client
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
# Mock user lookup - new-user-1 doesn't exist
@ -2440,6 +2473,7 @@ async def test_create_user_grants_admin_when_in_scim_admin_group(mocker, monkeyp
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
@ -2485,6 +2519,7 @@ async def test_create_user_keeps_default_when_not_in_scim_admin_group(mocker, mo
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
@ -2543,6 +2578,7 @@ async def test_update_user_demotes_admin_when_removed_from_scim_admin_group(mock
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user)
@ -2602,6 +2638,7 @@ async def test_update_user_does_not_force_role_when_scim_admin_group_unset(mocke
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user)
@ -2662,6 +2699,7 @@ async def test_update_user_demotes_when_default_params_lack_user_role(mocker, mo
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user)
@ -2721,6 +2759,7 @@ async def test_patch_user_demotes_admin_when_removed_from_scim_admin_group(mocke
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user)
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
@ -2788,6 +2827,7 @@ async def test_patch_user_grants_admin_by_team_display_name(mocker, monkeypatch)
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=updated_user)
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
@ -2835,6 +2875,7 @@ def _scim_admin_prisma(mocker, *, user_teams):
prisma = mocker.MagicMock()
prisma.db = mocker.MagicMock()
prisma.replica_db = prisma.db
prisma.db.litellm_usertable = mocker.MagicMock()
prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user)
prisma.db.litellm_usertable.find_many = AsyncMock(return_value=())
@ -2931,6 +2972,7 @@ async def test_update_group_recomputes_roles_for_changed_members(mocker):
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
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.update = AsyncMock(return_value=existing_team)
@ -2989,6 +3031,7 @@ async def test_patch_group_recomputes_roles_for_changed_members(mocker):
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
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.update = AsyncMock(return_value=existing_team)
@ -3043,6 +3086,7 @@ async def test_delete_group_recomputes_roles_for_members(mocker):
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
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()
@ -3075,6 +3119,7 @@ async def test_handle_existing_user_by_email_applies_role_when_admin_group_set(m
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
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={"user_id": "new-user-id"})
@ -3115,6 +3160,7 @@ async def test_handle_existing_user_by_email_leaves_role_when_admin_group_unset(
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
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={"user_id": "new-user-id"})
@ -3170,6 +3216,7 @@ async def test_create_user_existing_email_upsert_demotes_when_admin_group_set(mo
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
@ -3218,6 +3265,7 @@ async def test_create_group_recomputes_roles_for_members(mocker):
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
@ -3275,6 +3323,7 @@ async def test_update_group_rename_recomputes_retained_members(mocker):
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
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.update = AsyncMock(return_value=existing_team)
@ -3330,6 +3379,7 @@ async def test_patch_group_rename_recomputes_retained_members(mocker):
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
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.update = AsyncMock(return_value=existing_team)
@ -3398,6 +3448,7 @@ async def test_process_group_patch_operations_add_retains_existing_members(mocke
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
# new-user already exists in the DB
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=(mocker.MagicMock(user_id="new-user"),))
@ -3439,6 +3490,7 @@ async def test_process_group_patch_operations_remove_uses_members_with_roles(moc
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=mocker.MagicMock(user_id="drop-user"))
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
@ -3466,6 +3518,7 @@ async def test_get_groups_reports_members_from_members_with_roles(mocker):
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[team])
mock_prisma_client.db.litellm_teamtable.count = AsyncMock(return_value=1)
@ -3496,6 +3549,7 @@ async def test_apply_group_patch_updates_does_not_write_legacy_members(mocker):
"""
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
updated = mocker.MagicMock()
mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=updated)
@ -3516,6 +3570,7 @@ 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.replica_db = mock_prisma_client.db
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()
@ -3687,6 +3742,7 @@ async def test_patch_group_add_applies_delta_and_keeps_concurrent_add(mocker):
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
side_effect=[snapshot_team, refreshed_team, final_team]
@ -3780,6 +3836,7 @@ async def test_patch_group_replace_stays_absolute_against_concurrent_roster(mock
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
side_effect=[snapshot_team, refreshed_team, final_team]
@ -3885,6 +3942,7 @@ async def test_process_group_patch_remove_filtered_path_without_value(mocker):
prisma_client = mocker.MagicMock()
prisma_client.db = mocker.MagicMock()
prisma_client.replica_db = prisma_client.db
prisma_client.db.litellm_usertable = mocker.MagicMock()
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=LiteLLM_UserTable(user_id="user-1"))
prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
@ -3915,6 +3973,7 @@ async def test_process_group_patch_add_filtered_path_without_value(mocker):
prisma_client = mocker.MagicMock()
prisma_client.db = mocker.MagicMock()
prisma_client.replica_db = prisma_client.db
prisma_client.db.litellm_usertable = mocker.MagicMock()
prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=(LiteLLM_UserTable(user_id="user-3"),))
@ -3948,6 +4007,7 @@ async def test_process_group_patch_replace_empty_value_does_not_use_path_filter(
prisma_client = mocker.MagicMock()
prisma_client.db = mocker.MagicMock()
prisma_client.replica_db = prisma_client.db
prisma_client.db.litellm_usertable = mocker.MagicMock()
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=LiteLLM_UserTable(user_id="user-1"))
prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
@ -3992,6 +4052,7 @@ def _member_resolution_prisma(
prisma_client = mocker.MagicMock()
prisma_client.db = mocker.MagicMock()
prisma_client.replica_db = prisma_client.db
prisma_client.db.litellm_usertable = mocker.MagicMock()
prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=user_row)
@ -5602,6 +5663,7 @@ async def test_get_groups_members_are_typed_as_users(mocker):
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[team])
mock_prisma_client.db.litellm_teamtable.count = AsyncMock(return_value=1)
@ -5642,6 +5704,7 @@ async def test_update_user_roster_add_failure_propagates_and_skips_teams_write(m
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.update = AsyncMock()
@ -5682,6 +5745,7 @@ async def test_patch_user_roster_remove_failure_propagates_and_skips_teams_write
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
mock_prisma_client.db.litellm_usertable.update = AsyncMock()
@ -5839,6 +5903,7 @@ async def test_update_group_roster_failure_propagates(mocker):
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
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.update = AsyncMock(return_value=existing_team)
@ -5956,6 +6021,7 @@ async def test_patch_group_404s_when_team_deleted_mid_request(mocker):
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db = mocker.MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_teamtable = mocker.MagicMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(side_effect=[snapshot_team, None, None])
mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=None)
@ -6033,6 +6099,7 @@ def _shadow_tenant_prisma(
prisma_client = mocker.MagicMock()
prisma_client.db = mocker.MagicMock()
prisma_client.replica_db = 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)

7
tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -1646,6 +1646,7 @@ async def test_pass_through_request_streamed_response_is_owned_by_the_caller():
b'event: response.completed\ndata: {"type": "response.completed", "response": {"id": "%s"}}\n\n'
) % (raw_id.encode(), raw_id.encode())
prisma_client = MagicMock()
prisma_client.replica_db = prisma_client.db
prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None)
prisma_client.db.litellm_managedobjecttable.upsert = AsyncMock(return_value=None)
prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None)
@ -2755,6 +2756,7 @@ async def test_filter_endpoints_by_team_allowed_routes_with_filter():
# Mock prisma client
mock_prisma_client = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_team = MagicMock()
mock_team.metadata = {"allowed_passthrough_routes": ["/api/allowed1", "/api/allowed2"]}
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team)
@ -2795,6 +2797,7 @@ async def test_filter_endpoints_by_team_allowed_routes_team_not_found():
# Mock prisma client to return None (team not found)
mock_prisma_client = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
# Call the function and expect HTTPException
@ -2829,6 +2832,7 @@ async def test_filter_endpoints_by_team_allowed_routes_no_metadata():
# Mock prisma client with team that has None metadata
mock_prisma_client = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_team = MagicMock()
mock_team.metadata = None
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team)
@ -2865,6 +2869,7 @@ async def test_filter_endpoints_by_team_allowed_routes_no_allowed_routes_key():
# Mock prisma client with team that has metadata but no allowed_passthrough_routes
mock_prisma_client = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_team = MagicMock()
mock_team.metadata = {"some_other_key": "some_value"}
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team)
@ -2901,6 +2906,7 @@ async def test_filter_endpoints_by_team_allowed_routes_empty_allowed_list():
# Mock prisma client with team that has empty allowed_passthrough_routes
mock_prisma_client = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_team = MagicMock()
mock_team.metadata = {"allowed_passthrough_routes": []}
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team)
@ -2941,6 +2947,7 @@ async def test_filter_endpoints_by_team_allowed_routes_partial_match():
# Mock prisma client with team that allows only 2 routes
mock_prisma_client = MagicMock()
mock_prisma_client.replica_db = mock_prisma_client.db
mock_team = MagicMock()
mock_team.metadata = {"allowed_passthrough_routes": ["/api/openai", "/api/azure"]}
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team)

1
tests/test_litellm/proxy/proxy_server/conftest.py
View file

@ -161,6 +161,7 @@ def mock_prisma() -> MagicMock:
"""
client_mock = MagicMock()
client_mock.db = MagicMock()
client_mock.replica_db = client_mock.db
client_mock.connect = AsyncMock()
client_mock.disconnect = AsyncMock()
client_mock.health_check = AsyncMock(return_value=True)

12
tests/test_litellm/proxy/proxy_server/test_proxy_config.py
View file

@ -3615,6 +3615,7 @@ async def test_ProxyConfig__add_router_settings_from_db_config_updates_router():
fake_router = MagicMock()
fake_router.update_settings = MagicMock()
fake_prisma = MagicMock()
fake_prisma.replica_db = fake_prisma.db
fake_prisma.db.litellm_config.find_first = AsyncMock(
return_value=SimpleNamespace(param_value={"timeout": 30, "retries": 2, "fallbacks": []})
)
@ -3658,6 +3659,7 @@ async def test_ProxyConfig_add_deployment_applies_db_router_settings(monkeypatch
fake_router = MagicMock()
fake_router.get_model_list = MagicMock(return_value=[])
fake_prisma = MagicMock()
fake_prisma.replica_db = fake_prisma.db
fake_prisma.db.litellm_config.find_first = AsyncMock(
return_value=SimpleNamespace(param_value={"routing_strategy": "latency-based-routing"})
)
@ -3716,6 +3718,7 @@ def _encrypted_credential_row(credential_name: str, api_key: str) -> dict[str, o
def _fake_prisma_with_encrypted_credential(credential_name: str, api_key: str) -> MagicMock:
fake_prisma = MagicMock()
fake_prisma.replica_db = fake_prisma.db
fake_prisma.db.litellm_credentialstable.find_many = AsyncMock(
return_value=[_encrypted_credential_row(credential_name, api_key)]
)
@ -3732,6 +3735,7 @@ async def test_ProxyConfig_add_deployment_loads_db_credentials_before_reconcilin
pc = ProxyConfig()
fake_prisma = MagicMock()
fake_prisma.replica_db = fake_prisma.db
fake_prisma.db.litellm_credentialstable.find_many = AsyncMock(return_value=[])
_stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma)
monkeypatch.setattr(proxy_server, "general_settings", {})
@ -4414,6 +4418,7 @@ async def test_ProxyConfig__init_agents_in_db_keeps_config_defined_agents(clean_
)
prisma_client = MagicMock()
prisma_client.replica_db = prisma_client.db
prisma_client.db.litellm_agentstable.find_many = AsyncMock(return_value=[_FakeAgentRow("db-id", "db-agent")])
await ProxyConfig()._init_agents_in_db(prisma_client=prisma_client)
@ -4450,6 +4455,7 @@ async def test_ProxyStartupEvent_jwt_auth_resolves_agent_claims_against_live_reg
)
elif agents_source == "db":
prisma_client = MagicMock()
prisma_client.replica_db = prisma_client.db
prisma_client.db.litellm_agentstable.find_many = AsyncMock(
return_value=[_FakeAgentRow("db-id", "loaded-agent")]
)
@ -4571,6 +4577,7 @@ async def test_ProxyConfig__init_guardrails_in_db_skips_only_the_unloadable_row(
monkeypatch.setitem(registry_module.guardrail_initializer_registry, "lit5367_ok", _initializer)
prisma_client = MagicMock()
prisma_client.replica_db = prisma_client.db
prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(
return_value=[
_db_guardrail_row("first", "lit5367_ok"),
@ -4601,6 +4608,7 @@ async def test_add_deployment_re_reads_ui_settings_so_other_pods_converge(monkey
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
prisma_client = MagicMock()
prisma_client.replica_db = prisma_client.db
prisma_client.db.litellm_config.find_many = AsyncMock(return_value=[])
prisma_client.db.litellm_config.find_first = AsyncMock(return_value=None)
prisma_client.db.litellm_credentialstable.find_many = AsyncMock(return_value=[])
@ -4628,6 +4636,7 @@ async def test_add_deployment_syncs_ui_settings_even_when_the_model_reconcile_fa
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
prisma_client = MagicMock()
prisma_client.replica_db = prisma_client.db
prisma_client.db.litellm_uisettings.find_unique = AsyncMock(
return_value=SimpleNamespace(ui_settings={"allow_agents_for_team_admins": True})
)
@ -4893,7 +4902,8 @@ async def test_model_refresh_updates_availability_catalog_and_retains_it_on_db_f
},
)
find_many = AsyncMock(side_effect=[[row], RuntimeError("database unavailable"), []])
client = SimpleNamespace(db=SimpleNamespace(litellm_proxymodeltable=SimpleNamespace(find_many=find_many)))
tables = SimpleNamespace(litellm_proxymodeltable=SimpleNamespace(find_many=find_many))
client = SimpleNamespace(db=tables, replica_db=tables)
assert pc.auto_router_db_catalog is None
assert await pc._get_models_from_db(client) == [row]
loaded = pc.auto_router_db_catalog

10
tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py
View file

@ -340,6 +340,7 @@ async def test_model_info_v1_unrestricted_key_hides_other_team_byok(monkeypatch)
router.get_model_access_groups.return_value = {}
prisma_client = MagicMock()
prisma_client.replica_db = prisma_client.db
caller_user_row = MagicMock()
caller_user_row.teams = ["team-abc-123"]
caller_user_row.model_dump.return_value = {
@ -396,6 +397,7 @@ async def test_model_info_v1_service_key_hides_all_team_byok(monkeypatch):
router.get_model_access_groups.return_value = {}
prisma_client = MagicMock()
prisma_client.replica_db = prisma_client.db
monkeypatch.setattr(ps, "user_model", None)
monkeypatch.setattr(ps, "llm_model_list", router.model_list)
@ -445,6 +447,7 @@ async def test_model_info_v1_team_key_sees_own_byok_regardless_of_user_lookup(
router.get_model_access_groups.return_value = {}
prisma_client = MagicMock()
prisma_client.replica_db = prisma_client.db
prisma_client.db.litellm_usertable.find_unique = find_unique
async def _populate(**kwargs):
@ -486,6 +489,7 @@ async def test_model_info_v1_user_team_membership_grants_byok(monkeypatch):
router.get_model_access_groups.return_value = {}
prisma_client = MagicMock()
prisma_client.replica_db = prisma_client.db
prisma_client.db.litellm_usertable.find_unique = AsyncMock(
return_value=MagicMock(teams=["team-other"])
)
@ -534,6 +538,7 @@ async def test_model_info_v1_populates_access_via_team_ids(monkeypatch):
router.get_model_ids.return_value = ["global-id-1"]
prisma_client = MagicMock()
prisma_client.replica_db = prisma_client.db
async def _fake_populate(**kwargs):
for model in kwargs["all_models"]:
@ -614,6 +619,7 @@ async def test_populate_team_access_gives_view_only_admin_full_admin_scope(monke
monkeypatch.setattr(ps, "get_all_team_models", get_all_team_models)
prisma_client = MagicMock()
prisma_client.replica_db = prisma_client.db
prisma_client.db.litellm_usertable.find_unique = AsyncMock(
return_value=LiteLLM_UserTable(user_id="viewer", teams=[], models=[])
)
@ -676,6 +682,7 @@ async def test_populate_team_access_grants_config_access_group_model():
"access_group_ids": [],
}
prisma_client = MagicMock()
prisma_client.replica_db = prisma_client.db
prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[team_db_object])
admin = UserAPIKeyAuth(user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[])
@ -1727,6 +1734,7 @@ async def test_populate_team_access_grants_empty_models_user_direct_access(monke
teams=[],
)
prisma_client = MagicMock()
prisma_client.replica_db = prisma_client.db
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row)
monkeypatch.setattr(ps, "get_all_team_models", AsyncMock(return_value={}))
@ -1829,6 +1837,7 @@ async def test_populate_team_access_hides_models_the_calling_key_cannot_call(mon
teams=[],
)
prisma_client = MagicMock()
prisma_client.replica_db = prisma_client.db
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row)
monkeypatch.setattr(ps, "get_all_team_models", AsyncMock(return_value={}))
@ -1875,6 +1884,7 @@ async def test_populate_team_access_grants_all_proxy_models_user_direct_access(
teams=[],
)
prisma_client = MagicMock()
prisma_client.replica_db = prisma_client.db
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row)
monkeypatch.setattr(ps, "get_all_team_models", AsyncMock(return_value={}))