basicmachines-co-basic-memory/tests/api/v2/test_search_router.py

"""Tests for v2 search router endpoints."""

from datetime import datetime, timezone

import pytest
from httpx import AsyncClient
from pathlib import Path

from basic_memory.deps.services import get_search_service_v2_external
from basic_memory.models import Project
from basic_memory.repository.search_index_row import SearchIndexRow
from basic_memory.repository.semantic_errors import (
    SemanticDependenciesMissingError,
    SemanticSearchDisabledError,
)


async def create_test_entity(
    test_project, entity_data, entity_repository, search_service, file_service
):
    """Helper to create an entity with file and index it."""
    # Create file
    test_content = f"# {entity_data['title']}\n\nTest content"
    file_path = Path(test_project.path) / entity_data["file_path"]
    file_path.parent.mkdir(parents=True, exist_ok=True)
    await file_service.write_file(file_path, test_content)

    # Create entity
    entity = await entity_repository.create(entity_data)

    # Index for search
    await search_service.index_entity(entity)

    return entity


@pytest.mark.asyncio
async def test_search_entities(
    client: AsyncClient,
    test_project: Project,
    v2_project_url: str,
    entity_repository,
    search_service,
    file_service,
):
    """Test searching for entities."""
    # Create a test entity
    entity_data = {
        "title": "Searchable Entity",
        "note_type": "note",
        "content_type": "text/markdown",
        "file_path": "searchable.md",
        "checksum": "search123",
    }
    await create_test_entity(
        test_project, entity_data, entity_repository, search_service, file_service
    )

    # Search for the entity
    response = await client.post(f"{v2_project_url}/search/", json={"text": "Searchable"})

    assert response.status_code == 200
    data = response.json()

    # Verify response structure
    assert "results" in data
    assert "current_page" in data
    assert "page_size" in data


@pytest.mark.asyncio
async def test_search_with_pagination(
    client: AsyncClient,
    test_project: Project,
    v2_project_url: str,
    entity_repository,
    search_service,
    file_service,
):
    """Test search with pagination parameters."""
    # Create multiple test entities
    for i in range(5):
        entity_data = {
            "title": f"Search Entity {i}",
            "note_type": "note",
            "content_type": "text/markdown",
            "file_path": f"search_{i}.md",
            "checksum": f"searchsum{i}",
        }
        await create_test_entity(
            test_project, entity_data, entity_repository, search_service, file_service
        )

    # Search with pagination
    response = await client.post(
        f"{v2_project_url}/search/",
        json={"text": "Search Entity"},
        params={"page": 1, "page_size": 3},
    )

    assert response.status_code == 200
    data = response.json()
    assert data["current_page"] == 1
    assert data["page_size"] == 3
    assert data["total"] == 5
    assert data["has_more"] is True

    response = await client.post(
        f"{v2_project_url}/search/",
        json={"text": "Search Entity"},
        params={"page": 2, "page_size": 3},
    )

    assert response.status_code == 200
    data = response.json()
    assert data["current_page"] == 2
    assert data["page_size"] == 3
    assert data["total"] == 5
    assert data["has_more"] is False
    assert len(data["results"]) == 2


@pytest.mark.asyncio
async def test_search_with_item_type_filter_returns_total(
    client: AsyncClient,
    test_project: Project,
    v2_project_url: str,
    entity_repository,
    search_service,
    file_service,
):
    """Metadata-only graph searches should include exact totals for pagination."""
    for i in range(5):
        entity_data = {
            "title": f"Structured Entity {i}",
            "note_type": "note",
            "content_type": "text/markdown",
            "file_path": f"structured_{i}.md",
            "checksum": f"structuredsum{i}",
        }
        await create_test_entity(
            test_project, entity_data, entity_repository, search_service, file_service
        )

    response = await client.post(
        f"{v2_project_url}/search/",
        json={"entity_types": ["entity"]},
        params={"page": 1, "page_size": 3},
    )

    assert response.status_code == 200
    data = response.json()
    assert data["total"] == 5
    assert data["has_more"] is True
    assert len(data["results"]) == 3


@pytest.mark.asyncio
async def test_search_by_permalink(
    client: AsyncClient,
    test_project: Project,
    v2_project_url: str,
    entity_repository,
    search_service,
    file_service,
):
    """Test searching by permalink."""
    # Create a test entity with permalink
    entity_data = {
        "title": "Permalink Search",
        "note_type": "note",
        "content_type": "text/markdown",
        "file_path": "permalink_search.md",
        "checksum": "perm123",
        "permalink": "permalink-search",
    }
    await create_test_entity(
        test_project, entity_data, entity_repository, search_service, file_service
    )

    # Search by permalink
    response = await client.post(
        f"{v2_project_url}/search/", json={"permalink": "permalink-search"}
    )

    assert response.status_code == 200
    data = response.json()
    assert "results" in data


