File size: 2,856 Bytes
35676b4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """tests/test_reranker.py — unit tests for storage/reranker.py.
We don't require the actual cross-encoder model to be downloaded; the
tests use a fake model object that scores by simple keyword overlap so
we can assert deterministic ordering changes.
"""
from __future__ import annotations
import pytest
from storage import reranker
class _FakeModel:
"""Minimal stand-in for sentence_transformers.CrossEncoder.
Returns higher scores for candidates whose text shares more words with the query.
"""
def predict(self, pairs):
scores = []
for query, text in pairs:
q_tokens = set(query.lower().split())
t_tokens = set(text.lower().split())
scores.append(float(len(q_tokens & t_tokens)))
return scores
@pytest.fixture(autouse=True)
def _reset(monkeypatch):
"""Reset the module-level model cache between tests."""
reranker.reset_for_test()
yield
reranker.reset_for_test()
def _patch_with_fake(monkeypatch, model=None):
monkeypatch.setattr(reranker, "_get_model", lambda: model or _FakeModel())
def test_passthrough_when_few_candidates(monkeypatch):
_patch_with_fake(monkeypatch)
candidates = [{"text": "a"}, {"text": "b"}]
out = reranker.rerank("anything", candidates, top_k=3)
assert out == candidates
def test_truncates_to_top_k(monkeypatch):
_patch_with_fake(monkeypatch)
candidates = [{"text": f"chunk {i}"} for i in range(10)]
out = reranker.rerank("query", candidates, top_k=3)
assert len(out) == 3
def test_relevant_chunk_promoted_to_top(monkeypatch):
"""Candidate with most query-token overlap should rank #1 after rerank."""
_patch_with_fake(monkeypatch)
candidates = [
{"text": "the quick brown fox"},
{"text": "completely unrelated content here"},
{"text": "gross margin pressure from input costs and tariff exposure"},
{"text": "share buybacks resumed this quarter"},
{"text": "another unrelated thing"},
]
query = "gross margin pressure tariff"
out = reranker.rerank(query, candidates, top_k=1)
assert out[0]["text"].startswith("gross margin pressure")
def test_falls_back_when_model_unavailable(monkeypatch):
"""If model loading fails, return original order truncated to top_k."""
monkeypatch.setattr(reranker, "_get_model", lambda: None)
candidates = [{"text": f"c{i}"} for i in range(5)]
out = reranker.rerank("q", candidates, top_k=2)
assert out == candidates[:2]
def test_falls_back_when_predict_raises(monkeypatch):
class _BadModel:
def predict(self, pairs):
raise RuntimeError("boom")
_patch_with_fake(monkeypatch, model=_BadModel())
candidates = [{"text": f"c{i}"} for i in range(5)]
out = reranker.rerank("q", candidates, top_k=2)
assert out == candidates[:2]
|