wehub-resource-sync
271 lines
9.7 KiB
Python
271 lines
9.7 KiB
Python
"""Tests for the hybrid search engine."""
|
|
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
from code_review_graph.graph import GraphStore
|
|
from code_review_graph.parser import NodeInfo
|
|
from code_review_graph.search import (
|
|
detect_query_kind_boost,
|
|
hybrid_search,
|
|
rebuild_fts_index,
|
|
rrf_merge,
|
|
)
|
|
|
|
|
|
class TestHybridSearch:
|
|
def setup_method(self):
|
|
self.tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
|
|
self.store = GraphStore(self.tmp.name)
|
|
self._seed_data()
|
|
|
|
def teardown_method(self):
|
|
self.store.close()
|
|
Path(self.tmp.name).unlink(missing_ok=True)
|
|
|
|
def _seed_data(self):
|
|
"""Seed test nodes into the graph store."""
|
|
nodes = [
|
|
NodeInfo(
|
|
kind="Function", name="get_users", file_path="api.py",
|
|
line_start=1, line_end=20, language="python",
|
|
params="(db: Session)", return_type="list[User]",
|
|
),
|
|
NodeInfo(
|
|
kind="Function", name="create_user", file_path="api.py",
|
|
line_start=25, line_end=40, language="python",
|
|
params="(name: str, email: str)", return_type="User",
|
|
),
|
|
NodeInfo(
|
|
kind="Class", name="UserService", file_path="services.py",
|
|
line_start=1, line_end=100, language="python",
|
|
),
|
|
NodeInfo(
|
|
kind="Function", name="authenticate", file_path="auth.py",
|
|
line_start=5, line_end=30, language="python",
|
|
params="(token: str)", return_type="bool",
|
|
),
|
|
NodeInfo(
|
|
kind="Type", name="UserResponse", file_path="models.py",
|
|
line_start=1, line_end=15, language="python",
|
|
),
|
|
]
|
|
for node in nodes:
|
|
node_id = self.store.upsert_node(node, file_hash="abc123")
|
|
# Set signature for functions
|
|
if node.kind == "Function":
|
|
sig = f"def {node.name}{node.params or '()'} -> {node.return_type or 'None'}"
|
|
self.store._conn.execute(
|
|
"UPDATE nodes SET signature = ? WHERE id = ?", (sig, node_id)
|
|
)
|
|
self.store._conn.commit()
|
|
|
|
# --- rebuild_fts_index ---
|
|
|
|
def test_rebuild_fts_index(self):
|
|
"""rebuild_fts_index returns the correct count of indexed rows."""
|
|
count = rebuild_fts_index(self.store)
|
|
assert count == 5
|
|
|
|
def test_rebuild_fts_index_idempotent(self):
|
|
"""Rebuilding twice gives the same count."""
|
|
count1 = rebuild_fts_index(self.store)
|
|
count2 = rebuild_fts_index(self.store)
|
|
assert count1 == count2
|
|
|
|
# --- FTS search by name ---
|
|
|
|
def test_fts_search_by_name(self):
|
|
"""FTS search finds a node by its name."""
|
|
rebuild_fts_index(self.store)
|
|
results = hybrid_search(self.store, "get_users")
|
|
assert len(results) > 0
|
|
names = [r["name"] for r in results]
|
|
assert "get_users" in names
|
|
|
|
# --- FTS search by signature ---
|
|
|
|
def test_fts_search_by_signature(self):
|
|
"""FTS search finds a node by content in its signature."""
|
|
rebuild_fts_index(self.store)
|
|
results = hybrid_search(self.store, "Session")
|
|
assert len(results) > 0
|
|
# get_users has "Session" in its signature
|
|
names = [r["name"] for r in results]
|
|
assert "get_users" in names
|
|
|
|
# --- Kind boosting ---
|
|
|
|
def test_kind_boost_pascal_case(self):
|
|
"""PascalCase query boosts Class kind > 1.0."""
|
|
boosts = detect_query_kind_boost("UserService")
|
|
assert "Class" in boosts
|
|
assert boosts["Class"] > 1.0
|
|
|
|
def test_kind_boost_snake_case(self):
|
|
"""snake_case query boosts Function kind > 1.0."""
|
|
boosts = detect_query_kind_boost("get_users")
|
|
assert "Function" in boosts
|
|
assert boosts["Function"] > 1.0
|
|
|
|
def test_kind_boost_dotted(self):
|
|
"""Dotted query boosts qualified name matches."""
|
|
boosts = detect_query_kind_boost("api.get_users")
|
|
assert "_qualified" in boosts
|
|
assert boosts["_qualified"] > 1.0
|
|
|
|
def test_kind_boost_empty(self):
|
|
"""Empty query returns no boosts."""
|
|
boosts = detect_query_kind_boost("")
|
|
assert boosts == {}
|
|
|
|
def test_kind_boost_all_uppercase(self):
|
|
"""ALL_CAPS should not trigger PascalCase boost."""
|
|
boosts = detect_query_kind_boost("HTTP_STATUS")
|
|
assert "Class" not in boosts
|
|
# But should trigger snake_case boost
|
|
assert "Function" in boosts
|
|
|
|
# --- RRF merge ---
|
|
|
|
def test_rrf_merge(self):
|
|
"""Node appearing in both lists ranks highest after RRF merge."""
|
|
list_a = [(1, 10.0), (2, 8.0), (3, 6.0)]
|
|
list_b = [(2, 9.0), (4, 7.0), (1, 5.0)]
|
|
|
|
merged = rrf_merge(list_a, list_b)
|
|
ids = [item_id for item_id, _ in merged]
|
|
|
|
# Items 1 and 2 appear in both lists, so they should be top-ranked
|
|
assert ids[0] in (1, 2)
|
|
assert ids[1] in (1, 2)
|
|
# ID 2 is rank 0+0 in list_b and rank 1 in list_a
|
|
# ID 1 is rank 0 in list_a and rank 2 in list_b
|
|
# So ID 2 should rank higher: 1/(60+1+1) + 1/(60+0+1) vs 1/(60+0+1) + 1/(60+2+1)
|
|
assert ids[0] == 2
|
|
|
|
def test_rrf_merge_single_list(self):
|
|
"""RRF merge with a single list preserves order."""
|
|
single = [(10, 5.0), (20, 3.0), (30, 1.0)]
|
|
merged = rrf_merge(single)
|
|
ids = [item_id for item_id, _ in merged]
|
|
assert ids == [10, 20, 30]
|
|
|
|
def test_rrf_merge_empty(self):
|
|
"""RRF merge with empty lists returns empty."""
|
|
merged = rrf_merge([], [])
|
|
assert merged == []
|
|
|
|
# --- Fallback to keyword search ---
|
|
|
|
def test_fallback_to_keyword(self):
|
|
"""Works without FTS index by falling back to keyword LIKE matching."""
|
|
# Do NOT rebuild FTS index — drop it if it exists
|
|
try:
|
|
self.store._conn.execute("DROP TABLE IF EXISTS nodes_fts")
|
|
self.store._conn.commit()
|
|
except Exception:
|
|
pass
|
|
|
|
results = hybrid_search(self.store, "authenticate")
|
|
assert len(results) > 0
|
|
names = [r["name"] for r in results]
|
|
assert "authenticate" in names
|
|
|
|
# --- Empty query ---
|
|
|
|
def test_empty_query_handled(self):
|
|
"""Empty query returns empty results without crashing."""
|
|
results = hybrid_search(self.store, "")
|
|
assert results == []
|
|
|
|
def test_whitespace_query_handled(self):
|
|
"""Whitespace-only query returns empty results."""
|
|
results = hybrid_search(self.store, " ")
|
|
assert results == []
|
|
|
|
# --- Return fields ---
|
|
|
|
def test_hybrid_search_returns_expected_fields(self):
|
|
"""All expected fields are present in search results."""
|
|
rebuild_fts_index(self.store)
|
|
results = hybrid_search(self.store, "get_users")
|
|
assert len(results) > 0
|
|
|
|
expected_fields = {
|
|
"name", "qualified_name", "kind", "file_path",
|
|
"line_start", "line_end", "language", "params",
|
|
"return_type", "signature", "score",
|
|
}
|
|
for result in results:
|
|
assert expected_fields.issubset(result.keys()), (
|
|
f"Missing fields: {expected_fields - result.keys()}"
|
|
)
|
|
|
|
# --- Kind filtering ---
|
|
|
|
def test_kind_filter(self):
|
|
"""Kind parameter filters results to only that kind."""
|
|
rebuild_fts_index(self.store)
|
|
results = hybrid_search(self.store, "User", kind="Class")
|
|
for r in results:
|
|
assert r["kind"] == "Class"
|
|
|
|
# --- Context file boosting ---
|
|
|
|
def test_context_file_boost(self):
|
|
"""Nodes in context_files get boosted above others."""
|
|
rebuild_fts_index(self.store)
|
|
|
|
# Search for "user" which matches multiple nodes
|
|
results_with_ctx = hybrid_search(
|
|
self.store, "user", context_files=["api.py"]
|
|
)
|
|
|
|
# Find get_users in both result sets
|
|
if results_with_ctx:
|
|
api_nodes = [r for r in results_with_ctx if r["file_path"] == "api.py"]
|
|
if api_nodes:
|
|
# api.py nodes should have a score boost
|
|
api_score = api_nodes[0]["score"]
|
|
assert api_score > 0
|
|
|
|
# --- Limit parameter ---
|
|
|
|
def test_limit_respected(self):
|
|
"""Search respects the limit parameter."""
|
|
rebuild_fts_index(self.store)
|
|
results = hybrid_search(self.store, "user", limit=2)
|
|
assert len(results) <= 2
|
|
|
|
# --- FTS5 injection safety ---
|
|
|
|
def test_fts_query_with_special_chars(self):
|
|
"""FTS5 special characters are safely handled."""
|
|
rebuild_fts_index(self.store)
|
|
# These should not crash — FTS5 operators like AND, OR, NOT, *, etc.
|
|
for dangerous_query in ['OR user', 'NOT thing', 'user*', '"user"', 'a AND b']:
|
|
results = hybrid_search(self.store, dangerous_query)
|
|
# Just assert no exception was raised
|
|
assert isinstance(results, list)
|
|
|
|
def test_fts_rebuild_is_atomic(self):
|
|
"""Regression test for #259: rebuild_fts_index must wrap the DROP +
|
|
CREATE + INSERT sequence in a single transaction so a crash between
|
|
DROP and CREATE cannot leave the DB without an FTS table."""
|
|
# Build, rebuild, then verify the table exists and is queryable.
|
|
rebuild_fts_index(self.store)
|
|
|
|
# Verify the FTS table exists and has rows.
|
|
conn = self.store._conn
|
|
count = conn.execute("SELECT count(*) FROM nodes_fts").fetchone()[0]
|
|
assert count > 0
|
|
|
|
# Rebuild again — must not raise and must leave the table intact.
|
|
new_count = rebuild_fts_index(self.store)
|
|
assert new_count == count
|
|
|
|
# Verify search still works after double-rebuild.
|
|
results = hybrid_search(self.store, "auth")
|
|
assert isinstance(results, list)
|