@pytest.mark.asyncio
async def test_search_by_title(
    client: AsyncClient,
    test_project: Project,
    v2_project_url: str,
    entity_repository,
    search_service,
    file_service,
):
    """Test searching by title."""
    # Create a test entity
    entity_data = {
        "title": "Unique Title For Search",
        "note_type": "note",
        "content_type": "text/markdown",
        "file_path": "unique_title.md",
        "checksum": "title123",
    }
    await create_test_entity(
        test_project, entity_data, entity_repository, search_service, file_service
    )

    # Search by title
    response = await client.post(f"{v2_project_url}/search/", json={"title": "Unique Title"})

    assert response.status_code == 200
    data = response.json()
    assert "results" in data


@pytest.mark.asyncio
async def test_search_with_type_filter(
    client: AsyncClient,
    test_project: Project,
    v2_project_url: str,
    entity_repository,
    search_service,
    file_service,
):
    """Test searching with entity type filter."""
    # Create test entities of different types
    for note_type in ["note", "document"]:
        entity_data = {
            "title": f"Type {note_type}",
            "note_type": note_type,
            "content_type": "text/markdown",
            "file_path": f"type_{note_type}.md",
            "checksum": f"type{note_type}",
        }
        await create_test_entity(
            test_project, entity_data, entity_repository, search_service, file_service
        )

    # Search with type filter
    response = await client.post(
        f"{v2_project_url}/search/", json={"text": "Type", "note_types": ["note"]}
    )

    assert response.status_code == 200
    data = response.json()
    assert "results" in data


@pytest.mark.asyncio
async def test_search_with_date_filter(
    client: AsyncClient,
    test_project: Project,
    v2_project_url: str,
    entity_repository,
    search_service,
    file_service,
):
    """Test searching with date filter."""
    # Create a test entity
    entity_data = {
        "title": "Date Filtered",
        "note_type": "note",
        "content_type": "text/markdown",
        "file_path": "date_filtered.md",
        "checksum": "date123",
    }
    await create_test_entity(
        test_project, entity_data, entity_repository, search_service, file_service
    )

    # Search with date filter
    response = await client.post(
        f"{v2_project_url}/search/",
        json={"text": "Date Filtered", "after_date": "2024-01-01T00:00:00Z"},
    )

    assert response.status_code == 200
    data = response.json()
    assert "results" in data


@pytest.mark.asyncio
async def test_search_empty_query(
    client: AsyncClient,
    test_project: Project,
    v2_project_url: str,
):
    """Test search with empty query."""
    response = await client.post(f"{v2_project_url}/search/", json={})

    # Empty query should still be valid (returns all)
    assert response.status_code in [200, 422]


@pytest.mark.asyncio
async def test_search_whitespace_text_is_treated_as_empty(
    client: AsyncClient,
    test_project: Project,
    v2_project_url: str,
    entity_repository,
    search_service,
    file_service,
):
    """Whitespace-only text should not become an unfiltered project-wide search."""
    entity_data = {
        "title": "Whitespace Regression Entity",
        "note_type": "note",
        "content_type": "text/markdown",
        "file_path": "whitespace_regression.md",
        "checksum": "whitespace123",
    }
    await create_test_entity(
        test_project, entity_data, entity_repository, search_service, file_service
    )

    response = await client.post(f"{v2_project_url}/search/", json={"text": "   "})

    assert response.status_code == 200
    data = response.json()
    assert data["total"] == 0
    assert data["has_more"] is False
    assert data["results"] == []


@pytest.mark.asyncio
async def test_search_invalid_project_id(
    client: AsyncClient,
):
    """Test searching with invalid project ID returns 404."""
    response = await client.post("/v2/projects/999999/search/", json={"text": "test"})

    assert response.status_code == 404


@pytest.mark.asyncio
async def test_reindex(
    client: AsyncClient,
    test_project: Project,
    v2_project_url: str,
):
    """Test reindexing search index."""
    response = await client.post(f"{v2_project_url}/search/reindex")

    assert response.status_code == 200
    data = response.json()

    # Verify response structure
    assert "status" in data
    assert data["status"] == "ok"
    assert "message" in data


