dp-zp-agent/tools/graphrag/tests/test_graph_rebuild_atomic.py

77 lines
2.9 KiB
Python

"""Actual Neo4j rollback after DELETE and partial replacement; isolated scopes."""
import os
import uuid
import pytest
from tools.graphrag import build_graphrag as builder
from tools.graphrag.graph_db import create_neo4j_driver, get_neo4j_settings
pytestmark = pytest.mark.skipif(os.getenv("RUN_NEO4J_TESTS") != "1", reason="Requires local Neo4j.")
@pytest.fixture
def managed_graph(monkeypatch):
scope = "atomic-rebuild-" + uuid.uuid4().hex
monkeypatch.setattr(builder, "GRAPH_SCOPE", scope)
settings = get_neo4j_settings()
with (
create_neo4j_driver(settings) as driver,
driver.session(database=settings.database) as session,
):
session.run(
'CREATE (:Person {id:$id, name:"Original", graph_scope:$scope})',
id=scope + "old",
scope=scope,
).consume()
session.run(
"CREATE (:IndexState {id:$scope, metadata_json:$metadata})",
scope=scope,
metadata='{"version":"old"}',
).consume()
try:
yield session, scope
finally:
session.run(
"MATCH (n) WHERE n.graph_scope=$scope DETACH DELETE n", scope=scope
).consume()
session.run("MATCH (s:IndexState {id:$scope}) DELETE s", scope=scope).consume()
def replacement(scope):
payload = {key: [] for key in [*builder.NODE_QUERIES, *builder.RELATIONSHIP_QUERIES]}
payload["people"] = [{"id": scope + "new", "name": "Replacement", "graph_scope": scope}]
payload["documents"] = [{"path": scope + "/p.md", "graph_scope": scope}]
payload["person_documents"] = [{"person_id": scope + "new", "document_path": scope + "/p.md"}]
return payload
def test_failed_replacement_restores_deleted_graph_and_version(managed_graph):
session, scope = managed_graph
payload = replacement(scope)
payload["documents"].append({"path": None, "graph_scope": scope})
with pytest.raises(Exception):
with session.begin_transaction() as tx:
builder.write_graph_transaction(tx, payload, version={"version": "new"})
tx.commit()
rows = session.run(
"MATCH (n) WHERE n.graph_scope=$scope RETURN n.name AS name", scope=scope
).data()
assert rows == [{"name": "Original"}]
assert builder.graph_version(session) == {"version": "old"}
def test_successful_replacement_commits_graph_and_version_together(managed_graph):
session, scope = managed_graph
payload = replacement(scope)
version = {"version": "new", "git_revision": "new-commit"}
for _ in range(2):
with session.begin_transaction() as tx:
builder.write_graph_transaction(tx, payload, version=version)
tx.commit()
assert builder.graph_version(session) == version
nodes, relationships = builder.graph_counts(session)
assert dict(nodes) == {"Person": 1, "Document": 1}
assert dict(relationships) == {"HAS_DOCUMENT": 1}