topoteretes--cognee/cognee/tests/unit/modules/visualization/test_preprocessor.py

"""Unit tests for the visualization preprocessor.

The preprocessor is the single place where Story-view fields are derived
from raw graph adapter output. These tests pin the contract:

  - Every known node type maps to a non-default stage.
  - ``visual_rank`` prefers stamped ``topological_rank`` (Phase 1a) and
    falls back to a fixed stage order when unset.
  - ``contains`` / ``is_a`` edges are classified ``structural``.
  - Edges between the same stage pair sharing a relation collapse into one
    ``bundle_key`` so the renderer can bundle them.
  - Provenance is exposed only when at least one provenance field is set.
  - Color-map / schema-graph shape matches what the existing JS renderer
    already reads.
"""

import pytest

from cognee.modules.visualization.preprocessor import (
    OTHER_ENTITY_TYPES_LABEL,
    SCHEMA_MAX_ENTITY_TYPES,
    STAGE_ORDER,
    PreprocessedGraph,
    preprocess,
)


def _alice_like_graph():
    """A small graph that mirrors the shape of the canonical Alice example:
    one document, two chunks, three entities of two types, and one summary."""
    nodes_data = [
        ("doc1", {"type": "TextDocument", "name": "alice.md", "topological_rank": 1}),
        (
            "c1",
            {
                "type": "DocumentChunk",
                "text": "Alice knows Bob.",
                "source_pipeline": "cognify_pipeline",
                "source_task": "extract_chunks_from_documents",
                "topological_rank": 2,
            },
        ),
        (
            "c2",
            {
                "type": "DocumentChunk",
                "text": "NLP is a subfield of CS.",
                "source_pipeline": "cognify_pipeline",
                "source_task": "extract_chunks_from_documents",
                "topological_rank": 2,
            },
        ),
        (
            "alice",
            {
                "type": "Entity",
                "name": "Alice",
                "source_pipeline": "cognify_pipeline",
                "source_task": "extract_graph_from_data",
                "topological_rank": 3,
            },
        ),
        (
            "bob",
            {
                "type": "Entity",
                "name": "Bob",
                "source_pipeline": "cognify_pipeline",
                "source_task": "extract_graph_from_data",
                "topological_rank": 3,
            },
        ),
        (
            "nlp",
            {
                "type": "Entity",
                "name": "NLP",
                "source_pipeline": "cognify_pipeline",
                "source_task": "extract_graph_from_data",
                "topological_rank": 3,
            },
        ),
        ("person", {"type": "EntityType", "name": "Person", "topological_rank": 4}),
        ("field", {"type": "EntityType", "name": "Field", "topological_rank": 4}),
        (
            "sum1",
            {
                "type": "TextSummary",
                "text": "Alice and Bob in NLP.",
                "topological_rank": 5,
            },
        ),
    ]
    edges_data = [
        ("doc1", "c1", "contains", {}),
        ("doc1", "c2", "contains", {}),
        ("c1", "alice", "contains", {}),
        ("c1", "bob", "contains", {}),
        ("c2", "nlp", "contains", {}),
        ("alice", "person", "is_a", {}),
        ("bob", "person", "is_a", {}),
        ("nlp", "field", "is_a", {}),
        ("alice", "bob", "knows", {"relationship_name": "knows"}),
        ("c1", "sum1", "made_from", {}),
    ]
    return (nodes_data, edges_data)


def test_preprocess_returns_preprocessed_graph_dataclass():
    result = preprocess(_alice_like_graph())
    assert isinstance(result, PreprocessedGraph)
    assert len(result.nodes) == 9
    assert len(result.links) == 10


def test_stage_assignment_for_known_types():
    result = preprocess(_alice_like_graph())
    stages = {n["id"]: n["stage"] for n in result.nodes}
    assert stages["doc1"] == "document"
    assert stages["c1"] == "chunk"
    assert stages["c2"] == "chunk"
    assert stages["alice"] == "entity"
    assert stages["bob"] == "entity"
    assert stages["nlp"] == "entity"
    assert stages["person"] == "type"
    assert stages["field"] == "type"
    assert stages["sum1"] == "summary"


def test_stage_falls_through_to_other_for_unknown_types():
    nodes_data = [("x1", {"type": "MysteryType"})]
    edges_data = []
    result = preprocess((nodes_data, edges_data))
    assert result.nodes[0]["stage"] == "other"


