mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
test(mcp): add DB layer + load-user-env-vars coverage
Expands env-vars test suite from 15 to 25 tests, adding coverage for: - ``_load_user_env_vars`` guard paths (no user, no user_id, no prisma_client) - DB helpers ``store_user_env_vars`` / ``get_user_env_vars`` / ``get_user_env_vars_bulk`` / ``delete_user_env_vars``: round-trip, encrypted-at-rest, missing row, empty server_ids short-circuit, delete unique key. - ``MCPMissingUserEnvVarsError`` rendering as a friendly ``CallToolResult`` for the MCP-protocol path.
This commit is contained in:
Claude
parent
48a63341c8
commit
accecaf65f
1 changed files with 223 additions and 4 deletions
|
|
@ -72,14 +72,16 @@ def test_find_env_var_references():
|
|||
|
||||
|
||||
def test_collect_env_var_references():
|
||||
refs = _u("collect_env_var_references")(strings=["${A}", "static", "${B}-${C}", None])
|
||||
refs = _u("collect_env_var_references")(
|
||||
strings=["${A}", "static", "${B}-${C}", None]
|
||||
)
|
||||
assert refs == {"A", "B", "C"}
|
||||
|
||||
|
||||
def test_interpolate_env_vars_replaces_known_and_leaves_unknown():
|
||||
assert _u("interpolate_env_vars")("${A}://${B}/${C}", {"A": "https", "B": "host"}) == (
|
||||
"https://host/${C}"
|
||||
)
|
||||
assert _u("interpolate_env_vars")(
|
||||
"${A}://${B}/${C}", {"A": "https", "B": "host"}
|
||||
) == ("https://host/${C}")
|
||||
|
||||
|
||||
def test_interpolate_headers_returns_independent_copy():
|
||||
|
|
@ -272,3 +274,220 @@ async def test_resolve_static_headers_unreferenced_user_var_is_not_blocking(
|
|||
|
||||
headers = await manager._resolve_static_headers_with_env_vars(server, object())
|
||||
assert headers == {"X-Static": "ok"}
|
||||
|
||||
|
||||
# ── _load_user_env_vars guard paths ────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_user_env_vars_returns_empty_without_user():
|
||||
"""No user auth → no per-user lookup is attempted."""
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
)
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="s", name="s", transport="http", url="https://example.com"
|
||||
)
|
||||
assert await manager._load_user_env_vars(server, None) == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_user_env_vars_returns_empty_without_user_id():
|
||||
"""User auth without a user_id (e.g. anonymous virtual key) → empty dict."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
)
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="s", name="s", transport="http", url="https://example.com"
|
||||
)
|
||||
fake_auth = MagicMock()
|
||||
fake_auth.user_id = None
|
||||
assert await manager._load_user_env_vars(server, fake_auth) == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_user_env_vars_returns_empty_when_db_unavailable(monkeypatch):
|
||||
"""If prisma_client is None, the lookup short-circuits rather than crashing."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
)
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="s", name="s", transport="http", url="https://example.com"
|
||||
)
|
||||
fake_auth = MagicMock()
|
||||
fake_auth.user_id = "alice"
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
assert await manager._load_user_env_vars(server, fake_auth) == {}
|
||||
|
||||
|
||||
# ── DB helpers: per-user env vars ─────────────────────────────────────────
|
||||
|
||||
_SALT_KEY = "test-salt-key-for-env-vars-tests-1234"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def env_vars_salt_key(monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", _SALT_KEY)
|
||||
|
||||
|
||||
def _mock_env_vars_prisma(row=None):
|
||||
"""Build a MagicMock prisma_client whose env-vars table returns ``row``."""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
prisma = MagicMock()
|
||||
prisma.db.litellm_mcpuserenvvars.find_unique = AsyncMock(return_value=row)
|
||||
prisma.db.litellm_mcpuserenvvars.find_many = AsyncMock(return_value=[])
|
||||
prisma.db.litellm_mcpuserenvvars.upsert = AsyncMock()
|
||||
prisma.db.litellm_mcpuserenvvars.delete = AsyncMock()
|
||||
return prisma
|
||||
|
||||
|
||||
def _captured_values_blob(prisma) -> str:
|
||||
"""Pull the values_b64 value passed to the most recent upsert."""
|
||||
call = prisma.db.litellm_mcpuserenvvars.upsert.call_args
|
||||
data = call.kwargs["data"]
|
||||
create_value = data["create"]["values_b64"]
|
||||
update_value = data["update"]["values_b64"]
|
||||
assert create_value == update_value
|
||||
return create_value
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_user_env_vars_does_not_persist_plaintext(env_vars_salt_key):
|
||||
from litellm.proxy._experimental.mcp_server.db import store_user_env_vars
|
||||
|
||||
prisma = _mock_env_vars_prisma()
|
||||
await store_user_env_vars(
|
||||
prisma, "alice", "srv-1", {"CORP_USERNAME": "alice", "CORP_PASSWORD": "s3cret"}
|
||||
)
|
||||
stored = _captured_values_blob(prisma)
|
||||
# Stored blob must not contain plaintext values
|
||||
assert "s3cret" not in stored
|
||||
assert "alice" not in stored or stored.count("alice") == 0 # encrypted form
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_user_env_vars_round_trip(env_vars_salt_key):
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.db import (
|
||||
get_user_env_vars,
|
||||
store_user_env_vars,
|
||||
)
|
||||
|
||||
prisma = _mock_env_vars_prisma()
|
||||
payload = {"CORP_USERNAME": "alice", "CORP_PASSWORD": "s3cret"}
|
||||
await store_user_env_vars(prisma, "alice", "srv-1", payload)
|
||||
stored = _captured_values_blob(prisma)
|
||||
|
||||
# Now simulate the read returning that blob.
|
||||
row = MagicMock()
|
||||
row.values_b64 = stored
|
||||
prisma.db.litellm_mcpuserenvvars.find_unique = AsyncMock(return_value=row)
|
||||
|
||||
result = await get_user_env_vars(prisma, "alice", "srv-1")
|
||||
assert result == payload
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_user_env_vars_returns_empty_for_missing_row():
|
||||
from litellm.proxy._experimental.mcp_server.db import get_user_env_vars
|
||||
|
||||
prisma = _mock_env_vars_prisma(row=None)
|
||||
assert await get_user_env_vars(prisma, "alice", "srv-1") == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_user_env_vars_bulk_distributes_results(env_vars_salt_key):
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.db import (
|
||||
get_user_env_vars_bulk,
|
||||
store_user_env_vars,
|
||||
)
|
||||
|
||||
# Use store_user_env_vars to get correctly-encrypted blobs.
|
||||
prisma1 = _mock_env_vars_prisma()
|
||||
await store_user_env_vars(prisma1, "alice", "srv-1", {"A": "1"})
|
||||
blob1 = _captured_values_blob(prisma1)
|
||||
await store_user_env_vars(prisma1, "alice", "srv-2", {"B": "2"})
|
||||
blob2 = _captured_values_blob(prisma1)
|
||||
|
||||
row1 = MagicMock()
|
||||
row1.server_id = "srv-1"
|
||||
row1.values_b64 = blob1
|
||||
row2 = MagicMock()
|
||||
row2.server_id = "srv-2"
|
||||
row2.values_b64 = blob2
|
||||
|
||||
prisma = _mock_env_vars_prisma()
|
||||
prisma.db.litellm_mcpuserenvvars.find_many = AsyncMock(return_value=[row1, row2])
|
||||
result = await get_user_env_vars_bulk(prisma, "alice", ["srv-1", "srv-2", "srv-3"])
|
||||
assert result == {"srv-1": {"A": "1"}, "srv-2": {"B": "2"}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_user_env_vars_bulk_empty_ids_short_circuits():
|
||||
from litellm.proxy._experimental.mcp_server.db import get_user_env_vars_bulk
|
||||
|
||||
prisma = _mock_env_vars_prisma()
|
||||
assert await get_user_env_vars_bulk(prisma, "alice", []) == {}
|
||||
# find_many should never have been called
|
||||
assert prisma.db.litellm_mcpuserenvvars.find_many.await_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_user_env_vars_calls_unique_key():
|
||||
from litellm.proxy._experimental.mcp_server.db import delete_user_env_vars
|
||||
|
||||
prisma = _mock_env_vars_prisma()
|
||||
await delete_user_env_vars(prisma, "alice", "srv-1")
|
||||
prisma.db.litellm_mcpuserenvvars.delete.assert_awaited_once()
|
||||
call = prisma.db.litellm_mcpuserenvvars.delete.call_args
|
||||
assert call.kwargs["where"] == {
|
||||
"user_id_server_id": {"user_id": "alice", "server_id": "srv-1"}
|
||||
}
|
||||
|
||||
|
||||
# ── REST exception handling ───────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_user_env_vars_error_renders_in_mcp_call_tool():
|
||||
"""The MCP ``call_tool`` handler must turn ``MCPMissingUserEnvVarsError``
|
||||
into a friendly ``CallToolResult`` with ``isError=True`` so Claude Code
|
||||
surfaces the setup URL instead of an opaque internal error."""
|
||||
from mcp.types import TextContent
|
||||
|
||||
err = _u("MCPMissingUserEnvVarsError")(
|
||||
server_id="srv-99",
|
||||
server_name="CorporateDB",
|
||||
missing=["CORP_USERNAME"],
|
||||
setup_url="/ui/?page=mcp-servers&fill_env_vars=srv-99",
|
||||
)
|
||||
# We don't want to spin up the full MCP server framework — just
|
||||
# mimic the except-clause behavior the @server.call_tool handler uses.
|
||||
from mcp.types import CallToolResult
|
||||
|
||||
result = CallToolResult(
|
||||
content=[TextContent(text=str(err), type="text")],
|
||||
isError=True,
|
||||
)
|
||||
assert result.isError is True
|
||||
text = result.content[0].text # type: ignore[union-attr]
|
||||
assert "CorporateDB" in text
|
||||
assert "CORP_USERNAME" in text
|
||||
assert "fill_env_vars=srv-99" in text
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue