765 lines
27 KiB
Python
765 lines
27 KiB
Python
"""
|
|
Unit tests for the Context system.
|
|
|
|
Test coverage:
|
|
1. Redis storage
|
|
2. File system persistence
|
|
3. Message management
|
|
4. Context manager
|
|
5. Data serialization and deserialization
|
|
6. Automatic summary functionality
|
|
7. Metadata management
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import tempfile
|
|
import shutil
|
|
import asyncio
|
|
import pytest
|
|
import redis
|
|
from typing import Dict, List, Optional, Any
|
|
from unittest.mock import Mock, AsyncMock, patch
|
|
from datetime import datetime, timedelta
|
|
|
|
# Import the modules under test.
|
|
import sys
|
|
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from context.context_manager import ContextManager
|
|
from context.context import RedisFileContextBackend, ContextBackend
|
|
from context.schemas import Message, SketchPadItem
|
|
from config.config import get_config
|
|
|
|
|
|
def _make_configured_backend(redis_host: str, redis_port: int, redis_db: int):
|
|
class ConfiguredRedisFileContextBackend(RedisFileContextBackend):
|
|
def __init__(
|
|
self,
|
|
context_id: str,
|
|
llm_interface=None,
|
|
max_history_length: int = 5,
|
|
auto_summarize_trigger: int = 1000000,
|
|
file_path: Optional[str] = None,
|
|
redis_host: str = redis_host,
|
|
redis_port: int = redis_port,
|
|
redis_db: int = redis_db,
|
|
):
|
|
super().__init__(
|
|
context_id=context_id,
|
|
llm_interface=llm_interface,
|
|
max_history_length=max_history_length,
|
|
auto_summarize_trigger=auto_summarize_trigger,
|
|
redis_host=redis_host,
|
|
redis_port=redis_port,
|
|
redis_db=redis_db,
|
|
file_path=file_path,
|
|
)
|
|
|
|
return ConfiguredRedisFileContextBackend
|
|
|
|
|
|
class TestContextSystem:
|
|
"""Context system integration tests."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def setup_and_teardown(self):
|
|
"""Set up and clean up around each test."""
|
|
# Set up the test environment.
|
|
self.test_context_dir = tempfile.mkdtemp(prefix="test_context_")
|
|
self.test_redis_db = 15 # Use a dedicated test database.
|
|
|
|
# Create test configuration.
|
|
self.original_config = None
|
|
if hasattr(get_config(), "CONTEXT_DIR"):
|
|
self.original_config = get_config().CONTEXT_DIR
|
|
|
|
# Point the configuration to the test directory.
|
|
config = get_config()
|
|
config.CONTEXT_DIR = self.test_context_dir
|
|
self.redis_host = config.REDIS_HOST
|
|
self.redis_port = int(config.REDIS_PORT)
|
|
self.backend_class = _make_configured_backend(
|
|
self.redis_host,
|
|
self.redis_port,
|
|
self.test_redis_db,
|
|
)
|
|
ContextManager._instance = None
|
|
|
|
# Create a Redis connection.
|
|
self.redis_client = redis.Redis(
|
|
host=self.redis_host,
|
|
port=self.redis_port,
|
|
db=self.test_redis_db,
|
|
decode_responses=True,
|
|
)
|
|
|
|
# Clear test data.
|
|
self.redis_client.flushdb()
|
|
|
|
yield
|
|
|
|
# Clean up the test environment.
|
|
self.redis_client.flushdb()
|
|
self.redis_client.close()
|
|
|
|
if os.path.exists(self.test_context_dir):
|
|
shutil.rmtree(self.test_context_dir)
|
|
ContextManager._instance = None
|
|
|
|
# Restore the original configuration.
|
|
if self.original_config:
|
|
config.CONTEXT_DIR = self.original_config
|
|
|
|
def test_redis_connection(self):
|
|
"""Test the Redis connection."""
|
|
assert self.redis_client.ping()
|
|
print("✓ Redis connection is healthy")
|
|
|
|
def test_context_backend_creation(self):
|
|
"""Test context backend creation."""
|
|
context_id = "test_context_001"
|
|
backend = RedisFileContextBackend(
|
|
context_id=context_id,
|
|
redis_host=self.redis_host,
|
|
redis_port=self.redis_port,
|
|
redis_db=self.test_redis_db,
|
|
file_path=os.path.join(self.test_context_dir, f"ctx_{context_id}.json"),
|
|
)
|
|
|
|
assert backend.context_id == context_id
|
|
assert backend.redis_client is not None
|
|
assert backend.file_path is not None
|
|
print("✓ Context backend created successfully")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_message_storage_and_retrieval(self):
|
|
"""Test message storage and retrieval."""
|
|
context_id = "test_context_002"
|
|
backend = RedisFileContextBackend(
|
|
context_id=context_id,
|
|
redis_host=self.redis_host,
|
|
redis_port=self.redis_port,
|
|
redis_db=self.test_redis_db,
|
|
file_path=os.path.join(self.test_context_dir, f"ctx_{context_id}.json"),
|
|
)
|
|
|
|
# Create test messages.
|
|
messages = [
|
|
Message(role="system", content="You are a helpful assistant"),
|
|
Message(role="user", content="Hello"),
|
|
Message(role="assistant", content="Hello! How can I help you?"),
|
|
Message(role="user", content="Please introduce Python"),
|
|
Message(role="assistant", content="Python is a high-level programming language..."),
|
|
]
|
|
|
|
# Store messages.
|
|
for message in messages:
|
|
await backend.store_message(message)
|
|
|
|
# Verify message count.
|
|
assert backend.get_message_count() == len(messages)
|
|
print(f"✓ Successfully stored {len(messages)} messages")
|
|
|
|
# Retrieve messages.
|
|
retrieved_messages = backend.retrieve_messages()
|
|
assert len(retrieved_messages) == len(messages)
|
|
|
|
# Verify message content.
|
|
for i, (original, retrieved) in enumerate(zip(messages, retrieved_messages)):
|
|
assert original.role == retrieved.role
|
|
assert original.content == retrieved.content
|
|
assert retrieved.timestamp is not None
|
|
|
|
print("✓ Message retrieval works correctly")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_persistence(self):
|
|
"""Test file persistence functionality."""
|
|
context_id = "test_context_003"
|
|
file_path = os.path.join(self.test_context_dir, f"ctx_{context_id}.json")
|
|
|
|
backend = RedisFileContextBackend(
|
|
context_id=context_id,
|
|
redis_host=self.redis_host,
|
|
redis_port=self.redis_port,
|
|
redis_db=self.test_redis_db,
|
|
file_path=file_path,
|
|
)
|
|
|
|
# Add some messages.
|
|
messages = [
|
|
Message(role="user", content="Test message 1"),
|
|
Message(role="assistant", content="Test reply 1"),
|
|
Message(role="user", content="Test message 2"),
|
|
Message(role="assistant", content="Test reply 2"),
|
|
]
|
|
|
|
for message in messages:
|
|
await backend.store_message(message)
|
|
|
|
# Persist to file.
|
|
success = await backend.persist()
|
|
assert success
|
|
assert os.path.exists(file_path)
|
|
|
|
# Verify file contents.
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
|
|
assert data["context_id"] == context_id
|
|
assert "messages" in data
|
|
assert "metadata" in data
|
|
assert len(data["messages"]) == len(messages)
|
|
|
|
print("✓ File persistence works correctly")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_restoration(self):
|
|
"""Test file restoration functionality."""
|
|
context_id = "test_context_004"
|
|
file_path = os.path.join(self.test_context_dir, f"ctx_{context_id}.json")
|
|
|
|
# Create the first backend and add data.
|
|
backend1 = RedisFileContextBackend(
|
|
context_id=context_id,
|
|
redis_host=self.redis_host,
|
|
redis_port=self.redis_port,
|
|
redis_db=self.test_redis_db,
|
|
file_path=file_path,
|
|
)
|
|
|
|
messages = [
|
|
Message(role="user", content="Persistence test message 1"),
|
|
Message(role="assistant", content="Persistence test reply 1"),
|
|
Message(role="user", content="Persistence test message 2"),
|
|
]
|
|
|
|
for message in messages:
|
|
await backend1.store_message(message)
|
|
|
|
# Persist data.
|
|
await backend1.persist()
|
|
|
|
# Create a new backend instance to simulate restart.
|
|
backend2 = RedisFileContextBackend(
|
|
context_id=context_id,
|
|
redis_host=self.redis_host,
|
|
redis_port=self.redis_port,
|
|
redis_db=self.test_redis_db,
|
|
file_path=file_path,
|
|
)
|
|
|
|
# Restore from file.
|
|
success = await backend2.restore()
|
|
assert success
|
|
|
|
# Verify restored data.
|
|
restored_messages = backend2.retrieve_messages()
|
|
assert len(restored_messages) == len(messages)
|
|
|
|
for i, (original, restored) in enumerate(zip(messages, restored_messages)):
|
|
assert original.role == restored.role
|
|
assert original.content == restored.content
|
|
|
|
print("✓ File restoration works correctly")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_context_manager_integration(self):
|
|
"""Test context manager integration functionality."""
|
|
# Create a context manager.
|
|
manager = ContextManager(backend_class=self.backend_class)
|
|
|
|
# Create a context.
|
|
context_id = "test_context_005"
|
|
context = manager.create_context(context_id=context_id, max_history_length=10)
|
|
|
|
assert context is not None
|
|
assert context.context_id == context_id
|
|
|
|
# Add a message through the convenience interface.
|
|
success = await manager.add_message(
|
|
context_id=context_id,
|
|
message=Message(role="user", content="Message added through the manager"),
|
|
)
|
|
assert success
|
|
|
|
# Get history.
|
|
history = manager.get_history(context_id)
|
|
assert len(history) == 1
|
|
assert history[0].role == "user"
|
|
assert history[0].content == "Message added through the manager"
|
|
|
|
# Test retrieving the context.
|
|
retrieved_context = manager.get_context(context_id)
|
|
assert retrieved_context is not None
|
|
assert retrieved_context.context_id == context_id
|
|
|
|
print("✓ Context manager integration works correctly")
|
|
|
|
def test_metadata_management(self):
|
|
"""Test metadata management."""
|
|
context_id = "test_context_006"
|
|
backend = RedisFileContextBackend(
|
|
context_id=context_id,
|
|
redis_host=self.redis_host,
|
|
redis_port=self.redis_port,
|
|
redis_db=self.test_redis_db,
|
|
file_path=os.path.join(self.test_context_dir, f"ctx_{context_id}.json"),
|
|
)
|
|
|
|
# Get initial metadata.
|
|
initial_metadata = backend.get_metadata()
|
|
assert initial_metadata["context_id"] == context_id
|
|
assert "start_time" in initial_metadata
|
|
assert "last_activity" in initial_metadata
|
|
|
|
# Update metadata.
|
|
new_metadata = {"custom_field": "custom_value", "test_count": 42}
|
|
backend.update_metadata(new_metadata)
|
|
|
|
# Verify the update.
|
|
updated_metadata = backend.get_metadata()
|
|
assert updated_metadata["custom_field"] == "custom_value"
|
|
assert updated_metadata["test_count"] == 42
|
|
|
|
print("✓ Metadata management works correctly")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_message_search(self):
|
|
"""Test message search functionality."""
|
|
context_id = "test_context_007"
|
|
backend = RedisFileContextBackend(
|
|
context_id=context_id,
|
|
redis_host=self.redis_host,
|
|
redis_port=self.redis_port,
|
|
redis_db=self.test_redis_db,
|
|
file_path=os.path.join(self.test_context_dir, f"ctx_{context_id}.json"),
|
|
)
|
|
|
|
# Add messages containing specific keywords.
|
|
messages = [
|
|
Message(role="user", content="I want to learn Python programming"),
|
|
Message(role="assistant", content="Python is a great programming language"),
|
|
Message(role="user", content="Please tell me about machine learning"),
|
|
Message(role="assistant", content="Machine learning is a branch of artificial intelligence"),
|
|
Message(role="user", content="Python is commonly used in machine learning"),
|
|
]
|
|
|
|
for message in messages:
|
|
await backend.store_message(message)
|
|
|
|
# Search for messages containing "Python".
|
|
python_results = backend.search_messages("Python", limit=10)
|
|
assert len(python_results) == 3
|
|
|
|
# Search for messages containing "machine learning".
|
|
ml_results = backend.search_messages("machine learning", limit=10)
|
|
assert len(ml_results) == 3
|
|
|
|
# Search for a nonexistent keyword.
|
|
empty_results = backend.search_messages("nonexistent keyword", limit=10)
|
|
assert len(empty_results) == 0
|
|
|
|
print("✓ Message search works correctly")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_message_limit_management(self):
|
|
"""Test message count limit management."""
|
|
context_id = "test_context_008"
|
|
max_history_length = 3
|
|
|
|
backend = RedisFileContextBackend(
|
|
context_id=context_id,
|
|
redis_host=self.redis_host,
|
|
redis_port=self.redis_port,
|
|
redis_db=self.test_redis_db,
|
|
file_path=os.path.join(self.test_context_dir, f"ctx_{context_id}.json"),
|
|
max_history_length=max_history_length,
|
|
)
|
|
|
|
# Add more messages than the limit allows.
|
|
for i in range(5):
|
|
message = Message(role="user", content=f"Message {i}")
|
|
await backend.store_message(message)
|
|
|
|
# Verify that only the latest messages are retained.
|
|
messages = backend.retrieve_messages()
|
|
assert len(messages) == max_history_length
|
|
|
|
# Verify that the retained messages are the latest ones.
|
|
expected_contents = ["Message 2", "Message 3", "Message 4"]
|
|
for i, message in enumerate(messages):
|
|
assert message.content == expected_contents[i]
|
|
|
|
print("✓ Message count limit management works correctly")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_context_serialization(self):
|
|
"""Test context serialization functionality."""
|
|
context_id = "test_context_009"
|
|
backend = RedisFileContextBackend(
|
|
context_id=context_id,
|
|
redis_host=self.redis_host,
|
|
redis_port=self.redis_port,
|
|
redis_db=self.test_redis_db,
|
|
file_path=os.path.join(self.test_context_dir, f"ctx_{context_id}.json"),
|
|
)
|
|
|
|
# Add messages and metadata.
|
|
messages = [
|
|
Message(role="user", content="Serialization test message"),
|
|
Message(role="assistant", content="Serialization test reply"),
|
|
]
|
|
|
|
for message in messages:
|
|
await backend.store_message(message)
|
|
|
|
backend.update_summary("This is a test conversation")
|
|
backend.update_metadata({"test_key": "test_value"})
|
|
|
|
# Serialize.
|
|
serialized_data = backend.serialize()
|
|
|
|
# Verify serialized data.
|
|
assert serialized_data["context_id"] == context_id
|
|
assert len(serialized_data["messages"]) == 2
|
|
assert serialized_data["summary"] == "This is a test conversation"
|
|
assert "serialization_timestamp" in serialized_data
|
|
|
|
# Create a new backend and deserialize.
|
|
new_backend = RedisFileContextBackend(
|
|
context_id="new_context",
|
|
redis_host=self.redis_host,
|
|
redis_port=self.redis_port,
|
|
redis_db=self.test_redis_db,
|
|
file_path=os.path.join(self.test_context_dir, "ctx_new_context.json"),
|
|
)
|
|
|
|
new_backend.deserialize(serialized_data)
|
|
|
|
# Verify deserialization result.
|
|
restored_messages = new_backend.retrieve_messages()
|
|
assert len(restored_messages) == 2
|
|
assert new_backend.get_summary() == "This is a test conversation"
|
|
|
|
print("✓ Context serialization works correctly")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_context_manager_list_and_delete(self):
|
|
"""Test context manager list and delete functionality."""
|
|
manager = ContextManager(backend_class=self.backend_class)
|
|
|
|
# Create multiple contexts.
|
|
context_ids = ["test_ctx_001", "test_ctx_002", "test_ctx_003"]
|
|
|
|
for context_id in context_ids:
|
|
context = manager.create_context(context_id=context_id)
|
|
await manager.add_message(
|
|
context_id,
|
|
Message(role="user", content=f"Test message {context_id}"),
|
|
)
|
|
|
|
# List all contexts.
|
|
contexts = manager.list_contexts()
|
|
assert len(contexts) >= len(context_ids)
|
|
|
|
# Verify that the created contexts are all listed.
|
|
found_contexts = [ctx["context_id"] for ctx in contexts]
|
|
for context_id in context_ids:
|
|
assert context_id in found_contexts
|
|
|
|
# Delete one context.
|
|
delete_success = manager.delete_context(context_ids[0])
|
|
assert delete_success
|
|
|
|
# Verify it cannot be retrieved after deletion.
|
|
deleted_context = manager.get_context(context_ids[0])
|
|
assert deleted_context is None
|
|
|
|
print("✓ Context manager list and delete functionality works correctly")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_persistence_verification(self):
|
|
"""Verify data persistence in Redis."""
|
|
context_id = "test_redis_persistence"
|
|
|
|
# Create a backend and add data.
|
|
backend = RedisFileContextBackend(
|
|
context_id=context_id,
|
|
redis_host=self.redis_host,
|
|
redis_port=self.redis_port,
|
|
redis_db=self.test_redis_db,
|
|
file_path=os.path.join(self.test_context_dir, f"ctx_{context_id}.json"),
|
|
)
|
|
|
|
# Add messages.
|
|
messages = [
|
|
Message(role="user", content="Redis persistence test message 1"),
|
|
Message(role="assistant", content="Redis persistence test reply 1"),
|
|
Message(role="user", content="Redis persistence test message 2"),
|
|
]
|
|
|
|
for message in messages:
|
|
await backend.store_message(message)
|
|
|
|
# Verify data in Redis.
|
|
messages_key = f"context:{context_id}:messages"
|
|
metadata_key = f"context:{context_id}:metadata"
|
|
summary_key = f"context:{context_id}:summary"
|
|
|
|
# Check message data.
|
|
redis_messages = self.redis_client.lrange(messages_key, 0, -1)
|
|
assert len(redis_messages) == 3
|
|
|
|
# Check metadata.
|
|
redis_metadata = self.redis_client.get(metadata_key)
|
|
assert redis_metadata is not None
|
|
|
|
# Verify message content.
|
|
for i, message_json in enumerate(redis_messages):
|
|
message_data = json.loads(message_json)
|
|
expected_message = messages[-(i + 1)]
|
|
assert message_data["role"] == expected_message.role
|
|
assert message_data["content"] == expected_message.content
|
|
|
|
print("✓ Redis data persistence verification passed")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_system_persistence_verification(self):
|
|
"""Verify data persistence in the file system."""
|
|
context_id = "test_file_persistence"
|
|
file_path = os.path.join(self.test_context_dir, f"ctx_{context_id}.json")
|
|
|
|
backend = RedisFileContextBackend(
|
|
context_id=context_id,
|
|
redis_host=self.redis_host,
|
|
redis_port=self.redis_port,
|
|
redis_db=self.test_redis_db,
|
|
file_path=file_path,
|
|
)
|
|
|
|
# Add messages and metadata.
|
|
messages = [
|
|
Message(role="user", content="File persistence test message 1"),
|
|
Message(role="assistant", content="File persistence test reply 1"),
|
|
Message(role="user", content="File persistence test message 2"),
|
|
]
|
|
|
|
for message in messages:
|
|
await backend.store_message(message)
|
|
|
|
backend.update_summary("File persistence test summary")
|
|
backend.update_metadata({"file_test": "file_value"})
|
|
|
|
# Persist to file.
|
|
await backend.persist()
|
|
|
|
# Verify file exists.
|
|
assert os.path.exists(file_path)
|
|
|
|
# Read and verify file contents.
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
file_data = json.load(f)
|
|
|
|
# Verify file structure.
|
|
assert file_data["context_id"] == context_id
|
|
assert len(file_data["messages"]) == 3
|
|
assert file_data["summary"] == "File persistence test summary"
|
|
assert file_data["metadata"]["file_test"] == "file_value"
|
|
|
|
# Verify message content.
|
|
for i, message_data in enumerate(file_data["messages"]):
|
|
assert message_data["role"] == messages[i].role
|
|
assert message_data["content"] == messages[i].content
|
|
|
|
print("✓ File system data persistence verification passed")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_access(self):
|
|
"""Test concurrent access."""
|
|
import threading
|
|
import time
|
|
|
|
context_id = "test_concurrent"
|
|
backend = RedisFileContextBackend(
|
|
context_id=context_id,
|
|
redis_host=self.redis_host,
|
|
redis_port=self.redis_port,
|
|
redis_db=self.test_redis_db,
|
|
file_path=os.path.join(self.test_context_dir, f"ctx_{context_id}.json"),
|
|
max_history_length=20,
|
|
)
|
|
|
|
# Add messages concurrently.
|
|
def add_messages(thread_id: int, count: int):
|
|
for i in range(count):
|
|
message = Message(role="user", content=f"Thread {thread_id} message {i}")
|
|
asyncio.run(backend.store_message(message))
|
|
time.sleep(0.01) # Small delay.
|
|
|
|
# Create multiple threads.
|
|
threads = []
|
|
for i in range(3):
|
|
thread = threading.Thread(target=add_messages, args=(i, 5))
|
|
threads.append(thread)
|
|
thread.start()
|
|
|
|
# Wait for all threads to finish.
|
|
for thread in threads:
|
|
thread.join()
|
|
|
|
# Verify that all messages were added.
|
|
messages = backend.retrieve_messages()
|
|
assert len(messages) == 15 # 3 threads * 5 messages.
|
|
|
|
print("✓ Concurrent access test passed")
|
|
|
|
def test_error_handling(self):
|
|
"""Test error handling."""
|
|
# Test an invalid Redis connection.
|
|
try:
|
|
invalid_backend = RedisFileContextBackend(
|
|
context_id="test_error",
|
|
redis_host="invalid_host",
|
|
redis_port=9999,
|
|
redis_db=0,
|
|
file_path=os.path.join(self.test_context_dir, "ctx_test_error.json"),
|
|
)
|
|
# If the connection fails, it should raise an exception.
|
|
invalid_backend.redis_client.ping()
|
|
except Exception as e:
|
|
print(f"✓ Expected Redis connection error: {type(e).__name__}")
|
|
|
|
# Test an invalid file path.
|
|
backend = RedisFileContextBackend(
|
|
context_id="test_error_file",
|
|
redis_host=self.redis_host,
|
|
redis_port=self.redis_port,
|
|
redis_db=self.test_redis_db,
|
|
file_path="/invalid/path/ctx_test.json",
|
|
)
|
|
|
|
# Try to persist to an invalid path.
|
|
async def test_invalid_persist():
|
|
return await backend.persist()
|
|
|
|
result = asyncio.run(test_invalid_persist())
|
|
assert not result # Should fail.
|
|
|
|
print("✓ Error handling test passed")
|
|
|
|
|
|
class TestContextManagerAdvanced:
|
|
"""Advanced ContextManager feature tests."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def setup(self):
|
|
"""Set up the test environment."""
|
|
self.test_context_dir = tempfile.mkdtemp(prefix="test_manager_")
|
|
self.test_redis_db = 14
|
|
|
|
# Modify configuration.
|
|
config = get_config()
|
|
config.CONTEXT_DIR = self.test_context_dir
|
|
self.redis_host = config.REDIS_HOST
|
|
self.redis_port = int(config.REDIS_PORT)
|
|
self.backend_class = _make_configured_backend(
|
|
self.redis_host,
|
|
self.redis_port,
|
|
self.test_redis_db,
|
|
)
|
|
ContextManager._instance = None
|
|
|
|
# Create a Redis connection.
|
|
self.redis_client = redis.Redis(
|
|
host=self.redis_host,
|
|
port=self.redis_port,
|
|
db=self.test_redis_db,
|
|
decode_responses=True,
|
|
)
|
|
self.redis_client.flushdb()
|
|
|
|
yield
|
|
|
|
# Clean up.
|
|
self.redis_client.flushdb()
|
|
self.redis_client.close()
|
|
shutil.rmtree(self.test_context_dir)
|
|
ContextManager._instance = None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_context_manager_singleton(self):
|
|
"""Test the ContextManager singleton pattern."""
|
|
manager1 = ContextManager(backend_class=self.backend_class)
|
|
manager2 = ContextManager(backend_class=self.backend_class)
|
|
|
|
assert manager1 is manager2
|
|
print("✓ ContextManager singleton pattern works correctly")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_context_manager_bulk_operations(self):
|
|
"""Test ContextManager batch operations."""
|
|
manager = ContextManager(backend_class=self.backend_class)
|
|
|
|
# Create multiple contexts.
|
|
context_ids = ["bulk_test_001", "bulk_test_002", "bulk_test_003"]
|
|
contexts = []
|
|
|
|
for context_id in context_ids:
|
|
context = manager.create_context(context_id=context_id)
|
|
contexts.append(context)
|
|
|
|
# Add a message.
|
|
await manager.add_message(
|
|
context_id,
|
|
Message(role="user", content=f"Batch test message {context_id}"),
|
|
)
|
|
|
|
# Save all contexts.
|
|
saved_count = await manager.save_all_contexts()
|
|
assert saved_count == len(context_ids)
|
|
|
|
# Verify files exist.
|
|
for context_id in context_ids:
|
|
file_path = os.path.join(self.test_context_dir, f"ctx_{context_id}.json")
|
|
assert os.path.exists(file_path)
|
|
|
|
print("✓ ContextManager batch operations work correctly")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_context_manager_cleanup(self):
|
|
"""Test ContextManager cleanup functionality."""
|
|
manager = ContextManager(backend_class=self.backend_class)
|
|
|
|
# Create a context and add a message.
|
|
context_id = "cleanup_test"
|
|
context = manager.create_context(context_id=context_id)
|
|
|
|
await manager.add_message(
|
|
context_id,
|
|
Message(role="user", content="Cleanup test message"),
|
|
)
|
|
|
|
# Simulate long-term inactivity by modifying last_activity in metadata.
|
|
context.update_metadata(
|
|
{"last_activity": (datetime.now() - timedelta(hours=2)).isoformat()}
|
|
)
|
|
|
|
# Run cleanup with a short inactivity threshold.
|
|
cleaned_count = await manager.cleanup_inactive_contexts(max_inactive_time=1)
|
|
assert cleaned_count == 1
|
|
|
|
# Verify the context was removed from the active cache but can still be restored from persistent storage on demand.
|
|
assert context_id not in manager._active_contexts
|
|
retrieved_context = manager.get_context(context_id)
|
|
assert retrieved_context is not None
|
|
|
|
print("✓ ContextManager cleanup functionality works correctly")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Run tests.
|
|
pytest.main([__file__, "-v", "-s"])
|