def test_visual_rank_uses_stamped_topological_rank():
    """Phase 1a stamps topological_rank in the pipeline. The preprocessor
    must use that real value when it's a positive integer."""
    result = preprocess(_alice_like_graph())
    by_id = {n["id"]: n for n in result.nodes}
    assert by_id["doc1"]["visual_rank"] == 1
    assert by_id["c1"]["visual_rank"] == 2
    assert by_id["alice"]["visual_rank"] == 3
    assert by_id["person"]["visual_rank"] == 4
    assert by_id["sum1"]["visual_rank"] == 5


def test_visual_rank_falls_back_when_topological_rank_zero_or_none():
    """Legacy graphs (pre-Phase-1a) have rank=0 or rank=None on every node.
    The preprocessor must produce a usable rank from the stage."""
    nodes_data = [
        ("d", {"type": "TextDocument", "topological_rank": 0}),
        ("c", {"type": "DocumentChunk"}),  # no rank at all
        ("e", {"type": "Entity", "topological_rank": None}),
    ]
    edges_data = [("d", "c", "contains", {}), ("c", "e", "contains", {})]
    result = preprocess((nodes_data, edges_data))
    by_id = {n["id"]: n for n in result.nodes}
    # Stage-order fallback: document=1, chunk=2, entity=3
    assert by_id["d"]["visual_rank"] == STAGE_ORDER.index("document") + 1
    assert by_id["c"]["visual_rank"] == STAGE_ORDER.index("chunk") + 1
    assert by_id["e"]["visual_rank"] == STAGE_ORDER.index("entity") + 1


def test_has_meaningful_topological_rank_flag():
    """The renderer reads this flag to decide whether to use real ranks
    or fall back to its own type-based scheme."""
    real = preprocess(_alice_like_graph())
    assert real.has_meaningful_topological_rank is True

    nodes_data = [("d", {"type": "TextDocument", "topological_rank": 0})]
    edges_data = []
    legacy = preprocess((nodes_data, edges_data))
    assert legacy.has_meaningful_topological_rank is False


def test_structural_edges_classified_correctly():
    result = preprocess(_alice_like_graph())
    by_relation = {
        (link["source"], link["target"], link["relation"]): link for link in result.links
    }

    for key in [
        ("doc1", "c1", "contains"),
        ("doc1", "c2", "contains"),
        ("c1", "alice", "contains"),
        ("alice", "person", "is_a"),
        ("c1", "sum1", "made_from"),
    ]:
        assert by_relation[key]["edge_class"] == "structural", f"{key} should be structural"

    assert by_relation[("alice", "bob", "knows")]["edge_class"] == "semantic"


def test_bundle_key_collapses_structural_edges_into_groups():
    """The Alice-like graph has 5 ``contains`` edges, but they fall into two
    bundles: doc->chunk (2 edges) and chunk->entity (3 edges).

    This proves the renderer can replace 5 lines with 2 ribbons."""
    result = preprocess(_alice_like_graph())
    contains_bundles = {k: v for k, v in result.bundles.items() if "|contains" in k}
    assert len(contains_bundles) == 2
    counts = sorted(contains_bundles.values())
    assert counts == [2, 3]


def test_provenance_present_only_when_fields_set():
    result = preprocess(_alice_like_graph())
    by_id = {n["id"]: n for n in result.nodes}
    # doc1 has no provenance fields in the fixture — section must be hidden
    assert "provenance" not in by_id["doc1"]
    # c1 has source_pipeline and source_task set
    assert by_id["c1"]["provenance"] == {
        "source_pipeline": "cognify_pipeline",
        "source_task": "extract_chunks_from_documents",
    }


def test_color_maps_have_expected_keys():
    result = preprocess(_alice_like_graph())
    assert set(result.color_maps.keys()) == {"task", "pipeline", "node_set", "user"}
    # pipeline color map should contain the one pipeline that's set
    assert "cognify_pipeline" in result.color_maps["pipeline"]
    # task color map should contain both tasks
    assert "extract_chunks_from_documents" in result.color_maps["task"]
    assert "extract_graph_from_data" in result.color_maps["task"]


def test_pipeline_stages_in_canonical_order():
    result = preprocess(_alice_like_graph())
    # Story-view spine: document, chunk, entity, type, summary
    assert result.pipeline_stages == ["document", "chunk", "entity", "type", "summary"]


def test_degree_count_matches_edge_count():
    result = preprocess(_alice_like_graph())
    by_id = {n["id"]: n for n in result.nodes}
    # doc1 -> c1, doc1 -> c2  => degree 2
    assert by_id["doc1"]["degree"] == 2
    # c1: doc1->c1, c1->alice, c1->bob, c1->sum1 => degree 4
    assert by_id["c1"]["degree"] == 4


