our-claude-skills/custom-skills/31-notion-organizer/code/scripts/test_notion_search.py

#!/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 test_enrich_extracts_title_and_properties():
    """Candidate with properties → title and properties extracted."""
    from unittest.mock import MagicMock
    import notion_search

    notion = MagicMock()
    notion.blocks.children.list.return_value = {"results": []}  # no body

    candidate = {
        "id": "page-a",
        "url": "https://notion.so/a",
        "properties": {
            "Name": {
                "type": "title",
                "title": [{"plain_text": "Test Page"}],
            },
            "Status": {
                "type": "status",
                "status": {"name": "Done"},
            },
            "Topic": {
                "type": "multi_select",
                "multi_select": [{"name": "AI"}, {"name": "MCP"}],
            },
        },
    }
    enriched = notion_search.enrich_candidates(notion, [candidate])
    _assert(len(enriched) == 1, "one enriched candidate")
    e = enriched[0]
    _assert(e["title"] == "Test Page", "title extracted from title property")
    _assert(e["properties"]["Status"] == "Done", "status flattened to name")
    _assert(e["properties"]["Topic"] == ["AI", "MCP"], "multi_select flattened to names list")


def test_enrich_extracts_first_paragraph_excerpt():
    """First paragraph block → excerpt (200-char max)."""
    from unittest.mock import MagicMock
    import notion_search

    notion = MagicMock()
    notion.blocks.children.list.return_value = {
        "results": [
            {
                "type": "paragraph",
                "paragraph": {
                    "rich_text": [{"plain_text": "This is the first paragraph of the page."}]
                },
            }
        ]
    }
    candidate = {
        "id": "page-a", "url": "https://notion.so/a",
        "properties": {"Name": {"type": "title", "title": [{"plain_text": "Test"}]}},
    }
    enriched = notion_search.enrich_candidates(notion, [candidate])
    _assert(enriched[0]["excerpt"] == "This is the first paragraph of the page.",
            "excerpt is first paragraph plain text")


def test_enrich_falls_back_to_empty_excerpt():
    """Page leads with image/table only → excerpt is empty string, no crash."""
    from unittest.mock import MagicMock
    import notion_search

    notion = MagicMock()
    notion.blocks.children.list.return_value = {
        "results": [
            {"type": "image", "image": {}},
            {"type": "divider", "divider": {}},
        ]
    }
    candidate = {
        "id": "page-a", "url": "https://notion.so/a",
        "properties": {"Name": {"type": "title", "title": [{"plain_text": "Test"}]}},
    }
    enriched = notion_search.enrich_candidates(notion, [candidate])
    _assert(enriched[0]["excerpt"] == "", "empty excerpt when no text-bearing block")


def test_enrich_truncates_long_excerpt():
    """Excerpt is capped at 200 chars."""
    from unittest.mock import MagicMock
    import notion_search

    notion = MagicMock()
    long_text = "x" * 500
    notion.blocks.children.list.return_value = {
        "results": [
            {"type": "paragraph", "paragraph": {"rich_text": [{"plain_text": long_text}]}}
        ]
    }
    candidate = {
        "id": "page-a", "url": "https://notion.so/a",
        "properties": {"Name": {"type": "title", "title": [{"plain_text": "Test"}]}},
    }
    enriched = notion_search.enrich_candidates(notion, [candidate])
    _assert(len(enriched[0]["excerpt"]) == 200, "excerpt truncated to 200 chars")


def test_enrich_keeps_falsy_but_meaningful_values():
    """checkbox=False and number=0 are meaningful, not 'empty' — must survive the filter."""
    from unittest.mock import MagicMock
    import notion_search

    notion = MagicMock()
    notion.blocks.children.list.return_value = {"results": []}

    candidate = {
        "id": "page-a", "url": "https://notion.so/a",
        "properties": {
            "Name": {"type": "title", "title": [{"plain_text": "Test"}]},
            "Active": {"type": "checkbox", "checkbox": False},
            "Score": {"type": "number", "number": 0},
        },
    }
    enriched = notion_search.enrich_candidates(notion, [candidate])
    _assert(enriched[0]["properties"].get("Active") is False,
            "checkbox=False kept (not filtered as 'empty')")
    _assert(enriched[0]["properties"].get("Score") == 0,
            "number=0 kept (not filtered as 'empty')")


