Files
AIRegulation-DocAnalysis/backend/tests/mcp/test_search_regulations_tool.py
T
wangweiandCopilot 49ee50c104 fix: harden MCP endpoint after code review
Critical: the MCP SDK auto-enables DNS-rebinding protection when its host
parameter is left at the 127.0.0.1 default, hard-coding a loopback-only Host
allow-list. Every remote client (the only deployment this feature targets) was
refused with HTTP 421 before auth or the tool ran. Now driven by a new
MCP_ALLOWED_HOSTS setting, with '*' as an explicit, logged opt-out.

Also bounds query/top_k to match AskRequest (top_k is amplified 4x downstream,
so an unbounded value was a resource-exhaustion vector), decodes the
Authorization header as latin-1 per the ASGI spec instead of raising a 500 on
malformed bytes, and returns WWW-Authenticate on 401 per RFC 7235.

Moves the psycopg2 import guard into backend/tests/conftest.py: duplicated
across four test modules, it only worked because of alphabetical collection
order, and any earlier-sorting package would have reintroduced a live
connection attempt against the production database.

Registers the mcp module in the authoritative backend architecture doc.

84 backend tests pass. Verified against a live server: allowed remote Host
returns a valid initialize result, unknown Host returns 421, missing token
returns 401 with WWW-Authenticate.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-07-29 17:11:54 +08:00

99 lines
3.3 KiB
Python

"""Unit tests for the search_regulations MCP tool function.
Mocks AgentConversationService so no real retrieval/LLM call happens —
verifies only the protocol-adapter contract: correct call shape in,
correct dict shape out.
"""
from __future__ import annotations
import asyncio
from dataclasses import dataclass
from unittest.mock import MagicMock, patch
@dataclass
class _FakeSource:
"""Minimal stand-in for a real Source dataclass (only __dict__ is used)."""
# A dataclass, not a MagicMock: the adapter serializes sources via
# source.__dict__, and a MagicMock's __dict__ is full of internal mock
# attributes, which would make the assertions meaningless.
doc_id: str
doc_title: str
score: float
@dataclass
class _FakeAnswerResult:
"""Minimal stand-in for AnswerResult — only .answer/.sources are read."""
answer: str
sources: list
def test_search_regulations_calls_agent_ask_without_session():
"""search_regulations must call ask() with no session_id (stateless search)."""
from app.mcp.server import search_regulations
fake_service = MagicMock()
fake_service.ask.return_value = (
None,
_FakeAnswerResult(answer="国六排放标准要求...", sources=[_FakeSource("doc-1", "国六标准", 0.9)]),
)
with patch("app.mcp.server.get_agent_conversation_service", return_value=fake_service):
search_regulations(query="国六排放标准最新要求", top_k=3)
fake_service.ask.assert_called_once_with(query="国六排放标准最新要求", top_k=3)
assert "session_id" not in fake_service.ask.call_args.kwargs
def test_search_regulations_shapes_response_dict():
"""The returned dict must expose 'answer' and 'sources' (list of plain dicts)."""
from app.mcp.server import search_regulations
fake_service = MagicMock()
fake_service.ask.return_value = (
None,
_FakeAnswerResult(answer="答案文本", sources=[_FakeSource("doc-2", "国标GB1589", 0.8)]),
)
with patch("app.mcp.server.get_agent_conversation_service", return_value=fake_service):
result = search_regulations(query="q")
assert result == {
"answer": "答案文本",
"sources": [{"doc_id": "doc-2", "doc_title": "国标GB1589", "score": 0.8}],
}
def test_search_regulations_default_top_k():
"""top_k defaults to 5 when the caller omits it."""
from app.mcp.server import search_regulations
fake_service = MagicMock()
fake_service.ask.return_value = (None, _FakeAnswerResult(answer="a", sources=[]))
with patch("app.mcp.server.get_agent_conversation_service", return_value=fake_service):
search_regulations(query="q")
assert fake_service.ask.call_args.kwargs["top_k"] == 5
def test_advertised_schema_bounds_top_k_and_query():
"""The advertised JSON schema must carry the same bounds as AskRequest.
Bounds declared via Annotated are what the SDK validates against and what
clients see, so asserting on the generated schema is the only way to catch
a regression that silently drops them.
"""
from app.mcp.server import mcp
schema = asyncio.run(mcp.list_tools())[0].input_schema["properties"]
assert schema["top_k"]["minimum"] == 1
assert schema["top_k"]["maximum"] == 20
assert schema["query"]["minLength"] == 1
assert schema["query"]["maxLength"] == 2000