@pytest.mark.asyncio
async def test_reindex_invalid_project_id(
    client: AsyncClient,
):
    """Test reindexing with invalid project ID returns 404."""
    response = await client.post("/v2/projects/999999/search/reindex")

    assert response.status_code == 404


@pytest.mark.asyncio
async def test_v2_search_endpoints_use_project_id_not_name(
    client: AsyncClient,
    test_project: Project,
):
    """Test that v2 search endpoints reject string project names."""
    # Try to use project name instead of ID - should fail
    response = await client.post(f"/v2/{test_project.name}/search/", json={"text": "test"})

    # FastAPI path validation should reject non-integer project_id
    assert response.status_code in [404, 422]


@pytest.mark.asyncio
async def test_search_router_returns_400_for_semantic_disabled(
    client: AsyncClient, app, v2_project_url
):
    """SemanticSearchDisabledError should map to HTTP 400 with detail."""

    class RaisingSearchService:
        async def search(self, *args, **kwargs):
            raise SemanticSearchDisabledError("Semantic search is disabled for this project.")

        async def count(self, *args, **kwargs):
            raise SemanticSearchDisabledError("Semantic search is disabled for this project.")

    app.dependency_overrides[get_search_service_v2_external] = lambda: RaisingSearchService()
    try:
        response = await client.post(
            f"{v2_project_url}/search/",
            json={"text": "semantic query", "retrieval_mode": "vector"},
        )
    finally:
        app.dependency_overrides.pop(get_search_service_v2_external, None)

    assert response.status_code == 400
    assert "Semantic search is disabled" in response.json()["detail"]


@pytest.mark.asyncio
async def test_search_router_returns_400_for_semantic_missing_deps(
    client: AsyncClient, app, v2_project_url
):
    """SemanticDependenciesMissingError should map to HTTP 400 with detail."""

    class RaisingSearchService:
        async def search(self, *args, **kwargs):
            raise SemanticDependenciesMissingError("Semantic dependencies are missing.")

        async def count(self, *args, **kwargs):
            raise SemanticDependenciesMissingError("Semantic dependencies are missing.")

    app.dependency_overrides[get_search_service_v2_external] = lambda: RaisingSearchService()
    try:
        response = await client.post(
            f"{v2_project_url}/search/",
            json={"text": "semantic query", "retrieval_mode": "hybrid"},
        )
    finally:
        app.dependency_overrides.pop(get_search_service_v2_external, None)

    assert response.status_code == 400
    assert "Semantic dependencies are missing" in response.json()["detail"]


@pytest.mark.asyncio
async def test_search_router_returns_400_for_invalid_vector_query(
    client: AsyncClient, app, v2_project_url
):
    """ValueError from repository validation should map to HTTP 400 with detail."""

    class RaisingSearchService:
        async def search(self, *args, **kwargs):
            raise ValueError("Vector retrieval requires a text query.")

        async def count(self, *args, **kwargs):
            raise ValueError("Vector retrieval requires a text query.")

    app.dependency_overrides[get_search_service_v2_external] = lambda: RaisingSearchService()
    try:
        response = await client.post(
            f"{v2_project_url}/search/",
            json={"title": "Root", "retrieval_mode": "vector"},
        )
    finally:
        app.dependency_overrides.pop(get_search_service_v2_external, None)

    assert response.status_code == 400
    assert "Vector retrieval requires a text query" in response.json()["detail"]


@pytest.mark.asyncio
async def test_semantic_search_uses_probe_pagination_without_count(
    client: AsyncClient,
    app,
    v2_project_url: str,
):
    """Semantic searches should not run an extra count query."""
    now = datetime.now(timezone.utc)
    fake_rows = [
        SearchIndexRow(
            project_id=1,
            id=row_id,
            type="entity",
            file_path=f"notes/semantic-{row_id}.md",
            created_at=now,
            updated_at=now,
            title=f"Semantic Result {row_id}",
            permalink=f"notes/semantic-{row_id}",
            score=1.0 - (row_id / 10),
        )
        for row_id in range(1, 4)
    ]

    class FakeSearchService:
        async def search(self, query, *, limit, offset):
            assert query.retrieval_mode.value == "vector"
            assert limit == 3
            assert offset == 0
            return fake_rows

        async def count(self, *args, **kwargs):
            raise AssertionError("semantic search must not run count")

    app.dependency_overrides[get_search_service_v2_external] = lambda: FakeSearchService()
    try:
        response = await client.post(
            f"{v2_project_url}/search/",
            json={"text": "semantic query", "retrieval_mode": "vector"},
            params={"page": 1, "page_size": 2},
        )
    finally:
        app.dependency_overrides.pop(get_search_service_v2_external, None)

    assert response.status_code == 200
    data = response.json()
    assert data["total"] == 0
    assert data["has_more"] is True
    assert len(data["results"]) == 2


