test_retrieval_extractive.py

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

from retrieval.extractive import build_extractive_answer


def _doc(content: str, source: str = "manual.pdf") -> Document:
    return Document(page_content=content, metadata={"source": source})


class TestBuildExtractiveAnswer:
    def test_returns_empty_for_no_docs(self):
        assert build_extractive_answer("how to set alarm", []) == ""

    def test_returns_empty_for_no_overlap(self):
        doc = _doc("completely unrelated information here without matching tokens.")
        result = build_extractive_answer("coolfire alarm setup", [doc])
        assert result == ""

    def test_picks_overlapping_sentence(self):
        doc = _doc(
            "Press S3 to enter Alarm mode. Then press S2 to edit the time. "
            "Use S1 to adjust the hour."
        )
        result = build_extractive_answer("how to set alarm", [doc])
        assert result != ""
        assert "Press" in result or "press" in result

    def test_action_keyword_bonus(self):
        doc = _doc(
            "The watch has a battery. Press S3 to set the alarm. "
            "The battery lasts 5 days."
        )
        result = build_extractive_answer("how to set alarm", [doc])
        # The "press" sentence should be ranked higher
        assert "Press" in result or "press" in result

    def test_max_three_lines(self):
        doc = _doc(
            "Press button 1 for alarm. Hold button 2 for time. "
            "Tap button 3 to confirm. Set mode with button 4. "
            "Connect Bluetooth via settings."
        )
        result = build_extractive_answer("how to set alarm", [doc])
        lines = [l for l in result.split("\n") if l.startswith("-")]
        assert len(lines) <= 3

    def test_deduplicates_sentences(self):
        same = "Press S3 to enter alarm mode."
        doc = _doc(f"{same}  {same}  {same}")
        result = build_extractive_answer("how to set alarm", [doc])
        lines = [l for l in result.split("\n") if l.startswith("-")]
        assert len(lines) == 1

    def test_empty_question(self):
        doc = _doc("Some text here.")
        assert build_extractive_answer("", [doc]) == ""