fleet-memory/hindsight-api-slim/tests/test_tracing_integration.py

"""
Integration tests for OpenTelemetry tracing with memory engine operations.

Tests that parent spans are correctly created for retain, consolidation, reflect,
and mental_model_refresh operations.
"""

from datetime import datetime, timezone
from unittest.mock import MagicMock, patch

import pytest


@pytest.mark.asyncio
@patch("hindsight_api.engine.memory_engine.create_operation_span")
async def test_retain_creates_parent_span(mock_create_span, memory, request_context):
    """Test that retain operation creates a parent span."""
    # Setup
    mock_span = MagicMock()
    mock_span.__enter__ = MagicMock(return_value=mock_span)
    mock_span.__exit__ = MagicMock(return_value=False)
    mock_create_span.return_value = mock_span

    bank_id = f"test-retain-{datetime.now(timezone.utc).timestamp()}"

    try:
        # Execute retain (automatically creates bank if needed)
        await memory.retain_async(
            bank_id=bank_id,
            content="Test memory for tracing",
            context="Test context",
            request_context=request_context,
        )

        # Verify parent span was created
        mock_create_span.assert_called()
        call_args = mock_create_span.call_args
        assert call_args[0][0] == "retain"  # operation name
        assert call_args[0][1] == bank_id  # bank_id

        # Verify span was used as context manager
        mock_span.__enter__.assert_called()
        mock_span.__exit__.assert_called()
    finally:
        # Cleanup
        await memory.delete_bank(bank_id, request_context=request_context)


@pytest.mark.asyncio
@patch("hindsight_api.engine.memory_engine.create_operation_span")
async def test_consolidation_creates_parent_span(mock_create_span, memory, request_context):
    """Test that consolidation operation creates a parent span."""
    # Setup
    mock_span = MagicMock()
    mock_span.__enter__ = MagicMock(return_value=mock_span)
    mock_span.__exit__ = MagicMock(return_value=False)
    mock_create_span.return_value = mock_span

    bank_id = f"test-consolidation-{datetime.now(timezone.utc).timestamp()}"

    try:
        # Execute consolidation (bank will be created automatically)
        await memory.run_consolidation(
            bank_id=bank_id,
            request_context=request_context,
        )

        # Verify parent span was created
        mock_create_span.assert_called()
        call_args = mock_create_span.call_args
        assert call_args[0][0] == "consolidation"
        assert call_args[0][1] == bank_id

        # Verify span was used as context manager
        mock_span.__enter__.assert_called()
        mock_span.__exit__.assert_called()
    finally:
        # Cleanup
        await memory.delete_bank(bank_id, request_context=request_context)


@pytest.mark.asyncio
@patch("hindsight_api.engine.memory_engine.create_operation_span")
async def test_reflect_creates_parent_span(mock_create_span, memory, request_context):
    """Test that reflect operation creates a parent span."""
    # Setup
    mock_span = MagicMock()
    mock_span.__enter__ = MagicMock(return_value=mock_span)
    mock_span.__exit__ = MagicMock(return_value=False)
    mock_create_span.return_value = mock_span

    bank_id = f"test-reflect-{datetime.now(timezone.utc).timestamp()}"

    try:
        # Add some memories first
        await memory.retain_async(
            bank_id=bank_id,
            content="Paris is the capital of France",
            context="Geography fact",
            request_context=request_context,
        )

        # Reset mock to clear retain call
        mock_create_span.reset_mock()

        # Execute reflect
        await memory.reflect_async(
            bank_id=bank_id,
            query="What is the capital of France?",
            request_context=request_context,
        )

        # Verify parent span was created
        mock_create_span.assert_called()
        call_args = mock_create_span.call_args
        assert call_args[0][0] == "reflect"
        assert call_args[0][1] == bank_id

        # Verify span was used as context manager
        mock_span.__enter__.assert_called()
        mock_span.__exit__.assert_called()
    finally:
        # Cleanup
        await memory.delete_bank(bank_id, request_context=request_context)


@pytest.mark.asyncio
@patch("hindsight_api.engine.memory_engine.create_operation_span")
async def test_retain_batch_creates_single_parent_span(mock_create_span, memory, request_context):
    """Test that batch retain creates one parent span for the entire batch."""
    # Setup
    mock_span = MagicMock()
    mock_span.__enter__ = MagicMock(return_value=mock_span)
    mock_span.__exit__ = MagicMock(return_value=False)
    mock_create_span.return_value = mock_span

    bank_id = f"test-batch-{datetime.now(timezone.utc).timestamp()}"

    try:
        # Execute batch retain with multiple items
        await memory.retain_batch_async(
            bank_id=bank_id,
            contents=[
                {"content": "Memory 1", "context": "Context 1"},
                {"content": "Memory 2", "context": "Context 2"},
                {"content": "Memory 3", "context": "Context 3"},
            ],
            request_context=request_context,
        )

        # Verify parent span was created only once for the entire batch
        assert mock_create_span.call_count == 1
        call_args = mock_create_span.call_args
        assert call_args[0][0] == "retain"
        assert call_args[0][1] == bank_id
    finally:
        # Cleanup
        await memory.delete_bank(bank_id, request_context=request_context)


@pytest.mark.asyncio
@patch("hindsight_api.tracing._tracing_enabled", False)
@patch("hindsight_api.engine.memory_engine.create_operation_span")
async def test_operations_work_when_tracing_disabled(mock_create_span, memory, request_context):
    """Test that operations work correctly when tracing is disabled."""
    # Setup - create_operation_span should return a no-op context manager
    from contextlib import nullcontext

    mock_create_span.return_value = nullcontext()

    bank_id = f"test-no-trace-{datetime.now(timezone.utc).timestamp()}"

    try:
        # All operations should work without errors
        await memory.retain_async(
            bank_id=bank_id,
            content="Test memory",
            request_context=request_context,
        )

        await memory.run_consolidation(
            bank_id=bank_id,
            request_context=request_context,
        )

        await memory.reflect_async(
            bank_id=bank_id,
            query="Test query",
            request_context=request_context,
        )

        # Verify no errors occurred and spans were attempted to be created
        assert mock_create_span.call_count >= 3
    finally:
        # Cleanup
        await memory.delete_bank(bank_id, request_context=request_context)