f36e2104d8
tests / ragflow_tests_infinity (push) Has been cancelled
tests / ragflow_tests_elasticsearch (push) Has been cancelled
sep-tests / ragflow_preflight (push) Has been cancelled
sep-tests / ragflow_tests_infinity (push) Has been cancelled
sep-tests / ragflow_tests_elasticsearch (push) Has been cancelled
tests / ragflow_preflight (push) Has been cancelled
143 lines
4.4 KiB
Python
143 lines
4.4 KiB
Python
#
|
|
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
"""Regression tests for Extractor._merge_graph_nodes concurrency bug.
|
|
|
|
The historical implementation iterated over ``graph.neighbors(node1)`` directly
|
|
while mutating ``graph`` in the loop body (``add_edge`` / ``remove_node``).
|
|
Under concurrent merges on overlapping neighbourhoods this raised
|
|
``RuntimeError: dictionary keys changed during iteration``.
|
|
|
|
The fix snapshots the neighbour list. These tests pin that behaviour so the
|
|
bug cannot silently regress.
|
|
"""
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
|
|
import networkx as nx
|
|
import pytest
|
|
|
|
from rag.graphrag.general.extractor import Extractor
|
|
from rag.graphrag.utils import GraphChange
|
|
|
|
|
|
def _stub_extractor() -> Extractor:
|
|
llm = SimpleNamespace(llm_name="test-llm", max_length=4096)
|
|
ext = Extractor.__new__(Extractor)
|
|
ext._llm = llm
|
|
ext._language = "English"
|
|
|
|
async def _summary(_name, desc, task_id=""):
|
|
return desc
|
|
|
|
ext._handle_entity_relation_summary = _summary # type: ignore[assignment]
|
|
return ext
|
|
|
|
|
|
def _make_node(graph: nx.Graph, name: str) -> None:
|
|
graph.add_node(
|
|
name,
|
|
description=f"desc-{name}",
|
|
source_id=[name],
|
|
entity_type="person",
|
|
)
|
|
|
|
|
|
def _make_edge(graph: nx.Graph, src: str, tgt: str) -> None:
|
|
graph.add_edge(
|
|
src,
|
|
tgt,
|
|
src_id=src,
|
|
tgt_id=tgt,
|
|
description=f"{src}->{tgt}",
|
|
weight=1.0,
|
|
keywords=[],
|
|
source_id=[src],
|
|
)
|
|
|
|
|
|
@pytest.mark.p1
|
|
@pytest.mark.asyncio
|
|
async def test_merge_graph_nodes_handles_dense_neighbourhood():
|
|
"""A node with many neighbours must merge cleanly without raising."""
|
|
graph = nx.Graph()
|
|
for name in ["A", "B"] + [f"N{i}" for i in range(20)]:
|
|
_make_node(graph, name)
|
|
for i in range(20):
|
|
_make_edge(graph, "A", f"N{i}")
|
|
_make_edge(graph, "B", f"N{i}")
|
|
|
|
ext = _stub_extractor()
|
|
change = GraphChange()
|
|
await ext._merge_graph_nodes(graph, ["A", "B"], change)
|
|
|
|
assert "B" not in graph.nodes
|
|
assert "A" in graph.nodes
|
|
# All 20 N* neighbours should still be connected to the surviving node A
|
|
assert set(graph.neighbors("A")) == {f"N{i}" for i in range(20)}
|
|
|
|
|
|
@pytest.mark.p1
|
|
@pytest.mark.asyncio
|
|
async def test_merge_graph_nodes_neighbours_are_snapshotted():
|
|
"""Regression: iterating graph.neighbors() must not explode if the
|
|
underlying adjacency dict is mutated during the loop."""
|
|
graph = nx.Graph()
|
|
for name in ["A", "B", "C", "D"]:
|
|
_make_node(graph, name)
|
|
# B and C share neighbour D, so merging {A, B} adds edge A-D while
|
|
# the neighbour iterator for B is live.
|
|
_make_edge(graph, "B", "C")
|
|
_make_edge(graph, "B", "D")
|
|
_make_edge(graph, "A", "D")
|
|
|
|
ext = _stub_extractor()
|
|
change = GraphChange()
|
|
await ext._merge_graph_nodes(graph, ["A", "B"], change)
|
|
|
|
assert "B" not in graph.nodes
|
|
assert graph.has_edge("A", "C")
|
|
assert graph.has_edge("A", "D")
|
|
|
|
|
|
@pytest.mark.p1
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_merges_do_not_raise_under_semaphore():
|
|
"""Two concurrent merges on overlapping neighbourhoods must succeed
|
|
when serialized (as entity_resolution now does via Semaphore(1))."""
|
|
graph = nx.Graph()
|
|
for name in ["A1", "A2", "B1", "B2", "X"]:
|
|
_make_node(graph, name)
|
|
_make_edge(graph, "A1", "X")
|
|
_make_edge(graph, "A2", "X")
|
|
_make_edge(graph, "B1", "X")
|
|
_make_edge(graph, "B2", "X")
|
|
|
|
ext = _stub_extractor()
|
|
change = GraphChange()
|
|
sem = asyncio.Semaphore(1)
|
|
|
|
async def merge(nodes):
|
|
async with sem:
|
|
await ext._merge_graph_nodes(graph, nodes, change)
|
|
|
|
await asyncio.gather(merge(["A1", "A2"]), merge(["B1", "B2"]))
|
|
|
|
assert "A2" not in graph.nodes and "B2" not in graph.nodes
|
|
# Both survivors must still share neighbour X
|
|
assert graph.has_edge("A1", "X")
|
|
assert graph.has_edge("B1", "X")
|