mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 06:12:49 +00:00
Alex
574f96341e
The backend import package is now docsgpt, the name it will carry on PyPI; application was far too generic to install into anyone's site-packages. git mv plus a mechanical rewrite of every import, dotted string and path reference: 734 Python files, the compose files, Dockerfile, workflows, docs, setup scripts, devcontainer, k8s manifests, vscode config, pytest and coverage config, .gitignore. Behaviour is unchanged. Kept for one release: - A top-level application package whose meta-path finder resolves application.x.y to the already-imported docsgpt.x.y object, so old imports and entry points (celery -A application.app.celery, uvicorn application.asgi:asgi_app) keep working with a FutureWarning. - Celery registers every application.* task name as an alias of its docsgpt.* task on start-up, so messages queued by the previous release still run. The redbeat key prefix moves to redbeat:docsgpt:v2: so schedule entries the previous release wrote are left unread instead of firing twice. The backend image builds from the repository root (docker build -f docsgpt/Dockerfile .) so it can ship the alias package; a root .dockerignore allow-lists docsgpt/ and application/ and keeps caches, local data, .env files, the sample index files and the Dockerfile out. Compose and the image workflows point at the new context.
463 lines
16 KiB
Python
463 lines
16 KiB
Python
"""Tests for docsgpt/api/user/tools/mcp.py using real PG."""
|
|
|
|
from contextlib import contextmanager
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
from flask import Flask
|
|
|
|
|
|
@pytest.fixture
|
|
def app():
|
|
return Flask(__name__)
|
|
|
|
|
|
@contextmanager
|
|
def _patch_db(conn):
|
|
@contextmanager
|
|
def _yield():
|
|
yield conn
|
|
|
|
with patch(
|
|
"docsgpt.api.user.tools.mcp.db_session", _yield
|
|
), patch(
|
|
"docsgpt.api.user.tools.mcp.db_readonly", _yield
|
|
):
|
|
yield
|
|
|
|
|
|
class TestSanitizeMcpTransport:
|
|
def test_defaults_to_auto(self):
|
|
from docsgpt.api.user.tools.mcp import _sanitize_mcp_transport
|
|
cfg = {}
|
|
got = _sanitize_mcp_transport(cfg)
|
|
assert got == "auto"
|
|
assert cfg["transport_type"] == "auto"
|
|
|
|
def test_accepts_supported_transports(self):
|
|
from docsgpt.api.user.tools.mcp import _sanitize_mcp_transport
|
|
for t in ("auto", "sse", "http"):
|
|
cfg = {"transport_type": t}
|
|
assert _sanitize_mcp_transport(cfg) == t
|
|
|
|
def test_strips_command_and_args(self):
|
|
from docsgpt.api.user.tools.mcp import _sanitize_mcp_transport
|
|
cfg = {"transport_type": "http", "command": "/bin/x", "args": ["a"]}
|
|
_sanitize_mcp_transport(cfg)
|
|
assert "command" not in cfg
|
|
assert "args" not in cfg
|
|
|
|
def test_unsupported_transport_raises(self):
|
|
from docsgpt.api.user.tools.mcp import _sanitize_mcp_transport
|
|
with pytest.raises(ValueError):
|
|
_sanitize_mcp_transport({"transport_type": "websocket"})
|
|
|
|
|
|
class TestExtractAuthCredentials:
|
|
def test_api_key_auth(self):
|
|
from docsgpt.api.user.tools.mcp import _extract_auth_credentials
|
|
got = _extract_auth_credentials({
|
|
"auth_type": "api_key",
|
|
"api_key": "secret",
|
|
"api_key_header": "X-API-Key",
|
|
})
|
|
assert got == {"api_key": "secret", "api_key_header": "X-API-Key"}
|
|
|
|
def test_bearer_auth(self):
|
|
from docsgpt.api.user.tools.mcp import _extract_auth_credentials
|
|
got = _extract_auth_credentials({
|
|
"auth_type": "bearer",
|
|
"bearer_token": "my-token",
|
|
})
|
|
assert got == {"bearer_token": "my-token"}
|
|
|
|
def test_basic_auth(self):
|
|
from docsgpt.api.user.tools.mcp import _extract_auth_credentials
|
|
got = _extract_auth_credentials({
|
|
"auth_type": "basic",
|
|
"username": "u", "password": "p",
|
|
})
|
|
assert got == {"username": "u", "password": "p"}
|
|
|
|
def test_none_auth_empty_creds(self):
|
|
from docsgpt.api.user.tools.mcp import _extract_auth_credentials
|
|
assert _extract_auth_credentials({"auth_type": "none"}) == {}
|
|
|
|
|
|
class TestValidateMcpServerUrl:
|
|
def test_empty_url_raises(self):
|
|
from docsgpt.api.user.tools.mcp import _validate_mcp_server_url
|
|
with pytest.raises(ValueError):
|
|
_validate_mcp_server_url({})
|
|
|
|
def test_missing_server_url(self):
|
|
from docsgpt.api.user.tools.mcp import _validate_mcp_server_url
|
|
with pytest.raises(ValueError):
|
|
_validate_mcp_server_url({"server_url": ""})
|
|
|
|
def test_ssrf_url_raises(self):
|
|
from docsgpt.api.user.tools.mcp import _validate_mcp_server_url
|
|
with pytest.raises(ValueError):
|
|
_validate_mcp_server_url({"server_url": "http://127.0.0.1"})
|
|
|
|
def test_valid_public_url_passes(self):
|
|
from docsgpt.api.user.tools.mcp import _validate_mcp_server_url
|
|
# Should not raise for a public-ish URL
|
|
try:
|
|
_validate_mcp_server_url({"server_url": "https://example.com/mcp"})
|
|
except ValueError as e:
|
|
# If SSRF rules reject example.com for some reason, accept that
|
|
if "Invalid" not in str(e):
|
|
raise
|
|
|
|
|
|
class TestTestMCPServerConfig:
|
|
def test_returns_401_unauthenticated(self, app):
|
|
from docsgpt.api.user.tools.mcp import TestMCPServerConfig
|
|
|
|
with app.test_request_context(
|
|
"/api/mcp_server/test", method="POST",
|
|
json={"config": {}},
|
|
):
|
|
from flask import request
|
|
request.decoded_token = None
|
|
response = TestMCPServerConfig().post()
|
|
assert response.status_code == 401
|
|
|
|
def test_returns_400_missing_config(self, app):
|
|
from docsgpt.api.user.tools.mcp import TestMCPServerConfig
|
|
|
|
with app.test_request_context(
|
|
"/api/mcp_server/test", method="POST", json={},
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = TestMCPServerConfig().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_unsupported_transport_returns_400(self, app):
|
|
from docsgpt.api.user.tools.mcp import TestMCPServerConfig
|
|
|
|
with app.test_request_context(
|
|
"/api/mcp_server/test", method="POST",
|
|
json={"config": {"transport_type": "websocket"}},
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = TestMCPServerConfig().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_missing_url_returns_400(self, app):
|
|
from docsgpt.api.user.tools.mcp import TestMCPServerConfig
|
|
|
|
with app.test_request_context(
|
|
"/api/mcp_server/test", method="POST",
|
|
json={"config": {"transport_type": "http"}},
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = TestMCPServerConfig().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_connection_success(self, app):
|
|
from docsgpt.api.user.tools.mcp import TestMCPServerConfig
|
|
|
|
fake_tool = MagicMock()
|
|
fake_tool.test_connection.return_value = {
|
|
"success": True, "message": "OK",
|
|
"tools_count": 3, "tools": ["a", "b", "c"],
|
|
}
|
|
|
|
with patch(
|
|
"docsgpt.api.user.tools.mcp.MCPTool",
|
|
return_value=fake_tool,
|
|
), app.test_request_context(
|
|
"/api/mcp_server/test", method="POST",
|
|
json={
|
|
"config": {
|
|
"transport_type": "http",
|
|
"server_url": "https://example.com/mcp",
|
|
"auth_type": "none",
|
|
},
|
|
},
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = TestMCPServerConfig().post()
|
|
assert response.status_code == 200
|
|
assert response.json["success"] is True
|
|
assert response.json["tools_count"] == 3
|
|
|
|
def test_connection_failure_returns_200_with_failure_message(self, app):
|
|
from docsgpt.api.user.tools.mcp import TestMCPServerConfig
|
|
|
|
fake_tool = MagicMock()
|
|
fake_tool.test_connection.return_value = {
|
|
"success": False, "message": "Cannot reach server",
|
|
}
|
|
|
|
with patch(
|
|
"docsgpt.api.user.tools.mcp.MCPTool",
|
|
return_value=fake_tool,
|
|
), app.test_request_context(
|
|
"/api/mcp_server/test", method="POST",
|
|
json={
|
|
"config": {
|
|
"transport_type": "http",
|
|
"server_url": "https://example.com/mcp",
|
|
},
|
|
},
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = TestMCPServerConfig().post()
|
|
assert response.status_code == 200
|
|
assert response.json["success"] is False
|
|
|
|
def test_oauth_required_returns_200(self, app):
|
|
from docsgpt.api.user.tools.mcp import TestMCPServerConfig
|
|
|
|
fake_tool = MagicMock()
|
|
fake_tool.test_connection.return_value = {
|
|
"success": False,
|
|
"requires_oauth": True,
|
|
"auth_url": "https://auth/ex",
|
|
}
|
|
|
|
with patch(
|
|
"docsgpt.api.user.tools.mcp.MCPTool",
|
|
return_value=fake_tool,
|
|
), app.test_request_context(
|
|
"/api/mcp_server/test", method="POST",
|
|
json={
|
|
"config": {
|
|
"transport_type": "http",
|
|
"server_url": "https://example.com/mcp",
|
|
"auth_type": "oauth",
|
|
},
|
|
},
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = TestMCPServerConfig().post()
|
|
assert response.status_code == 200
|
|
assert response.json["requires_oauth"] is True
|
|
|
|
def test_unexpected_exception_returns_500(self, app):
|
|
from docsgpt.api.user.tools.mcp import TestMCPServerConfig
|
|
|
|
with patch(
|
|
"docsgpt.api.user.tools.mcp.MCPTool",
|
|
side_effect=RuntimeError("boom"),
|
|
), app.test_request_context(
|
|
"/api/mcp_server/test", method="POST",
|
|
json={
|
|
"config": {
|
|
"transport_type": "http",
|
|
"server_url": "https://example.com/mcp",
|
|
},
|
|
},
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = TestMCPServerConfig().post()
|
|
assert response.status_code == 500
|
|
|
|
|
|
class TestMCPServerSave:
|
|
def test_returns_401_unauthenticated(self, app):
|
|
from docsgpt.api.user.tools.mcp import MCPServerSave
|
|
|
|
with app.test_request_context(
|
|
"/api/mcp_server/save", method="POST",
|
|
json={"displayName": "n", "config": {}},
|
|
):
|
|
from flask import request
|
|
request.decoded_token = None
|
|
response = MCPServerSave().post()
|
|
assert response.status_code == 401
|
|
|
|
def test_returns_400_missing_fields(self, app):
|
|
from docsgpt.api.user.tools.mcp import MCPServerSave
|
|
|
|
with app.test_request_context(
|
|
"/api/mcp_server/save", method="POST", json={},
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = MCPServerSave().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_unsupported_transport_returns_400(self, app):
|
|
from docsgpt.api.user.tools.mcp import MCPServerSave
|
|
|
|
with app.test_request_context(
|
|
"/api/mcp_server/save", method="POST",
|
|
json={
|
|
"displayName": "Srv",
|
|
"config": {"transport_type": "bogus"},
|
|
},
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = MCPServerSave().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_missing_server_url_returns_400(self, app):
|
|
from docsgpt.api.user.tools.mcp import MCPServerSave
|
|
|
|
with app.test_request_context(
|
|
"/api/mcp_server/save", method="POST",
|
|
json={"displayName": "Srv", "config": {"transport_type": "http"}},
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = MCPServerSave().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_oauth_missing_task_id_returns_400(self, app):
|
|
from docsgpt.api.user.tools.mcp import MCPServerSave
|
|
|
|
with app.test_request_context(
|
|
"/api/mcp_server/save", method="POST",
|
|
json={
|
|
"displayName": "Srv",
|
|
"config": {
|
|
"transport_type": "http",
|
|
"server_url": "https://example.com/mcp",
|
|
"auth_type": "oauth",
|
|
},
|
|
},
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": "u"}
|
|
response = MCPServerSave().post()
|
|
assert response.status_code == 400
|
|
|
|
def test_creates_mcp_tool_successfully(self, app, pg_conn):
|
|
from docsgpt.api.user.tools.mcp import MCPServerSave
|
|
|
|
user = "u-mcp-save"
|
|
fake_tool = MagicMock()
|
|
fake_tool.discover_tools.return_value = {"tools": ["t1"]}
|
|
fake_tool.get_actions_metadata.return_value = [{"name": "t1"}]
|
|
|
|
with _patch_db(pg_conn), patch(
|
|
"docsgpt.api.user.tools.mcp.MCPTool",
|
|
return_value=fake_tool,
|
|
), app.test_request_context(
|
|
"/api/mcp_server/save", method="POST",
|
|
json={
|
|
"displayName": "My MCP",
|
|
"config": {
|
|
"transport_type": "http",
|
|
"server_url": "https://example.com/mcp",
|
|
"auth_type": "none",
|
|
},
|
|
"status": True,
|
|
},
|
|
):
|
|
from flask import request
|
|
request.decoded_token = {"sub": user}
|
|
response = MCPServerSave().post()
|
|
assert response.status_code in (200, 201)
|
|
|
|
|
|
class TestMCPOAuthCallback:
|
|
def test_error_param_redirects_error(self, app):
|
|
from docsgpt.api.user.tools.mcp import MCPOAuthCallback
|
|
|
|
with app.test_request_context(
|
|
"/api/mcp_server/oauth_callback?error=access_denied"
|
|
):
|
|
response = MCPOAuthCallback().get()
|
|
assert response.status_code == 302
|
|
assert "status=error" in response.location
|
|
|
|
def test_missing_code_or_state_redirects_error(self, app):
|
|
from docsgpt.api.user.tools.mcp import MCPOAuthCallback
|
|
|
|
with app.test_request_context(
|
|
"/api/mcp_server/oauth_callback"
|
|
):
|
|
response = MCPOAuthCallback().get()
|
|
assert response.status_code == 302
|
|
|
|
def test_success_redirects_success(self, app):
|
|
from docsgpt.api.user.tools.mcp import MCPOAuthCallback
|
|
|
|
fake_redis = MagicMock()
|
|
fake_manager = MagicMock()
|
|
fake_manager.handle_oauth_callback.return_value = True
|
|
|
|
with patch(
|
|
"docsgpt.api.user.tools.mcp.get_redis_instance",
|
|
return_value=fake_redis,
|
|
), patch(
|
|
"docsgpt.api.user.tools.mcp.MCPOAuthManager",
|
|
return_value=fake_manager,
|
|
), app.test_request_context(
|
|
"/api/mcp_server/oauth_callback?code=c1&state=s1"
|
|
):
|
|
response = MCPOAuthCallback().get()
|
|
assert response.status_code == 302
|
|
assert "status=success" in response.location
|
|
|
|
def test_manager_failure_redirects_error(self, app):
|
|
from docsgpt.api.user.tools.mcp import MCPOAuthCallback
|
|
|
|
fake_redis = MagicMock()
|
|
fake_manager = MagicMock()
|
|
fake_manager.handle_oauth_callback.return_value = False
|
|
|
|
with patch(
|
|
"docsgpt.api.user.tools.mcp.get_redis_instance",
|
|
return_value=fake_redis,
|
|
), patch(
|
|
"docsgpt.api.user.tools.mcp.MCPOAuthManager",
|
|
return_value=fake_manager,
|
|
), app.test_request_context(
|
|
"/api/mcp_server/oauth_callback?code=c1&state=s1"
|
|
):
|
|
response = MCPOAuthCallback().get()
|
|
assert response.status_code == 302
|
|
assert "status=error" in response.location
|
|
|
|
def test_no_redis_redirects_error(self, app):
|
|
from docsgpt.api.user.tools.mcp import MCPOAuthCallback
|
|
|
|
with patch(
|
|
"docsgpt.api.user.tools.mcp.get_redis_instance",
|
|
return_value=None,
|
|
), app.test_request_context(
|
|
"/api/mcp_server/oauth_callback?code=c&state=s"
|
|
):
|
|
response = MCPOAuthCallback().get()
|
|
assert response.status_code == 302
|
|
assert "Redis" in response.location or "status=error" in response.location
|
|
|
|
def test_exception_redirects_error(self, app):
|
|
from docsgpt.api.user.tools.mcp import MCPOAuthCallback
|
|
|
|
with patch(
|
|
"docsgpt.api.user.tools.mcp.get_redis_instance",
|
|
side_effect=RuntimeError("boom"),
|
|
), app.test_request_context(
|
|
"/api/mcp_server/oauth_callback?code=c&state=s"
|
|
):
|
|
response = MCPOAuthCallback().get()
|
|
assert response.status_code == 302
|
|
|
|
|
|
class TestMCPAuthStatus:
|
|
def test_returns_401_unauthenticated(self, app):
|
|
from docsgpt.api.user.tools.mcp import MCPAuthStatus
|
|
|
|
with app.test_request_context(
|
|
"/api/mcp_server/auth_status"
|
|
):
|
|
from flask import request
|
|
request.decoded_token = None
|
|
response = MCPAuthStatus().get()
|
|
assert response.status_code == 401
|