From c41946a9d21c5d7318dd6bd10f80efb63587ee3c Mon Sep 17 00:00:00 2001
From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Date: Tue, 29 Sep 2026 15:41:53 -0700
Subject: [PATCH] test(mcp): stub the writer replica the fresh-policy reads use

reload_admitted_user now passes check_db_only through to get_user_object,
so the user row is read from writer_db. Point the mocks at the replica the
code actually reads and give each parametrized case its own user id.
---
 .../_experimental/mcp_server/test_discoverable_endpoints.py   | 2 +-
 .../_experimental/mcp_server/test_proxy_api_credentials.py    | 4 ++--
 .../management_endpoints/test_mcp_management_endpoints.py     | 4 ++--
 3 files changed, 5 insertions(+), 5 deletions(-)

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 b1f0b3fa67e..4e27ec134d4 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
@@ -7847,7 +7847,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.db.litellm_usertable.find_unique = AsyncMock(
+    prisma.writer_db.litellm_usertable.find_unique = AsyncMock(
         return_value=LiteLLM_UserTable(user_id="fresh-jwt-user", teams=["team-a"])
     )
     proxy_globals.user_api_key_cache = cache
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py
index ed3e5f48516..8484dfdde72 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_proxy_api_credentials.py
@@ -139,7 +139,7 @@ async def test_mint_reads_the_users_teams_from_the_database_not_a_stale_cached_r
         key="stale-cache-user", value=_user(user_id="stale-cache-user", teams=[]), model_type=LiteLLM_UserTable
     )
     prisma = MagicMock()
-    prisma.db.litellm_usertable.find_unique = AsyncMock(
+    prisma.writer_db.litellm_usertable.find_unique = AsyncMock(
         return_value=_user(user_id="stale-cache-user", teams=["team-a"])
     )
     monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
@@ -164,7 +164,7 @@ async def test_mint_refuses_a_user_scim_deactivated_after_the_cache_last_saw_the
         key="deactivated-user", value=_user(user_id="deactivated-user", teams=["team-a"]), model_type=LiteLLM_UserTable
     )
     prisma = MagicMock()
-    prisma.db.litellm_usertable.find_unique = AsyncMock(
+    prisma.writer_db.litellm_usertable.find_unique = AsyncMock(
         return_value=_user(user_id="deactivated-user", teams=["team-a"], metadata={"scim_active": False})
     )
     monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
index 160cf8be4e0..f5fc5ae24d4 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
@@ -7679,7 +7679,7 @@ class TestConnectedAppViewAnnotation:
 
         flags = {server.server_id: server.connected_app_reachable for server in result}
         assert flags == {"server-1": True, "server-2": False}
-        reload_mock.assert_awaited_once_with("test_user_id")
+        reload_mock.assert_awaited_once_with("test_user_id", requires_fresh_policy=False)
         mock_manager.get_allowed_mcp_servers.assert_awaited_once_with(admitted_auth)
 
     @pytest.mark.asyncio
@@ -9653,7 +9653,7 @@ class TestMCPServerResolutionCharacterization:
         server_id: str,
     ) -> tuple[MagicMock, MCPServerManager, UserAPIKeyAuth]:
         team_id: Final = UI_SESSION_TOKEN_TEAM_ID if grant_route == "direct user object_permission" else "lit3974_team"
-        user_id: Final = "lit3974_direct_user"
+        user_id: Final = f"{server_id}:{grant_route}:user"
         key_permission: Final = LiteLLM_ObjectPermissionTable(
             object_permission_id=f"lit3974_{grant_route}_key_permission",
             mcp_servers=None,