test_graph_store.py

"""Direct GraphStore unit tests (things the pipeline exercises indirectly)."""

from __future__ import annotations

from pathlib import Path

from sciogen.ir import Edge, EdgeKind, Node, NodeKind
from sciogen.stores.graph import GraphStore


def _node(node_id: str, kind: NodeKind, file: str) -> Node:
    return Node(
        id=node_id, kind=kind, name=node_id, file=file, language="python",
        line_start=0, line_end=0, params=[], return_type="", complexity=0,
        extends=[], implements=[], docstring="", fields=[], size=0,
        last_modified=0.0, package="", parent="", embedding_text="",
    )


def test_prune_keeps_referenced_placeholder_and_drops_orphans(tmp_path: Path):
    """A placeholder that is only a CALLS *target* (incoming/bwd edge) must be
    kept; a placeholder no edge touches must be dropped. Regression for the
    prune crash on large graphs where an undirected degree count undercounted
    incoming edges and DELETE was attempted on a still-connected node.
    """
    gs = GraphStore(tmp_path / "g.kuzu")
    real = _node("src/a.py:foo", NodeKind.FUNCTION, "src/a.py")
    referenced = _node("external:numpy.array", NodeKind.EXTERNAL, "")
    orphan = _node("external:dead.symbol", NodeKind.EXTERNAL, "")

    gs.begin()
    gs.upsert_nodes([real, referenced, orphan])
    gs.create_edges([Edge(real.id, referenced.id, EdgeKind.CALLS, 0.4)])
    gs.prune_orphan_placeholders()
    gs.commit()

    ids = {n["id"] for n in gs.all_nodes()}
    assert referenced.id in ids       # kept: it is a real edge target
    assert orphan.id not in ids       # pruned: nothing references it
    assert real.id in ids
    # the edge into the kept placeholder survived
    assert (real.id, referenced.id, EdgeKind.CALLS.value, 0.4) in gs.all_edges()
    gs.close()


def test_prune_is_noop_without_placeholders(tmp_path: Path):
    gs = GraphStore(tmp_path / "g.kuzu")
    a = _node("src/a.py:foo", NodeKind.FUNCTION, "src/a.py")
    b = _node("src/b.py:bar", NodeKind.FUNCTION, "src/b.py")
    gs.begin()
    gs.upsert_nodes([a, b])
    gs.create_edges([Edge(a.id, b.id, EdgeKind.CALLS, 1.0)])
    gs.prune_orphan_placeholders()
    gs.commit()
    assert {n["id"] for n in gs.all_nodes()} == {a.id, b.id}
    gs.close()