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