mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-03 22:13:01 +00:00
Alex
e351f45d88
* feat: SSE notification system
Adds a per-user SSE pipe (GET /api/events) plus a per-message
chat-stream reconnect endpoint (GET /api/messages/<id>/events).
Backend substrate:
- application/events/ — durable journal (Redis Streams) + live
pub/sub for user-scoped events, with publish_user_event() as
the worker-side entrypoint.
- application/streaming/ — broadcast_channel for pub/sub fanout
and event_replay for the per-message snapshot+tail path.
- application/storage/db/repositories/message_events.py +
alembic 0007 — Postgres journal for chat-stream events.
- application/worker.py — ingest/reingest/remote/connector/
attachment/mcp_oauth tasks publish queued/progress/completed/
failed envelopes alongside their existing status updates.
Frontend client:
- frontend/src/events/ — connect/reconnect, Last-Event-ID cursor,
backoff with jitter. Each tab runs its own connection; no
cross-tab dedup (future work).
- frontend/src/notifications/ — recentEvents ring, cursor
tracking, tool-approval toast.
- frontend/src/upload/uploadSlice.ts — extraReducers for
source.ingest.* and attachment.* events.
Coverage: 132 SSE tests across events substrate, replay, journal,
routes, and worker publishes.
* refactor(attachments): remove polling, SSE-only
frontend/src/components/MessageInput.tsx no longer runs a 2s
setInterval against getTaskStatus for every processing
attachment. The attachment.* SSE reducers in uploadSlice.ts are
now the sole driver of attachment state transitions.
* feat(connector): consume source.ingest.* SSE, remove polling
frontend/src/components/ConnectorTree.tsx now mirrors FileTree's
slice-walking pattern: it watches notifications.recentEvents
for source.ingest.{completed,failed} envelopes matching the
sync's source id, and no longer polls /task_status every 2s.
* refactor(source-ingest): remove polling, SSE-only
frontend/src/upload/Upload.tsx and
frontend/src/components/FileTree.tsx no longer run getTaskStatus
polling fallbacks. The source.ingest.* SSE reducers in
uploadSlice.ts and FileTree's slice walk are now the sole
drivers of upload/reingest state transitions.
* refactor(mcp-oauth): carry authorization_url in SSE, remove polling
application/worker.py::mcp_oauth now publishes
authorization_url on the mcp.oauth.awaiting_redirect envelope.
frontend/src/modals/MCPServerModal.tsx consumes it from SSE
instead of polling /oauth_status/<task_id> every 1s.
The URL is generated inside DocsGPTOAuth.redirect_handler when
the FastMCP client triggers OAuth. The worker now plumbs a
publish callback through tool_config -> MCPTool -> DocsGPTOAuth
so the awaiting_redirect publish fires from inside the handler
at the exact point the URL becomes known. The legacy Redis
mcp_oauth_status setex writes and the GET
/api/mcp_server/oauth_status/<task_id> endpoint are kept as
belt-and-suspenders; nothing in the frontend reads them now.
* feat(source-ingest): plumb limited flag through SSE for token-cap UX
application/worker.py::ingest_worker and remote_worker now publish
``limited: bool`` on the source.ingest.completed envelope.
uploadSlice routes ``payload.limited === true`` to a failed status
with a ``tokenLimitReached`` flag, and UploadToast surfaces the
translated tokenLimit i18n string. No worker code path sets
limited=true today; this is a forward-looking contract so when
token-cap detection lands, the UX is already wired.
* refactor(mcp-oauth): read status from SSE journal, drop polling endpoint
MCPOAuthManager.get_oauth_status now walks the per-user SSE Streams
journal (user:{user_id}:stream) for the latest mcp.oauth.* envelope
matching the task id, returning the status string derived from the
event type suffix and the payload fields. The worker is the single
source of truth — its publish_user_event calls write the same
record the SSE client receives live.
Removed:
- /api/mcp_server/oauth_status/<task_id> route in
application/api/user/tools/mcp.py
- mcp_oauth_status worker function and mcp_oauth_status_task Celery
wrapper
- All mcp_oauth_status:{task_id} Redis setex writes (4 in mcp_oauth,
2 in DocsGPTOAuth.redirect_handler / callback_handler)
- The update_status closure in mcp_oauth that wrote the polling
payload
Tests updated:
- get_oauth_status now takes (task_id, user_id); new coverage walks
a fake xrevrange response for the completed envelope, the no-match
case, and a Redis-down case
- Removed TestMCPOAuthStatus route tests and TestMcpOauthStatusTask
celery-wrapper test
- Removed the two oauth_status methods from the integration runner
mcp_oauth:auth_url/state/code/error Redis keys remain — they are
the OAuth flow's own state (not the dropped polling payload).
* chore(mcp-oauth): delete orphaned getMCPOAuthStatus client
The /api/mcp_server/oauth_status/<task_id> endpoint was removed in
the prior commit; the corresponding userService method and the
MCP_OAUTH_STATUS endpoint constant had no remaining callers in the
frontend, so they're deleted along with it.
* fix(events): drop live publish when journal write fails
application/events/publisher.py returned an envelope to live
pubsub subscribers even when the XADD to the durable journal
failed. The envelope had no ``id`` field, which bypassed the SSE
route's dedup floor and broke ``Last-Event-ID`` semantics for any
reconnecting client.
Best-effort delivery means dropping consistently, not delivering
inconsistent state. Now: if the journal write fails the publisher
returns None and skips the live publish entirely.
* fix(notifications): dedupe sseEventReceived against immediate dupes
Snapshot replay + live tail can both deliver the same id when the
live pubsub frame and the replay XRANGE overlap. The route's own
dedup floor catches the common case, but consumers walking
``recentEvents`` (FileTree, ConnectorTree, MCPServerModal,
ToolApprovalToast) would otherwise act on the same envelope
twice when a duplicate slipped through.
Belt-and-suspenders: short-circuit when the most recent id in
the ring matches the incoming one.
* fix(events): skip replay budget INCR when no snapshot work possible
_allow_replay incremented the per-user counter on every
/api/events GET, including no-op connects from a fresh client
with no cursor against an empty backlog. React StrictMode dev
double-mounts plus a few tabs trivially tripped the default
30-per-60s budget on idle reconnects.
XLEN pre-check: when last_event_id is None and the user stream
is empty, the connect can't do snapshot work — return True
without INCR. Cursor-bearing connects still INCR unconditionally
(probing the cursor's relationship to stream contents would
require a redundant XRANGE).
* fix(streaming): tighten journal contract + recover from seq collisions
Two related fixes to application/streaming/message_journal.py.
1. record_event now rejects non-dict payloads at the gate. The
live path (base.py::_emit) wrapped non-dicts as
{"value": payload}; the replay path in event_replay synthesized
{"type": event_type}. A reconnecting client would receive a
different envelope than the one originally streamed. Now both
paths see byte-identical envelopes because non-dicts can't be
journaled at all. The corresponding event_replay fallback is
replaced with a warn-and-skip for any legacy rows.
2. record_event handles IntegrityError on (message_id, sequence_no)
collisions by reading latest_sequence_no and retrying once with
latest+1. The most likely cause is a stale seq seed on a
continuation retry where the route read MAX(seq) from a
separate connection before another writer committed past it.
Previously the error was swallowed and the event silently
dropped from the journal; now it lands at the next available
seq. The live pubsub publish uses the materialised seq so the
journal row and the live frame agree.
* perf(streaming): batch message_events INSERTs per stream
complete_stream previously opened a fresh db_session() per yielded
event, doing one Postgres INSERT + commit per chunk on the WSGI
thread. Streaming answers emit ~100s of answer chunks per response,
so the route was paying ~100 PG roundtrips per stream serialized on
commit latency.
New BatchedJournalWriter in application/streaming/message_journal.py
accumulates rows per stream and flushes on three triggers:
- size: buffer reaches 16 entries
- time: 100ms elapsed since the last flush
- lifecycle: close() at end-of-stream
Live pubsub publishes still fire synchronously per record(), so
subscribers see events in real time — only the durable journal write
is amortized. On bulk INSERT IntegrityError the writer falls back to
per-row record() with the existing seq+1 retry so a single colliding
seq doesn't drop the rest of the batch.
complete_stream wires journal_writer.close() into every exit path
(happy end, tool-approval-paused end, GeneratorExit, error handler)
so the terminal event is committed before the generator returns —
otherwise a reconnecting client could snapshot up to the last flush
boundary and live-tail waiting for an end that's still in memory.
Repository gets bulk_record() — one SQLAlchemy executemany INSERT
for the bulk path. All-or-nothing on collision (Postgres aborts the
whole batch); the writer's per-row fallback handles recovery.
* chore(upload): drop dead UploadTask.lastEventAt field
The lastEventAt field on UploadTask had no remaining consumers — the
matching Attachment.lastEventAt was cleaned up earlier. Remove the
field declaration and the slice write site.
* chore(frontend): drop orphaned getTaskStatus client
After the polling-removal sweep no caller in frontend/src/ references
userService.getTaskStatus or endpoints.USER.TASK_STATUS. The backend
route /api/task_status itself stays — agents, webhooks, e2e specs,
and the public docs still depend on it.
* docs(repo): remove stale planning docs from repo root
notification-channel-design.md, plan.md, and reminder-tool-design.md
were leftover Claude planning artifacts from the SSE substrate work
that landed accidentally. CLAUDE.md prohibits creating planning docs
unless asked — delete them.
* docs(message-events): clarify repo vs wrapper payload contract
MessageEventsRepository.record accepts any JSONB-compatible value; the
streaming wrapper record_event tightens this to dicts only because the
live and replay paths reconstruct non-dict payloads differently. Spell
the split out so the next reader of the repo method doesn't assume the
wrapper's contract applies here.
* refactor(events): raise on malformed stream id instead of lex fallback
stream_id_compare's lex-fallback branch was a footgun: a malformed id
that sorts lex-greater than a real one would pin live-tail dedup
forever, dropping every subsequent legitimate event silently. Both
current callers in application/api/events/routes.py pre-validate
inputs against _STREAM_ID_RE before calling, so changing the function
to raise ValueError is a no-op on the happy path and turns the future-
caller footgun into a loud failure.
* test(tasks): cover cleanup_message_events task body
Adds skipped-when-no-POSTGRES_URI and happy-path coverage for the
Celery janitor. The skipped path returns the documented short-circuit
shape without touching the repo. The happy path seeds a backdated
row, runs the task against the pg_conn fixture, and asserts the
retention window's row is deleted while in-window rows survive.
Mirrors the TestCleanupPendingToolState pattern.
* fix(notifications): treat /c/new as no current conversation
useMatch('/c/:conversationId') treats the literal URL /c/new as a
real conversation id, so the toast suppression check confused
'user is on /c/new' with 'user is on the conversation needing
approval'. Explicit guard: when the matched id is 'new', fall
through to the no-match case so approval toasts still surface.
* docs(events): enumerate publish_user_event None-return paths
The function returns Optional[str] today, with None conflating five
distinct outcomes (missing args / push disabled / unserialisable /
Redis down / XADD failed). Every current call site is fire-and-
forget and ignores the return, so the right move is to document the
five cases rather than promote to an enum return — keeps the API
small while making the diagnostic surface (logs) obvious. If a
future caller needs to react differently per reason, promote then.
* refactor(sources): move source-id derivation out of worker module
application/api/user/sources/upload.py imported _derive_source_id
from application.worker — pulling the entire Celery worker module
into the API process at import time just for a two-line helper.
Move DOCSGPT_INGEST_NAMESPACE and the derivation function to a
new application/storage/db/source_ids.py module that both layers
can import without that dependency edge. worker.py re-exports the
old names (_derive_source_id, DOCSGPT_INGEST_NAMESPACE) for
backward-compatible imports from tests and any other in-tree
callers; new code should import from the new module directly.
* fix(cache): enable Redis health_check_interval to surface half-open TCP
Without health_check_interval, a half-open TCP socket (NAT silently
dropped state, ELB idle-close) can leave pubsub.get_message hanging
past the SSE generator's keepalive cadence — the kernel never
surfaces the dead socket because no payload is in flight. Setting
health_check_interval=10 makes redis-py ping every 10s when
otherwise idle, so the next get_message after the dead window
raises and the SSE loop falls into its reconnect path instead of
silently freezing on the user.
* chore(events): rename attachment.processing.progress to attachment.progress
The event-type taxonomy was inconsistent: source ingest emits
source.ingest.progress (three segments) while attachments emitted
attachment.processing.progress (four segments). Drops the
.processing. infix for parity. Worker publish sites, the slice
reducer's match, and the worker tests all flip together.
No external consumers — the event type is purely internal between
the publisher and the in-tab slice; safe to rename in one commit.
* feat: events cleanup
* fix: better docs
* fix: e2e tests
2661 lines
93 KiB
Python
2661 lines
93 KiB
Python
"""Comprehensive tests for application/agents/tools/mcp_tool.py
|
|
|
|
Covers: MCPTool init, cache key generation, transport creation, tool formatting,
|
|
result formatting, execute_action, discover_tools, test_connection,
|
|
get_actions_metadata, DocsGPTOAuth, NonInteractiveOAuth, DBTokenStorage,
|
|
MCPOAuthManager.
|
|
"""
|
|
|
|
import asyncio
|
|
import concurrent.futures
|
|
import json
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
|
|
# ---- Fixtures to isolate module-level side effects ----
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _patch_mcp_globals(monkeypatch):
|
|
"""Patch module-level cache to avoid real connections.
|
|
|
|
MongoDB is no longer used at module level; DBTokenStorage now backs
|
|
onto the ``connector_sessions`` Postgres repository. The cache patch
|
|
is still required to avoid hitting real Redis.
|
|
"""
|
|
import sys
|
|
|
|
if "application.agents.tools.mcp_tool" in sys.modules:
|
|
mcp_mod = sys.modules["application.agents.tools.mcp_tool"]
|
|
else:
|
|
mock_tasks = MagicMock()
|
|
monkeypatch.setitem(sys.modules, "application.api.user.tasks", mock_tasks)
|
|
import application.agents.tools.mcp_tool as mcp_mod
|
|
|
|
monkeypatch.setattr(mcp_mod, "_mcp_clients_cache", {})
|
|
# Bypass DNS-resolving URL validation for tests using fake hostnames.
|
|
monkeypatch.setattr(mcp_mod, "validate_url", lambda u, **kw: u)
|
|
|
|
|
|
@pytest.fixture
|
|
def mcp_config():
|
|
return {
|
|
"server_url": "https://mcp.example.com/api",
|
|
"transport_type": "http",
|
|
"auth_type": "none",
|
|
"timeout": 10,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def bearer_config():
|
|
return {
|
|
"server_url": "https://mcp.example.com/api",
|
|
"transport_type": "http",
|
|
"auth_type": "bearer",
|
|
"auth_credentials": {"bearer_token": "tok_123"},
|
|
"timeout": 10,
|
|
}
|
|
|
|
|
|
def _make_tool(config, **kwargs):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
with patch.object(MCPTool, "_setup_client"):
|
|
return MCPTool(config, **kwargs)
|
|
|
|
|
|
# =====================================================================
|
|
# MCPTool Initialization
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestMCPToolInit:
|
|
|
|
def test_basic_init(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
assert tool.server_url == "https://mcp.example.com/api"
|
|
assert tool.transport_type == "http"
|
|
assert tool.auth_type == "none"
|
|
assert tool.timeout == 10
|
|
assert tool.available_tools == []
|
|
|
|
def test_bearer_auth_credentials(self, bearer_config):
|
|
tool = _make_tool(bearer_config)
|
|
assert tool.auth_credentials["bearer_token"] == "tok_123"
|
|
|
|
def test_no_server_url_skips_setup(self):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
with patch.object(MCPTool, "_setup_client") as mock_setup:
|
|
MCPTool({"server_url": "", "auth_type": "none"})
|
|
mock_setup.assert_not_called()
|
|
|
|
def test_oauth_skips_setup(self):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
with patch.object(MCPTool, "_setup_client") as mock_setup:
|
|
MCPTool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "oauth",
|
|
})
|
|
mock_setup.assert_not_called()
|
|
|
|
def test_encrypted_credentials_decryption(self):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
with patch.object(MCPTool, "_setup_client"), \
|
|
patch("application.agents.tools.mcp_tool.decrypt_credentials",
|
|
return_value={"bearer_token": "decrypted_tok"}):
|
|
tool = MCPTool(
|
|
{
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "bearer",
|
|
"encrypted_credentials": "enc_data",
|
|
},
|
|
user_id="user1",
|
|
)
|
|
assert tool.auth_credentials == {"bearer_token": "decrypted_tok"}
|
|
|
|
def test_query_mode_default_false(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
assert tool.query_mode is False
|
|
|
|
def test_query_mode_explicit_true(self):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "none",
|
|
"query_mode": True,
|
|
})
|
|
assert tool.query_mode is True
|
|
|
|
def test_custom_headers_stored(self):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "none",
|
|
"headers": {"X-Custom": "val"},
|
|
})
|
|
assert tool.custom_headers == {"X-Custom": "val"}
|
|
|
|
def test_rejects_metadata_ip(self, monkeypatch):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
from application.core.url_validation import validate_url as real_validate_url
|
|
import application.agents.tools.mcp_tool as mcp_mod
|
|
|
|
monkeypatch.setattr(mcp_mod, "validate_url", real_validate_url)
|
|
with pytest.raises(ValueError, match="Invalid MCP server URL"):
|
|
MCPTool(config={"server_url": "http://169.254.169.254/latest/meta-data", "auth_type": "none"})
|
|
|
|
def test_rejects_localhost(self, monkeypatch):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
from application.core.url_validation import validate_url as real_validate_url
|
|
import application.agents.tools.mcp_tool as mcp_mod
|
|
|
|
monkeypatch.setattr(mcp_mod, "validate_url", real_validate_url)
|
|
with pytest.raises(ValueError, match="Invalid MCP server URL"):
|
|
MCPTool(config={"server_url": "http://localhost:8080/mcp", "auth_type": "none"})
|
|
|
|
def test_rejects_private_ip(self, monkeypatch):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
from application.core.url_validation import validate_url as real_validate_url
|
|
import application.agents.tools.mcp_tool as mcp_mod
|
|
|
|
monkeypatch.setattr(mcp_mod, "validate_url", real_validate_url)
|
|
with pytest.raises(ValueError, match="Invalid MCP server URL"):
|
|
MCPTool(config={"server_url": "http://10.0.0.1/mcp", "auth_type": "none"})
|
|
|
|
def test_accepts_public_url(self):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com/api",
|
|
"auth_type": "none",
|
|
})
|
|
assert tool.server_url == "https://mcp.example.com/api"
|
|
|
|
def test_empty_server_url_allowed(self):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
with patch.object(MCPTool, "_setup_client"):
|
|
tool = MCPTool(config={"server_url": "", "auth_type": "none"})
|
|
assert tool.server_url == ""
|
|
|
|
|
|
# =====================================================================
|
|
# Redirect URI Resolution
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestResolveRedirectUri:
|
|
|
|
def test_configured_redirect_uri_used(self):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "none",
|
|
"redirect_uri": "https://my.app/callback/",
|
|
})
|
|
assert tool.redirect_uri == "https://my.app/callback"
|
|
|
|
def test_fallback_to_settings(self, monkeypatch):
|
|
from application.core.settings import settings
|
|
|
|
monkeypatch.setattr(settings, "API_URL", "https://api.docsgpt.co")
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "none",
|
|
})
|
|
assert "/api/mcp_server/callback" in tool.redirect_uri
|
|
|
|
|
|
# =====================================================================
|
|
# Cache Key Generation
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestGenerateCacheKey:
|
|
|
|
def test_none_auth(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
assert "none" in tool._cache_key
|
|
assert "mcp.example.com" in tool._cache_key
|
|
|
|
def test_bearer_auth(self, bearer_config):
|
|
tool = _make_tool(bearer_config)
|
|
assert "bearer:" in tool._cache_key
|
|
|
|
def test_api_key_auth(self):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "api_key",
|
|
"auth_credentials": {"api_key": "sk-test12345678"},
|
|
})
|
|
assert "apikey:" in tool._cache_key
|
|
|
|
def test_basic_auth(self):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "basic",
|
|
"auth_credentials": {"username": "user1", "password": "pass"},
|
|
})
|
|
assert "basic:user1" in tool._cache_key
|
|
|
|
def test_oauth_auth_includes_scopes(self):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "oauth",
|
|
"oauth_scopes": ["read", "write"],
|
|
})
|
|
assert "oauth:" in tool._cache_key
|
|
assert "read,write" in tool._cache_key
|
|
|
|
def test_bearer_empty_token(self):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "bearer",
|
|
"auth_credentials": {},
|
|
})
|
|
assert "bearer:none" in tool._cache_key
|
|
|
|
def test_api_key_empty(self):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "api_key",
|
|
"auth_credentials": {},
|
|
})
|
|
assert "apikey:none" in tool._cache_key
|
|
|
|
|
|
# =====================================================================
|
|
# Transport Creation
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestCreateTransport:
|
|
|
|
def test_http_transport(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
transport = tool._create_transport()
|
|
assert transport is not None
|
|
|
|
def test_sse_transport(self):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com/sse",
|
|
"transport_type": "sse",
|
|
"auth_type": "none",
|
|
})
|
|
transport = tool._create_transport()
|
|
assert transport is not None
|
|
|
|
def test_auto_detects_sse(self):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com/sse",
|
|
"transport_type": "auto",
|
|
"auth_type": "none",
|
|
})
|
|
transport = tool._create_transport()
|
|
assert transport is not None
|
|
|
|
def test_auto_defaults_to_http(self):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com/api",
|
|
"transport_type": "auto",
|
|
"auth_type": "none",
|
|
})
|
|
transport = tool._create_transport()
|
|
assert transport is not None
|
|
|
|
def test_stdio_transport_disabled(self):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"transport_type": "stdio",
|
|
"auth_type": "none",
|
|
})
|
|
with pytest.raises(ValueError, match="STDIO transport is disabled"):
|
|
tool._create_transport()
|
|
|
|
def test_api_key_header_injected(self):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"transport_type": "http",
|
|
"auth_type": "api_key",
|
|
"auth_credentials": {
|
|
"api_key": "sk-test",
|
|
"api_key_header": "X-Custom-Key",
|
|
},
|
|
})
|
|
transport = tool._create_transport()
|
|
assert transport is not None
|
|
|
|
def test_basic_auth_header_injected(self):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"transport_type": "http",
|
|
"auth_type": "basic",
|
|
"auth_credentials": {"username": "user", "password": "pass"},
|
|
})
|
|
transport = tool._create_transport()
|
|
assert transport is not None
|
|
|
|
def test_unknown_transport_type_defaults_to_http(self):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"transport_type": "grpc",
|
|
"auth_type": "none",
|
|
})
|
|
transport = tool._create_transport()
|
|
assert transport is not None
|
|
|
|
|
|
# =====================================================================
|
|
# Format Tools
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestFormatTools:
|
|
|
|
def test_format_list_of_dicts(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
result = tool._format_tools([{"name": "tool1", "description": "desc"}])
|
|
assert len(result) == 1
|
|
assert result[0]["name"] == "tool1"
|
|
|
|
def test_format_tools_with_name_attribute(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
mock_tool = MagicMock()
|
|
mock_tool.name = "my_tool"
|
|
mock_tool.description = "A tool"
|
|
mock_tool.inputSchema = {"type": "object", "properties": {}}
|
|
|
|
result = tool._format_tools([mock_tool])
|
|
assert len(result) == 1
|
|
assert result[0]["name"] == "my_tool"
|
|
assert result[0]["inputSchema"] == {"type": "object", "properties": {}}
|
|
|
|
def test_format_tools_with_model_dump(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
mock_tool = MagicMock(spec=[])
|
|
mock_tool.model_dump = MagicMock(
|
|
return_value={"name": "dumped", "description": "from dump"}
|
|
)
|
|
# Ensure no "name" attribute
|
|
result = tool._format_tools([mock_tool])
|
|
assert len(result) == 1
|
|
assert result[0]["name"] == "dumped"
|
|
|
|
def test_format_tools_fallback_str(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
|
|
class BareTool:
|
|
def __str__(self):
|
|
return "bare_tool"
|
|
|
|
result = tool._format_tools([BareTool()])
|
|
assert len(result) == 1
|
|
assert result[0]["name"] == "bare_tool"
|
|
assert result[0]["description"] == ""
|
|
|
|
def test_format_tools_response_object(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
resp = MagicMock()
|
|
resp.tools = [{"name": "t1", "description": "d1"}]
|
|
|
|
result = tool._format_tools(resp)
|
|
assert len(result) == 1
|
|
|
|
def test_format_tools_without_input_schema(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
mock_tool = MagicMock()
|
|
mock_tool.name = "simple"
|
|
mock_tool.description = "no schema"
|
|
del mock_tool.inputSchema
|
|
|
|
result = tool._format_tools([mock_tool])
|
|
assert "inputSchema" not in result[0]
|
|
|
|
def test_format_empty(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
assert tool._format_tools([]) == []
|
|
assert tool._format_tools("unexpected") == []
|
|
|
|
|
|
# =====================================================================
|
|
# Format Result
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestFormatResult:
|
|
|
|
def test_format_result_with_text_content(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
mock_result = MagicMock()
|
|
text_item = MagicMock()
|
|
text_item.text = "Hello"
|
|
del text_item.data
|
|
mock_result.content = [text_item]
|
|
mock_result.isError = False
|
|
|
|
result = tool._format_result(mock_result)
|
|
assert result["content"][0]["type"] == "text"
|
|
assert result["content"][0]["text"] == "Hello"
|
|
assert result["isError"] is False
|
|
|
|
def test_format_result_with_data_content(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
mock_result = MagicMock()
|
|
data_item = MagicMock()
|
|
del data_item.text
|
|
data_item.data = {"key": "value"}
|
|
mock_result.content = [data_item]
|
|
mock_result.isError = False
|
|
|
|
result = tool._format_result(mock_result)
|
|
assert result["content"][0]["type"] == "data"
|
|
assert result["content"][0]["data"] == {"key": "value"}
|
|
|
|
def test_format_result_unknown_content_type(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
mock_result = MagicMock()
|
|
unknown_item = MagicMock()
|
|
del unknown_item.text
|
|
del unknown_item.data
|
|
mock_result.content = [unknown_item]
|
|
mock_result.isError = False
|
|
|
|
result = tool._format_result(mock_result)
|
|
assert result["content"][0]["type"] == "unknown"
|
|
|
|
def test_format_raw_result(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
raw = {"key": "value"}
|
|
assert tool._format_result(raw) == raw
|
|
|
|
|
|
# =====================================================================
|
|
# Execute Action
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestExecuteAction:
|
|
|
|
def test_no_server_raises(self):
|
|
tool = _make_tool({"server_url": "", "auth_type": "none"})
|
|
with pytest.raises(Exception, match="No MCP server configured"):
|
|
tool.execute_action("test_action")
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_successful_execute(self, mock_run, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = MagicMock()
|
|
mock_run.return_value = {"key": "value"}
|
|
|
|
result = tool.execute_action("test_action", param1="val1")
|
|
|
|
mock_run.assert_called_once_with("call_tool", "test_action", param1="val1")
|
|
assert result == {"key": "value"}
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_empty_kwargs_cleaned(self, mock_run, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = MagicMock()
|
|
mock_run.return_value = {}
|
|
|
|
tool.execute_action("test", param1="", param2=None, param3="real")
|
|
|
|
call_kwargs = mock_run.call_args[1]
|
|
assert "param1" not in call_kwargs
|
|
assert "param2" not in call_kwargs
|
|
assert call_kwargs["param3"] == "real"
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_auth_error_retries_for_non_oauth(self, mock_run, mcp_config):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = MagicMock()
|
|
|
|
mock_run.side_effect = [
|
|
Exception("401 Unauthorized"),
|
|
{"key": "retry_ok"},
|
|
]
|
|
|
|
with patch.object(MCPTool, "_setup_client"):
|
|
result = tool.execute_action("act")
|
|
assert result == {"key": "retry_ok"}
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_auth_error_raises_for_oauth(self, mock_run):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "oauth",
|
|
})
|
|
tool._client = MagicMock()
|
|
mock_run.side_effect = Exception("401 Unauthorized")
|
|
|
|
with pytest.raises(Exception, match="OAuth session expired"):
|
|
tool.execute_action("act")
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_non_auth_error_raises(self, mock_run, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = MagicMock()
|
|
mock_run.side_effect = Exception("Something weird happened")
|
|
|
|
with pytest.raises(Exception, match="Failed to execute action"):
|
|
tool.execute_action("act")
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_no_client_calls_setup(self, mock_run, mcp_config):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = None
|
|
mock_run.return_value = {"ok": True}
|
|
|
|
with patch.object(MCPTool, "_setup_client") as mock_setup:
|
|
tool.execute_action("act")
|
|
mock_setup.assert_called()
|
|
|
|
|
|
# =====================================================================
|
|
# Discover Tools
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestDiscoverTools:
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_discover_tools_success(self, mock_run, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = MagicMock()
|
|
mock_run.return_value = [{"name": "t1", "description": "d1"}]
|
|
|
|
result = tool.discover_tools()
|
|
assert len(result) == 1
|
|
assert result[0]["name"] == "t1"
|
|
|
|
def test_discover_tools_no_server_url(self):
|
|
tool = _make_tool({"server_url": "", "auth_type": "none"})
|
|
assert tool.discover_tools() == []
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_discover_tools_error(self, mock_run, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = MagicMock()
|
|
mock_run.side_effect = Exception("connection lost")
|
|
|
|
with pytest.raises(Exception, match="Failed to discover tools"):
|
|
tool.discover_tools()
|
|
|
|
|
|
# =====================================================================
|
|
# Test Connection
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestTestConnection:
|
|
|
|
def test_no_server_url(self):
|
|
tool = _make_tool({"server_url": "", "auth_type": "none"})
|
|
result = tool.test_connection()
|
|
assert result["success"] is False
|
|
assert "No server URL" in result["message"]
|
|
|
|
def test_invalid_scheme(self):
|
|
tool = _make_tool({"server_url": "ftp://bad.com", "auth_type": "none"})
|
|
result = tool.test_connection()
|
|
assert result["success"] is False
|
|
assert "Invalid URL scheme" in result["message"]
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_regular_connection_success(self, mock_run, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = MagicMock()
|
|
mock_run.side_effect = [
|
|
None, # ping
|
|
[{"name": "t1", "description": "d1"}], # list_tools
|
|
]
|
|
|
|
result = tool.test_connection()
|
|
assert result["success"] is True
|
|
assert result["tools_count"] == 1
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_regular_connection_ping_fails_tools_work(self, mock_run, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = MagicMock()
|
|
mock_run.side_effect = [
|
|
Exception("ping failed"), # ping
|
|
[{"name": "t1", "description": "d1"}], # list_tools
|
|
]
|
|
|
|
result = tool.test_connection()
|
|
assert result["success"] is True
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_regular_connection_both_fail(self, mock_run, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = MagicMock()
|
|
mock_run.side_effect = [
|
|
Exception("ping failed"),
|
|
Exception("tools failed"),
|
|
]
|
|
|
|
result = tool.test_connection()
|
|
assert result["success"] is False
|
|
|
|
def test_client_init_failure(self):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
tool = _make_tool({
|
|
"server_url": "https://good.example.com",
|
|
"auth_type": "none",
|
|
})
|
|
tool._client = None
|
|
|
|
with patch.object(MCPTool, "_setup_client", side_effect=Exception("init fail")):
|
|
result = tool.test_connection()
|
|
assert result["success"] is False
|
|
assert "Client init failed" in result["message"]
|
|
|
|
|
|
# =====================================================================
|
|
# Map Error
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestMapError:
|
|
|
|
def test_timeout_error(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
err = tool._map_error("test", concurrent.futures.TimeoutError())
|
|
assert "Timed out" in str(err)
|
|
|
|
def test_connection_refused(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
err = tool._map_error("test", ConnectionRefusedError())
|
|
assert "Connection refused" in str(err)
|
|
|
|
def test_403_forbidden(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
err = tool._map_error("test", Exception("403 Forbidden"))
|
|
assert "Access denied" in str(err)
|
|
|
|
def test_401_unauthorized(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
err = tool._map_error("test", Exception("401 Unauthorized"))
|
|
assert "Authentication failed" in str(err)
|
|
|
|
def test_econnrefused_pattern(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
err = tool._map_error("test", Exception("ECONNREFUSED error"))
|
|
assert "Connection refused" in str(err)
|
|
|
|
def test_ssl_error(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
err = tool._map_error("test", Exception("SSL certificate verify failed"))
|
|
assert "SSL" in str(err)
|
|
|
|
def test_unknown_error_passthrough(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
original = RuntimeError("something weird")
|
|
err = tool._map_error("test", original)
|
|
assert err is original
|
|
|
|
|
|
# =====================================================================
|
|
# Get Actions Metadata
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestGetActionsMetadata:
|
|
|
|
def test_empty_tools(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool.available_tools = []
|
|
assert tool.get_actions_metadata() == []
|
|
|
|
def test_tools_with_input_schema(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool.available_tools = [
|
|
{
|
|
"name": "search",
|
|
"description": "Search things",
|
|
"inputSchema": {
|
|
"type": "object",
|
|
"properties": {"query": {"type": "string"}},
|
|
"required": ["query"],
|
|
"additionalProperties": False,
|
|
"description": "Search params",
|
|
},
|
|
}
|
|
]
|
|
meta = tool.get_actions_metadata()
|
|
assert len(meta) == 1
|
|
assert meta[0]["name"] == "search"
|
|
assert "query" in meta[0]["parameters"]["properties"]
|
|
assert meta[0]["parameters"]["additionalProperties"] is False
|
|
|
|
def test_tools_without_schema(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool.available_tools = [{"name": "ping", "description": "Ping"}]
|
|
meta = tool.get_actions_metadata()
|
|
assert meta[0]["parameters"]["properties"] == {}
|
|
|
|
def test_tools_with_flat_schema(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool.available_tools = [
|
|
{
|
|
"name": "flat",
|
|
"description": "Flat schema",
|
|
"inputSchema": {"query": {"type": "string"}},
|
|
}
|
|
]
|
|
meta = tool.get_actions_metadata()
|
|
assert "query" in meta[0]["parameters"]["properties"]
|
|
|
|
def test_tools_with_alternate_schema_keys(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool.available_tools = [
|
|
{
|
|
"name": "alt",
|
|
"description": "Alt",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"x": {"type": "number"}},
|
|
},
|
|
}
|
|
]
|
|
meta = tool.get_actions_metadata()
|
|
assert "x" in meta[0]["parameters"]["properties"]
|
|
|
|
def test_config_requirements(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
reqs = tool.get_config_requirements()
|
|
assert "server_url" in reqs
|
|
assert "auth_type" in reqs
|
|
assert reqs["server_url"]["required"] is True
|
|
assert "timeout" in reqs
|
|
|
|
|
|
# =====================================================================
|
|
# Setup Client (caching)
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSetupClient:
|
|
|
|
def test_setup_client_caches_client(self, mcp_config):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
tool = MCPTool.__new__(MCPTool)
|
|
tool.config = mcp_config
|
|
tool.server_url = mcp_config["server_url"]
|
|
tool.transport_type = "http"
|
|
tool.auth_type = "none"
|
|
tool.timeout = 10
|
|
tool.custom_headers = {}
|
|
tool.auth_credentials = {}
|
|
tool.oauth_scopes = []
|
|
tool.oauth_task_id = None
|
|
tool.oauth_client_name = "DocsGPT-MCP"
|
|
tool.redirect_uri = "https://example.com/callback"
|
|
tool.query_mode = False
|
|
tool._cache_key = "test_cache_key"
|
|
tool._client = None
|
|
tool.available_tools = []
|
|
|
|
mock_client = MagicMock()
|
|
with patch.object(MCPTool, "_create_transport", return_value=MagicMock()), \
|
|
patch("application.agents.tools.mcp_tool.Client", return_value=mock_client):
|
|
tool._setup_client()
|
|
assert tool._client is mock_client
|
|
|
|
|
|
# =====================================================================
|
|
# MCPOAuthManager
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestMCPOAuthManager:
|
|
|
|
def test_handle_callback_success(self):
|
|
from application.agents.tools.mcp_tool import MCPOAuthManager
|
|
|
|
mock_redis = MagicMock()
|
|
manager = MCPOAuthManager(mock_redis)
|
|
|
|
result = manager.handle_oauth_callback(state="abc123", code="auth_code")
|
|
assert result is True
|
|
mock_redis.setex.assert_called()
|
|
|
|
def test_handle_callback_no_redis(self):
|
|
from application.agents.tools.mcp_tool import MCPOAuthManager
|
|
|
|
manager = MCPOAuthManager(None)
|
|
result = manager.handle_oauth_callback(state="abc", code="code")
|
|
assert result is False
|
|
|
|
def test_handle_callback_no_state(self):
|
|
from application.agents.tools.mcp_tool import MCPOAuthManager
|
|
|
|
mock_redis = MagicMock()
|
|
manager = MCPOAuthManager(mock_redis)
|
|
result = manager.handle_oauth_callback(state="", code="code")
|
|
assert result is False
|
|
|
|
def test_handle_callback_error(self):
|
|
from application.agents.tools.mcp_tool import MCPOAuthManager
|
|
|
|
mock_redis = MagicMock()
|
|
manager = MCPOAuthManager(mock_redis)
|
|
|
|
result = manager.handle_oauth_callback(
|
|
state="abc", code="", error="access_denied"
|
|
)
|
|
assert result is False
|
|
|
|
def test_get_oauth_status_no_task(self):
|
|
from application.agents.tools.mcp_tool import MCPOAuthManager
|
|
|
|
manager = MCPOAuthManager(MagicMock())
|
|
result = manager.get_oauth_status("", "alice")
|
|
assert result["status"] == "not_started"
|
|
|
|
def test_get_oauth_status_no_user(self):
|
|
from application.agents.tools.mcp_tool import MCPOAuthManager
|
|
|
|
manager = MCPOAuthManager(MagicMock())
|
|
result = manager.get_oauth_status("task123", "")
|
|
# Without a user id we can't address the per-user SSE journal,
|
|
# so the manager refuses rather than scanning every user's
|
|
# stream.
|
|
assert result["status"] == "not_found"
|
|
|
|
def test_get_oauth_status_reads_completed_envelope_from_journal(self):
|
|
"""The manager walks the user's SSE Streams journal and
|
|
surfaces the latest ``mcp.oauth.*`` envelope for the task.
|
|
|
|
Verifies the full polling-contract surface: ``status`` is
|
|
derived from the event-type suffix, and the completed
|
|
payload's ``tools`` / ``tools_count`` fields are passed
|
|
through unchanged so ``mcp.py``'s ``connect_mcp`` can use them
|
|
without further plumbing.
|
|
"""
|
|
from application.agents.tools.mcp_tool import MCPOAuthManager
|
|
|
|
completed_envelope = json.dumps(
|
|
{
|
|
"type": "mcp.oauth.completed",
|
|
"scope": {"kind": "mcp_oauth", "id": "task123"},
|
|
"payload": {
|
|
"task_id": "task123",
|
|
"tools": [{"name": "t1", "description": "d"}],
|
|
"tools_count": 1,
|
|
},
|
|
}
|
|
).encode("utf-8")
|
|
# xrevrange yields newest-first; an unrelated older entry
|
|
# should be skipped on the way to the matching one.
|
|
unrelated_envelope = json.dumps(
|
|
{
|
|
"type": "source.ingest.completed",
|
|
"scope": {"kind": "source", "id": "abc"},
|
|
"payload": {},
|
|
}
|
|
).encode("utf-8")
|
|
mock_redis = MagicMock()
|
|
mock_redis.xrevrange.return_value = [
|
|
(b"1735682400000-0", {b"event": completed_envelope}),
|
|
(b"1735682300000-0", {b"event": unrelated_envelope}),
|
|
]
|
|
|
|
manager = MCPOAuthManager(mock_redis)
|
|
result = manager.get_oauth_status("task123", "alice")
|
|
|
|
assert result["status"] == "completed"
|
|
assert result["tools"] == [{"name": "t1", "description": "d"}]
|
|
assert result["tools_count"] == 1
|
|
# The stream key is scoped per-user — never global.
|
|
mock_redis.xrevrange.assert_called_once()
|
|
call_args = mock_redis.xrevrange.call_args
|
|
assert "user:alice:stream" in call_args.args
|
|
# Scan window must cover the full bounded stream so a flood of
|
|
# concurrent source-ingest events between popup-completed and
|
|
# Save can't push the OAuth envelope out of view.
|
|
from application.core.settings import settings
|
|
|
|
assert call_args.kwargs.get("count") >= settings.EVENTS_STREAM_MAXLEN
|
|
|
|
def test_get_oauth_status_scan_covers_events_stream_maxlen(self):
|
|
"""Regression: count must scale with ``EVENTS_STREAM_MAXLEN``
|
|
so the OAuth completion envelope is reachable even after
|
|
concurrent source-ingest events flood the user stream."""
|
|
from application.agents.tools.mcp_tool import MCPOAuthManager
|
|
from application.core.settings import settings
|
|
|
|
mock_redis = MagicMock()
|
|
mock_redis.xrevrange.return_value = []
|
|
|
|
manager = MCPOAuthManager(mock_redis)
|
|
manager.get_oauth_status("task123", "alice")
|
|
|
|
call_args = mock_redis.xrevrange.call_args
|
|
assert call_args.kwargs.get("count") >= settings.EVENTS_STREAM_MAXLEN
|
|
|
|
def test_get_oauth_status_returns_not_found_when_no_match(self):
|
|
from application.agents.tools.mcp_tool import MCPOAuthManager
|
|
|
|
mock_redis = MagicMock()
|
|
# Stream is non-empty but has nothing matching this task.
|
|
unrelated_envelope = json.dumps(
|
|
{
|
|
"type": "mcp.oauth.completed",
|
|
"scope": {"kind": "mcp_oauth", "id": "other-task"},
|
|
"payload": {"task_id": "other-task"},
|
|
}
|
|
).encode("utf-8")
|
|
mock_redis.xrevrange.return_value = [
|
|
(b"1735682400000-0", {b"event": unrelated_envelope}),
|
|
]
|
|
|
|
manager = MCPOAuthManager(mock_redis)
|
|
result = manager.get_oauth_status("task123", "alice")
|
|
|
|
assert result["status"] == "not_found"
|
|
|
|
|
|
# =====================================================================
|
|
# DBTokenStorage
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestDBTokenStorage:
|
|
"""Covers the repository-backed DBTokenStorage post-PG migration.
|
|
|
|
Round-trip tests use the ephemeral ``pg_conn`` fixture and patch
|
|
``db_session``/``db_readonly`` in ``application.agents.tools.mcp_tool``
|
|
so the real INSERT/SELECT SQL runs. This is the shape that caught the
|
|
``server_url`` NULL-column regression — ``get_tokens`` only succeeds
|
|
if ``set_tokens`` populated the scalar column.
|
|
"""
|
|
|
|
@staticmethod
|
|
def _patch_db(monkeypatch, pg_conn):
|
|
from contextlib import contextmanager
|
|
|
|
import application.agents.tools.mcp_tool as mcp_mod
|
|
|
|
@contextmanager
|
|
def _yield():
|
|
yield pg_conn
|
|
|
|
monkeypatch.setattr(mcp_mod, "db_session", _yield, raising=False)
|
|
monkeypatch.setattr(mcp_mod, "db_readonly", _yield, raising=False)
|
|
# mcp_tool imports db_session/db_readonly *inside* the helper
|
|
# methods, so also patch the origin module they come from.
|
|
import application.storage.db.session as session_mod
|
|
|
|
monkeypatch.setattr(session_mod, "db_session", _yield)
|
|
monkeypatch.setattr(session_mod, "db_readonly", _yield)
|
|
|
|
def test_get_base_url(self):
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
|
|
assert (
|
|
DBTokenStorage.get_base_url("https://mcp.example.com/api/v1")
|
|
== "https://mcp.example.com"
|
|
)
|
|
|
|
def test_get_base_url_with_port(self):
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
|
|
assert (
|
|
DBTokenStorage.get_base_url("http://localhost:8080/path")
|
|
== "http://localhost:8080"
|
|
)
|
|
|
|
def test_pg_provider_includes_base_url(self):
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
|
|
storage = DBTokenStorage(
|
|
server_url="https://mcp.example.com/api",
|
|
user_id="user1",
|
|
)
|
|
assert storage._pg_provider() == "mcp:https://mcp.example.com"
|
|
|
|
def test_serialize_client_info(self):
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
|
|
storage = DBTokenStorage(
|
|
server_url="https://mcp.example.com",
|
|
user_id="user1",
|
|
)
|
|
info = {"redirect_uris": ["https://example.com/cb"]}
|
|
result = storage._serialize_client_info(info)
|
|
assert result["redirect_uris"] == ["https://example.com/cb"]
|
|
|
|
def test_get_tokens_none_when_no_row(self, monkeypatch, pg_conn):
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
|
|
self._patch_db(monkeypatch, pg_conn)
|
|
storage = DBTokenStorage(
|
|
server_url="https://mcp.example.com/api",
|
|
user_id="user-empty",
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
result = loop.run_until_complete(storage.get_tokens())
|
|
assert result is None
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_set_then_get_tokens_roundtrip(self, monkeypatch, pg_conn):
|
|
"""Regression: ``set_tokens`` must populate the scalar
|
|
``server_url`` column so ``get_tokens`` (which goes through
|
|
``get_by_user_and_server_url``) can resolve the row."""
|
|
from mcp.shared.auth import OAuthToken
|
|
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
|
|
self._patch_db(monkeypatch, pg_conn)
|
|
storage = DBTokenStorage(
|
|
server_url="https://mcp.example.com/api/v1",
|
|
user_id="user-roundtrip",
|
|
)
|
|
|
|
tokens_in = OAuthToken(
|
|
access_token="at-abc",
|
|
token_type="Bearer",
|
|
expires_in=3600,
|
|
refresh_token="rt-xyz",
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
loop.run_until_complete(storage.set_tokens(tokens_in))
|
|
tokens_out = loop.run_until_complete(storage.get_tokens())
|
|
finally:
|
|
loop.close()
|
|
|
|
assert tokens_out is not None
|
|
assert tokens_out.access_token == "at-abc"
|
|
assert tokens_out.refresh_token == "rt-xyz"
|
|
|
|
def test_set_tokens_populates_scalar_server_url(
|
|
self, monkeypatch, pg_conn,
|
|
):
|
|
"""Direct check that the fix writes ``server_url`` to the scalar
|
|
column, not only into the JSONB blob."""
|
|
from mcp.shared.auth import OAuthToken
|
|
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
from application.storage.db.repositories.connector_sessions import (
|
|
ConnectorSessionsRepository,
|
|
)
|
|
|
|
self._patch_db(monkeypatch, pg_conn)
|
|
base_url = "https://mcp.scalar.example.com"
|
|
storage = DBTokenStorage(
|
|
server_url=f"{base_url}/api",
|
|
user_id="user-scalar",
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
loop.run_until_complete(
|
|
storage.set_tokens(
|
|
OAuthToken(access_token="at", token_type="Bearer"),
|
|
),
|
|
)
|
|
finally:
|
|
loop.close()
|
|
|
|
row = ConnectorSessionsRepository(pg_conn).get_by_user_and_server_url(
|
|
"user-scalar", base_url,
|
|
)
|
|
assert row is not None
|
|
assert row["server_url"] == base_url
|
|
# ``server_url`` must NOT be duplicated inside the JSONB blob.
|
|
session_data = row["session_data"] or {}
|
|
assert "server_url" not in session_data
|
|
assert session_data.get("tokens", {}).get("access_token") == "at"
|
|
|
|
def test_clear_removes_row(self, monkeypatch, pg_conn):
|
|
from mcp.shared.auth import OAuthToken
|
|
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
from application.storage.db.repositories.connector_sessions import (
|
|
ConnectorSessionsRepository,
|
|
)
|
|
|
|
self._patch_db(monkeypatch, pg_conn)
|
|
base_url = "https://mcp.clear.example.com"
|
|
storage = DBTokenStorage(
|
|
server_url=f"{base_url}/api",
|
|
user_id="user-clear",
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
loop.run_until_complete(
|
|
storage.set_tokens(
|
|
OAuthToken(access_token="at", token_type="Bearer"),
|
|
),
|
|
)
|
|
loop.run_until_complete(storage.clear())
|
|
finally:
|
|
loop.close()
|
|
|
|
assert (
|
|
ConnectorSessionsRepository(pg_conn).get_by_user_and_server_url(
|
|
"user-clear", base_url,
|
|
)
|
|
is None
|
|
)
|
|
|
|
|
|
# =====================================================================
|
|
# NonInteractiveOAuth
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.skip(reason="OAuth class signatures changed post-PG migration (db kwarg removed); needs rewrite")
|
|
@pytest.mark.unit
|
|
class TestNonInteractiveOAuth:
|
|
|
|
def test_redirect_handler_raises(self):
|
|
from application.agents.tools.mcp_tool import NonInteractiveOAuth
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_collection.find_one.return_value = None
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
oauth = NonInteractiveOAuth(
|
|
mcp_url="https://mcp.example.com/api",
|
|
scopes=["read"],
|
|
redis_client=None,
|
|
redirect_uri="https://example.com/callback",
|
|
db=mock_db,
|
|
user_id="user1",
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
with pytest.raises(Exception, match="OAuth session expired"):
|
|
loop.run_until_complete(
|
|
oauth.redirect_handler("https://auth.example.com/authorize?state=x")
|
|
)
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_callback_handler_raises(self):
|
|
from application.agents.tools.mcp_tool import NonInteractiveOAuth
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_collection.find_one.return_value = None
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
oauth = NonInteractiveOAuth(
|
|
mcp_url="https://mcp.example.com/api",
|
|
scopes=["read"],
|
|
redis_client=None,
|
|
redirect_uri="https://example.com/callback",
|
|
db=mock_db,
|
|
user_id="user1",
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
with pytest.raises(Exception, match="OAuth session expired"):
|
|
loop.run_until_complete(oauth.callback_handler())
|
|
finally:
|
|
loop.close()
|
|
|
|
|
|
# =====================================================================
|
|
# Run Async Operation
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestRunAsyncOperation:
|
|
|
|
def test_run_in_new_loop(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = MagicMock()
|
|
|
|
async def fake_execute(*args, **kwargs):
|
|
return "ok"
|
|
|
|
with patch.object(tool, "_execute_with_client", side_effect=fake_execute):
|
|
result = tool._run_in_new_loop("ping")
|
|
assert result == "ok"
|
|
|
|
|
|
# =====================================================================
|
|
# Resolve Redirect URI (additional coverage)
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestResolveRedirectUriExtended:
|
|
|
|
def test_mcp_oauth_redirect_uri_setting(self, monkeypatch):
|
|
from application.core.settings import settings
|
|
|
|
monkeypatch.setattr(settings, "MCP_OAUTH_REDIRECT_URI", "https://custom.redirect/callback/")
|
|
# Ensure no configured redirect_uri in config
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "none",
|
|
})
|
|
assert tool.redirect_uri == "https://custom.redirect/callback"
|
|
|
|
def test_connector_redirect_base_uri_setting(self, monkeypatch):
|
|
from application.core.settings import settings
|
|
|
|
monkeypatch.setattr(settings, "MCP_OAUTH_REDIRECT_URI", None, raising=False)
|
|
monkeypatch.setattr(
|
|
settings, "CONNECTOR_REDIRECT_BASE_URI",
|
|
"https://connector.example.com/some/path",
|
|
)
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "none",
|
|
})
|
|
assert tool.redirect_uri == "https://connector.example.com/api/mcp_server/callback"
|
|
|
|
def test_connector_redirect_base_uri_invalid_url(self, monkeypatch):
|
|
from application.core.settings import settings
|
|
|
|
monkeypatch.setattr(settings, "MCP_OAUTH_REDIRECT_URI", None, raising=False)
|
|
# Provide a base URI that has no scheme
|
|
monkeypatch.setattr(
|
|
settings, "CONNECTOR_REDIRECT_BASE_URI", "no-scheme-url",
|
|
)
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "none",
|
|
})
|
|
# Falls through to API_URL fallback
|
|
assert "/api/mcp_server/callback" in tool.redirect_uri
|
|
|
|
|
|
# =====================================================================
|
|
# _setup_client additional coverage (cache expiry, OAuth branches)
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSetupClientExtended:
|
|
|
|
def test_cache_hit_returns_cached_client(self):
|
|
import application.agents.tools.mcp_tool as mcp_mod
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
tool = MCPTool.__new__(MCPTool)
|
|
tool.config = {"server_url": "https://mcp.example.com", "auth_type": "none"}
|
|
tool.server_url = "https://mcp.example.com"
|
|
tool.transport_type = "http"
|
|
tool.auth_type = "none"
|
|
tool.timeout = 10
|
|
tool.custom_headers = {}
|
|
tool.auth_credentials = {}
|
|
tool.oauth_scopes = []
|
|
tool.oauth_task_id = None
|
|
tool.oauth_client_name = "DocsGPT-MCP"
|
|
tool.redirect_uri = "https://example.com/callback"
|
|
tool.query_mode = False
|
|
tool._cache_key = "cache_hit_test_key"
|
|
tool._client = None
|
|
tool.available_tools = []
|
|
|
|
cached_client = MagicMock()
|
|
mcp_mod._mcp_clients_cache["cache_hit_test_key"] = {
|
|
"client": cached_client,
|
|
"created_at": __import__("time").time(),
|
|
}
|
|
|
|
tool._setup_client()
|
|
assert tool._client is cached_client
|
|
|
|
def test_expired_cache_creates_new_client(self):
|
|
import application.agents.tools.mcp_tool as mcp_mod
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
tool = MCPTool.__new__(MCPTool)
|
|
tool.config = {"server_url": "https://mcp.example.com", "auth_type": "none"}
|
|
tool.server_url = "https://mcp.example.com"
|
|
tool.transport_type = "http"
|
|
tool.auth_type = "none"
|
|
tool.timeout = 10
|
|
tool.custom_headers = {}
|
|
tool.auth_credentials = {}
|
|
tool.oauth_scopes = []
|
|
tool.oauth_task_id = None
|
|
tool.oauth_client_name = "DocsGPT-MCP"
|
|
tool.redirect_uri = "https://example.com/callback"
|
|
tool.query_mode = False
|
|
tool._cache_key = "expired_cache_key"
|
|
tool._client = None
|
|
tool.available_tools = []
|
|
|
|
old_client = MagicMock()
|
|
mcp_mod._mcp_clients_cache["expired_cache_key"] = {
|
|
"client": old_client,
|
|
"created_at": __import__("time").time() - 600,
|
|
}
|
|
|
|
new_client = MagicMock()
|
|
with patch.object(MCPTool, "_create_transport", return_value=MagicMock()), \
|
|
patch("application.agents.tools.mcp_tool.Client", return_value=new_client):
|
|
tool._setup_client()
|
|
assert tool._client is new_client
|
|
assert "expired_cache_key" not in mcp_mod._mcp_clients_cache or \
|
|
mcp_mod._mcp_clients_cache["expired_cache_key"]["client"] is new_client
|
|
|
|
def test_setup_client_oauth_query_mode(self):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
tool = MCPTool.__new__(MCPTool)
|
|
tool.config = {"server_url": "https://mcp.example.com", "auth_type": "oauth"}
|
|
tool.server_url = "https://mcp.example.com"
|
|
tool.transport_type = "http"
|
|
tool.auth_type = "oauth"
|
|
tool.timeout = 10
|
|
tool.custom_headers = {}
|
|
tool.auth_credentials = {}
|
|
tool.oauth_scopes = ["read"]
|
|
tool.oauth_task_id = None
|
|
tool.oauth_client_name = "DocsGPT-MCP"
|
|
tool.redirect_uri = "https://example.com/callback"
|
|
tool.query_mode = True
|
|
tool._cache_key = "oauth_qm_key"
|
|
tool._client = None
|
|
tool.available_tools = []
|
|
tool.user_id = "user1"
|
|
|
|
mock_client = MagicMock()
|
|
with patch.object(MCPTool, "_create_transport", return_value=MagicMock()), \
|
|
patch("application.agents.tools.mcp_tool.Client", return_value=mock_client), \
|
|
patch("application.agents.tools.mcp_tool.get_redis_instance", return_value=MagicMock()), \
|
|
patch("application.agents.tools.mcp_tool.NonInteractiveOAuth"):
|
|
tool._setup_client()
|
|
assert tool._client is mock_client
|
|
|
|
def test_setup_client_oauth_interactive_mode(self):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
tool = MCPTool.__new__(MCPTool)
|
|
tool.config = {"server_url": "https://mcp.example.com", "auth_type": "oauth"}
|
|
tool.server_url = "https://mcp.example.com"
|
|
tool.transport_type = "http"
|
|
tool.auth_type = "oauth"
|
|
tool.timeout = 10
|
|
tool.custom_headers = {}
|
|
tool.auth_credentials = {}
|
|
tool.oauth_scopes = ["read"]
|
|
tool.oauth_task_id = "task123"
|
|
tool.oauth_client_name = "DocsGPT-MCP"
|
|
tool.redirect_uri = "https://example.com/callback"
|
|
tool.query_mode = False
|
|
tool._cache_key = "oauth_interactive_key"
|
|
tool._client = None
|
|
tool.available_tools = []
|
|
tool.user_id = "user1"
|
|
tool.oauth_redirect_publish = None
|
|
|
|
mock_client = MagicMock()
|
|
with patch.object(MCPTool, "_create_transport", return_value=MagicMock()), \
|
|
patch("application.agents.tools.mcp_tool.Client", return_value=mock_client), \
|
|
patch("application.agents.tools.mcp_tool.get_redis_instance", return_value=MagicMock()), \
|
|
patch("application.agents.tools.mcp_tool.DocsGPTOAuth"):
|
|
tool._setup_client()
|
|
assert tool._client is mock_client
|
|
|
|
def test_setup_client_bearer_auth(self):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
tool = MCPTool.__new__(MCPTool)
|
|
tool.config = {"server_url": "https://mcp.example.com", "auth_type": "bearer"}
|
|
tool.server_url = "https://mcp.example.com"
|
|
tool.transport_type = "http"
|
|
tool.auth_type = "bearer"
|
|
tool.timeout = 10
|
|
tool.custom_headers = {}
|
|
tool.auth_credentials = {"bearer_token": "my_token"}
|
|
tool.oauth_scopes = []
|
|
tool.oauth_task_id = None
|
|
tool.oauth_client_name = "DocsGPT-MCP"
|
|
tool.redirect_uri = "https://example.com/callback"
|
|
tool.query_mode = False
|
|
tool._cache_key = "bearer_setup_key"
|
|
tool._client = None
|
|
tool.available_tools = []
|
|
|
|
mock_client = MagicMock()
|
|
with patch.object(MCPTool, "_create_transport", return_value=MagicMock()), \
|
|
patch("application.agents.tools.mcp_tool.Client", return_value=mock_client), \
|
|
patch("application.agents.tools.mcp_tool.BearerAuth") as mock_bearer_auth:
|
|
tool._setup_client()
|
|
mock_bearer_auth.assert_called_once_with("my_token")
|
|
assert tool._client is mock_client
|
|
|
|
|
|
# =====================================================================
|
|
# _execute_with_client async coverage
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestExecuteWithClient:
|
|
|
|
@staticmethod
|
|
def _make_async_client():
|
|
"""Create a mock client that supports async context manager."""
|
|
from unittest.mock import AsyncMock as AM
|
|
|
|
mock_client = MagicMock()
|
|
mock_client.__aenter__ = AM(return_value=mock_client)
|
|
mock_client.__aexit__ = AM(return_value=None)
|
|
return mock_client
|
|
|
|
def test_ping_operation(self, mcp_config):
|
|
from unittest.mock import AsyncMock as AM
|
|
|
|
tool = _make_tool(mcp_config)
|
|
mock_client = self._make_async_client()
|
|
mock_client.ping = AM(return_value="pong")
|
|
tool._client = mock_client
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
result = loop.run_until_complete(tool._execute_with_client("ping"))
|
|
assert result == "pong"
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_list_tools_operation(self, mcp_config):
|
|
from unittest.mock import AsyncMock as AM
|
|
|
|
tool = _make_tool(mcp_config)
|
|
mock_client = self._make_async_client()
|
|
mock_client.list_tools = AM(return_value=[{"name": "t1", "description": "d1"}])
|
|
tool._client = mock_client
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
result = loop.run_until_complete(tool._execute_with_client("list_tools"))
|
|
assert len(result) == 1
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_call_tool_operation(self, mcp_config):
|
|
from unittest.mock import AsyncMock as AM
|
|
|
|
tool = _make_tool(mcp_config)
|
|
mock_client = self._make_async_client()
|
|
mock_client.call_tool = AM(return_value={"result": "called my_action"})
|
|
tool._client = mock_client
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
result = loop.run_until_complete(
|
|
tool._execute_with_client("call_tool", "my_action", key="val")
|
|
)
|
|
assert result == {"result": "called my_action"}
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_list_resources_operation(self, mcp_config):
|
|
from unittest.mock import AsyncMock as AM
|
|
|
|
tool = _make_tool(mcp_config)
|
|
mock_client = self._make_async_client()
|
|
mock_client.list_resources = AM(return_value=["r1", "r2"])
|
|
tool._client = mock_client
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
result = loop.run_until_complete(tool._execute_with_client("list_resources"))
|
|
assert result == ["r1", "r2"]
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_list_prompts_operation(self, mcp_config):
|
|
from unittest.mock import AsyncMock as AM
|
|
|
|
tool = _make_tool(mcp_config)
|
|
mock_client = self._make_async_client()
|
|
mock_client.list_prompts = AM(return_value=["p1"])
|
|
tool._client = mock_client
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
result = loop.run_until_complete(tool._execute_with_client("list_prompts"))
|
|
assert result == ["p1"]
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_unknown_operation_raises(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
mock_client = self._make_async_client()
|
|
tool._client = mock_client
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
with pytest.raises(Exception, match="Unknown operation"):
|
|
loop.run_until_complete(tool._execute_with_client("bogus_op"))
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_no_client_raises(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = None
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
with pytest.raises(Exception, match="not initialized"):
|
|
loop.run_until_complete(tool._execute_with_client("ping"))
|
|
finally:
|
|
loop.close()
|
|
|
|
|
|
# =====================================================================
|
|
# _run_async_operation (error mapping path)
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestRunAsyncOperationExtended:
|
|
|
|
def test_error_is_mapped_and_raised(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = MagicMock()
|
|
|
|
with patch.object(tool, "_run_in_new_loop", side_effect=ConnectionRefusedError()):
|
|
with pytest.raises(Exception, match="Connection refused"):
|
|
tool._run_async_operation("ping")
|
|
|
|
def test_inside_running_loop_uses_thread_pool(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = MagicMock()
|
|
|
|
# Simulate being inside a running event loop
|
|
with patch("asyncio.get_running_loop", return_value=MagicMock()), \
|
|
patch("concurrent.futures.ThreadPoolExecutor") as mock_tp:
|
|
mock_future = MagicMock()
|
|
mock_future.result.return_value = "thread_result"
|
|
mock_executor = MagicMock()
|
|
mock_executor.__enter__ = MagicMock(return_value=mock_executor)
|
|
mock_executor.__exit__ = MagicMock(return_value=False)
|
|
mock_executor.submit.return_value = mock_future
|
|
mock_tp.return_value = mock_executor
|
|
|
|
result = tool._run_async_operation("ping")
|
|
assert result == "thread_result"
|
|
|
|
|
|
# =====================================================================
|
|
# test_connection additional coverage
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestTestConnectionExtended:
|
|
|
|
def test_url_parse_exception(self):
|
|
"""Test that an unparseable URL returns failure."""
|
|
tool = _make_tool({"server_url": "://bad", "auth_type": "none"})
|
|
result = tool.test_connection()
|
|
assert result["success"] is False
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_no_tools_and_no_ping_fails(self, mock_run, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = MagicMock()
|
|
# ping succeeds but discover_tools returns empty
|
|
mock_run.side_effect = [
|
|
None, # ping ok
|
|
]
|
|
with patch.object(tool, "discover_tools", return_value=[]):
|
|
result = tool.test_connection()
|
|
# ping_ok is True but tools is empty, should still succeed
|
|
assert result["success"] is True
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_ping_fails_no_tools_fails(self, mock_run, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = MagicMock()
|
|
mock_run.side_effect = [
|
|
Exception("ping failed"),
|
|
]
|
|
with patch.object(tool, "discover_tools", return_value=[]):
|
|
result = tool.test_connection()
|
|
assert result["success"] is False
|
|
assert "ping failed" in result["message"]
|
|
|
|
def test_oauth_connection_with_valid_tokens(self):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "oauth",
|
|
"oauth_scopes": ["read"],
|
|
})
|
|
tool.user_id = "user1"
|
|
tool._client = MagicMock()
|
|
|
|
mock_token = MagicMock()
|
|
mock_token.access_token = "valid_token"
|
|
|
|
with patch("application.agents.tools.mcp_tool.DBTokenStorage") as mock_storage_cls:
|
|
mock_storage = MagicMock()
|
|
|
|
async def fake_get_tokens():
|
|
return mock_token
|
|
|
|
mock_storage.get_tokens = fake_get_tokens
|
|
mock_storage_cls.return_value = mock_storage
|
|
|
|
with patch.object(tool, "discover_tools", return_value=[{"name": "t1", "description": "d1"}]), \
|
|
patch.object(MCPTool, "_setup_client"):
|
|
result = tool.test_connection()
|
|
assert result["success"] is True
|
|
assert result["tools_count"] == 1
|
|
|
|
def test_oauth_connection_with_expired_tokens_starts_task(self):
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "oauth",
|
|
"oauth_scopes": ["read"],
|
|
})
|
|
tool.user_id = "user1"
|
|
tool._client = MagicMock()
|
|
|
|
with patch("application.agents.tools.mcp_tool.DBTokenStorage") as mock_storage_cls:
|
|
mock_storage = MagicMock()
|
|
|
|
async def fake_get_tokens():
|
|
return None
|
|
|
|
mock_storage.get_tokens = fake_get_tokens
|
|
mock_storage_cls.return_value = mock_storage
|
|
|
|
mock_task_result = MagicMock()
|
|
mock_task_result.id = "task_abc"
|
|
with patch("application.agents.tools.mcp_tool.mcp_oauth_task") as mock_task:
|
|
mock_task.delay.return_value = mock_task_result
|
|
result = tool.test_connection()
|
|
assert result["success"] is False
|
|
assert result["requires_oauth"] is True
|
|
assert result["task_id"] == "task_abc"
|
|
|
|
def test_oauth_connection_token_validation_fails(self):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "oauth",
|
|
"oauth_scopes": ["read"],
|
|
})
|
|
tool.user_id = "user1"
|
|
tool._client = MagicMock()
|
|
|
|
mock_token = MagicMock()
|
|
mock_token.access_token = "expired_token"
|
|
|
|
with patch("application.agents.tools.mcp_tool.DBTokenStorage") as mock_storage_cls:
|
|
mock_storage = MagicMock()
|
|
|
|
async def fake_get_tokens():
|
|
return mock_token
|
|
|
|
mock_storage.get_tokens = fake_get_tokens
|
|
mock_storage_cls.return_value = mock_storage
|
|
|
|
mock_task_result = MagicMock()
|
|
mock_task_result.id = "task_retry"
|
|
with patch.object(tool, "discover_tools", side_effect=Exception("401 Unauthorized")), \
|
|
patch.object(MCPTool, "_setup_client"), \
|
|
patch("application.agents.tools.mcp_tool.mcp_oauth_task") as mock_task:
|
|
mock_task.delay.return_value = mock_task_result
|
|
result = tool.test_connection()
|
|
assert result["success"] is False
|
|
assert result["requires_oauth"] is True
|
|
|
|
|
|
# =====================================================================
|
|
# execute_action extended (format_result path)
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestExecuteActionExtended:
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_execute_formats_result(self, mock_run, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = MagicMock()
|
|
mock_result = MagicMock()
|
|
text_item = MagicMock()
|
|
text_item.text = "result text"
|
|
del text_item.data
|
|
mock_result.content = [text_item]
|
|
mock_result.isError = False
|
|
mock_run.return_value = mock_result
|
|
|
|
result = tool.execute_action("test_action", query="hello")
|
|
assert result["content"][0]["type"] == "text"
|
|
assert result["isError"] is False
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_execute_auth_retry_second_attempt_fails(self, mock_run):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
tool = _make_tool({
|
|
"server_url": "https://mcp.example.com",
|
|
"auth_type": "bearer",
|
|
"auth_credentials": {"bearer_token": "tok"},
|
|
})
|
|
tool._client = MagicMock()
|
|
|
|
mock_run.side_effect = Exception("401 Unauthorized")
|
|
|
|
with patch.object(MCPTool, "_setup_client"):
|
|
with pytest.raises(Exception, match="failed after re-auth attempt"):
|
|
tool.execute_action("act")
|
|
|
|
|
|
# =====================================================================
|
|
# discover_tools extended
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestDiscoverToolsExtended:
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_discover_tools_no_client_calls_setup(self, mock_run, mcp_config):
|
|
from application.agents.tools.mcp_tool import MCPTool
|
|
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = None
|
|
mock_run.return_value = [{"name": "t1", "description": "d1"}]
|
|
|
|
with patch.object(MCPTool, "_setup_client") as mock_setup:
|
|
result = tool.discover_tools()
|
|
mock_setup.assert_called_once()
|
|
assert len(result) == 1
|
|
|
|
|
|
# =====================================================================
|
|
# _test_regular_connection extended
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestRegularConnectionExtended:
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_regular_connection_message_format(self, mock_run, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = MagicMock()
|
|
mock_run.side_effect = [None] # ping ok
|
|
with patch.object(tool, "discover_tools", return_value=[
|
|
{"name": "single_tool", "description": "only one"},
|
|
]):
|
|
result = tool.test_connection()
|
|
assert result["success"] is True
|
|
assert "1 tool" in result["message"]
|
|
# Singular form for 1 tool
|
|
assert "tools" not in result["message"]
|
|
|
|
@patch("application.agents.tools.mcp_tool.MCPTool._run_async_operation")
|
|
def test_regular_connection_multiple_tools(self, mock_run, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._client = MagicMock()
|
|
mock_run.side_effect = [None] # ping ok
|
|
with patch.object(tool, "discover_tools", return_value=[
|
|
{"name": "t1", "description": "d1"},
|
|
{"name": "t2", "description": "d2"},
|
|
]):
|
|
result = tool.test_connection()
|
|
assert result["success"] is True
|
|
assert "2 tools" in result["message"]
|
|
assert len(result["tools"]) == 2
|
|
|
|
|
|
# =====================================================================
|
|
# DocsGPTOAuth extended
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.skip(reason="OAuth class signatures changed post-PG migration (db kwarg removed); needs rewrite")
|
|
@pytest.mark.unit
|
|
class TestDocsGPTOAuthExtended:
|
|
|
|
def test_process_auth_url_success(self):
|
|
from application.agents.tools.mcp_tool import DocsGPTOAuth
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_collection.find_one.return_value = None
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
oauth = DocsGPTOAuth(
|
|
mcp_url="https://mcp.example.com/api",
|
|
scopes=["read", "write"],
|
|
redis_client=MagicMock(),
|
|
redirect_uri="https://example.com/callback",
|
|
task_id="task1",
|
|
db=mock_db,
|
|
user_id="user1",
|
|
)
|
|
|
|
url, state = oauth._process_auth_url(
|
|
"https://auth.example.com/authorize?state=abc123&client_id=xyz"
|
|
)
|
|
assert state == "abc123"
|
|
assert url == "https://auth.example.com/authorize?state=abc123&client_id=xyz"
|
|
|
|
def test_process_auth_url_no_state(self):
|
|
from application.agents.tools.mcp_tool import DocsGPTOAuth
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_collection.find_one.return_value = None
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
oauth = DocsGPTOAuth(
|
|
mcp_url="https://mcp.example.com/api",
|
|
scopes="read",
|
|
redis_client=MagicMock(),
|
|
redirect_uri="https://example.com/callback",
|
|
db=mock_db,
|
|
user_id="user1",
|
|
)
|
|
|
|
with pytest.raises(Exception, match="Failed to process auth URL"):
|
|
oauth._process_auth_url("https://auth.example.com/authorize?client_id=xyz")
|
|
|
|
def test_redirect_handler_stores_in_redis(self):
|
|
from application.agents.tools.mcp_tool import DocsGPTOAuth
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_collection.find_one.return_value = None
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
mock_redis = MagicMock()
|
|
|
|
oauth = DocsGPTOAuth(
|
|
mcp_url="https://mcp.example.com/api",
|
|
scopes=["read"],
|
|
redis_client=mock_redis,
|
|
redirect_uri="https://example.com/callback",
|
|
task_id="task1",
|
|
db=mock_db,
|
|
user_id="user1",
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
loop.run_until_complete(
|
|
oauth.redirect_handler(
|
|
"https://auth.example.com/authorize?state=mystate&code=123"
|
|
)
|
|
)
|
|
finally:
|
|
loop.close()
|
|
|
|
assert oauth.auth_url == "https://auth.example.com/authorize?state=mystate&code=123"
|
|
assert oauth.extracted_state == "mystate"
|
|
# Redis setex should have been called for auth_url and status
|
|
assert mock_redis.setex.call_count >= 2
|
|
|
|
def test_redirect_handler_no_redis(self):
|
|
from application.agents.tools.mcp_tool import DocsGPTOAuth
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_collection.find_one.return_value = None
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
oauth = DocsGPTOAuth(
|
|
mcp_url="https://mcp.example.com/api",
|
|
scopes=["read"],
|
|
redis_client=None,
|
|
redirect_uri="https://example.com/callback",
|
|
db=mock_db,
|
|
user_id="user1",
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
loop.run_until_complete(
|
|
oauth.redirect_handler(
|
|
"https://auth.example.com/authorize?state=s1"
|
|
)
|
|
)
|
|
finally:
|
|
loop.close()
|
|
|
|
assert oauth.extracted_state == "s1"
|
|
|
|
def test_callback_handler_no_redis_raises(self):
|
|
from application.agents.tools.mcp_tool import DocsGPTOAuth
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_collection.find_one.return_value = None
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
oauth = DocsGPTOAuth(
|
|
mcp_url="https://mcp.example.com/api",
|
|
scopes=["read"],
|
|
redis_client=None,
|
|
redirect_uri="https://example.com/callback",
|
|
db=mock_db,
|
|
user_id="user1",
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
with pytest.raises(Exception, match="Redis client or state not configured"):
|
|
loop.run_until_complete(oauth.callback_handler())
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_callback_handler_receives_code(self):
|
|
from application.agents.tools.mcp_tool import DocsGPTOAuth
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_collection.find_one.return_value = None
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
mock_redis = MagicMock()
|
|
# First get returns the code
|
|
mock_redis.get.return_value = b"auth_code_123"
|
|
|
|
oauth = DocsGPTOAuth(
|
|
mcp_url="https://mcp.example.com/api",
|
|
scopes=["read"],
|
|
redis_client=mock_redis,
|
|
redirect_uri="https://example.com/callback",
|
|
task_id="task1",
|
|
db=mock_db,
|
|
user_id="user1",
|
|
)
|
|
oauth.extracted_state = "mystate"
|
|
oauth.auth_url = "https://auth.example.com/authorize"
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
code, state = loop.run_until_complete(oauth.callback_handler())
|
|
assert code == "auth_code_123"
|
|
assert state == "mystate"
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_callback_handler_receives_error(self):
|
|
from application.agents.tools.mcp_tool import DocsGPTOAuth
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_collection.find_one.return_value = None
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
mock_redis = MagicMock()
|
|
# First get for code returns None, second get for error returns error
|
|
mock_redis.get.side_effect = [None, b"access_denied"]
|
|
|
|
oauth = DocsGPTOAuth(
|
|
mcp_url="https://mcp.example.com/api",
|
|
scopes=["read"],
|
|
redis_client=mock_redis,
|
|
redirect_uri="https://example.com/callback",
|
|
db=mock_db,
|
|
user_id="user1",
|
|
)
|
|
oauth.extracted_state = "mystate"
|
|
oauth.auth_url = "https://auth.example.com/authorize"
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
with pytest.raises(Exception, match="OAuth error: access_denied"):
|
|
loop.run_until_complete(oauth.callback_handler())
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_init_scopes_as_string(self):
|
|
from application.agents.tools.mcp_tool import DocsGPTOAuth
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_collection.find_one.return_value = None
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
oauth = DocsGPTOAuth(
|
|
mcp_url="https://mcp.example.com/api",
|
|
scopes="read write",
|
|
redis_client=MagicMock(),
|
|
redirect_uri="https://example.com/callback",
|
|
db=mock_db,
|
|
user_id="user1",
|
|
)
|
|
assert oauth.server_base_url == "https://mcp.example.com"
|
|
|
|
|
|
# =====================================================================
|
|
# DBTokenStorage extended
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.skip(reason="DBTokenStorage signature changed post-PG migration; needs repo-based rewrite")
|
|
@pytest.mark.unit
|
|
class TestDBTokenStorageExtended:
|
|
|
|
def test_get_tokens_with_valid_data(self):
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_collection.find_one.return_value = {
|
|
"tokens": {
|
|
"access_token": "at_123",
|
|
"token_type": "bearer",
|
|
}
|
|
}
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
storage = DBTokenStorage(
|
|
server_url="https://mcp.example.com",
|
|
user_id="user1",
|
|
db_client=mock_db,
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
result = loop.run_until_complete(storage.get_tokens())
|
|
assert result is not None
|
|
assert result.access_token == "at_123"
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_get_tokens_with_invalid_data(self):
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_collection.find_one.return_value = {
|
|
"tokens": {"bad_field": "bad_value"}
|
|
}
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
storage = DBTokenStorage(
|
|
server_url="https://mcp.example.com",
|
|
user_id="user1",
|
|
db_client=mock_db,
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
result = loop.run_until_complete(storage.get_tokens())
|
|
assert result is None
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_set_tokens(self):
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
from mcp.shared.auth import OAuthToken
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
storage = DBTokenStorage(
|
|
server_url="https://mcp.example.com",
|
|
user_id="user1",
|
|
db_client=mock_db,
|
|
)
|
|
|
|
token = OAuthToken(access_token="new_token", token_type="bearer")
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
loop.run_until_complete(storage.set_tokens(token))
|
|
mock_collection.update_one.assert_called_once()
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_get_client_info_none(self):
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_collection.find_one.return_value = None
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
storage = DBTokenStorage(
|
|
server_url="https://mcp.example.com",
|
|
user_id="user1",
|
|
db_client=mock_db,
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
result = loop.run_until_complete(storage.get_client_info())
|
|
assert result is None
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_get_client_info_no_client_info_key(self):
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_collection.find_one.return_value = {"tokens": {}}
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
storage = DBTokenStorage(
|
|
server_url="https://mcp.example.com",
|
|
user_id="user1",
|
|
db_client=mock_db,
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
result = loop.run_until_complete(storage.get_client_info())
|
|
assert result is None
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_get_client_info_with_valid_data(self):
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_collection.find_one.return_value = {
|
|
"client_info": {
|
|
"client_id": "cid123",
|
|
"redirect_uris": ["https://example.com/callback"],
|
|
}
|
|
}
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
storage = DBTokenStorage(
|
|
server_url="https://mcp.example.com",
|
|
user_id="user1",
|
|
db_client=mock_db,
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
result = loop.run_until_complete(storage.get_client_info())
|
|
assert result is not None
|
|
assert result.client_id == "cid123"
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_get_client_info_redirect_uri_mismatch(self):
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_collection.find_one.return_value = {
|
|
"client_info": {
|
|
"client_id": "cid123",
|
|
"redirect_uris": ["https://old.example.com/callback"],
|
|
}
|
|
}
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
storage = DBTokenStorage(
|
|
server_url="https://mcp.example.com",
|
|
user_id="user1",
|
|
db_client=mock_db,
|
|
expected_redirect_uri="https://new.example.com/callback",
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
result = loop.run_until_complete(storage.get_client_info())
|
|
assert result is None
|
|
mock_collection.update_one.assert_called_once()
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_get_client_info_invalid_data(self):
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_collection.find_one.return_value = {
|
|
"client_info": {"invalid_key": "value"}
|
|
}
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
storage = DBTokenStorage(
|
|
server_url="https://mcp.example.com",
|
|
user_id="user1",
|
|
db_client=mock_db,
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
result = loop.run_until_complete(storage.get_client_info())
|
|
assert result is None
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_set_client_info(self):
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
from mcp.shared.auth import OAuthClientInformationFull
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
storage = DBTokenStorage(
|
|
server_url="https://mcp.example.com",
|
|
user_id="user1",
|
|
db_client=mock_db,
|
|
)
|
|
|
|
client_info = OAuthClientInformationFull(
|
|
client_id="cid123",
|
|
redirect_uris=["https://example.com/callback"],
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
loop.run_until_complete(storage.set_client_info(client_info))
|
|
mock_collection.update_one.assert_called_once()
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_clear_all(self):
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
|
|
mock_db = MagicMock()
|
|
mock_collection = MagicMock()
|
|
mock_db.__getitem__ = MagicMock(return_value=mock_collection)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
loop.run_until_complete(DBTokenStorage.clear_all(mock_db))
|
|
mock_collection.delete_many.assert_called_once_with({})
|
|
finally:
|
|
loop.close()
|
|
|
|
def test_serialize_client_info_without_redirect_uris(self):
|
|
from application.agents.tools.mcp_tool import DBTokenStorage
|
|
|
|
mock_db = MagicMock()
|
|
storage = DBTokenStorage(
|
|
server_url="https://mcp.example.com",
|
|
user_id="user1",
|
|
db_client=mock_db,
|
|
)
|
|
info = {"client_name": "test"}
|
|
result = storage._serialize_client_info(info)
|
|
assert result == {"client_name": "test"}
|
|
|
|
|
|
# =====================================================================
|
|
# MCPOAuthManager extended
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestMCPOAuthManagerExtended:
|
|
|
|
def test_handle_callback_redis_setex_for_state(self):
|
|
from application.agents.tools.mcp_tool import MCPOAuthManager
|
|
|
|
mock_redis = MagicMock()
|
|
manager = MCPOAuthManager(mock_redis)
|
|
|
|
result = manager.handle_oauth_callback(state="s1", code="c1")
|
|
assert result is True
|
|
# Should call setex for code and state
|
|
assert mock_redis.setex.call_count == 2
|
|
|
|
def test_handle_callback_with_error_stores_error(self):
|
|
from application.agents.tools.mcp_tool import MCPOAuthManager
|
|
|
|
mock_redis = MagicMock()
|
|
manager = MCPOAuthManager(mock_redis)
|
|
|
|
result = manager.handle_oauth_callback(
|
|
state="s1", code="", error="invalid_scope"
|
|
)
|
|
assert result is False
|
|
# Should store error in redis
|
|
mock_redis.setex.assert_called()
|
|
|
|
def test_get_oauth_status_returns_not_found_on_redis_error(self):
|
|
"""A failure inside ``xrevrange`` (Redis down, network blip) is
|
|
swallowed — the manager returns ``not_found`` so the caller
|
|
can present a clean "OAuth failed, try again" message rather
|
|
than a 500.
|
|
"""
|
|
from application.agents.tools.mcp_tool import MCPOAuthManager
|
|
|
|
mock_redis = MagicMock()
|
|
mock_redis.xrevrange.side_effect = Exception("Redis went away")
|
|
|
|
manager = MCPOAuthManager(mock_redis)
|
|
result = manager.get_oauth_status("task123", "alice")
|
|
|
|
assert result["status"] == "not_found"
|
|
|
|
|
|
# =====================================================================
|
|
# get_actions_metadata extended
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestGetActionsMetadataExtended:
|
|
|
|
def test_tools_with_schema_key(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool.available_tools = [
|
|
{
|
|
"name": "schema_tool",
|
|
"description": "Uses schema key",
|
|
"schema": {
|
|
"type": "object",
|
|
"properties": {"q": {"type": "string"}},
|
|
"required": ["q"],
|
|
},
|
|
}
|
|
]
|
|
meta = tool.get_actions_metadata()
|
|
assert len(meta) == 1
|
|
assert "q" in meta[0]["parameters"]["properties"]
|
|
|
|
def test_tools_with_input_schema_key(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool.available_tools = [
|
|
{
|
|
"name": "is_tool",
|
|
"description": "Uses input_schema key",
|
|
"input_schema": {
|
|
"type": "object",
|
|
"properties": {"x": {"type": "number"}},
|
|
},
|
|
}
|
|
]
|
|
meta = tool.get_actions_metadata()
|
|
assert "x" in meta[0]["parameters"]["properties"]
|
|
|
|
def test_multiple_tools(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool.available_tools = [
|
|
{"name": "a", "description": "da"},
|
|
{"name": "b", "description": "db", "inputSchema": {
|
|
"type": "object",
|
|
"properties": {"p": {"type": "string"}},
|
|
}},
|
|
]
|
|
meta = tool.get_actions_metadata()
|
|
assert len(meta) == 2
|
|
assert meta[0]["name"] == "a"
|
|
assert meta[0]["parameters"]["properties"] == {}
|
|
assert "p" in meta[1]["parameters"]["properties"]
|
|
|
|
|
|
# =====================================================================
|
|
# Coverage gap tests (lines 207-210, 288-293, 346-347, 416-417, 620)
|
|
# =====================================================================
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestMCPToolGaps:
|
|
|
|
def test_create_transport_stdio_raises(self, mcp_config):
|
|
"""Cover line 199-200: stdio transport raises ValueError."""
|
|
mcp_config["transport_type"] = "stdio"
|
|
tool = _make_tool(mcp_config)
|
|
|
|
with pytest.raises(ValueError, match="STDIO transport is disabled"):
|
|
tool._create_transport()
|
|
|
|
def test_run_in_new_loop(self, mcp_config):
|
|
"""Cover lines 288-293: _run_in_new_loop creates a new event loop."""
|
|
tool = _make_tool(mcp_config)
|
|
|
|
async def dummy_operation(*args, **kwargs):
|
|
return "result"
|
|
|
|
tool._execute_with_client = dummy_operation
|
|
result = tool._run_in_new_loop("test_op")
|
|
assert result == "result"
|
|
|
|
def test_execute_action_auth_error_retry(self, mcp_config):
|
|
"""Cover lines 346-347: auth error detection in execute_action."""
|
|
tool = _make_tool(mcp_config)
|
|
tool.available_tools = [{"name": "test_action"}]
|
|
tool.auth_type = "bearer"
|
|
|
|
call_count = 0
|
|
|
|
def mock_run_async(operation, action_name, **kwargs):
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count == 1:
|
|
raise Exception("401 unauthorized")
|
|
return MagicMock(content=[MagicMock(text="success")])
|
|
|
|
tool._run_async_operation = mock_run_async
|
|
tool._setup_client = MagicMock()
|
|
|
|
result = tool.execute_action("test_action")
|
|
assert "success" in str(result)
|
|
|
|
def test_test_connection_invalid_url(self, mcp_config):
|
|
"""Cover lines 416-417: test_connection with invalid URL."""
|
|
tool = _make_tool(mcp_config)
|
|
tool.server_url = "not a url at all"
|
|
tool._client = None
|
|
|
|
result = tool.test_connection()
|
|
assert result["success"] is False
|
|
|
|
def test_get_config_requirements_has_username(self, mcp_config):
|
|
"""Cover line 620: config requirements include username field."""
|
|
tool = _make_tool(mcp_config)
|
|
config = tool.get_config_requirements()
|
|
assert "username" in config
|
|
assert config["username"]["description"] == "Username for basic authentication"
|
|
assert config["username"]["depends_on"] == {"auth_type": "basic"}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Coverage — additional uncovered lines: 207-210, 288-293, 346-347, 416-417, 620
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestMCPToolTransportCreation:
|
|
|
|
def test_unknown_transport_defaults_to_http(self, mcp_config):
|
|
"""Cover line 207-210 and 212: unknown transport type defaults to StreamableHttpTransport."""
|
|
mcp_config["transport_type"] = "unknown_protocol"
|
|
tool = _make_tool(mcp_config)
|
|
transport = tool._create_transport()
|
|
# Should be StreamableHttpTransport (the default fallback)
|
|
assert transport is not None
|
|
|
|
def test_sse_transport_creation(self, mcp_config):
|
|
"""Cover lines 201-203: SSE transport creation."""
|
|
mcp_config["transport_type"] = "sse"
|
|
tool = _make_tool(mcp_config)
|
|
transport = tool._create_transport()
|
|
assert transport is not None
|
|
|
|
def test_stdio_transport_raises(self, mcp_config):
|
|
"""Cover line 200: stdio transport is disabled."""
|
|
mcp_config["transport_type"] = "stdio"
|
|
tool = _make_tool(mcp_config)
|
|
with pytest.raises(ValueError, match="STDIO transport is disabled"):
|
|
tool._create_transport()
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestMCPToolRunAsyncOperation:
|
|
|
|
def test_run_async_operation_maps_error(self, mcp_config):
|
|
"""Cover lines 288-293: _run_async_operation exception mapped."""
|
|
tool = _make_tool(mcp_config)
|
|
|
|
async def bad_execute(op, *a, **kw):
|
|
raise ConnectionRefusedError("refused")
|
|
|
|
tool._execute_with_client = bad_execute
|
|
tool._client = MagicMock()
|
|
|
|
with pytest.raises(Exception, match="Connection refused"):
|
|
tool._run_async_operation("ping")
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestMCPToolExecuteActionAuth:
|
|
|
|
def test_execute_action_oauth_auth_error(self, mcp_config):
|
|
"""Cover lines 346-347: OAuth auth error raises specific message."""
|
|
mcp_config["auth_type"] = "oauth"
|
|
tool = _make_tool(mcp_config)
|
|
|
|
def bad_run(*a, **kw):
|
|
raise Exception("401 Unauthorized")
|
|
|
|
tool._run_async_operation = bad_run
|
|
tool._client = MagicMock()
|
|
|
|
with pytest.raises(Exception, match="OAuth session expired"):
|
|
tool.execute_action("test_action")
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestMCPToolTestConnectionInvalidScheme:
|
|
|
|
def test_test_connection_ftp_scheme_invalid(self, mcp_config):
|
|
"""Cover lines 416-417: test_connection with invalid URL scheme."""
|
|
tool = _make_tool(mcp_config)
|
|
tool.server_url = "ftp://invalid.example.com"
|
|
tool._client = None
|
|
|
|
result = tool.test_connection()
|
|
assert result["success"] is False
|
|
assert "scheme" in result["message"].lower() or "Invalid" in result["message"]
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestMCPToolConfigRequirements:
|
|
|
|
def test_get_config_requirements_has_password(self, mcp_config):
|
|
"""Cover line 620+: config requirements include password field."""
|
|
tool = _make_tool(mcp_config)
|
|
config = tool.get_config_requirements()
|
|
assert "password" in config
|
|
assert config["password"]["secret"] is True
|
|
assert config["password"]["depends_on"] == {"auth_type": "basic"}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Additional coverage for mcp_tool.py
|
|
# Lines: 207-210 (stdio transport), 288-293 (_map_error + _run_in_new_loop),
|
|
# 346-347 (execute_action error handling), 620 (config requirements password)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestMCPToolStdioTransport:
|
|
"""Cover line 199-200: stdio transport raises ValueError."""
|
|
|
|
def test_create_stdio_transport_raises(self, mcp_config):
|
|
mcp_config["transport_type"] = "stdio"
|
|
tool = _make_tool(mcp_config)
|
|
tool.transport_type = "stdio"
|
|
|
|
with pytest.raises(ValueError, match="STDIO transport is disabled"):
|
|
tool._create_transport()
|
|
|
|
def test_create_sse_transport(self, mcp_config):
|
|
"""Cover line 201-203: SSE transport."""
|
|
mcp_config["transport_type"] = "sse"
|
|
tool = _make_tool(mcp_config)
|
|
tool.transport_type = "sse"
|
|
transport = tool._create_transport()
|
|
# Should return an SSETransport
|
|
assert transport is not None
|
|
|
|
def test_create_unknown_transport_defaults_to_http(self, mcp_config):
|
|
"""Cover line 211-212: unknown transport defaults to StreamableHttpTransport."""
|
|
tool = _make_tool(mcp_config)
|
|
tool.transport_type = "unknown_transport"
|
|
transport = tool._create_transport()
|
|
assert transport is not None
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestMCPToolRunInNewLoop:
|
|
"""Cover lines 288-293: _run_in_new_loop and _map_error."""
|
|
|
|
def test_run_in_new_loop(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
|
|
async def mock_execute(*args, **kwargs):
|
|
return "loop_result"
|
|
|
|
tool._execute_with_client = mock_execute
|
|
result = tool._run_in_new_loop("list_tools")
|
|
assert result == "loop_result"
|
|
|
|
def test_map_error_timeout(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
from asyncio import TimeoutError as AsyncTimeout
|
|
|
|
err = tool._map_error("test_op", AsyncTimeout("timed out"))
|
|
assert isinstance(err, Exception)
|
|
assert "timeout" in str(err).lower() or "timed out" in str(err).lower()
|
|
|
|
def test_map_error_generic(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
err = tool._map_error("test_op", RuntimeError("something broke"))
|
|
assert isinstance(err, Exception)
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestMCPToolExecuteActionErrorHandling:
|
|
"""Cover lines 346-347: execute_action non-auth error."""
|
|
|
|
def test_execute_action_generic_error(self, mcp_config):
|
|
tool = _make_tool(mcp_config)
|
|
tool._run_async_operation = MagicMock(
|
|
side_effect=RuntimeError("generic failure")
|
|
)
|
|
with pytest.raises(Exception, match="Failed to execute action"):
|
|
tool.execute_action("some_action", key="value")
|