mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
test(proxy): stub the writer reads and restore the MCP manager singleton
Fresh-policy user lookups read writer_db, so the team and rest-endpoint mocks stubbed a replica the code no longer reads, and the dashboard session fake still had the pre-kwarg signature. The manager reload also rebound global_mcp_server_manager in every MCP module without restoring it, leaking an empty manager into later files.
This commit is contained in:
parent
da39768d63
commit
b096456de7
3 changed files with 21 additions and 7 deletions
|
|
@ -224,6 +224,20 @@ def enable_eager_mcp_oauth_discovery(monkeypatch):
|
|||
monkeypatch.setenv("LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP", "1")
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def restore_mcp_manager_singleton():
|
||||
"""``_reload_mcp_manager_module`` rebinds ``global_mcp_server_manager`` in every MCP module, so
|
||||
without this the next test file inherits a manager that has none of its servers registered."""
|
||||
bound: Final = tuple(
|
||||
(module, module.global_mcp_server_manager)
|
||||
for name, module in tuple(sys.modules.items())
|
||||
if name.startswith("litellm.proxy._experimental.mcp_server.") and hasattr(module, "global_mcp_server_manager")
|
||||
)
|
||||
yield
|
||||
for module, manager in bound:
|
||||
module.global_mcp_server_manager = manager
|
||||
|
||||
|
||||
class TestMCPServerManager:
|
||||
"""Test MCP Server Manager stdio functionality"""
|
||||
|
||||
|
|
|
|||
|
|
@ -1311,7 +1311,7 @@ class TestListToolsRestAPI:
|
|||
session_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user")
|
||||
admitted_auth = UserAPIKeyAuth(user_id="grant-user", org_id="admitted-org")
|
||||
|
||||
async def fake_reload(user_id):
|
||||
async def fake_reload(user_id, *, requires_fresh_policy=False):
|
||||
assert user_id == "grant-user"
|
||||
return admitted_auth
|
||||
|
||||
|
|
|
|||
|
|
@ -4355,7 +4355,7 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs(
|
|||
prisma_client.db.litellm_teamtable.find_many = AsyncMock(side_effect=find_many)
|
||||
prisma_client.db.litellm_teamtable.count = AsyncMock(side_effect=count)
|
||||
prisma_client.db.litellm_verificationtoken.group_by = AsyncMock(return_value=[])
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
prisma_client.writer_db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=LiteLLM_UserTable(
|
||||
user_id="org_admin_user",
|
||||
teams=["team_in_org_A", "team_in_org_B"],
|
||||
|
|
@ -4394,11 +4394,11 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs(
|
|||
assert await list_teams(None) == own_view
|
||||
assert await list_teams("org_admin_user", search="team_in_org_B") == ["team_in_org_B"]
|
||||
assert await list_teams("other_user") == ["other_team_in_org_A"]
|
||||
prisma_client.db.litellm_usertable.find_unique.assert_awaited_with(
|
||||
prisma_client.writer_db.litellm_usertable.find_unique.assert_awaited_with(
|
||||
where={"user_id": "org_admin_user"}, include={"organization_memberships": True}
|
||||
)
|
||||
|
||||
prisma_client.db.litellm_usertable.find_unique.side_effect = RuntimeError("db down")
|
||||
prisma_client.writer_db.litellm_usertable.find_unique.side_effect = RuntimeError("db down")
|
||||
with pytest.raises(ValueError, match="db down"):
|
||||
await list_teams("org_admin_user")
|
||||
|
||||
|
|
@ -15813,7 +15813,7 @@ async def test_get_team_spend_by_user_team_admin_sees_every_member(mock_db_clien
|
|||
alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), [])
|
||||
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha])
|
||||
mock_db_client.db.query_raw = AsyncMock(return_value=[])
|
||||
mock_db_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=_team_spend_by_user_caller("alice", ["team-alpha"])
|
||||
)
|
||||
|
||||
|
|
@ -15835,7 +15835,7 @@ async def test_get_team_spend_by_user_plain_member_only_sees_own_row(mock_db_cli
|
|||
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha])
|
||||
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock_db_client.db.query_raw = AsyncMock(return_value=[_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2)])
|
||||
mock_db_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=_team_spend_by_user_caller("bob", ["team-alpha"])
|
||||
)
|
||||
|
||||
|
|
@ -15856,7 +15856,7 @@ async def test_get_team_spend_by_user_member_of_other_team_gets_404(mock_db_clie
|
|||
|
||||
caller = UserAPIKeyAuth(user_id="bob", user_role=LitellmUserRoles.INTERNAL_USER)
|
||||
mock_db_client.db.query_raw = AsyncMock(return_value=[])
|
||||
mock_db_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=_team_spend_by_user_caller("bob", ["team-alpha"])
|
||||
)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue