#!/usr/bin/env python3 """ Test the new SketchManager implementation. """ import asyncio import json import os import tempfile import shutil from typing import Optional from context.sketch_manager import get_sketch_manager, SketchManager from context.sketch_pad import SketchPadBackend, RedisFileSketchPadBackend from context.schemas import SketchPadListItem, SketchPadStatistics class MockSketchPadBackend(SketchPadBackend): """Mock SketchPad backend for tests.""" def __init__(self, sketch_pad_id: str, file_path: Optional[str] = None): self.sketch_pad_id = sketch_pad_id self.file_path = file_path self._storage = {} self._lock = asyncio.Lock() async def set_item( self, key: str, value: any, ttl: Optional[int] = None, summary: Optional[str] = None, tags: Optional[set] = None, ) -> str: from context.schemas import SketchPadItem from datetime import datetime item = SketchPadItem( value=value, timestamp=datetime.now(), summary=summary, tags=tags or set(), ) self._storage[key] = item return key def get_item(self, key: str): return self._storage.get(key) def get_value(self, key: str): item = self.get_item(key) return item.value if item else None def search_by_tags(self, tags: set, match_all: bool = False): results = [] for key, item in self._storage.items(): if match_all: if tags.issubset(item.tags): results.append((key, item)) else: if tags.intersection(item.tags): results.append((key, item)) return results def search_by_content(self, query: str, limit: int = 5): results = [] query_lower = query.lower() for key, item in self._storage.items(): content = str(item.value) + (item.summary or "") if query_lower in content.lower(): results.append((key, item)) if len(results) >= limit: break return results def delete(self, key: str) -> bool: if key in self._storage: del self._storage[key] return True return False def exists(self, key: str) -> bool: return key in self._storage def keys(self, pattern: Optional[str] = None) -> list[str]: if pattern: return [k for k in self._storage.keys() if pattern in k] return list(self._storage.keys()) def clear(self) -> None: self._storage.clear() def serialize(self) -> dict[str, any]: return { "sketch_pad_id": self.sketch_pad_id, "items": {k: v.model_dump() for k, v in self._storage.items()}, "serialization_timestamp": "2024-01-01T00:00:00", } def deserialize(self, data: dict[str, any]) -> None: from context.schemas import SketchPadItem if "items" in data: for key, item_data in data["items"].items(): try: item = SketchPadItem(**item_data) self._storage[key] = item except Exception as e: print(f"Warning: Failed to deserialize item {key}: {e}") def persist(self) -> None: if self.file_path: data = self.serialize() os.makedirs(os.path.dirname(self.file_path), exist_ok=True) with open(self.file_path, "w", encoding="utf-8") as f: json.dump(data, f, ensure_ascii=False, indent=2) def restore(self) -> None: if self.file_path and os.path.exists(self.file_path): with open(self.file_path, "r", encoding="utf-8") as f: data = json.load(f) self.deserialize(data) def get_statistics(self) -> SketchPadStatistics: return SketchPadStatistics( total_items=len(self._storage), max_items=len(self._storage), items_with_summary=sum( 1 for item in self._storage.values() if item.summary ), total_accesses=sum(item.access_count for item in self._storage.values()), popular_tags={}, content_types={}, avg_access_per_item=0.0, memory_usage_percent=0.0, ) def list_items(self, include_value: bool = False) -> list[SketchPadListItem]: return [ SketchPadListItem( key=key, summary=item.summary, timestamp=item.timestamp.isoformat(), tags=sorted(item.tags), content_type=item.content_type, access_count=item.access_count, content_hash=item.content_hash, value=item.value if include_value else None, ) for key, item in self._storage.items() ] async def test_sketch_manager(): """Test basic SketchManager functionality.""" SketchManager._instance = None # Create a temporary directory. temp_dir = tempfile.mkdtemp() try: # Create a SketchManager that uses the mock backend. manager = SketchManager(backend_class=MockSketchPadBackend) manager.sketch_dir = temp_dir print("=== Test Basic SketchManager Functionality ===") # Test creating a SketchPad. sketch_pad = manager.create_sketch_pad(sketch_id="test_sketch") print(f"✓ Created SketchPad: {sketch_pad.sketch_pad_id}") # Test setting an item. key = await manager.set_item( "test_sketch", "key1", "value1", tags={"test", "demo"} ) print(f"✓ Set item: {key}") # Test retrieving an item. item = manager.get_item("test_sketch", "key1") print(f"✓ Retrieved item: {item.value if item else None}") # Test retrieving a value. value = manager.get_value("test_sketch", "key1") print(f"✓ Retrieved value: {value}") # Test tag-based search. results = manager.search_by_tags("test_sketch", {"test"}) print(f"✓ Tag search: found {len(results)} items") # Test content-based search. results = manager.search_by_content("test_sketch", "value") print(f"✓ Content search: found {len(results)} items") # Test deleting an item. success = manager.delete_item("test_sketch", "key1") print(f"✓ Deleted item: {success}") # Test retrieving a SketchPad. retrieved_pad = manager.get_sketch_pad("test_sketch") print(f"✓ Retrieved SketchPad: {retrieved_pad is not None}") # Test listing SketchPads. sketches = manager.list_sketch_pads() print(f"✓ Listed SketchPads: found {len(sketches)}") # Test saving a SketchPad. success = manager.save_sketch_pad("test_sketch") print(f"✓ Saved SketchPad: {success}") # Test deleting a SketchPad. success = manager.delete_sketch_pad("test_sketch") print(f"✓ Deleted SketchPad: {success}") print("\n=== All Tests Passed ===") finally: # Clean up the temporary directory. shutil.rmtree(temp_dir) async def test_global_manager(): """Test the global SketchManager.""" SketchManager._instance = None print("\n=== Test Global SketchManager ===") # Get the global manager. manager = get_sketch_manager() print(f"✓ Retrieved global manager: {type(manager)}") # Test creating a SketchPad using the Redis backend. try: sketch_pad = manager.create_sketch_pad(sketch_id="global_test") print(f"✓ Created global SketchPad: {sketch_pad.sketch_pad_id}") # Test setting an item. key = await manager.set_item( "global_test", "global_key", "global_value", tags={"global"} ) print(f"✓ Set global item: {key}") # Test retrieving a value. value = manager.get_value("global_test", "global_key") print(f"✓ Retrieved global value: {value}") # Clean up. manager.delete_sketch_pad("global_test") print("✓ Cleaned up global test data") except Exception as e: print(f"⚠ Global manager test failed (Redis may be required): {e}") print("=== Global Manager Test Complete ===") async def main(): """Main test function.""" print("Starting tests for the new SketchManager implementation...") await test_sketch_manager() await test_global_manager() print("\nAll tests complete!") if __name__ == "__main__": asyncio.run(main())