test_retrieval_fusion.py

"""
Unit tests for retrieval/fusion.py — pure functions, no network required.
"""
import pytest
from langchain_core.documents import Document

from retrieval.fusion import (
    doc_identity_key,
    reciprocal_rank_fusion,
    merge_retrieved_docs,
    prioritize_non_image_docs,
    rank_text_docs,
)


def _doc(content: str, source: str = "test.pdf", page=None, source_type: str = "pdf") -> Document:
    meta = {"source": source, "source_type": source_type}
    if page is not None:
        meta["page"] = page
    return Document(page_content=content, metadata=meta)


class TestDocIdentityKey:
    def test_uses_source_page_content(self):
        doc = _doc("hello world", source="a.pdf", page=1)
        key = doc_identity_key(doc)
        assert key == ("a.pdf", 1, "hello world")

    def test_truncates_content(self):
        long_content = "x" * 300
        doc = _doc(long_content, source="b.pdf")
        key = doc_identity_key(doc)
        assert len(key[2]) == 200


class TestReciprocalRankFusion:
    def test_presence_in_multiple_lists_boosts_score(self):
        """A doc appearing in both lists should outscore a doc in only one list."""
        doc_shared = _doc("shared content")
        doc_vector_only = _doc("vector only")

        # vector list has both; keyword list only has shared
        list1 = [doc_vector_only, doc_shared]
        list2 = [doc_shared]
        merged = reciprocal_rank_fusion([list1, list2])
        # shared appears at rank 2 in list1 AND rank 1 in list2; should beat
        # vector_only which appears at rank 1 in list1 only
        keys = [doc_identity_key(d) for d in merged]
        assert keys[0] == doc_identity_key(doc_shared)

    def test_limit_respected(self):
        docs = [_doc(f"content {i}") for i in range(10)]
        merged = reciprocal_rank_fusion([docs], limit=3)
        assert len(merged) == 3

    def test_deduplication(self):
        doc = _doc("same content")
        merged = reciprocal_rank_fusion([[doc], [doc]])
        assert len(merged) == 1


class TestMergeRetrievedDocs:
    def test_uses_rrf_when_both_lists_nonempty(self):
        vec = [_doc("vec A"), _doc("vec B")]
        kw = [_doc("kw A"), _doc("kw B")]
        merged = merge_retrieved_docs(vec, kw, limit=4)
        assert len(merged) <= 4
        # Should include docs from both lists
        contents = {d.page_content for d in merged}
        assert "vec A" in contents
        assert "kw A" in contents

    def test_falls_back_to_vector_only(self):
        vec = [_doc("vec A"), _doc("vec B")]
        merged = merge_retrieved_docs(vec, [], limit=2)
        assert len(merged) == 2

    def test_deduplication(self):
        doc = _doc("same")
        merged = merge_retrieved_docs([doc, doc], [], limit=10)
        assert len(merged) == 1


class TestPrioritizeNonImageDocs:
    def test_text_before_images(self):
        img_doc = _doc("image text", source_type="image")
        pdf_doc = _doc("pdf text", source_type="pdf")
        result = prioritize_non_image_docs([img_doc, pdf_doc], limit=2)
        assert result[0].metadata["source_type"] == "pdf"

    def test_limit(self):
        docs = [_doc(f"doc {i}") for i in range(5)]
        result = prioritize_non_image_docs(docs, limit=3)
        assert len(result) == 3

    def test_empty(self):
        assert prioritize_non_image_docs([], limit=5) == []


class TestRankTextDocs:
    def test_source_match_ranks_higher(self):
        question = "coolfire watch"
        matching = _doc("general text", source="coolfire watch.pdf")
        non_matching = _doc("another document", source="other.pdf")
        ranked = rank_text_docs(question, [non_matching, matching])
        assert ranked[0].metadata["source"] == "coolfire watch.pdf"

    def test_content_match_contributes(self):
        question = "handyspring reminder"
        doc_with = _doc("handyspring drink reminder filter bottle", source="x.pdf")
        doc_without = _doc("completely unrelated text here", source="y.pdf")
        ranked = rank_text_docs(question, [doc_without, doc_with])
        assert ranked[0].page_content == "handyspring drink reminder filter bottle"

    def test_no_query_tokens_returns_unchanged(self):
        docs = [_doc("text A"), _doc("text B")]
        result = rank_text_docs("", docs)
        assert result == docs