/cognee/tests/unit/infrastructure/engine/test_identity_fields.py

import logging
from uuid import UUID, uuid4

from cognee.infrastructure.engine.models.DataPoint import DataPoint


class PersonWithIdentity(DataPoint):
    name: str
    metadata: dict = {"index_fields": ["name"], "identity_fields": ["name"]}


class DepartmentWithIdentity(DataPoint):
    name: str
    metadata: dict = {"index_fields": ["name"], "identity_fields": ["name"]}


class MultiFieldIdentity(DataPoint):
    first_name: str
    last_name: str
    metadata: dict = {
        "index_fields": ["first_name"],
        "identity_fields": ["first_name", "last_name"],
    }


class NoIdentityFields(DataPoint):
    name: str
    metadata: dict = {"index_fields": ["name"]}


class PartialIdentity(DataPoint):
    name: str
    age: int = 0
    metadata: dict = {"index_fields": ["name"], "identity_fields": ["name", "age"]}


class TestSameValuesSameUUID:
    def test_same_name_produces_same_id(self):
        p1 = PersonWithIdentity(name="John")
        p2 = PersonWithIdentity(name="John")
        assert p1.id == p2.id

    def test_deterministic_across_calls(self):
        ids = [PersonWithIdentity(name="Alice").id for _ in range(10)]
        assert len(set(ids)) == 1


class TestDifferentValuesDifferentUUID:
    def test_different_names_produce_different_ids(self):
        p1 = PersonWithIdentity(name="John")
        p2 = PersonWithIdentity(name="Jane")
        assert p1.id != p2.id


class TestCrossTypeSafety:
    def test_same_name_different_class_different_id(self):
        person = PersonWithIdentity(name="Engineering")
        department = DepartmentWithIdentity(name="Engineering")
        assert person.id != department.id


class TestExplicitIdOverride:
    def test_explicit_id_takes_precedence(self):
        explicit_id = uuid4()
        p = PersonWithIdentity(id=explicit_id, name="John")
        assert p.id == explicit_id

    def test_explicit_id_differs_from_generated(self):
        explicit_id = uuid4()
        p_explicit = PersonWithIdentity(id=explicit_id, name="John")
        p_generated = PersonWithIdentity(name="John")
        assert p_explicit.id != p_generated.id


class TestDefaultFieldsIncludedInIdentity:
    def test_default_field_included_in_identity(self):
        """Fields with defaults are available via model_dump() and produce deterministic IDs."""
        p1 = PartialIdentity(name="John")
        p2 = PartialIdentity(name="John")
        assert p1.id == p2.id
        assert p1.id.version == 5

    def test_different_default_override_produces_different_id(self):
        p1 = PartialIdentity(name="John", age=25)
        p2 = PartialIdentity(name="John", age=30)
        assert p1.id != p2.id

    def test_truly_missing_field(self):
        """A class where identity_fields references a non-existent field."""

        class BadIdentity(DataPoint):
            name: str
            metadata: dict = {
                "index_fields": ["name"],
                "identity_fields": ["name", "nonexistent"],
            }

        b1 = BadIdentity(name="John")
        b2 = BadIdentity(name="John")
        # nonexistent is not in model_dump(), so falls back to UUID4 - different each time
        assert b1.id != b2.id


class TestNoIdentityFieldsBackwardCompat:
    def test_no_identity_fields_produces_random_uuid(self):
        n1 = NoIdentityFields(name="John")
        n2 = NoIdentityFields(name="John")
        assert n1.id != n2.id

    def test_base_datapoint_no_identity_fields(self):
        dp1 = DataPoint()
        dp2 = DataPoint()
        assert dp1.id != dp2.id


class TestMultiFieldIdentity:
    def test_same_multi_fields_same_id(self):
        m1 = MultiFieldIdentity(first_name="John", last_name="Doe")
        m2 = MultiFieldIdentity(first_name="John", last_name="Doe")
        assert m1.id == m2.id

    def test_different_multi_fields_different_id(self):
        m1 = MultiFieldIdentity(first_name="John", last_name="Doe")
        m2 = MultiFieldIdentity(first_name="John", last_name="Smith")
        assert m1.id != m2.id

    def test_field_order_matters(self):
        """first_name='A', last_name='B' should differ from first_name='B', last_name='A'."""
        m1 = MultiFieldIdentity(first_name="A", last_name="B")
        m2 = MultiFieldIdentity(first_name="B", last_name="A")
        assert m1.id != m2.id


