test: await the now-async _apply_message and justify new internal patches for the TQ008 gate

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-08-25 22:03:18 +00:00
parent d45c4f9d25
commit b219354539
4 changed files with 21 additions and 21 deletions

View file

@ -1592,7 +1592,7 @@ async def test_cleared_m2m_row_can_no_longer_mint_a_replacement_token():
)
token_endpoint = AsyncMock()
with patch(
with patch( # test-quality-ok: fakes the token-endpoint HTTP client to prove no mint request leaves the process
"litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client",
return_value=token_endpoint,
):

View file

@ -5250,12 +5250,12 @@ class TestMCPServerManager:
invalidated: list[str] = []
with (
patch.object(
patch.object( # test-quality-ok: records which server ids the broadcast handler routes to the process-wide manager singleton
global_mcp_server_manager,
"invalidate_local_upstream_m2m_tokens",
new=AsyncMock(side_effect=lambda server_id: invalidated.append(server_id)),
),
patch.object(global_mcp_server_manager, "refresh_local_server_from_db", new=AsyncMock()),
patch.object(global_mcp_server_manager, "refresh_local_server_from_db", new=AsyncMock()), # test-quality-ok: keeps the process-wide manager singleton from touching the DB in this routing test
):
await apply_mcp_upstream_m2m_invalidation("mcp:per_user_token:alice:srv-1")
assert invalidated == []
@ -5289,14 +5289,14 @@ class TestMCPServerManager:
applied: list[LiteLLM_MCPServerTable] = []
with (
patch.object(global_mcp_server_manager, "invalidate_local_upstream_m2m_tokens", new=AsyncMock()),
patch.object(
patch.object(global_mcp_server_manager, "invalidate_local_upstream_m2m_tokens", new=AsyncMock()), # test-quality-ok: quiets the process-wide manager singleton's cache drop; the DB re-read is what this test pins
patch.object( # test-quality-ok: records the registry refresh applied to the process-wide manager singleton
global_mcp_server_manager,
"update_server",
new=AsyncMock(side_effect=lambda server: applied.append(server)),
),
patch.object(mcp_db, "get_mcp_server", new=AsyncMock(return_value=revoked_row)),
patch("litellm.proxy.proxy_server.prisma_client", new=MagicMock()),
patch.object(mcp_db, "get_mcp_server", new=AsyncMock(return_value=revoked_row)), # test-quality-ok: stands in for the DB row read; no live prisma in this unit test
patch("litellm.proxy.proxy_server.prisma_client", new=MagicMock()), # test-quality-ok: satisfies the prisma-configured guard; no live DB in this unit test
):
await apply_mcp_upstream_m2m_invalidation(mcp_oauth2_mint_invalidation_key("srv-1"))

View file

@ -7239,7 +7239,7 @@ async def test_invalidate_team_member_spend_state_broadcasts_the_spend_counter_t
additional_in_memory_caches=(remote_spend_counter_in_memory_cache,),
)
for cache_key in ("spend:team_member:user-1:team-1", "spend_db_floor:spend:team_member:user-1:team-1"):
remote_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API
await remote_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API
{"type": "message", "data": _published_message_for(cache_key)}
)
@ -7304,7 +7304,7 @@ async def test_invalidate_team_member_spend_state_self_delivered_broadcast_does_
additional_in_memory_caches=(local_spend_counter_cache.in_memory_cache,),
)
for _, message in published:
own_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API
await own_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API
{"type": "message", "data": message}
)

View file

@ -6783,19 +6783,19 @@ def _patched_clear_oauth_token_endpoint(clear_mock, purge_mock, manager, revoke_
actually emptied rather than only what it returned."""
manager.revoke_upstream_m2m_tokens = revoke_mock
with (
patch(
patch( # test-quality-ok: stands in for the prisma-backed store; no live DB in this unit test
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=_make_prisma_client(),
),
patch(
patch( # test-quality-ok: records the DB clear so the test can assert which stores the endpoint emptied
"litellm.proxy.management_endpoints.mcp_management_endpoints.clear_mcp_server_oauth_token",
new=clear_mock,
),
patch(
patch( # test-quality-ok: records the per-user purge so the test can assert which stores the endpoint emptied
"litellm.proxy.management_endpoints.mcp_management_endpoints.purge_user_oauth_credentials_for_server",
new=purge_mock,
),
patch(
patch( # test-quality-ok: injects a recording manager for the module-global singleton
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
manager,
),
@ -6916,19 +6916,19 @@ async def test_clear_mcp_server_oauth_token_endpoint_drives_a_real_purge_for_the
manager.revoke_upstream_m2m_tokens = AsyncMock()
with (
patch(
patch( # test-quality-ok: stands in for the prisma-backed store; no live DB in this unit test
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=prisma,
),
patch(
patch( # test-quality-ok: stubs the DB clear; the per-user purge path is what this test pins
"litellm.proxy.management_endpoints.mcp_management_endpoints.clear_mcp_server_oauth_token",
new=AsyncMock(return_value=ClearedMCPServerOAuthToken(server=server, had_token=False, cleared_client_credentials=False)),
),
patch(
patch( # test-quality-ok: injects a recording manager for the module-global singleton
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
manager,
),
patch(
patch( # test-quality-ok: injects the same recording manager where the purge helper resolves the singleton
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager",
manager,
),
@ -7066,11 +7066,11 @@ async def test_clear_mcp_server_oauth_token_endpoint_rejects_non_admin():
clear_mock = AsyncMock()
with (
patch(
patch( # test-quality-ok: stands in for the prisma-backed store; no live DB in this unit test
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=_make_prisma_client(),
),
patch(
patch( # test-quality-ok: records the DB clear to prove the admin check rejects before any clearing
"litellm.proxy.management_endpoints.mcp_management_endpoints.clear_mcp_server_oauth_token",
new=clear_mock,
),
@ -7095,11 +7095,11 @@ async def test_clear_mcp_server_oauth_token_endpoint_404s_for_unknown_server():
)
with (
patch(
patch( # test-quality-ok: stands in for the prisma-backed store; no live DB in this unit test
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=_make_prisma_client(),
),
patch(
patch( # test-quality-ok: stubs the DB clear returning no row, the unknown-server shape this test pins
"litellm.proxy.management_endpoints.mcp_management_endpoints.clear_mcp_server_oauth_token",
new=AsyncMock(return_value=None),
),