Files
medax_pipeline/tests/unit/test_dummy_node_reference_resolution.py
2026-09-08 10:59:05 +02:00

130 lines
5.0 KiB
Python

"""Tests for the dummy-node cross-batch reference mechanism.
The batching design: when a batch references a target that lives in another
batch, a *dummy* target node is created carrying the real target's id. On the
node side dummies are skipped (``node_generator`` drops ``label == 'dummy'``),
but the *edge* is still emitted with the real endpoint id. BioCypher appends
node/edge CSV part files after every batch, and a single ``neo4j-admin
import`` at the end stitches edges to nodes by id. So as long as the real
target lands in *some* batch, the edge resolves at import time.
These tests model that by emitting nodes/edges per batch and accumulating the
result, mimicking the appended CSVs that the final import consumes.
"""
import networkx as nx
import import_fhir_to_nx_diGraph as pipeline
def _emit_batch(graph):
"""Emit one batch the way run_biocypher does: nodes first, then edges."""
nodes = {nid: label for (nid, label, _props) in pipeline.node_generator(graph)}
edges = [
(eid, src, tgt, label)
for (eid, src, tgt, label, _attrs) in pipeline.edge_generator(graph)
]
return nodes, edges
def _accumulate(batches):
"""Union of all batches' CSV part files, as the final import would see."""
all_nodes = {}
all_edges = []
for graph in batches:
nodes, edges = _emit_batch(graph)
all_nodes.update(nodes) # neo4j --skip-duplicate-nodes: union by id
all_edges.extend(edges)
return all_nodes, all_edges
def _batch_with_dummy_target():
"""Batch 1: real Observation O1 referencing P1, which is elsewhere.
``edge_label`` mirrors what ``process_references.py`` always sets on a
real dummy node (the referenced resource type, e.g. "Patient").
"""
g = nx.DiGraph()
g.add_node("obs1", label="observation", unique_id="O1")
g.add_node("dummy_p1", label="dummy", unique_id="P1", edge_label="Patient")
g.add_edge("obs1", "dummy_p1", id="e1")
return g
def _batch_with_real_target():
"""Batch 2: the real Patient node P1 that the dummy stood in for."""
g = nx.DiGraph()
g.add_node("pat1", label="patient", unique_id="P1")
return g
def _unbatched_graph():
"""Everything in one graph: the reference target is the real node."""
g = nx.DiGraph()
g.add_node("obs1", label="observation", unique_id="O1")
g.add_node("pat1", label="patient", unique_id="P1")
g.add_edge("obs1", "pat1", id="e1")
return g
def test_dummy_target_skipped_as_node_but_edge_survives():
nodes, edges = _emit_batch(_batch_with_dummy_target())
assert set(nodes) == {"O1"} # dummy P1 not written as a node
assert len(edges) == 1
_, src, tgt, _ = edges[0]
assert (src, tgt) == ("O1", "P1") # edge keeps the real endpoint id
def test_cross_batch_edge_endpoint_supplied_by_later_batch():
nodes, edges = _accumulate(
[_batch_with_dummy_target(), _batch_with_real_target()]
)
# Real P1 arrives in batch 2; dummy was skipped. Every edge endpoint now
# corresponds to a real node, so the import connects them.
assert set(nodes) == {"O1", "P1"}
endpoints = {s for _, s, _, _ in edges} | {t for _, _, t, _ in edges}
assert endpoints <= set(nodes)
def test_batched_connectivity_matches_unbatched():
batched_nodes, batched_edges = _accumulate(
[_batch_with_dummy_target(), _batch_with_real_target()]
)
full_nodes, full_edges = _emit_batch(_unbatched_graph())
assert set(batched_nodes) == set(full_nodes)
# Connectivity (endpoint pairs) is identical, ignoring relationship type.
assert {(s, t) for _, s, t, _ in batched_edges} == {
(s, t) for _, s, t, _ in full_edges
}
def test_within_batch_edge_has_correct_relationship_label():
# When both endpoints are real in the same batch, the type is correct.
(_, _, _, label, _), = pipeline.edge_generator(_unbatched_graph())
assert label == "Observation_to_Patient"
def test_batched_and_unbatched_produce_identical_edges():
# The full equivalence check: batching must not change the OUTCOME, and
# that includes relationship *types*, not just connectivity. We compare
# the complete edge set (endpoints + relationship label) of a batched run
# against the unbatched run. Since the dummy node carries the real
# target's resource type via `edge_label`, the relationship label is
# derived correctly even while the target is still a dummy, so no
# cross-batch '..._to_Dummy' relationship should ever appear.
_, batched_edges = _accumulate(
[_batch_with_dummy_target(), _batch_with_real_target()]
)
_, full_edges = _emit_batch(_unbatched_graph())
def relationships(edges):
return {(src, tgt, label) for _eid, src, tgt, label in edges}
# No dummy-typed relationship should be produced at all.
assert not any(
label.endswith("_to_Dummy") for _e, _s, _t, label in batched_edges
)
# And the batched edge set must match the unbatched one exactly.
assert relationships(batched_edges) == relationships(full_edges)