python/tests/test_networkx.py
"""
Real tests for ``KnutClient.to_networkx`` — no live FalkorDB required.
These exercise the *actual* method (not a re-implementation of the NetworkX
calls, as ``test_proto.py::TestNetworkXIntegration`` does). A fixture client
serves canned Arrow tables in place of the Flight SQL fetch, so the whole
production path runs for real:
to_networkx → query_df → query (fixture Arrow Table) → .to_pandas()
→ _build_networkx → nx.from_pandas_edgelist / set_node_attributes
Only the gRPC/DoGet hop is stubbed; the Arrow→pandas→NetworkX conversion under
test is the genuine article. The fixture mirrors ``examples/social_graph.py``
(Alice/Bob/Carol/Dave), so the example's shortest-path claim is checked here.
"""
from __future__ import annotations
import networkx as nx
import pandas as pd
import pyarrow as pa
import pytest
from knut import KnutClient
# ── Fixture graph — same shape as examples/social_graph.py ─────────────────────
VERTEX_QUERY = "MATCH (n:Person) RETURN n._gar_id AS _gar_id, n.name, n.age"
EDGE_QUERY = (
"MATCH (a:Person)-[r:KNOWS]->(b:Person) "
"RETURN a._gar_id AS src, b._gar_id AS dst, r.since"
)
# Alice(0) → Bob(1), Alice(0) → Carol(2), Bob(1) → Dave(3), Carol(2) → Dave(3)
_VERTICES = pa.table(
{
"_gar_id": pa.array([0, 1, 2, 3], type=pa.int64()),
"name": pa.array(["Alice", "Bob", "Carol", "Dave"], type=pa.string()),
"age": pa.array([30, 25, 35, 28], type=pa.int64()),
}
)
_EDGES = pa.table(
{
"src": pa.array([0, 0, 1, 2], type=pa.int64()),
"dst": pa.array([1, 2, 3, 3], type=pa.int64()),
"since": pa.array(
["2020-01-01", "2019-06-15", "2021-03-10", "2018-11-01"],
type=pa.string(),
),
}
)
class FixtureClient(KnutClient):
"""
A ``KnutClient`` whose Flight SQL fetch is replaced by a fixture lookup.
We deliberately bypass ``KnutClient.__init__`` (which would build a real
``FlightClient``) and override :meth:`query` — the single seam every read
goes through — to return canned Arrow tables keyed by the Cypher string.
Everything above ``query`` (``query_df.to_pandas``, ``to_networkx``,
``_build_networkx``) is the real, unmodified code.
"""
def __init__(self, tables: dict[str, pa.Table]) -> None: # noqa: D107
self._tables = tables
self.calls: list[str] = []
def query(self, cypher: str) -> pa.Table: # type: ignore[override]
self.calls.append(cypher)
try:
return self._tables[cypher]
except KeyError:
raise AssertionError(f"unexpected query: {cypher!r}") from None
@pytest.fixture
def client() -> FixtureClient:
return FixtureClient({VERTEX_QUERY: _VERTICES, EDGE_QUERY: _EDGES})
# ── The real thing: KnutClient.to_networkx end-to-end ──────────────────────────
def test_to_networkx_builds_directed_graph(client: FixtureClient) -> None:
G = client.to_networkx(vertex_query=VERTEX_QUERY, edge_query=EDGE_QUERY)
# It really went through the Flight-SQL seam (both queries fetched).
assert client.calls == [EDGE_QUERY, VERTEX_QUERY]
# Directedness — DiGraph by default.
assert isinstance(G, nx.DiGraph)
assert G.is_directed()
# Node + edge counts.
assert G.number_of_nodes() == 4
assert G.number_of_edges() == 4
def test_node_attributes_preserved(client: FixtureClient) -> None:
G = client.to_networkx(vertex_query=VERTEX_QUERY, edge_query=EDGE_QUERY)
assert G.nodes[0]["name"] == "Alice"
assert G.nodes[0]["age"] == 30
assert G.nodes[3]["name"] == "Dave"
assert G.nodes[3]["age"] == 28
# Every node carries both attribute columns.
for _node, attrs in G.nodes(data=True):
assert set(attrs) == {"name", "age"}
def test_edge_attributes_preserved(client: FixtureClient) -> None:
G = client.to_networkx(vertex_query=VERTEX_QUERY, edge_query=EDGE_QUERY)
assert G[0][1]["since"] == "2020-01-01"
assert G[0][2]["since"] == "2019-06-15"
assert G[1][3]["since"] == "2021-03-10"
assert G[2][3]["since"] == "2018-11-01"
def test_direction_is_not_symmetric(client: FixtureClient) -> None:
"""RED-when-broken: a DiGraph must not invent the reverse edge."""
G = client.to_networkx(vertex_query=VERTEX_QUERY, edge_query=EDGE_QUERY)
assert G.has_edge(0, 1)
assert not G.has_edge(1, 0)
assert G.has_edge(2, 3)
assert not G.has_edge(3, 2)
def test_shortest_path_matches_example(client: FixtureClient) -> None:
"""Mirrors social_graph.py: Alice(0) → Dave(3) over the directed graph."""
G = client.to_networkx(vertex_query=VERTEX_QUERY, edge_query=EDGE_QUERY)
path = nx.shortest_path(G, source=0, target=3)
assert path[0] == 0 and path[-1] == 3
assert len(path) == 3 # 0 → {1 or 2} → 3
assert path[1] in (1, 2)
# Dave (sink) is reachable from Alice; the reverse is not (directedness).
assert nx.has_path(G, 0, 3)
assert not nx.has_path(G, 3, 0)
# In-degree centrality: Dave and Carol both receive links; Alice receives none.
centrality = nx.in_degree_centrality(G)
assert centrality[0] == 0.0
assert centrality[3] > 0.0
def test_undirected_option(client: FixtureClient) -> None:
G = client.to_networkx(
vertex_query=VERTEX_QUERY, edge_query=EDGE_QUERY, directed=False
)
assert isinstance(G, nx.Graph)
assert not G.is_directed()
# Undirected: the reverse edge is now traversable.
assert G.has_edge(1, 0)
assert nx.has_path(G, 3, 0)
def test_custom_column_names() -> None:
"""to_networkx honours non-default id/src/dst column names."""
vertices = pa.table(
{
"vid": pa.array([10, 20], type=pa.int64()),
"label": pa.array(["x", "y"], type=pa.string()),
}
)
edges = pa.table(
{
"from": pa.array([10], type=pa.int64()),
"to": pa.array([20], type=pa.int64()),
}
)
c = FixtureClient({"V": vertices, "E": edges})
G = c.to_networkx(
vertex_query="V",
edge_query="E",
node_id_col="vid",
src_col="from",
dst_col="to",
)
assert G.number_of_nodes() == 2
assert G.has_edge(10, 20)
assert G.nodes[10]["label"] == "x"
# ── The pure seam directly, from fixture DataFrames ────────────────────────────
def test_build_networkx_pure_seam() -> None:
"""_build_networkx builds the same graph straight from DataFrames (no I/O)."""
G = KnutClient._build_networkx(
_VERTICES.to_pandas(),
_EDGES.to_pandas(),
)
assert G.number_of_nodes() == 4
assert G.number_of_edges() == 4
assert G.nodes[2]["name"] == "Carol"
assert G[0][1]["since"] == "2020-01-01"
def test_edges_with_no_attributes() -> None:
"""An edge frame with only src/dst yields edges with empty attr dicts."""
edges = pd.DataFrame({"src": [0, 1], "dst": [1, 2]})
vertices = pd.DataFrame({"_gar_id": [0, 1, 2], "name": ["a", "b", "c"]})
G = KnutClient._build_networkx(vertices, edges)
assert G.number_of_edges() == 2
assert G[0][1] == {} # no edge attributes
assert G.nodes[1]["name"] == "b"