def test_label_priority_marks_documents_and_types_always():
    result = preprocess(_alice_like_graph())
    by_id = {n["id"]: n for n in result.nodes}
    # Documents and entity-types are landmarks; always labeled in Key mode
    assert by_id["doc1"]["label_priority"] is True
    assert by_id["person"]["label_priority"] is True
    assert by_id["field"]["label_priority"] is True


def test_edge_class_counts_summed():
    result = preprocess(_alice_like_graph())
    # 5 contains + 3 is_a + 1 made_from = 9 structural, 1 knows = 1 semantic
    assert result.edge_classes["structural"] == 9
    assert result.edge_classes["semantic"] == 1


def test_handles_3tuple_edges_without_edge_info():
    """Some adapters may yield 3-tuple edges (no edge_info dict)."""
    nodes_data = [("a", {"type": "Entity"}), ("b", {"type": "Entity"})]
    edges_data = [("a", "b", "knows")]  # 3-tuple
    result = preprocess((nodes_data, edges_data))
    assert len(result.links) == 1
    assert result.links[0]["edge_class"] == "semantic"


def test_node_color_preserved_from_type_map():
    """The preprocessor's TYPE_COLOR_MAP drives node colors in both the
    canvas and the legend swatches. Pin the canonical four so a color
    palette change doesn't silently break the visual encoding."""
    result = preprocess(_alice_like_graph())
    by_id = {n["id"]: n for n in result.nodes}
    assert by_id["alice"]["color"] == "#6510F4"  # Entity
    assert by_id["person"]["color"] == "#D5C2FF"  # EntityType
    assert by_id["c1"]["color"] == "#0DFF00"  # DocumentChunk
    assert (
        by_id["doc1"]["color"] == "#A550FF"
    )  # TextDocument (was default gray before Phase 1 polish)


def test_ontology_valid_overrides_color():
    """Ontology-grounded nodes get a distinct fill — it must differ from the
    unknown-type fallback gray so ontology matches stand apart visually."""
    nodes_data = [
        ("e", {"type": "Entity", "name": "X", "ontology_valid": True}),
        ("u", {"type": "MysteryType", "name": "Y"}),
    ]
    edges_data = []
    result = preprocess((nodes_data, edges_data))
    assert result.nodes[0]["color"] == "#FF5CA8"
    assert result.nodes[0]["color"] != result.nodes[1]["color"]


def test_schema_graph_falls_back_to_type_graph_when_no_schema_nodes():
    result = preprocess(_alice_like_graph())
    assert "nodes" in result.schema_graph
    assert "links" in result.schema_graph
    # Type-graph fallback emits one GraphNodeType per distinct semantic type.
    # Entity instances resolve to their EntityType via is_a, so "type:Entity"
    # is replaced by the resolved "type:Person" / "type:Field".
    type_node_ids = {
        n["id"] for n in result.schema_graph["nodes"] if n.get("type") == "GraphNodeType"
    }
    assert "type:TextDocument" in type_node_ids
    assert "type:DocumentChunk" in type_node_ids
    assert "type:Entity" not in type_node_ids
    assert "type:Person" in type_node_ids


def test_schema_graph_uses_schema_nodes_when_present():
    nodes_data = [
        (
            "users",
            {
                "type": "SchemaTable",
                "name": "users",
                "columns": '[{"name": "id", "type": "uuid"}]',
                "primary_key": "id",
            },
        ),
        (
            "posts",
            {
                "type": "SchemaTable",
                "name": "posts",
                "columns": '[{"name": "id", "type": "uuid"}]',
                "primary_key": "id",
            },
        ),
    ]
    edges_data = [("posts", "users", "has_relationship", {"relationship_name": "foreign_key"})]
    result = preprocess((nodes_data, edges_data))
    schema_node_types = {n.get("type") for n in result.schema_graph["nodes"]}
    assert "SchemaTable" in schema_node_types


def test_schema_type_nodes_resolve_semantic_types_via_is_a():
    """Entity instances collapse into their EntityType semantic types
    (Person/Field) via the is_a edge, so the literal "Entity" never appears."""
    result = preprocess(_alice_like_graph())
    type_nodes = {
        n["name"]: n for n in result.schema_graph["nodes"] if n["type"] == "GraphNodeType"
    }

    assert "Entity" not in type_nodes
    assert "Person" in type_nodes
    assert "Field" in type_nodes
    assert type_nodes["Person"]["instance_count"] == 2
    assert type_nodes["Field"]["instance_count"] == 1


