#!/usr/bin/env python3 """Tests for notion_search.py — run with `python3 test_notion_search.py`.""" import sys from pathlib import Path sys.path.insert(0, str(Path(__file__).parent)) def _assert(cond, msg): if not cond: print(f" ✗ FAIL: {msg}") raise SystemExit(1) print(f" ✓ {msg}") def test_call_claude_dispatches_to_sdk_when_available(): """When SDK is available and ANTHROPIC_API_KEY is set, use SDK path.""" import os from unittest.mock import patch import _search_llm with patch.object(_search_llm, "_have_anthropic_sdk", return_value=True), \ patch.dict(os.environ, {"ANTHROPIC_API_KEY": "sk-fake"}), \ patch.object(_search_llm, "_call_via_sdk", return_value="sdk-response") as sdk_mock, \ patch.object(_search_llm, "_call_via_cli", return_value="cli-response") as cli_mock: result = _search_llm.call_claude("hello") _assert(result == "sdk-response", "SDK path returned its response") _assert(sdk_mock.called, "SDK path was invoked") _assert(not cli_mock.called, "CLI path was NOT invoked") def test_call_claude_falls_back_to_cli_when_no_sdk(): """When SDK is missing but CLI is available, use CLI path.""" from unittest.mock import patch import _search_llm with patch.object(_search_llm, "_have_anthropic_sdk", return_value=False), \ patch.object(_search_llm, "_have_claude_cli", return_value=True), \ patch.object(_search_llm, "_call_via_cli", return_value="cli-response") as cli_mock: result = _search_llm.call_claude("hello") _assert(result == "cli-response", "CLI path returned its response") _assert(cli_mock.called, "CLI path was invoked") def test_call_claude_raises_when_neither_available(): """Both clients missing → RuntimeError with setup hint.""" from unittest.mock import patch import _search_llm with patch.object(_search_llm, "_have_anthropic_sdk", return_value=False), \ patch.object(_search_llm, "_have_claude_cli", return_value=False): try: _search_llm.call_claude("hello") _assert(False, "should have raised RuntimeError") except RuntimeError as exc: _assert("ANTHROPIC_API_KEY" in str(exc) or "claude" in str(exc).lower(), "error mentions setup options") def test_cache_miss_returns_none(): import tempfile from pathlib import Path import _search_cache with tempfile.TemporaryDirectory() as tmpdir: result = _search_cache.cache_get("query", ["id1", "id2"], cache_dir=Path(tmpdir)) _assert(result is None, "cache miss returns None") def test_cache_hit_returns_cached_results(): import tempfile from pathlib import Path import _search_cache with tempfile.TemporaryDirectory() as tmpdir: cache_dir = Path(tmpdir) results = [{"id": "abc", "title": "Test"}] _search_cache.cache_put("query", ["id1", "id2"], results, cache_dir=cache_dir) hit = _search_cache.cache_get("query", ["id1", "id2"], cache_dir=cache_dir) _assert(hit == results, "cache hit returns the stored results") def test_cache_miss_on_different_candidates(): import tempfile from pathlib import Path import _search_cache with tempfile.TemporaryDirectory() as tmpdir: cache_dir = Path(tmpdir) _search_cache.cache_put("query", ["id1", "id2"], [{"x": 1}], cache_dir=cache_dir) miss = _search_cache.cache_get("query", ["id1", "id3"], cache_dir=cache_dir) _assert(miss is None, "different candidate set hashes to different key") def test_cache_miss_after_ttl_expires(): import tempfile from pathlib import Path import _search_cache with tempfile.TemporaryDirectory() as tmpdir: cache_dir = Path(tmpdir) _search_cache.cache_put("query", ["id1"], [{"x": 1}], cache_dir=cache_dir) # TTL=0 means immediately expired miss = _search_cache.cache_get("query", ["id1"], cache_dir=cache_dir, ttl_seconds=0) _assert(miss is None, "expired cache returns None") def test_cache_handles_corrupted_file(): import tempfile import hashlib from pathlib import Path import _search_cache with tempfile.TemporaryDirectory() as tmpdir: cache_dir = Path(tmpdir) cache_dir.mkdir(exist_ok=True) # Write garbage to the expected key path key = hashlib.sha256("query|id1".encode()).hexdigest() (cache_dir / f"{key}.json").write_text("not json {{{") result = _search_cache.cache_get("query", ["id1"], cache_dir=cache_dir) _assert(result is None, "corrupted cache file returns None") def test_expand_query_returns_variants_with_original_first(): """LLM mock returns variants; expansion includes original verbatim as first item.""" import notion_search fake_llm = lambda prompt, **kw: '["AI agents", "multi-agent systems", "autonomous LLMs"]' variants = notion_search.expand_query("AI agents", llm_caller=fake_llm) _assert(variants[0] == "AI agents", "original query is first variant") _assert(len(variants) >= 2, "at least 2 variants returned") _assert("multi-agent systems" in variants, "LLM-suggested variant included") def test_expand_query_dedupes_and_caps(): """Duplicates removed; total ≤ max_variants.""" import notion_search fake_llm = lambda prompt, **kw: '["AI", "AI", "agents", "agents", "systems", "tools", "models"]' variants = notion_search.expand_query("AI", llm_caller=fake_llm, max_variants=3) _assert(len(variants) <= 3, "capped at max_variants=3") _assert(len(variants) == len(set(variants)), "no duplicates in result") def test_expand_query_falls_back_on_invalid_json(): """LLM returns prose instead of JSON → fall back to [original].""" import notion_search fake_llm = lambda prompt, **kw: "I'm sorry, I can't help with that." variants = notion_search.expand_query("AI agents", llm_caller=fake_llm) _assert(variants == ["AI agents"], "non-JSON response falls back to original only") def test_expand_query_falls_back_on_llm_exception(): """LLM raises → fall back to [original].""" import notion_search def boom(prompt, **kw): raise RuntimeError("API error") variants = notion_search.expand_query("AI agents", llm_caller=boom) _assert(variants == ["AI agents"], "LLM exception falls back to original only") def test_search_unions_and_dedupes_workspace(): """Two variants returning overlapping pages produce a deduped candidate set.""" from unittest.mock import MagicMock import notion_search notion = MagicMock() notion.search.side_effect = [ {"results": [ {"id": "page-a", "object": "page", "url": "https://notion.so/a"}, {"id": "page-b", "object": "page", "url": "https://notion.so/b"}, ]}, {"results": [ {"id": "page-b", "object": "page", "url": "https://notion.so/b"}, {"id": "page-c", "object": "page", "url": "https://notion.so/c"}, ]}, ] candidates = notion_search.search_candidates( notion, ["query1", "query2"], databases=None ) ids = [c["id"] for c in candidates] _assert(set(ids) == {"page-a", "page-b", "page-c"}, "union of pages from both variants") _assert(len(ids) == 3, "duplicate page-b deduped") def test_search_caps_candidates_at_30(): """If total candidates exceed 30, cap to 30.""" from unittest.mock import MagicMock import notion_search notion = MagicMock() notion.search.return_value = { "results": [ {"id": f"page-{i:02d}", "object": "page", "url": f"https://notion.so/{i}"} for i in range(40) ] } candidates = notion_search.search_candidates(notion, ["query"], databases=None) _assert(len(candidates) == 30, f"capped to 30 (got {len(candidates)})") def test_search_uses_data_source_query_when_databases_specified(): """--databases flag → per-DB data_sources.query instead of workspace search.""" from unittest.mock import MagicMock, patch import notion_search notion = MagicMock() notion.data_sources.query.return_value = { "results": [{"id": "page-a", "url": "https://notion.so/a"}] } with patch("notion_search.compat.resolve_data_source_id", side_effect=lambda c, ds: ds): candidates = notion_search.search_candidates( notion, ["query"], databases=["db-id-1"] ) _assert(notion.data_sources.query.called, "data_sources.query was called") _assert(not notion.search.called, "workspace search NOT called when databases specified") _assert(len(candidates) == 1, "candidate from data source returned") def test_search_filters_only_pages(): """search() may return databases too; we keep only object='page'.""" from unittest.mock import MagicMock import notion_search notion = MagicMock() notion.search.return_value = {"results": [ {"id": "page-a", "object": "page", "url": "https://notion.so/a"}, {"id": "db-1", "object": "database", "url": "https://notion.so/db1"}, ]} candidates = notion_search.search_candidates(notion, ["query"], databases=None) ids = [c["id"] for c in candidates] _assert("page-a" in ids, "page kept") _assert("db-1" not in ids, "database object filtered out") def test_search_passes_property_filter_to_data_sources_query(): """--filter JSON is passed through to data_sources.query as the `filter` parameter.""" from unittest.mock import MagicMock, patch import notion_search notion = MagicMock() notion.data_sources.query.return_value = {"results": []} prop_filter = {"property": "Status", "status": {"equals": "Done"}} with patch("notion_search.compat.resolve_data_source_id", side_effect=lambda c, ds: ds): notion_search.search_candidates( notion, ["query"], databases=["db-id-1"], prop_filter=prop_filter, ) _assert(notion.data_sources.query.called, "data_sources.query was called") call_kwargs = notion.data_sources.query.call_args.kwargs _assert(call_kwargs.get("filter") == prop_filter, "prop_filter passed through as `filter` keyword argument") def run_all(): tests = [ test_call_claude_dispatches_to_sdk_when_available, test_call_claude_falls_back_to_cli_when_no_sdk, test_call_claude_raises_when_neither_available, test_cache_miss_returns_none, test_cache_hit_returns_cached_results, test_cache_miss_on_different_candidates, test_cache_miss_after_ttl_expires, test_cache_handles_corrupted_file, test_expand_query_returns_variants_with_original_first, test_expand_query_dedupes_and_caps, test_expand_query_falls_back_on_invalid_json, test_expand_query_falls_back_on_llm_exception, test_search_unions_and_dedupes_workspace, test_search_caps_candidates_at_30, test_search_uses_data_source_query_when_databases_specified, test_search_filters_only_pages, test_search_passes_property_filter_to_data_sources_query, ] for t in tests: print(f"\n{t.__name__}") t() print("\n" + "=" * 50) print(f"✅ All {len(tests)} tests passed") print("=" * 50) if __name__ == "__main__": run_all()