271 lines
8.5 KiB
Python
271 lines
8.5 KiB
Python
#!/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())
|