Testing Graph Learning Code: Unit Tests That Catch Real Mistakes
GNN research code is notoriously undertested. Experiments run, results look
plausible, and there is no mechanism to detect when a change breaks something
that was previously correct. This tutorial covers the specific tests that catch
the errors most common in graph learning code: shape mismatches, construction
failures, determinism violations, and edge cases in graph topology.
Why GNN Code Needs Its Own Test Patterns
Standard unit testing advice — test inputs and outputs, test edge cases — applies
here, but graph learning code has additional failure modes:
- Shape convention errors: The wrong transpose turns a correct-looking
tensor into wrong-semantics input - Device-split bugs: Model weights on GPU, data on CPU — silent on
construction, loud in the forward pass - Aggregation correctness: A scatter over wrong dimensions produces right
shape, wrong values - Batch offset errors:
GraphBatchedge index offsets are wrong for some
graph sizes - Non-determinism hiding bugs: A test passes most of the time but fails
occasionally due to floating-point non-determinism
Test patterns designed for these failure modes are different from typical
ML tests.
Test 1: Construction Validation Tests
Test that invalid constructions raise errors and valid ones succeed:
import pytest
import torch
from tgraphx import Graph
def make_valid_graph(N=10, E=20, d=8, de=4):
return dict(
node_features=torch.randn(N, d),
edge_index=torch.randint(0, N, (2, E)),
edge_features=torch.randn(E, de),
node_labels=torch.randint(0, 3, (N,)),
)
def test_valid_construction():
kwargs = make_valid_graph()
g = Graph(**kwargs)
assert g.node_features.shape == (10, 8)
assert g.edge_index.shape == (2, 20)
def test_wrong_edge_index_orientation():
kwargs = make_valid_graph()
kwargs['edge_index'] = torch.randint(0, 10, (20, 2)) # [E, 2] not [2, E]
with pytest.raises(Exception):
Graph(**kwargs)
def test_edge_count_mismatch():
kwargs = make_valid_graph()
kwargs['edge_features'] = torch.randn(25, 4) # 25 rows, but E=20
with pytest.raises(Exception):
Graph(**kwargs)
def test_node_label_count_mismatch():
kwargs = make_valid_graph()
kwargs['node_labels'] = torch.randint(0, 3, (15,)) # 15, but N=10
with pytest.raises(Exception):
Graph(**kwargs)
These tests document the expected behavior of the constructor and catch
regressions if the validation logic changes.
Test 2: Device Consistency Tests
def test_to_device_moves_all_tensors():
device = "cpu" # use CPU for portability in CI
kwargs = make_valid_graph()
g = Graph(**kwargs).to(device)
assert g.node_features.device.type == device
assert g.edge_index.device.type == device
assert g.edge_features.device.type == device
assert g.node_labels.device.type == device
def test_gpu_if_available():
if not torch.cuda.is_available():
pytest.skip("CUDA not available")
kwargs = make_valid_graph()
g = Graph(**kwargs).to("cuda")
assert g.node_features.is_cuda
assert g.edge_index.is_cuda
Test 3: Shape Propagation Through Layers
The most important test for any GNN layer: does the output shape match what
the documentation says?
from tgraphx.layers import TensorGraphSAGELayer
from tgraphx import Graph
def test_sage_layer_output_shape():
N, E, in_dim, out_dim = 15, 30, 32, 64
g = Graph(
node_features=torch.randn(N, in_dim),
edge_index=torch.randint(0, N, (2, E)),
edge_features=torch.randn(E, 8),
node_labels=torch.randint(0, 3, (N,)),
)
layer = TensorGraphSAGELayer(in_dim, out_dim)
out = layer(g)
# Output should be [N, out_dim]
assert out.shape == (N, out_dim), (
f"Expected ({N}, {out_dim}), got {out.shape}"
)
Write one of these tests for each layer you use. Run them as part of your
experiment setup to catch API changes across TGraphX versions.
Test 4: Determinism Tests
Non-determinism in GNN code comes from CUDA operations, random initialization,
and data loading order. Tests that verify determinism help catch subtle bugs
where results change between runs:
from tgraphx.reproducibility import set_seed
def make_model_output(seed, g):
from tgraphx.layers import TensorGraphSAGELayer
set_seed(seed, deterministic=True)
layer = TensorGraphSAGELayer(8, 16)
with torch.no_grad():
return layer(g).clone()
def test_determinism_with_same_seed():
kwargs = make_valid_graph(d=8)
g = Graph(**kwargs)
out1 = make_model_output(seed=42, g=g)
out2 = make_model_output(seed=42, g=g)
assert torch.allclose(out1, out2), "Same seed should produce same output"
def test_different_seeds_differ():
kwargs = make_valid_graph(d=8)
g = Graph(**kwargs)
out1 = make_model_output(seed=42, g=g)
out2 = make_model_output(seed=99, g=g)
assert not torch.allclose(out1, out2), "Different seeds should differ"
Test 5: GraphBatch Correctness
Test that GraphBatch.from_graphs correctly offsets edge indices and
concatenates features:
from tgraphx import GraphBatch
def test_graph_batch_node_count():
g1 = Graph(**make_valid_graph(N=5, E=8))
g2 = Graph(**make_valid_graph(N=7, E=12))
g3 = Graph(**make_valid_graph(N=4, E=6))
batch = GraphBatch.from_graphs([g1, g2, g3])
assert batch.node_features.shape[0] == 5 + 7 + 4 # N total
def test_graph_batch_edge_index_valid():
g1 = Graph(**make_valid_graph(N=5, E=8))
g2 = Graph(**make_valid_graph(N=7, E=12))
batch = GraphBatch.from_graphs([g1, g2])
N_total = batch.node_features.shape[0]
# All edge indices should be within [0, N_total)
assert batch.edge_index.min() >= 0
assert batch.edge_index.max() < N_total
def test_graph_batch_no_cross_edges():
# Edges from g1 should only reference g1 nodes (indices 0..N1-1)
# Edges from g2 should only reference g2 nodes (indices N1..N1+N2-1)
g1 = Graph(**make_valid_graph(N=5, E=8))
g2 = Graph(**make_valid_graph(N=7, E=12))
batch = GraphBatch.from_graphs([g1, g2])
# g1 edges: first 8 columns of edge_index
g1_edges = batch.edge_index[:, :8]
assert g1_edges.max() < 5, "g1 edges reference g2 nodes — offset bug"
# g2 edges: next 12 columns
g2_edges = batch.edge_index[:, 8:]
assert g2_edges.min() >= 5, "g2 edges reference g1 nodes — offset bug"
assert g2_edges.max() < 12, "g2 edges out of range"
Test 6: Graph Builder Shapes
from tgraphx.graph_builders import build_knn_graph, build_grid_graph
def test_knn_graph_shape():
node_features = torch.randn(20, 32)
edge_index = build_knn_graph(node_features, k=4)
assert edge_index.shape[0] == 2, "edge_index must be [2, E]"
assert edge_index.shape[1] <= 20 * 4 # at most k edges per node
def test_grid_graph_shape():
H, W = 4, 5
g = build_grid_graph(H, W)
assert g.node_features.shape[0] == H * W
assert g.edge_index.shape[0] == 2
Test 7: Overfit Test (Integration)
The 5-sample overfit test is a valuable integration test that catches bugs
in the full training pipeline that unit tests miss:
def test_model_can_overfit_small_batch():
from tgraphx import GraphBatch
from tgraphx.layers import TensorGraphSAGELayer
import torch.nn as nn, torch.optim as optim
graphs = [Graph(**make_valid_graph(N=8, E=15, d=16)) for _ in range(5)]
batch = GraphBatch.from_graphs(graphs)
class TinyGNN(nn.Module):
def __init__(self):
super().__init__()
self.layer = TensorGraphSAGELayer(16, 3)
def forward(self, g):
return self.layer(g)
set_seed(42)
model = TinyGNN()
opt = optim.Adam(model.parameters(), lr=0.05)
crit = nn.CrossEntropyLoss()
for _ in range(200):
opt.zero_grad()
out = model(batch)
loss = crit(out, batch.node_labels)
loss.backward()
opt.step()
# Should overfit completely on 5 small graphs
final_loss = loss.item()
assert final_loss < 0.1, f"Failed to overfit: loss={final_loss:.4f}"
Running the Tests
Organize tests in a tests/ directory alongside your project code:
project/
├── model.py
├── data.py
└── tests/
├── test_graph_construction.py
├── test_layers.py
├── test_batch.py
└── test_integration.py
Run with pytest:
pytest tests/ -v
These tests are fast (CPU, small graphs) and should run in under 30 seconds.
Run them before every training job and after every code change.