"""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 ]