tests/test_retriever.py

"""
Tests for DocMind Hybrid Retriever — RRF fusion logic.
"""

import sys
import os

# Add project root to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from pipeline.retriever import rrf_fusion


class TestRRFFusion:
    """Test Reciprocal Rank Fusion implementation."""

    def test_basic_fusion(self):
        """Two lists with overlapping items should sum scores."""
        bm25 = [("c1", 5.0), ("c2", 4.0), ("c3", 3.0)]
        dense = [("c2", 0.9), ("c3", 0.8), ("c4", 0.7)]

        result = rrf_fusion(bm25, dense, k=60, top_k=3)

        # c2 should be first (appears in both lists)
        ids = [cid for cid, _ in result]
        assert ids[0] == "c2", f"Expected c2 first, got {ids[0]}"
        assert len(result) == 3

    def test_no_overlap(self):
        """Disjoint lists should include items from both."""
        bm25 = [("a1", 5.0), ("a2", 4.0)]
        dense = [("b1", 0.9), ("b2", 0.8)]

        result = rrf_fusion(bm25, dense, k=60, top_k=4)

        ids = {cid for cid, _ in result}
        assert ids == {"a1", "a2", "b1", "b2"}

    def test_complete_overlap(self):
        """Identical lists should give higher scores than single-list items."""
        items = [("c1", 5.0), ("c2", 4.0)]
        result = rrf_fusion(items, items, k=60, top_k=2)

        # Each item appears in both lists — scores should be doubled
        for _, score in result:
            # score = 2 * 1/(60+rank), which is > 1/(60+rank)
            assert score > 1.0 / 62  # More than a single-list entry

    def test_empty_inputs(self):
        """Empty input lists should return empty results."""
        result = rrf_fusion([], [], k=60, top_k=8)
        assert result == []

    def test_one_empty_list(self):
        """One empty list should still return items from the other."""
        bm25 = [("c1", 5.0), ("c2", 4.0)]
        result = rrf_fusion(bm25, [], k=60, top_k=2)

        assert len(result) == 2
        ids = [cid for cid, _ in result]
        assert "c1" in ids

    def test_top_k_limit(self):
        """Should respect the top_k parameter."""
        bm25 = [(f"c{i}", float(10 - i)) for i in range(10)]
        dense = [(f"c{i}", float(10 - i) / 10) for i in range(10)]

        result = rrf_fusion(bm25, dense, k=60, top_k=3)
        assert len(result) == 3

    def test_score_ordering(self):
        """Results should be sorted by descending RRF score."""
        bm25 = [("c1", 5.0), ("c2", 4.0), ("c3", 3.0)]
        dense = [("c1", 0.9), ("c2", 0.8), ("c3", 0.7)]

        result = rrf_fusion(bm25, dense, k=60, top_k=3)

        scores = [score for _, score in result]
        assert scores == sorted(scores, reverse=True), "Scores should be descending"

    def test_rrf_score_formula(self):
        """Verify the exact RRF score computation."""
        bm25 = [("c1", 5.0)]  # rank 1
        dense = [("c1", 0.9)]  # rank 1

        result = rrf_fusion(bm25, dense, k=60, top_k=1)

        expected_score = 1.0 / (60 + 1) + 1.0 / (60 + 1)
        assert abs(result[0][1] - expected_score) < 1e-10, \
            f"Expected {expected_score}, got {result[0][1]}"


# ── Run tests ───────────────────────────────────────────────────────

if __name__ == "__main__":
    test = TestRRFFusion()
    test_methods = [m for m in dir(test) if m.startswith("test_")]

    passed = 0
    failed = 0
    for method_name in test_methods:
        try:
            getattr(test, method_name)()
            print(f"  [PASS] {method_name}")
            passed += 1
        except AssertionError as e:
            print(f"  [FAIL] {method_name}: {e}")
            failed += 1
        except Exception as e:
            print(f"  [ERROR] {method_name}: {type(e).__name__}: {e}")
            failed += 1

    print(f"\n{'='*40}")
    print(f"Results: {passed} passed, {failed} failed")