"""Telemetry coverage for the v2 knowledge router.""" from __future__ import annotations import importlib from contextlib import contextmanager from datetime import datetime, timezone from types import SimpleNamespace import pytest from fastapi import BackgroundTasks, Response from basic_memory.schemas.base import Entity from basic_memory.schemas.request import EditEntityRequest knowledge_router_module = importlib.import_module("basic_memory.api.v2.routers.knowledge_router") def _capture_spans(): spans: list[tuple[str, dict]] = [] @contextmanager def fake_span(name: str, **attrs): spans.append((name, attrs)) yield return spans, fake_span def _fake_entity(*, external_id: str = "entity-123", file_path: str = "notes/test.md"): now = datetime.now(timezone.utc) return SimpleNamespace( external_id=external_id, id=1, title="Telemetry Entity", note_type="note", content_type="text/markdown", permalink="notes/test", file_path=file_path, entity_metadata=None, observations=[], relations=[], created_at=now, updated_at=now, created_by=None, last_updated_by=None, ) def _assert_names_in_order(names: list[str], expected: list[str]) -> None: cursor = 0 for expected_name in expected: cursor = names.index(expected_name, cursor) + 1 @pytest.mark.asyncio async def test_create_entity_emits_root_and_nested_spans(monkeypatch) -> None: spans, fake_span = _capture_spans() monkeypatch.setattr(knowledge_router_module.telemetry, "span", fake_span) entity = _fake_entity() class FakeEntityService: async def create_entity(self, data): return entity class FakeSearchService: async def index_entity(self, entity): return None class FakeTaskScheduler: def schedule(self, *args, **kwargs): return None class FakeFileService: async def read_file_content(self, path): return "telemetry content" result = await knowledge_router_module.create_entity( project_id="project-123", data=Entity( title="Telemetry Entity", directory="notes", note_type="note", content_type="text/markdown", content="telemetry content", ), background_tasks=BackgroundTasks(), entity_service=FakeEntityService(), search_service=FakeSearchService(), task_scheduler=FakeTaskScheduler(), file_service=FakeFileService(), app_config=SimpleNamespace(semantic_search_enabled=False), fast=False, ) assert result.content == "telemetry content" _assert_names_in_order( [name for name, _ in spans], [ "api.request.knowledge.create_entity", "api.knowledge.create_entity.write_entity", "api.knowledge.create_entity.search_index", "api.knowledge.create_entity.vector_sync", "api.knowledge.create_entity.read_content", ], ) @pytest.mark.asyncio async def test_update_entity_emits_root_and_nested_spans(monkeypatch) -> None: spans, fake_span = _capture_spans() monkeypatch.setattr(knowledge_router_module.telemetry, "span", fake_span) entity = _fake_entity() class FakeEntityService: async def update_entity(self, existing, data): return entity class FakeSearchService: async def index_entity(self, entity): return None class FakeEntityRepository: async def get_by_external_id(self, external_id): return entity class FakeTaskScheduler: def schedule(self, *args, **kwargs): return None class FakeFileService: async def read_file_content(self, path): return "updated telemetry content" response = Response() result = await knowledge_router_module.update_entity_by_id( data=Entity( title="Telemetry Entity", directory="notes", note_type="note", content_type="text/markdown", content="updated telemetry content", ), response=response, background_tasks=BackgroundTasks(), project_id="project-123", entity_service=FakeEntityService(), search_service=FakeSearchService(), entity_repository=FakeEntityRepository(), task_scheduler=FakeTaskScheduler(), file_service=FakeFileService(), app_config=SimpleNamespace(semantic_search_enabled=False), entity_id=entity.external_id, fast=False, ) assert result.content == "updated telemetry content" _assert_names_in_order( [name for name, _ in spans], [ "api.request.knowledge.update_entity", "api.knowledge.update_entity.load_entity", "api.knowledge.update_entity.write_entity", "api.knowledge.update_entity.search_index", "api.knowledge.update_entity.vector_sync", "api.knowledge.update_entity.read_content", ], ) @pytest.mark.asyncio async def test_edit_entity_emits_root_and_nested_spans(monkeypatch) -> None: spans, fake_span = _capture_spans() monkeypatch.setattr(knowledge_router_module.telemetry, "span", fake_span) entity = _fake_entity() class FakeEntityService: async def edit_entity(self, **kwargs): return entity class FakeSearchService: async def index_entity(self, entity): return None class FakeEntityRepository: async def get_by_external_id(self, external_id): return entity class FakeTaskScheduler: def schedule(self, *args, **kwargs): return None class FakeFileService: async def read_file_content(self, path): return "edited telemetry content" result = await knowledge_router_module.edit_entity_by_id( data=EditEntityRequest(operation="append", content="edited telemetry content"), background_tasks=BackgroundTasks(), project_id="project-123", entity_service=FakeEntityService(), search_service=FakeSearchService(), entity_repository=FakeEntityRepository(), task_scheduler=FakeTaskScheduler(), file_service=FakeFileService(), app_config=SimpleNamespace(semantic_search_enabled=False), entity_id=entity.external_id, fast=False, ) assert result.content == "edited telemetry content" _assert_names_in_order( [name for name, _ in spans], [ "api.request.knowledge.edit_entity", "api.knowledge.edit_entity.load_entity", "api.knowledge.edit_entity.write_entity", "api.knowledge.edit_entity.search_index", "api.knowledge.edit_entity.vector_sync", "api.knowledge.edit_entity.read_content", ], )