mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
136 lines
4.1 KiB
Python
136 lines
4.1 KiB
Python
"""Base test class for memory-related tests to reduce redundancy."""
|
|
|
|
from datetime import datetime
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
from conftest import ToolTestHelper, create_mock_ctx
|
|
from fastmcp import FastMCP
|
|
|
|
# Import guard for optional hanzo_memory dependency
|
|
try:
|
|
from hanzo_memory.models import Memory
|
|
|
|
HANZO_MEMORY_AVAILABLE = True
|
|
except ImportError:
|
|
HANZO_MEMORY_AVAILABLE = False
|
|
Memory = None # type: ignore
|
|
|
|
# Skip entire module if hanzo_memory is not available
|
|
pytestmark = pytest.mark.skipif(
|
|
not HANZO_MEMORY_AVAILABLE, reason="hanzo_memory package not installed"
|
|
)
|
|
|
|
# Only import these if hanzo_memory is available
|
|
if HANZO_MEMORY_AVAILABLE:
|
|
from hanzo_tools.memory import (
|
|
CreateMemoriesTool,
|
|
DeleteMemoriesTool,
|
|
RecallMemoriesTool,
|
|
UpdateMemoriesTool,
|
|
)
|
|
|
|
|
|
class MemoryTestBase:
|
|
"""Base class for memory test cases with common fixtures and utilities."""
|
|
|
|
@pytest.fixture
|
|
def tool_helper(self):
|
|
"""Provide ToolTestHelper for tests."""
|
|
return ToolTestHelper
|
|
|
|
@pytest.fixture
|
|
def mock_ctx(self):
|
|
"""Create mock context for tool calls."""
|
|
return create_mock_ctx()
|
|
|
|
@pytest.fixture
|
|
def mock_memory(self):
|
|
"""Create a standard mock memory object."""
|
|
return Memory(
|
|
memory_id="test_123",
|
|
user_id="test_user",
|
|
project_id="test_project",
|
|
content="Test memory content",
|
|
metadata={"type": "statement"},
|
|
importance=1.0,
|
|
created_at=datetime.fromisoformat("2024-01-01T00:00:00"),
|
|
updated_at=datetime.fromisoformat("2024-01-01T00:00:00"),
|
|
embedding=[0.1] * 1536,
|
|
)
|
|
|
|
@pytest.fixture
|
|
def mock_memory_service(self):
|
|
"""Create a mock memory service with standard responses."""
|
|
with patch(
|
|
"hanzo_memory.services.memory.get_memory_service"
|
|
) as mock_get_service:
|
|
mock_service = Mock()
|
|
|
|
# Set up standard responses
|
|
mock_service.create_memory.return_value = Mock(
|
|
memory_id="mem_123",
|
|
user_id="test_user",
|
|
project_id="test_project",
|
|
content="Created memory",
|
|
metadata={},
|
|
importance=1.0,
|
|
)
|
|
|
|
mock_service.update_memory.return_value = Mock(
|
|
memory_id="mem_123", content="Updated memory"
|
|
)
|
|
|
|
mock_service.delete_memory.return_value = None
|
|
|
|
mock_service.search_memories.return_value = [
|
|
Mock(
|
|
memory_id="mem_123",
|
|
content="Found memory",
|
|
importance=0.9,
|
|
metadata={},
|
|
)
|
|
]
|
|
|
|
mock_get_service.return_value = mock_service
|
|
yield mock_service
|
|
|
|
@pytest.fixture
|
|
def mcp_server(self):
|
|
"""Create a FastMCP server instance."""
|
|
return FastMCP("test-server")
|
|
|
|
@pytest.fixture
|
|
def permission_manager(self):
|
|
"""Create a permission manager with /tmp allowed."""
|
|
from hanzo_mcp.security.permissions import PermissionManager
|
|
|
|
pm = PermissionManager()
|
|
pm.add_allowed_path("/tmp")
|
|
return pm
|
|
|
|
def assert_tool_registration(self, tools, expected_count=9):
|
|
"""Assert that the expected number of memory tools are registered."""
|
|
assert len(tools) == expected_count
|
|
tool_types = {type(tool) for tool in tools}
|
|
expected_types = {
|
|
CreateMemoriesTool,
|
|
UpdateMemoriesTool,
|
|
DeleteMemoriesTool,
|
|
RecallMemoriesTool,
|
|
}
|
|
# Check that at least the core tools are present
|
|
assert expected_types.issubset(tool_types)
|
|
|
|
def create_memory_tool_params(self, **overrides):
|
|
"""Create standard parameters for memory tool calls."""
|
|
params = {
|
|
"user_id": "test_user",
|
|
"project_id": "test_project",
|
|
"content": "Test content",
|
|
"metadata": {"type": "test"},
|
|
"importance": 1.0,
|
|
}
|
|
params.update(overrides)
|
|
return params
|