basicmachines-co-basic-memory/tests/services/test_relation_service.py

"""Tests for RelationService."""

import pytest
import pytest_asyncio
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from basic_memory.models import Entity as EntityModel
from basic_memory.repository.relation_repository import RelationRepository
from basic_memory.schemas import Entity, Relation
from basic_memory.services.relation_service import RelationService


@pytest_asyncio.fixture
async def relation_repository(
    session_maker: async_sessionmaker[AsyncSession],
) -> RelationRepository:
    """Create a RelationRepository instance."""
    return RelationRepository(session_maker)


@pytest_asyncio.fixture
async def relation_service(relation_repository: RelationRepository) -> RelationService:
    """Create RelationService with repository."""
    return RelationService(relation_repository)


@pytest_asyncio.fixture
async def test_entities(
    session_maker: async_sessionmaker[AsyncSession],
) -> tuple[EntityModel, EntityModel]:
    """Create two test entities."""
    async with session_maker() as session:
        entity1 = EntityModel(
            id="test/test_entity_1",
            name="test_entity_1",
            entity_type="test",
            description="Test entity 1",
        )
        entity2 = EntityModel(
            id="test/test_entity_2",
            name="test_entity_2",
            entity_type="test",
            description="Test entity 2",
        )
        session.add_all([entity1, entity2])
        await session.commit()
        return entity1, entity2


@pytest.mark.asyncio
async def test_create_relation(
    relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel]
):
    """Test creating a basic relation between two entities."""
    entity1, entity2 = test_entities

    relation_data = Relation(from_id=entity1.id, to_id=entity2.id, relation_type="test_relation")

    relation = await relation_service.create_relation(relation_data)

    assert relation.from_id == entity1.id
    assert relation.to_id == entity2.id
    assert relation.relation_type == "test_relation"


@pytest.mark.asyncio
async def test_create_relations(
    relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel]
):
    """Test creating a basic relation between two entities."""
    entity1, entity2 = test_entities

    relation_data = [
        Relation(from_id=entity1.id, to_id=entity2.id, relation_type="type_0"),
        Relation(from_id=entity1.id, to_id=entity2.id, relation_type="type_1"),
    ]

    relations = await relation_service.create_relations(relation_data)
    assert len(relations) == 2
    relation0 = relations[0]
    assert relation0.from_id == entity1.id
    assert relation0.to_id == entity2.id
    assert relation0.relation_type == "type_0"

    relation1 = relations[1]
    assert relation1.from_id == entity1.id
    assert relation1.to_id == entity2.id
    assert relation1.relation_type == "type_1"


@pytest.mark.asyncio
async def test_create_relation_with_context(
    relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel]
):
    """Test creating a relation with context information."""
    entity1, entity2 = test_entities

    relation_data = Relation(
        from_id=entity1.id, to_id=entity2.id, relation_type="test_relation", context="test context"
    )

    relation = await relation_service.create_relation(relation_data)

    assert relation.context == "test context"


@pytest.mark.asyncio
async def test_delete_relation(
    relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel]
):
    """Test deleting a relation between entities."""
    entity1, entity2 = test_entities

    # Create a relation first
    relation_data = Relation(from_id=entity1.id, to_id=entity2.id, relation_type="test_relation")
    await relation_service.create_relation(relation_data)

    # Create Entity schema instances for delete_relation call
    from_entity = Entity(
        id=entity1.id,
        name=entity1.name,
        entity_type=entity1.entity_type,
        description=entity1.description,
        observations=[],
        relations=[],
    )
    to_entity = Entity(
        id=entity2.id,
        name=entity2.name,
        entity_type=entity2.entity_type,
        description=entity2.description,
        observations=[],
        relations=[],
    )

    # Delete the relation
    result = await relation_service.delete_relation(from_entity, to_entity, "test_relation")

    assert result is True


@pytest.mark.asyncio
async def test_delete_nonexistent_relation(
    relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel]
):
    """Test trying to delete a relation that doesn't exist."""
    entity1, entity2 = test_entities

    from_entity = Entity(
        id=entity1.id,
        name=entity1.name,
        entity_type=entity1.entity_type,
        description=entity1.description,
        observations=[],
        relations=[],
    )
    to_entity = Entity(
        id=entity2.id,
        name=entity2.name,
        entity_type=entity2.entity_type,
        description=entity2.description,
        observations=[],
        relations=[],
    )

    result = await relation_service.delete_relation(from_entity, to_entity, "nonexistent_relation")

    assert result is False


@pytest.mark.asyncio
async def test_delete_relations_by_criteria(
    relation_service: RelationService, test_entities: tuple[EntityModel, EntityModel]
):
    """Test deleting relations by criteria."""
    entity1, entity2 = test_entities

    # Create test relations
    await relation_service.create_relation(
        Relation(from_id=entity1.id, to_id=entity2.id, relation_type="relation1")
    )
    await relation_service.create_relation(
        Relation(from_id=entity1.id, to_id=entity2.id, relation_type="relation2")
    )

    # Delete relations matching criteria
    result = await relation_service.delete_relations(
        [{"from_id": entity1.id, "to_id": entity2.id, "relation_type": "relation1"}]
    )

    assert result is True