class TestStringNormalization:
    def test_case_insensitive(self):
        p1 = PersonWithIdentity(name="John")
        p2 = PersonWithIdentity(name="JOHN")
        assert p1.id == p2.id

    def test_spaces_normalized(self):
        p1 = PersonWithIdentity(name="John Doe")
        p2 = PersonWithIdentity(name="John_Doe")
        assert p1.id == p2.id

    def test_apostrophes_removed(self):
        p1 = PersonWithIdentity(name="O'Brien")
        p2 = PersonWithIdentity(name="OBrien")
        assert p1.id == p2.id


class TestIdIsUUID5:
    def test_generated_id_is_valid_uuid(self):
        p = PersonWithIdentity(name="Test")
        assert isinstance(p.id, UUID)
        # UUID5 has version == 5
        assert p.id.version == 5

    def test_no_identity_id_is_uuid4(self):
        n = NoIdentityFields(name="Test")
        assert isinstance(n.id, UUID)
        assert n.id.version == 4


class TestInheritanceWarning:
    def test_subclass_dropping_identity_fields_logs_warning(self, caplog):
        """Subclass that overrides metadata without identity_fields triggers a warning."""

        class Parent(DataPoint):
            name: str
            metadata: dict = {"index_fields": ["name"], "identity_fields": ["name"]}

        class Child(Parent):
            name: str
            metadata: dict = {"index_fields": ["name"]}  # forgot identity_fields

        with caplog.at_level(logging.WARNING):
            c = Child(name="Test")
        assert "drops identity_fields" in caplog.text
        # Should fall back to UUID4 since no identity_fields on Child
        assert c.id.version == 4


class TestIdForMatchesIdentityFields:
    """id_for(...) and the identity_fields-derived instance id must be the SAME
    value — they share one implementation (id_for) precisely so they cannot drift.
    """

    def test_id_for_matches_single_field_instance(self):
        assert PersonWithIdentity.id_for("John") == PersonWithIdentity(name="John").id

    def test_id_for_matches_after_normalization(self):
        assert PersonWithIdentity.id_for("John Doe") == PersonWithIdentity(name="John_Doe").id

    def test_id_for_matches_multi_field_in_declaration_order(self):
        assert (
            MultiFieldIdentity.id_for("John", "Doe")
            == MultiFieldIdentity(first_name="John", last_name="Doe").id
        )

    def test_id_for_namespaced_by_class(self):
        # Same input, different classes -> different ids (class supplies the namespace).
        assert PersonWithIdentity.id_for("Engineering") != DepartmentWithIdentity.id_for(
            "Engineering"
        )


class TestNormalizationMatchesGenerateNodeId:
    """The identity normalization must stay byte-for-byte aligned with the legacy
    ``generate_node_id`` normalization, otherwise the graph id migration can no
    longer recompute historical ids from a node's stored name.
    """

    def test_normalization_equivalence(self):
        from uuid import NAMESPACE_OID, uuid5

        from cognee.infrastructure.engine.utils.generate_node_id import generate_node_id

        for raw in ["John", "John Doe", "O'Brien", "New York City", "ALL CAPS", "mixed Case's"]:
            normalized = DataPoint._normalize_identity_value(raw)
            # generate_node_id hashes the bare normalized value (no class prefix).
            assert generate_node_id(raw) == uuid5(NAMESPACE_OID, normalized)


class TestTypeFieldPreserved:
    def test_type_is_class_name(self):
        p = PersonWithIdentity(name="John")
        assert p.type == "PersonWithIdentity"

    def test_type_not_affected_by_identity_fields(self):
        n = NoIdentityFields(name="John")
        assert n.type == "NoIdentityFields"