"""Tests proving upsert_entity_from_markdown optimizations. Verifies that: 1. Redundant get_by_file_path call is eliminated (entity passed directly) 2. Final reload uses find_by_ids (PK lookup) instead of get_by_file_path (string lookup) 3. Telemetry sub-spans are emitted for each DB phase 4. Correctness is preserved for create, update, and edit flows """ from __future__ import annotations from datetime import datetime, timezone from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock import pytest from basic_memory.markdown.schemas import ( EntityFrontmatter, EntityMarkdown, Observation as MarkdownObservation, Relation as MarkdownRelation, ) from basic_memory.schemas import Entity as EntitySchema from basic_memory.services.entity_service import EntityService # --- Helpers --- def _make_markdown( title: str = "Test Entity", observations: list | None = None, relations: list | None = None, ) -> EntityMarkdown: frontmatter = EntityFrontmatter(metadata={"title": title, "type": "note"}) return EntityMarkdown( frontmatter=frontmatter, observations=observations or [], relations=relations or [], created=datetime.now(timezone.utc), modified=datetime.now(timezone.utc), ) # --- Optimization 1: No redundant get_by_file_path in update_entity_relations --- @pytest.mark.asyncio async def test_upsert_update_does_not_refetch_entity(entity_service: EntityService, monkeypatch): """update_entity_relations should NOT call get_by_file_path — entity is passed directly.""" # Create an entity first entity = await entity_service.create_entity( EntitySchema( title="Refetch Test", directory="notes", note_type="note", content="# Refetch Test\n\n## Observations\n- [fact] some fact", ) ) # Spy on get_by_file_path calls original_get_by_file_path = entity_service.repository.get_by_file_path call_count = 0 async def spy_get_by_file_path(*args, **kwargs): nonlocal call_count call_count += 1 return await original_get_by_file_path(*args, **kwargs) monkeypatch.setattr(entity_service.repository, "get_by_file_path", spy_get_by_file_path) # Run upsert with is_new=False — this calls update_entity_and_observations + update_entity_relations markdown = _make_markdown( title="Refetch Test", observations=[MarkdownObservation(content="updated fact", category="fact")], ) await entity_service.upsert_entity_from_markdown(Path(entity.file_path), markdown, is_new=False) # update_entity_and_observations calls get_by_file_path once (to load the entity) # update_entity_relations should NOT call it at all (entity passed directly) assert call_count == 1, ( f"Expected 1 get_by_file_path call (in update_entity_and_observations only), " f"got {call_count}. update_entity_relations should not re-fetch." ) # --- Optimization 2: Final reload uses find_by_ids (PK) not get_by_file_path --- @pytest.mark.asyncio async def test_update_entity_relations_uses_pk_reload(entity_service: EntityService, monkeypatch): """update_entity_relations should use find_by_ids for the final reload, not get_by_file_path.""" entity = await entity_service.create_entity( EntitySchema( title="PK Reload Test", directory="notes", note_type="note", content="# PK Reload Test", ) ) # Spy on find_by_ids calls original_find_by_ids = entity_service.repository.find_by_ids find_by_ids_calls = [] async def spy_find_by_ids(ids): find_by_ids_calls.append(ids) return await original_find_by_ids(ids) monkeypatch.setattr(entity_service.repository, "find_by_ids", spy_find_by_ids) markdown = _make_markdown(title="PK Reload Test") await entity_service.upsert_entity_from_markdown(Path(entity.file_path), markdown, is_new=False) # update_entity_relations should call find_by_ids once with the entity's PK assert len(find_by_ids_calls) == 1 assert find_by_ids_calls[0] == [entity.id] @pytest.mark.asyncio async def test_create_or_update_entity_uses_lightweight_exact_resolution( entity_service: EntityService, monkeypatch ): """create_or_update_entity should use strict lookups without eager relation loading.""" schema = EntitySchema( title="Create Or Update", directory="notes", note_type="note", content="# Create Or Update", ) sentinel_entity = SimpleNamespace(file_path="notes/existing.md") resolve_calls: list[tuple[str, dict]] = [] async def fake_resolve_link(link_text: str, **kwargs): resolve_calls.append((link_text, kwargs)) if link_text == schema.file_path: return None return sentinel_entity monkeypatch.setattr(entity_service.link_resolver, "resolve_link", fake_resolve_link) monkeypatch.setattr(entity_service, "update_entity", AsyncMock(return_value=sentinel_entity)) entity, is_new = await entity_service.create_or_update_entity(schema) assert entity is sentinel_entity assert is_new is False assert resolve_calls == [ (schema.file_path, {"strict": True, "load_relations": False}), (schema.permalink, {"strict": True, "load_relations": False}), ] @pytest.mark.asyncio async def test_upsert_with_relations_uses_lightweight_exact_resolution( entity_service: EntityService, monkeypatch ): """Relation target resolution should skip eager loading during upsert.""" target = await entity_service.create_entity( EntitySchema( title="Lightweight Target", directory="notes", note_type="note", content="# Lightweight Target", ) ) source = await entity_service.create_entity( EntitySchema( title="Lightweight Source", directory="notes", note_type="note", content="# Lightweight Source", ) ) resolve_calls: list[tuple[str, dict]] = [] async def fake_resolve_link(link_text: str, **kwargs): resolve_calls.append((link_text, kwargs)) return target monkeypatch.setattr(entity_service.link_resolver, "resolve_link", fake_resolve_link) markdown = _make_markdown( title="Lightweight Source", relations=[MarkdownRelation(type="links_to", target="Lightweight Target")], ) await entity_service.upsert_entity_from_markdown(Path(source.file_path), markdown, is_new=False) assert resolve_calls == [ ("Lightweight Target", {"strict": True, "load_relations": False}), ] @pytest.mark.asyncio async def test_upsert_can_defer_relation_target_resolution( entity_service: EntityService, monkeypatch ): """Cloud one-file indexing can store unresolved relation rows for later repair.""" await entity_service.create_entity( EntitySchema( title="Deferred Target", directory="notes", note_type="note", content="# Deferred Target", ) ) source = await entity_service.create_entity( EntitySchema( title="Deferred Source", directory="notes", note_type="note", content="# Deferred Source", ) ) resolve_link = AsyncMock(side_effect=AssertionError("relation lookup should be deferred")) monkeypatch.setattr(entity_service.link_resolver, "resolve_link", resolve_link) markdown = _make_markdown( title="Deferred Source", relations=[MarkdownRelation(type="links_to", target="Deferred Target")], ) updated = await entity_service.upsert_entity_from_markdown( Path(source.file_path), markdown, is_new=False, resolve_relations=False, ) resolve_link.assert_not_awaited() outgoing = updated.outgoing_relations assert len(outgoing) == 1 assert outgoing[0].to_id is None assert outgoing[0].to_name == "Deferred Target" @pytest.mark.asyncio async def test_upsert_deferred_relation_resolution_keeps_self_links_resolved( entity_service: EntityService, monkeypatch ): """Deferred relation mode should still resolve self-links without a target lookup.""" source = await entity_service.create_entity( EntitySchema( title="Deferred Self", directory="notes", note_type="note", content="# Deferred Self", ) ) resolve_link = AsyncMock(side_effect=AssertionError("relation lookup should be deferred")) monkeypatch.setattr(entity_service.link_resolver, "resolve_link", resolve_link) markdown = _make_markdown( title="Deferred Self", relations=[MarkdownRelation(type="links_to", target="Deferred Self")], ) updated = await entity_service.upsert_entity_from_markdown( Path(source.file_path), markdown, is_new=False, resolve_relations=False, ) resolve_link.assert_not_awaited() outgoing = updated.outgoing_relations assert len(outgoing) == 1 assert outgoing[0].to_id == source.id assert outgoing[0].to_name == "Deferred Self" @pytest.mark.asyncio async def test_upsert_deferred_relation_resolution_does_not_guess_duplicate_titles( entity_service: EntityService, monkeypatch ): """Deferred relation mode should not treat duplicate titles as guaranteed self-links.""" source = await entity_service.create_entity( EntitySchema( title="Duplicate Title", directory="notes/source", note_type="note", content="# Duplicate Title", ) ) await entity_service.create_entity( EntitySchema( title="Duplicate Title", directory="notes/other", note_type="note", content="# Duplicate Title", ) ) resolve_link = AsyncMock(side_effect=AssertionError("relation lookup should be deferred")) monkeypatch.setattr(entity_service.link_resolver, "resolve_link", resolve_link) markdown = _make_markdown( title="Duplicate Title", relations=[MarkdownRelation(type="links_to", target="Duplicate Title")], ) updated = await entity_service.upsert_entity_from_markdown( Path(source.file_path), markdown, is_new=False, resolve_relations=False, ) resolve_link.assert_not_awaited() outgoing = updated.outgoing_relations assert len(outgoing) == 1 assert outgoing[0].to_id is None assert outgoing[0].to_name == "Duplicate Title" # --- Correctness: full round-trip --- @pytest.mark.asyncio async def test_upsert_update_preserves_observations(entity_service: EntityService): """After upsert (update path), observations should be correctly replaced.""" entity = await entity_service.create_entity( EntitySchema( title="Obs Test", directory="notes", note_type="note", content="# Obs Test\n\n## Observations\n- [fact] original fact", ) ) assert len(entity.observations) == 1 markdown = _make_markdown( title="Obs Test", observations=[ MarkdownObservation(content="new fact 1", category="fact"), MarkdownObservation(content="new fact 2", category="idea"), ], ) updated = await entity_service.upsert_entity_from_markdown( Path(entity.file_path), markdown, is_new=False ) assert updated.id == entity.id assert len(updated.observations) == 2 obs_contents = {o.content for o in updated.observations} assert obs_contents == {"new fact 1", "new fact 2"} @pytest.mark.asyncio async def test_upsert_update_preserves_relations(entity_service: EntityService): """After upsert (update path), relations should be correctly replaced.""" target = await entity_service.create_entity( EntitySchema( title="Relation Target", directory="notes", note_type="note", content="# Relation Target", ) ) source = await entity_service.create_entity( EntitySchema( title="Relation Source", directory="notes", note_type="note", content="# Relation Source\n\n## Relations\n- links_to [[Relation Target]]", ) ) assert len(source.relations) == 1 markdown = _make_markdown( title="Relation Source", relations=[MarkdownRelation(type="references", target="Relation Target")], ) updated = await entity_service.upsert_entity_from_markdown( Path(source.file_path), markdown, is_new=False ) assert updated.id == source.id # Old relation replaced with new one outgoing = [r for r in updated.relations if r.from_id == source.id] assert len(outgoing) == 1 assert outgoing[0].relation_type == "references" assert outgoing[0].to_id == target.id @pytest.mark.asyncio async def test_upsert_create_path_works(entity_service: EntityService): """The is_new=True path should still work correctly.""" markdown = _make_markdown( title="Create Path Test", observations=[MarkdownObservation(content="a fact", category="fact")], ) result = await entity_service.upsert_entity_from_markdown( Path("notes/create-path-test.md"), markdown, is_new=True ) assert result.title == "Create Path Test" assert len(result.observations) == 1 assert result.observations[0].content == "a fact" @pytest.mark.asyncio async def test_edit_entity_end_to_end(entity_service: EntityService): """Full edit_entity flow uses optimized upsert and returns correct entity.""" entity = await entity_service.create_entity( EntitySchema( title="Edit E2E", directory="notes", note_type="note", content="# Edit E2E\n\nOriginal content.", ) ) updated = await entity_service.edit_entity( entity.file_path, operation="append", content="\n\n## Observations\n- [fact] appended fact", ) assert updated.id == entity.id assert len(updated.observations) == 1 assert updated.observations[0].content == "appended fact" # Checksum should be set (not None) after edit completes assert updated.checksum is not None @pytest.mark.asyncio async def test_edit_entity_uses_lightweight_identifier_resolution( entity_service: EntityService, monkeypatch ): """edit_entity should resolve the target note without eager relation loading.""" entity = await entity_service.create_entity( EntitySchema( title="Edit Lightweight", directory="notes", note_type="note", content="# Edit Lightweight\n\nOriginal content.", ) ) original_resolve_link = entity_service.link_resolver.resolve_link resolve_calls: list[tuple[str, dict]] = [] async def spy_resolve_link(link_text: str, **kwargs): resolve_calls.append((link_text, kwargs)) return await original_resolve_link(link_text, **kwargs) monkeypatch.setattr(entity_service.link_resolver, "resolve_link", spy_resolve_link) await entity_service.edit_entity( entity.file_path, operation="append", content="\n\nNo relation changes here.", ) assert resolve_calls[0] == ( entity.file_path, {"strict": True, "load_relations": False}, )