forked from zovos/vk_hackathon
56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
"""Unit tests for search/aggregation.py"""
|
|
import sys
|
|
import os
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
|
|
|
os.environ.setdefault("EMBEDDINGS_DENSE_URL", "http://localhost/embed")
|
|
os.environ.setdefault("RERANKER_URL", "http://localhost/rerank")
|
|
os.environ.setdefault("QDRANT_URL", "http://localhost:6333")
|
|
os.environ.setdefault("API_KEY", "test-key")
|
|
|
|
from search.aggregation import aggregate_message_ids
|
|
from search.config import TOP_K
|
|
|
|
|
|
def _point(message_ids: list[str]):
|
|
"""Fake qdrant point with payload."""
|
|
class FakePoint:
|
|
payload = {"metadata": {"message_ids": message_ids}}
|
|
return FakePoint()
|
|
|
|
|
|
class TestAggregateMessageIds:
|
|
def test_empty_inputs(self):
|
|
result = aggregate_message_ids([], [])
|
|
assert result == []
|
|
|
|
def test_basic_dedup(self):
|
|
head = [_point(["m1", "m2"]), _point(["m2", "m3"])]
|
|
result = aggregate_message_ids(head, [])
|
|
assert result.count("m2") == 1
|
|
|
|
def test_head_before_tail(self):
|
|
head = [_point(["head_msg"])]
|
|
tail = [_point(["tail_msg"])]
|
|
result = aggregate_message_ids(head, tail)
|
|
assert result.index("head_msg") < result.index("tail_msg")
|
|
|
|
def test_top_k_limit(self):
|
|
# Create enough points to exceed TOP_K
|
|
points = [_point([f"m{i}"]) for i in range(TOP_K + 20)]
|
|
result = aggregate_message_ids(points, [])
|
|
assert len(result) <= TOP_K
|
|
|
|
def test_cross_point_dedup(self):
|
|
head = [_point(["shared"]), _point(["shared", "unique"])]
|
|
result = aggregate_message_ids(head, [])
|
|
assert result.count("shared") == 1
|
|
assert "unique" in result
|
|
|
|
def test_tail_fills_after_head(self):
|
|
head = [_point(["h1"])]
|
|
tail = [_point(["t1"]), _point(["t2"])]
|
|
result = aggregate_message_ids(head, tail)
|
|
assert "h1" in result
|
|
assert "t1" in result
|
|
assert "t2" in result
|