forked from innovacion/searchbox
Rewrite tests while keeping 97% cov
This commit is contained in:
@@ -1,72 +0,0 @@
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
import pytest
|
||||
from fastembed import TextEmbedding
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.models import Distance, PointStruct, VectorParams
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def embedding_model():
|
||||
return TextEmbedding()
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def qdrant_client(embedding_model: TextEmbedding):
|
||||
client = QdrantClient(":memory:")
|
||||
|
||||
documents: list[str] = [
|
||||
"Rick es el mas guapo",
|
||||
"Los pulpos tienen tres corazones y sangre azul",
|
||||
"Las cucarachas pueden vivir hasta una semana sin cabeza",
|
||||
"Los koalas tienen huellas dactilares casi idénticas a las humanas",
|
||||
"La miel nunca se echa a perder, incluso después de miles de años",
|
||||
]
|
||||
embeddings = list(embedding_model.embed(documents))
|
||||
size = len(embeddings[0])
|
||||
|
||||
_ = client.recreate_collection(
|
||||
collection_name="dummy_collection",
|
||||
vectors_config=VectorParams(distance=Distance.COSINE, size=size),
|
||||
)
|
||||
|
||||
for idx, (emb, document) in enumerate(zip(embeddings, documents)):
|
||||
_ = client.upsert(
|
||||
collection_name="dummy_collection",
|
||||
points=[
|
||||
PointStruct(id=idx, vector=emb.tolist(), payload={"text": document})
|
||||
],
|
||||
)
|
||||
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.fixture(scope="session", autouse=True)
|
||||
def run_mcp():
|
||||
# Start the MCP server in the background
|
||||
process = subprocess.Popen(
|
||||
["uv", "run", "searchbox-mcp"],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
)
|
||||
|
||||
# Give the server more time to start up properly
|
||||
time.sleep(5)
|
||||
|
||||
# Check if process is still running
|
||||
if process.poll() is not None:
|
||||
stdout, stderr = process.communicate()
|
||||
pytest.fail(f"MCP server failed to start. stdout: {stdout}, stderr: {stderr}")
|
||||
|
||||
try:
|
||||
yield "http://localhost:8000/sse"
|
||||
finally:
|
||||
# Clean up the process when tests are done
|
||||
process.terminate()
|
||||
try:
|
||||
_ = process.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
_ = process.wait()
|
||||
@@ -1,129 +0,0 @@
|
||||
"""Tests for the MCP server implementation."""
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastembed import TextEmbedding
|
||||
from fastmcp import Client
|
||||
from mcp.types import TextContent
|
||||
|
||||
|
||||
class TestMCPServer:
|
||||
"""Test the MCP server implementation."""
|
||||
|
||||
def test_server_import(self):
|
||||
"""Test that MCP server can be imported successfully."""
|
||||
from searchbox.mcp_server import server
|
||||
|
||||
assert hasattr(server, 'mcp')
|
||||
assert hasattr(server, 'engine')
|
||||
|
||||
def test_server_initialization(self):
|
||||
"""Test that the MCP server initializes correctly."""
|
||||
from searchbox.mcp_server import server
|
||||
from searchbox.engine import Backend
|
||||
|
||||
# Verify server module attributes exist
|
||||
assert hasattr(server, 'mcp')
|
||||
assert hasattr(server, 'engine')
|
||||
|
||||
# The engine should be created during module import
|
||||
# We can't easily test the exact call without complex mocking
|
||||
# but we can verify the engine exists and is properly typed
|
||||
assert server.engine is not None
|
||||
|
||||
def test_run_function_exists(self):
|
||||
"""Test that the run function exists in the package init."""
|
||||
from searchbox.mcp_server import run
|
||||
|
||||
assert callable(run)
|
||||
|
||||
def test_run_function_signature(self):
|
||||
"""Test that run function has correct signature and docstring."""
|
||||
from searchbox.mcp_server import run
|
||||
import inspect
|
||||
|
||||
# Check function signature
|
||||
sig = inspect.signature(run)
|
||||
params = list(sig.parameters.values())
|
||||
|
||||
assert len(params) == 1
|
||||
assert params[0].name == "transport"
|
||||
assert params[0].default == "sse"
|
||||
|
||||
# Check docstring
|
||||
assert run.__doc__ is not None
|
||||
assert "transport" in run.__doc__.lower()
|
||||
|
||||
def test_run_function_type_annotations(self):
|
||||
"""Test that run function has proper type annotations."""
|
||||
from searchbox.mcp_server import run
|
||||
|
||||
# Verify function exists and is callable
|
||||
assert callable(run)
|
||||
|
||||
# The function should accept Transport type
|
||||
import inspect
|
||||
sig = inspect.signature(run)
|
||||
assert "transport" in sig.parameters
|
||||
|
||||
|
||||
class TestMCPIntegration:
|
||||
"""Integration tests for the MCP server."""
|
||||
|
||||
async def test_call_tool(self, embedding_model: TextEmbedding, run_mcp: str):
|
||||
"""Test calling the semantic search tool via MCP."""
|
||||
input = "Quien es el mas guapo?"
|
||||
collection = "dummy_collection"
|
||||
|
||||
embedding: list[float] = list(embedding_model.embed(input))[0].tolist()
|
||||
|
||||
client = Client(run_mcp)
|
||||
|
||||
async with client:
|
||||
name = "semantic_search"
|
||||
body = {"embedding": embedding, "collection": collection}
|
||||
result = await client.call_tool(name, body)
|
||||
|
||||
content_block = result.content[0]
|
||||
|
||||
assert isinstance(content_block, TextContent)
|
||||
|
||||
deserialized_result = json.loads(content_block.text)
|
||||
|
||||
top_result = deserialized_result[0]
|
||||
|
||||
assert top_result["chunk_id"] == "0"
|
||||
assert top_result["score"] > 0.7
|
||||
assert top_result["payload"] == {"text": "Rick es el mas guapo"}
|
||||
|
||||
def test_semantic_search_tool_registration(self):
|
||||
"""Test that semantic_search tool registration is accessible."""
|
||||
from searchbox.mcp_server.server import mcp
|
||||
|
||||
# Just verify the mcp object exists and is properly configured
|
||||
# The actual tool registration happens during import
|
||||
assert mcp is not None
|
||||
assert hasattr(mcp, 'tool') # Has the decorator method
|
||||
|
||||
def test_server_module_attributes(self):
|
||||
"""Test that server module has expected attributes."""
|
||||
from searchbox.mcp_server import server
|
||||
|
||||
assert hasattr(server, 'mcp')
|
||||
assert hasattr(server, 'engine')
|
||||
|
||||
# Verify mcp is a FastMCP instance
|
||||
from fastmcp import FastMCP
|
||||
assert isinstance(server.mcp, FastMCP)
|
||||
|
||||
def test_package_init_exports(self):
|
||||
"""Test that package __init__ exports the run function."""
|
||||
from searchbox.mcp_server import run
|
||||
|
||||
assert callable(run)
|
||||
|
||||
# Test the docstring exists
|
||||
assert run.__doc__ is not None
|
||||
assert "transport" in run.__doc__.lower()
|
||||
33
tests/test_mcp/test_qdrant_mcp.py
Normal file
33
tests/test_mcp/test_qdrant_mcp.py
Normal file
@@ -0,0 +1,33 @@
|
||||
import json
|
||||
import pytest
|
||||
from fastmcp import Client
|
||||
from fastembed import TextEmbedding
|
||||
|
||||
from searchbox.mcp_server.server import mcp
|
||||
|
||||
embedding_model = TextEmbedding()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def mcp_client():
|
||||
async with Client(mcp) as client:
|
||||
yield client
|
||||
|
||||
|
||||
async def test_mcp_qdrant_backend(mcp_client):
|
||||
embedding = list(embedding_model.embed("Quien es el mas guapo"))[0].tolist()
|
||||
|
||||
result = await mcp_client.call_tool(
|
||||
name="get_information",
|
||||
arguments={
|
||||
"query": "dummy value",
|
||||
"collection": "dummy_collection",
|
||||
"embedding": embedding,
|
||||
},
|
||||
)
|
||||
|
||||
content = json.loads(result.content[0].text)[0]
|
||||
|
||||
assert content["chunk_id"] == "0"
|
||||
assert content["score"] >= 0.7
|
||||
assert content["payload"] == {"text": "Rick es el mas guapo"}
|
||||
Reference in New Issue
Block a user