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 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,
    ]
    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()