def test_schema_type_nodes_carry_bounded_deterministic_samples():
    result = preprocess(_alice_like_graph())
    type_nodes = {
        n["name"]: n for n in result.schema_graph["nodes"] if n["type"] == "GraphNodeType"
    }

    person = type_nodes["Person"]
    # Alice has degree 3 (contains, is_a, knows), Bob has degree 3
    # (contains, is_a, knows). Tie on degree breaks to name order: Alice, Bob.
    assert person["samples"] == ["Alice", "Bob"]
    assert person["sample_size"] == 2
    assert type_nodes["Field"]["samples"] == ["NLP"]

    # Samples never exceed the per-type cap regardless of instance count.
    for node in type_nodes.values():
        assert node["sample_size"] <= 5
        assert len(node["samples"]) == node["sample_size"]


def test_schema_type_nodes_carry_full_relationship_distribution():
    result = preprocess(_alice_like_graph())
    type_nodes = {
        n["name"]: n for n in result.schema_graph["nodes"] if n["type"] == "GraphNodeType"
    }

    person_rels = {
        (r["relation"], r["to_type"]): r["count"] for r in type_nodes["Person"]["relationships"]
    }
    # Alice + Bob both is_a the EntityType node; alice knows bob (Person -> Person).
    assert person_rels[("is_a", "EntityType")] == 2
    assert person_rels[("knows", "Person")] == 1

    # DocumentChunk contains Person twice (alice, bob), Field once (nlp).
    chunk_rels = {
        (r["relation"], r["to_type"]): r["count"]
        for r in type_nodes["DocumentChunk"]["relationships"]
    }
    assert chunk_rels[("contains", "Person")] == 2
    assert chunk_rels[("contains", "Field")] == 1


def _many_entity_types_graph(num_types):
    """num_types semantic entity types with strictly descending instance counts:
    Type00 has num_types instances, Type01 has num_types-1, ... down to 1."""
    nodes_data = []
    edges_data = []
    for i in range(num_types):
        type_name = f"Type{i:02d}"
        type_id = f"etype{i}"
        nodes_data.append((type_id, {"type": "EntityType", "name": type_name}))
        for j in range(num_types - i):
            entity_id = f"e{i}_{j}"
            nodes_data.append((entity_id, {"type": "Entity", "name": f"{type_name}_inst{j}"}))
            edges_data.append((entity_id, type_id, "is_a", {}))
    return (nodes_data, edges_data)


def test_entity_type_long_tail_rolls_up_into_other_entities():
    """Beyond SCHEMA_MAX_ENTITY_TYPES, the tail of semantic entity types
    collapses into one rollup card so the Entity column stays bounded."""
    num_types = SCHEMA_MAX_ENTITY_TYPES + 3
    result = preprocess(_many_entity_types_graph(num_types))
    type_nodes = {
        n["name"]: n for n in result.schema_graph["nodes"] if n["type"] == "GraphNodeType"
    }

    semantic_cards = [name for name in type_nodes if name.startswith("Type")] + (
        [OTHER_ENTITY_TYPES_LABEL] if OTHER_ENTITY_TYPES_LABEL in type_nodes else []
    )
    assert len(semantic_cards) == SCHEMA_MAX_ENTITY_TYPES

    # Top types keep their own cards; the smallest types are rolled up.
    assert "Type00" in type_nodes
    assert f"Type{num_types - 1:02d}" not in type_nodes

    rollup = type_nodes[OTHER_ENTITY_TYPES_LABEL]
    assert rollup["rollup"] is True
    tail_size = num_types - (SCHEMA_MAX_ENTITY_TYPES - 1)
    assert len(rollup["rolled_up_types"]) == tail_size
    # Tail of descending counts ends at 1: tail_size + (tail_size-1) + ... + 1.
    assert rollup["instance_count"] == tail_size * (tail_size + 1) // 2
    # Rollup keeps the same rank as the kept entity-type cards (one column).
    assert rollup["rank"] == type_nodes["Type00"]["rank"]
    # The lead field announces the rollup.
    assert any(f["name"] == "entity types" for f in rollup["fields"])

    # Pair-relationship nodes never reference a rolled-up type name.
    rolled_names = {t["name"] for t in rollup["rolled_up_types"]}
    for node in result.schema_graph["nodes"]:
        if node["type"] == "GraphRelationshipType":
            assert node["source_type"] not in rolled_names
            assert node["target_type"] not in rolled_names

    # Instance drill-down still reaches the rolled-up instances.
    instances = result.schema_graph["instances_by_type"][OTHER_ENTITY_TYPES_LABEL]
    assert len(instances) == rollup["instance_count"]


