learningcircuit--local-deep-research/tests/web_search_engines/test_engine_groups.py

"""Tests for the search-engine selector grouping/ordering logic."""

from local_deep_research.web_search_engines.engine_groups import (
    FAVORITES_GROUP_KEY,
    SEARCH_ENGINE_GROUPS,
    classify_engine_group,
    effective_group,
    group_icon,
    group_label,
    group_order,
    is_collection_engine,
)


class TestIsCollectionEngine:
    def test_library_aggregate_is_collection(self):
        assert is_collection_engine("library") is True

    def test_per_collection_pseudo_engine(self):
        assert is_collection_engine("collection_42") is True

    def test_regular_engine_is_not_collection(self):
        assert is_collection_engine("arxiv") is False
        assert is_collection_engine("tavily") is False


class TestClassifyEngineGroup:
    def test_collections_win_over_category(self):
        # A collection is still "collections" even though it is Local RAG.
        assert (
            classify_engine_group("collection_1", "Local RAG", False)
            == "collections"
        )
        assert (
            classify_engine_group("library", "Local RAG", False)
            == "collections"
        )

    def test_academic(self):
        # Free and keyed academic engines both land in academic.
        assert classify_engine_group("arxiv", "Scientific", False) == "academic"
        assert (
            classify_engine_group("nasa_ads", "Scientific", True) == "academic"
        )

    def test_local_rag_non_collection(self):
        assert (
            classify_engine_group("elasticsearch", "Local RAG", False)
            == "local_rag"
        )

    def test_books_code_news_have_own_bands_regardless_of_key(self):
        assert classify_engine_group("gutenberg", "Books", False) == "books"
        # github is Code AND requires a key -> still the Code band, not api_key.
        assert classify_engine_group("github", "Code", True) == "code"
        assert classify_engine_group("wikinews", "News", False) == "news"

    def test_generic_web_split_by_api_key(self):
        assert classify_engine_group("ddg", "Web Search", False) == "no_api_key"
        assert classify_engine_group("tavily", "Web Search", True) == "api_key"
        # The default "Search" category also splits by key.
        assert (
            classify_engine_group("wikipedia", "Search", False) == "no_api_key"
        )


class TestEffectiveGroup:
    def test_favorite_overrides_category(self):
        assert effective_group("academic", True) == FAVORITES_GROUP_KEY

    def test_non_favorite_keeps_base(self):
        assert effective_group("academic", False) == "academic"


class TestGroupOrdering:
    def test_favorites_sorts_first(self):
        assert group_order("favorites") == 0

    def test_full_band_order_is_as_specified(self):
        expected = [
            "favorites",
            "collections",
            "academic",
            "local_rag",
            "books",
            "code",
            "news",
            "no_api_key",
            "api_key",
        ]
        assert [g.key for g in SEARCH_ENGINE_GROUPS] == expected
        # group_order is strictly increasing in that sequence.
        orders = [group_order(key) for key in expected]
        assert orders == sorted(orders)
        assert orders == list(range(len(expected)))

    def test_unknown_group_sorts_last(self):
        assert group_order("nope") == len(SEARCH_ENGINE_GROUPS)


class TestGroupLabelAndIcon:
    def test_known_labels(self):
        assert group_label("no_api_key") == "No API key"
        assert group_label("api_key") == "API key"
        assert group_label("favorites") == "Favorites"

    def test_unknown_label_falls_back_to_key(self):
        assert group_label("mystery") == "mystery"

    def test_icons(self):
        assert group_icon("favorites") == "⭐"
        assert group_icon("mystery") == ""


class TestEndToEndOrdering:
    """A representative engine set sorts into the intended band order."""

    def test_engines_sort_into_bands(self):
        # (engine_id, category, requires_api_key, is_favorite)
        engines = [
            ("tavily", "Web Search", True, False),  # api_key
            ("arxiv", "Scientific", False, False),  # academic
            ("collection_7", "Local RAG", False, False),  # collections
            ("ddg", "Web Search", False, False),  # no_api_key
            ("github", "Code", True, False),  # code
            ("elasticsearch", "Local RAG", False, False),  # local_rag
            ("brave", "Web Search", True, True),  # favorite (overlay)
            ("gutenberg", "Books", False, False),  # books
            ("wikinews", "News", False, False),  # news
        ]
        options = []
        for engine_id, category, key, fav in engines:
            base = classify_engine_group(engine_id, category, key)
            shown = effective_group(base, fav)
            options.append(
                {
                    "value": engine_id,
                    "label": engine_id,
                    "group_order": group_order(shown),
                }
            )
        options.sort(key=lambda x: (x["group_order"], x["label"].lower()))
        ordered_ids = [o["value"] for o in options]
        assert ordered_ids == [
            "brave",  # favorites
            "collection_7",  # collections
            "arxiv",  # academic
            "elasticsearch",  # local_rag
            "gutenberg",  # books
            "github",  # code
            "wikinews",  # news
            "ddg",  # no_api_key
            "tavily",  # api_key
        ]