setup tests
This commit is contained in:
@@ -0,0 +1,150 @@
|
||||
"""Tests for EPUB search reranking."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from python.ebook_search.config import EbookSearchConfig, RerankConfig, load_rerank_config
|
||||
from python.ebook_search.rerank import rerank_chunks
|
||||
from python.ebook_search.search import SearchResult, apply_rerank, skip_rerank
|
||||
|
||||
|
||||
def candidates() -> list[SearchResult]:
|
||||
return [
|
||||
SearchResult(chunk_id=1, text="alpha", source_title="A", score=0.9),
|
||||
SearchResult(chunk_id=2, text="beta", source_title="B", score=0.8),
|
||||
SearchResult(chunk_id=3, text="gamma", source_title="C", score=0.7),
|
||||
]
|
||||
|
||||
|
||||
def rerank_response(payload: dict[str, object] | None = None, *, content: bytes | None = None) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
200,
|
||||
content=content,
|
||||
json=payload,
|
||||
request=httpx.Request("POST", "http://rerank.test/rerank"),
|
||||
)
|
||||
|
||||
|
||||
def test_config_defaults_keep_reranking_optional(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("EBOOK_SEARCH_RERANK_ENABLED", raising=False)
|
||||
monkeypatch.delenv("EBOOK_SEARCH_RERANK_BASE_URL", raising=False)
|
||||
monkeypatch.delenv("EBOOK_SEARCH_RERANK_MODEL", raising=False)
|
||||
monkeypatch.delenv("EBOOK_SEARCH_RERANK_CANDIDATES", raising=False)
|
||||
monkeypatch.delenv("EBOOK_SEARCH_RERANK_TIMEOUT_SECONDS", raising=False)
|
||||
|
||||
config = load_rerank_config()
|
||||
|
||||
assert config.enabled is False
|
||||
assert config.base_url == "http://192.168.90.25:8001"
|
||||
assert config.model == "qwen3-reranker-06b"
|
||||
assert config.candidates == 24
|
||||
assert config.timeout_seconds == 30
|
||||
|
||||
|
||||
def test_reranking_disabled_returns_original_fused_order() -> None:
|
||||
config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=2)
|
||||
|
||||
response = skip_rerank("query", candidates(), config)
|
||||
|
||||
assert response.rank_label == "Hybrid"
|
||||
assert [result.chunk_id for result in response.results] == [1, 2]
|
||||
|
||||
|
||||
def test_reranking_enabled_reorders_candidates(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def fake_post(_url: str, *, json: dict[str, object], timeout: float) -> httpx.Response:
|
||||
assert timeout == 30
|
||||
assert json == {
|
||||
"model": "qwen3-reranker-06b",
|
||||
"query": "query",
|
||||
"documents": ["alpha", "beta", "gamma"],
|
||||
}
|
||||
return rerank_response(
|
||||
{
|
||||
"results": [
|
||||
{"index": 0, "relevance_score": 0.1},
|
||||
{"index": 1, "relevance_score": 0.9},
|
||||
{"index": 2, "relevance_score": 0.4},
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
monkeypatch.setattr(httpx, "post", fake_post)
|
||||
|
||||
results = rerank_chunks("query", candidates(), RerankConfig())
|
||||
|
||||
assert [result.chunk_id for result in results] == [2, 1, 3]
|
||||
assert [round(result.score, 3) for result in results] == [0.45, 0.1, 0.0]
|
||||
assert [result.rerank_score for result in results] == [0.9, 0.1, 0.4]
|
||||
|
||||
|
||||
def test_reranking_cannot_ignore_hybrid_score(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
candidates = [
|
||||
SearchResult(chunk_id=1, text="strong hybrid", source_title="A", score=1.0),
|
||||
SearchResult(chunk_id=2, text="weak hybrid", source_title="B", score=0.1),
|
||||
]
|
||||
|
||||
def fake_post(_url: str, **_kwargs: object) -> httpx.Response:
|
||||
return rerank_response(
|
||||
{
|
||||
"results": [
|
||||
{"index": 0, "relevance_score": 0.7},
|
||||
{"index": 1, "relevance_score": 1.0},
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
monkeypatch.setattr(httpx, "post", fake_post)
|
||||
|
||||
results = rerank_chunks("query", candidates, RerankConfig())
|
||||
|
||||
assert [result.chunk_id for result in results] == [1, 2]
|
||||
assert results[0].score == 0.7
|
||||
assert results[1].score == 0.0
|
||||
assert results[1].rerank_score == 1.0
|
||||
|
||||
|
||||
def test_vllm_rerank_timeout_raises(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def fake_rerank_chunks(
|
||||
_query: str,
|
||||
_candidates: list[SearchResult],
|
||||
_config: RerankConfig,
|
||||
) -> list[SearchResult]:
|
||||
message = "timeout"
|
||||
raise httpx.TimeoutException(message)
|
||||
|
||||
monkeypatch.setattr("python.ebook_search.search.rerank_chunks", fake_rerank_chunks)
|
||||
config = EbookSearchConfig(rerank=RerankConfig(enabled=True), top_k=2)
|
||||
|
||||
with pytest.raises(httpx.TimeoutException, match="timeout"):
|
||||
apply_rerank("query", candidates(), config)
|
||||
|
||||
|
||||
def test_malformed_vllm_rerank_json_does_not_crash_search(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def fake_post(_url: str, **_kwargs: object) -> httpx.Response:
|
||||
return rerank_response(content=b"not-json")
|
||||
|
||||
monkeypatch.setattr(httpx, "post", fake_post)
|
||||
|
||||
results = rerank_chunks("query", candidates()[:1], RerankConfig())
|
||||
|
||||
assert results[0].score == 0.0
|
||||
|
||||
|
||||
def test_vllm_rerank_scores_are_clamped(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def fake_post(_url: str, **_kwargs: object) -> httpx.Response:
|
||||
return rerank_response(
|
||||
{
|
||||
"results": [
|
||||
{"index": 0, "relevance_score": -1},
|
||||
{"index": 1, "relevance_score": 2},
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
monkeypatch.setattr(httpx, "post", fake_post)
|
||||
|
||||
results = rerank_chunks("query", candidates()[:2], RerankConfig())
|
||||
|
||||
assert [result.rerank_score for result in results] == [0.0, 1.0]
|
||||
Reference in New Issue
Block a user