basicmachines-co-basic-memory/tests/test_memory_service.py

"""Tests for the MemoryService class."""
import pytest
from basic_memory.services import MemoryService
from basic_memory.fileio import read_entity_file
from basic_memory.schemas import CreateEntityRequest, CreateRelationsRequest, AddObservationsRequest, RelationRequest

test_entities_data = [
    {
        "name": "Test_Entity_1",
        "entity_type": "test",
        "observations": ["Observation 1.1", "Observation 1.2"]
    },
    {
        "name": "Test_Entity_2",
        "entity_type": "test",
        "observations": ["Observation 2.1", "Observation 2.2"]
    }
]

@pytest.mark.asyncio
async def test_create_entities(memory_service: MemoryService):
    """Should create multiple entities in parallel with their observations."""

    entity_input = CreateEntityRequest.model_validate({"entities": test_entities_data})
    entities = await memory_service.create_entities(entity_input.entities)

    # Verify the SQLAlchemy models were created
    assert len(entities) == 2

    # Check first entity
    assert entities[0].name == "Test_Entity_1"
    assert entities[0].entity_type == "test"
    assert len(entities[0].observations) == 2
    assert entities[0].observations[0].content == "Observation 1.1"
    assert entities[0].observations[1].content == "Observation 1.2"

    # Check second entity
    assert entities[1].name == "Test_Entity_2"
    assert entities[1].entity_type == "test"
    assert len(entities[1].observations) == 2
    assert entities[1].observations[0].content == "Observation 2.1"
    assert entities[1].observations[1].content == "Observation 2.2"

    # Verify files were created (returns Pydantic Entity)
    file_entity1 = await read_entity_file(memory_service.entities_path, entities[0].id)
    file_entity2 = await read_entity_file(memory_service.entities_path, entities[1].id)
    
    assert file_entity1.name == "Test_Entity_1"
    assert file_entity2.name == "Test_Entity_2"

@pytest.mark.asyncio
async def test_add_observations(memory_service: MemoryService):
    """Should add observations to an existing entity."""
    entity_input = CreateEntityRequest.model_validate({"entities": test_entities_data})
    entities = await memory_service.create_entities([entity_input.entities[0]])
    entity = entities[0]

    # Create observations input
    observations_data = {
        "entity_id": entity.id,
        "observations": [
            "New observation 1",
            "New observation 2"
        ]
    }

    # Add observations - returns List[models.Observation]
    observation_input = AddObservationsRequest.model_validate(observations_data)
    added_observations = await memory_service.add_observations(observation_input)

    # Check the SQLAlchemy model results
    assert len(added_observations) == 2
    assert added_observations[0].content == "New observation 1"
    assert added_observations[0].context is None
    assert added_observations[1].content == "New observation 2" 
    assert added_observations[1].context is None

    # Verify file was updated - returns Pydantic Entity
    updated_entity = await read_entity_file(memory_service.entities_path, entity.id)
    assert len(updated_entity.observations) == 4  # 2 original + 2 new
    #assert updated_entity.observations[2] == "New observation 1"
    #assert updated_entity.observations[3] == "New observation 2"

    # Verify database - returns SQLAlchemy Entity
    db_entity = await memory_service.entity_service.get_entity(entity.id)
    assert len(db_entity.observations) == 4

@pytest.mark.asyncio
async def test_add_observations_nonexistent_entity(memory_service: MemoryService):
    """Should raise an appropriate error when adding observations to a non-existent entity."""
    observations_data = {
        "entity_id": "nonexistent-id",
        "observations": ["Test observation"]
    }
    
    with pytest.raises(Exception) as exc:  # We might want to define a specific error type
        observation_input = AddObservationsRequest.model_validate(observations_data)
        await memory_service.add_observations(observation_input)

@pytest.mark.asyncio
async def test_create_relations(memory_service: MemoryService):
    """Should create relations between entities and update both filesystem and database."""
    entity_input = CreateEntityRequest.model_validate({"entities": test_entities_data})
    entities = await memory_service.create_entities(entity_input.entities)
    entity1, entity2 = entities

    # Create test relations data using actual entity IDs
    test_relations_data = [
        {
            "from_id": entity1.id,
            "to_id": entity2.id,
            "relation_type": "connects_to"
        },
        {
            "from_id": entity2.id,
            "to_id": entity1.id,
            "relation_type": "references",
            "context": "test context"
        }
    ]

    # Create relations - returns List[models.Relation]
    input_args = CreateRelationsRequest.model_validate({"relations": test_relations_data})
    relations = await memory_service.create_relations(input_args.relations)

    # Verify SQLAlchemy Relation models were created
    assert len(relations) == 2

    # Check first relation
    assert relations[0].from_id == entity1.id
    assert relations[0].to_id == entity2.id
    assert relations[0].relation_type == "connects_to"
    assert relations[0].context is None

    # Check second relation
    assert relations[1].from_id == entity2.id
    assert relations[1].to_id == entity1.id
    assert relations[1].relation_type == "references"
    assert relations[1].context == "test context"

    # Read updated entities from filesystem - returns Pydantic Entities
    updated_entity1 = await read_entity_file(memory_service.entities_path, entity1.id)
    updated_entity2 = await read_entity_file(memory_service.entities_path, entity2.id)

    # Verify relations were added to entity1 in filesystem
    assert len(updated_entity1.relations) == 1
    assert updated_entity1.relations[0].from_id == entity1.id
    assert updated_entity1.relations[0].to_id == entity2.id
    assert updated_entity1.relations[0].relation_type == "connects_to"

    # Verify relations were added to entity2 in filesystem
    assert len(updated_entity2.relations) == 1
    assert updated_entity2.relations[0].from_id == entity2.id
    assert updated_entity2.relations[0].to_id == entity1.id
    assert updated_entity2.relations[0].relation_type == "references"
    assert updated_entity2.relations[0].context == "test context"

    # Now verify database state - get SQLAlchemy Entity models
    db_entity1 = await memory_service.entity_service.get_entity(entity1.id)
    db_entity2 = await memory_service.entity_service.get_entity(entity2.id)

    # Entity 1 should have one outgoing relation to entity 2
    assert len(db_entity1.outgoing_relations) == 1
    outgoing = db_entity1.outgoing_relations[0]
    assert outgoing.from_id == entity1.id
    assert outgoing.to_id == entity2.id
    assert outgoing.relation_type == "connects_to"
    assert outgoing.context is None

    # Entity 1 should have one incoming relation from entity 2
    assert len(db_entity1.incoming_relations) == 1
    incoming = db_entity1.incoming_relations[0]
    assert incoming.from_id == entity2.id
    assert incoming.to_id == entity1.id
    assert incoming.relation_type == "references"
    assert incoming.context == "test context"

    # Verify the same for entity 2 (reversed)
    assert len(db_entity2.outgoing_relations) == 1
    assert len(db_entity2.incoming_relations) == 1

@pytest.mark.asyncio
async def test_create_relations_with_invalid_entity_id(memory_service: MemoryService):
    """Should raise an appropriate error when trying to create relations with non-existent entity IDs."""
    # Create one entity - returns SQLAlchemy Entity
    entity_input = CreateEntityRequest.model_validate({"entities": test_entities_data})
    entities = await memory_service.create_entities([entity_input.entities[0]])
    entity1 = entities[0]

    # Try to create relation with non-existent entity ID
    bad_relation = {
        "from_id": entity1.id,
        "to_id": "nonexistent-id",
        "relation_type": "connects_to"
    }
    
    with pytest.raises(Exception) as exc:  # We might want to define a specific error type
        await memory_service.create_relations([RelationRequest.model_validate(bad_relation)])