96 lines
2.8 KiB
Python
96 lines
2.8 KiB
Python
import pytest
|
|
import pytest_asyncio
|
|
from datetime import datetime, timezone
|
|
import chromadb
|
|
from chromadb.config import Settings
|
|
|
|
from src.processor.dto import EnrichedNewsItemDTO
|
|
from src.storage.chroma_store import ChromaStore
|
|
|
|
@pytest_asyncio.fixture
|
|
async def chroma_store():
|
|
# Use EphemeralClient for in-memory testing
|
|
client = chromadb.EphemeralClient(Settings(allow_reset=True))
|
|
client.reset()
|
|
store = ChromaStore(client=client, collection_name="test_top_ranked_collection")
|
|
yield store
|
|
client.reset()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_top_ranked_sorting(chroma_store: ChromaStore):
|
|
"""
|
|
Test that get_top_ranked returns items sorted by relevance_score in descending order.
|
|
"""
|
|
# 1. Arrange - create items with various relevance scores
|
|
items = [
|
|
EnrichedNewsItemDTO(
|
|
title=f"News {score}",
|
|
url=f"https://example.com/{score}",
|
|
content_text=f"Content for news with score {score}",
|
|
source="Source",
|
|
timestamp=datetime.now(timezone.utc),
|
|
relevance_score=score,
|
|
summary_ru=f"Сводка {score}",
|
|
anomalies_detected=[],
|
|
category="Tech"
|
|
) for score in [5, 10, 2, 8, 1]
|
|
]
|
|
|
|
for item in items:
|
|
await chroma_store.store(item)
|
|
|
|
# 2. Act
|
|
results = await chroma_store.get_top_ranked(limit=10)
|
|
|
|
# 3. Assert
|
|
assert len(results) == 5
|
|
scores = [r.relevance_score for r in results]
|
|
# Should be [10, 8, 5, 2, 1]
|
|
assert scores == [10, 8, 5, 2, 1]
|
|
assert results[0].title == "News 10"
|
|
assert results[-1].title == "News 1"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_top_ranked_limit(chroma_store: ChromaStore):
|
|
"""
|
|
Test that get_top_ranked respects the limit parameter.
|
|
"""
|
|
# 1. Arrange
|
|
items = [
|
|
EnrichedNewsItemDTO(
|
|
title=f"News {i}",
|
|
url=f"https://example.com/{i}",
|
|
content_text=f"Content {i}",
|
|
source="Source",
|
|
timestamp=datetime.now(timezone.utc),
|
|
relevance_score=i,
|
|
summary_ru=f"Сводка {i}",
|
|
anomalies_detected=[],
|
|
category="Tech"
|
|
) for i in range(1, 11) # 10 items
|
|
]
|
|
|
|
for item in items:
|
|
await chroma_store.store(item)
|
|
|
|
# 2. Act
|
|
limit_5 = await chroma_store.get_top_ranked(limit=5)
|
|
limit_2 = await chroma_store.get_top_ranked(limit=2)
|
|
|
|
# 3. Assert
|
|
assert len(limit_5) == 5
|
|
assert len(limit_2) == 2
|
|
assert limit_5[0].relevance_score == 10
|
|
assert limit_5[4].relevance_score == 6
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_top_ranked_empty_store(chroma_store: ChromaStore):
|
|
"""
|
|
Test that get_top_ranked returns an empty list if store is empty.
|
|
"""
|
|
# 1. Act
|
|
results = await chroma_store.get_top_ranked(limit=10)
|
|
|
|
# 2. Assert
|
|
assert results == []
|