@pytest.mark.asyncio
async def test_search_has_more_when_more_results_exist(
    client: AsyncClient,
    test_project: Project,
    v2_project_url: str,
    entity_repository,
    search_service,
    file_service,
):
    """has_more should be True when there are more results beyond the current page."""
    # Create enough entities to exceed a small page_size
    for i in range(4):
        entity_data = {
            "title": f"HasMore Entity {i}",
            "note_type": "note",
            "content_type": "text/markdown",
            "file_path": f"hasmore_{i}.md",
            "checksum": f"hasmore{i}",
        }
        await create_test_entity(
            test_project, entity_data, entity_repository, search_service, file_service
        )

    # Request page_size=2 — with 4 entities, has_more should be True
    response = await client.post(
        f"{v2_project_url}/search/",
        json={"text": "HasMore Entity"},
        params={"page": 1, "page_size": 2},
    )
    assert response.status_code == 200
    data = response.json()
    assert "has_more" in data
    assert data["has_more"] is True
    assert len(data["results"]) == 2


@pytest.mark.asyncio
async def test_search_has_more_false_on_last_page(
    client: AsyncClient,
    test_project: Project,
    v2_project_url: str,
    entity_repository,
    search_service,
    file_service,
):
    """has_more should be False when all results fit on the current page."""
    entity_data = {
        "title": "Solo Search Entity",
        "note_type": "note",
        "content_type": "text/markdown",
        "file_path": "solo_search.md",
        "checksum": "solo123",
    }
    await create_test_entity(
        test_project, entity_data, entity_repository, search_service, file_service
    )

    response = await client.post(
        f"{v2_project_url}/search/",
        json={"text": "Solo Search Entity"},
        params={"page": 1, "page_size": 10},
    )
    assert response.status_code == 200
    data = response.json()
    assert data["has_more"] is False


@pytest.mark.asyncio
async def test_search_result_includes_matched_chunk(
    client: AsyncClient,
    app,
    v2_project_url: str,
):
    """matched_chunk field appears in search API JSON when set on SearchIndexRow."""
    now = datetime.now(timezone.utc)
    fake_row = SearchIndexRow(
        project_id=1,
        id=42,
        type="entity",
        file_path="notes/pricing.md",
        created_at=now,
        updated_at=now,
        title="Pricing Notes",
        permalink="notes/pricing",
        content_snippet="# Pricing Notes\n\n- [pricing] Team plan is $9/mo per seat",
        score=0.85,
        matched_chunk_text="- [pricing] Team plan is $9/mo per seat",
    )

    class FakeSearchService:
        async def search(self, *args, **kwargs):
            return [fake_row]

        async def count(self, *args, **kwargs):
            return 1

    app.dependency_overrides[get_search_service_v2_external] = lambda: FakeSearchService()
    try:
        response = await client.post(
            f"{v2_project_url}/search/",
            json={"text": "pricing"},
        )
    finally:
        app.dependency_overrides.pop(get_search_service_v2_external, None)

    assert response.status_code == 200
    data = response.json()
    assert len(data["results"]) == 1
    result = data["results"][0]
    assert result["matched_chunk"] == "- [pricing] Team plan is $9/mo per seat"


@pytest.mark.asyncio
async def test_search_result_omits_matched_chunk_when_none(
    client: AsyncClient,
    app,
    v2_project_url: str,
):
    """matched_chunk field is null when not set (FTS-only results)."""
    now = datetime.now(timezone.utc)
    fake_row = SearchIndexRow(
        project_id=1,
        id=43,
        type="entity",
        file_path="notes/general.md",
        created_at=now,
        updated_at=now,
        title="General Notes",
        permalink="notes/general",
        content_snippet="# General Notes\n\nSome content here",
        score=0.7,
    )

    class FakeSearchService:
        async def search(self, *args, **kwargs):
            return [fake_row]

        async def count(self, *args, **kwargs):
            return 1

    app.dependency_overrides[get_search_service_v2_external] = lambda: FakeSearchService()
    try:
        response = await client.post(
            f"{v2_project_url}/search/",
            json={"text": "general"},
        )
    finally:
        app.dependency_overrides.pop(get_search_service_v2_external, None)

    assert response.status_code == 200
    data = response.json()
    assert len(data["results"]) == 1
    result = data["results"][0]
    assert result["matched_chunk"] is None