def test_rerank_orders_by_score_descending():
    """LLM returns scores; output sorted score-desc, top N."""
    import notion_search

    candidates = [
        {"id": "p1", "url": "u1", "title": "Page 1", "properties": {}, "excerpt": "x"},
        {"id": "p2", "url": "u2", "title": "Page 2", "properties": {}, "excerpt": "x"},
        {"id": "p3", "url": "u3", "title": "Page 3", "properties": {}, "excerpt": "x"},
    ]

    # LLM returns: p3 most relevant, then p1, then p2
    fake_llm = lambda prompt, **kw: '''[
        {"index": 2, "score": 0.95, "why": "best match"},
        {"index": 0, "score": 0.7, "why": "ok match"},
        {"index": 1, "score": 0.3, "why": "weak match"}
    ]'''

    ranked = notion_search.rerank("query", candidates, llm_caller=fake_llm, limit=3)
    _assert(ranked[0]["id"] == "p3", "highest score first")
    _assert(ranked[0]["relevance"] == 0.95, "score attached to result")
    _assert(ranked[0]["snippet"] == "best match", "why text attached as snippet")
    _assert(ranked[1]["id"] == "p1", "second place correct")
    _assert(ranked[2]["id"] == "p2", "lowest score last")


def test_rerank_respects_limit():
    """limit=2 → returns top 2 only."""
    import notion_search

    candidates = [
        {"id": f"p{i}", "url": f"u{i}", "title": f"Page {i}", "properties": {}, "excerpt": "x"}
        for i in range(5)
    ]
    fake_llm = lambda prompt, **kw: '''[
        {"index": 0, "score": 0.9, "why": "x"},
        {"index": 1, "score": 0.8, "why": "x"},
        {"index": 2, "score": 0.7, "why": "x"},
        {"index": 3, "score": 0.6, "why": "x"},
        {"index": 4, "score": 0.5, "why": "x"}
    ]'''
    ranked = notion_search.rerank("query", candidates, llm_caller=fake_llm, limit=2)
    _assert(len(ranked) == 2, "exactly limit results returned")


def test_rerank_falls_back_on_parse_error():
    """Non-JSON response → return candidates in input order, unranked (no scores)."""
    import notion_search

    candidates = [
        {"id": "p1", "url": "u1", "title": "Page 1", "properties": {}, "excerpt": "x"},
        {"id": "p2", "url": "u2", "title": "Page 2", "properties": {}, "excerpt": "x"},
    ]
    fake_llm = lambda prompt, **kw: "I cannot help with that."
    ranked = notion_search.rerank("query", candidates, llm_caller=fake_llm, limit=10)
    _assert(len(ranked) == 2, "all candidates returned in fallback")
    _assert(ranked[0]["id"] == "p1", "fallback preserves input order")
    _assert(ranked[0]["relevance"] is None, "no score in fallback")
    _assert(ranked[0]["snippet"] is None, "no snippet in fallback")


def test_rerank_falls_back_on_llm_exception():
    """LLM raises → fallback unranked."""
    import notion_search

    candidates = [{"id": "p1", "url": "u1", "title": "Page 1", "properties": {}, "excerpt": "x"}]

    def boom(prompt, **kw):
        raise RuntimeError("API down")

    ranked = notion_search.rerank("query", candidates, llm_caller=boom, limit=10)
    _assert(len(ranked) == 1, "candidate returned despite LLM failure")
    _assert(ranked[0]["relevance"] is None, "no score in fallback")


def test_pipeline_search_returns_json_serializable_results():
    """End-to-end pipeline returns results matching the spec's JSON schema."""
    from unittest.mock import MagicMock
    import json as json_mod
    import notion_search

    notion = MagicMock()
    notion.search.return_value = {
        "results": [
            {
                "id": "page-a", "object": "page", "url": "https://notion.so/a",
                "properties": {
                    "Name": {"type": "title", "title": [{"plain_text": "AI Agents Page"}]},
                    "Status": {"type": "status", "status": {"name": "Done"}},
                },
            }
        ]
    }
    notion.blocks.children.list.return_value = {
        "results": [
            {"type": "paragraph",
             "paragraph": {"rich_text": [{"plain_text": "Notes about AI agents."}]}}
        ]
    }

    expand_llm = lambda prompt, **kw: '["AI agents"]'
    rerank_llm = lambda prompt, **kw: '[{"index": 0, "score": 0.9, "why": "match"}]'

    results = notion_search.run_search(
        notion, "AI agents",
        expand_llm=expand_llm, rerank_llm=rerank_llm,
        limit=10, use_cache=False,
    )

    _assert(len(results) == 1, "one result")
    r = results[0]
    _assert(r["id"] == "page-a", "id present")
    _assert(r["title"] == "AI Agents Page", "title extracted")
    _assert(r["relevance"] == 0.9, "relevance score from rerank")
    _assert(r["properties"]["Status"] == "Done", "property included")
    _assert(r["excerpt"] == "Notes about AI agents.", "excerpt included")
    # Schema must be JSON-serializable
    json_mod.dumps(results)


