#!/usr/bin/env python3
"""
Test the initialization flow fix for QdrantUnifiedMiddleware.
This test validates that:
1. The middleware initialization properly calls the background reindexing setup
2. on_call_tool uses the correct initialization method
3. Backward compatibility is maintained
"""
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from middleware.qdrant_middleware import QdrantUnifiedMiddleware
class TestQdrantInitializationFlow:
"""Test the corrected initialization flow in QdrantUnifiedMiddleware."""
@pytest.fixture
def mock_middleware(self):
"""Create a QdrantUnifiedMiddleware with mocked dependencies."""
with (
patch(
"middleware.qdrant_middleware.QdrantClientManager"
) as mock_client_manager_class,
patch(
"middleware.qdrant_middleware.QdrantStorageManager"
) as mock_storage_manager_class,
patch(
"middleware.qdrant_middleware.QdrantSearchManager"
) as mock_search_manager_class,
patch(
"middleware.qdrant_middleware.QdrantResourceHandler"
) as mock_resource_handler_class,
):
# Create mock instances
mock_client_manager = AsyncMock()
mock_client_manager.is_initialized = False
mock_client_manager.is_available = True
mock_client_manager.initialize.return_value = True
mock_client_manager_class.return_value = mock_client_manager
mock_storage_manager = AsyncMock()
mock_storage_manager_class.return_value = mock_storage_manager
mock_search_manager = AsyncMock()
mock_search_manager_class.return_value = mock_search_manager
mock_resource_handler = AsyncMock()
mock_resource_handler_class.return_value = mock_resource_handler
# Create middleware
middleware = QdrantUnifiedMiddleware(
qdrant_host="localhost", qdrant_port=6333, enabled=True
)
# Store references to mocks for test access
middleware._mock_client_manager = mock_client_manager
middleware._mock_storage_manager = mock_storage_manager
return middleware
@pytest.mark.asyncio
async def test_initialize_middleware_and_reindexing_calls_background_reindexing(
self, mock_middleware
):
"""Test that the new initialization method properly starts background reindexing."""
with patch.object(
mock_middleware, "_start_background_reindexing", new_callable=AsyncMock
) as mock_start_reindexing:
# Call the new initialization method
result = await mock_middleware.initialize_middleware_and_reindexing()
# Verify results
assert result is True
assert mock_middleware._initialized is True
mock_middleware._mock_client_manager.initialize.assert_called_once()
mock_start_reindexing.assert_called_once()
@pytest.mark.asyncio
async def test_backward_compatibility_initialize_method(self, mock_middleware):
"""Test that the legacy initialize() method still works and delegates properly."""
with patch.object(
mock_middleware,
"initialize_middleware_and_reindexing",
new_callable=AsyncMock,
) as mock_init_full:
mock_init_full.return_value = True
# Call the legacy method
result = await mock_middleware.initialize()
# Verify it delegates to the new method
assert result is True
mock_init_full.assert_called_once()
@pytest.mark.asyncio
async def test_on_call_tool_uses_correct_initialization(self, mock_middleware):
"""Test that on_call_tool calls the middleware initialization instead of client manager directly."""
# Mock context and call_next
mock_context = MagicMock()
mock_context.message.name = "test_tool"
mock_context.message.arguments = {}
mock_call_next = AsyncMock()
mock_call_next.return_value = {"result": "success"}
# Mock the initialization method
with patch.object(
mock_middleware,
"initialize_middleware_and_reindexing",
new_callable=AsyncMock,
) as mock_init_full:
mock_init_full.return_value = True
# Mock other dependencies
with (
patch(
"middleware.qdrant_middleware.get_session_context",
return_value="test_session",
),
patch.object(
mock_middleware.storage_manager,
"_store_response_with_params",
new_callable=AsyncMock,
),
):
# Set client as not initialized to trigger initialization
mock_middleware._mock_client_manager.is_initialized = False
# Call on_call_tool
result = await mock_middleware.on_call_tool(
mock_context, mock_call_next
)
# Verify the middleware initialization was called
mock_init_full.assert_called_once()
# Verify the tool was called and response returned
mock_call_next.assert_called_once()
assert result == {"result": "success"}
@pytest.mark.asyncio
async def test_ensure_initialization_uses_correct_method(self, mock_middleware):
"""Test that _ensure_initialization calls the middleware initialization."""
with (
patch.object(
mock_middleware,
"initialize_middleware_and_reindexing",
new_callable=AsyncMock,
) as mock_init_full,
patch("asyncio.create_task") as mock_create_task,
):
# Set up for triggering initialization
mock_middleware._early_init_started = False
# Call _ensure_initialization
await mock_middleware._ensure_initialization("test_context")
# Verify it started the task with the correct method
assert mock_middleware._early_init_started is True
mock_create_task.assert_called_once()
# Get the task argument and verify it's the right method
task_arg = mock_create_task.call_args[0][0]
# The task should be a coroutine for the middleware initialization
assert hasattr(task_arg, "__await__") # It's a coroutine
def test_initialization_variables_properly_set(self, mock_middleware):
"""Test that initialization variables are properly set in __init__."""
# Verify reindexing control variables are set
assert hasattr(mock_middleware, "_early_init_started")
assert hasattr(mock_middleware, "_reindexing_task")
assert hasattr(mock_middleware, "_reindexing_enabled")
assert mock_middleware._early_init_started is False
assert mock_middleware._reindexing_task is None
assert mock_middleware._reindexing_enabled is True
@pytest.mark.asyncio
async def test_reindexing_not_started_when_disabled(self, mock_middleware):
"""Test that reindexing is not started when disabled."""
with patch.object(
mock_middleware, "_start_background_reindexing", new_callable=AsyncMock
) as mock_start_reindexing:
# Disable reindexing
mock_middleware._reindexing_enabled = False
# Call initialization
result = await mock_middleware.initialize_middleware_and_reindexing()
# Verify reindexing was not started
assert result is True
mock_start_reindexing.assert_not_called()
@pytest.mark.asyncio
async def test_reindexing_not_started_on_client_init_failure(self, mock_middleware):
"""Test that reindexing is not started when client initialization fails."""
with patch.object(
mock_middleware, "_start_background_reindexing", new_callable=AsyncMock
) as mock_start_reindexing:
# Make client initialization fail
mock_middleware._mock_client_manager.initialize.return_value = False
# Call initialization
result = await mock_middleware.initialize_middleware_and_reindexing()
# Verify reindexing was not started
assert result is False
assert mock_middleware._initialized is False
mock_start_reindexing.assert_not_called()
if __name__ == "__main__":
pytest.main([__file__, "-v"])
