wehub-resource-sync
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
445 lines
15 KiB
Python
445 lines
15 KiB
Python
# Copyright (C) 2026 Microsoft
|
|
# Licensed under the MIT License
|
|
|
|
"""Tests for the finalize_graph streaming functions.
|
|
|
|
Covers _build_degree_map, finalize_entities, finalize_relationships,
|
|
and the orchestrating finalize_graph function.
|
|
"""
|
|
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from graphrag.data_model.schemas import (
|
|
ENTITIES_FINAL_COLUMNS,
|
|
RELATIONSHIPS_FINAL_COLUMNS,
|
|
)
|
|
from graphrag.index.operations.finalize_entities import finalize_entities
|
|
from graphrag.index.operations.finalize_relationships import (
|
|
finalize_relationships,
|
|
)
|
|
from graphrag.index.workflows.finalize_graph import (
|
|
_build_degree_map,
|
|
finalize_graph,
|
|
)
|
|
from graphrag_storage.tables.table import Table
|
|
|
|
|
|
class FakeTable(Table):
|
|
"""In-memory table that supports async iteration and write collection.
|
|
|
|
Rows passed to write() are collected in ``written`` for assertions.
|
|
Each call to ``__aiter__`` resets the read cursor so the table can
|
|
be iterated multiple times, matching the real CSVTable behavior
|
|
under truncate mode.
|
|
"""
|
|
|
|
def __init__(self, rows: list[dict[str, Any]] | None = None) -> None:
|
|
self._rows = list(rows or [])
|
|
self._index = 0
|
|
self.written: list[dict[str, Any]] = []
|
|
|
|
def __aiter__(self):
|
|
"""Return an async iterator over the seed rows."""
|
|
self._index = 0
|
|
return self
|
|
|
|
async def __anext__(self) -> dict[str, Any]:
|
|
"""Yield the next row or stop."""
|
|
if self._index >= len(self._rows):
|
|
raise StopAsyncIteration
|
|
row = dict(self._rows[self._index])
|
|
self._index += 1
|
|
return row
|
|
|
|
async def length(self) -> int:
|
|
"""Return number of seed rows."""
|
|
return len(self._rows)
|
|
|
|
async def has(self, row_id: str) -> bool:
|
|
"""Check if a row with the given ID exists in seed rows."""
|
|
return any(r.get("id") == row_id for r in self._rows)
|
|
|
|
async def write(self, row: dict[str, Any]) -> None:
|
|
"""Collect written rows for test assertions."""
|
|
self.written.append(row)
|
|
|
|
async def close(self) -> None:
|
|
"""No-op."""
|
|
|
|
|
|
def _make_entity_row(
|
|
title: str,
|
|
entity_type: str = "ENTITY",
|
|
description: str = "",
|
|
frequency: int = 1,
|
|
) -> dict[str, Any]:
|
|
"""Build a minimal entity row matching pre-finalization shape."""
|
|
return {
|
|
"title": title,
|
|
"type": entity_type,
|
|
"description": description,
|
|
"frequency": frequency,
|
|
"text_unit_ids": ["tu1"],
|
|
}
|
|
|
|
|
|
def _make_relationship_row(
|
|
source: str,
|
|
target: str,
|
|
weight: float = 1.0,
|
|
description: str = "",
|
|
) -> dict[str, Any]:
|
|
"""Build a minimal relationship row matching pre-finalization shape."""
|
|
return {
|
|
"source": source,
|
|
"target": target,
|
|
"weight": weight,
|
|
"description": description,
|
|
"text_unit_ids": ["tu1"],
|
|
}
|
|
|
|
|
|
class TestBuildDegreeMap:
|
|
"""Tests for the streaming _build_degree_map helper."""
|
|
|
|
async def test_simple_triangle(self):
|
|
"""Three nodes forming a triangle should each have degree 2."""
|
|
table = FakeTable([
|
|
_make_relationship_row("A", "B"),
|
|
_make_relationship_row("B", "C"),
|
|
_make_relationship_row("A", "C"),
|
|
])
|
|
result = await _build_degree_map(table)
|
|
assert result == {"A": 2, "B": 2, "C": 2}
|
|
|
|
async def test_star_topology(self):
|
|
"""Hub connected to four leaves should have degree 4; leaves degree 1."""
|
|
table = FakeTable([
|
|
_make_relationship_row("hub", "a"),
|
|
_make_relationship_row("hub", "b"),
|
|
_make_relationship_row("hub", "c"),
|
|
_make_relationship_row("hub", "d"),
|
|
])
|
|
result = await _build_degree_map(table)
|
|
assert result["hub"] == 4
|
|
for leaf in ("a", "b", "c", "d"):
|
|
assert result[leaf] == 1
|
|
|
|
async def test_duplicate_edges_deduplicated(self):
|
|
"""Duplicate (A, B) edges should be counted only once."""
|
|
table = FakeTable([
|
|
_make_relationship_row("A", "B"),
|
|
_make_relationship_row("A", "B"),
|
|
])
|
|
result = await _build_degree_map(table)
|
|
assert result == {"A": 1, "B": 1}
|
|
|
|
async def test_reversed_duplicate_edges_deduplicated(self):
|
|
"""(A, B) and (B, A) are the same undirected edge."""
|
|
table = FakeTable([
|
|
_make_relationship_row("A", "B"),
|
|
_make_relationship_row("B", "A"),
|
|
])
|
|
result = await _build_degree_map(table)
|
|
assert result == {"A": 1, "B": 1}
|
|
|
|
async def test_empty_table(self):
|
|
"""Empty relationship table should produce empty degree map."""
|
|
table = FakeTable([])
|
|
result = await _build_degree_map(table)
|
|
assert result == {}
|
|
|
|
async def test_single_edge(self):
|
|
"""One edge yields degree 1 for both endpoints."""
|
|
table = FakeTable([_make_relationship_row("X", "Y")])
|
|
result = await _build_degree_map(table)
|
|
assert result == {"X": 1, "Y": 1}
|
|
|
|
async def test_disconnected_components(self):
|
|
"""Two separate components computed correctly."""
|
|
table = FakeTable([
|
|
_make_relationship_row("A", "B"),
|
|
_make_relationship_row("C", "D"),
|
|
])
|
|
result = await _build_degree_map(table)
|
|
assert result == {"A": 1, "B": 1, "C": 1, "D": 1}
|
|
|
|
|
|
class TestFinalizeEntities:
|
|
"""Tests for stream-finalize entity rows."""
|
|
|
|
async def test_enriches_with_degree(self):
|
|
"""Entities should receive degree from the degree map."""
|
|
table = FakeTable([
|
|
_make_entity_row("A"),
|
|
_make_entity_row("B"),
|
|
])
|
|
degree_map = {"A": 3, "B": 1}
|
|
await finalize_entities(table, degree_map)
|
|
|
|
assert len(table.written) == 2
|
|
assert table.written[0]["degree"] == 3
|
|
assert table.written[1]["degree"] == 1
|
|
|
|
async def test_missing_degree_defaults_to_zero(self):
|
|
"""Entity not in degree map should get degree 0."""
|
|
table = FakeTable([_make_entity_row("UNKNOWN")])
|
|
degree_map = {"A": 5}
|
|
await finalize_entities(table, degree_map)
|
|
|
|
assert len(table.written) == 1
|
|
assert table.written[0]["degree"] == 0
|
|
|
|
async def test_deduplicates_by_title(self):
|
|
"""Duplicate titles should be skipped."""
|
|
table = FakeTable([
|
|
_make_entity_row("A"),
|
|
_make_entity_row("A"),
|
|
_make_entity_row("B"),
|
|
])
|
|
degree_map = {"A": 1, "B": 2}
|
|
await finalize_entities(table, degree_map)
|
|
|
|
assert len(table.written) == 2
|
|
titles = [r["title"] for r in table.written]
|
|
assert titles == ["A", "B"]
|
|
|
|
async def test_skips_empty_title(self):
|
|
"""Rows with empty or missing title should be skipped."""
|
|
table = FakeTable([
|
|
_make_entity_row(""),
|
|
_make_entity_row("A"),
|
|
])
|
|
degree_map = {"A": 1}
|
|
await finalize_entities(table, degree_map)
|
|
|
|
assert len(table.written) == 1
|
|
assert table.written[0]["title"] == "A"
|
|
|
|
async def test_assigns_sequential_human_readable_ids(self):
|
|
"""human_readable_id should be 0-based sequential."""
|
|
table = FakeTable([
|
|
_make_entity_row("A"),
|
|
_make_entity_row("B"),
|
|
_make_entity_row("C"),
|
|
])
|
|
degree_map = {"A": 1, "B": 1, "C": 1}
|
|
await finalize_entities(table, degree_map)
|
|
|
|
ids = [r["human_readable_id"] for r in table.written]
|
|
assert ids == [0, 1, 2]
|
|
|
|
async def test_assigns_unique_ids(self):
|
|
"""Each entity should get a unique UUID id."""
|
|
table = FakeTable([
|
|
_make_entity_row("A"),
|
|
_make_entity_row("B"),
|
|
])
|
|
degree_map = {"A": 1, "B": 1}
|
|
await finalize_entities(table, degree_map)
|
|
|
|
row_ids = [r["id"] for r in table.written]
|
|
assert len(set(row_ids)) == 2
|
|
|
|
async def test_output_columns_match_schema(self):
|
|
"""Written rows should contain exactly the ENTITIES_FINAL_COLUMNS."""
|
|
table = FakeTable([_make_entity_row("A")])
|
|
degree_map = {"A": 1}
|
|
await finalize_entities(table, degree_map)
|
|
|
|
assert set(table.written[0].keys()) == set(ENTITIES_FINAL_COLUMNS)
|
|
|
|
async def test_returns_sample_rows_up_to_five(self):
|
|
"""Should return at most 5 sample rows."""
|
|
rows = [_make_entity_row(f"E{i}") for i in range(8)]
|
|
table = FakeTable(rows)
|
|
degree_map = {f"E{i}": 1 for i in range(8)}
|
|
samples = await finalize_entities(table, degree_map)
|
|
|
|
assert len(samples) == 5
|
|
assert len(table.written) == 8
|
|
|
|
async def test_empty_table(self):
|
|
"""Empty input should produce no output."""
|
|
table = FakeTable([])
|
|
degree_map = {}
|
|
samples = await finalize_entities(table, degree_map)
|
|
|
|
assert samples == []
|
|
assert table.written == []
|
|
|
|
|
|
class TestFinalizeRelationships:
|
|
"""Tests for stream-finalize relationship rows."""
|
|
|
|
async def test_enriches_with_combined_degree(self):
|
|
"""combined_degree should be sum of source and target degrees."""
|
|
table = FakeTable([_make_relationship_row("A", "B")])
|
|
degree_map = {"A": 3, "B": 2}
|
|
await finalize_relationships(table, degree_map)
|
|
|
|
assert len(table.written) == 1
|
|
assert table.written[0]["combined_degree"] == 5
|
|
|
|
async def test_missing_degree_defaults_to_zero(self):
|
|
"""Nodes not in degree map contribute 0 to combined_degree."""
|
|
table = FakeTable([_make_relationship_row("X", "Y")])
|
|
degree_map = {"X": 4}
|
|
await finalize_relationships(table, degree_map)
|
|
|
|
assert table.written[0]["combined_degree"] == 4
|
|
|
|
async def test_deduplicates_by_source_target(self):
|
|
"""Duplicate (source, target) pairs should be skipped."""
|
|
table = FakeTable([
|
|
_make_relationship_row("A", "B"),
|
|
_make_relationship_row("A", "B"),
|
|
_make_relationship_row("B", "C"),
|
|
])
|
|
degree_map = {"A": 1, "B": 2, "C": 1}
|
|
await finalize_relationships(table, degree_map)
|
|
|
|
assert len(table.written) == 2
|
|
pairs = [(r["source"], r["target"]) for r in table.written]
|
|
assert pairs == [("A", "B"), ("B", "C")]
|
|
|
|
async def test_reversed_pair_not_deduplicated(self):
|
|
"""(A,B) and (B,A) are treated as distinct directed edges."""
|
|
table = FakeTable([
|
|
_make_relationship_row("A", "B"),
|
|
_make_relationship_row("B", "A"),
|
|
])
|
|
degree_map = {"A": 1, "B": 1}
|
|
await finalize_relationships(table, degree_map)
|
|
|
|
assert len(table.written) == 2
|
|
|
|
async def test_assigns_sequential_human_readable_ids(self):
|
|
"""human_readable_id should be 0-based sequential."""
|
|
table = FakeTable([
|
|
_make_relationship_row("A", "B"),
|
|
_make_relationship_row("B", "C"),
|
|
])
|
|
degree_map = {"A": 1, "B": 2, "C": 1}
|
|
await finalize_relationships(table, degree_map)
|
|
|
|
ids = [r["human_readable_id"] for r in table.written]
|
|
assert ids == [0, 1]
|
|
|
|
async def test_assigns_unique_ids(self):
|
|
"""Each relationship should get a unique UUID id."""
|
|
table = FakeTable([
|
|
_make_relationship_row("A", "B"),
|
|
_make_relationship_row("B", "C"),
|
|
])
|
|
degree_map = {"A": 1, "B": 2, "C": 1}
|
|
await finalize_relationships(table, degree_map)
|
|
|
|
row_ids = [r["id"] for r in table.written]
|
|
assert len(set(row_ids)) == 2
|
|
|
|
async def test_output_columns_match_schema(self):
|
|
"""Written rows should contain exactly RELATIONSHIPS_FINAL_COLUMNS."""
|
|
table = FakeTable([_make_relationship_row("A", "B")])
|
|
degree_map = {"A": 1, "B": 1}
|
|
await finalize_relationships(table, degree_map)
|
|
|
|
assert set(table.written[0].keys()) == set(RELATIONSHIPS_FINAL_COLUMNS)
|
|
|
|
async def test_returns_sample_rows_up_to_five(self):
|
|
"""Should return at most 5 sample rows."""
|
|
rows = [_make_relationship_row(f"S{i}", f"T{i}") for i in range(8)]
|
|
table = FakeTable(rows)
|
|
degree_map = {f"S{i}": 1 for i in range(8)} | {f"T{i}": 1 for i in range(8)}
|
|
samples = await finalize_relationships(table, degree_map)
|
|
|
|
assert len(samples) == 5
|
|
assert len(table.written) == 8
|
|
|
|
async def test_empty_table(self):
|
|
"""Empty input should produce no output."""
|
|
table = FakeTable([])
|
|
degree_map = {}
|
|
samples = await finalize_relationships(table, degree_map)
|
|
|
|
assert samples == []
|
|
assert table.written == []
|
|
|
|
|
|
class TestFinalizeGraph:
|
|
"""Tests for the orchestrating finalize_graph function."""
|
|
|
|
async def test_produces_entities_and_relationships_keys(self):
|
|
"""Result dict should have 'entities' and 'relationships' keys."""
|
|
entities_table = FakeTable([_make_entity_row("A")])
|
|
relationships_table = FakeTable([_make_relationship_row("A", "B")])
|
|
result = await finalize_graph(entities_table, relationships_table)
|
|
|
|
assert "entities" in result
|
|
assert "relationships" in result
|
|
|
|
async def test_degree_flows_through_to_entities(self):
|
|
"""Entity degree should reflect computed edge degrees."""
|
|
entities_table = FakeTable([
|
|
_make_entity_row("A"),
|
|
_make_entity_row("B"),
|
|
_make_entity_row("C"),
|
|
])
|
|
relationships_table = FakeTable([
|
|
_make_relationship_row("A", "B"),
|
|
_make_relationship_row("A", "C"),
|
|
])
|
|
await finalize_graph(entities_table, relationships_table)
|
|
|
|
degree_by_title = {r["title"]: r["degree"] for r in entities_table.written}
|
|
assert degree_by_title["A"] == 2
|
|
assert degree_by_title["B"] == 1
|
|
assert degree_by_title["C"] == 1
|
|
|
|
async def test_combined_degree_flows_through_to_relationships(self):
|
|
"""Relationship combined_degree should be sum of endpoint degrees."""
|
|
entities_table = FakeTable([
|
|
_make_entity_row("A"),
|
|
_make_entity_row("B"),
|
|
])
|
|
relationships_table = FakeTable([
|
|
_make_relationship_row("A", "B"),
|
|
])
|
|
await finalize_graph(entities_table, relationships_table)
|
|
|
|
assert len(relationships_table.written) == 1
|
|
assert relationships_table.written[0]["combined_degree"] == 2
|
|
|
|
async def test_empty_graph(self):
|
|
"""Empty tables should produce empty results."""
|
|
entities_table = FakeTable([])
|
|
relationships_table = FakeTable([])
|
|
result = await finalize_graph(entities_table, relationships_table)
|
|
|
|
assert result == {"entities": [], "relationships": []}
|
|
|
|
@pytest.mark.parametrize(
|
|
("entity_count", "relationship_count"),
|
|
[
|
|
(3, 2),
|
|
(10, 15),
|
|
],
|
|
ids=["small", "medium"],
|
|
)
|
|
async def test_all_rows_written(self, entity_count: int, relationship_count: int):
|
|
"""All unique entities and relationships should be written."""
|
|
entity_rows = [_make_entity_row(f"E{i}") for i in range(entity_count)]
|
|
relationship_rows = [
|
|
_make_relationship_row(f"E{i}", f"E{(i + 1) % entity_count}")
|
|
for i in range(relationship_count)
|
|
]
|
|
entities_table = FakeTable(entity_rows)
|
|
relationships_table = FakeTable(relationship_rows)
|
|
|
|
await finalize_graph(entities_table, relationships_table)
|
|
|
|
assert len(entities_table.written) == entity_count
|
|
unique_rel_pairs = {(r["source"], r["target"]) for r in relationship_rows}
|
|
assert len(relationships_table.written) == len(unique_rel_pairs)
|