test(ebook): cover protected phrases and migrate suite to async
Add test_protected_phrases.py covering phrase-matching behavior in the RAG engine, and update the existing ebook_search tests to use the async SQLAlchemy engine/session (create_async_engine, AsyncSession) and async HTTP paths.
This commit is contained in:
@@ -2,10 +2,11 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from threading import Event
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from python.ebook_search.config import EbookSearchConfig, RerankConfig
|
||||
from python.ebook_search.search import SearchResult, search_ebooks
|
||||
@@ -14,18 +15,18 @@ if TYPE_CHECKING:
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
|
||||
def test_search_ebooks_runs_vector_and_bm25_in_parallel(mocker: MockerFixture) -> None:
|
||||
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
|
||||
async def test_search_ebooks_runs_vector_and_bm25_in_parallel(mocker: MockerFixture) -> None:
|
||||
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
||||
vector_started = Event()
|
||||
bm25_started = Event()
|
||||
received_engines: list[object] = []
|
||||
|
||||
def fake_vector_candidates(received_engine, query, _config):
|
||||
async def fake_vector_candidates(received_engine, _client, query, _config):
|
||||
"""Return vector candidates after confirming BM25 has started."""
|
||||
received_engines.append(received_engine)
|
||||
assert query == "what is parallel"
|
||||
vector_started.set()
|
||||
assert bm25_started.wait(timeout=2)
|
||||
assert await asyncio.to_thread(bm25_started.wait, 2)
|
||||
return [SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)]
|
||||
|
||||
def fake_bm25_candidates(query, _config):
|
||||
@@ -39,7 +40,7 @@ def test_search_ebooks_runs_vector_and_bm25_in_parallel(mocker: MockerFixture) -
|
||||
mocker.patch("python.ebook_search.search.bm25_candidates", side_effect=fake_bm25_candidates)
|
||||
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||
|
||||
response = search_ebooks(engine, "what is parallel", config)
|
||||
response = await search_ebooks(engine, mocker.Mock(), "what is parallel", config)
|
||||
|
||||
timings = {step.name: step for step in response.timings}
|
||||
assert [result.chunk_id for result in response.results] == [1, 2]
|
||||
@@ -47,3 +48,26 @@ def test_search_ebooks_runs_vector_and_bm25_in_parallel(mocker: MockerFixture) -
|
||||
assert timings["BM25 search"].counts_toward_total is False
|
||||
assert timings["Hybrid retrieval"].counts_toward_total is True
|
||||
assert received_engines == [engine]
|
||||
|
||||
|
||||
async def test_search_ebooks_skips_phrase_matching_when_disabled(mocker: MockerFixture) -> None:
|
||||
"""Phrase matching can be disabled for one search request."""
|
||||
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
||||
mocker.patch(
|
||||
"python.ebook_search.search.vector_candidates",
|
||||
return_value=[SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)],
|
||||
)
|
||||
mocker.patch("python.ebook_search.search.bm25_candidates", return_value=[])
|
||||
detect_mock = mocker.patch("python.ebook_search.search.query_phrase_matches")
|
||||
boost_mock = mocker.patch("python.ebook_search.search.apply_phrase_mention_boosts")
|
||||
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||
|
||||
response = await search_ebooks(engine, mocker.Mock(), "what is parallel", config, phrase_matching=False)
|
||||
|
||||
timing_names = {step.name for step in response.timings}
|
||||
assert [result.chunk_id for result in response.results] == [1]
|
||||
assert response.phrase_matches == ()
|
||||
assert "Protected phrase detection skipped" in timing_names
|
||||
assert "Phrase mention boost skipped" in timing_names
|
||||
detect_mock.assert_not_called()
|
||||
boost_mock.assert_not_called()
|
||||
|
||||
Reference in New Issue
Block a user