topoteretes--cognee/cognee/tests/unit/migration/test_migration.py

"""Unit tests for the memory-migration module (COGX, sources, loader, formats).

All tests are pure: no databases, no LLM calls.
"""

import asyncio
import json
import xml.etree.ElementTree as ET
from datetime import datetime, timezone

from cognee.modules.migration.cogx import (
    COGX_VERSION,
    COGXArchiveWriter,
    COGXDocument,
    COGXEntity,
    COGXEpisode,
    COGXFact,
    COGXMemory,
    COGXMemoryBlock,
    COGXRawNode,
    COGXTurn,
    parse_timestamp,
    read_archive,
    read_manifest,
)
from cognee.modules.migration.formats import write_cypher, write_graphml, write_json
from cognee.modules.migration import loader
from cognee.modules.migration.loader import record_data_id, translate_records
from cognee.modules.migration.sources import (
    COGXArchiveSource,
    GraphitiSource,
    LettaSource,
    Mem0Source,
)


def collect(source):
    async def _collect():
        return [record async for record in source.records()]

    return asyncio.run(_collect())


class TestParseTimestamp:
    def test_iso_string(self):
        parsed = parse_timestamp("2024-03-01T12:00:00Z")
        assert parsed == datetime(2024, 3, 1, 12, 0, tzinfo=timezone.utc)

    def test_epoch_seconds_and_milliseconds(self):
        seconds = parse_timestamp(1709294400)
        milliseconds = parse_timestamp(1709294400000)
        assert seconds == milliseconds

    def test_epoch_micro_and_nanoseconds(self):
        seconds = parse_timestamp(1709294400)
        assert parse_timestamp(1709294400000000) == seconds
        assert parse_timestamp(1709294400000000000) == seconds

    def test_invalid_values(self):
        assert parse_timestamp(None) is None
        assert parse_timestamp("not-a-date") is None
        # Out-of-range epochs return None instead of raising.
        assert parse_timestamp(float("inf")) is None


class TestCOGXArchive:
    def _sample_records(self):
        return [
            COGXDocument(external_system="test", external_id="d1", content="hello"),
            COGXEpisode(
                external_system="test",
                external_id="e1",
                turns=[COGXTurn(role="user", content="hi")],
            ),
            COGXEntity(external_system="test", external_id="n1", name="Alice"),
            COGXFact(
                external_system="test",
                external_id="f1",
                subject_ref="n1",
                predicate="works_at",
                object_ref="Acme",
                valid_at=datetime(2024, 1, 1, tzinfo=timezone.utc),
            ),
            COGXMemory(external_system="test", external_id="m1", content="likes tea"),
            COGXMemoryBlock(external_system="test", external_id="b1", label="persona", value="x"),
        ]

    def test_roundtrip(self, tmp_path):
        archive_dir = tmp_path / "archive"
        records = self._sample_records()
        with COGXArchiveWriter(archive_dir, source_system="test") as writer:
            for record in records:
                writer.write(record)

        loaded = list(read_archive(archive_dir))
        assert {r.kind for r in loaded} == {r.kind for r in records}
        assert len(loaded) == len(records)

        fact = next(r for r in loaded if r.kind == "fact")
        assert fact.predicate == "works_at"
        assert fact.valid_at == datetime(2024, 1, 1, tzinfo=timezone.utc)

        manifest = read_manifest(archive_dir)
        assert manifest.source_system == "test"
        assert manifest.counts["fact"] == 1

    def test_archive_source_streams_records(self, tmp_path):
        archive_dir = tmp_path / "archive"
        with COGXArchiveWriter(archive_dir, source_system="test") as writer:
            for record in self._sample_records():
                writer.write(record)

        source = COGXArchiveSource(archive_dir)
        assert source.source_system == "test"
        assert source.mode == "preserve"
        assert len(collect(source)) == 6

    def test_raw_nodes_roundtrip(self, tmp_path):
        archive_dir = tmp_path / "archive"
        raw = {"id": "11111111-1111-1111-1111-111111111111", "type": "EntityType", "name": "Person"}
        with COGXArchiveWriter(archive_dir, source_system="test") as writer:
            writer.write(self._sample_records()[2])
            writer.write_raw_node(raw)

        loaded = list(read_archive(archive_dir))
        raw_nodes = [record for record in loaded if record.kind == "raw_node"]
        assert len(raw_nodes) == 1
        assert isinstance(raw_nodes[0], COGXRawNode)
        assert raw_nodes[0].properties == raw
        assert read_manifest(archive_dir).counts["raw_node"] == 1

    def test_writer_starts_from_clean_slate(self, tmp_path):
        archive_dir = tmp_path / "archive"
        for _ in range(2):
            with COGXArchiveWriter(archive_dir, source_system="test") as writer:
                writer.write(self._sample_records()[2])
                writer.write_raw_node({"id": "raw-1", "type": "NodeSet"})

        assert len((archive_dir / "entities.jsonl").read_text().splitlines()) == 1
        assert len((archive_dir / "nodes.jsonl").read_text().splitlines()) == 1
        manifest = read_manifest(archive_dir)
        assert manifest.counts == {"entity": 1, "raw_node": 1}

    def test_future_major_version_rejected(self, tmp_path):
        archive_dir = tmp_path / "archive"
        with COGXArchiveWriter(archive_dir, source_system="test") as writer:
            writer.write(self._sample_records()[2])
        manifest_path = archive_dir / "manifest.json"
        manifest = json.loads(manifest_path.read_text())
        manifest["cogx_version"] = "99.0"
        manifest_path.write_text(json.dumps(manifest))

        try:
            COGXArchiveSource(archive_dir)
            raise AssertionError("expected ValueError")
        except ValueError as error:
            assert "99.0" in str(error)
            assert COGX_VERSION in str(error)


