basicmachines-co-basic-memory/tests/mcp/test_tool_add_observations.py

"""Tests for add_observations MCP tool."""

import pytest

from basic_memory.mcp.tools.knowledge import create_entities, add_observations
from basic_memory.schemas.base import ObservationCategory, Entity
from basic_memory.schemas.request import (
    CreateEntityRequest,
    AddObservationsRequest,
    ObservationCreate,
)


@pytest.mark.asyncio
async def test_add_basic_observation(client):
    """Test adding a single observation with default category."""
    # First create an entity to add observations to
    entity_request = CreateEntityRequest(entities=[Entity(title="TestEntity", entity_type="test")])
    result = await create_entities(entity_request)
    entity_id = result.entities[0].permalink

    # Add an observation
    request = AddObservationsRequest(
        permalink=entity_id, observations=[ObservationCreate(content="Test observation")]
    )
    updated = await add_observations(request)

    # Verify the observation was added
    assert len(updated.observations) == 1
    obs = updated.observations[0]
    assert obs.content == "Test observation"
    assert obs.category == ObservationCategory.NOTE  # Default category


@pytest.mark.asyncio
async def test_add_categorized_observations(client):
    """Test adding observations with different categories."""
    # Create test entity
    entity_request = CreateEntityRequest(entities=[Entity(title="TestEntity", entity_type="test")])
    result = await create_entities(entity_request)
    entity_id = result.entities[0].permalink

    # Add observations with different categories
    request = AddObservationsRequest(
        permalink=entity_id,
        observations=[
            ObservationCreate(
                content="Implementation uses SQLite", category=ObservationCategory.TECH
            ),
            ObservationCreate(
                content="Chose SQLite for simplicity", category=ObservationCategory.DESIGN
            ),
            ObservationCreate(
                content="Supports atomic operations", category=ObservationCategory.FEATURE
            ),
        ],
    )
    updated = await add_observations(request)

    assert len(updated.observations) == 3

    # Find and verify each observation by category
    tech_obs = next(o for o in updated.observations if o.category == ObservationCategory.TECH)
    design_obs = next(o for o in updated.observations if o.category == ObservationCategory.DESIGN)
    feature_obs = next(o for o in updated.observations if o.category == ObservationCategory.FEATURE)

    assert tech_obs.content == "Implementation uses SQLite"
    assert design_obs.content == "Chose SQLite for simplicity"
    assert feature_obs.content == "Supports atomic operations"


@pytest.mark.asyncio
async def test_add_observations_with_context(client):
    """Test adding observations with shared context."""
    # Create test entity
    entity_request = CreateEntityRequest(entities=[Entity(title="TestEntity", entity_type="test")])
    result = await create_entities(entity_request)
    entity_id = result.entities[0].permalink

    # Add observations with context
    shared_context = "Design meeting 2024-12-25"
    request = AddObservationsRequest(
        permalink=entity_id,
        context=shared_context,
        observations=[
            ObservationCreate(
                content="Decided on file format", category=ObservationCategory.DESIGN
            ),
            ObservationCreate(content="Will use markdown", category=ObservationCategory.TECH),
        ],
    )
    updated = await add_observations(request)

    assert len(updated.observations) == 2
    for obs in updated.observations:
        # Note: context handling depends on our schema - might need adjustment
        if hasattr(obs, "context"):
            assert obs.context == shared_context


@pytest.mark.asyncio
async def test_add_observations_preserves_existing(client):
    """Test that adding observations preserves existing ones."""
    # Create entity with initial observation
    entity_request = CreateEntityRequest(
        entities=[
            Entity(title="TestEntity", entity_type="test", observations=["Initial observation"])
        ]
    )
    result = await create_entities(entity_request)
    entity_id = result.entities[0].permalink

    # Add new observations
    request = AddObservationsRequest(
        permalink=entity_id,
        observations=[
            ObservationCreate(content="New observation", category=ObservationCategory.TECH)
        ],
    )
    updated = await add_observations(request)

    # Should have both observations
    assert len(updated.observations) == 2
    contents = {obs.content for obs in updated.observations}
    assert "Initial observation" in contents
    assert "New observation" in contents


@pytest.mark.asyncio
async def test_add_multiple_observations_same_category(client):
    """Test adding multiple observations in the same category."""
    # Create test entity
    entity_request = CreateEntityRequest(entities=[Entity(title="TestEntity", entity_type="test")])
    result = await create_entities(entity_request)
    entity_id = result.entities[0].permalink

    # Add multiple tech observations
    tech_observations = ["Uses async/await", "Implements SQLite backend", "Handles UTF-8 encoding"]

    request = AddObservationsRequest(
        permalink=entity_id,
        observations=[
            ObservationCreate(content=obs, category=ObservationCategory.TECH)
            for obs in tech_observations
        ],
    )
    updated = await add_observations(request)

    # Verify all observations were added with correct category
    assert len(updated.observations) == 3
    for obs in updated.observations:
        assert obs.category == ObservationCategory.TECH
        assert obs.content in tech_observations


@pytest.mark.asyncio
async def test_add_observation_to_nonexistent_entity(client):
    """Test adding observations to a non-existent entity fails."""
    # Create request for non-existent entity
    request = AddObservationsRequest(
        permalink="test/nonexistent",
        observations=[
            ObservationCreate(content="This should fail", category=ObservationCategory.NOTE)
        ],
    )

    # Should fail because entity doesn't exist
    with pytest.raises(Exception):  # Adjust exception type based on your error handling
        await add_observations(request)