6b7e6b44f1
Python Build and Type Check / python-ci (ubuntu-latest, 3.11) (push) Has been cancelled
Python Build and Type Check / python-ci (ubuntu-latest, 3.13) (push) Has been cancelled
Python Build and Type Check / python-ci (windows-latest, 3.11) (push) Has been cancelled
Python Build and Type Check / python-ci (windows-latest, 3.13) (push) Has been cancelled
Python Integration Tests / python-ci (ubuntu-latest, 3.13) (push) Has been cancelled
Python Integration Tests / python-ci (windows-latest, 3.13) (push) Has been cancelled
Python Notebook Tests / python-ci (ubuntu-latest, 3.13) (push) Has been cancelled
Python Notebook Tests / python-ci (windows-latest, 3.13) (push) Has been cancelled
Python Smoke Tests / python-ci (ubuntu-latest, 3.13) (push) Has been cancelled
Python Smoke Tests / python-ci (windows-latest, 3.13) (push) Has been cancelled
Python Unit Tests / python-ci (ubuntu-latest, 3.13) (push) Has been cancelled
Python Unit Tests / python-ci (windows-latest, 3.13) (push) Has been cancelled
gh-pages / build (push) Has been cancelled
Python Publish (pypi) / Upload release to PyPI (push) Has been cancelled
Spellcheck / spellcheck (push) Has been cancelled
296 lines
9.9 KiB
Python
296 lines
9.9 KiB
Python
# Copyright (C) 2026 Microsoft
|
|
|
|
"""Tests for the cluster_graph operation.
|
|
|
|
These tests pin down the behavior of cluster_graph and its internal
|
|
_compute_leiden_communities function so that refactoring (vectorizing
|
|
iterrows, reducing copies, etc.) can be verified against known output.
|
|
"""
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
from graphrag.index.operations.cluster_graph import (
|
|
Communities,
|
|
cluster_graph,
|
|
)
|
|
|
|
|
|
def _make_edges(
|
|
rows: list[tuple[str, str, float]],
|
|
) -> pd.DataFrame:
|
|
"""Build a minimal relationships DataFrame from (source, target, weight)."""
|
|
return pd.DataFrame([{"source": s, "target": t, "weight": w} for s, t, w in rows])
|
|
|
|
|
|
def _node_sets(clusters: Communities) -> list[set[str]]:
|
|
"""Extract sorted-by-level list of node sets from cluster output."""
|
|
return [set(nodes) for _, _, _, nodes in clusters]
|
|
|
|
|
|
# -------------------------------------------------------------------
|
|
# Basic clustering
|
|
# -------------------------------------------------------------------
|
|
|
|
|
|
class TestClusterGraphBasic:
|
|
"""Verify basic clustering on small synthetic graphs."""
|
|
|
|
def test_single_triangle(self):
|
|
"""A single triangle should produce one community at level 0."""
|
|
edges = _make_edges([("X", "Y", 1.0), ("X", "Z", 1.0), ("Y", "Z", 1.0)])
|
|
clusters = cluster_graph(edges, max_cluster_size=10, use_lcc=False, seed=42)
|
|
|
|
assert len(clusters) == 1
|
|
level, _cid, parent, nodes = clusters[0]
|
|
assert level == 0
|
|
assert parent == -1
|
|
assert set(nodes) == {"X", "Y", "Z"}
|
|
|
|
def test_two_disconnected_cliques(self):
|
|
"""Two disconnected triangles should produce two communities."""
|
|
edges = _make_edges([
|
|
("A", "B", 1.0),
|
|
("A", "C", 1.0),
|
|
("B", "C", 1.0),
|
|
("D", "E", 1.0),
|
|
("D", "F", 1.0),
|
|
("E", "F", 1.0),
|
|
])
|
|
clusters = cluster_graph(edges, max_cluster_size=10, use_lcc=False, seed=42)
|
|
|
|
assert len(clusters) == 2
|
|
node_sets = _node_sets(clusters)
|
|
assert {"A", "B", "C"} in node_sets
|
|
assert {"D", "E", "F"} in node_sets
|
|
for level, _, parent, _ in clusters:
|
|
assert level == 0
|
|
assert parent == -1
|
|
|
|
def test_lcc_filters_to_largest_component(self):
|
|
"""With use_lcc=True, only the largest connected component is kept."""
|
|
edges = _make_edges([
|
|
("A", "B", 1.0),
|
|
("A", "C", 1.0),
|
|
("B", "C", 1.0),
|
|
("D", "E", 1.0),
|
|
("D", "F", 1.0),
|
|
("E", "F", 1.0),
|
|
])
|
|
clusters = cluster_graph(edges, max_cluster_size=10, use_lcc=True, seed=42)
|
|
|
|
assert len(clusters) == 1
|
|
all_nodes = set(clusters[0][3])
|
|
assert len(all_nodes) == 3
|
|
|
|
|
|
# -------------------------------------------------------------------
|
|
# Edge normalization
|
|
# -------------------------------------------------------------------
|
|
|
|
|
|
class TestEdgeNormalization:
|
|
"""Verify that direction normalization and deduplication work."""
|
|
|
|
def test_reversed_edges_produce_same_result(self):
|
|
"""Reversing all edge directions should yield identical clusters."""
|
|
forward = _make_edges([
|
|
("A", "B", 1.0),
|
|
("A", "C", 1.0),
|
|
("B", "C", 1.0),
|
|
("D", "E", 1.0),
|
|
("D", "F", 1.0),
|
|
("E", "F", 1.0),
|
|
])
|
|
backward = _make_edges([
|
|
("B", "A", 1.0),
|
|
("C", "A", 1.0),
|
|
("C", "B", 1.0),
|
|
("E", "D", 1.0),
|
|
("F", "D", 1.0),
|
|
("F", "E", 1.0),
|
|
])
|
|
clusters_fwd = cluster_graph(
|
|
forward, max_cluster_size=10, use_lcc=False, seed=42
|
|
)
|
|
clusters_bwd = cluster_graph(
|
|
backward, max_cluster_size=10, use_lcc=False, seed=42
|
|
)
|
|
|
|
assert _node_sets(clusters_fwd) == _node_sets(clusters_bwd)
|
|
|
|
def test_duplicate_edges_are_deduped(self):
|
|
"""A→B and B→A should be treated as one edge after normalization."""
|
|
edges = _make_edges([
|
|
("A", "B", 1.0),
|
|
("B", "A", 2.0),
|
|
("A", "C", 1.0),
|
|
("B", "C", 1.0),
|
|
])
|
|
clusters = cluster_graph(edges, max_cluster_size=10, use_lcc=False, seed=42)
|
|
|
|
assert len(clusters) == 1
|
|
assert set(clusters[0][3]) == {"A", "B", "C"}
|
|
|
|
def test_missing_weight_defaults_to_one(self):
|
|
"""Edges without a weight column should default to weight 1.0."""
|
|
edges = pd.DataFrame({
|
|
"source": ["A", "A", "B"],
|
|
"target": ["B", "C", "C"],
|
|
})
|
|
clusters = cluster_graph(edges, max_cluster_size=10, use_lcc=False, seed=42)
|
|
|
|
assert len(clusters) == 1
|
|
assert set(clusters[0][3]) == {"A", "B", "C"}
|
|
|
|
|
|
# -------------------------------------------------------------------
|
|
# Determinism
|
|
# -------------------------------------------------------------------
|
|
|
|
|
|
class TestDeterminism:
|
|
"""Verify that seeding produces reproducible results."""
|
|
|
|
def test_same_seed_same_result(self):
|
|
"""Identical seed should yield identical output."""
|
|
edges = _make_edges([
|
|
("A", "B", 1.0),
|
|
("A", "C", 1.0),
|
|
("B", "C", 1.0),
|
|
("D", "E", 1.0),
|
|
("D", "F", 1.0),
|
|
("E", "F", 1.0),
|
|
])
|
|
c1 = cluster_graph(edges, max_cluster_size=10, use_lcc=False, seed=123)
|
|
c2 = cluster_graph(edges, max_cluster_size=10, use_lcc=False, seed=123)
|
|
|
|
assert c1 == c2
|
|
|
|
def test_does_not_mutate_input(self):
|
|
"""cluster_graph should not modify the input DataFrame."""
|
|
edges = _make_edges([
|
|
("A", "B", 1.0),
|
|
("A", "C", 1.0),
|
|
("B", "C", 1.0),
|
|
])
|
|
original = edges.copy()
|
|
cluster_graph(edges, max_cluster_size=10, use_lcc=False, seed=42)
|
|
|
|
pd.testing.assert_frame_equal(edges, original)
|
|
|
|
|
|
# -------------------------------------------------------------------
|
|
# Output structure
|
|
# -------------------------------------------------------------------
|
|
|
|
|
|
class TestOutputStructure:
|
|
"""Verify the shape and types of the Communities output."""
|
|
|
|
def test_output_tuple_structure(self):
|
|
"""Each entry should be (level, community_id, parent, node_list)."""
|
|
edges = _make_edges([("A", "B", 1.0), ("A", "C", 1.0), ("B", "C", 1.0)])
|
|
clusters = cluster_graph(edges, max_cluster_size=10, use_lcc=False, seed=42)
|
|
|
|
for entry in clusters:
|
|
assert len(entry) == 4
|
|
level, cid, parent, nodes = entry
|
|
assert isinstance(level, int)
|
|
assert isinstance(cid, int)
|
|
assert isinstance(parent, int)
|
|
assert isinstance(nodes, list)
|
|
assert all(isinstance(n, str) for n in nodes)
|
|
|
|
def test_level_zero_has_parent_minus_one(self):
|
|
"""All level-0 clusters should have parent == -1."""
|
|
edges = _make_edges([
|
|
("A", "B", 1.0),
|
|
("A", "C", 1.0),
|
|
("B", "C", 1.0),
|
|
("D", "E", 1.0),
|
|
("D", "F", 1.0),
|
|
("E", "F", 1.0),
|
|
])
|
|
clusters = cluster_graph(edges, max_cluster_size=10, use_lcc=False, seed=42)
|
|
|
|
for level, _, parent, _ in clusters:
|
|
if level == 0:
|
|
assert parent == -1
|
|
|
|
def test_all_nodes_covered_at_each_level(self):
|
|
"""At any given level, the union of all community nodes should
|
|
equal exactly the set of all nodes in the graph for that level."""
|
|
edges = _make_edges([
|
|
("A", "B", 1.0),
|
|
("A", "C", 1.0),
|
|
("B", "C", 1.0),
|
|
("D", "E", 1.0),
|
|
("D", "F", 1.0),
|
|
("E", "F", 1.0),
|
|
])
|
|
clusters = cluster_graph(edges, max_cluster_size=10, use_lcc=False, seed=42)
|
|
|
|
levels: dict[int, set[str]] = {}
|
|
for level, _, _, nodes in clusters:
|
|
levels.setdefault(level, set()).update(nodes)
|
|
|
|
all_nodes = {"A", "B", "C", "D", "E", "F"}
|
|
for level, covered_nodes in levels.items():
|
|
assert covered_nodes == all_nodes, (
|
|
f"Level {level}: expected {all_nodes}, got {covered_nodes}"
|
|
)
|
|
|
|
|
|
# -------------------------------------------------------------------
|
|
# Real test data (golden file regression)
|
|
# -------------------------------------------------------------------
|
|
|
|
|
|
class TestClusterGraphRealData:
|
|
"""Regression tests using the shared test fixture data."""
|
|
|
|
@pytest.fixture
|
|
def relationships(self) -> pd.DataFrame:
|
|
"""Load the test relationships fixture."""
|
|
return pd.read_parquet("tests/verbs/data/relationships.parquet")
|
|
|
|
def test_cluster_count(self, relationships: pd.DataFrame):
|
|
"""Pin the expected number of clusters from the fixture data."""
|
|
clusters = cluster_graph(
|
|
relationships,
|
|
max_cluster_size=10,
|
|
use_lcc=True,
|
|
seed=0xDEADBEEF,
|
|
)
|
|
assert len(clusters) == 122
|
|
|
|
def test_level_distribution(self, relationships: pd.DataFrame):
|
|
"""Pin the expected number of clusters per level."""
|
|
clusters = cluster_graph(
|
|
relationships,
|
|
max_cluster_size=10,
|
|
use_lcc=True,
|
|
seed=0xDEADBEEF,
|
|
)
|
|
from collections import Counter
|
|
|
|
level_counts = Counter(c[0] for c in clusters)
|
|
assert level_counts == {0: 23, 1: 65, 2: 32, 3: 2}
|
|
|
|
def test_level_zero_nodes_sample(self, relationships: pd.DataFrame):
|
|
"""Spot-check a few known nodes in level-0 clusters."""
|
|
clusters = cluster_graph(
|
|
relationships,
|
|
max_cluster_size=10,
|
|
use_lcc=True,
|
|
seed=0xDEADBEEF,
|
|
)
|
|
level_0 = [c for c in clusters if c[0] == 0]
|
|
all_level_0_nodes = set()
|
|
for _, _, _, nodes in level_0:
|
|
all_level_0_nodes.update(nodes)
|
|
|
|
assert "SCROOGE" in all_level_0_nodes
|
|
assert "ABRAHAM" in all_level_0_nodes
|
|
assert "JACOB MARLEY" in all_level_0_nodes
|