test_streamable_http.py•4.49 kB
from __future__ import annotations
import contextlib
import socket
import threading
import time
from collections.abc import Iterator
import pytest
import uvicorn
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client
from uniprot_mcp.http_app import app as asgi_app
@contextlib.contextmanager
def run_uvicorn(app) -> Iterator[str]:
"""Run the ASGI app in a background thread for the duration of a test."""
sock = socket.socket()
sock.bind(("127.0.0.1", 0))
host, port = sock.getsockname()
sock.close()
config = uvicorn.Config(app, host=host, port=port, log_level="error")
server = uvicorn.Server(config=config)
thread = threading.Thread(target=server.run, daemon=True)
thread.start()
try:
# Wait until the server has started.
while not server.started:
time.sleep(0.01)
yield f"http://{host}:{port}"
finally:
server.should_exit = True
thread.join(timeout=5)
@pytest.mark.asyncio
async def test_streamable_http_handshake_and_tools(monkeypatch, load_fixture):
async def fake_fetch_entry_json(client, accession, *, fields=None, version=None):
return load_fixture("entry_reviewed.json")
async def fake_search_json(
client,
*,
query,
size,
reviewed_only=False,
fields=None,
sort=None,
include_isoform=None,
):
return load_fixture("search_results.json")
async def fake_start_id_mapping(client, *, from_db, to_db, ids):
return "fake-job"
async def fake_get_mapping_status(client, job_id):
return {"status": "FINISHED"}
async def fake_get_mapping_results(client, job_id):
return load_fixture("mapping_result.json")
monkeypatch.setattr(
"uniprot_mcp.adapters.uniprot_client.fetch_entry_json", fake_fetch_entry_json
)
monkeypatch.setattr("uniprot_mcp.server.fetch_entry_json", fake_fetch_entry_json)
monkeypatch.setattr("uniprot_mcp.adapters.uniprot_client.search_json", fake_search_json)
monkeypatch.setattr("uniprot_mcp.server.search_json", fake_search_json)
monkeypatch.setattr(
"uniprot_mcp.adapters.uniprot_client.start_id_mapping", fake_start_id_mapping
)
monkeypatch.setattr("uniprot_mcp.server.start_id_mapping", fake_start_id_mapping)
monkeypatch.setattr(
"uniprot_mcp.adapters.uniprot_client.get_mapping_status",
fake_get_mapping_status,
)
monkeypatch.setattr("uniprot_mcp.server.get_mapping_status", fake_get_mapping_status)
monkeypatch.setattr(
"uniprot_mcp.adapters.uniprot_client.get_mapping_results",
fake_get_mapping_results,
)
monkeypatch.setattr("uniprot_mcp.server.get_mapping_results", fake_get_mapping_results)
with run_uvicorn(asgi_app) as base_url:
endpoint = f"{base_url}/mcp"
async with streamablehttp_client(endpoint) as (read, write, _):
async with ClientSession(read, write) as session:
await session.initialize()
tools = await session.list_tools()
tool_names = {tool.name for tool in tools.tools}
assert {"fetch_entry", "search_uniprot", "map_ids"}.issubset(tool_names)
entry = await session.call_tool("fetch_entry", {"accession": "P12345"})
assert entry.structuredContent is not None
assert entry.structuredContent["accession"] == "P12345"
search = await session.call_tool(
"search_uniprot",
{"query": "kinase", "size": 2, "reviewed_only": True},
)
assert search.structuredContent is not None
hits = (
search.structuredContent.get("result")
if isinstance(search.structuredContent, dict)
else search.structuredContent
)
assert isinstance(hits, list)
assert len(hits) == 2
mapping = await session.call_tool(
"map_ids",
{
"from_db": "UniProtKB_AC-ID",
"to_db": "Ensembl",
"ids": ["P12345"],
},
)
assert mapping.structuredContent is not None
results = mapping.structuredContent.get("results", {})
assert "P12345" in results