class TestMem0Source:
    def test_plain_list(self):
        memories = collect(
            Mem0Source(
                [
                    {
                        "id": "abc",
                        "memory": "Alice works at Acme",
                        "user_id": "u1",
                        "categories": ["work"],
                        "created_at": "2024-01-01T00:00:00Z",
                    }
                ]
            )
        )
        assert len(memories) == 1
        assert memories[0].kind == "memory"
        assert memories[0].external_id == "abc"
        assert memories[0].scope.user_id == "u1"
        assert memories[0].categories == ["work"]

    def test_results_wrapper_and_text_key(self):
        memories = collect(Mem0Source({"results": [{"id": "1", "text": "remembers things"}]}))
        assert len(memories) == 1
        assert memories[0].content == "remembers things"

    def test_export_file(self, tmp_path):
        export_file = tmp_path / "mem0.json"
        export_file.write_text(json.dumps([{"id": "1", "memory": "from file"}]))
        memories = collect(Mem0Source(export_file))
        assert memories[0].content == "from file"

    def test_skips_items_without_content(self):
        assert collect(Mem0Source([{"id": "1"}, {"id": "2", "memory": "kept"}])) != []


class TestLettaSource:
    def test_agent_file(self):
        agent_file = {
            "agents": [
                {
                    "name": "assistant",
                    "core_memory": [
                        {"id": "blk1", "label": "persona", "value": "I am helpful", "limit": 2000},
                        {"label": "human", "value": ""},
                    ],
                    "messages": [
                        {"role": "user", "content": "hello", "created_at": "2024-01-01T00:00:00Z"},
                        {"role": "assistant", "content": [{"type": "text", "text": "hi there"}]},
                        {"role": "system", "content": "ignored"},
                    ],
                    "archival_memory": [{"id": "p1", "text": "archived note"}],
                }
            ]
        }
        records = collect(LettaSource(agent_file))
        kinds = [record.kind for record in records]
        assert kinds.count("memory_block") == 1  # empty block skipped
        assert kinds.count("episode") == 1
        assert kinds.count("document") == 1

        episode = next(r for r in records if r.kind == "episode")
        assert len(episode.turns) == 2  # system message excluded
        assert episode.scope.agent_id == "assistant"

        block = next(r for r in records if r.kind == "memory_block")
        assert block.label == "persona"
        assert block.limit == 2000


class TestZepSource:
    def test_graphiti_export(self):
        export = {
            "episodes": [
                {
                    "uuid": "ep1",
                    "name": "chat",
                    "content": "Alice said she moved to Berlin",
                    "created_at": "2024-02-01T00:00:00Z",
                    "group_id": "g1",
                }
            ],
            "nodes": [
                {
                    "uuid": "n1",
                    "name": "Alice",
                    "labels": ["Entity", "Person"],
                    "summary": "A person",
                },
                {"uuid": "n2", "name": "Berlin"},
            ],
            "edges": [
                {
                    "uuid": "f1",
                    "source_node_uuid": "n1",
                    "target_node_uuid": "n2",
                    "name": "LIVES_IN",
                    "fact": "Alice lives in Berlin",
                    "valid_at": "2024-02-01T00:00:00Z",
                    "invalid_at": None,
                    "episodes": ["ep1"],
                }
            ],
        }
        records = collect(GraphitiSource(export))
        assert [record.kind for record in records] == ["episode", "entity", "entity", "fact"]

        entity = next(r for r in records if r.kind == "entity")
        assert entity.entity_type == "Person"

        fact = records[-1]
        assert fact.subject_ref == "n1"
        assert fact.fact_text == "Alice lives in Berlin"
        assert fact.valid_at is not None
        assert fact.invalid_at is None
        assert fact.provenance == ["ep1"]

    def test_default_mode_is_hybrid(self):
        assert GraphitiSource({"nodes": []}).mode == "hybrid"


