cadSet/CADDesigner-Code/test/test_context_simple.py

#!/usr/bin/env python3
"""
Context system simplified test script.

Focuses on validating core features:
1. Redis storage
2. File persistence
3. Message management
"""

import os
import sys
import json
import tempfile
import shutil
import asyncio
import redis
import traceback
from typing import Dict, List, Optional, Any

# 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 import RedisFileContextBackend
from context.schemas import Message


class SimpleContextTester:
    """Simplified Context system tester."""
    
    def __init__(self):
        """Initialize the tester."""
        self.test_context_dir = tempfile.mkdtemp(prefix="test_context_")
        self.test_redis_db = 12  # Use a dedicated test database.
        
        # 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_basic_message_storage(self) -> bool:
        """Test basic message storage functionality."""
        try:
            context_id = "test_basic_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 a test message.
            message = Message(role="user", content="Test message")
            
            # Store the message.
            await backend.store_message(message)
            
            # Verify message count.
            count = backend.get_message_count()
            assert count == 1, f"Expected 1 message, got {count}"
            
            # Retrieve messages.
            messages = backend.retrieve_messages()
            assert len(messages) == 1, f"Expected 1 message, got {len(messages)}"
            
            # Verify message content.
            retrieved_message = messages[0]
            assert retrieved_message.role == "user"
            assert retrieved_message.content == "Test message"
            
            self.log_test("Basic message storage", True, "Successfully stored and retrieved the message")
            return True
            
        except Exception as e:
            error_msg = f"Basic message storage test failed: {e}\n{traceback.format_exc()}"
            self.log_test("Basic message storage", False, error_msg)
            return False
    
    async def test_file_persistence_simple(self) -> bool:
        """Test simple file persistence functionality."""
        try:
            context_id = "test_file_simple"
            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 one message.
            message = Message(role="user", content="Persistence test message")
            await backend.store_message(message)
            
            # Persist to file.
            success = await backend.persist()
            assert success, "Persistence failed"
            assert os.path.exists(file_path), f"File does not exist: {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']) == 1
            assert data['messages'][0]['content'] == "Persistence test message"
            
            self.log_test("File persistence", True, "Successfully persisted to the file system")
            return True
            
        except Exception as e:
            error_msg = f"File persistence test failed: {e}\n{traceback.format_exc()}"
            self.log_test("File persistence", False, error_msg)
            return False
    
    async def test_redis_data_verification(self) -> bool:
        """Verify data in Redis."""
        try:
            context_id = "test_redis_verify"
            
            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 a message.
            message = Message(role="user", content="Redis verification message")
            await backend.store_message(message)
            
            # Ensure metadata is updated in Redis.
            backend.update_metadata({"test_verification": "true"})
            
            # 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) == 1, f"Redis should contain 1 message, got {len(redis_messages)}"
            
            # Check metadata.
            redis_metadata = self.redis_client.get(metadata_key)
            assert redis_metadata is not None, "Redis should contain metadata"
            
            # Verify message content.
            message_data = json.loads(redis_messages[0])
            assert message_data['role'] == "user"
            assert message_data['content'] == "Redis verification message"
            
            # Verify metadata contents.
            metadata_data = json.loads(redis_metadata)
            assert metadata_data['context_id'] == context_id
            assert metadata_data['test_verification'] == "true"
            
            self.log_test("Redis data verification", True, "Redis data storage works correctly")
            return True
            
        except Exception as e:
            error_msg = f"Redis data verification failed: {e}\n{traceback.format_exc()}"
            self.log_test("Redis data verification", False, error_msg)
            return False
    
    async def test_metadata_management(self) -> bool:
        """Test metadata management."""
        try:
            context_id = "test_metadata"
            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")
            )
            
            # 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
            
            self.log_test("Metadata management", True, "Metadata management works correctly")
            return True
            
        except Exception as e:
            error_msg = f"Metadata management test failed: {e}\n{traceback.format_exc()}"
            self.log_test("Metadata management", False, error_msg)
            return False
    
    async def test_message_search(self) -> bool:
        """Test message search functionality."""
        try:
            context_id = "test_search"
            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 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")
            ]
            
            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) == 2, f"Expected 2 messages containing Python, got {len(python_results)}"
            
            # Search for messages containing "machine learning".
            ml_results = backend.search_messages("machine learning", limit=10)
            assert len(ml_results) == 1, f"Expected 1 message containing machine learning, got {len(ml_results)}"
            
            self.log_test("Message search", True, "Message search works correctly")
            return True
            
        except Exception as e:
            error_msg = f"Message search test failed: {e}\n{traceback.format_exc()}"
            self.log_test("Message search", False, error_msg)
            return False
    
    async def run_all_tests(self) -> bool:
        """Run all tests."""
        print("=" * 60)
        print("Starting simplified Context system tests")
        print("=" * 60)
        
        tests = [
            ("Redis connection", self.test_redis_connection),
            ("Basic message storage", self.test_basic_message_storage),
            ("File persistence", self.test_file_persistence_simple),
            ("Redis data verification", self.test_redis_data_verification),
            ("Metadata management", self.test_metadata_management),
            ("Message search", self.test_message_search),
        ]
        
        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:
                error_msg = f"Test exception: {e}\n{traceback.format_exc()}"
                self.log_test(test_name, False, error_msg)
                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 core feature verification succeeded:")
            print("  ✓ Redis storage works correctly")
            print("  ✓ File system persistence works correctly")
            print("  ✓ Message management works correctly")
            print("  ✓ Metadata management works correctly")
            print("  ✓ Message search works correctly")
        else:
            print("\n❌ Some tests failed")
        
        return all_passed


async def main():
    """Main function."""
    print("Context System Simplified 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 = SimpleContextTester()
    try:
        success = await tester.run_all_tests()
        return 0 if success else 1
    finally:
        tester.cleanup()


if __name__ == "__main__":
    exit(asyncio.run(main()))