lib/extractor/tests/models/test_entity.py

import uuid
from uuid import UUID

from ntr_text_fragmentation.models.linker_entity import LinkerEntity
from tests.custom_entity import CustomEntity


class TestEntity:
    """Набор тестов для проверки класса LinkerEntity."""

    def test_create_entity(self):
        """Тест создания объекта LinkerEntity с заданными параметрами."""
        entity_id = uuid.uuid4()
        name = "Тестовая сущность"
        text = "Это текст тестовой сущности"
        in_search_text = "Текст для поиска"
        metadata = {"key": "value"}

        entity = LinkerEntity(
            id=entity_id,
            name=name,
            text=text,
            in_search_text=in_search_text,
            metadata=metadata
        )

        assert entity.id == entity_id
        assert entity.name == name
        assert entity.text == text
        assert entity.in_search_text == in_search_text
        assert entity.metadata == metadata
        assert entity.type == "Entity"  # Проверка значения по умолчанию
        assert entity.source_id is None  # Проверка опциональных полей
        assert entity.target_id is None
        assert entity.number_in_relation is None
        assert entity.is_link() is False  # Не является связью
        
    def test_create_entity_with_link_fields(self):
        """Тест создания объекта LinkerEntity с полями связи."""
        entity_id = uuid.uuid4()
        source_id = uuid.uuid4()
        target_id = uuid.uuid4()
        
        entity = LinkerEntity(
            id=entity_id,
            name="Тестовая связывающая сущность",
            text="Это текст связывающей сущности",
            in_search_text="Текст для поиска",
            metadata={"key": "value"},
            source_id=source_id,
            target_id=target_id,
            number_in_relation=1
        )
        
        assert entity.id == entity_id
        assert entity.source_id == source_id
        assert entity.target_id == target_id
        assert entity.number_in_relation == 1
        assert entity.is_link() is True  # Является связью

    def test_entity_fixture(self, sample_entity):
        """Тест использования фикстуры с предустановленными значениями."""
        assert sample_entity.id == UUID('12345678-1234-5678-1234-567812345678')
        assert sample_entity.name == "Тестовая сущность"
        assert sample_entity.text == "Текст тестовой сущности"
        assert sample_entity.in_search_text is None  # По умолчанию None
        assert sample_entity.metadata == {"test_key": "test_value"}
        assert sample_entity.type == "Entity"
        assert sample_entity.is_link() is False

    def test_auto_id_generation(self):
        """Тест автоматической генерации ID, если он не указан."""
        entity = LinkerEntity(
            id=None,
            name="Тест",
            text="Текст",
            metadata={}
        )

        assert entity.id is not None
        assert isinstance(entity.id, UUID)
        
    def test_invalid_link_fields(self):
        """Тест создания сущности с некорректными полями связи."""
        # Пробуем создать сущность только с source_id
        try:
            LinkerEntity(
                id=uuid.uuid4(),
                name="Некорректная сущность",
                text="Текст некорректной сущности",
                metadata={},
                source_id=uuid.uuid4()
            )
            assert False, "Создание сущности только с source_id должно вызывать исключение"
        except ValueError:
            pass
            
        # Пробуем создать сущность только с target_id
        try:
            LinkerEntity(
                id=uuid.uuid4(),
                name="Некорректная сущность",
                text="Текст некорректной сущности",
                metadata={},
                target_id=uuid.uuid4()
            )
            assert False, "Создание сущности только с target_id должно вызывать исключение"
        except ValueError:
            pass

    def test_custom_type(self):
        """Тест использования пользовательского типа."""
        custom_type = "CustomEntity"
        entity = LinkerEntity(
            id=uuid.uuid4(),
            name="Тест",
            text="Текст",
            metadata={},
            type=custom_type
        )

        assert entity.type == custom_type

    def test_to_string(self):
        """Тест метода __str__."""
        # Тест стандартного вывода
        entity = LinkerEntity(
            id=uuid.uuid4(),
            name="Тест",
            text="Текст",
            metadata={}
        )
        expected_string = "Тест: Текст"
        assert str(entity) == expected_string
        
        # Тест с использованием in_search_text
        entity_with_search = LinkerEntity(
            id=uuid.uuid4(),
            name="Тест",
            text="Текст",
            in_search_text="Текст для поиска",
            metadata={}
        )
        assert str(entity_with_search) == "Текст для поиска"

    def test_compare_same_entities(self):
        """Тест сравнения одинаковых сущностей."""
        entity_id = uuid.uuid4()
        name = "Сущность"
        text = "Текст"
        entity_type = "TestEntity"

        entity1 = LinkerEntity(
            id=entity_id,
            name=name,
            text=text,
            in_search_text="Текст для поиска 1",
            metadata={"a": 1},
            type=entity_type
        )

        entity2 = LinkerEntity(
            id=entity_id,
            name=name,
            text=text,
            in_search_text="Текст для поиска 2",  # Этот параметр не учитывается в сравнении
            metadata={"b": 2},                    # Этот параметр не учитывается в сравнении
            type=entity_type
        )

        assert entity1 == entity2

    def test_compare_different_entities(self, sample_entity):
        """Тест сравнения разных сущностей."""
        # Проверка с другим ID
        different_id_entity = LinkerEntity(
            id=uuid.uuid4(),
            name=sample_entity.name,
            text=sample_entity.text,
            metadata=sample_entity.metadata
        )
        assert sample_entity != different_id_entity

        # Проверка с другим именем
        different_name_entity = LinkerEntity(
            id=sample_entity.id,
            name="Другое имя",
            text=sample_entity.text,
            metadata=sample_entity.metadata
        )
        assert sample_entity != different_name_entity

        # Проверка с другим текстом
        different_text_entity = LinkerEntity(
            id=sample_entity.id,
            name=sample_entity.name,
            text="Другой текст",
            metadata=sample_entity.metadata
        )
        assert sample_entity != different_text_entity

        # Проверка с другим типом
        different_type_entity = LinkerEntity(
            id=sample_entity.id,
            name=sample_entity.name,
            text=sample_entity.text,
            metadata=sample_entity.metadata,
            type="ДругойТип"
        )
        assert sample_entity != different_type_entity

    def test_compare_with_other_class(self, sample_entity):
        """Тест сравнения с объектом другого класса."""
        # Создаем объект другого класса
        class OtherClass:
            pass

        other = OtherClass()
        assert sample_entity != other

    def test_serialize(self, sample_custom_entity):
        """Тест метода serialize для кастомного класса-наследника Entity."""
        # Сериализуем объект 
        serialized = sample_custom_entity.serialize()
        
        # Проверяем, что сериализованный объект - это базовый Entity
        assert isinstance(serialized, LinkerEntity)
        assert serialized.id == sample_custom_entity.id
        assert serialized.name == "Тестовый кастомный объект"
        assert serialized.text == "Текст кастомного объекта"
        # Проверяем, что тип соответствует имени класса согласно новой логике
        assert serialized.type == "CustomEntity"
        
        # Проверяем, что кастомные поля автоматически сохранены в метаданных
        assert "_custom_field1" in serialized.metadata
        assert "_custom_field2" in serialized.metadata
        assert serialized.metadata["_custom_field1"] == "custom_value"
        assert serialized.metadata["_custom_field2"] == 42
        assert serialized.metadata["original_key"] == "original_value"
        
    def test_deserialize(self):
        """Тест метода deserialize для кастомного класса-наследника Entity."""
        # Создаем базовый Entity с метаданными, как будто это сериализованный CustomEntity
        entity_id = uuid.uuid4()
        serialized_entity = LinkerEntity(
            id=entity_id,
            name="Тестовый кастомный объект",
            text="Текст кастомного объекта",
            in_search_text="Текст для поиска",
            metadata={
                "_custom_field1": "custom_value",
                "_custom_field2": 42,
                "original_key": "original_value"
            },
            type="CustomEntity"  # Используем имя класса при сериализации
        )
        
        # Десериализуем объект
        template = CustomEntity(
            id=uuid.uuid4(),  # Не важно, будет заменено
            name="",
            text="",
            in_search_text=None,
            metadata={},
            custom_field1="",
            custom_field2=0
        )
        deserialized = template.deserialize(serialized_entity)
        
        # Проверяем, что десериализованный объект корректно восстановил все поля
        assert isinstance(deserialized, CustomEntity)
        assert deserialized.id == entity_id
        assert deserialized.name == "Тестовый кастомный объект"
        assert deserialized.text == "Текст кастомного объекта"
        assert deserialized.in_search_text == "Текст для поиска"
        assert deserialized.custom_field1 == "custom_value"
        assert deserialized.custom_field2 == 42
        assert deserialized.metadata == {"original_key": "original_value"}
        
    def test_serialize_deserialize_roundtrip(self, sample_custom_entity):
        """Тест полного цикла сериализации и десериализации."""
        # Полный цикл сериализация -> десериализация
        serialized = sample_custom_entity.serialize()
        deserialized = sample_custom_entity.deserialize(serialized)
        
        # Проверяем, что объект после полного цикла идентичен исходному
        assert deserialized == sample_custom_entity
        assert deserialized.id == sample_custom_entity.id
        assert deserialized.name == sample_custom_entity.name
        assert deserialized.text == sample_custom_entity.text
        assert deserialized.in_search_text == sample_custom_entity.in_search_text
        assert deserialized.metadata == sample_custom_entity.metadata
        assert deserialized.type == sample_custom_entity.type
        assert deserialized.custom_field1 == sample_custom_entity.custom_field1
        assert deserialized.custom_field2 == sample_custom_entity.custom_field2
        
    def test_static_deserialize_method(self, sample_custom_entity):
        """Тест статического метода deserialize класса LinkerEntity."""
        # CustomEntity уже зарегистрирован через декоратор @register_entity
        
        # Сериализуем объект 
        serialized = sample_custom_entity.serialize()
        
        # Десериализуем через статический метод LinkerEntity.deserialize
        deserialized = LinkerEntity.deserialize(serialized)
        
        # Проверяем, что десериализация прошла правильно
        assert isinstance(deserialized, CustomEntity)
        assert deserialized == sample_custom_entity
        assert deserialized.id == sample_custom_entity.id
        assert deserialized.name == sample_custom_entity.name
        assert deserialized.text == sample_custom_entity.text
        assert deserialized.custom_field1 == sample_custom_entity.custom_field1
        assert deserialized.custom_field2 == sample_custom_entity.custom_field2