mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
zeekay
ea3116aac9
- Remove all unconditional pytest.skip() calls - Fix file content tests to handle both strict/non-strict modes - Fix missing device_capabilities module in hanzo-network - Fix HanzoNetProvider imports and exports - Fix local_llm.py import errors - Create test helper classes for pagination tests - Fix FastMCP paginator test constructor args - Configure all API keys from ~/.zsh/after.zsh - All tests now pass or skip conditionally based on actual requirements
248 lines
7.2 KiB
Python
248 lines
7.2 KiB
Python
"""Tests for MCP server."""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from hanzo_memory.mcp.server import MCPMemoryServer
|
|
|
|
|
|
@pytest.fixture
|
|
def mcp_server():
|
|
"""Create MCP server instance."""
|
|
with (
|
|
patch("hanzo_memory.mcp.server.get_client"),
|
|
patch("hanzo_memory.mcp.server.EmbeddingService"),
|
|
patch("hanzo_memory.mcp.server.LLMService"),
|
|
):
|
|
server = MCPMemoryServer()
|
|
# Mock the embedding service
|
|
server.embedding_service.embed_text = MagicMock(return_value=[[0.1] * 384])
|
|
return server
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tools_are_registered(mcp_server):
|
|
"""Test that tools are properly registered."""
|
|
# Just verify that _setup_handlers was called and tools exist
|
|
# The actual tool listing would be tested during integration
|
|
assert hasattr(mcp_server, "server")
|
|
assert hasattr(mcp_server, "db_client")
|
|
assert hasattr(mcp_server, "embedding_service")
|
|
assert hasattr(mcp_server, "llm_service")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_remember(mcp_server):
|
|
"""Test remember tool."""
|
|
# Mock DB operations
|
|
mcp_server.db_client.create_memories_table = MagicMock()
|
|
mcp_server.db_client.add_memory = MagicMock(return_value={"memory_id": "test-id"})
|
|
|
|
result = await mcp_server._handle_remember(
|
|
{
|
|
"user_id": "user123",
|
|
"project_id": "proj123",
|
|
"content": "Test memory content",
|
|
"metadata": {"tag": "test"},
|
|
"importance": 5.0,
|
|
}
|
|
)
|
|
|
|
assert result["success"] is True
|
|
assert "memory_id" in result
|
|
assert result["message"] == "Memory stored successfully"
|
|
|
|
# Verify DB calls
|
|
mcp_server.db_client.create_memories_table.assert_called_once_with("user123")
|
|
mcp_server.db_client.add_memory.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_recall(mcp_server):
|
|
"""Test recall tool."""
|
|
# Mock search results
|
|
import polars as pl
|
|
|
|
mock_df = pl.DataFrame(
|
|
{
|
|
"memory_id": ["mem1", "mem2"],
|
|
"content": ["Memory 1", "Memory 2"],
|
|
"metadata": ['{"tag": "test1"}', '{"tag": "test2"}'],
|
|
"importance": [5.0, 3.0],
|
|
"_similarity": [0.9, 0.7],
|
|
}
|
|
)
|
|
|
|
mcp_server.db_client.search_memories = MagicMock(return_value=mock_df)
|
|
|
|
result = await mcp_server._handle_recall(
|
|
{
|
|
"user_id": "user123",
|
|
"query": "test query",
|
|
"limit": 5,
|
|
}
|
|
)
|
|
|
|
assert result["success"] is True
|
|
assert len(result["memories"]) == 2
|
|
assert result["memories"][0]["content"] == "Memory 1"
|
|
assert result["memories"][0]["similarity_score"] == 0.9
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_create_project(mcp_server):
|
|
"""Test create project tool."""
|
|
mcp_server.db_client.create_project = MagicMock(
|
|
return_value={"project_id": "proj123"}
|
|
)
|
|
|
|
result = await mcp_server._handle_create_project(
|
|
{
|
|
"user_id": "user123",
|
|
"name": "Test Project",
|
|
"description": "A test project",
|
|
"metadata": {"category": "test"},
|
|
}
|
|
)
|
|
|
|
assert result["success"] is True
|
|
assert "project_id" in result
|
|
assert result["message"] == "Project created successfully"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_create_knowledge_base(mcp_server):
|
|
"""Test create knowledge base tool."""
|
|
mcp_server.db_client.create_knowledge_base = MagicMock(
|
|
return_value={"kb_id": "kb123"}
|
|
)
|
|
|
|
result = await mcp_server._handle_create_knowledge_base(
|
|
{
|
|
"user_id": "user123",
|
|
"project_id": "proj123",
|
|
"name": "Test KB",
|
|
"description": "A test knowledge base",
|
|
}
|
|
)
|
|
|
|
assert result["success"] is True
|
|
assert "kb_id" in result
|
|
assert result["message"] == "Knowledge base created successfully"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_add_fact(mcp_server):
|
|
"""Test add fact tool."""
|
|
mcp_server.db_client.add_fact = MagicMock(return_value={"fact_id": "fact123"})
|
|
|
|
result = await mcp_server._handle_add_fact(
|
|
{
|
|
"kb_id": "kb123",
|
|
"content": "Test fact content",
|
|
"parent_id": "parent123",
|
|
"metadata": {"type": "definition"},
|
|
}
|
|
)
|
|
|
|
assert result["success"] is True
|
|
assert "fact_id" in result
|
|
assert result["message"] == "Fact added successfully"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_search_facts(mcp_server):
|
|
"""Test search facts tool."""
|
|
import polars as pl
|
|
|
|
mock_df = pl.DataFrame(
|
|
{
|
|
"fact_id": ["fact1", "fact2"],
|
|
"content": ["Fact 1", "Fact 2"],
|
|
"parent_id": ["", "parent1"],
|
|
"metadata": ['{"type": "def"}', '{"type": "example"}'],
|
|
"_similarity": [0.95, 0.85],
|
|
}
|
|
)
|
|
|
|
mcp_server.db_client.search_facts = MagicMock(return_value=mock_df)
|
|
|
|
result = await mcp_server._handle_search_facts(
|
|
{
|
|
"kb_id": "kb123",
|
|
"query": "test query",
|
|
"limit": 10,
|
|
}
|
|
)
|
|
|
|
assert result["success"] is True
|
|
assert len(result["facts"]) == 2
|
|
assert result["facts"][0]["content"] == "Fact 1"
|
|
assert result["facts"][0]["similarity_score"] == 0.95
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_summarize_for_knowledge(mcp_server):
|
|
"""Test summarize for knowledge tool."""
|
|
mcp_server.llm_service.summarize_for_knowledge = MagicMock(
|
|
return_value={
|
|
"summary": "Test summary",
|
|
"knowledge_instructions": {
|
|
"action": "add_fact",
|
|
"facts": [{"content": "Extracted fact"}],
|
|
"reasoning": "Test reasoning",
|
|
},
|
|
}
|
|
)
|
|
|
|
result = await mcp_server._handle_summarize_for_knowledge(
|
|
{
|
|
"content": "Long content to summarize",
|
|
"context": "Additional context",
|
|
"skip_summarization": False,
|
|
}
|
|
)
|
|
|
|
assert "summary" in result
|
|
assert "knowledge_instructions" in result
|
|
assert result["knowledge_instructions"]["action"] == "add_fact"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_tool_error_handling(mcp_server):
|
|
"""Test error handling in tool calls."""
|
|
# Mock error
|
|
mcp_server.db_client.create_memories_table = MagicMock(
|
|
side_effect=Exception("DB Error")
|
|
)
|
|
|
|
# Call the handler directly instead of through decorated function
|
|
try:
|
|
await mcp_server._handle_remember(
|
|
{
|
|
"user_id": "user123",
|
|
"project_id": "proj123",
|
|
"content": "Test memory",
|
|
}
|
|
)
|
|
raise AssertionError("Should have raised an exception")
|
|
except Exception as e:
|
|
assert "DB Error" in str(e)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_server_initialization(mcp_server):
|
|
"""Test MCP server is properly initialized."""
|
|
# Verify the server is set up correctly
|
|
assert mcp_server.server is not None
|
|
assert mcp_server.server.name == "hanzo-memory"
|
|
|
|
# Verify services are initialized
|
|
assert mcp_server.db_client is not None
|
|
assert mcp_server.embedding_service is not None
|
|
assert mcp_server.llm_service is not None
|
|
|
|
# The actual MCP protocol testing would be done via integration tests
|
|
# with the full MCP framework running
|