77 lines
2.9 KiB
Python
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}
|