class TestTranslateRecords:
    def _records(self):
        return [
            COGXMemory(external_system="mem0", external_id="m1", content="likes tea"),
            COGXEpisode(
                external_system="zep",
                external_id="e1",
                turns=[
                    COGXTurn(
                        role="user",
                        content="hello",
                        occurred_at=datetime(2024, 1, 1, tzinfo=timezone.utc),
                    )
                ],
            ),
            COGXEntity(
                external_system="zep",
                external_id="n1",
                name="Alice",
                entity_type="Person",
                description="A person",
                aliases=["Alice Smith"],
            ),
            COGXFact(
                external_system="zep",
                external_id="f1",
                subject_ref="n1",
                predicate="lives_in",
                object_ref="Berlin",
                fact_text="Alice lives in Berlin",
                valid_at=datetime(2024, 2, 1, tzinfo=timezone.utc),
            ),
        ]

    def test_re_derive_mode(self):
        result = translate_records(self._records(), "re-derive")
        assert result.cognify_data_items is True
        assert result.graph_batches == []
        # memory + episode + entity digest + fact digest
        assert len(result.data_items) == 4
        assert all(item.data_id is not None for item in result.data_items)
        # Episode transcript carries the timestamp.
        episode_item = result.data_items[1]
        assert "2024-01-01" in episode_item.data
        # Fact digest carries validity qualifiers.
        fact_digest = result.data_items[3]
        assert "valid from 2024-02-01" in fact_digest.data

    def test_deterministic_data_ids(self):
        record = COGXMemory(external_system="mem0", external_id="m1", content="x")
        assert record_data_id(record) == record_data_id(record)
        first = translate_records(self._records(), "re-derive").data_items[0]
        second = translate_records(self._records(), "re-derive").data_items[0]
        assert first.data_id == second.data_id

    def test_preserve_mode(self):
        result = translate_records(self._records(), "preserve")
        assert result.cognify_data_items is False
        # Textual records still become (raw) data items.
        assert len(result.data_items) == 2
        assert len(result.graph_batches) == 1

        batch = result.graph_batches[0]
        node_names = {getattr(node, "name", None) for node in batch["nodes"]}
        # Alice, Berlin (created from unresolved ref), and the Person entity type.
        assert {"Alice", "Berlin", "Person"} <= node_names

        assert len(batch["edges"]) == 1
        source_id, target_id, relationship, properties = batch["edges"][0]
        assert relationship == "lives_in"
        assert properties["edge_text"] == "Alice lives in Berlin"
        assert properties["valid_at"].startswith("2024-02-01")
        assert properties["source_system"] == "zep"

        alice = next(node for node in batch["nodes"] if getattr(node, "name", "") == "Alice")
        assert "Also known as: Alice Smith" in alice.description
        assert alice.id == source_id

    def test_entity_merge_by_name(self):
        records = [
            COGXEntity(external_system="a", external_id="x1", name="Alice"),
            COGXEntity(external_system="b", external_id="x2", name="Alice"),
        ]
        batch = translate_records(records, "preserve").graph_batches[0]
        entities = [node for node in batch["nodes"] if type(node).__name__ == "Entity"]
        assert len(entities) == 1

    def test_entity_merge_combines_descriptions(self):
        records = [
            COGXEntity(external_system="a", external_id="x1", name="Alice", description="Founder"),
            COGXEntity(
                external_system="b",
                external_id="x2",
                name="Alice",
                description="Engineer",
                aliases=["Alice Smith"],
                entity_type="Person",
            ),
        ]
        batch = translate_records(records, "preserve").graph_batches[0]
        alice = next(node for node in batch["nodes"] if getattr(node, "name", "") == "Alice")
        assert "Founder" in alice.description
        assert "Engineer" in alice.description
        assert "Alice Smith" in alice.description
        assert alice.is_a is not None and alice.is_a.name == "Person"

    def test_raw_nodes_resolve_uuid_fact_refs(self):
        chunk_id = "0c113fd0-1111-2222-3333-444444444444"
        records = [
            COGXRawNode(properties={"id": chunk_id, "type": "DocumentChunk", "text": "hello"}),
            COGXEntity(
                external_system="cognee",
                external_id="55555555-5555-5555-5555-555555555555",
                name="Alice",
            ),
            COGXFact(
                external_system="cognee",
                external_id="f1",
                subject_ref=chunk_id,
                predicate="contains",
                object_ref="55555555-5555-5555-5555-555555555555",
            ),
        ]
        result = translate_records(records, "preserve")
        assert result.skipped_facts == 0
        batch = result.graph_batches[0]
        # The raw node was rehydrated with its original id, not fabricated.
        assert any(str(node.id) == chunk_id for node in batch["nodes"])
        assert not any(getattr(node, "name", "") == chunk_id for node in batch["nodes"])
        assert len(batch["edges"]) == 1
        assert str(batch["edges"][0][0]) == chunk_id

    def test_unresolved_uuid_fact_refs_are_skipped(self):
        records = [
            COGXEntity(external_system="z", external_id="n1", name="Alice"),
            COGXFact(
                external_system="z",
                external_id="f1",
                subject_ref="n1",
                predicate="mentions",
                object_ref="99999999-9999-9999-9999-999999999999",
            ),
        ]
        result = translate_records(records, "preserve")
        assert result.skipped_facts == 1
        batch = result.graph_batches[0]
        assert batch["edges"] == []
        # No Entity literally named by the dangling UUID.
        node_names = {getattr(node, "name", None) for node in batch["nodes"]}
        assert "99999999-9999-9999-9999-999999999999" not in node_names

    def test_graph_batches_are_bounded(self, monkeypatch):
        monkeypatch.setattr(loader, "BATCH_NODE_TARGET", 3)
        records = [
            COGXEntity(external_system="z", external_id=f"n{i}", name=f"Entity {i}")
            for i in range(8)
        ] + [
            COGXFact(
                external_system="z",
                external_id="f-cross",
                subject_ref="n0",
                predicate="knows",
                object_ref="n7",
            )
        ]
        result = translate_records(records, "preserve")
        assert len(result.graph_batches) > 1
        assert all(len(batch["nodes"]) <= 4 for batch in result.graph_batches)
        # The cross-batch fact's batch contains both endpoints.
        edge_batch = next(batch for batch in result.graph_batches if batch["edges"])
        source_id, target_id, _, _ = edge_batch["edges"][0]
        node_ids = {node.id for node in edge_batch["nodes"]}
        assert source_id in node_ids and target_id in node_ids