def test_pipeline_no_rerank_skips_llm():
    """--no-rerank flag → results have null relevance, no rerank LLM call."""
    from unittest.mock import MagicMock
    import notion_search

    notion = MagicMock()
    notion.search.return_value = {
        "results": [
            {"id": "p1", "object": "page", "url": "u1",
             "properties": {"Name": {"type": "title", "title": [{"plain_text": "Page 1"}]}}}
        ]
    }
    notion.blocks.children.list.return_value = {"results": []}

    expand_llm = lambda prompt, **kw: '["query"]'

    rerank_calls = []
    def rerank_llm(prompt, **kw):
        rerank_calls.append(prompt)
        return "[]"

    results = notion_search.run_search(
        notion, "query",
        expand_llm=expand_llm, rerank_llm=rerank_llm,
        no_rerank=True, use_cache=False,
    )
    _assert(len(rerank_calls) == 0, "rerank LLM not called when --no-rerank")
    _assert(results[0]["relevance"] is None, "no relevance score when no rerank")


def test_pipeline_no_expand_uses_single_query():
    """--no-expand → expand_llm not called, Notion gets only original query."""
    from unittest.mock import MagicMock
    import notion_search

    notion = MagicMock()
    notion.search.return_value = {"results": []}

    expand_calls = []
    def expand_llm(prompt, **kw):
        expand_calls.append(prompt)
        return '["should not be used"]'

    rerank_llm = lambda prompt, **kw: "[]"

    notion_search.run_search(
        notion, "exact term",
        expand_llm=expand_llm, rerank_llm=rerank_llm,
        no_expand=True, use_cache=False,
    )
    _assert(len(expand_calls) == 0, "expand LLM not called when --no-expand")
    _assert(notion.search.call_count == 1, "Notion search called exactly once")


def test_pipeline_uses_cache_on_second_call():
    """Two identical pipeline calls → second one skips rerank LLM."""
    import tempfile
    from pathlib import Path
    from unittest.mock import MagicMock
    import notion_search

    notion = MagicMock()
    notion.search.return_value = {
        "results": [
            {"id": "p1", "object": "page", "url": "u1",
             "properties": {"Name": {"type": "title", "title": [{"plain_text": "Page 1"}]}}}
        ]
    }
    notion.blocks.children.list.return_value = {"results": []}

    expand_llm = lambda prompt, **kw: '["query"]'

    rerank_calls = []
    def rerank_llm(prompt, **kw):
        rerank_calls.append(prompt)
        return '[{"index": 0, "score": 0.5, "why": "x"}]'

    with tempfile.TemporaryDirectory() as tmpdir:
        cache_dir = Path(tmpdir)

        notion_search.run_search(
            notion, "query", expand_llm=expand_llm, rerank_llm=rerank_llm,
            use_cache=True, cache_dir=cache_dir,
        )
        notion_search.run_search(
            notion, "query", expand_llm=expand_llm, rerank_llm=rerank_llm,
            use_cache=True, cache_dir=cache_dir,
        )

    _assert(len(rerank_calls) == 1, "second call skipped rerank LLM (cache hit)")


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,
        test_enrich_extracts_title_and_properties,
        test_enrich_extracts_first_paragraph_excerpt,
        test_enrich_falls_back_to_empty_excerpt,
        test_enrich_truncates_long_excerpt,
        test_enrich_keeps_falsy_but_meaningful_values,
        test_rerank_orders_by_score_descending,
        test_rerank_respects_limit,
        test_rerank_falls_back_on_parse_error,
        test_rerank_falls_back_on_llm_exception,
        test_pipeline_search_returns_json_serializable_results,
        test_pipeline_no_rerank_skips_llm,
        test_pipeline_no_expand_uses_single_query,
        test_pipeline_uses_cache_on_second_call,
    ]
    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()