diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index e4f274497be..6b92e2908ce 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -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 diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 38988d65c04..4a3272bb17b 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -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", diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 74bd67efaf2..a7baa742590 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -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"} diff --git a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py index 0fc91748af2..e65ea4a65b6 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py @@ -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 diff --git a/tests/test_litellm/integrations/test_responses_background_cost.py b/tests/test_litellm/integrations/test_responses_background_cost.py index 0d4218f2137..de8aa4ef624 100644 --- a/tests/test_litellm/integrations/test_responses_background_cost.py +++ b/tests/test_litellm/integrations/test_responses_background_cost.py @@ -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 diff --git a/tests/test_litellm/llms/base_llm/files/test_litellm_db_storage_backend.py b/tests/test_litellm/llms/base_llm/files/test_litellm_db_storage_backend.py index fabcb340a48..44adea3bb9c 100644 --- a/tests/test_litellm/llms/base_llm/files/test_litellm_db_storage_backend.py +++ b/tests/test_litellm/llms/base_llm/files/test_litellm_db_storage_backend.py @@ -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 diff --git a/tests/test_litellm/llms/base_llm/files/test_storage_backend_factory.py b/tests/test_litellm/llms/base_llm/files/test_storage_backend_factory.py index 945691c5b98..ad578ad29e0 100644 --- a/tests/test_litellm/llms/base_llm/files/test_storage_backend_factory.py +++ b/tests/test_litellm/llms/base_llm/files/test_storage_backend_factory.py @@ -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) diff --git a/tests/test_litellm/llms/base_llm/test_base_managed_resource.py b/tests/test_litellm/llms/base_llm/test_base_managed_resource.py index a6c518616f4..31f32e766c0 100644 --- a/tests/test_litellm/llms/base_llm/test_base_managed_resource.py +++ b/tests/test_litellm/llms/base_llm/test_base_managed_resource.py @@ -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 = ( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index f9a0075e530..e77b1e9595f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -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) diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index f8b9043a23f..ae5a048af28 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -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) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index bf3b9aed234..2a7fcced670 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -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() diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index dbcf622bbb1..1b25e6657cd 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -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) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 7a64d5f2218..b349ed7ceb5 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -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) diff --git a/tests/test_litellm/proxy/proxy_server/conftest.py b/tests/test_litellm/proxy/proxy_server/conftest.py index ae1b42363ef..8bebcac5641 100644 --- a/tests/test_litellm/proxy/proxy_server/conftest.py +++ b/tests/test_litellm/proxy/proxy_server/conftest.py @@ -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) diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 7e198bc9131..0c94f4c7aaa 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -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 diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index baa032f75e6..70cee175d4a 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -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={}))