311 lines
12 KiB
Python
311 lines
12 KiB
Python
from __future__ import annotations
|
||
|
||
from types import SimpleNamespace
|
||
|
||
import pytest
|
||
import pytest_asyncio
|
||
from fastapi import FastAPI
|
||
from httpx import ASGITransport, AsyncClient
|
||
from sqlalchemy import JSON, select
|
||
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
||
|
||
from app.api.mcp_api import routes as mcp_routes
|
||
from app.core.database import get_db
|
||
from app.dependencies.auth import get_current_user
|
||
from app.models.base import Base
|
||
from app.models.integration import Integration
|
||
from app.services.mcp_config import McpConfigService
|
||
|
||
|
||
def _use_sqlite_json_types(*tables) -> None:
|
||
for table in tables:
|
||
for column in table.columns:
|
||
if column.type.__class__.__name__ in {"JSONB", "ARRAY", "Vector"}:
|
||
column.type = JSON()
|
||
|
||
|
||
@pytest_asyncio.fixture
|
||
async def db_session():
|
||
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
||
tables = [Integration.__table__]
|
||
_use_sqlite_json_types(*tables)
|
||
async with engine.begin() as conn:
|
||
await conn.run_sync(lambda sync_conn: Base.metadata.create_all(sync_conn, tables=tables))
|
||
maker = async_sessionmaker(engine, expire_on_commit=False)
|
||
async with maker() as session:
|
||
yield session
|
||
await engine.dispose()
|
||
|
||
|
||
def _build_app(db_session, *, role: str, tenant_id: int = 1) -> FastAPI:
|
||
app = FastAPI()
|
||
app.include_router(mcp_routes.router)
|
||
|
||
async def _fake_get_db():
|
||
yield db_session
|
||
|
||
app.dependency_overrides[get_db] = _fake_get_db
|
||
app.dependency_overrides[get_current_user] = lambda: SimpleNamespace(
|
||
id=7,
|
||
tenant_id=tenant_id,
|
||
role=role,
|
||
username="tester",
|
||
email="tester@example.com",
|
||
name="Tester",
|
||
)
|
||
return app
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_mcp_admin_api_requires_manager_role(db_session) -> None:
|
||
app = _build_app(db_session, role="AI")
|
||
transport = ASGITransport(app=app)
|
||
|
||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||
response = await client.get("/mcp/admin/catalog")
|
||
|
||
assert response.status_code == 403
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_mcp_admin_catalog_and_server_lifecycle(db_session, monkeypatch) -> None:
|
||
async def _fake_probe(self, *, template, params): # noqa: ANN001
|
||
_ = (self, template, params)
|
||
return {
|
||
"status": "ready",
|
||
"checked_at": "2026-04-15T12:01:00+08:00",
|
||
"message": "PMFlow backend can reach Microsoft Learn MCP.",
|
||
"tool_names": ["microsoft_docs_search", "microsoft_docs_fetch"],
|
||
}
|
||
|
||
monkeypatch.setattr(McpConfigService, "_probe_backend_runtime", _fake_probe)
|
||
app = _build_app(db_session, role="ADMIN")
|
||
transport = ASGITransport(app=app)
|
||
|
||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||
catalog_response = await client.get("/mcp/admin/catalog")
|
||
create_response = await client.post(
|
||
"/mcp/admin/servers",
|
||
json={
|
||
"name": "microsoft-learn",
|
||
"display_name": "Microsoft Learn MCP",
|
||
"template_id": "microsoft_learn",
|
||
"params": {},
|
||
},
|
||
)
|
||
|
||
created = create_response.json()["data"]
|
||
list_response = await client.get("/mcp/admin/servers")
|
||
validate_response = await client.post(f"/mcp/admin/servers/{created['id']}/validate")
|
||
update_response = await client.put(
|
||
f"/mcp/admin/servers/{created['id']}",
|
||
json={
|
||
"display_name": "Microsoft Learn Docs",
|
||
"params": {},
|
||
},
|
||
)
|
||
confirm_response = await client.post(f"/mcp/admin/servers/{created['id']}/confirm-client-apply")
|
||
|
||
assert catalog_response.status_code == 200
|
||
assert catalog_response.json()["code"] == 0
|
||
assert catalog_response.json()["data"]["recommended_template_id"] == "microsoft_learn"
|
||
assert catalog_response.json()["data"]["items"][0]["execution_mode"] == "backend_direct"
|
||
assert any(item["id"] == "duckduckgo" for item in catalog_response.json()["data"]["items"])
|
||
|
||
assert create_response.status_code == 200
|
||
assert create_response.json()["code"] == 0
|
||
assert created["execution_mode"] == "backend_direct"
|
||
assert created["client_apply_state"] == "not_required"
|
||
assert created["runtime_probe"]["status"] == "unknown"
|
||
|
||
assert list_response.status_code == 200
|
||
assert list_response.json()["data"]["items"][0]["name"] == "microsoft-learn"
|
||
|
||
assert validate_response.status_code == 200
|
||
assert validate_response.json()["code"] == 0
|
||
assert validate_response.json()["data"]["is_valid"] is True
|
||
assert validate_response.json()["data"]["preview"]["add_command"] == "无需本机命令。保存后由 PMFlow 后端直接连接 Microsoft Learn MCP。"
|
||
assert validate_response.json()["data"]["preview"]["verify_command"] == "点击“重新校验”后，PMFlow 会从后端执行 MCP initialize 和 tools/list。"
|
||
assert validate_response.json()["data"]["runtime_probe"]["status"] == "ready"
|
||
assert validate_response.json()["data"]["runtime_probe"]["tool_names"] == ["microsoft_docs_search", "microsoft_docs_fetch"]
|
||
|
||
assert update_response.status_code == 200
|
||
assert update_response.json()["code"] == 0
|
||
assert update_response.json()["data"]["display_name"] == "Microsoft Learn Docs"
|
||
|
||
assert confirm_response.status_code == 400
|
||
assert confirm_response.json()["code"] == 400
|
||
assert "后端直连" in confirm_response.json()["message"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_mcp_admin_servers_are_tenant_scoped(db_session) -> None:
|
||
tenant1_app = _build_app(db_session, role="ADMIN", tenant_id=1)
|
||
tenant2_app = _build_app(db_session, role="ADMIN", tenant_id=2)
|
||
|
||
async with AsyncClient(transport=ASGITransport(app=tenant1_app), base_url="http://test") as client:
|
||
create_response = await client.post(
|
||
"/mcp/admin/servers",
|
||
json={
|
||
"name": "playwright",
|
||
"display_name": "Playwright MCP",
|
||
"template_id": "playwright",
|
||
"params": {"output_dir": "~/.codex/mcp/playwright-output"},
|
||
},
|
||
)
|
||
|
||
async with AsyncClient(transport=ASGITransport(app=tenant2_app), base_url="http://test") as client:
|
||
list_response = await client.get("/mcp/admin/servers")
|
||
|
||
assert create_response.status_code == 200
|
||
assert create_response.json()["code"] == 0
|
||
assert list_response.status_code == 200
|
||
assert list_response.json()["code"] == 0
|
||
assert list_response.json()["data"]["items"] == []
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_mcp_admin_github_template_rejects_raw_token_and_only_persists_env_var_name(db_session) -> None:
|
||
app = _build_app(db_session, role="ADMIN")
|
||
transport = ASGITransport(app=app)
|
||
|
||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||
invalid_response = await client.post(
|
||
"/mcp/admin/servers",
|
||
json={
|
||
"name": "github",
|
||
"display_name": "GitHub MCP",
|
||
"template_id": "github",
|
||
"params": {
|
||
"token_env_var_name": "GITHUB_PAT",
|
||
"access_token": "ghp_secret_should_not_be_saved",
|
||
},
|
||
},
|
||
)
|
||
valid_response = await client.post(
|
||
"/mcp/admin/servers",
|
||
json={
|
||
"name": "github",
|
||
"display_name": "GitHub MCP",
|
||
"template_id": "github",
|
||
"params": {
|
||
"token_env_var_name": "GITHUB_PAT",
|
||
},
|
||
},
|
||
)
|
||
|
||
row = (
|
||
await db_session.execute(
|
||
select(Integration).where(
|
||
Integration.tenant_id == 1,
|
||
Integration.type == "mcp_template_center",
|
||
)
|
||
)
|
||
).scalar_one()
|
||
server = row.config["servers"][0]
|
||
|
||
assert invalid_response.status_code == 400
|
||
assert invalid_response.json()["code"] == 400
|
||
assert "token" in invalid_response.json()["message"].lower()
|
||
|
||
assert valid_response.status_code == 200
|
||
assert valid_response.json()["code"] == 0
|
||
assert valid_response.json()["data"]["execution_mode"] == "local_bridge"
|
||
assert server["params"]["token_env_var_name"] == "GITHUB_PAT"
|
||
assert "access_token" not in server["params"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_mcp_admin_filesystem_requires_allowed_paths(db_session) -> None:
|
||
app = _build_app(db_session, role="ADMIN")
|
||
transport = ASGITransport(app=app)
|
||
|
||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||
response = await client.post(
|
||
"/mcp/admin/servers",
|
||
json={
|
||
"name": "filesystem",
|
||
"display_name": "Filesystem MCP",
|
||
"template_id": "filesystem",
|
||
"params": {},
|
||
},
|
||
)
|
||
|
||
assert response.status_code == 400
|
||
assert response.json()["code"] == 400
|
||
assert "allowed_paths" in response.json()["message"]
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_local_bridge_template_uses_manual_instructions_and_can_be_confirmed(db_session) -> None:
|
||
app = _build_app(db_session, role="ADMIN")
|
||
transport = ASGITransport(app=app)
|
||
|
||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||
create_response = await client.post(
|
||
"/mcp/admin/servers",
|
||
json={
|
||
"name": "filesystem",
|
||
"display_name": "Filesystem MCP",
|
||
"template_id": "filesystem",
|
||
"params": {
|
||
"allowed_paths": ["/Users/yuuuu/Documents"],
|
||
},
|
||
},
|
||
)
|
||
created = create_response.json()["data"]
|
||
validate_response = await client.post(f"/mcp/admin/servers/{created['id']}/validate")
|
||
confirm_response = await client.post(f"/mcp/admin/servers/{created['id']}/confirm-client-apply")
|
||
|
||
assert create_response.status_code == 200
|
||
assert created["execution_mode"] == "local_bridge"
|
||
assert created["client_apply_state"] == "pending_local_apply"
|
||
assert "codex mcp add" not in created["preview"]["add_command"].lower()
|
||
assert "local bridge" in created["preview"]["add_command"].lower()
|
||
|
||
assert validate_response.status_code == 200
|
||
assert validate_response.json()["data"]["runtime_probe"]["status"] == "local_bridge_required"
|
||
|
||
assert confirm_response.status_code == 200
|
||
assert confirm_response.json()["data"]["client_apply_state"] == "user_confirmed"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_duckduckgo_template_uses_backend_direct_preview_without_extra_params(db_session, monkeypatch) -> None:
|
||
async def _fake_probe(self, *, template, params): # noqa: ANN001
|
||
_ = (self, params)
|
||
assert template.id == "duckduckgo"
|
||
return {
|
||
"status": "ready",
|
||
"checked_at": "2026-04-16T20:01:00+08:00",
|
||
"message": "PMFlow backend can directly reach DuckDuckGo search runtime.",
|
||
"tool_names": ["ddg-text-search", "ddg-news-search"],
|
||
}
|
||
|
||
monkeypatch.setattr(McpConfigService, "_probe_backend_runtime", _fake_probe)
|
||
app = _build_app(db_session, role="ADMIN")
|
||
transport = ASGITransport(app=app)
|
||
|
||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||
create_response = await client.post(
|
||
"/mcp/admin/servers",
|
||
json={
|
||
"name": "duckduckgo",
|
||
"display_name": "DuckDuckGo MCP",
|
||
"template_id": "duckduckgo",
|
||
"params": {},
|
||
},
|
||
)
|
||
created = create_response.json()["data"]
|
||
validate_response = await client.post(f"/mcp/admin/servers/{created['id']}/validate")
|
||
|
||
assert create_response.status_code == 200
|
||
assert created["execution_mode"] == "backend_direct"
|
||
assert created["client_apply_state"] == "not_required"
|
||
assert created["params"] == {}
|
||
assert "无需本机命令" in created["preview"]["add_command"]
|
||
assert "ddg-text-search" in created["preview"]["expected_output"]
|
||
|
||
assert validate_response.status_code == 200
|
||
assert validate_response.json()["data"]["runtime_probe"]["status"] == "ready"
|