#!/usr/bin/env python3
"""
Test suite for improved Qdrant unified tools with FastMCP2 conventions.
Tests the new search_qdrant and fetch_qdrant tools with proper type validation.
"""
import asyncio
import sys
from pathlib import Path
# Add the project root to Python path
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))
# Test imports
try:
from middleware.qdrant_middleware import (
QdrantUnifiedMiddleware,
setup_enhanced_qdrant_tools,
)
# Try importing query parsing functions - they may be in a different module
try:
from middleware.qdrant_core.search import (
_extract_service_from_tool,
_parse_unified_query,
)
except ImportError:
# Create mock functions for testing if not available
def _parse_unified_query(query):
return {"capability": "general_search", "confidence": 0.8}
def _extract_service_from_tool(tool_name):
if "gmail" in tool_name:
return "gmail"
elif "drive" in tool_name:
return "drive"
elif "calendar" in tool_name:
return "calendar"
else:
return "unknown"
try:
from middleware.qdrant_types import (
QdrantDocumentMetadata,
QdrantFetchResponse,
QdrantSearchResponse,
QdrantSearchResultItem,
)
except ImportError:
# Create mock types for testing
QdrantSearchResponse = dict
QdrantFetchResponse = dict
QdrantSearchResultItem = dict
QdrantDocumentMetadata = dict
print("⚠️ Using mock types for testing")
print("✅ All imports successful")
except ImportError as e:
print(f"❌ Import failed: {e}")
sys.exit(1)
# Test data
TEST_QUERIES = [
# Overview/Analytics queries
("overview", "overview"),
("analytics dashboard", "overview"),
("tool usage stats", "overview"),
# Service history queries
("gmail last week", "service_history"),
("service:drive documents", "service_history"),
("recent calendar events", "service_history"),
# General search queries
("email collaboration documents", "general_search"),
("document creation workflow", "general_search"),
("user authentication errors", "general_search"),
]
SERVICE_EXTRACTION_TESTS = [
("search_gmail_messages", "gmail"),
("get_drive_file_content", "drive"),
("list_calendar_events", "calendar"),
("create_docs", "docs"),
("update_sheets_data", "sheets"),
("send_chat_message", "chat"),
("list_photos_albums", "photos"),
("create_form", "forms"),
("export_presentation", "slides"),
("unknown_tool_name", "unknown"),
]
def test_query_parser():
"""Test the enhanced query parser functionality."""
print("\n🔍 Testing Query Parser")
passed = 0
failed = 0
for query, expected_capability in TEST_QUERIES:
try:
result = _parse_unified_query(query)
actual_capability = result.get("capability")
confidence = result.get("confidence", 0.0)
if actual_capability == expected_capability:
print(
f" ✅ '{query}' -> {actual_capability} (confidence: {confidence:.2f})"
)
passed += 1
else:
print(
f" ❌ '{query}' -> Expected: {expected_capability}, Got: {actual_capability}"
)
failed += 1
except Exception as e:
print(f" ❌ '{query}' -> Error: {e}")
failed += 1
print(f"\nQuery Parser Results: {passed} passed, {failed} failed")
return failed == 0
def test_service_extraction():
"""Test service extraction from tool names."""
print("\n🏷️ Testing Service Extraction")
passed = 0
failed = 0
for tool_name, expected_service in SERVICE_EXTRACTION_TESTS:
try:
actual_service = _extract_service_from_tool(tool_name)
if actual_service == expected_service:
print(f" ✅ '{tool_name}' -> {actual_service}")
passed += 1
else:
print(
f" ❌ '{tool_name}' -> Expected: {expected_service}, Got: {actual_service}"
)
failed += 1
except Exception as e:
print(f" ❌ '{tool_name}' -> Error: {e}")
failed += 1
print(f"\nService Extraction Results: {passed} passed, {failed} failed")
return failed == 0
def test_type_definitions():
"""Test that TypedDict definitions are properly structured."""
print("\n📋 Testing Type Definitions")
try:
# Test QdrantSearchResponse structure
search_response = QdrantSearchResponse(
results=[
QdrantSearchResultItem(
id="123e4567-e89b-12d3-a456-426614174000",
title="Test Result",
url="qdrant://test_collection?123e4567-e89b-12d3-a456-426614174000",
)
],
query="test query",
query_type="general_search",
total_results=1,
collection_name="test_collection",
)
print(" ✅ QdrantSearchResponse structure valid")
# Test QdrantFetchResponse structure
fetch_response = QdrantFetchResponse(
id="123e4567-e89b-12d3-a456-426614174001",
title="Test Document",
text="Test document content",
url="qdrant://test_collection?123e4567-e89b-12d3-a456-426614174001",
metadata=QdrantDocumentMetadata(
tool_name="test_tool",
service="test_service",
service_display_name="Test Service",
user_email="test@example.com",
timestamp="2024-01-01T00:00:00Z",
response_type="dict",
arguments_count=2,
payload_type="tool_response",
collection_name="test_collection",
point_id="123e4567-e89b-12d3-a456-426614174001",
),
found=True,
collection_name="test_collection",
)
print(" ✅ QdrantFetchResponse structure valid")
print(" ✅ QdrantDocumentMetadata structure valid")
return True
except Exception as e:
print(f" ❌ Type definition error: {e}")
return False
def test_middleware_initialization():
"""Test middleware initialization without requiring actual Qdrant connection."""
print("\n⚙️ Testing Middleware Initialization")
try:
# Create middleware with mock settings
middleware = QdrantUnifiedMiddleware(
qdrant_host="localhost",
qdrant_port=6333,
collection_name="test_collection",
enabled=True,
auto_discovery=False, # Disable auto-discovery for testing
)
print(" ✅ Middleware instance created")
print(f" ✅ Collection name: {middleware.config.collection_name}")
print(f" ✅ Host: {middleware.config.host}")
print(f" ✅ Enabled: {middleware.config.enabled}")
# Test reindexing control attributes
print(f" ✅ Reindexing enabled: {middleware._reindexing_enabled}")
print(
f" ✅ Reindexing task initialized: {middleware._reindexing_task is None}"
)
return True
except Exception as e:
print(f" ❌ Middleware initialization error: {e}")
return False
def test_reindexing_capabilities():
"""Test reindexing method availability and configuration."""
print("\n🔧 Testing Reindexing Capabilities")
try:
# Create middleware instance
middleware = QdrantUnifiedMiddleware(
collection_name="test_collection", enabled=False # Disable for testing
)
# Check that reindexing methods exist
reindexing_methods = [
"trigger_immediate_reindexing",
"get_collection_health_status",
"stop_background_reindexing",
"_start_background_reindexing",
]
for method_name in reindexing_methods:
if hasattr(middleware, method_name):
method = getattr(middleware, method_name)
if callable(method):
print(f" ✅ {method_name} method available")
else:
print(f" ❌ {method_name} exists but not callable")
return False
else:
print(f" ❌ {method_name} method missing")
return False
# Check storage manager reindexing methods
storage_methods = [
"reindex_collection",
"cleanup_stale_data",
"schedule_background_reindexing",
"_analyze_collection_health",
]
for method_name in storage_methods:
if hasattr(middleware.storage_manager, method_name):
print(f" ✅ Storage manager {method_name} method available")
else:
print(f" ❌ Storage manager {method_name} method missing")
return False
# Check client manager optimization methods
client_methods = [
"optimize_collection_performance",
"rebuild_collection_completely",
]
for method_name in client_methods:
if hasattr(middleware.client_manager, method_name):
print(f" ✅ Client manager {method_name} method available")
else:
print(f" ❌ Client manager {method_name} method missing")
return False
print(" ✅ All reindexing capabilities verified")
return True
except Exception as e:
print(f" ❌ Reindexing capabilities test error: {e}")
return False
def test_index_configuration_validation():
"""Test that index configurations are valid for Qdrant."""
print("\n📋 Testing Index Configuration Validation")
try:
from middleware.qdrant_core.lazy_imports import get_qdrant_imports
# Test lazy import of enhanced index classes
_, qdrant_models = get_qdrant_imports()
# Verify all required index classes are available
required_index_classes = [
"KeywordIndexParams",
"KeywordIndexType",
"IntegerIndexParams",
"IntegerIndexType",
"BoolIndexParams",
"BoolIndexType",
"DatetimeIndexParams",
"DatetimeIndexType",
]
for class_name in required_index_classes:
if class_name in qdrant_models:
print(f" ✅ {class_name} available in lazy imports")
else:
print(f" ❌ {class_name} missing from lazy imports")
return False
# Test creating index configurations (without actually connecting to Qdrant)
try:
# Test tenant-optimized keyword index
tenant_idx = qdrant_models["KeywordIndexParams"](
type=qdrant_models["KeywordIndexType"].KEYWORD,
is_tenant=True,
on_disk=False,
)
print(" ✅ Tenant-optimized keyword index configuration valid")
# Test principal time-based integer index
time_idx = qdrant_models["IntegerIndexParams"](
type=qdrant_models["IntegerIndexType"].INTEGER,
lookup=False,
range=True,
is_principal=True,
)
print(" ✅ Principal time-based integer index configuration valid")
# Test boolean index
bool_idx = qdrant_models["BoolIndexParams"](
type=qdrant_models["BoolIndexType"].BOOL, on_disk=True
)
print(" ✅ Boolean index configuration valid")
# Test datetime index
dt_idx = qdrant_models["DatetimeIndexParams"](
type=qdrant_models["DatetimeIndexType"].DATETIME, on_disk=True
)
print(" ✅ Datetime index configuration valid")
except Exception as e:
print(f" ❌ Index configuration creation failed: {e}")
return False
print(" ✅ All index configurations validated successfully")
return True
except Exception as e:
print(f" ❌ Index configuration validation error: {e}")
return False
async def test_tool_setup():
"""Test that tools can be set up without errors."""
print("\n🔧 Testing Tool Setup")
try:
# Create mock FastMCP instance
class MockMCP:
def __init__(self):
self.tools = {}
def tool(self, name=None, description=None, tags=None, annotations=None):
def decorator(func):
self.tools[name or func.__name__] = {
"function": func,
"description": description,
"tags": tags,
"annotations": annotations,
}
return func
return decorator
# Create mock middleware
middleware = QdrantUnifiedMiddleware(
collection_name="test_collection", enabled=False # Disable for testing
)
mock_mcp = MockMCP()
# Setup tools
setup_enhanced_qdrant_tools(mock_mcp, middleware)
print(f" ✅ {len(mock_mcp.tools)} tools registered")
for tool_name, tool_info in mock_mcp.tools.items():
desc = (
tool_info["description"][:60] + "..."
if len(tool_info["description"]) > 60
else tool_info["description"]
)
print(f" - {tool_name}: {desc}")
# Verify expected tools exist
expected_tools = ["search_qdrant", "fetch_qdrant"]
for expected_tool in expected_tools:
if expected_tool in mock_mcp.tools:
print(f" ✅ {expected_tool} tool registered correctly")
else:
print(f" ❌ {expected_tool} tool missing")
return False
# Test backwards compatibility - legacy tools should still exist
legacy_tools = [
"search_tool_history",
"get_tool_analytics",
"get_response_details",
]
legacy_found = 0
for legacy_tool in legacy_tools:
if legacy_tool in mock_mcp.tools:
print(f" ✅ {legacy_tool} legacy tool present")
legacy_found += 1
if legacy_found == len(legacy_tools):
print(
f" ✅ All {legacy_found} legacy tools preserved for backward compatibility"
)
else:
print(f" ⚠️ Only {legacy_found}/{len(legacy_tools)} legacy tools found")
return True
except Exception as e:
print(f" ❌ Tool setup error: {e}")
return False
def test_url_format():
"""Test that URLs follow the specified format."""
print("\n🔗 Testing URL Format")
try:
collection_name = "test_collection"
point_id = "123e4567-e89b-12d3-a456-426614174002"
expected_format = f"qdrant://{collection_name}?{point_id}"
print(f" ✅ URL format: {expected_format}")
# Test parsing
if expected_format.startswith("qdrant://"):
parts = expected_format.replace("qdrant://", "").split("?")
if len(parts) == 2:
parsed_collection = parts[0]
parsed_point_id = parts[1]
if parsed_collection == collection_name and parsed_point_id == point_id:
print(
f" ✅ URL parsing successful: collection={parsed_collection}, point={parsed_point_id}"
)
return True
else:
print(
f" ❌ URL parsing failed: expected collection={collection_name}, point={point_id}"
)
return False
else:
print(" ❌ URL format invalid: wrong number of parts")
return False
else:
print(" ❌ URL format invalid: missing qdrant:// prefix")
return False
except Exception as e:
print(f" ❌ URL format test error: {e}")
return False
def test_fastmcp2_compliance():
"""Test FastMCP2 conventions compliance."""
print("\n📝 Testing FastMCP2 Compliance")
try:
# Import check
try:
from tools.common_types import UserGoogleEmail
print(" ✅ UserGoogleEmail import successful")
except ImportError as e:
print(f" ⚠️ UserGoogleEmail import failed: {e}")
# Test that our types are structured dictionaries, not classes
from middleware.qdrant_types import QdrantFetchResponse, QdrantSearchResponse
# These should be TypedDict, so we can create them as dict
search_resp = {
"results": [],
"query": "test",
"query_type": "general_search",
"total_results": 0,
"collection_name": "test",
}
# Should work as TypedDict
typed_search_resp: QdrantSearchResponse = search_resp
print(" ✅ QdrantSearchResponse is properly typed")
fetch_resp = {
"id": "123e4567-e89b-12d3-a456-426614174003",
"title": "test",
"text": "test",
"url": "qdrant://test_collection?123e4567-e89b-12d3-a456-426614174003",
"metadata": {
"tool_name": "test",
"service": "test",
"service_display_name": "test",
"user_email": "test",
"timestamp": "test",
"response_type": "test",
"arguments_count": 0,
"payload_type": "test",
"collection_name": "test",
"point_id": "123e4567-e89b-12d3-a456-426614174003",
},
"found": True,
"collection_name": "test",
}
typed_fetch_resp: QdrantFetchResponse = fetch_resp
print(" ✅ QdrantFetchResponse is properly typed")
return True
except Exception as e:
print(f" ❌ FastMCP2 compliance error: {e}")
return False
async def run_all_tests():
"""Run all test functions."""
print("🚀 Starting Improved Qdrant Tools Test Suite")
print("=" * 60)
test_results = []
# Run synchronous tests
test_results.append(("Query Parser", test_query_parser()))
test_results.append(("Service Extraction", test_service_extraction()))
test_results.append(("Type Definitions", test_type_definitions()))
test_results.append(("Middleware Init", test_middleware_initialization()))
test_results.append(("URL Format", test_url_format()))
test_results.append(("FastMCP2 Compliance", test_fastmcp2_compliance()))
test_results.append(("Reindexing Capabilities", test_reindexing_capabilities()))
test_results.append(
("Index Configuration Validation", test_index_configuration_validation())
)
# Run async tests
test_results.append(("Tool Setup", await test_tool_setup()))
# Print summary
print("\n" + "=" * 60)
print("📊 TEST SUMMARY")
print("=" * 60)
passed = 0
failed = 0
for test_name, result in test_results:
status = "✅ PASS" if result else "❌ FAIL"
print(f"{status} {test_name}")
if result:
passed += 1
else:
failed += 1
print(f"\nTotal: {passed + failed} tests")
print(f"Passed: {passed}")
print(f"Failed: {failed}")
if failed == 0:
print("\n🎉 All tests passed! The improved Qdrant tools are ready.")
print("\n🔧 New Tools Available:")
print(" - search_qdrant: Enhanced search with FastMCP2 types")
print(" - fetch_qdrant: Enhanced fetch with structured metadata")
print("\n📋 Features:")
print(" - Proper TypedDict return types for FastMCP schema generation")
print(" - Qdrant-specific URL format: qdrant://collection?point_id")
print(" - Service metadata integration with icons")
print(" - Enhanced query parsing with 4 core capabilities")
print(" - Backward compatibility with legacy tools")
print("\n🚀 Reindexing Capabilities:")
print(" - Intelligent background reindexing with adaptive scheduling")
print(" - Tenant-optimized indexes for multi-user collections")
print(" - Principal time-based indexes for efficient range queries")
print(" - Collection health monitoring and automatic optimization")
print(" - Parameterized indexes with lookup/range configuration")
print(" - On-disk vs in-memory index placement optimization")
return True
else:
print(f"\n⚠️ {failed} test(s) failed. Please review the issues above.")
return False
if __name__ == "__main__":
success = asyncio.run(run_all_tests())
sys.exit(0 if success else 1)
