Jerry
451 lines
16 KiB
Python
451 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Context system basic feature validation script.
|
|
|
|
Validates core features directly without relying on pytest:
|
|
1. Redis storage
|
|
2. File persistence
|
|
3. Message management
|
|
4. Context manager
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import json
|
|
import tempfile
|
|
import shutil
|
|
import asyncio
|
|
import redis
|
|
from typing import Dict, List, Optional, Any
|
|
from datetime import datetime, timedelta
|
|
|
|
# Add the project root directory to the Python path.
|
|
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
sys.path.insert(0, project_root)
|
|
|
|
from context.context_manager import ContextManager
|
|
from context.context import RedisFileContextBackend
|
|
from context.schemas import Message
|
|
from config.config import get_config
|
|
|
|
|
|
class ContextSystemTester:
|
|
"""Context system tester."""
|
|
|
|
def __init__(self):
|
|
"""Initialize the tester."""
|
|
self.test_context_dir = tempfile.mkdtemp(prefix="test_context_")
|
|
self.test_redis_db = 13 # Use a dedicated test database.
|
|
|
|
# Point the configuration to the test directory.
|
|
config = get_config()
|
|
config.CONTEXT_DIR = self.test_context_dir
|
|
|
|
# Create a Redis connection.
|
|
self.redis_client = redis.Redis(db=self.test_redis_db, decode_responses=True)
|
|
self.redis_client.flushdb()
|
|
|
|
self.test_results = []
|
|
|
|
def cleanup(self):
|
|
"""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)
|
|
|
|
def log_test(self, test_name: str, success: bool, message: str = ""):
|
|
"""Record a test result."""
|
|
status = "✓ PASS" if success else "✗ FAIL"
|
|
print(f"{status} {test_name}: {message}")
|
|
self.test_results.append({
|
|
"test": test_name,
|
|
"success": success,
|
|
"message": message
|
|
})
|
|
|
|
def test_redis_connection(self) -> bool:
|
|
"""Test the Redis connection."""
|
|
try:
|
|
assert self.redis_client.ping()
|
|
self.log_test("Redis connection", True, "Redis service is healthy")
|
|
return True
|
|
except Exception as e:
|
|
self.log_test("Redis connection", False, f"Redis connection failed: {e}")
|
|
return False
|
|
|
|
async def test_message_storage(self) -> bool:
|
|
"""Test message storage functionality."""
|
|
try:
|
|
context_id = "test_message_storage"
|
|
backend = RedisFileContextBackend(
|
|
context_id=context_id,
|
|
redis_host="localhost",
|
|
redis_port=6379,
|
|
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?")
|
|
]
|
|
|
|
# Store messages.
|
|
for message in messages:
|
|
await backend.store_message(message)
|
|
|
|
# Verify message count.
|
|
assert backend.get_message_count() == len(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
|
|
|
|
self.log_test("Message storage", True, f"Successfully stored and retrieved {len(messages)} messages")
|
|
return True
|
|
|
|
except Exception as e:
|
|
self.log_test("Message storage", False, f"Message storage test failed: {e}")
|
|
return False
|
|
|
|
async def test_file_persistence(self) -> bool:
|
|
"""Test file persistence functionality."""
|
|
try:
|
|
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="localhost",
|
|
redis_port=6379,
|
|
redis_db=self.test_redis_db,
|
|
file_path=file_path
|
|
)
|
|
|
|
# Add messages.
|
|
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 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 len(data['messages']) == len(messages)
|
|
|
|
self.log_test("File persistence", True, "Successfully persisted to the file system")
|
|
return True
|
|
|
|
except Exception as e:
|
|
self.log_test("File persistence", False, f"File persistence test failed: {e}")
|
|
return False
|
|
|
|
async def test_file_restoration(self) -> bool:
|
|
"""Test file restoration functionality."""
|
|
try:
|
|
context_id = "test_file_restoration"
|
|
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="localhost",
|
|
redis_port=6379,
|
|
redis_db=self.test_redis_db,
|
|
file_path=file_path
|
|
)
|
|
|
|
messages = [
|
|
Message(role="user", content="Restoration test message 1"),
|
|
Message(role="assistant", content="Restoration test reply 1"),
|
|
Message(role="user", content="Restoration test message 2")
|
|
]
|
|
|
|
for message in messages:
|
|
await backend1.store_message(message)
|
|
|
|
# Persist data.
|
|
await backend1.persist()
|
|
|
|
# Clear Redis data to simulate a restart.
|
|
self.redis_client.flushdb()
|
|
|
|
# Create a new backend instance to simulate restart.
|
|
backend2 = RedisFileContextBackend(
|
|
context_id=context_id,
|
|
redis_host="localhost",
|
|
redis_port=6379,
|
|
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
|
|
|
|
self.log_test("File restoration", True, "Successfully restored data from file")
|
|
return True
|
|
|
|
except Exception as e:
|
|
self.log_test("File restoration", False, f"File restoration test failed: {e}")
|
|
return False
|
|
|
|
async def test_context_manager(self) -> bool:
|
|
"""Test the context manager."""
|
|
try:
|
|
manager = ContextManager(backend_class=RedisFileContextBackend)
|
|
|
|
# Create a context.
|
|
context_id = "test_manager_context"
|
|
context = manager.create_context(context_id=context_id)
|
|
|
|
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,
|
|
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
|
|
|
|
self.log_test("Context manager", True, "Context manager works correctly")
|
|
return True
|
|
|
|
except Exception as e:
|
|
self.log_test("Context manager", False, f"Context manager test failed: {e}")
|
|
return False
|
|
|
|
async def test_redis_persistence_verification(self) -> bool:
|
|
"""Verify data persistence in Redis."""
|
|
try:
|
|
context_id = "test_redis_verification"
|
|
|
|
backend = RedisFileContextBackend(
|
|
context_id=context_id,
|
|
redis_host="localhost",
|
|
redis_port=6379,
|
|
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 verification message 1"),
|
|
Message(role="assistant", content="Redis verification reply 1"),
|
|
Message(role="user", content="Redis verification 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"
|
|
|
|
# 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)
|
|
assert message_data['role'] == messages[i].role
|
|
assert message_data['content'] == messages[i].content
|
|
|
|
self.log_test("Redis persistence verification", True, "Redis data storage works correctly")
|
|
return True
|
|
|
|
except Exception as e:
|
|
self.log_test("Redis persistence verification", False, f"Redis persistence verification failed: {e}")
|
|
return False
|
|
|
|
async def test_concurrent_access(self) -> bool:
|
|
"""Test concurrent access."""
|
|
try:
|
|
import threading
|
|
import time
|
|
|
|
context_id = "test_concurrent"
|
|
backend = RedisFileContextBackend(
|
|
context_id=context_id,
|
|
redis_host="localhost",
|
|
redis_port=6379,
|
|
redis_db=self.test_redis_db,
|
|
file_path=os.path.join(self.test_context_dir, f"ctx_{context_id}.json")
|
|
)
|
|
|
|
# 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}")
|
|
# Use a new event loop.
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(backend.store_message(message))
|
|
finally:
|
|
loop.close()
|
|
time.sleep(0.01)
|
|
|
|
# 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.
|
|
|
|
self.log_test("Concurrent access", True, "Concurrent access is safe")
|
|
return True
|
|
|
|
except Exception as e:
|
|
self.log_test("Concurrent access", False, f"Concurrent access test failed: {e}")
|
|
return False
|
|
|
|
async def run_all_tests(self) -> bool:
|
|
"""Run all tests."""
|
|
print("=" * 60)
|
|
print("Starting Context system tests")
|
|
print("=" * 60)
|
|
|
|
tests = [
|
|
("Redis connection", self.test_redis_connection),
|
|
("Message storage", self.test_message_storage),
|
|
("File persistence", self.test_file_persistence),
|
|
("File restoration", self.test_file_restoration),
|
|
("Context manager", self.test_context_manager),
|
|
("Redis persistence verification", self.test_redis_persistence_verification),
|
|
("Concurrent access", self.test_concurrent_access),
|
|
]
|
|
|
|
all_passed = True
|
|
|
|
for test_name, test_func in tests:
|
|
print(f"\nRunning test: {test_name}")
|
|
print("-" * 40)
|
|
|
|
try:
|
|
if asyncio.iscoroutinefunction(test_func):
|
|
result = await test_func()
|
|
else:
|
|
result = test_func()
|
|
|
|
if not result:
|
|
all_passed = False
|
|
|
|
except Exception as e:
|
|
self.log_test(test_name, False, f"Test exception: {e}")
|
|
all_passed = False
|
|
|
|
# Output test summary.
|
|
print("\n" + "=" * 60)
|
|
print("Test Result Summary")
|
|
print("=" * 60)
|
|
|
|
passed_count = sum(1 for result in self.test_results if result["success"])
|
|
total_count = len(self.test_results)
|
|
|
|
for result in self.test_results:
|
|
status = "✓" if result["success"] else "✗"
|
|
print(f"{status} {result['test']}: {result['message']}")
|
|
|
|
print(f"\nTotal: {passed_count}/{total_count} tests passed")
|
|
|
|
if all_passed:
|
|
print("\n🎉 All tests passed!")
|
|
print("Context system feature verification succeeded:")
|
|
print(" ✓ Redis storage works correctly")
|
|
print(" ✓ File system persistence works correctly")
|
|
print(" ✓ Message management works correctly")
|
|
print(" ✓ Context manager works correctly")
|
|
print(" ✓ Concurrent access is safe")
|
|
else:
|
|
print("\n❌ Some tests failed")
|
|
|
|
return all_passed
|
|
|
|
|
|
async def main():
|
|
"""Main function."""
|
|
print("Context System Basic Feature Validation")
|
|
print("=" * 40)
|
|
|
|
# Check Redis connection.
|
|
try:
|
|
r = redis.Redis(host="localhost", port=6379, decode_responses=True)
|
|
r.ping()
|
|
r.close()
|
|
print("✓ Redis service is available")
|
|
except Exception as e:
|
|
print(f"❌ Redis service is unavailable: {e}")
|
|
print("Please ensure the Redis service is running")
|
|
print("Startup commands:")
|
|
print(" macOS: brew services start redis")
|
|
print(" Linux: sudo systemctl start redis")
|
|
print(" Windows: redis-server")
|
|
return 1
|
|
|
|
# Run tests.
|
|
tester = ContextSystemTester()
|
|
try:
|
|
success = await tester.run_all_tests()
|
|
return 0 if success else 1
|
|
finally:
|
|
tester.cleanup()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
exit(asyncio.run(main()))
|