- Add ADR 001 for Hybrid Search Architecture - Implement Phase 1 (Exact Match) and Phase 2 (Semantic Fallback) in ChromaStore - Wrap blocking ChromaDB calls in asyncio.to_thread - Update IVectorStore interface to support category filtering and thresholds - Add comprehensive tests for hybrid search logic
72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
import pytest
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
from aiogram.types import Message
|
|
from aiogram.filters import CommandObject
|
|
from src.bot.handlers import get_router
|
|
from src.processor.dto import EnrichedNewsItemDTO
|
|
from datetime import datetime, timezone
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_latest_command_with_category_passing():
|
|
# Arrange
|
|
storage = MagicMock()
|
|
storage.get_latest = AsyncMock(return_value=[])
|
|
processor = MagicMock()
|
|
|
|
message = MagicMock(spec=Message)
|
|
message.answer = AsyncMock()
|
|
command = CommandObject(command="latest", args="Tech")
|
|
|
|
# We need to call the handler directly or via the router
|
|
# For simplicity, let's call the handler function if it was exported,
|
|
# but it's defined inside get_router.
|
|
# Let's extract the handler from the router.
|
|
router = get_router(storage, processor, "123")
|
|
|
|
# Find the handler for /latest
|
|
handler = None
|
|
for observer in router.message.handlers:
|
|
if "latest" in str(observer.callback):
|
|
handler = observer.callback
|
|
break
|
|
|
|
assert handler is not None
|
|
|
|
# Act
|
|
await handler(message, command)
|
|
|
|
# Assert
|
|
# Verify that storage.get_latest was called with the category
|
|
storage.get_latest.assert_called_once_with(limit=10, category="Tech")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_command_with_threshold():
|
|
# Arrange
|
|
storage = MagicMock()
|
|
storage.search = AsyncMock(return_value=[])
|
|
processor = MagicMock()
|
|
|
|
message = MagicMock(spec=Message)
|
|
message.answer = AsyncMock()
|
|
command = CommandObject(command="search", args="AI News")
|
|
|
|
router = get_router(storage, processor, "123")
|
|
|
|
handler = None
|
|
for observer in router.message.handlers:
|
|
if "search" in str(observer.callback):
|
|
handler = observer.callback
|
|
break
|
|
|
|
assert handler is not None
|
|
|
|
# Act
|
|
await handler(message, command)
|
|
|
|
# Assert
|
|
# Verify that storage.search was called with a threshold
|
|
args, kwargs = storage.search.call_args
|
|
assert kwargs["query"] == "AI News"
|
|
assert "threshold" in kwargs
|
|
assert kwargs["threshold"] < 1.0 # Should have some threshold
|