"""Tests for the reusable batch indexing executor.""" from __future__ import annotations import asyncio from pathlib import Path from textwrap import dedent from unittest.mock import AsyncMock import pytest from sqlalchemy import text from basic_memory.file_utils import remove_frontmatter from basic_memory.indexing import ( BatchIndexer, IndexFrontmatterUpdate, IndexFrontmatterWriteResult, IndexInputFile, ) from basic_memory.schemas import Entity as EntitySchema from basic_memory.services.exceptions import SyncFatalError class _TestFileWriter: """Adapt the real FileService for batch indexer tests.""" def __init__(self, file_service) -> None: self.file_service = file_service async def write_frontmatter( self, update: IndexFrontmatterUpdate ) -> IndexFrontmatterWriteResult: result = await self.file_service.update_frontmatter_with_result( update.path, update.metadata ) return IndexFrontmatterWriteResult(checksum=result.checksum, content=result.content) async def _create_file(path: Path, content: str | bytes) -> None: path.parent.mkdir(parents=True, exist_ok=True) if isinstance(content, bytes): path.write_bytes(content) else: path.write_text(content) async def _load_input(file_service, path: str) -> IndexInputFile: metadata = await file_service.get_file_metadata(path) return IndexInputFile( path=path, size=metadata.size, checksum=await file_service.compute_checksum(path), content_type=file_service.content_type(path), last_modified=metadata.modified_at, created_at=metadata.created_at, content=await file_service.read_file_bytes(path), ) def _make_batch_indexer( app_config, entity_service, entity_repository, relation_repository, search_service, file_service ) -> BatchIndexer: return BatchIndexer( app_config=app_config, entity_service=entity_service, entity_repository=entity_repository, relation_repository=relation_repository, search_service=search_service, file_writer=_TestFileWriter(file_service), ) @pytest.mark.asyncio async def test_batch_indexer_parses_markdown_with_parallel_path( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, project_config, ): path_one = "notes/one.md" path_two = "notes/two.md" await _create_file( project_config.home / path_one, dedent( """ --- title: One type: note --- # One """ ).strip(), ) await _create_file( project_config.home / path_two, dedent( """ --- title: Two type: note --- # Two """ ).strip(), ) files = { path_one: await _load_input(file_service, path_one), path_two: await _load_input(file_service, path_two), } batch_indexer = _make_batch_indexer( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, ) original_parse = entity_service.entity_parser.parse_markdown_content in_flight = 0 max_in_flight = 0 async def spy_parse(*args, **kwargs): nonlocal in_flight, max_in_flight in_flight += 1 max_in_flight = max(max_in_flight, in_flight) await asyncio.sleep(0.05) try: return await original_parse(*args, **kwargs) finally: in_flight -= 1 entity_service.entity_parser.parse_markdown_content = spy_parse try: result = await batch_indexer.index_files( files, max_concurrent=2, parse_max_concurrent=2, ) finally: entity_service.entity_parser.parse_markdown_content = original_parse assert max_in_flight >= 2 assert len(result.indexed) == 2 assert result.errors == [] @pytest.mark.asyncio async def test_batch_indexer_creates_entities_with_parallel_path( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, project_config, ): path_one = "notes/alpha.md" path_two = "notes/beta.md" await _create_file( project_config.home / path_one, dedent( """ --- title: Alpha type: note --- # Alpha """ ).strip(), ) await _create_file( project_config.home / path_two, dedent( """ --- title: Beta type: note --- # Beta """ ).strip(), ) files = { path_one: await _load_input(file_service, path_one), path_two: await _load_input(file_service, path_two), } batch_indexer = _make_batch_indexer( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, ) original_upsert = entity_service.upsert_entity_from_markdown in_flight = 0 max_in_flight = 0 async def spy_upsert(*args, **kwargs): nonlocal in_flight, max_in_flight in_flight += 1 max_in_flight = max(max_in_flight, in_flight) await asyncio.sleep(0.05) try: return await original_upsert(*args, **kwargs) finally: in_flight -= 1 entity_service.upsert_entity_from_markdown = spy_upsert try: result = await batch_indexer.index_files( files, max_concurrent=2, parse_max_concurrent=2, ) finally: entity_service.upsert_entity_from_markdown = original_upsert assert max_in_flight >= 2 assert len(result.indexed) == 2 assert result.errors == [] @pytest.mark.asyncio async def test_batch_indexer_returns_original_markdown_content_when_no_frontmatter_rewrite( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, project_config, ): app_config.disable_permalinks = True path = "notes/original.md" original_content = dedent( """ --- title: Original type: note --- # Original """ ).strip() await _create_file(project_config.home / path, original_content) files = {path: await _load_input(file_service, path)} batch_indexer = _make_batch_indexer( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, ) result = await batch_indexer.index_files( files, max_concurrent=1, parse_max_concurrent=1, ) # Trigger: Windows persists CRLF for text writes even when the test literal uses LF. # Why: this assertion cares about "no rewrite happened", not about forcing one newline # convention across platforms. # Outcome: compare against the exact markdown text stored on disk for this file. persisted_content = (project_config.home / path).read_bytes().decode("utf-8") assert result.errors == [] assert len(result.indexed) == 1 assert result.indexed[0].markdown_content == persisted_content @pytest.mark.asyncio async def test_batch_indexer_indexes_non_markdown_files( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, project_config, ): pdf_path = "assets/doc.pdf" image_path = "assets/image.png" await _create_file(project_config.home / pdf_path, b"%PDF-1.4 test") await _create_file(project_config.home / image_path, b"\x89PNG\r\n\x1a\nrest") files = { pdf_path: await _load_input(file_service, pdf_path), image_path: await _load_input(file_service, image_path), } batch_indexer = _make_batch_indexer( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, ) result = await batch_indexer.index_files( files, max_concurrent=2, parse_max_concurrent=2, ) assert {indexed.path for indexed in result.indexed} == {pdf_path, image_path} assert all(indexed.markdown_content is None for indexed in result.indexed) pdf_entity = await entity_repository.get_by_file_path(pdf_path) image_entity = await entity_repository.get_by_file_path(image_path) assert pdf_entity is not None assert pdf_entity.content_type == "application/pdf" assert image_entity is not None assert image_entity.content_type == "image/png" @pytest.mark.asyncio async def test_batch_indexer_resolves_relations_and_refreshes_search( app_config, entity_service, entity_repository, relation_repository, search_service, search_repository, file_service, project_config, ): source_path = "notes/source.md" target_path = "notes/target.md" await _create_file( project_config.home / source_path, dedent( """ --- title: Source type: note --- # Source - depends_on [[Target]] """ ).strip(), ) await _create_file( project_config.home / target_path, dedent( """ --- title: Target type: note --- # Target """ ).strip(), ) files = { source_path: await _load_input(file_service, source_path), target_path: await _load_input(file_service, target_path), } batch_indexer = _make_batch_indexer( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, ) result = await batch_indexer.index_files( files, max_concurrent=2, parse_max_concurrent=2, ) source = await entity_repository.get_by_file_path(source_path) target = await entity_repository.get_by_file_path(target_path) assert source is not None assert target is not None assert len(source.outgoing_relations) == 1 assert source.outgoing_relations[0].to_id == target.id assert result.relations_unresolved == 0 assert result.search_indexed == 2 relation_rows = await search_repository.execute_query( text( "SELECT COUNT(*) FROM search_index " "WHERE entity_id = :entity_id AND type = 'relation' AND to_id IS NOT NULL" ), {"entity_id": source.id}, ) assert relation_rows.scalar_one() == 1 @pytest.mark.asyncio async def test_batch_indexer_assigns_unique_permalinks_for_batch_local_conflicts( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, project_config, ): path_one = "notes/basic memory bug.md" path_two = "notes/basic-memory-bug.md" await _create_file( project_config.home / path_one, dedent( """ --- title: Basic Memory Bug type: note --- # Basic Memory Bug """ ).strip(), ) await _create_file( project_config.home / path_two, dedent( """ --- title: Basic Memory Bug Report type: note --- # Basic Memory Bug Report """ ).strip(), ) files = { path_one: await _load_input(file_service, path_one), path_two: await _load_input(file_service, path_two), } original_contents = { path: file.content.decode("utf-8") for path, file in files.items() if file.content is not None } batch_indexer = _make_batch_indexer( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, ) result = await batch_indexer.index_files( files, max_concurrent=2, parse_max_concurrent=2, ) assert result.errors == [] indexed_by_path = {indexed.path: indexed for indexed in result.indexed} assert indexed_by_path[path_one].markdown_content is not None assert indexed_by_path[path_two].markdown_content is not None assert indexed_by_path[path_one].markdown_content != original_contents[path_one] assert indexed_by_path[path_two].markdown_content != original_contents[path_two] assert indexed_by_path[path_one].markdown_content == await file_service.read_file_content( path_one ) assert indexed_by_path[path_two].markdown_content == await file_service.read_file_content( path_two ) entities = await entity_repository.find_all() assert len(entities) == 2 permalinks = [entity.permalink for entity in entities if entity.permalink] assert len(set(permalinks)) == 2 @pytest.mark.asyncio async def test_batch_indexer_uses_parsed_markdown_body_for_malformed_frontmatter_delimiters( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, project_config, ): app_config.disable_permalinks = True app_config.ensure_frontmatter_on_sync = False path = "notes/malformed.md" malformed_content = dedent( """ --- this is not valid frontmatter # Malformed Frontmatter The parser should still index this file. """ ).strip() await _create_file(project_config.home / path, malformed_content) files = {path: await _load_input(file_service, path)} batch_indexer = _make_batch_indexer( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, ) result = await batch_indexer.index_files( files, max_concurrent=1, parse_max_concurrent=1, ) # Trigger: malformed frontmatter should pass through without normalization. # Why: Windows can still surface that unchanged file with CRLF line endings. # Outcome: compare the indexed markdown to the persisted file content, not the LF # test literal used to create it. persisted_content = (project_config.home / path).read_bytes().decode("utf-8") assert result.errors == [] assert len(result.indexed) == 1 assert result.indexed[0].markdown_content == persisted_content entity = await entity_repository.get_by_file_path(path) assert entity is not None @pytest.mark.asyncio async def test_batch_indexer_re_raises_fatal_sync_errors( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, ): batch_indexer = _make_batch_indexer( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, ) async def fatal_worker(path: str) -> str: raise SyncFatalError(f"fatal batch failure for {path}") with pytest.raises(SyncFatalError, match="fatal batch failure"): await batch_indexer._run_bounded( ["notes/fatal.md"], limit=1, worker=fatal_worker, ) @pytest.mark.asyncio async def test_batch_indexer_index_markdown_file_rewrites_permalink_after_repository_conflict( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, project_config, monkeypatch, ): existing = await entity_service.create_entity_with_content( EntitySchema( title="Existing Note", directory="notes", content="# Existing Note\n\nOriginal content.\n", ) ) conflicting_permalink = existing.entity.permalink assert conflicting_permalink is not None path = "notes/race.md" await _create_file( project_config.home / path, dedent( f"""\ --- title: Race Note type: note permalink: {conflicting_permalink} --- # Race Note Body content. """ ), ) async def stale_permalink(*args, **kwargs) -> str: return conflicting_permalink batch_indexer = _make_batch_indexer( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, ) monkeypatch.setattr(entity_service, "resolve_permalink", stale_permalink) indexed = await batch_indexer.index_markdown_file( await _load_input(file_service, path), index_search=False, ) persisted_content = await file_service.read_file_content(path) assert indexed.permalink == f"{conflicting_permalink}-1" assert indexed.markdown_content == persisted_content @pytest.mark.asyncio async def test_batch_indexer_index_markdown_file_can_defer_relation_resolution( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, project_config, monkeypatch, ): await entity_service.create_entity_with_content( EntitySchema( title="Deferred Target", directory="notes", content="# Deferred Target\n", ) ) path = "notes/deferred-source.md" await _create_file( project_config.home / path, dedent( """ --- title: Deferred Source type: note --- # Deferred Source - links_to [[Deferred Target]] """ ), ) resolve_link = AsyncMock(side_effect=AssertionError("relation lookup should be deferred")) monkeypatch.setattr(entity_service.link_resolver, "resolve_link", resolve_link) batch_indexer = _make_batch_indexer( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, ) await batch_indexer.index_markdown_file( await _load_input(file_service, path), index_search=False, resolve_relations=False, ) resolve_link.assert_not_awaited() source = await entity_repository.get_by_file_path(path) assert source is not None assert len(source.outgoing_relations) == 1 assert source.outgoing_relations[0].to_id is None assert source.outgoing_relations[0].to_name == "Deferred Target" @pytest.mark.asyncio async def test_batch_indexer_uses_strict_link_resolution_for_deferred_relations( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, project_config, monkeypatch, ): """Regression: batch indexer's deferred relation resolution must call resolve_link with strict=True. Mirror of sync_service.resolve_forward_references. Fuzzy fallback in the deferred path silently fills in to_id from BM25/ts_rank results, polluting the graph with confidently-wrong edges. Entity-creation already uses strict=True; this is the other deferred path. """ path = "notes/source.md" await _create_file( project_config.home / path, dedent( """ --- title: Source type: note --- # Source - links_to [[never-resolves-target]] """ ), ) batch_indexer = _make_batch_indexer( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, ) original_resolve_link = entity_service.link_resolver.resolve_link seen_strict: list[object] = [] async def spy_resolve_link(*args, **kwargs): seen_strict.append(kwargs.get("strict", False)) return await original_resolve_link(*args, **kwargs) monkeypatch.setattr(entity_service.link_resolver, "resolve_link", spy_resolve_link) await batch_indexer.index_files( {path: await _load_input(file_service, path)}, max_concurrent=1, ) assert seen_strict, "batch indexer did not invoke link_resolver.resolve_link" assert all(strict is True for strict in seen_strict), ( f"Deferred resolution must call resolve_link(strict=True). Observed: {seen_strict!r}" ) # The unresolvable relation stayed unresolved. source = await entity_repository.get_by_file_path(path) assert source is not None assert len(source.outgoing_relations) == 1 assert source.outgoing_relations[0].to_id is None assert source.outgoing_relations[0].to_name == "never-resolves-target" @pytest.mark.asyncio async def test_batch_indexer_strips_frontmatter_from_search_content_when_body_is_empty( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, project_config, monkeypatch, ): path = "notes/frontmatter-only.md" await _create_file( project_config.home / path, dedent( """ --- title: Frontmatter Only type: note status: draft --- """ ).strip(), ) index_entity_data = AsyncMock() monkeypatch.setattr(search_service, "index_entity_data", index_entity_data) batch_indexer = _make_batch_indexer( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, ) await batch_indexer.index_markdown_file( await _load_input(file_service, path), index_search=True ) persisted_content = await file_service.read_file_content(path) entity = await entity_repository.get_by_file_path(path) assert entity is not None index_entity_data.assert_awaited_once() await_args = index_entity_data.await_args assert await_args is not None args, kwargs = await_args assert args[0].id == entity.id assert kwargs["content"] == remove_frontmatter(persisted_content) @pytest.mark.asyncio async def test_batch_indexer_does_not_inject_frontmatter_when_sync_enforcement_is_disabled( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, project_config, monkeypatch, ): app_config.ensure_frontmatter_on_sync = False created = await entity_service.create_entity_with_content( EntitySchema( title="Frontmatterless", directory="notes", content="# Frontmatterless\n\nOriginal content.\n", ) ) path = created.entity.file_path assert path is not None existing_permalink = created.entity.permalink assert existing_permalink is not None original_content = "# Frontmatterless\n\nBody content.\n" await _create_file(project_config.home / path, original_content) original_writer = file_service.update_frontmatter_with_result frontmatter_writer = AsyncMock(side_effect=original_writer) monkeypatch.setattr(file_service, "update_frontmatter_with_result", frontmatter_writer) batch_indexer = _make_batch_indexer( app_config, entity_service, entity_repository, relation_repository, search_service, file_service, ) indexed = await batch_indexer.index_markdown_file( await _load_input(file_service, path), index_search=False, ) # Trigger: Windows persists CRLF for text files even when the test literal uses LF. # Why: this assertion cares about preserving a frontmatterless file, not about newline style. # Outcome: compare against the exact content stored on disk after sync. persisted_content = (project_config.home / path).read_bytes().decode("utf-8") entity = await entity_repository.get_by_file_path(path) assert entity is not None assert entity.permalink == existing_permalink assert frontmatter_writer.await_count == 0 assert indexed.markdown_content == persisted_content assert (await file_service.read_file_bytes(path)).decode("utf-8") == persisted_content