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

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