chore: restore original directory structure (project under code-review-graph-main/)
This commit is contained in:
AuraK Developer
@@ -0,0 +1,341 @@
|
||||
"""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.tmp.close() # release the handle before GraphStore reopens it on Windows
|
||||
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)
|
||||
|
||||
# --- _out_mode tracking ---
|
||||
|
||||
def test_out_mode_fts_only(self):
|
||||
"""_out_mode is 'fts' when only FTS contributes (no embeddings)."""
|
||||
rebuild_fts_index(self.store)
|
||||
out: list[str] = []
|
||||
results = hybrid_search(self.store, "authenticate", _out_mode=out)
|
||||
assert out == ["fts"]
|
||||
assert len(results) > 0
|
||||
|
||||
def test_out_mode_keyword(self):
|
||||
"""_out_mode is 'keyword' when FTS table is absent and no embeddings."""
|
||||
self.store._conn.execute("DROP TABLE IF EXISTS nodes_fts")
|
||||
self.store._conn.commit()
|
||||
out: list[str] = []
|
||||
results = hybrid_search(self.store, "authenticate", _out_mode=out)
|
||||
assert out == ["keyword"]
|
||||
assert len(results) > 0
|
||||
|
||||
def test_out_mode_keyword_no_results(self):
|
||||
"""_out_mode is 'none' when keyword fallback also returns 0 results."""
|
||||
self.store._conn.execute("DROP TABLE IF EXISTS nodes_fts")
|
||||
self.store._conn.commit()
|
||||
out: list[str] = []
|
||||
results = hybrid_search(self.store, "xyzzy_nonexistent_abc123", _out_mode=out)
|
||||
assert results == []
|
||||
assert out == ["none"]
|
||||
|
||||
def test_out_mode_semantic(self, monkeypatch):
|
||||
"""_out_mode is 'semantic' when only embeddings contribute."""
|
||||
import code_review_graph.search as search_mod
|
||||
|
||||
node_id = self.store._conn.execute(
|
||||
"SELECT id FROM nodes WHERE name = 'authenticate'"
|
||||
).fetchone()[0]
|
||||
|
||||
def fake_emb(store, query, limit=50, model=None, provider=None):
|
||||
return [(node_id, 0.9)]
|
||||
|
||||
monkeypatch.setattr(search_mod, "_embedding_search", fake_emb)
|
||||
out: list[str] = []
|
||||
results = hybrid_search(self.store, "authenticate", _out_mode=out)
|
||||
assert out == ["semantic"]
|
||||
assert len(results) > 0
|
||||
|
||||
def test_out_mode_hybrid(self, monkeypatch):
|
||||
"""_out_mode is 'hybrid' when both FTS and embeddings contribute."""
|
||||
import code_review_graph.search as search_mod
|
||||
|
||||
rebuild_fts_index(self.store)
|
||||
node_id = self.store._conn.execute(
|
||||
"SELECT id FROM nodes WHERE name = 'authenticate'"
|
||||
).fetchone()[0]
|
||||
|
||||
def fake_emb(store, query, limit=50, model=None, provider=None):
|
||||
return [(node_id, 0.9)]
|
||||
|
||||
monkeypatch.setattr(search_mod, "_embedding_search", fake_emb)
|
||||
out: list[str] = []
|
||||
results = hybrid_search(self.store, "authenticate", _out_mode=out)
|
||||
assert out == ["hybrid"]
|
||||
assert len(results) > 0
|
||||
|
||||
def test_out_mode_empty_query(self):
|
||||
"""_out_mode is 'none' for empty queries (no search ran)."""
|
||||
out: list[str] = []
|
||||
results = hybrid_search(self.store, "", _out_mode=out)
|
||||
assert results == []
|
||||
assert out == ["none"]
|
||||
|
||||
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)
|
||||
Reference in New Issue
Block a user