test_retrieval_query.py

"""
Unit tests for retrieval/query.py — all pure functions, no network required.
"""
import pytest
from retrieval.query import (
    normalize_for_match,
    is_procedural_question,
    is_image_intent,
    infer_source_hints,
    source_matches_hints,
    path_matches_hints,
    rewrite_query_for_retrieval,
)


class TestNormalizeForMatch:
    def test_lower_and_strip_punctuation(self):
        assert normalize_for_match("Hello, World!") == "hello world"

    def test_collapses_whitespace(self):
        assert normalize_for_match("  a   b  ") == "a b"

    def test_empty_string(self):
        assert normalize_for_match("") == ""

    def test_none(self):
        assert normalize_for_match(None) == ""

    def test_numbers_preserved(self):
        # normalize_for_match lowercases; "S1" becomes "s1" and digits are kept
        assert "s1" in normalize_for_match("S1 button")


class TestIsProceduralQuestion:
    def test_how_to(self):
        assert is_procedural_question("how to set alarm")

    def test_connect(self):
        assert is_procedural_question("how do I connect Bluetooth?")

    def test_setup(self):
        assert is_procedural_question("setup the watch")

    def test_non_procedural(self):
        assert not is_procedural_question("what is the battery life?")

    def test_alarm(self):
        assert is_procedural_question("set alarm")


class TestIsImageIntent:
    def test_photo_keyword(self):
        assert is_image_intent("show me the photo of the watch")

    def test_image_keyword(self):
        assert is_image_intent("show me an image")

    def test_jpg_keyword(self):
        assert is_image_intent("do you have a jpg of the coolfire?")

    def test_picture_keyword(self):
        assert is_image_intent("picture of the keychain")

    def test_show_without_image(self):
        # "show" without image keyword and without procedural context
        assert is_image_intent("show coolfire watch")

    def test_show_with_procedural_not_image(self):
        # "show me how to set" is procedural, not image
        assert not is_image_intent("show me how to set alarm")

    def test_normal_question(self):
        assert not is_image_intent("what is the warranty?")


class TestInferSourceHints:
    def test_coolfire(self):
        hints = infer_source_hints("coolfire watch alarm")
        assert any("coolfire" in h for h in hints)

    def test_handyspring(self):
        hints = infer_source_hints("handyspring water bottle")
        assert any("handyspring" in h for h in hints)

    def test_handyspring_typo(self):
        # The rewrite function fixes the typo; hints are inferred from rewritten query
        q = rewrite_query_for_retrieval("handysrping drink reminder")
        hints = infer_source_hints(q)
        assert any("handyspring" in h for h in hints)

    def test_keychain(self):
        hints = infer_source_hints("talking keychain time")
        assert any("key chain" in h or "keychain" in h for h in hints)

    def test_bluetooth_talking_watch(self):
        hints = infer_source_hints("how to use bluetooth talking watch")
        assert any("five senses" in h for h in hints)

    def test_no_hints_for_generic(self):
        hints = infer_source_hints("what is the warranty")
        assert hints == []


class TestSourceMatchesHints:
    def test_match(self):
        assert source_matches_hints("handyspring reminder.pdf", ["handyspring"])

    def test_no_match(self):
        assert not source_matches_hints("coolfire watch.pdf", ["handyspring"])

    def test_empty_hints(self):
        assert not source_matches_hints("anything.pdf", [])


class TestPathMatchesHints:
    def test_match_by_normalized_name(self):
        assert path_matches_hints("/path/to/handyspring drink reminder.jpg", ["handyspring"])

    def test_compact_match(self):
        # "keychain" matches "key chain" via compact comparison
        assert path_matches_hints("/path/FIVE SENSES talking key chain.jpg", ["key chain"])

    def test_no_match(self):
        assert not path_matches_hints("/path/coolfire watch.jpg", ["handyspring"])


class TestRewriteQueryForRetrieval:
    def test_fixes_handysrping(self):
        result = rewrite_query_for_retrieval("handysrping bottle")
        assert "handysrping" not in result
        assert "handyspring" in result

    def test_expands_coolfire(self):
        result = rewrite_query_for_retrieval("coolfire alarm")
        assert "vibration watch" in result or "vibrating watch" in result

    def test_expands_keychain(self):
        result = rewrite_query_for_retrieval("talking keychain time")
        assert "key chain" in result

    def test_empty_query(self):
        assert rewrite_query_for_retrieval("") == ""

    def test_no_duplicates(self):
        result = rewrite_query_for_retrieval("coolfire coolfire watch")
        tokens = result.split()
        assert len(tokens) == len(set(tokens))