def test_entity_types_under_cap_are_not_rolled_up():
    result = preprocess(_alice_like_graph())
    names = {n["name"] for n in result.schema_graph["nodes"] if n["type"] == "GraphNodeType"}
    assert OTHER_ENTITY_TYPES_LABEL not in names

    result_at_cap = preprocess(_many_entity_types_graph(SCHEMA_MAX_ENTITY_TYPES))
    names_at_cap = {
        n["name"] for n in result_at_cap.schema_graph["nodes"] if n["type"] == "GraphNodeType"
    }
    assert OTHER_ENTITY_TYPES_LABEL not in names_at_cap
    assert sum(1 for name in names_at_cap if name.startswith("Type")) == SCHEMA_MAX_ENTITY_TYPES


class TestUnnamedNodeFallbacks:
    UUID_NAME = "13e52fce-2d52-4a8b-9f01-aabbccddeeff"
    HASH_NAME = "a" * 64

    def test_uuid_name_gets_readable_placeholder(self):
        nodes_data = [("n1", {"type": "Entity", "name": self.UUID_NAME})]
        result = preprocess((nodes_data, []))
        node = result.nodes[0]
        assert node["name"].startswith("Unnamed Entity")
        assert self.UUID_NAME not in node["name"]
        assert node["is_unnamed"] is True

    def test_hash_fallback_fields_are_skipped(self):
        nodes_data = [
            ("n1", {"type": "DocumentChunk", "text": self.HASH_NAME, "description": "real text"})
        ]
        result = preprocess((nodes_data, []))
        assert result.nodes[0]["name"] == "real text"

    def test_unnamed_nodes_never_get_label_priority(self):
        # Documents are always-label landmarks — unless they are unnamed.
        nodes_data = [
            ("d1", {"type": "TextDocument", "name": self.UUID_NAME}),
            ("d2", {"type": "TextDocument", "name": "alice.md"}),
        ]
        result = preprocess((nodes_data, []))
        by_name_priority = {n["is_unnamed"]: n["label_priority"] for n in result.nodes}
        assert by_name_priority[True] is False
        assert by_name_priority[False] is True

    def test_regular_names_untouched(self):
        result = preprocess(_alice_like_graph())
        names = {n["name"] for n in result.nodes}
        assert "Alice" in names
        assert not any(name.startswith("Unnamed ") for name in names)


def test_empty_graph_does_not_crash():
    result = preprocess(([], []))
    assert result.nodes == []
    assert result.links == []
    assert result.has_meaningful_topological_rank is False
    assert result.pipeline_stages == []


def test_provenance_index_indexes_only_nodes_with_provenance():
    result = preprocess(_alice_like_graph())
    # c1 has provenance; doc1 doesn't
    assert "c1" in result.provenance_index
    assert "doc1" not in result.provenance_index


def test_operation_layer_maps_operations_to_present_types():
    """The transformation impact-layer emits operation nodes + typed links only
    for catalog operations that touch a type present in the graph, expanding
    'Entity' to the semantic entity types."""
    nodes = [
        ("d1", {"type": "TextDocument", "name": "a.txt"}),
        ("p1", {"type": "Entity", "name": "Carlos"}),
        ("t_person", {"type": "EntityType", "name": "Person"}),
    ]
    edges = [("p1", "t_person", "is_a", {})]
    schema = preprocess((nodes, edges)).schema_graph

    op_ids = {o["id"] for o in schema["operations"]}
    assert "op:cognify" in op_ids

    cognify_targets = {
        (link["target"], link["effect"])
        for link in schema["operation_links"]
        if link["source"] == "op:cognify"
    }
    assert ("type:TextDocument", "produces") in cognify_targets
    assert ("type:Person", "produces") in cognify_targets  # Entity → semantic type
    assert ("type:EntityType", "produces") in cognify_targets

    # An operation whose only targets are absent (Rule) is filtered out entirely.
    assert "op:coding_rule_associations" not in op_ids

    # Modify effects are surfaced (feedback weighting touches entity types).
    feedback_targets = {
        (link["target"], link["effect"])
        for link in schema["operation_links"]
        if link["source"] == "op:apply_feedback_weights"
    }
    assert ("type:Person", "modifies") in feedback_targets


if __name__ == "__main__":
    pytest.main([__file__, "-v"])