class TestFormatEmitters:
    NODES = [
        ("id-1", {"type": "Entity", "name": "Alice", "description": "A person"}),
        ("id-2", {"type": "Entity", "name": "Berlin", "tags": ["city", "capital"]}),
    ]
    EDGES = [("id-1", "id-2", "lives_in", {"edge_text": "Alice lives in Berlin"})]

    def test_json(self, tmp_path):
        destination = tmp_path / "graph.json"
        write_json(self.NODES, self.EDGES, destination)
        payload = json.loads(destination.read_text())
        assert len(payload["nodes"]) == 2
        assert payload["edges"][0]["relationship_name"] == "lives_in"
        assert payload["edges"][0]["edge_text"] == "Alice lives in Berlin"

    def test_graphml_is_valid_xml(self, tmp_path):
        destination = tmp_path / "graph.graphml"
        write_graphml(self.NODES, self.EDGES, destination)
        root = ET.parse(destination).getroot()
        namespace = "{http://graphml.graphdrawing.org/xmlns}"
        graph = root.find(f"{namespace}graph")
        assert len(graph.findall(f"{namespace}node")) == 2
        assert len(graph.findall(f"{namespace}edge")) == 1

    def test_cypher(self, tmp_path):
        destination = tmp_path / "graph.cypher"
        write_cypher(self.NODES, self.EDGES, destination)
        script = destination.read_text()
        # Indexed shared label keeps edge MATCH clauses off AllNodesScan plans.
        assert "CREATE INDEX IF NOT EXISTS FOR (n:CogneeNode) ON (n.id);" in script
        assert 'MERGE (n:CogneeNode {id: "id-1"}) SET n:`Entity`' in script
        assert "MERGE (a)-[r:`lives_in`]->(b)" in script
        # Non-scalar property serialized as JSON string, not raw repr.
        assert '"[\\"city\\", \\"capital\\"]"' in script


class TestImportModeValidation:
    def test_unknown_mode_rejected(self):
        try:
            Mem0Source([], mode="yolo")
            raise AssertionError("expected ValueError")
        except ValueError as error:
            assert "yolo" in str(error)