116 lines
3.6 KiB
Python
116 lines
3.6 KiB
Python
"""Unit tests for search/query_builder.py (pure logic only, no HTTP)"""
|
|
import sys
|
|
import os
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
|
|
|
# Stub env vars before importing search modules
|
|
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.schemas import Entities, Question
|
|
from search.query_builder import (
|
|
build_primary_query,
|
|
build_extra_dense_queries,
|
|
build_sparse_query,
|
|
build_entity_tokens,
|
|
)
|
|
|
|
|
|
def _q(**kwargs) -> Question:
|
|
defaults = dict(text="default question")
|
|
defaults.update(kwargs)
|
|
return Question(**defaults)
|
|
|
|
|
|
class TestBuildPrimaryQuery:
|
|
def test_uses_search_text_over_text(self):
|
|
q = _q(text="original", search_text="refined query")
|
|
assert build_primary_query(q) == "refined query"
|
|
|
|
def test_fallback_to_text(self):
|
|
q = _q(text="fallback text", search_text="")
|
|
assert build_primary_query(q) == "fallback text"
|
|
|
|
def test_strips_whitespace(self):
|
|
q = _q(text=" trimmed ")
|
|
assert build_primary_query(q) == "trimmed"
|
|
|
|
def test_collapses_internal_spaces(self):
|
|
q = _q(text="too many spaces")
|
|
result = build_primary_query(q)
|
|
assert " " not in result
|
|
|
|
|
|
class TestBuildExtraDenseQueries:
|
|
def test_no_extras_when_none(self):
|
|
q = _q(text="q")
|
|
assert build_extra_dense_queries(q) == []
|
|
|
|
def test_includes_variants(self):
|
|
q = _q(text="q", variants=["var1", "var2"])
|
|
extras = build_extra_dense_queries(q)
|
|
assert "var1" in extras
|
|
assert "var2" in extras
|
|
|
|
def test_includes_hyde(self):
|
|
q = _q(text="q", hyde=["hypothetical answer"])
|
|
extras = build_extra_dense_queries(q)
|
|
assert "hypothetical answer" in extras
|
|
|
|
def test_skips_empty_strings(self):
|
|
q = _q(text="q", variants=["", " ", "valid"])
|
|
extras = build_extra_dense_queries(q)
|
|
assert "" not in extras
|
|
assert " " not in extras
|
|
assert "valid" in extras
|
|
|
|
|
|
class TestBuildSparseQuery:
|
|
def test_uses_keywords_when_present(self):
|
|
q = _q(text="question", keywords=["go", "golang", "performance"])
|
|
result = build_sparse_query(q)
|
|
assert "go" in result
|
|
assert "golang" in result
|
|
|
|
def test_fallback_to_primary_when_no_keywords(self):
|
|
q = _q(text="fallback question", search_text="refined")
|
|
result = build_sparse_query(q)
|
|
assert result == "refined"
|
|
|
|
def test_empty_keywords_fallback(self):
|
|
q = _q(text="my question", keywords=[])
|
|
result = build_sparse_query(q)
|
|
assert result == "my question"
|
|
|
|
|
|
class TestBuildEntityTokens:
|
|
def test_no_entities(self):
|
|
q = _q(text="q")
|
|
assert build_entity_tokens(q) == []
|
|
|
|
def test_people_extracted(self):
|
|
q = _q(text="q", entities=Entities(people=["Alice", "Bob"]))
|
|
tokens = build_entity_tokens(q)
|
|
assert "Alice" in tokens
|
|
assert "Bob" in tokens
|
|
|
|
def test_all_entity_fields(self):
|
|
q = _q(
|
|
text="q",
|
|
entities=Entities(
|
|
people=["Alice"],
|
|
emails=["alice@corp.com"],
|
|
documents=["report.pdf"],
|
|
names=["Project X"],
|
|
links=["https://example.com"],
|
|
),
|
|
)
|
|
tokens = build_entity_tokens(q)
|
|
assert len(tokens) == 5
|
|
|
|
def test_strips_whitespace(self):
|
|
q = _q(text="q", entities=Entities(people=[" Alice "]))
|
|
tokens = build_entity_tokens(q)
|
|
assert "Alice" in tokens
|