From 60b7f299950ac51cb3d43d24f3e643e7b0c6e201 Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Sun, 28 Jun 2026 02:17:07 +0000 Subject: [PATCH 01/12] test: configure pytest discovery in pyproject Add [tool.pytest.ini_options] with testpaths so a bare `pytest` collects the suite; test configuration previously lived only in the tox command. Co-Authored-By: Claude --- pyproject.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 0ae221149..b713692c5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,3 +7,6 @@ build-backend = "setuptools.build_meta" # and pip versions skew and ship a broken mixed install. Listing any unsafe-package resets # pip-tools' defaults, so setuptools and pip are repeated to keep them excluded too. unsafe-package = ["numpy", "setuptools", "pip"] + +[tool.pytest.ini_options] +testpaths = ["pychunkedgraph/tests"] From 336169128ef1bdbdd8b6f8ca5778372ffc7261ab Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Sun, 28 Jun 2026 02:18:04 +0000 Subject: [PATCH 02/12] test: drop unused gen_graph_with_edges fixture The fixture duplicated gen_graph with file-backed edge/component I/O but had no callers. Co-Authored-By: Claude --- pychunkedgraph/tests/conftest.py | 69 -------------------------------- 1 file changed, 69 deletions(-) diff --git a/pychunkedgraph/tests/conftest.py b/pychunkedgraph/tests/conftest.py index 73245f842..9d6dd6fdb 100644 --- a/pychunkedgraph/tests/conftest.py +++ b/pychunkedgraph/tests/conftest.py @@ -204,75 +204,6 @@ def fin(): return partial(_cgraph, request) -@pytest.fixture(scope="function", params=["bigtable", "hbase"]) -def gen_graph_with_edges(request, tmp_path, bigtable_emulator, hbase_emulator): - """Like gen_graph but with real edge/component I/O via local filesystem (file:// protocol).""" - backend = request.param - - def _cgraph(request, n_layers=10, atomic_chunk_bounds: np.ndarray = np.array([])): - edges_dir = f"file://{tmp_path}/edges" - components_dir = f"file://{tmp_path}/components" - - if backend == "bigtable": - backend_client = { - "TYPE": "bigtable", - "CONFIG": { - "ADMIN": True, - "READ_ONLY": False, - "PROJECT": "IGNORE_ENVIRONMENT_PROJECT", - "INSTANCE": "emulated_instance", - "CREDENTIALS": credentials.AnonymousCredentials(), - "MAX_ROW_KEY_COUNT": 1000, - }, - } - else: - backend_client = { - "TYPE": "hbase", - "CONFIG": { - "BASE_URL": f"http://127.0.0.1:{hbase_emulator}", - "MAX_ROW_KEY_COUNT": 1000, - }, - } - - config = { - "data_source": { - "EDGES": edges_dir, - "COMPONENTS": components_dir, - "WATERSHED": "gs://microns-seunglab/minnie65/ws_minnie65_0", - }, - "graph_config": { - "CHUNK_SIZE": [512, 512, 64], - "FANOUT": 2, - "SPATIAL_BITS": 10, - "ID_PREFIX": "", - "ROOT_LOCK_EXPIRY": timedelta(seconds=5), - }, - "backend_client": backend_client, - "ingest_config": {}, - } - - meta, _, client_info, _ = bootstrap("test", config=config) - graph = ChunkedGraph(graph_id="test", meta=meta, client_info=client_info) - # No mock_edges - use real I/O via file:// protocol - graph.meta._ws_cv = CloudVolumeMock() - graph.meta._ws_info_d = mock_ws_info() - graph.meta.ws_ts_scale = lambda mip=0: TensorStoreMock() - graph.meta.layer_count = n_layers - graph.meta.layer_chunk_bounds = get_layer_chunk_bounds( - n_layers, atomic_chunk_bounds=atomic_chunk_bounds - ) - - graph.create() - - def fin(): - _delete_test_table(graph, backend) - - request.addfinalizer(fin) - return graph - - return partial(_cgraph, request) - - @pytest.fixture(scope="function") def gen_graph_simplequerytest(request, gen_graph): """ From 5b16dce0615e78307f84a516e8af76f0ff09ea3a Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Sun, 28 Jun 2026 02:18:57 +0000 Subject: [PATCH 03/12] test: extract fake_timestamp helper for the shared past timestamp Replace the duplicated `datetime.now(UTC) - timedelta(days=10)` with one helper so the "safely old" edit/lineage timestamp has a single source; rename local vars that collided with the helper name. Co-Authored-By: Claude --- .../tests/graph/test_analysis_pathing.py | 19 ++- pychunkedgraph/tests/graph/test_cache.py | 5 +- .../tests/graph/test_chunkedgraph_extended.py | 58 ++++---- .../tests/graph/test_edits_extended.py | 5 +- .../tests/graph/test_graph_build.py | 15 +-- pychunkedgraph/tests/graph/test_history.py | 10 +- pychunkedgraph/tests/graph/test_lineage.py | 12 +- pychunkedgraph/tests/graph/test_locks.py | 50 +++---- pychunkedgraph/tests/graph/test_merge.py | 125 +++++++++--------- pychunkedgraph/tests/graph/test_mincut.py | 39 +++--- pychunkedgraph/tests/graph/test_misc.py | 14 +- pychunkedgraph/tests/graph/test_operation.py | 39 +++--- pychunkedgraph/tests/graph/test_root_lock.py | 11 +- .../tests/graph/test_segmenthistory.py | 20 +-- pychunkedgraph/tests/graph/test_split.py | 123 +++++++++-------- .../tests/graph/test_stale_edges.py | 87 ++++++------ pychunkedgraph/tests/graph/test_subgraph.py | 9 +- .../tests/graph/test_sv_lookup_main.py | 6 +- pychunkedgraph/tests/graph/test_undo_redo.py | 11 +- pychunkedgraph/tests/helpers.py | 6 + .../tests/ingest/test_ingest_cross_edges.py | 21 ++- .../tests/ingest/test_ingest_parent_layer.py | 7 +- pychunkedgraph/tests/meshing/test_setup.py | 4 +- 23 files changed, 344 insertions(+), 352 deletions(-) diff --git a/pychunkedgraph/tests/graph/test_analysis_pathing.py b/pychunkedgraph/tests/graph/test_analysis_pathing.py index 31cf2d30a..d0a7f8a6b 100644 --- a/pychunkedgraph/tests/graph/test_analysis_pathing.py +++ b/pychunkedgraph/tests/graph/test_analysis_pathing.py @@ -1,6 +1,5 @@ """Tests for pychunkedgraph.graph.analysis.pathing""" -from datetime import datetime, timedelta, UTC from math import inf from unittest.mock import MagicMock @@ -15,14 +14,14 @@ compute_rough_coordinate_path, ) -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...ingest.create.parent_layer import add_parent_chunk class TestGetFirstSharedParent: def _build_graph(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, @@ -58,7 +57,7 @@ def test_same_root(self, gen_graph): def test_different_roots_returns_none(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() # Create two disconnected chunks create_chunk( @@ -85,7 +84,7 @@ def test_different_roots_returns_none(self, gen_graph): class TestGetChildrenAtLayer: def test_basic(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, @@ -106,7 +105,7 @@ def test_basic(self, gen_graph): def test_allow_lower_layers(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, @@ -307,7 +306,7 @@ class TestGetChildrenAtLayerEdgeCases: def test_children_at_layer_2_with_multiple_svs(self, gen_graph): """Query children at layer 2 when root has multiple SVs in same chunk.""" graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, @@ -334,7 +333,7 @@ def test_children_at_layer_2_with_multiple_svs(self, gen_graph): def test_children_at_intermediate_layer(self, gen_graph): """Query children at layer 3 from root at layer 4.""" graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, @@ -364,7 +363,7 @@ def test_children_at_intermediate_layer(self, gen_graph): def test_children_allow_lower_layers_with_cross_chunk(self, gen_graph): """Query with allow_lower_layers=True should include layer<=target.""" graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, @@ -386,7 +385,7 @@ def test_children_at_layer_from_l2_node(self, gen_graph): """Querying children at layer 2 from a layer 2 node should return the node itself or its layer-2 children (which is itself).""" graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, diff --git a/pychunkedgraph/tests/graph/test_cache.py b/pychunkedgraph/tests/graph/test_cache.py index e4c68d887..587541b78 100644 --- a/pychunkedgraph/tests/graph/test_cache.py +++ b/pychunkedgraph/tests/graph/test_cache.py @@ -1,13 +1,12 @@ """Tests for pychunkedgraph.graph.cache""" -from datetime import datetime, timedelta, UTC import numpy as np import pytest from pychunkedgraph.graph.cache import CacheService, update -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...ingest.create.parent_layer import add_parent_chunk @@ -29,7 +28,7 @@ def _build_simple_graph(self, gen_graph): from math import inf graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, diff --git a/pychunkedgraph/tests/graph/test_chunkedgraph_extended.py b/pychunkedgraph/tests/graph/test_chunkedgraph_extended.py index 2241608de..5e6d4a06e 100644 --- a/pychunkedgraph/tests/graph/test_chunkedgraph_extended.py +++ b/pychunkedgraph/tests/graph/test_chunkedgraph_extended.py @@ -6,7 +6,7 @@ import numpy as np import pytest -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...ingest.create.parent_layer import add_parent_chunk from ...graph.operation import GraphEditOperation, MergeOperation, SplitOperation from ...graph.exceptions import PreconditionError @@ -16,7 +16,7 @@ class TestChunkedGraphExtended: def _build_graph(self, gen_graph): """Build a simple multi-chunk graph.""" graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() # Chunk A: sv 0, 1 connected create_chunk( @@ -117,7 +117,7 @@ def _build_and_merge(self, gen_graph): """Build a single-chunk graph with two disconnected SVs and merge them.""" atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -214,7 +214,7 @@ def _build_two_sv_graph(self, gen_graph): """Build a 2-layer graph with two disconnected SVs.""" atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -268,7 +268,7 @@ def _build_two_sv_graph(self, gen_graph): """Build a 2-layer graph with two disconnected SVs.""" atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -312,7 +312,7 @@ class TestGetRootsExtended: def _build_cross_chunk(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -408,7 +408,7 @@ class TestGetChildrenExtended: def _build_graph(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -488,7 +488,7 @@ class TestIsLatestRootsExtended: def _build_and_merge(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -538,7 +538,7 @@ class TestGetNodeTimestampsExtended: def _build_graph(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -606,7 +606,7 @@ class TestGetOperationIdsExtended: def test_get_operation_ids_no_ops(self, gen_graph): """get_operation_ids on a node with no operations returns empty dict.""" graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0)], @@ -627,7 +627,7 @@ class TestGetSingleLeafMultipleExtended: def _build_graph(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -699,7 +699,7 @@ class TestGetL2ChildrenExtended: def _build_graph(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -768,7 +768,7 @@ class TestGetChunkLayersExtended: def _build_graph(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0)], @@ -841,7 +841,7 @@ class TestGetAtomicCrossEdgesExtended: def _build_graph(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -890,7 +890,7 @@ def test_get_atomic_cross_edges_multiple_l2(self, gen_graph): def test_get_atomic_cross_edges_no_cross(self, gen_graph): """get_atomic_cross_edges for an L2 node with no cross edges.""" graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0)], @@ -913,7 +913,7 @@ class TestGetAllParentsDictExtended: def _build_graph(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -961,7 +961,7 @@ class TestMiscMethods: def _build_graph(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -1043,7 +1043,7 @@ class TestIsLatestRootsAfterMerge: def _build_and_merge(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -1081,7 +1081,7 @@ class TestGetSubgraphNodesOnly: def _build_graph(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -1137,7 +1137,7 @@ class TestGetSubgraphEdgesOnly: def _build_graph(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -1186,7 +1186,7 @@ def _build_and_merge(self, gen_graph): """Build graph with two disconnected SVs, merge them, return old and new roots.""" atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -1243,7 +1243,7 @@ class TestGetChunkCoordinatesMultiple: def _build_graph(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -1329,7 +1329,7 @@ class TestParentChunkIdMethods: def _build_graph(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -1402,7 +1402,7 @@ class TestReadChunkEdges: def test_read_chunk_edges_returns_dict(self, gen_graph): """read_chunk_edges should return a dict (possibly empty for gs:// edges source).""" graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0)], @@ -1435,7 +1435,7 @@ class TestGetProofreadRootIds: def test_get_proofread_root_ids_no_ops(self, gen_graph): """get_proofread_root_ids with no operations should return empty arrays.""" graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0)], @@ -1454,7 +1454,7 @@ def test_get_proofread_root_ids_after_merge(self, gen_graph): """get_proofread_root_ids after a merge should return the old and new roots.""" atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -1486,7 +1486,7 @@ def _build_connected_graph(self, gen_graph): """Build a 2-layer graph with two connected SVs.""" atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -1552,7 +1552,7 @@ def test_get_earliest_timestamp_after_merge(self, gen_graph): """After creating a graph and performing a merge, get_earliest_timestamp should return a valid datetime.""" atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[ @@ -1576,7 +1576,7 @@ def test_get_earliest_timestamp_after_merge(self, gen_graph): def test_get_earliest_timestamp_stamped_by_root_build(self, gen_graph): """Building the root layer stamps earliest_ts strictly above every root's ts.""" graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0)], diff --git a/pychunkedgraph/tests/graph/test_edits_extended.py b/pychunkedgraph/tests/graph/test_edits_extended.py index 34be9c473..576cb657f 100644 --- a/pychunkedgraph/tests/graph/test_edits_extended.py +++ b/pychunkedgraph/tests/graph/test_edits_extended.py @@ -1,6 +1,5 @@ """Tests for pychunkedgraph.graph.edits - extended coverage""" -from datetime import datetime, timedelta, UTC from math import inf import numpy as np @@ -9,7 +8,7 @@ from pychunkedgraph.graph.edits import flip_ids from pychunkedgraph.graph import basetypes -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...ingest.create.parent_layer import add_parent_chunk @@ -35,7 +34,7 @@ def test_basic(self, gen_graph): from pychunkedgraph.graph.edits import _init_old_hierarchy graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, diff --git a/pychunkedgraph/tests/graph/test_graph_build.py b/pychunkedgraph/tests/graph/test_graph_build.py index e773d1af3..f53cbede4 100644 --- a/pychunkedgraph/tests/graph/test_graph_build.py +++ b/pychunkedgraph/tests/graph/test_graph_build.py @@ -1,10 +1,9 @@ -from datetime import datetime, timedelta, UTC from math import inf import numpy as np import pytest -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...graph import attributes, basetypes, serializers from ...ingest.create.parent_layer import add_parent_chunk @@ -376,12 +375,12 @@ def test_double_chunk_creation(self, gen_graph): cg = gen_graph(n_layers=4, atomic_chunk_bounds=atomic_chunk_bounds) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk B @@ -389,28 +388,28 @@ def test_double_chunk_creation(self, gen_graph): cg, vertices=[to_label(cg, 1, 1, 0, 0, 1)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) add_parent_chunk( cg, 3, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) add_parent_chunk( cg, 3, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) add_parent_chunk( cg, 4, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) diff --git a/pychunkedgraph/tests/graph/test_history.py b/pychunkedgraph/tests/graph/test_history.py index 845271a06..bb2573ce5 100644 --- a/pychunkedgraph/tests/graph/test_history.py +++ b/pychunkedgraph/tests/graph/test_history.py @@ -3,7 +3,7 @@ import numpy as np import pytest -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...graph import ChunkedGraph from ...graph.lineage import lineage_graph, get_root_id_history from ...graph.misc import get_delta_roots @@ -16,25 +16,25 @@ class TestGraphHistory: @pytest.mark.timeout(120) def test_cut_merge_history(self, gen_graph): cg: ChunkedGraph = gen_graph(n_layers=3) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) add_parent_chunk( cg, 3, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) diff --git a/pychunkedgraph/tests/graph/test_lineage.py b/pychunkedgraph/tests/graph/test_lineage.py index 3f2211f6c..4a8dd3a01 100644 --- a/pychunkedgraph/tests/graph/test_lineage.py +++ b/pychunkedgraph/tests/graph/test_lineage.py @@ -18,7 +18,7 @@ ) from pychunkedgraph.graph import attributes -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...ingest.create.parent_layer import add_parent_chunk @@ -28,7 +28,7 @@ def _build_and_merge(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], @@ -154,7 +154,7 @@ def _build_graph_with_two_merges(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() from ..helpers import create_chunk, to_label create_chunk( @@ -225,7 +225,7 @@ def _build_and_merge(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() from ..helpers import create_chunk, to_label create_chunk( @@ -277,7 +277,7 @@ def _build_and_merge(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() from ..helpers import create_chunk, to_label create_chunk( @@ -346,7 +346,7 @@ def _build_graph_with_two_merges(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, diff --git a/pychunkedgraph/tests/graph/test_locks.py b/pychunkedgraph/tests/graph/test_locks.py index 0d82262a9..1fcec5912 100644 --- a/pychunkedgraph/tests/graph/test_locks.py +++ b/pychunkedgraph/tests/graph/test_locks.py @@ -1,7 +1,6 @@ import threading import time from time import sleep -from datetime import datetime, timedelta, UTC import numpy as np import pytest @@ -11,6 +10,7 @@ create_chunk, make_cg_with_row_key_lock_registry, to_label, + fake_timestamp, ) from ...graph import attributes, exceptions from ...graph.locks import ( @@ -42,12 +42,12 @@ def test_lock_unlock(self, gen_graph): cg = gen_graph(n_layers=3) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk B @@ -55,14 +55,14 @@ def test_lock_unlock(self, gen_graph): cg, vertices=[to_label(cg, 1, 1, 0, 0, 1)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) add_parent_chunk( cg, 3, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) @@ -108,12 +108,12 @@ def test_lock_expiration(self, gen_graph): cg = gen_graph(n_layers=3) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk B @@ -121,14 +121,14 @@ def test_lock_expiration(self, gen_graph): cg, vertices=[to_label(cg, 1, 1, 0, 0, 1)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) add_parent_chunk( cg, 3, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) @@ -176,12 +176,12 @@ def test_lock_renew(self, gen_graph): cg = gen_graph(n_layers=3) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk B @@ -189,14 +189,14 @@ def test_lock_renew(self, gen_graph): cg, vertices=[to_label(cg, 1, 1, 0, 0, 1)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) add_parent_chunk( cg, 3, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) @@ -228,12 +228,12 @@ def test_lock_merge_lock_old_id(self, gen_graph): cg = gen_graph(n_layers=3) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk B @@ -241,14 +241,14 @@ def test_lock_merge_lock_old_id(self, gen_graph): cg, vertices=[to_label(cg, 1, 1, 0, 0, 1)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) add_parent_chunk( cg, 3, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) @@ -294,12 +294,12 @@ def test_indefinite_lock(self, gen_graph): cg = gen_graph(n_layers=3) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk B @@ -307,14 +307,14 @@ def test_indefinite_lock(self, gen_graph): cg, vertices=[to_label(cg, 1, 1, 0, 0, 1)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) add_parent_chunk( cg, 3, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) @@ -367,12 +367,12 @@ def test_indefinite_lock_with_normal_lock_expiration(self, gen_graph): cg = gen_graph(n_layers=3) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk B @@ -380,14 +380,14 @@ def test_indefinite_lock_with_normal_lock_expiration(self, gen_graph): cg, vertices=[to_label(cg, 1, 1, 0, 0, 1)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) add_parent_chunk( cg, 3, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) diff --git a/pychunkedgraph/tests/graph/test_merge.py b/pychunkedgraph/tests/graph/test_merge.py index 73925c240..d481e6184 100644 --- a/pychunkedgraph/tests/graph/test_merge.py +++ b/pychunkedgraph/tests/graph/test_merge.py @@ -1,11 +1,10 @@ -from datetime import datetime, timedelta, UTC from math import inf from warnings import warn import numpy as np import pytest -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...graph import ChunkedGraph from ...graph import serializers from ...ingest.create.parent_layer import add_parent_chunk @@ -28,12 +27,12 @@ def test_merge_pair_same_chunk(self, gen_graph): cg = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Merge @@ -68,12 +67,12 @@ def test_merge_pair_neighboring_chunks(self, gen_graph): cg = gen_graph(n_layers=3) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk B @@ -81,14 +80,14 @@ def test_merge_pair_neighboring_chunks(self, gen_graph): cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) add_parent_chunk( cg, 3, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) @@ -124,12 +123,12 @@ def test_merge_pair_disconnected_chunks(self, gen_graph): cg = gen_graph(n_layers=5) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk Z @@ -137,35 +136,35 @@ def test_merge_pair_disconnected_chunks(self, gen_graph): cg, vertices=[to_label(cg, 1, 7, 7, 7, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) add_parent_chunk( cg, 3, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) add_parent_chunk( cg, 3, [3, 3, 3], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) add_parent_chunk( cg, 4, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) add_parent_chunk( cg, 5, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) @@ -207,12 +206,12 @@ def test_merge_pair_already_connected(self, gen_graph): cg = gen_graph(n_layers=2) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) res_old = cg.client.read_all_rows() @@ -250,7 +249,7 @@ def test_merge_triple_chain_to_full_circle_same_chunk(self, gen_graph): cg = gen_graph(n_layers=2) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[ @@ -262,7 +261,7 @@ def test_merge_triple_chain_to_full_circle_same_chunk(self, gen_graph): (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 2), 0.5), (to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2), 0.5), ], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Merge @@ -287,7 +286,7 @@ def test_merge_triple_chain_to_full_circle_neighboring_chunks(self, gen_graph): cg = gen_graph(n_layers=3) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], @@ -295,7 +294,7 @@ def test_merge_triple_chain_to_full_circle_neighboring_chunks(self, gen_graph): (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5), (to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 1, 0, 0, 0), inf), ], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk B @@ -303,14 +302,14 @@ def test_merge_triple_chain_to_full_circle_neighboring_chunks(self, gen_graph): cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), inf)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) add_parent_chunk( cg, 3, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) @@ -336,7 +335,7 @@ def test_merge_triple_chain_to_full_circle_disconnected_chunks(self, gen_graph): cg = gen_graph(n_layers=5) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], @@ -344,7 +343,7 @@ def test_merge_triple_chain_to_full_circle_disconnected_chunks(self, gen_graph): (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5), (to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 7, 7, 7, 0), inf), ], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk B @@ -352,14 +351,14 @@ def test_merge_triple_chain_to_full_circle_disconnected_chunks(self, gen_graph): cg, vertices=[to_label(cg, 1, 7, 7, 7, 0)], edges=[(to_label(cg, 1, 7, 7, 7, 0), to_label(cg, 1, 0, 0, 0, 1), inf)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 3, [3, 3, 3], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 4, [1, 1, 1], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 5, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 3, [3, 3, 3], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 4, [1, 1, 1], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 5, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # Merge new_root_ids = cg.add_edges( @@ -395,12 +394,12 @@ def test_merge_same_node(self, gen_graph): cg = gen_graph(n_layers=2) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) res_old = cg.client.read_all_rows() @@ -428,12 +427,12 @@ def test_merge_pair_abstract_nodes(self, gen_graph): cg = gen_graph(n_layers=3) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk B @@ -441,10 +440,10 @@ def test_merge_pair_abstract_nodes(self, gen_graph): cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) res_old = cg.client.read_all_rows() res_old.consume_all() @@ -542,7 +541,7 @@ def test_cross_edges(self, gen_graph): cg = gen_graph(n_layers=5) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], @@ -550,7 +549,7 @@ def test_cross_edges(self, gen_graph): (to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 1, 0, 0, 0), inf), (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), inf), ], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk B @@ -561,20 +560,20 @@ def test_cross_edges(self, gen_graph): (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), inf), (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 1), inf), ], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk C create_chunk( cg, vertices=[to_label(cg, 1, 2, 0, 0, 0)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 3, [1, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 5, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 3, [1, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 5, [0, 0, 0], time_stamp=fake_ts, n_threads=1) new_roots = cg.add_edges( "Jane Doe", @@ -604,24 +603,24 @@ def test_merge_creates_skip_connection(self, gen_graph): """ cg = gen_graph(n_layers=5) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 7, 7, 7, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 3, [3, 3, 3], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 5, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 3, [3, 3, 3], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 5, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # Before merge: verify both nodes have root at layer 5 root1_pre = cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) @@ -654,24 +653,24 @@ def test_merge_multi_layer_hierarchy_correctness(self, gen_graph): """ cg = gen_graph(n_layers=5) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 7, 7, 7, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 3, [3, 3, 3], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 5, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 3, [3, 3, 3], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 5, [0, 0, 0], time_stamp=fake_ts, n_threads=1) result = cg.add_edges( "Jane Doe", @@ -707,21 +706,21 @@ def test_merge_no_skip_when_siblings_exist(self, gen_graph): """ cg = gen_graph(n_layers=3) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # Merge result = cg.add_edges( diff --git a/pychunkedgraph/tests/graph/test_mincut.py b/pychunkedgraph/tests/graph/test_mincut.py index 8ef6ba239..53d8355db 100644 --- a/pychunkedgraph/tests/graph/test_mincut.py +++ b/pychunkedgraph/tests/graph/test_mincut.py @@ -1,10 +1,9 @@ -from datetime import datetime, timedelta, UTC from math import inf import numpy as np import pytest -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...graph import exceptions from ...ingest.create.parent_layer import add_parent_chunk @@ -26,12 +25,12 @@ def test_cut_regular_link(self, gen_graph): cg = gen_graph(n_layers=3) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk B @@ -39,14 +38,14 @@ def test_cut_regular_link(self, gen_graph): cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) add_parent_chunk( cg, 3, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) @@ -97,12 +96,12 @@ def test_cut_no_link(self, gen_graph): cg = gen_graph(n_layers=3) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk B @@ -110,14 +109,14 @@ def test_cut_no_link(self, gen_graph): cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) add_parent_chunk( cg, 3, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) @@ -158,12 +157,12 @@ def test_cut_old_link(self, gen_graph): cg = gen_graph(n_layers=3) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk B @@ -171,14 +170,14 @@ def test_cut_old_link(self, gen_graph): cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) add_parent_chunk( cg, 3, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) cg.remove_edges( @@ -226,12 +225,12 @@ def test_cut_indivisible_link(self, gen_graph): cg = gen_graph(n_layers=3) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), inf)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Preparation: Build Chunk B @@ -239,14 +238,14 @@ def test_cut_indivisible_link(self, gen_graph): cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), inf)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) add_parent_chunk( cg, 3, [0, 0, 0], - time_stamp=fake_timestamp, + time_stamp=fake_ts, n_threads=1, ) @@ -288,7 +287,7 @@ def test_mincut_disrespects_sources_or_sinks(self, gen_graph): """ cg = gen_graph(n_layers=2) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[ @@ -302,7 +301,7 @@ def test_mincut_disrespects_sources_or_sinks(self, gen_graph): (to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2), 3), (to_label(cg, 1, 0, 0, 0, 2), to_label(cg, 1, 0, 0, 0, 3), 10), ], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Mincut diff --git a/pychunkedgraph/tests/graph/test_misc.py b/pychunkedgraph/tests/graph/test_misc.py index e51502be8..60786a84a 100644 --- a/pychunkedgraph/tests/graph/test_misc.py +++ b/pychunkedgraph/tests/graph/test_misc.py @@ -16,14 +16,14 @@ from pychunkedgraph.graph.edges import Edges from pychunkedgraph.graph.types import Agglomeration -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...ingest.create.parent_layer import add_parent_chunk class TestGetLatestRoots: def test_basic(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, @@ -41,7 +41,7 @@ def test_basic(self, gen_graph): def test_with_timestamp(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, @@ -64,7 +64,7 @@ def test_basic(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], @@ -90,7 +90,7 @@ def test_after_merge(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() sv0 = to_label(graph, 1, 0, 0, 0, 0) sv1 = to_label(graph, 1, 0, 0, 0, 1) @@ -132,7 +132,7 @@ def test_empty_when_no_operations(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0)], @@ -271,7 +271,7 @@ def test_returns_numpy_array_after_merge(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() sv0 = to_label(graph, 1, 0, 0, 0, 0) sv1 = to_label(graph, 1, 0, 0, 0, 1) diff --git a/pychunkedgraph/tests/graph/test_operation.py b/pychunkedgraph/tests/graph/test_operation.py index fa916ae0d..0a271d8d7 100644 --- a/pychunkedgraph/tests/graph/test_operation.py +++ b/pychunkedgraph/tests/graph/test_operation.py @@ -5,13 +5,12 @@ -- all using real graph operations through the BigTable emulator. """ -from datetime import datetime, timedelta, UTC from math import inf import numpy as np import pytest -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...graph import attributes from ...graph.operation import ( GraphEditOperation, @@ -32,7 +31,7 @@ def _build_two_sv_disconnected(gen_graph): """2-layer graph, two disconnected SVs in the same chunk.""" cg = gen_graph(n_layers=2, atomic_chunk_bounds=np.array([1, 1, 1])) - ts = datetime.now(UTC) - timedelta(days=10) + ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], @@ -45,7 +44,7 @@ def _build_two_sv_disconnected(gen_graph): def _build_two_sv_connected(gen_graph): """2-layer graph, two connected SVs in the same chunk.""" cg = gen_graph(n_layers=2, atomic_chunk_bounds=np.array([1, 1, 1])) - ts = datetime.now(UTC) - timedelta(days=10) + ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], @@ -60,7 +59,7 @@ def _build_two_sv_connected(gen_graph): def _build_cross_chunk(gen_graph): """4-layer graph with cross-chunk edges suitable for MulticutOperation.""" cg = gen_graph(n_layers=4) - ts = datetime.now(UTC) - timedelta(days=10) + ts = fake_timestamp() sv0 = to_label(cg, 1, 0, 0, 0, 0) sv1 = to_label(cg, 1, 0, 0, 0, 1) create_chunk( @@ -94,21 +93,21 @@ class TestOperationFromLogRecord: def merged_graph(self, gen_graph): """Build a simple 2-chunk graph and perform a merge, returning (cg, operation_id).""" cg = gen_graph(n_layers=3) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # Split first to get two separate roots split_result = cg.remove_edges( @@ -175,21 +174,21 @@ class TestOperationInversion: def split_and_merge_ops(self, gen_graph): """Build graph, split, merge -- return (cg, merge_op_id, split_op_id).""" cg = gen_graph(n_layers=3) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) split_result = cg.remove_edges( "test_user", @@ -233,21 +232,21 @@ class TestUndoRedoChainResolution: def graph_with_undo(self, gen_graph): """Build graph, perform split, then undo -- return (cg, split_op_id, undo_result).""" cg = gen_graph(n_layers=3) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # Split split_result = cg.remove_edges( @@ -572,7 +571,7 @@ class TestUndoRedoExecute: def _build_connected_cross_chunk(self, gen_graph): """Build a 3-layer graph with between-chunk edge -- suitable for split+undo.""" cg = gen_graph(n_layers=3) - ts = datetime.now(UTC) - timedelta(days=10) + ts = fake_timestamp() sv0 = to_label(cg, 1, 0, 0, 0, 0) sv1 = to_label(cg, 1, 1, 0, 0, 0) create_chunk( @@ -754,7 +753,7 @@ class TestUndoRedoLogRecordTypes: def _build_and_split(self, gen_graph): """Build a cross-chunk graph and split it -- suitable for undo/redo.""" cg = gen_graph(n_layers=3) - ts = datetime.now(UTC) - timedelta(days=10) + ts = fake_timestamp() sv0 = to_label(cg, 1, 0, 0, 0, 0) sv1 = to_label(cg, 1, 1, 0, 0, 0) create_chunk( @@ -867,7 +866,7 @@ class TestUndoEdgeValidation: def _build_connected_cross_chunk(self, gen_graph): """Build a 3-layer graph with between-chunk edge suitable for split+undo.""" cg = gen_graph(n_layers=3) - ts = datetime.now(UTC) - timedelta(days=10) + ts = fake_timestamp() sv0 = to_label(cg, 1, 0, 0, 0, 0) sv1 = to_label(cg, 1, 1, 0, 0, 0) create_chunk( diff --git a/pychunkedgraph/tests/graph/test_root_lock.py b/pychunkedgraph/tests/graph/test_root_lock.py index ea0e5a21d..b96a28eb6 100644 --- a/pychunkedgraph/tests/graph/test_root_lock.py +++ b/pychunkedgraph/tests/graph/test_root_lock.py @@ -3,12 +3,11 @@ Tests lock acquisition, release, and behavior on operation failure. """ -from datetime import datetime, timedelta, UTC import numpy as np import pytest -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...graph import exceptions from ...graph.locks import RootLock from ...ingest.create.parent_layer import add_parent_chunk @@ -19,21 +18,21 @@ class TestRootLock: def simple_graph(self, gen_graph): """Build a 2-chunk graph with a single edge, return (cg, root_id).""" cg = gen_graph(n_layers=3) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) root_id = cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) return cg, root_id diff --git a/pychunkedgraph/tests/graph/test_segmenthistory.py b/pychunkedgraph/tests/graph/test_segmenthistory.py index 2d158b73f..6c2ac5340 100644 --- a/pychunkedgraph/tests/graph/test_segmenthistory.py +++ b/pychunkedgraph/tests/graph/test_segmenthistory.py @@ -14,7 +14,7 @@ from pychunkedgraph.graph import attributes -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...ingest.create.parent_layer import add_parent_chunk @@ -23,7 +23,7 @@ def _build_and_merge(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], @@ -327,7 +327,7 @@ def test_empty_graph(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) # Create a chunk with vertices but perform no edits - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0)], @@ -342,7 +342,7 @@ def test_basic(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], @@ -372,7 +372,7 @@ class TestMergeLog: def _build_and_merge(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], @@ -426,7 +426,7 @@ class TestPastOperationIdsExtended: def _build_and_merge(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], @@ -467,7 +467,7 @@ class TestPastFutureIdMappingExtended: def _build_and_merge(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], @@ -511,7 +511,7 @@ class TestMergeSplitHistory: def _build_merge_and_split(self, gen_graph): atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], @@ -595,7 +595,7 @@ def test_collect_edited_sv_ids_no_edits(self, gen_graph): """collect_edited_sv_ids returns empty array when no edits exist for a root.""" atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0)], @@ -613,7 +613,7 @@ def test_change_log_summary_no_operations(self, gen_graph): """change_log_summary with no operations should show zero splits/merges.""" atomic_chunk_bounds = np.array([1, 1, 1]) graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0)], diff --git a/pychunkedgraph/tests/graph/test_split.py b/pychunkedgraph/tests/graph/test_split.py index 42dc7cee6..c6a33dfb1 100644 --- a/pychunkedgraph/tests/graph/test_split.py +++ b/pychunkedgraph/tests/graph/test_split.py @@ -1,11 +1,10 @@ -from datetime import datetime, timedelta, UTC from math import inf from warnings import warn import numpy as np import pytest -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...graph import ChunkedGraph from ...graph import exceptions from ...graph.misc import get_latest_roots @@ -27,12 +26,12 @@ def test_split_pair_same_chunk(self, gen_graph): cg: ChunkedGraph = gen_graph(n_layers=2) # Preparation: Build Chunk A - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Split @@ -64,11 +63,11 @@ def test_split_pair_same_chunk(self, gen_graph): # verify old state cg.cache = None assert cg.get_root( - to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_timestamp - ) == cg.get_root(to_label(cg, 1, 0, 0, 0, 1), time_stamp=fake_timestamp) + to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_ts + ) == cg.get_root(to_label(cg, 1, 0, 0, 0, 1), time_stamp=fake_ts) leaves = np.unique( cg.get_subgraph( - [cg.get_root(to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_timestamp)], + [cg.get_root(to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_ts)], leaves_only=True, ) ) @@ -77,7 +76,7 @@ def test_split_pair_same_chunk(self, gen_graph): assert to_label(cg, 1, 0, 0, 0, 1) in leaves assert len(get_latest_roots(cg)) == 2 - assert len(get_latest_roots(cg, fake_timestamp)) == 1 + assert len(get_latest_roots(cg, fake_ts)) == 1 def test_split_nonexisting_edge(self, gen_graph): """ @@ -89,7 +88,7 @@ def test_split_nonexisting_edge(self, gen_graph): └─────┘ └─────┘ """ cg = gen_graph(n_layers=2) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], @@ -97,7 +96,7 @@ def test_split_nonexisting_edge(self, gen_graph): (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5), (to_label(cg, 1, 0, 0, 0, 2), to_label(cg, 1, 0, 0, 0, 1), 0.5), ], - timestamp=fake_timestamp, + timestamp=fake_ts, ) new_root_ids = cg.remove_edges( "Jane Doe", @@ -118,20 +117,20 @@ def test_split_pair_neighboring_chunks(self, gen_graph): └─────┴─────┘ └─────┴─────┘ """ cg: ChunkedGraph = gen_graph(n_layers=3) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 1.0)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 1.0)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) new_root_ids = cg.remove_edges( "Jane Doe", source_ids=to_label(cg, 1, 1, 0, 0, 0), @@ -159,11 +158,11 @@ def test_split_pair_neighboring_chunks(self, gen_graph): # verify old state assert cg.get_root( - to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_timestamp - ) == cg.get_root(to_label(cg, 1, 1, 0, 0, 0), time_stamp=fake_timestamp) + to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_ts + ) == cg.get_root(to_label(cg, 1, 1, 0, 0, 0), time_stamp=fake_ts) leaves = np.unique( cg.get_subgraph( - [cg.get_root(to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_timestamp)], + [cg.get_root(to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_ts)], leaves_only=True, ) ) @@ -171,7 +170,7 @@ def test_split_pair_neighboring_chunks(self, gen_graph): assert to_label(cg, 1, 0, 0, 0, 0) in leaves assert to_label(cg, 1, 1, 0, 0, 0) in leaves assert len(get_latest_roots(cg)) == 2 - assert len(get_latest_roots(cg, fake_timestamp)) == 1 + assert len(get_latest_roots(cg, fake_ts)) == 1 @pytest.mark.timeout(30) def test_split_verify_cross_chunk_edges(self, gen_graph): @@ -184,7 +183,7 @@ def test_split_verify_cross_chunk_edges(self, gen_graph): └─────┴─────┴─────┘ └─────┴─────┴─────┘ """ cg: ChunkedGraph = gen_graph(n_layers=4) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 1)], @@ -192,18 +191,18 @@ def test_split_verify_cross_chunk_edges(self, gen_graph): (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 2, 0, 0, 0), inf), (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 1), 0.5), ], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 2, 0, 0, 0)], edges=[(to_label(cg, 1, 2, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), inf)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 3, [1, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 3, [1, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) assert cg.get_root(to_label(cg, 1, 1, 0, 0, 0)) == cg.get_root( to_label(cg, 1, 1, 0, 0, 1) @@ -238,7 +237,7 @@ def test_split_verify_cross_chunk_edges(self, gen_graph): ) assert len(get_latest_roots(cg)) == 2 - assert len(get_latest_roots(cg, fake_timestamp)) == 1 + assert len(get_latest_roots(cg, fake_ts)) == 1 @pytest.mark.timeout(30) def test_split_verify_loop(self, gen_graph): @@ -251,7 +250,7 @@ def test_split_verify_loop(self, gen_graph): └─────┴────────┴─────┘ └─────┴────────┴─────┘ """ cg: ChunkedGraph = gen_graph(n_layers=4) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[ @@ -266,7 +265,7 @@ def test_split_verify_loop(self, gen_graph): (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 2), 0.5), (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 3), 0.5), ], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, @@ -276,12 +275,12 @@ def test_split_verify_loop(self, gen_graph): (to_label(cg, 1, 2, 0, 0, 1), to_label(cg, 1, 1, 0, 0, 1), inf), (to_label(cg, 1, 2, 0, 0, 1), to_label(cg, 1, 2, 0, 0, 0), 0.5), ], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 3, [1, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 3, [1, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) assert cg.get_root(to_label(cg, 1, 1, 0, 0, 0)) == cg.get_root( to_label(cg, 1, 1, 0, 0, 1) @@ -307,7 +306,7 @@ def test_split_verify_loop(self, gen_graph): assert len(new_root_ids) == 2 assert len(get_latest_roots(cg)) == 3 - assert len(get_latest_roots(cg, fake_timestamp)) == 1 + assert len(get_latest_roots(cg, fake_ts)) == 1 @pytest.mark.timeout(30) def test_split_pair_already_disconnected(self, gen_graph): @@ -320,12 +319,12 @@ def test_split_pair_already_disconnected(self, gen_graph): └─────┘ └─────┘ """ cg: ChunkedGraph = gen_graph(n_layers=2) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) res_old = cg.client.read_all_rows() res_old.consume_all() @@ -358,7 +357,7 @@ def test_split_full_circle_to_triple_chain_same_chunk(self, gen_graph): └─────┘ └─────┘ """ cg: ChunkedGraph = gen_graph(n_layers=2) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[ @@ -371,7 +370,7 @@ def test_split_full_circle_to_triple_chain_same_chunk(self, gen_graph): (to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2), 0.5), (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.3), ], - timestamp=fake_timestamp, + timestamp=fake_ts, ) new_root_ids = cg.remove_edges( "Jane Doe", @@ -390,11 +389,11 @@ def test_split_full_circle_to_triple_chain_same_chunk(self, gen_graph): # verify old state old_root_id = cg.get_root( - to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_timestamp + to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_ts ) assert new_root_ids[0] != old_root_id assert len(get_latest_roots(cg)) == 1 - assert len(get_latest_roots(cg, fake_timestamp)) == 1 + assert len(get_latest_roots(cg, fake_ts)) == 1 @pytest.mark.timeout(30) def test_split_full_circle_to_triple_chain_neighboring_chunks(self, gen_graph): @@ -407,7 +406,7 @@ def test_split_full_circle_to_triple_chain_neighboring_chunks(self, gen_graph): └─────┴─────┘ └─────┴─────┘ """ cg: ChunkedGraph = gen_graph(n_layers=3) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], @@ -416,7 +415,7 @@ def test_split_full_circle_to_triple_chain_neighboring_chunks(self, gen_graph): (to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 1, 0, 0, 0), 0.5), (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.3), ], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, @@ -425,9 +424,9 @@ def test_split_full_circle_to_triple_chain_neighboring_chunks(self, gen_graph): (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5), (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.3), ], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) new_root_ids = cg.remove_edges( "Jane Doe", @@ -446,11 +445,11 @@ def test_split_full_circle_to_triple_chain_neighboring_chunks(self, gen_graph): # verify old state old_root_id = cg.get_root( - to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_timestamp + to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_ts ) assert new_root_ids[0] != old_root_id assert len(get_latest_roots(cg)) == 1 - assert len(get_latest_roots(cg, fake_timestamp)) == 1 + assert len(get_latest_roots(cg, fake_ts)) == 1 @pytest.mark.timeout(30) def test_split_same_node(self, gen_graph): @@ -463,12 +462,12 @@ def test_split_same_node(self, gen_graph): └─────┘ """ cg: ChunkedGraph = gen_graph(n_layers=2) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) res_old = cg.client.read_all_rows() @@ -493,21 +492,21 @@ def test_split_pair_abstract_nodes(self, gen_graph): """ cg: ChunkedGraph = gen_graph(n_layers=3) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) res_old = cg.client.read_all_rows() res_old.consume_all() with pytest.raises((exceptions.PreconditionError, AssertionError)): @@ -601,21 +600,21 @@ def test_split_multi_layer_hierarchy_correctness(self, gen_graph): └─────┴─────┘ """ cg = gen_graph(n_layers=4) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) result = cg.remove_edges( "Jane Doe", @@ -650,12 +649,12 @@ def test_split_creates_isolated_components_with_skip_connections(self, gen_graph └─────┴─────┴─────┘ """ cg = gen_graph(n_layers=4) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, @@ -664,17 +663,17 @@ def test_split_creates_isolated_components_with_skip_connections(self, gen_graph (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5), (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 2, 0, 0, 0), inf), ], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 2, 0, 0, 0)], edges=[(to_label(cg, 1, 2, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), inf)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 3, [1, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 3, [1, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # All three should share a root before split root_pre = cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) diff --git a/pychunkedgraph/tests/graph/test_stale_edges.py b/pychunkedgraph/tests/graph/test_stale_edges.py index 35a6a3fa7..9ac1aa08b 100644 --- a/pychunkedgraph/tests/graph/test_stale_edges.py +++ b/pychunkedgraph/tests/graph/test_stale_edges.py @@ -4,12 +4,11 @@ operations through the BigTable emulator. """ -from datetime import datetime, timedelta, UTC import numpy as np import pytest -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...graph.edges.stale import get_stale_nodes, get_new_nodes from ...ingest.create.parent_layer import add_parent_chunk @@ -28,21 +27,21 @@ def test_stale_nodes_detected_after_split(self, gen_graph): └─────┴─────┘ """ cg = gen_graph(n_layers=3) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # Get old parents before edit old_root = cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) @@ -71,21 +70,21 @@ def test_no_stale_nodes_for_current_ids(self, gen_graph): └─────┴─────┘ """ cg = gen_graph(n_layers=3) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # Split cg.remove_edges( @@ -115,21 +114,21 @@ def test_get_new_nodes_resolves_to_correct_layer(self, gen_graph): └─────┴─────┘ """ cg = gen_graph(n_layers=3) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # Get L2 parent of SV 1 before edit sv1 = to_label(cg, 1, 0, 0, 0, 0) @@ -160,31 +159,31 @@ def test_no_stale_nodes_in_unaffected_region(self, gen_graph): └─────┴─────┴─────┘ """ cg = gen_graph(n_layers=4) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Chunk C - isolated node, not connected to A or B create_chunk( cg, vertices=[to_label(cg, 1, 2, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 3, [1, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 3, [1, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # Get the isolated node's root before edit isolated_root = cg.get_root(to_label(cg, 1, 2, 0, 0, 0)) @@ -214,16 +213,16 @@ def test_get_new_nodes_returns_self_for_non_stale(self, gen_graph): └─────┘ """ cg = gen_graph(n_layers=4) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) sv = to_label(cg, 1, 0, 0, 0, 0) l2_parent = cg.get_parent(sv) @@ -244,21 +243,21 @@ def test_get_stale_nodes_empty_for_fresh_graph(self, gen_graph): └─────┴─────┘ """ cg = gen_graph(n_layers=3) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) root = cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) l2_0 = cg.get_parent(to_label(cg, 1, 0, 0, 0, 0)) @@ -281,21 +280,21 @@ def test_get_new_nodes_multiple_svs(self, gen_graph): └─────┴─────┘ """ cg = gen_graph(n_layers=3) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) sv1 = to_label(cg, 1, 0, 0, 0, 0) sv2 = to_label(cg, 1, 1, 0, 0, 0) @@ -320,15 +319,15 @@ def test_get_new_nodes_with_duplicate_svs(self, gen_graph): └─────┘ """ cg = gen_graph(n_layers=3) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) sv = to_label(cg, 1, 0, 0, 0, 0) svs = np.array([sv, sv, sv], dtype=np.uint64) @@ -352,7 +351,7 @@ def test_get_stale_nodes_with_l2_ids_after_merge(self, gen_graph): """ atomic_chunk_bounds = np.array([1, 1, 1]) cg = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() sv0 = to_label(cg, 1, 0, 0, 0, 0) sv1 = to_label(cg, 1, 0, 0, 0, 1) @@ -361,7 +360,7 @@ def test_get_stale_nodes_with_l2_ids_after_merge(self, gen_graph): cg, vertices=[sv0, sv1], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) # Get L2 parents before merge (each SV has its own L2 parent) @@ -393,14 +392,14 @@ def test_get_stale_nodes_returns_numpy_array(self, gen_graph): """ atomic_chunk_bounds = np.array([1, 1, 1]) cg = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() sv0 = to_label(cg, 1, 0, 0, 0, 0) create_chunk( cg, vertices=[sv0], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) root = cg.get_root(sv0) @@ -419,17 +418,17 @@ def test_get_new_nodes_at_root_layer(self, gen_graph): └─────┘ """ cg = gen_graph(n_layers=4) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() sv = to_label(cg, 1, 0, 0, 0, 0) create_chunk( cg, vertices=[sv], edges=[], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) + add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) root = cg.get_root(sv) root_layer = cg.get_chunk_layer(root) diff --git a/pychunkedgraph/tests/graph/test_subgraph.py b/pychunkedgraph/tests/graph/test_subgraph.py index 9ed062b35..1725a5813 100644 --- a/pychunkedgraph/tests/graph/test_subgraph.py +++ b/pychunkedgraph/tests/graph/test_subgraph.py @@ -1,6 +1,5 @@ """Tests for pychunkedgraph.graph.subgraph""" -from datetime import datetime, timedelta, UTC from math import inf import numpy as np @@ -8,14 +7,14 @@ from pychunkedgraph.graph.subgraph import SubgraphProgress, get_subgraph_nodes -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...ingest.create.parent_layer import add_parent_chunk class TestSubgraphProgress: def test_init(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0)], @@ -36,7 +35,7 @@ def test_init(self, gen_graph): def test_serializable_keys(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0)], @@ -61,7 +60,7 @@ def test_serializable_keys(self, gen_graph): class TestGetSubgraphNodes: def _build_graph(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, diff --git a/pychunkedgraph/tests/graph/test_sv_lookup_main.py b/pychunkedgraph/tests/graph/test_sv_lookup_main.py index c536d01fa..e8be1d5af 100644 --- a/pychunkedgraph/tests/graph/test_sv_lookup_main.py +++ b/pychunkedgraph/tests/graph/test_sv_lookup_main.py @@ -17,7 +17,7 @@ from pychunkedgraph.graph.sv_lookup import resolve_supervoxels_at_coords from pychunkedgraph.graph.sv_lookup import main as sv_lookup_main -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...ingest.create.parent_layer import add_parent_chunk UTC = timezone.utc @@ -62,7 +62,7 @@ def _build_two_sv_graph(gen_graph): with these SVs at coordinates (0,0,0), (1,0,0), (2,0,0) respectively. """ graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() sv0 = to_label(graph, 1, 0, 0, 0, 0) sv1 = to_label(graph, 1, 0, 0, 0, 1) @@ -107,7 +107,7 @@ def _build_two_root_graph(gen_graph): and sv1 under root_b. graph.meta._ws_cv is seeded with the three SVs. """ graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() sv0 = to_label(graph, 1, 0, 0, 0, 0) sv1 = to_label(graph, 1, 0, 0, 0, 1) diff --git a/pychunkedgraph/tests/graph/test_undo_redo.py b/pychunkedgraph/tests/graph/test_undo_redo.py index 63708d3fa..3858312f0 100644 --- a/pychunkedgraph/tests/graph/test_undo_redo.py +++ b/pychunkedgraph/tests/graph/test_undo_redo.py @@ -4,12 +4,11 @@ operations through the BigTable emulator. """ -from datetime import datetime, timedelta, UTC import numpy as np import pytest -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...ingest.create.parent_layer import add_parent_chunk @@ -25,21 +24,21 @@ def two_chunk_graph(self, gen_graph): └─────┴─────┘ """ cg = gen_graph(n_layers=3) - fake_timestamp = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) create_chunk( cg, vertices=[to_label(cg, 1, 1, 0, 0, 0)], edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_timestamp, + timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_timestamp, n_threads=1) + add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) return cg @pytest.mark.timeout(30) diff --git a/pychunkedgraph/tests/helpers.py b/pychunkedgraph/tests/helpers.py index 50ee0a737..449058729 100644 --- a/pychunkedgraph/tests/helpers.py +++ b/pychunkedgraph/tests/helpers.py @@ -1,4 +1,5 @@ import threading +from datetime import datetime, timedelta, UTC from functools import reduce from unittest.mock import MagicMock @@ -10,6 +11,11 @@ from ..ingest.create.atomic_layer import add_atomic_chunk +def fake_timestamp(): + """A timestamp safely in the past, for edits/lineage tests needing an old parent_ts.""" + return datetime.now(UTC) - timedelta(days=10) + + class CloudVolumeBounds(object): def __init__(self, bounds=[[0, 0, 0], [0, 0, 0]]): self._bounds = np.array(bounds) diff --git a/pychunkedgraph/tests/ingest/test_ingest_cross_edges.py b/pychunkedgraph/tests/ingest/test_ingest_cross_edges.py index d65b767f4..fef5c3a3f 100644 --- a/pychunkedgraph/tests/ingest/test_ingest_cross_edges.py +++ b/pychunkedgraph/tests/ingest/test_ingest_cross_edges.py @@ -1,6 +1,5 @@ """Tests for pychunkedgraph.ingest.create.cross_edges""" -from datetime import datetime, timedelta, UTC from math import inf import numpy as np @@ -14,7 +13,7 @@ ) from pychunkedgraph.graph import basetypes -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...ingest.create.parent_layer import add_parent_chunk @@ -69,7 +68,7 @@ class TestGetChildrenChunkCrossEdges: def test_no_cross_edges(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, @@ -84,7 +83,7 @@ def test_no_cross_edges(self, gen_graph): def test_with_cross_edges(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, @@ -111,7 +110,7 @@ def test_no_atomic_chunks_returns_empty(self, gen_graph): """When the chunk coordinate is out of bounds, get_touching_atomic_chunks returns empty and the function returns early with an empty list.""" cg = gen_graph(n_layers=3, atomic_chunk_bounds=np.array([1, 1, 1])) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, @@ -133,7 +132,7 @@ def test_basic_cross_edges(self, gen_graph): """A 4-layer graph with cross-chunk connected SVs returns cross edges when called with use_threads=False.""" cg = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() # Chunk A (0,0,0): sv 0 connected cross-chunk to chunk B create_chunk( @@ -175,7 +174,7 @@ class TestGetChildrenChunkCrossEdgesAdditional: def test_multiple_cross_edges(self, gen_graph): """Multiple SVs with cross-chunk edges should all be found.""" cg = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() # Chunk A: two SVs, each cross-chunk connected create_chunk( @@ -213,7 +212,7 @@ def test_cross_edges_layer4(self, gen_graph): L4 [0,0,0] has L3 children [0,0,0] (x=0,1) and [1,0,0] (x=2,3). Touching face is at L2 x=1 and x=2.""" cg = gen_graph(n_layers=5) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() # SV at L1 [1,0,0] - on the right boundary of L3 [0,0,0] create_chunk( @@ -256,7 +255,7 @@ def test_no_threads_with_cross_edges(self, gen_graph): chunks of L3 [0,0,0], which includes L2 at x=0 with AtomicCrossChunkEdge[3]. """ cg = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() # SV at L1 [0,0,0] with cross edge to [2,0,0] (layer-3 cross edge) create_chunk( @@ -291,7 +290,7 @@ def test_no_threads_with_cross_edges(self, gen_graph): def test_no_threads_empty_chunk(self, gen_graph): """use_threads=False with out-of-bounds chunk should return empty dict.""" cg = gen_graph(n_layers=3, atomic_chunk_bounds=np.array([1, 1, 1])) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, @@ -312,7 +311,7 @@ def test_no_threads_empty_chunk(self, gen_graph): def test_no_cross_edges_returns_empty(self, gen_graph): """When chunks have no cross edges at the relevant layers, result is empty.""" cg = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( cg, diff --git a/pychunkedgraph/tests/ingest/test_ingest_parent_layer.py b/pychunkedgraph/tests/ingest/test_ingest_parent_layer.py index cdcd2ce5a..290f27250 100644 --- a/pychunkedgraph/tests/ingest/test_ingest_parent_layer.py +++ b/pychunkedgraph/tests/ingest/test_ingest_parent_layer.py @@ -1,19 +1,18 @@ """Tests for pychunkedgraph.ingest.create.parent_layer""" -from datetime import datetime, timedelta, UTC from math import inf import numpy as np import pytest -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp from ...ingest.create.parent_layer import add_parent_chunk class TestAddParentChunk: def test_single_thread(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, @@ -35,7 +34,7 @@ def test_single_thread(self, gen_graph): def test_multi_chunk(self, gen_graph): graph = gen_graph(n_layers=4) - fake_ts = datetime.now(UTC) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, diff --git a/pychunkedgraph/tests/meshing/test_setup.py b/pychunkedgraph/tests/meshing/test_setup.py index d089d1ce0..b31d54517 100644 --- a/pychunkedgraph/tests/meshing/test_setup.py +++ b/pychunkedgraph/tests/meshing/test_setup.py @@ -2,7 +2,7 @@ from ...ingest.create.parent_layer import add_parent_chunk from ...meshing.setup import derive_initial_ts -from ..helpers import create_chunk, to_label +from ..helpers import create_chunk, to_label, fake_timestamp def test_derive_initial_ts_uses_stamped_boundary(gen_graph): @@ -16,7 +16,7 @@ def test_derive_initial_ts_uses_stamped_boundary(gen_graph): def test_derive_initial_ts_after_root_build(gen_graph): """Root-layer build stamps earliest_ts; derive consumes it (fails if it stays unset).""" graph = gen_graph(n_layers=4) - fake_ts = datetime.now(timezone.utc) - timedelta(days=10) + fake_ts = fake_timestamp() create_chunk( graph, vertices=[to_label(graph, 1, 0, 0, 0, 0)], From ac0060dc3e3bc2ae194e630c4643ce1f11a397ee Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Sun, 28 Jun 2026 17:26:23 +0000 Subject: [PATCH 04/12] test: add build_graph helper deriving the parent hierarchy A test graph is determined by its atomic chunks; build_graph takes only that topology and builds the full (derivable) parent hierarchy at one timestamp, replacing hand-listed add_parent_chunk scaffolding. Co-Authored-By: Claude --- pychunkedgraph/tests/graph/test_cache.py | 36 +++++--------- pychunkedgraph/tests/graph/test_subgraph.py | 55 +++++---------------- pychunkedgraph/tests/helpers.py | 24 +++++++++ 3 files changed, 47 insertions(+), 68 deletions(-) diff --git a/pychunkedgraph/tests/graph/test_cache.py b/pychunkedgraph/tests/graph/test_cache.py index 587541b78..f0cb2ce34 100644 --- a/pychunkedgraph/tests/graph/test_cache.py +++ b/pychunkedgraph/tests/graph/test_cache.py @@ -1,13 +1,13 @@ """Tests for pychunkedgraph.graph.cache""" +from math import inf import numpy as np import pytest from pychunkedgraph.graph.cache import CacheService, update -from ..helpers import create_chunk, to_label, fake_timestamp -from ...ingest.create.parent_layer import add_parent_chunk +from ..helpers import to_label, build_graph class TestUpdate: @@ -25,31 +25,17 @@ def test_many_to_one(self): class TestCacheService: def _build_simple_graph(self, gen_graph): """Build a simple 2-chunk graph with 2 SVs per chunk.""" - from math import inf - - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - timestamp=fake_ts, - ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), + cg, _ = build_graph( + gen_graph, + chunks=[ + ( + [(0, 0, 0, 0), (0, 0, 0, 1)], + [((0, 0, 0, 0), (0, 0, 0, 1), 0.5), ((0, 0, 0, 0), (1, 0, 0, 0), inf)], + ), + ([(1, 0, 0, 0)], [((1, 0, 0, 0), (0, 0, 0, 0), inf)]), ], - timestamp=fake_ts, ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - return graph + return cg def test_len(self, gen_graph): graph = self._build_simple_graph(gen_graph) diff --git a/pychunkedgraph/tests/graph/test_subgraph.py b/pychunkedgraph/tests/graph/test_subgraph.py index 1725a5813..b405e0fb0 100644 --- a/pychunkedgraph/tests/graph/test_subgraph.py +++ b/pychunkedgraph/tests/graph/test_subgraph.py @@ -7,22 +7,12 @@ from pychunkedgraph.graph.subgraph import SubgraphProgress, get_subgraph_nodes -from ..helpers import create_chunk, to_label, fake_timestamp -from ...ingest.create.parent_layer import add_parent_chunk +from ..helpers import to_label, build_graph class TestSubgraphProgress: def test_init(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, - ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) + graph, _ = build_graph(gen_graph, chunks=[([(0, 0, 0, 0)], [])]) root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) progress = SubgraphProgress( @@ -34,16 +24,7 @@ def test_init(self, gen_graph): assert not progress.done_processing() def test_serializable_keys(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, - ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) + graph, _ = build_graph(gen_graph, chunks=[([(0, 0, 0, 0)], [])]) root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) progress = SubgraphProgress( @@ -59,29 +40,17 @@ def test_serializable_keys(self, gen_graph): class TestGetSubgraphNodes: def _build_graph(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - timestamp=fake_ts, - ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), + cg, _ = build_graph( + gen_graph, + chunks=[ + ( + [(0, 0, 0, 0), (0, 0, 0, 1)], + [((0, 0, 0, 0), (0, 0, 0, 1), 0.5), ((0, 0, 0, 0), (1, 0, 0, 0), inf)], + ), + ([(1, 0, 0, 0)], [((1, 0, 0, 0), (0, 0, 0, 0), inf)]), ], - timestamp=fake_ts, ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - return graph + return cg def test_single_node(self, gen_graph): graph = self._build_graph(gen_graph) diff --git a/pychunkedgraph/tests/helpers.py b/pychunkedgraph/tests/helpers.py index 449058729..512f719c8 100644 --- a/pychunkedgraph/tests/helpers.py +++ b/pychunkedgraph/tests/helpers.py @@ -9,6 +9,7 @@ from ..graph.edges import EDGE_TYPES from ..graph import basetypes from ..ingest.create.atomic_layer import add_atomic_chunk +from ..ingest.create.parent_layer import add_parent_chunk def fake_timestamp(): @@ -162,6 +163,29 @@ def get_layer_chunk_bounds( return layer_bounds_d +def build_graph(gen_graph, n_layers=4, chunks=(), *, timestamp=None, atomic_chunk_bounds=None): + """Build a test graph: atomic chunks plus the full parent hierarchy they imply, at one ts. + + A vertex is an (x, y, z, seg) atomic coordinate; an edge is ((x,y,z,seg), (x,y,z,seg), + affinity); chunks is a list of (vertices, edges). Parents are derived. Returns (cg, ts). + """ + bounds = np.array([]) if atomic_chunk_bounds is None else atomic_chunk_bounds + cg = gen_graph(n_layers=n_layers, atomic_chunk_bounds=bounds) + ts = fake_timestamp() if timestamp is None else timestamp + atomic_coords = set() + for vertices, edges in chunks: + verts = [to_label(cg, 1, *v) for v in vertices] + labeled = [(to_label(cg, 1, *a), to_label(cg, 1, *b), aff) for a, b, aff in edges] + create_chunk(cg, vertices=verts, edges=labeled, timestamp=ts) + atomic_coords.update(v[:3] for v in vertices) + fanout = cg.meta.graph_config.FANOUT + for layer in range(3, n_layers + 1): + coords = {tuple(np.array(c) // fanout ** (layer - 2)) for c in atomic_coords} + for coord in sorted(coords): + add_parent_chunk(cg, layer, list(coord), time_stamp=ts, n_threads=1) + return cg, ts + + class RowKeyLockRegistry: """Thread-safe in-memory stand-in for kvdbclient's row-key lock API. From 9eb24de935e43627b9d671085c02b45d01e3bcab Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Sun, 28 Jun 2026 18:33:24 +0000 Subject: [PATCH 05/12] test: build_graph spec uses named SV coordinates Supervoxels are named and given as SV(x, y, z, seg) with zeros defaulting away, so build specs carry no repeated raw coordinate tuples; build_graph returns a BuiltGraph(cg, sv, ts) namedtuple and callers reference sv by name. --- pychunkedgraph/tests/graph/test_cache.py | 85 +++++++++------------ pychunkedgraph/tests/graph/test_subgraph.py | 54 ++++++------- pychunkedgraph/tests/helpers.py | 37 +++++---- 3 files changed, 81 insertions(+), 95 deletions(-) diff --git a/pychunkedgraph/tests/graph/test_cache.py b/pychunkedgraph/tests/graph/test_cache.py index f0cb2ce34..1a964ac09 100644 --- a/pychunkedgraph/tests/graph/test_cache.py +++ b/pychunkedgraph/tests/graph/test_cache.py @@ -3,11 +3,10 @@ from math import inf import numpy as np -import pytest from pychunkedgraph.graph.cache import CacheService, update -from ..helpers import to_label, build_graph +from ..helpers import SV, build_graph class TestUpdate: @@ -24,38 +23,31 @@ def test_many_to_one(self): class TestCacheService: def _build_simple_graph(self, gen_graph): - """Build a simple 2-chunk graph with 2 SVs per chunk.""" - cg, _ = build_graph( + """Build a simple 2-chunk graph with 2 SVs in the first chunk.""" + return build_graph( gen_graph, - chunks=[ - ( - [(0, 0, 0, 0), (0, 0, 0, 1)], - [((0, 0, 0, 0), (0, 0, 0, 1), 0.5), ((0, 0, 0, 0), (1, 0, 0, 0), inf)], - ), - ([(1, 0, 0, 0)], [((1, 0, 0, 0), (0, 0, 0, 0), inf)]), - ], + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, + edges=[("a0", "a1", 0.5), ("a0", "b", inf)], ) - return cg def test_len(self, gen_graph): - graph = self._build_simple_graph(gen_graph) - cache = CacheService(graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - cache.parent(sv) + g = self._build_simple_graph(gen_graph) + cache = CacheService(g.cg) + cache.parent(g.sv["a0"]) assert len(cache) >= 1 def test_clear(self, gen_graph): - graph = self._build_simple_graph(gen_graph) - cache = CacheService(graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - cache.parent(sv) + g = self._build_simple_graph(gen_graph) + cache = CacheService(g.cg) + cache.parent(g.sv["a0"]) cache.clear() assert len(cache) == 0 def test_parent_miss_then_hit(self, gen_graph): - graph = self._build_simple_graph(gen_graph) - cache = CacheService(graph) - sv = to_label(graph, 1, 0, 0, 0, 0) + g = self._build_simple_graph(gen_graph) + cache = CacheService(g.cg) + sv = g.sv["a0"] # First call is a miss parent1 = cache.parent(sv) @@ -67,9 +59,9 @@ def test_parent_miss_then_hit(self, gen_graph): assert parent1 == parent2 def test_children_backfills_parent(self, gen_graph): - graph = self._build_simple_graph(gen_graph) - cache = CacheService(graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) + g = self._build_simple_graph(gen_graph) + cache = CacheService(g.cg) + root = g.cg.get_root(g.sv["a0"]) children = cache.children(root) assert len(children) > 0 # Children should be backfilled as parents @@ -77,9 +69,9 @@ def test_children_backfills_parent(self, gen_graph): assert child in cache.parents_cache def test_get_stats(self, gen_graph): - graph = self._build_simple_graph(gen_graph) - cache = CacheService(graph) - sv = to_label(graph, 1, 0, 0, 0, 0) + g = self._build_simple_graph(gen_graph) + cache = CacheService(g.cg) + sv = g.sv["a0"] cache.parent(sv) cache.parent(sv) stats = cache.get_stats() @@ -88,44 +80,37 @@ def test_get_stats(self, gen_graph): assert "hit_rate" in stats["parents"] def test_reset_stats(self, gen_graph): - graph = self._build_simple_graph(gen_graph) - cache = CacheService(graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - cache.parent(sv) + g = self._build_simple_graph(gen_graph) + cache = CacheService(g.cg) + cache.parent(g.sv["a0"]) cache.reset_stats() assert cache.stats["parents"]["hits"] == 0 assert cache.stats["parents"]["misses"] == 0 def test_parents_multiple_empty(self, gen_graph): - graph = self._build_simple_graph(gen_graph) - cache = CacheService(graph) + g = self._build_simple_graph(gen_graph) + cache = CacheService(g.cg) result = cache.parents_multiple(np.array([], dtype=np.uint64)) assert len(result) == 0 def test_parents_multiple(self, gen_graph): - graph = self._build_simple_graph(gen_graph) - cache = CacheService(graph) - svs = np.array( - [ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - dtype=np.uint64, - ) + g = self._build_simple_graph(gen_graph) + cache = CacheService(g.cg) + svs = np.array([g.sv["a0"], g.sv["a1"]], dtype=np.uint64) result = cache.parents_multiple(svs) assert len(result) == 2 def test_children_multiple(self, gen_graph): - graph = self._build_simple_graph(gen_graph) - cache = CacheService(graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) + g = self._build_simple_graph(gen_graph) + cache = CacheService(g.cg) + root = g.cg.get_root(g.sv["a0"]) result = cache.children_multiple(np.array([root], dtype=np.uint64)) assert root in result def test_children_multiple_flatten(self, gen_graph): - graph = self._build_simple_graph(gen_graph) - cache = CacheService(graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) + g = self._build_simple_graph(gen_graph) + cache = CacheService(g.cg) + root = g.cg.get_root(g.sv["a0"]) result = cache.children_multiple( np.array([root], dtype=np.uint64), flatten=True ) diff --git a/pychunkedgraph/tests/graph/test_subgraph.py b/pychunkedgraph/tests/graph/test_subgraph.py index b405e0fb0..ca897518d 100644 --- a/pychunkedgraph/tests/graph/test_subgraph.py +++ b/pychunkedgraph/tests/graph/test_subgraph.py @@ -3,20 +3,18 @@ from math import inf import numpy as np -import pytest from pychunkedgraph.graph.subgraph import SubgraphProgress, get_subgraph_nodes -from ..helpers import to_label, build_graph +from ..helpers import SV, build_graph class TestSubgraphProgress: def test_init(self, gen_graph): - graph, _ = build_graph(gen_graph, chunks=[([(0, 0, 0, 0)], [])]) - - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) + g = build_graph(gen_graph, n_layers=4, supervoxels={"a": SV()}) + root = g.cg.get_root(g.sv["a"]) progress = SubgraphProgress( - graph.meta, + g.cg.meta, node_ids=[root], return_layers=[2], serializable=False, @@ -24,11 +22,10 @@ def test_init(self, gen_graph): assert not progress.done_processing() def test_serializable_keys(self, gen_graph): - graph, _ = build_graph(gen_graph, chunks=[([(0, 0, 0, 0)], [])]) - - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) + g = build_graph(gen_graph, n_layers=4, supervoxels={"a": SV()}) + root = g.cg.get_root(g.sv["a"]) progress = SubgraphProgress( - graph.meta, + g.cg.meta, node_ids=[root], return_layers=[2], serializable=True, @@ -40,41 +37,36 @@ def test_serializable_keys(self, gen_graph): class TestGetSubgraphNodes: def _build_graph(self, gen_graph): - cg, _ = build_graph( + return build_graph( gen_graph, - chunks=[ - ( - [(0, 0, 0, 0), (0, 0, 0, 1)], - [((0, 0, 0, 0), (0, 0, 0, 1), 0.5), ((0, 0, 0, 0), (1, 0, 0, 0), inf)], - ), - ([(1, 0, 0, 0)], [((1, 0, 0, 0), (0, 0, 0, 0), inf)]), - ], + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, + edges=[("a0", "a1", 0.5), ("a0", "b", inf)], ) - return cg def test_single_node(self, gen_graph): - graph = self._build_graph(gen_graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - result = get_subgraph_nodes(graph, root) + g = self._build_graph(gen_graph) + root = g.cg.get_root(g.sv["a0"]) + result = get_subgraph_nodes(g.cg, root) assert isinstance(result, dict) assert 2 in result def test_return_flattened(self, gen_graph): - graph = self._build_graph(gen_graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - result = get_subgraph_nodes(graph, root, return_flattened=True) + g = self._build_graph(gen_graph) + root = g.cg.get_root(g.sv["a0"]) + result = get_subgraph_nodes(g.cg, root, return_flattened=True) assert isinstance(result, np.ndarray) assert len(result) > 0 def test_multiple_nodes(self, gen_graph): - graph = self._build_graph(gen_graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - result = get_subgraph_nodes(graph, [root]) + g = self._build_graph(gen_graph) + root = g.cg.get_root(g.sv["a0"]) + result = get_subgraph_nodes(g.cg, [root]) assert root in result def test_serializable(self, gen_graph): - graph = self._build_graph(gen_graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - result = get_subgraph_nodes(graph, root, serializable=True) + g = self._build_graph(gen_graph) + root = g.cg.get_root(g.sv["a0"]) + result = get_subgraph_nodes(g.cg, root, serializable=True) # Keys should be layer ints, values should be arrays assert isinstance(result, dict) diff --git a/pychunkedgraph/tests/helpers.py b/pychunkedgraph/tests/helpers.py index 512f719c8..ba0f8d133 100644 --- a/pychunkedgraph/tests/helpers.py +++ b/pychunkedgraph/tests/helpers.py @@ -1,4 +1,5 @@ import threading +from collections import namedtuple from datetime import datetime, timedelta, UTC from functools import reduce from unittest.mock import MagicMock @@ -163,27 +164,35 @@ def get_layer_chunk_bounds( return layer_bounds_d -def build_graph(gen_graph, n_layers=4, chunks=(), *, timestamp=None, atomic_chunk_bounds=None): - """Build a test graph: atomic chunks plus the full parent hierarchy they imply, at one ts. +SV = namedtuple("SV", ["x", "y", "z", "seg"], defaults=(0, 0, 0, 0)) +BuiltGraph = namedtuple("BuiltGraph", ["cg", "sv", "ts"]) - A vertex is an (x, y, z, seg) atomic coordinate; an edge is ((x,y,z,seg), (x,y,z,seg), - affinity); chunks is a list of (vertices, edges). Parents are derived. Returns (cg, ts). + +def build_graph(gen_graph, n_layers, supervoxels, edges=(), *, timestamp=None, atomic_chunk_bounds=None): + """Build a test graph from named supervoxels and edges; parents derived, at one ts. + + supervoxels maps a name to its (x, y, z, seg) atomic coordinate; edges are + (name, name, affinity). Returns BuiltGraph(cg, sv, ts); sv maps each name to its node id. """ bounds = np.array([]) if atomic_chunk_bounds is None else atomic_chunk_bounds cg = gen_graph(n_layers=n_layers, atomic_chunk_bounds=bounds) ts = fake_timestamp() if timestamp is None else timestamp - atomic_coords = set() - for vertices, edges in chunks: - verts = [to_label(cg, 1, *v) for v in vertices] - labeled = [(to_label(cg, 1, *a), to_label(cg, 1, *b), aff) for a, b, aff in edges] - create_chunk(cg, vertices=verts, edges=labeled, timestamp=ts) - atomic_coords.update(v[:3] for v in vertices) + sv = {name: to_label(cg, 1, *coord) for name, coord in supervoxels.items()} + chunk = {name: tuple(coord[:3]) for name, coord in supervoxels.items()} + members = {} + for name in supervoxels: + members.setdefault(chunk[name], []).append(sv[name]) + for coord, labels in members.items(): + chunk_edges = [ + (sv[a], sv[b], aff) for a, b, aff in edges if coord in (chunk[a], chunk[b]) + ] + create_chunk(cg, vertices=labels, edges=chunk_edges, timestamp=ts) fanout = cg.meta.graph_config.FANOUT for layer in range(3, n_layers + 1): - coords = {tuple(np.array(c) // fanout ** (layer - 2)) for c in atomic_coords} - for coord in sorted(coords): - add_parent_chunk(cg, layer, list(coord), time_stamp=ts, n_threads=1) - return cg, ts + pcoords = {tuple(np.array(c) // fanout ** (layer - 2)) for c in members} + for pcoord in sorted(pcoords): + add_parent_chunk(cg, layer, list(pcoord), time_stamp=ts, n_threads=1) + return BuiltGraph(cg, sv, ts) class RowKeyLockRegistry: From 49675ff2bfceec577fd95acc3fba01f9adb379d1 Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Sun, 28 Jun 2026 18:53:32 +0000 Subject: [PATCH 06/12] test: build_graph returns BuiltGraph(cg, sv) Drop the unused ts from the result; callers unpack `cg, sv = build_graph(...)` and reference supervoxels by name. --- pychunkedgraph/tests/graph/test_cache.py | 69 ++++++++++----------- pychunkedgraph/tests/graph/test_subgraph.py | 38 ++++++------ pychunkedgraph/tests/helpers.py | 4 +- 3 files changed, 55 insertions(+), 56 deletions(-) diff --git a/pychunkedgraph/tests/graph/test_cache.py b/pychunkedgraph/tests/graph/test_cache.py index 1a964ac09..00e6b17d7 100644 --- a/pychunkedgraph/tests/graph/test_cache.py +++ b/pychunkedgraph/tests/graph/test_cache.py @@ -22,8 +22,7 @@ def test_many_to_one(self): class TestCacheService: - def _build_simple_graph(self, gen_graph): - """Build a simple 2-chunk graph with 2 SVs in the first chunk.""" + def _build(self, gen_graph): return build_graph( gen_graph, n_layers=4, @@ -32,36 +31,36 @@ def _build_simple_graph(self, gen_graph): ) def test_len(self, gen_graph): - g = self._build_simple_graph(gen_graph) - cache = CacheService(g.cg) - cache.parent(g.sv["a0"]) + cg, sv = self._build(gen_graph) + cache = CacheService(cg) + cache.parent(sv["a0"]) assert len(cache) >= 1 def test_clear(self, gen_graph): - g = self._build_simple_graph(gen_graph) - cache = CacheService(g.cg) - cache.parent(g.sv["a0"]) + cg, sv = self._build(gen_graph) + cache = CacheService(cg) + cache.parent(sv["a0"]) cache.clear() assert len(cache) == 0 def test_parent_miss_then_hit(self, gen_graph): - g = self._build_simple_graph(gen_graph) - cache = CacheService(g.cg) - sv = g.sv["a0"] + cg, sv = self._build(gen_graph) + cache = CacheService(cg) + a0 = sv["a0"] # First call is a miss - parent1 = cache.parent(sv) + parent1 = cache.parent(a0) assert cache.stats["parents"]["misses"] == 1 # Second call is a hit - parent2 = cache.parent(sv) + parent2 = cache.parent(a0) assert cache.stats["parents"]["hits"] == 1 assert parent1 == parent2 def test_children_backfills_parent(self, gen_graph): - g = self._build_simple_graph(gen_graph) - cache = CacheService(g.cg) - root = g.cg.get_root(g.sv["a0"]) + cg, sv = self._build(gen_graph) + cache = CacheService(cg) + root = cg.get_root(sv["a0"]) children = cache.children(root) assert len(children) > 0 # Children should be backfilled as parents @@ -69,48 +68,48 @@ def test_children_backfills_parent(self, gen_graph): assert child in cache.parents_cache def test_get_stats(self, gen_graph): - g = self._build_simple_graph(gen_graph) - cache = CacheService(g.cg) - sv = g.sv["a0"] - cache.parent(sv) - cache.parent(sv) + cg, sv = self._build(gen_graph) + cache = CacheService(cg) + a0 = sv["a0"] + cache.parent(a0) + cache.parent(a0) stats = cache.get_stats() assert "parents" in stats assert stats["parents"]["total"] == 2 assert "hit_rate" in stats["parents"] def test_reset_stats(self, gen_graph): - g = self._build_simple_graph(gen_graph) - cache = CacheService(g.cg) - cache.parent(g.sv["a0"]) + cg, sv = self._build(gen_graph) + cache = CacheService(cg) + cache.parent(sv["a0"]) cache.reset_stats() assert cache.stats["parents"]["hits"] == 0 assert cache.stats["parents"]["misses"] == 0 def test_parents_multiple_empty(self, gen_graph): - g = self._build_simple_graph(gen_graph) - cache = CacheService(g.cg) + cg, sv = self._build(gen_graph) + cache = CacheService(cg) result = cache.parents_multiple(np.array([], dtype=np.uint64)) assert len(result) == 0 def test_parents_multiple(self, gen_graph): - g = self._build_simple_graph(gen_graph) - cache = CacheService(g.cg) - svs = np.array([g.sv["a0"], g.sv["a1"]], dtype=np.uint64) + cg, sv = self._build(gen_graph) + cache = CacheService(cg) + svs = np.array([sv["a0"], sv["a1"]], dtype=np.uint64) result = cache.parents_multiple(svs) assert len(result) == 2 def test_children_multiple(self, gen_graph): - g = self._build_simple_graph(gen_graph) - cache = CacheService(g.cg) - root = g.cg.get_root(g.sv["a0"]) + cg, sv = self._build(gen_graph) + cache = CacheService(cg) + root = cg.get_root(sv["a0"]) result = cache.children_multiple(np.array([root], dtype=np.uint64)) assert root in result def test_children_multiple_flatten(self, gen_graph): - g = self._build_simple_graph(gen_graph) - cache = CacheService(g.cg) - root = g.cg.get_root(g.sv["a0"]) + cg, sv = self._build(gen_graph) + cache = CacheService(cg) + root = cg.get_root(sv["a0"]) result = cache.children_multiple( np.array([root], dtype=np.uint64), flatten=True ) diff --git a/pychunkedgraph/tests/graph/test_subgraph.py b/pychunkedgraph/tests/graph/test_subgraph.py index ca897518d..ccfddf104 100644 --- a/pychunkedgraph/tests/graph/test_subgraph.py +++ b/pychunkedgraph/tests/graph/test_subgraph.py @@ -11,10 +11,10 @@ class TestSubgraphProgress: def test_init(self, gen_graph): - g = build_graph(gen_graph, n_layers=4, supervoxels={"a": SV()}) - root = g.cg.get_root(g.sv["a"]) + cg, sv = build_graph(gen_graph, n_layers=4, supervoxels={"a": SV()}) + root = cg.get_root(sv["a"]) progress = SubgraphProgress( - g.cg.meta, + cg.meta, node_ids=[root], return_layers=[2], serializable=False, @@ -22,10 +22,10 @@ def test_init(self, gen_graph): assert not progress.done_processing() def test_serializable_keys(self, gen_graph): - g = build_graph(gen_graph, n_layers=4, supervoxels={"a": SV()}) - root = g.cg.get_root(g.sv["a"]) + cg, sv = build_graph(gen_graph, n_layers=4, supervoxels={"a": SV()}) + root = cg.get_root(sv["a"]) progress = SubgraphProgress( - g.cg.meta, + cg.meta, node_ids=[root], return_layers=[2], serializable=True, @@ -36,7 +36,7 @@ def test_serializable_keys(self, gen_graph): class TestGetSubgraphNodes: - def _build_graph(self, gen_graph): + def _build(self, gen_graph): return build_graph( gen_graph, n_layers=4, @@ -45,28 +45,28 @@ def _build_graph(self, gen_graph): ) def test_single_node(self, gen_graph): - g = self._build_graph(gen_graph) - root = g.cg.get_root(g.sv["a0"]) - result = get_subgraph_nodes(g.cg, root) + cg, sv = self._build(gen_graph) + root = cg.get_root(sv["a0"]) + result = get_subgraph_nodes(cg, root) assert isinstance(result, dict) assert 2 in result def test_return_flattened(self, gen_graph): - g = self._build_graph(gen_graph) - root = g.cg.get_root(g.sv["a0"]) - result = get_subgraph_nodes(g.cg, root, return_flattened=True) + cg, sv = self._build(gen_graph) + root = cg.get_root(sv["a0"]) + result = get_subgraph_nodes(cg, root, return_flattened=True) assert isinstance(result, np.ndarray) assert len(result) > 0 def test_multiple_nodes(self, gen_graph): - g = self._build_graph(gen_graph) - root = g.cg.get_root(g.sv["a0"]) - result = get_subgraph_nodes(g.cg, [root]) + cg, sv = self._build(gen_graph) + root = cg.get_root(sv["a0"]) + result = get_subgraph_nodes(cg, [root]) assert root in result def test_serializable(self, gen_graph): - g = self._build_graph(gen_graph) - root = g.cg.get_root(g.sv["a0"]) - result = get_subgraph_nodes(g.cg, root, serializable=True) + cg, sv = self._build(gen_graph) + root = cg.get_root(sv["a0"]) + result = get_subgraph_nodes(cg, root, serializable=True) # Keys should be layer ints, values should be arrays assert isinstance(result, dict) diff --git a/pychunkedgraph/tests/helpers.py b/pychunkedgraph/tests/helpers.py index ba0f8d133..0ca36d535 100644 --- a/pychunkedgraph/tests/helpers.py +++ b/pychunkedgraph/tests/helpers.py @@ -165,7 +165,7 @@ def get_layer_chunk_bounds( SV = namedtuple("SV", ["x", "y", "z", "seg"], defaults=(0, 0, 0, 0)) -BuiltGraph = namedtuple("BuiltGraph", ["cg", "sv", "ts"]) +BuiltGraph = namedtuple("BuiltGraph", ["cg", "sv"]) def build_graph(gen_graph, n_layers, supervoxels, edges=(), *, timestamp=None, atomic_chunk_bounds=None): @@ -192,7 +192,7 @@ def build_graph(gen_graph, n_layers, supervoxels, edges=(), *, timestamp=None, a pcoords = {tuple(np.array(c) // fanout ** (layer - 2)) for c in members} for pcoord in sorted(pcoords): add_parent_chunk(cg, layer, list(pcoord), time_stamp=ts, n_threads=1) - return BuiltGraph(cg, sv, ts) + return BuiltGraph(cg, sv) class RowKeyLockRegistry: From 03a77d9020ee53b5ec86944adcf9a9fd4f514227 Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Sun, 28 Jun 2026 21:12:38 +0000 Subject: [PATCH 07/12] test: migrate graph-test builders to the build_graph factory Replace the per-fixture gen_graph + create_chunk + add_parent_chunk + to_label scaffolding with named-supervoxel build_graph specs. The parent hierarchy is derived from the atomic chunks rather than hand-listed, and readable SV() coordinates replace the repeated raw (x, y, z, seg) tuples. Co-Authored-By: Claude --- .../tests/graph/test_analysis_pathing.py | 498 ++----- .../tests/graph/test_chunkedgraph_extended.py | 1177 ++++++----------- pychunkedgraph/tests/graph/test_lineage.py | 184 +-- pychunkedgraph/tests/graph/test_operation.py | 304 ++--- .../tests/graph/test_segmenthistory.py | 184 +-- .../tests/graph/test_sv_lookup_main.py | 75 +- 6 files changed, 786 insertions(+), 1636 deletions(-) diff --git a/pychunkedgraph/tests/graph/test_analysis_pathing.py b/pychunkedgraph/tests/graph/test_analysis_pathing.py index d0a7f8a6b..5585902a5 100644 --- a/pychunkedgraph/tests/graph/test_analysis_pathing.py +++ b/pychunkedgraph/tests/graph/test_analysis_pathing.py @@ -14,110 +14,64 @@ compute_rough_coordinate_path, ) -from ..helpers import create_chunk, to_label, fake_timestamp -from ...ingest.create.parent_layer import add_parent_chunk +from ..helpers import SV, build_graph class TestGetFirstSharedParent: def _build_graph(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, + edges=[("a0", "a1", 0.5), ("a0", "b", inf)], ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - ], - timestamp=fake_ts, - ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - return graph def test_same_root(self, gen_graph): - graph = self._build_graph(gen_graph) - sv0 = to_label(graph, 1, 0, 0, 0, 0) - sv1 = to_label(graph, 1, 1, 0, 0, 0) - parent = get_first_shared_parent(graph, sv0, sv1) + cg, sv = self._build_graph(gen_graph) + sv0 = sv["a0"] + sv1 = sv["b"] + parent = get_first_shared_parent(cg, sv0, sv1) assert parent is not None # The shared parent should be an ancestor of both SVs - root = graph.get_root(sv0) + root = cg.get_root(sv0) # Verify the shared parent is on the path to root - assert graph.get_root(parent) == root + assert cg.get_root(parent) == root def test_different_roots_returns_none(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - # Create two disconnected chunks - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, - ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "b": SV(x=1)}, ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - - sv0 = to_label(graph, 1, 0, 0, 0, 0) - sv1 = to_label(graph, 1, 1, 0, 0, 0) - parent = get_first_shared_parent(graph, sv0, sv1) + sv0 = sv["a0"] + sv1 = sv["b"] + parent = get_first_shared_parent(cg, sv0, sv1) assert parent is None class TestGetChildrenAtLayer: def test_basic(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - ], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1)}, + edges=[("a0", "a1", 0.5)], ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - children = get_children_at_layer(graph, root, 2) + root = cg.get_root(sv["a0"]) + children = get_children_at_layer(cg, root, 2) assert len(children) > 0 for child in children: - assert graph.get_chunk_layer(child) == 2 + assert cg.get_chunk_layer(child) == 2 def test_allow_lower_layers(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV()}, ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - children = get_children_at_layer(graph, root, 2, allow_lower_layers=True) + root = cg.get_root(sv["a0"]) + children = get_children_at_layer(cg, root, 2, allow_lower_layers=True) assert len(children) > 0 @@ -127,48 +81,19 @@ def _build_3chunk_graph(self, gen_graph): A:sv0 -- B:sv0 -- C:sv0 """ - graph = gen_graph(n_layers=4) - - # Chunk A: sv0 connected to B:sv0 via cross-chunk edge - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - ) - - # Chunk B: sv0 connected to A:sv0 and C:sv0 via cross-chunk edges - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 2, 0, 0, 0), inf), - ], - ) - - # Chunk C: sv0 connected to B:sv0 via cross-chunk edge - create_chunk( - graph, - vertices=[to_label(graph, 1, 2, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 2, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "b": SV(x=1), "c": SV(x=2)}, + edges=[("a0", "b", inf), ("b", "c", inf)], ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 3, [1, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - - return graph - def test_basic(self, gen_graph): """get_lvl2_edge_list should return edges between L2 IDs for a connected root.""" - graph = self._build_3chunk_graph(gen_graph) + cg, sv = self._build_3chunk_graph(gen_graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - edges = get_lvl2_edge_list(graph, root) + root = cg.get_root(sv["a0"]) + edges = get_lvl2_edge_list(cg, root) # There should be at least 2 edges: A_l2--B_l2 and B_l2--C_l2 assert edges.shape[0] >= 2 @@ -177,22 +102,17 @@ def test_basic(self, gen_graph): # All edge IDs should be L2 nodes (layer 2) for edge in edges: for node_id in edge: - assert graph.get_chunk_layer(node_id) == 2 + assert cg.get_chunk_layer(node_id) == 2 def test_single_chunk_no_cross_edges(self, gen_graph): """A single isolated chunk should produce no L2 edges.""" - graph = gen_graph(n_layers=4) - - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV()}, ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - edges = get_lvl2_edge_list(graph, root) + root = cg.get_root(sv["a0"]) + edges = get_lvl2_edge_list(cg, root) assert edges.shape[0] == 0 @@ -203,50 +123,24 @@ def _build_3chunk_graph(self, gen_graph): A:sv0 -- B:sv0 -- C:sv0 """ - graph = gen_graph(n_layers=4) - - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - ) - - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 2, 0, 0, 0), inf), - ], - ) - - create_chunk( - graph, - vertices=[to_label(graph, 1, 2, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 2, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "b": SV(x=1), "c": SV(x=2)}, + edges=[("a0", "b", inf), ("b", "c", inf)], ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 3, [1, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - - return graph - def test_path_between_endpoints(self, gen_graph): """find_l2_shortest_path should return a path from source to target L2 IDs.""" - graph = self._build_3chunk_graph(gen_graph) + cg, sv = self._build_3chunk_graph(gen_graph) # Get L2 parents of the supervoxels - sv_a = to_label(graph, 1, 0, 0, 0, 0) - sv_c = to_label(graph, 1, 2, 0, 0, 0) - l2_a = graph.get_parent(sv_a) - l2_c = graph.get_parent(sv_c) + sv_a = sv["a0"] + sv_c = sv["c"] + l2_a = cg.get_parent(sv_a) + l2_c = cg.get_parent(sv_c) - path = find_l2_shortest_path(graph, l2_a, l2_c) + path = find_l2_shortest_path(cg, l2_a, l2_c) assert path is not None assert len(path) == 3 # A_l2 -> B_l2 -> C_l2 @@ -255,18 +149,18 @@ def test_path_between_endpoints(self, gen_graph): assert path[-1] == l2_c # All nodes in path should be layer 2 for node_id in path: - assert graph.get_chunk_layer(node_id) == 2 + assert cg.get_chunk_layer(node_id) == 2 def test_adjacent_l2_ids(self, gen_graph): """find_l2_shortest_path between directly connected L2 IDs should return length 2 path.""" - graph = self._build_3chunk_graph(gen_graph) + cg, sv = self._build_3chunk_graph(gen_graph) - sv_a = to_label(graph, 1, 0, 0, 0, 0) - sv_b = to_label(graph, 1, 1, 0, 0, 0) - l2_a = graph.get_parent(sv_a) - l2_b = graph.get_parent(sv_b) + sv_a = sv["a0"] + sv_b = sv["b"] + l2_a = cg.get_parent(sv_a) + l2_b = cg.get_parent(sv_b) - path = find_l2_shortest_path(graph, l2_a, l2_b) + path = find_l2_shortest_path(cg, l2_a, l2_b) assert path is not None assert len(path) == 2 @@ -275,28 +169,18 @@ def test_adjacent_l2_ids(self, gen_graph): def test_disconnected_returns_none(self, gen_graph): """find_l2_shortest_path should return None when L2 IDs belong to different roots.""" - graph = gen_graph(n_layers=4) - # Create two disconnected chunks - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], - ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[], + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "b": SV(x=1)}, ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - - sv_a = to_label(graph, 1, 0, 0, 0, 0) - sv_b = to_label(graph, 1, 1, 0, 0, 0) - l2_a = graph.get_parent(sv_a) - l2_b = graph.get_parent(sv_b) + sv_a = sv["a0"] + sv_b = sv["b"] + l2_a = cg.get_parent(sv_a) + l2_b = cg.get_parent(sv_b) - path = find_l2_shortest_path(graph, l2_a, l2_b) + path = find_l2_shortest_path(cg, l2_a, l2_b) assert path is None @@ -305,102 +189,59 @@ class TestGetChildrenAtLayerEdgeCases: def test_children_at_layer_2_with_multiple_svs(self, gen_graph): """Query children at layer 2 when root has multiple SVs in same chunk.""" - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - to_label(graph, 1, 0, 0, 0, 2), - ], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - (to_label(graph, 1, 0, 0, 0, 1), to_label(graph, 1, 0, 0, 0, 2), 0.5), - ], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1), "a2": SV(seg=2)}, + edges=[("a0", "a1", 0.5), ("a1", "a2", 0.5)], ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - children = get_children_at_layer(graph, root, 2) + root = cg.get_root(sv["a0"]) + children = get_children_at_layer(cg, root, 2) assert len(children) > 0 for child in children: - assert graph.get_chunk_layer(child) == 2 + assert cg.get_chunk_layer(child) == 2 def test_children_at_intermediate_layer(self, gen_graph): """Query children at layer 3 from root at layer 4.""" - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - timestamp=fake_ts, - ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - ], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "b": SV(x=1)}, + edges=[("a0", "b", inf)], ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - children = get_children_at_layer(graph, root, 3) + root = cg.get_root(sv["a0"]) + children = get_children_at_layer(cg, root, 3) assert len(children) > 0 for child in children: - assert graph.get_chunk_layer(child) == 3 + assert cg.get_chunk_layer(child) == 3 def test_children_allow_lower_layers_with_cross_chunk(self, gen_graph): """Query with allow_lower_layers=True should include layer<=target.""" - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV()}, ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) + root = cg.get_root(sv["a0"]) # Ask for layer 3 with allow_lower_layers=True - children = get_children_at_layer(graph, root, 3, allow_lower_layers=True) + children = get_children_at_layer(cg, root, 3, allow_lower_layers=True) assert len(children) > 0 for child in children: - assert graph.get_chunk_layer(child) <= 3 + assert cg.get_chunk_layer(child) <= 3 def test_children_at_layer_from_l2_node(self, gen_graph): """Querying children at layer 2 from a layer 2 node should return the node itself or its layer-2 children (which is itself).""" - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV()}, ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - - sv = to_label(graph, 1, 0, 0, 0, 0) - l2 = graph.get_parent(sv) + sv0 = sv["a0"] + l2 = cg.get_parent(sv0) # From l2, get children at layer 2 (with allow_lower=True since # the children of an L2 node are SVs at layer 1) - children = get_children_at_layer(graph, l2, 2, allow_lower_layers=True) + children = get_children_at_layer(cg, l2, 2, allow_lower_layers=True) assert len(children) > 0 @@ -409,52 +250,28 @@ class TestGetLvl2EdgeListWithBbox: def _build_3chunk_graph(self, gen_graph): """Build a graph with 3 chunks connected linearly.""" - graph = gen_graph(n_layers=4) - - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 2, 0, 0, 0), inf), - ], - ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 2, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 2, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "b": SV(x=1), "c": SV(x=2)}, + edges=[("a0", "b", inf), ("b", "c", inf)], ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 3, [1, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - - return graph - def test_lvl2_edge_list_with_bbox(self, gen_graph): """get_lvl2_edge_list with a bbox should return edges within the bbox.""" - graph = self._build_3chunk_graph(gen_graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) + cg, sv = self._build_3chunk_graph(gen_graph) + root = cg.get_root(sv["a0"]) # Use a large bbox that encompasses everything bbox = np.array([[0, 0, 0], [2048, 2048, 256]]) - edges = get_lvl2_edge_list(graph, root, bbox=bbox) + edges = get_lvl2_edge_list(cg, root, bbox=bbox) # Should have edges assert edges.shape[1] == 2 # All IDs should be L2 nodes for edge in edges: for node_id in edge: - assert graph.get_chunk_layer(node_id) == 2 + assert cg.get_chunk_layer(node_id) == 2 class TestFindL2ShortestPathEdgeCases: @@ -462,49 +279,19 @@ class TestFindL2ShortestPathEdgeCases: def test_path_through_chain(self, gen_graph): """find_l2_shortest_path through a 4-chunk chain should return correct length.""" - graph = gen_graph(n_layers=4) - # Build a 4-chunk chain: A(0,0,0)--B(1,0,0)--C(2,0,0)--D(3,0,0) - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 2, 0, 0, 0), inf), - ], - ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 2, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 2, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - (to_label(graph, 1, 2, 0, 0, 0), to_label(graph, 1, 3, 0, 0, 0), inf), - ], + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "b": SV(x=1), "c": SV(x=2), "d": SV(x=3)}, + edges=[("a0", "b", inf), ("b", "c", inf), ("c", "d", inf)], ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 3, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 3, 0, 0, 0), to_label(graph, 1, 2, 0, 0, 0), inf), - ], - ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 3, [1, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - - sv_a = to_label(graph, 1, 0, 0, 0, 0) - sv_d = to_label(graph, 1, 3, 0, 0, 0) - l2_a = graph.get_parent(sv_a) - l2_d = graph.get_parent(sv_d) + sv_a = sv["a0"] + sv_d = sv["d"] + l2_a = cg.get_parent(sv_a) + l2_d = cg.get_parent(sv_d) - path = find_l2_shortest_path(graph, l2_a, l2_d) + path = find_l2_shortest_path(cg, l2_a, l2_d) assert path is not None assert len(path) == 4 # A_l2 -> B_l2 -> C_l2 -> D_l2 assert path[0] == l2_a @@ -516,40 +303,27 @@ class TestComputeRoughCoordinatePath: def test_basic_coordinate_path(self, gen_graph): """compute_rough_coordinate_path should return a list of float32 3D coordinates.""" - graph = gen_graph(n_layers=4) - - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "b": SV(x=1)}, + edges=[("a0", "b", inf)], ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - ], - ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - - sv_a = to_label(graph, 1, 0, 0, 0, 0) - sv_b = to_label(graph, 1, 1, 0, 0, 0) - l2_a = graph.get_parent(sv_a) - l2_b = graph.get_parent(sv_b) + sv_a = sv["a0"] + sv_b = sv["b"] + l2_a = cg.get_parent(sv_a) + l2_b = cg.get_parent(sv_b) - path = find_l2_shortest_path(graph, l2_a, l2_b) + path = find_l2_shortest_path(cg, l2_a, l2_b) assert path is not None # Mock cv methods that CloudVolumeMock doesn't have mock_cv = MagicMock() mock_cv.mip_voxel_offset = MagicMock(return_value=np.array([0, 0, 0])) mock_cv.mip_resolution = MagicMock(return_value=np.array([1, 1, 1])) - graph.meta._ws_cv = mock_cv + cg.meta._ws_cv = mock_cv - coordinate_path = compute_rough_coordinate_path(graph, path) + coordinate_path = compute_rough_coordinate_path(cg, path) assert len(coordinate_path) == len(path) for coord in coordinate_path: assert isinstance(coord, np.ndarray) diff --git a/pychunkedgraph/tests/graph/test_chunkedgraph_extended.py b/pychunkedgraph/tests/graph/test_chunkedgraph_extended.py index 5e6d4a06e..99202ecc6 100644 --- a/pychunkedgraph/tests/graph/test_chunkedgraph_extended.py +++ b/pychunkedgraph/tests/graph/test_chunkedgraph_extended.py @@ -6,8 +6,7 @@ import numpy as np import pytest -from ..helpers import create_chunk, to_label, fake_timestamp -from ...ingest.create.parent_layer import add_parent_chunk +from ..helpers import SV, build_graph from ...graph.operation import GraphEditOperation, MergeOperation, SplitOperation from ...graph.exceptions import PreconditionError @@ -15,124 +14,95 @@ class TestChunkedGraphExtended: def _build_graph(self, gen_graph): """Build a simple multi-chunk graph.""" - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - - # Chunk A: sv 0, 1 connected - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - timestamp=fake_ts, + # Chunk A: sv 0, 1 connected; Chunk B: sv 0 connected cross-chunk to A + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, + edges=[("a0", "a1", 0.5), ("a0", "b", inf)], ) - # Chunk B: sv 0 connected cross-chunk to A - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - ], - timestamp=fake_ts, - ) - - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - return graph - def test_is_root_true(self, gen_graph): - graph = self._build_graph(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - root = graph.get_root(sv) - assert graph.is_root(root) + cg, sv = self._build_graph(gen_graph) + root = cg.get_root(sv["a0"]) + assert cg.is_root(root) def test_is_root_false(self, gen_graph): - graph = self._build_graph(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - assert not graph.is_root(sv) + cg, sv = self._build_graph(gen_graph) + assert not cg.is_root(sv["a0"]) def test_get_parents_raw_only(self, gen_graph): - graph = self._build_graph(gen_graph) + cg, sv = self._build_graph(gen_graph) svs = np.array( [ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), + sv["a0"], + sv["a1"], ], dtype=np.uint64, ) - parents = graph.get_parents(svs, raw_only=True) + parents = cg.get_parents(svs, raw_only=True) assert len(parents) == 2 # Parents should be L2 IDs for p in parents: - assert graph.get_chunk_layer(p) == 2 + assert cg.get_chunk_layer(p) == 2 def test_get_parents_fail_to_zero(self, gen_graph): - graph = self._build_graph(gen_graph) + cg, sv = self._build_graph(gen_graph) # Non-existent ID should return 0 with fail_to_zero bad_id = np.uint64(99999999) - result = graph.get_parents( + result = cg.get_parents( np.array([bad_id], dtype=np.uint64), fail_to_zero=True ) assert result[0] == 0 def test_get_children_flatten(self, gen_graph): - graph = self._build_graph(gen_graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - children = graph.get_children([root], flatten=True) + cg, sv = self._build_graph(gen_graph) + root = cg.get_root(sv["a0"]) + children = cg.get_children([root], flatten=True) assert isinstance(children, np.ndarray) assert len(children) > 0 def test_is_latest_roots(self, gen_graph): - graph = self._build_graph(gen_graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - result = graph.is_latest_roots(np.array([root], dtype=np.uint64)) + cg, sv = self._build_graph(gen_graph) + root = cg.get_root(sv["a0"]) + result = cg.is_latest_roots(np.array([root], dtype=np.uint64)) assert result[0] def test_get_node_timestamps(self, gen_graph): - graph = self._build_graph(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - root = graph.get_root(sv) - ts = graph.get_node_timestamps(np.array([root]), return_numpy=False) + cg, sv = self._build_graph(gen_graph) + root = cg.get_root(sv["a0"]) + ts = cg.get_node_timestamps(np.array([root]), return_numpy=False) assert len(ts) == 1 def test_get_earliest_timestamp(self, gen_graph): - graph = self._build_graph(gen_graph) + cg, sv = self._build_graph(gen_graph) # no ops -> the boundary stamped by the root-layer build - assert isinstance(graph.get_earliest_timestamp(), datetime) + assert isinstance(cg.get_earliest_timestamp(), datetime) def test_get_l2children(self, gen_graph): - graph = self._build_graph(gen_graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - l2_children = graph.get_l2children(np.array([root], dtype=np.uint64)) + cg, sv = self._build_graph(gen_graph) + root = cg.get_root(sv["a0"]) + l2_children = cg.get_l2children(np.array([root], dtype=np.uint64)) assert len(l2_children) > 0 for child in l2_children: - assert graph.get_chunk_layer(child) == 2 + assert cg.get_chunk_layer(child) == 2 # --- helpers for edit-based tests --- def _build_and_merge(self, gen_graph): """Build a single-chunk graph with two disconnected SVs and merge them.""" - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) - result = graph.add_edges( + result = cg.add_edges( "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], + [sv["a0"], sv["a1"]], affinities=[0.3], ) - return graph, result.new_root_ids[0], result + return cg, result.new_root_ids[0], result @pytest.mark.timeout(30) def test_get_operation_ids(self, gen_graph): @@ -148,28 +118,36 @@ def test_get_operation_ids(self, gen_graph): @pytest.mark.timeout(30) def test_get_single_leaf_multiple(self, gen_graph): """get_single_leaf_multiple for an L2 node should return an L1 supervoxel.""" - graph, new_root, _ = self._build_and_merge(gen_graph) + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, + ) + result = cg.add_edges( + "TestUser", + [sv["a0"], sv["a1"]], + affinities=[0.3], + ) + new_root = result.new_root_ids[0] # The new_root in n_layers=2 is actually L2 - assert graph.get_chunk_layer(new_root) == 2 - leaves = graph.get_single_leaf_multiple(np.array([new_root], dtype=np.uint64)) + assert cg.get_chunk_layer(new_root) == 2 + leaves = cg.get_single_leaf_multiple(np.array([new_root], dtype=np.uint64)) assert len(leaves) == 1 - assert graph.get_chunk_layer(leaves[0]) == 1 + assert cg.get_chunk_layer(leaves[0]) == 1 # The returned leaf should be one of our two SVs - sv0 = to_label(graph, 1, 0, 0, 0, 0) - sv1 = to_label(graph, 1, 0, 0, 0, 1) - assert leaves[0] in [sv0, sv1] + assert leaves[0] in [sv["a0"], sv["a1"]] @pytest.mark.timeout(30) def test_get_atomic_cross_edges(self, gen_graph): """get_atomic_cross_edges for an L2 node with cross-chunk connections.""" - graph = self._build_graph(gen_graph) - sv_a0 = to_label(graph, 1, 0, 0, 0, 0) + cg, sv = self._build_graph(gen_graph) # Get the L2 parent of sv_a0 - parent = graph.get_parents(np.array([sv_a0], dtype=np.uint64), raw_only=True)[0] - assert graph.get_chunk_layer(parent) == 2 + parent = cg.get_parents(np.array([sv["a0"]], dtype=np.uint64), raw_only=True)[0] + assert cg.get_chunk_layer(parent) == 2 - result = graph.get_atomic_cross_edges([parent]) + result = cg.get_atomic_cross_edges([parent]) assert parent in result # Should have at least one layer of cross edges assert isinstance(result[parent], dict) @@ -177,26 +155,34 @@ def test_get_atomic_cross_edges(self, gen_graph): @pytest.mark.timeout(30) def test_get_cross_chunk_edges_raw(self, gen_graph): """get_cross_chunk_edges with raw_only=True should return cross edges.""" - graph = self._build_graph(gen_graph) - sv_a0 = to_label(graph, 1, 0, 0, 0, 0) + cg, sv = self._build_graph(gen_graph) # Get the L2 parent - parent = graph.get_parents(np.array([sv_a0], dtype=np.uint64), raw_only=True)[0] - assert graph.get_chunk_layer(parent) == 2 + parent = cg.get_parents(np.array([sv["a0"]], dtype=np.uint64), raw_only=True)[0] + assert cg.get_chunk_layer(parent) == 2 - result = graph.get_cross_chunk_edges([parent], raw_only=True) + result = cg.get_cross_chunk_edges([parent], raw_only=True) assert parent in result assert isinstance(result[parent], dict) @pytest.mark.timeout(30) def test_get_parents_not_current(self, gen_graph): """get_parents with current=False should return list of (parent, timestamp) tuples.""" - graph, new_root, _ = self._build_and_merge(gen_graph) - sv0 = to_label(graph, 1, 0, 0, 0, 0) + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, + ) + cg.add_edges( + "TestUser", + [sv["a0"], sv["a1"]], + affinities=[0.3], + ) # current=False returns list of lists of (value, timestamp) pairs - parents = graph.get_parents( - np.array([sv0], dtype=np.uint64), raw_only=True, current=False + parents = cg.get_parents( + np.array([sv["a0"]], dtype=np.uint64), raw_only=True, current=False ) assert len(parents) == 1 # Each element is a list of (parent_value, timestamp) tuples @@ -212,52 +198,45 @@ class TestFromLogRecord: def _build_two_sv_graph(self, gen_graph): """Build a 2-layer graph with two disconnected SVs.""" - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) - return graph @pytest.mark.timeout(30) def test_merge_from_log(self, gen_graph): """After a merge, from_log_record should return a MergeOperation.""" - graph = self._build_two_sv_graph(gen_graph) - result = graph.add_edges( + cg, sv = self._build_two_sv_graph(gen_graph) + result = cg.add_edges( "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], + [sv["a0"], sv["a1"]], affinities=[0.3], ) - log, ts = graph.client.read_log_entry(result.operation_id) - op = GraphEditOperation.from_log_record(graph, log) + log, ts = cg.client.read_log_entry(result.operation_id) + op = GraphEditOperation.from_log_record(cg, log) assert isinstance(op, MergeOperation) @pytest.mark.timeout(30) def test_split_from_log(self, gen_graph): """After a split, from_log_record should return a SplitOperation.""" - graph = self._build_two_sv_graph(gen_graph) + cg, sv = self._build_two_sv_graph(gen_graph) # First merge so the SVs belong to the same root - graph.add_edges( + cg.add_edges( "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], + [sv["a0"], sv["a1"]], affinities=[0.3], ) # Now split them - split_result = graph.remove_edges( + split_result = cg.remove_edges( "TestUser", - source_ids=to_label(graph, 1, 0, 0, 0, 0), - sink_ids=to_label(graph, 1, 0, 0, 0, 1), + source_ids=sv["a0"], + sink_ids=sv["a1"], mincut=False, ) - log, ts = graph.client.read_log_entry(split_result.operation_id) - op = GraphEditOperation.from_log_record(graph, log) + log, ts = cg.client.read_log_entry(split_result.operation_id) + op = GraphEditOperation.from_log_record(cg, log) assert isinstance(op, SplitOperation) @@ -266,41 +245,33 @@ class TestCheckIds: def _build_two_sv_graph(self, gen_graph): """Build a 2-layer graph with two disconnected SVs.""" - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) - return graph @pytest.mark.timeout(30) def test_source_equals_sink_raises(self, gen_graph): """MergeOperation with source==sink should raise PreconditionError (self-loop).""" - graph = self._build_two_sv_graph(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) + cg, sv = self._build_two_sv_graph(gen_graph) with pytest.raises(PreconditionError): - graph.add_edges( + cg.add_edges( "TestUser", - [sv, sv], + [sv["a0"], sv["a0"]], affinities=[0.3], ) @pytest.mark.timeout(30) def test_nonexistent_supervoxel_raises(self, gen_graph): """Using a supervoxel ID that doesn't exist should raise an error.""" - graph = self._build_two_sv_graph(gen_graph) - sv_real = to_label(graph, 1, 0, 0, 0, 0) + cg, sv = self._build_two_sv_graph(gen_graph) + sv_real = sv["a0"] # Use a layer-2 ID as a fake "supervoxel", which fails the layer check - sv_fake = to_label(graph, 2, 0, 0, 0, 99) + sv_fake = cg.get_node_id(np.uint64(99), layer=2, x=0, y=0, z=0) with pytest.raises(Exception): - graph.add_edges( + cg.add_edges( "TestUser", [sv_real, sv_fake], affinities=[0.3], @@ -311,144 +282,98 @@ class TestGetRootsExtended: """Tests for get_roots with stop_layer and ceil parameters (lines 380-461).""" def _build_cross_chunk(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, + edges=[("a0", "a1", 0.5), ("a0", "b", inf)], ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - ], - timestamp=fake_ts, - ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 3, [1, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - return graph, fake_ts @pytest.mark.timeout(30) def test_get_roots_with_stop_layer(self, gen_graph): """get_roots with stop_layer should return IDs at that layer.""" - graph, _ = self._build_cross_chunk(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) + cg, sv = self._build_cross_chunk(gen_graph) # Stop at layer 3 instead of going to root (layer 4) - result = graph.get_roots(np.array([sv], dtype=np.uint64), stop_layer=3) + result = cg.get_roots(np.array([sv["a0"]], dtype=np.uint64), stop_layer=3) assert len(result) == 1 - assert graph.get_chunk_layer(result[0]) == 3 + assert cg.get_chunk_layer(result[0]) == 3 @pytest.mark.timeout(30) def test_get_roots_with_stop_layer_and_ceil_false(self, gen_graph): """get_roots with stop_layer and ceil=False should not exceed stop_layer.""" - graph, _ = self._build_cross_chunk(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - result = graph.get_roots( - np.array([sv], dtype=np.uint64), stop_layer=3, ceil=False + cg, sv = self._build_cross_chunk(gen_graph) + result = cg.get_roots( + np.array([sv["a0"]], dtype=np.uint64), stop_layer=3, ceil=False ) assert len(result) == 1 - assert graph.get_chunk_layer(result[0]) <= 3 + assert cg.get_chunk_layer(result[0]) <= 3 @pytest.mark.timeout(30) def test_get_roots_multiple_svs(self, gen_graph): """get_roots with multiple SVs should return root for each.""" - graph, _ = self._build_cross_chunk(gen_graph) + cg, sv = self._build_cross_chunk(gen_graph) svs = np.array( [ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - to_label(graph, 1, 1, 0, 0, 0), + sv["a0"], + sv["a1"], + sv["b"], ], dtype=np.uint64, ) - roots = graph.get_roots(svs) + roots = cg.get_roots(svs) assert len(roots) == 3 # All should reach the top layer for r in roots: - assert graph.get_chunk_layer(r) == 4 + assert cg.get_chunk_layer(r) == 4 @pytest.mark.timeout(30) def test_get_roots_already_at_stop_layer(self, gen_graph): """get_roots for a node already at stop_layer should return it unchanged.""" - graph, _ = self._build_cross_chunk(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - root = graph.get_root(sv) + cg, sv = self._build_cross_chunk(gen_graph) + root = cg.get_root(sv["a0"]) # root is at layer 4; asking for stop_layer=4 should return it - result = graph.get_roots(np.array([root], dtype=np.uint64), stop_layer=4) + result = cg.get_roots(np.array([root], dtype=np.uint64), stop_layer=4) assert result[0] == root @pytest.mark.timeout(30) def test_get_roots_fail_to_zero(self, gen_graph): """get_roots with a zero ID and fail_to_zero should keep it as zero.""" - graph, _ = self._build_cross_chunk(gen_graph) - result = graph.get_roots(np.array([0], dtype=np.uint64), fail_to_zero=True) + cg, sv = self._build_cross_chunk(gen_graph) + result = cg.get_roots(np.array([0], dtype=np.uint64), fail_to_zero=True) assert result[0] == 0 @pytest.mark.timeout(30) def test_get_root_stop_layer_ceil_false(self, gen_graph): """get_root (singular) with stop_layer and ceil=False.""" - graph, _ = self._build_cross_chunk(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - result = graph.get_root(sv, stop_layer=3, ceil=False) - assert graph.get_chunk_layer(result) <= 3 + cg, sv = self._build_cross_chunk(gen_graph) + result = cg.get_root(sv["a0"], stop_layer=3, ceil=False) + assert cg.get_chunk_layer(result) <= 3 class TestGetChildrenExtended: """Tests for get_children with flatten=True and edge cases (lines 271-296).""" def _build_graph(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, + edges=[("a0", "a1", 0.5), ("a0", "b", inf)], ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - ], - timestamp=fake_ts, - ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 3, [1, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - return graph @pytest.mark.timeout(30) def test_get_children_flatten_multiple(self, gen_graph): """get_children with multiple node IDs and flatten=True returns flat array.""" - graph = self._build_graph(gen_graph) - sv_a0 = to_label(graph, 1, 0, 0, 0, 0) - sv_b0 = to_label(graph, 1, 1, 0, 0, 0) + cg, sv = self._build_graph(gen_graph) - parent_a = graph.get_parents(np.array([sv_a0], dtype=np.uint64), raw_only=True)[ + parent_a = cg.get_parents(np.array([sv["a0"]], dtype=np.uint64), raw_only=True)[ 0 ] - parent_b = graph.get_parents(np.array([sv_b0], dtype=np.uint64), raw_only=True)[ + parent_b = cg.get_parents(np.array([sv["b"]], dtype=np.uint64), raw_only=True)[ 0 ] - children = graph.get_children([parent_a, parent_b], flatten=True) + children = cg.get_children([parent_a, parent_b], flatten=True) assert isinstance(children, np.ndarray) # Should contain at least sv_a0, sv_a1, sv_b0 assert len(children) >= 3 @@ -456,28 +381,26 @@ def test_get_children_flatten_multiple(self, gen_graph): @pytest.mark.timeout(30) def test_get_children_flatten_empty(self, gen_graph): """get_children with flatten=True on empty list returns empty array.""" - graph = self._build_graph(gen_graph) - children = graph.get_children([], flatten=True) + cg, sv = self._build_graph(gen_graph) + children = cg.get_children([], flatten=True) assert isinstance(children, np.ndarray) assert len(children) == 0 @pytest.mark.timeout(30) def test_get_children_dict(self, gen_graph): """get_children without flatten returns a dict.""" - graph = self._build_graph(gen_graph) - sv_a0 = to_label(graph, 1, 0, 0, 0, 0) - parent = graph.get_parents(np.array([sv_a0], dtype=np.uint64), raw_only=True)[0] - children_d = graph.get_children([parent]) + cg, sv = self._build_graph(gen_graph) + parent = cg.get_parents(np.array([sv["a0"]], dtype=np.uint64), raw_only=True)[0] + children_d = cg.get_children([parent]) assert isinstance(children_d, dict) assert parent in children_d @pytest.mark.timeout(30) def test_get_children_scalar(self, gen_graph): """get_children with a scalar node_id returns an array.""" - graph = self._build_graph(gen_graph) - sv_a0 = to_label(graph, 1, 0, 0, 0, 0) - parent = graph.get_parents(np.array([sv_a0], dtype=np.uint64), raw_only=True)[0] - children = graph.get_children(parent, raw_only=True) + cg, sv = self._build_graph(gen_graph) + parent = cg.get_parents(np.array([sv["a0"]], dtype=np.uint64), raw_only=True)[0] + children = cg.get_children(parent, raw_only=True) assert isinstance(children, np.ndarray) assert len(children) >= 1 @@ -486,30 +409,24 @@ class TestIsLatestRootsExtended: """Tests for is_latest_roots (lines 524-544).""" def _build_and_merge(self, gen_graph): - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) # Get the initial roots - root0 = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - root1 = graph.get_root(to_label(graph, 1, 0, 0, 0, 1)) + root0 = cg.get_root(sv["a0"]) + root1 = cg.get_root(sv["a1"]) # Merge - result = graph.add_edges( + result = cg.add_edges( "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], + [sv["a0"], sv["a1"]], affinities=[0.3], ) new_root = result.new_root_ids[0] - return graph, root0, root1, new_root + return cg, root0, root1, new_root @pytest.mark.timeout(30) def test_is_latest_roots_after_merge(self, gen_graph): @@ -537,29 +454,19 @@ class TestGetNodeTimestampsExtended: """Tests for get_node_timestamps (lines 773-800).""" def _build_graph(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - ], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1)}, + edges=[("a0", "a1", 0.5)], ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - return graph @pytest.mark.timeout(30) def test_get_node_timestamps_return_numpy(self, gen_graph): """get_node_timestamps with return_numpy=True should return numpy array.""" - graph = self._build_graph(gen_graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - ts = graph.get_node_timestamps( + cg, sv = self._build_graph(gen_graph) + root = cg.get_root(sv["a0"]) + ts = cg.get_node_timestamps( np.array([root], dtype=np.uint64), return_numpy=True ) assert isinstance(ts, np.ndarray) @@ -568,9 +475,9 @@ def test_get_node_timestamps_return_numpy(self, gen_graph): @pytest.mark.timeout(30) def test_get_node_timestamps_return_list(self, gen_graph): """get_node_timestamps with return_numpy=False should return a list.""" - graph = self._build_graph(gen_graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - ts = graph.get_node_timestamps( + cg, sv = self._build_graph(gen_graph) + root = cg.get_root(sv["a0"]) + ts = cg.get_node_timestamps( np.array([root], dtype=np.uint64), return_numpy=False ) assert isinstance(ts, list) @@ -579,8 +486,8 @@ def test_get_node_timestamps_return_list(self, gen_graph): @pytest.mark.timeout(30) def test_get_node_timestamps_empty(self, gen_graph): """get_node_timestamps with nonexistent nodes should handle gracefully.""" - graph = self._build_graph(gen_graph) - ts = graph.get_node_timestamps( + cg, sv = self._build_graph(gen_graph) + ts = cg.get_node_timestamps( np.array([np.uint64(99999999)], dtype=np.uint64), return_numpy=True ) # Should either return empty or return a fallback timestamp @@ -605,18 +512,13 @@ class TestGetOperationIdsExtended: @pytest.mark.timeout(30) def test_get_operation_ids_no_ops(self, gen_graph): """get_operation_ids on a node with no operations returns empty dict.""" - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV()}, ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - result = graph.get_operation_ids([root]) + root = cg.get_root(sv["a0"]) + result = cg.get_operation_ids([root]) # No operations => root may not be in result, or have empty list if root in result: assert isinstance(result[root], list) @@ -626,138 +528,93 @@ class TestGetSingleLeafMultipleExtended: """Tests for get_single_leaf_multiple (lines 1044-1062).""" def _build_graph(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - timestamp=fake_ts, - ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - ], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, + edges=[("a0", "a1", 0.5), ("a0", "b", inf)], ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 3, [1, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - return graph @pytest.mark.timeout(30) def test_get_single_leaf_from_root(self, gen_graph): """get_single_leaf_multiple from a root (layer 4) should drill down to layer 1.""" - graph = self._build_graph(gen_graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - assert graph.get_chunk_layer(root) == 4 - leaves = graph.get_single_leaf_multiple(np.array([root], dtype=np.uint64)) + cg, sv = self._build_graph(gen_graph) + root = cg.get_root(sv["a0"]) + assert cg.get_chunk_layer(root) == 4 + leaves = cg.get_single_leaf_multiple(np.array([root], dtype=np.uint64)) assert len(leaves) == 1 - assert graph.get_chunk_layer(leaves[0]) == 1 + assert cg.get_chunk_layer(leaves[0]) == 1 @pytest.mark.timeout(30) def test_get_single_leaf_from_l2(self, gen_graph): """get_single_leaf_multiple from L2 node should return one of its SV children.""" - graph = self._build_graph(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - parent = graph.get_parents(np.array([sv], dtype=np.uint64), raw_only=True)[0] - assert graph.get_chunk_layer(parent) == 2 - leaves = graph.get_single_leaf_multiple(np.array([parent], dtype=np.uint64)) + cg, sv = self._build_graph(gen_graph) + parent = cg.get_parents(np.array([sv["a0"]], dtype=np.uint64), raw_only=True)[0] + assert cg.get_chunk_layer(parent) == 2 + leaves = cg.get_single_leaf_multiple(np.array([parent], dtype=np.uint64)) assert len(leaves) == 1 - assert graph.get_chunk_layer(leaves[0]) == 1 + assert cg.get_chunk_layer(leaves[0]) == 1 @pytest.mark.timeout(30) def test_get_single_leaf_multiple_nodes(self, gen_graph): """get_single_leaf_multiple with multiple node IDs should return one leaf each.""" - graph = self._build_graph(gen_graph) - sv_a = to_label(graph, 1, 0, 0, 0, 0) - sv_b = to_label(graph, 1, 1, 0, 0, 0) - parent_a = graph.get_parents(np.array([sv_a], dtype=np.uint64), raw_only=True)[ + cg, sv = self._build_graph(gen_graph) + parent_a = cg.get_parents(np.array([sv["a0"]], dtype=np.uint64), raw_only=True)[ 0 ] - parent_b = graph.get_parents(np.array([sv_b], dtype=np.uint64), raw_only=True)[ + parent_b = cg.get_parents(np.array([sv["b"]], dtype=np.uint64), raw_only=True)[ 0 ] - leaves = graph.get_single_leaf_multiple( + leaves = cg.get_single_leaf_multiple( np.array([parent_a, parent_b], dtype=np.uint64) ) assert len(leaves) == 2 for leaf in leaves: - assert graph.get_chunk_layer(leaf) == 1 + assert cg.get_chunk_layer(leaf) == 1 class TestGetL2ChildrenExtended: """Tests for get_l2children (lines 1079-1092).""" def _build_graph(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, + edges=[("a0", "a1", 0.5), ("a0", "b", inf)], ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - ], - timestamp=fake_ts, - ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 3, [1, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - return graph @pytest.mark.timeout(30) def test_get_l2children_from_root(self, gen_graph): """get_l2children from a root should return all L2 children.""" - graph = self._build_graph(gen_graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - l2_children = graph.get_l2children(np.array([root], dtype=np.uint64)) + cg, sv = self._build_graph(gen_graph) + root = cg.get_root(sv["a0"]) + l2_children = cg.get_l2children(np.array([root], dtype=np.uint64)) assert isinstance(l2_children, np.ndarray) assert len(l2_children) >= 2 for child in l2_children: - assert graph.get_chunk_layer(child) == 2 + assert cg.get_chunk_layer(child) == 2 @pytest.mark.timeout(30) def test_get_l2children_from_l3(self, gen_graph): """get_l2children from an L3 node should return L2 children.""" - graph = self._build_graph(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) + cg, sv = self._build_graph(gen_graph) # Get L3 parent - root = graph.get_root(sv, stop_layer=3) - assert graph.get_chunk_layer(root) == 3 - l2_children = graph.get_l2children(np.array([root], dtype=np.uint64)) + root = cg.get_root(sv["a0"], stop_layer=3) + assert cg.get_chunk_layer(root) == 3 + l2_children = cg.get_l2children(np.array([root], dtype=np.uint64)) assert len(l2_children) >= 1 for child in l2_children: - assert graph.get_chunk_layer(child) == 2 + assert cg.get_chunk_layer(child) == 2 @pytest.mark.timeout(30) def test_get_l2children_from_l2(self, gen_graph): """get_l2children from an L2 node drills down to its children, which are L1 - so no L2 children are found; result is empty.""" - graph = self._build_graph(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - parent = graph.get_parents(np.array([sv], dtype=np.uint64), raw_only=True)[0] - assert graph.get_chunk_layer(parent) == 2 - l2_children = graph.get_l2children(np.array([parent], dtype=np.uint64)) + cg, sv = self._build_graph(gen_graph) + parent = cg.get_parents(np.array([sv["a0"]], dtype=np.uint64), raw_only=True)[0] + assert cg.get_chunk_layer(parent) == 2 + l2_children = cg.get_l2children(np.array([parent], dtype=np.uint64)) # L2 nodes only have L1 (SV) children, so no L2 descendants found assert isinstance(l2_children, np.ndarray) assert len(l2_children) == 0 @@ -767,38 +624,21 @@ class TestGetChunkLayersExtended: """Tests for get_chunk_layers and related helpers (line 951-952, 946).""" def _build_graph(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - timestamp=fake_ts, - ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - ], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "b": SV(x=1)}, + edges=[("a0", "b", inf)], ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 3, [1, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - return graph @pytest.mark.timeout(30) def test_get_chunk_layers_multiple(self, gen_graph): """get_chunk_layers for nodes at different layers.""" - graph = self._build_graph(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - parent_l2 = graph.get_parents(np.array([sv], dtype=np.uint64), raw_only=True)[0] - root = graph.get_root(sv) - layers = graph.get_chunk_layers( - np.array([sv, parent_l2, root], dtype=np.uint64) + cg, sv = self._build_graph(gen_graph) + parent_l2 = cg.get_parents(np.array([sv["a0"]], dtype=np.uint64), raw_only=True)[0] + root = cg.get_root(sv["a0"]) + layers = cg.get_chunk_layers( + np.array([sv["a0"], parent_l2, root], dtype=np.uint64) ) assert layers[0] == 1 assert layers[1] == 2 @@ -807,28 +647,24 @@ def test_get_chunk_layers_multiple(self, gen_graph): @pytest.mark.timeout(30) def test_get_segment_id_limit(self, gen_graph): """get_segment_id_limit should return a valid limit.""" - graph = self._build_graph(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - limit = graph.get_segment_id_limit(sv) + cg, sv = self._build_graph(gen_graph) + limit = cg.get_segment_id_limit(sv["a0"]) assert limit > 0 @pytest.mark.timeout(30) def test_get_chunk_coordinates(self, gen_graph): """get_chunk_coordinates should return the chunk coordinates of a node.""" - graph = self._build_graph(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - coords = graph.get_chunk_coordinates(sv) + cg, sv = self._build_graph(gen_graph) + coords = cg.get_chunk_coordinates(sv["a0"]) assert len(coords) == 3 np.testing.assert_array_equal(coords, [0, 0, 0]) @pytest.mark.timeout(30) def test_get_chunk_layers_and_coordinates(self, gen_graph): """get_chunk_layers_and_coordinates returns layers and coords together.""" - graph = self._build_graph(gen_graph) - sv_a = to_label(graph, 1, 0, 0, 0, 0) - sv_b = to_label(graph, 1, 1, 0, 0, 0) - layers, coords = graph.get_chunk_layers_and_coordinates( - np.array([sv_a, sv_b], dtype=np.uint64) + cg, sv = self._build_graph(gen_graph) + layers, coords = cg.get_chunk_layers_and_coordinates( + np.array([sv["a0"], sv["b"]], dtype=np.uint64) ) assert len(layers) == 2 assert layers[0] == 1 @@ -840,46 +676,24 @@ class TestGetAtomicCrossEdgesExtended: """Tests for get_atomic_cross_edges (lines 315-336).""" def _build_graph(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - timestamp=fake_ts, - ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - ], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, + edges=[("a0", "a1", 0.5), ("a0", "b", inf)], ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 3, [1, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - return graph @pytest.mark.timeout(30) def test_get_atomic_cross_edges_multiple_l2(self, gen_graph): """get_atomic_cross_edges with multiple L2 IDs.""" - graph = self._build_graph(gen_graph) - sv_a = to_label(graph, 1, 0, 0, 0, 0) - sv_b = to_label(graph, 1, 1, 0, 0, 0) - parent_a = graph.get_parents(np.array([sv_a], dtype=np.uint64), raw_only=True)[ + cg, sv = self._build_graph(gen_graph) + parent_a = cg.get_parents(np.array([sv["a0"]], dtype=np.uint64), raw_only=True)[ 0 ] - parent_b = graph.get_parents(np.array([sv_b], dtype=np.uint64), raw_only=True)[ + parent_b = cg.get_parents(np.array([sv["b"]], dtype=np.uint64), raw_only=True)[ 0 ] - result = graph.get_atomic_cross_edges([parent_a, parent_b]) + result = cg.get_atomic_cross_edges([parent_a, parent_b]) assert parent_a in result assert parent_b in result # At least one should have cross edges @@ -889,19 +703,13 @@ def test_get_atomic_cross_edges_multiple_l2(self, gen_graph): @pytest.mark.timeout(30) def test_get_atomic_cross_edges_no_cross(self, gen_graph): """get_atomic_cross_edges for an L2 node with no cross edges.""" - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV()}, ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - sv = to_label(graph, 1, 0, 0, 0, 0) - parent = graph.get_parents(np.array([sv], dtype=np.uint64), raw_only=True)[0] - result = graph.get_atomic_cross_edges([parent]) + parent = cg.get_parents(np.array([sv["a0"]], dtype=np.uint64), raw_only=True)[0] + result = cg.get_atomic_cross_edges([parent]) assert parent in result assert isinstance(result[parent], dict) # No cross edges @@ -912,29 +720,18 @@ class TestGetAllParentsDictExtended: """Tests for get_all_parents_dict and get_all_parents_dict_multiple.""" def _build_graph(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - ], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1)}, + edges=[("a0", "a1", 0.5)], ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - return graph @pytest.mark.timeout(30) def test_get_all_parents_dict(self, gen_graph): """get_all_parents_dict returns a dict mapping layer -> parent.""" - graph = self._build_graph(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - d = graph.get_all_parents_dict(sv) + cg, sv = self._build_graph(gen_graph) + d = cg.get_all_parents_dict(sv["a0"]) assert isinstance(d, dict) # Should have entries for layers 2, 3, 4 assert 2 in d @@ -943,96 +740,82 @@ def test_get_all_parents_dict(self, gen_graph): @pytest.mark.timeout(30) def test_get_all_parents_dict_multiple(self, gen_graph): """get_all_parents_dict_multiple for multiple SVs.""" - graph = self._build_graph(gen_graph) - sv0 = to_label(graph, 1, 0, 0, 0, 0) - sv1 = to_label(graph, 1, 0, 0, 0, 1) - result = graph.get_all_parents_dict_multiple( - np.array([sv0, sv1], dtype=np.uint64) + cg, sv = self._build_graph(gen_graph) + result = cg.get_all_parents_dict_multiple( + np.array([sv["a0"], sv["a1"]], dtype=np.uint64) ) - assert sv0 in result - assert sv1 in result + assert sv["a0"] in result + assert sv["a1"] in result # Both should have parents at layer 2 - assert 2 in result[sv0] - assert 2 in result[sv1] + assert 2 in result[sv["a0"]] + assert 2 in result[sv["a1"]] class TestMiscMethods: """Tests for misc ChunkedGraph methods.""" def _build_graph(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - ], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1)}, + edges=[("a0", "a1", 0.5)], ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - return graph @pytest.mark.timeout(30) def test_get_serialized_info(self, gen_graph): """get_serialized_info should return a dict with graph_id.""" - graph = self._build_graph(gen_graph) - info = graph.get_serialized_info() + cg, sv = self._build_graph(gen_graph) + info = cg.get_serialized_info() assert isinstance(info, dict) assert "graph_id" in info @pytest.mark.timeout(30) def test_get_chunk_id(self, gen_graph): """get_chunk_id should return a valid chunk id.""" - graph = self._build_graph(gen_graph) - chunk_id = graph.get_chunk_id(layer=2, x=0, y=0, z=0) + cg, sv = self._build_graph(gen_graph) + chunk_id = cg.get_chunk_id(layer=2, x=0, y=0, z=0) assert chunk_id > 0 - assert graph.get_chunk_layer(chunk_id) == 2 + assert cg.get_chunk_layer(chunk_id) == 2 @pytest.mark.timeout(30) def test_get_node_id(self, gen_graph): """get_node_id should construct node IDs correctly.""" - graph = self._build_graph(gen_graph) - node_id = graph.get_node_id(np.uint64(1), layer=1, x=0, y=0, z=0) + cg, sv = self._build_graph(gen_graph) + node_id = cg.get_node_id(np.uint64(1), layer=1, x=0, y=0, z=0) assert node_id > 0 - assert graph.get_chunk_layer(node_id) == 1 + assert cg.get_chunk_layer(node_id) == 1 @pytest.mark.timeout(30) def test_get_segment_id(self, gen_graph): """get_segment_id should extract segment id from node id.""" - graph = self._build_graph(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 5) - seg_id = graph.get_segment_id(sv) + cg, sv = self._build_graph(gen_graph) + node = cg.get_node_id(np.uint64(5), layer=1, x=0, y=0, z=0) + seg_id = cg.get_segment_id(node) assert seg_id == 5 @pytest.mark.timeout(30) def test_get_parent_chunk_id(self, gen_graph): """get_parent_chunk_id should return the chunk id of the parent layer.""" - graph = self._build_graph(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - parent_chunk = graph.get_parent_chunk_id(sv) - assert graph.get_chunk_layer(parent_chunk) == 2 + cg, sv = self._build_graph(gen_graph) + parent_chunk = cg.get_parent_chunk_id(sv["a0"]) + assert cg.get_chunk_layer(parent_chunk) == 2 @pytest.mark.timeout(30) def test_get_children_chunk_ids(self, gen_graph): """get_children_chunk_ids should return chunk IDs one layer below.""" - graph = self._build_graph(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - root = graph.get_root(sv) + cg, sv = self._build_graph(gen_graph) + root = cg.get_root(sv["a0"]) # root is at layer 4; children chunks should be at layer 3 - children_chunks = graph.get_children_chunk_ids(root) + children_chunks = cg.get_children_chunk_ids(root) for cc in children_chunks: - assert graph.get_chunk_layer(cc) == 3 + assert cg.get_chunk_layer(cc) == 3 @pytest.mark.timeout(30) def test_get_cross_chunk_edges_empty(self, gen_graph): """get_cross_chunk_edges with empty node_ids should return empty dict.""" - graph = self._build_graph(gen_graph) - result = graph.get_cross_chunk_edges([], raw_only=True) + cg, sv = self._build_graph(gen_graph) + result = cg.get_cross_chunk_edges([], raw_only=True) assert isinstance(result, dict) assert len(result) == 0 @@ -1041,28 +824,22 @@ class TestIsLatestRootsAfterMerge: """Test is_latest_roots after a merge operation (lines 524-539, 689-701).""" def _build_and_merge(self, gen_graph): - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) - old_root0 = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - old_root1 = graph.get_root(to_label(graph, 1, 0, 0, 0, 1)) + old_root0 = cg.get_root(sv["a0"]) + old_root1 = cg.get_root(sv["a1"]) - result = graph.add_edges( + result = cg.add_edges( "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], + [sv["a0"], sv["a1"]], affinities=[0.3], ) new_root = result.new_root_ids[0] - return graph, old_root0, old_root1, new_root + return cg, old_root0, old_root1, new_root @pytest.mark.timeout(30) def test_is_latest_roots_after_merge(self, gen_graph): @@ -1080,53 +857,33 @@ class TestGetSubgraphNodesOnly: """Test get_subgraph with nodes_only=True (lines 602-613).""" def _build_graph(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, + edges=[("a0", "a1", 0.5), ("a0", "b", inf)], ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - ], - timestamp=fake_ts, - ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 3, [1, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - return graph @pytest.mark.timeout(30) def test_get_subgraph_nodes_only(self, gen_graph): """get_subgraph with nodes_only=True should return layer->node_ids dict.""" - graph = self._build_graph(gen_graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - result = graph.get_subgraph(root, nodes_only=True) + cg, sv = self._build_graph(gen_graph) + root = cg.get_root(sv["a0"]) + result = cg.get_subgraph(root, nodes_only=True) # Result should be a dict with layer 2 by default assert isinstance(result, dict) assert 2 in result l2_nodes = result[2] assert len(l2_nodes) >= 2 for node in l2_nodes: - assert graph.get_chunk_layer(node) == 2 + assert cg.get_chunk_layer(node) == 2 @pytest.mark.timeout(30) def test_get_subgraph_nodes_only_multiple_layers(self, gen_graph): """get_subgraph with nodes_only=True and return_layers=[2,3].""" - graph = self._build_graph(gen_graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - result = graph.get_subgraph(root, nodes_only=True, return_layers=[2, 3]) + cg, sv = self._build_graph(gen_graph) + root = cg.get_root(sv["a0"]) + result = cg.get_subgraph(root, nodes_only=True, return_layers=[2, 3]) assert isinstance(result, dict) # Should have entries for layer 2 and/or 3 assert 2 in result or 3 in result @@ -1136,39 +893,19 @@ class TestGetSubgraphEdgesOnly: """Test get_subgraph with edges_only=True.""" def _build_graph(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - timestamp=fake_ts, - ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - ], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, + edges=[("a0", "a1", 0.5), ("a0", "b", inf)], ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 3, [1, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - return graph @pytest.mark.timeout(30) def test_get_subgraph_edges_only(self, gen_graph): """get_subgraph with edges_only=True should return edges.""" - graph = self._build_graph(gen_graph) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - result = graph.get_subgraph(root, edges_only=True) + cg, sv = self._build_graph(gen_graph) + root = cg.get_root(sv["a0"]) + result = cg.get_subgraph(root, edges_only=True) # edges_only returns Edges from get_l2_agglomerations # It should be a tuple of Edges or similar iterable assert result is not None @@ -1184,28 +921,22 @@ class TestIsLatestRootsDetailed: def _build_and_merge(self, gen_graph): """Build graph with two disconnected SVs, merge them, return old and new roots.""" - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) - old_root0 = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - old_root1 = graph.get_root(to_label(graph, 1, 0, 0, 0, 1)) + old_root0 = cg.get_root(sv["a0"]) + old_root1 = cg.get_root(sv["a1"]) - result = graph.add_edges( + result = cg.add_edges( "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], + [sv["a0"], sv["a1"]], affinities=[0.3], ) new_root = result.new_root_ids[0] - return graph, old_root0, old_root1, new_root + return cg, old_root0, old_root1, new_root @pytest.mark.timeout(30) def test_is_latest_roots_correct(self, gen_graph): @@ -1242,50 +973,28 @@ class TestGetChunkCoordinatesMultiple: """Tests for get_chunk_coordinates_multiple with same/different layer assertions.""" def _build_graph(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, + edges=[("a0", "a1", 0.5), ("a0", "b", inf)], ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - ], - timestamp=fake_ts, - ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 3, [1, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - return graph @pytest.mark.timeout(30) def test_same_layer(self, gen_graph): """get_chunk_coordinates_multiple with L2 node IDs should return correct coordinates.""" - graph = self._build_graph(gen_graph) - sv_a = to_label(graph, 1, 0, 0, 0, 0) - sv_b = to_label(graph, 1, 1, 0, 0, 0) + cg, sv = self._build_graph(gen_graph) # Get L2 parents - parent_a = graph.get_parents(np.array([sv_a], dtype=np.uint64), raw_only=True)[ + parent_a = cg.get_parents(np.array([sv["a0"]], dtype=np.uint64), raw_only=True)[ 0 ] - parent_b = graph.get_parents(np.array([sv_b], dtype=np.uint64), raw_only=True)[ + parent_b = cg.get_parents(np.array([sv["b"]], dtype=np.uint64), raw_only=True)[ 0 ] - assert graph.get_chunk_layer(parent_a) == 2 - assert graph.get_chunk_layer(parent_b) == 2 + assert cg.get_chunk_layer(parent_a) == 2 + assert cg.get_chunk_layer(parent_b) == 2 - coords = graph.get_chunk_coordinates_multiple( + coords = cg.get_chunk_coordinates_multiple( np.array([parent_a, parent_b], dtype=np.uint64) ) assert coords.shape == (2, 3) @@ -1296,26 +1005,25 @@ def test_same_layer(self, gen_graph): @pytest.mark.timeout(30) def test_different_layers_raises(self, gen_graph): """get_chunk_coordinates_multiple with nodes at different layers should raise.""" - graph = self._build_graph(gen_graph) - sv_a = to_label(graph, 1, 0, 0, 0, 0) - parent_l2 = graph.get_parents(np.array([sv_a], dtype=np.uint64), raw_only=True)[ + cg, sv = self._build_graph(gen_graph) + parent_l2 = cg.get_parents(np.array([sv["a0"]], dtype=np.uint64), raw_only=True)[ 0 ] - root = graph.get_root(sv_a) + root = cg.get_root(sv["a0"]) - assert graph.get_chunk_layer(parent_l2) == 2 - assert graph.get_chunk_layer(root) == 4 + assert cg.get_chunk_layer(parent_l2) == 2 + assert cg.get_chunk_layer(root) == 4 with pytest.raises(AssertionError, match="must be same layer"): - graph.get_chunk_coordinates_multiple( + cg.get_chunk_coordinates_multiple( np.array([parent_l2, root], dtype=np.uint64) ) @pytest.mark.timeout(30) def test_empty_array(self, gen_graph): """get_chunk_coordinates_multiple with empty array should return empty result.""" - graph = self._build_graph(gen_graph) - coords = graph.get_chunk_coordinates_multiple(np.array([], dtype=np.uint64)) + cg, sv = self._build_graph(gen_graph) + coords = cg.get_chunk_coordinates_multiple(np.array([], dtype=np.uint64)) assert len(coords) == 0 @@ -1328,63 +1036,40 @@ class TestParentChunkIdMethods: """Tests for get_parent_chunk_id_multiple and get_parent_chunk_ids.""" def _build_graph(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), - ], - timestamp=fake_ts, - ) - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), - ], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, + edges=[("a0", "a1", 0.5), ("a0", "b", inf)], ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 3, [1, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - return graph @pytest.mark.timeout(30) def test_get_parent_chunk_id_multiple(self, gen_graph): """get_parent_chunk_id_multiple should return parent chunk IDs for all nodes.""" - graph = self._build_graph(gen_graph) - sv_a = to_label(graph, 1, 0, 0, 0, 0) - sv_b = to_label(graph, 1, 1, 0, 0, 0) + cg, sv = self._build_graph(gen_graph) # Get L2 parents - parent_a = graph.get_parents(np.array([sv_a], dtype=np.uint64), raw_only=True)[ + parent_a = cg.get_parents(np.array([sv["a0"]], dtype=np.uint64), raw_only=True)[ 0 ] - parent_b = graph.get_parents(np.array([sv_b], dtype=np.uint64), raw_only=True)[ + parent_b = cg.get_parents(np.array([sv["b"]], dtype=np.uint64), raw_only=True)[ 0 ] - parent_chunks = graph.get_parent_chunk_id_multiple( + parent_chunks = cg.get_parent_chunk_id_multiple( np.array([parent_a, parent_b], dtype=np.uint64) ) assert len(parent_chunks) == 2 for pc in parent_chunks: - assert graph.get_chunk_layer(pc) == 3 + assert cg.get_chunk_layer(pc) == 3 @pytest.mark.timeout(30) def test_get_parent_chunk_ids(self, gen_graph): """get_parent_chunk_ids should return all parent chunk IDs up the hierarchy.""" - graph = self._build_graph(gen_graph) - sv = to_label(graph, 1, 0, 0, 0, 0) - parent_chunk_ids = graph.get_parent_chunk_ids(sv) + cg, sv = self._build_graph(gen_graph) + parent_chunk_ids = cg.get_parent_chunk_ids(sv["a0"]) # Should have parent chunk IDs for layers 2, 3, 4 assert len(parent_chunk_ids) >= 2 - layers = [graph.get_chunk_layer(pc) for pc in parent_chunk_ids] + layers = [cg.get_chunk_layer(pc) for pc in parent_chunk_ids] # Layers should be ascending (from layer 2 up) for i in range(len(layers) - 1): assert layers[i] < layers[i + 1] @@ -1401,22 +1086,16 @@ class TestReadChunkEdges: @pytest.mark.timeout(30) def test_read_chunk_edges_returns_dict(self, gen_graph): """read_chunk_edges should return a dict (possibly empty for gs:// edges source).""" - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV()}, ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - sv = to_label(graph, 1, 0, 0, 0, 0) - parent = graph.get_parents(np.array([sv], dtype=np.uint64), raw_only=True)[0] + parent = cg.get_parents(np.array([sv["a0"]], dtype=np.uint64), raw_only=True)[0] # read_chunk_edges uses io.edges.get_chunk_edges which reads from GCS/file. # With gs:// edges source and no actual files, it should raise or return empty. try: - result = graph.read_chunk_edges(np.array([parent], dtype=np.uint64)) + result = cg.read_chunk_edges(np.array([parent], dtype=np.uint64)) assert isinstance(result, dict) except Exception: # Expected: GCS access will fail in test env @@ -1434,42 +1113,31 @@ class TestGetProofreadRootIds: @pytest.mark.timeout(30) def test_get_proofread_root_ids_no_ops(self, gen_graph): """get_proofread_root_ids with no operations should return empty arrays.""" - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV()}, ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - old_roots, new_roots = graph.get_proofread_root_ids() + old_roots, new_roots = cg.get_proofread_root_ids() assert len(old_roots) == 0 assert len(new_roots) == 0 @pytest.mark.timeout(30) def test_get_proofread_root_ids_after_merge(self, gen_graph): """get_proofread_root_ids after a merge should return the old and new roots.""" - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) - result = graph.add_edges( + result = cg.add_edges( "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], + [sv["a0"], sv["a1"]], affinities=[0.3], ) - old_roots, new_roots = graph.get_proofread_root_ids() + old_roots, new_roots = cg.get_proofread_root_ids() assert len(new_roots) >= 1 assert result.new_root_ids[0] in new_roots @@ -1484,34 +1152,26 @@ class TestRemoveEdgesShim: def _build_connected_graph(self, gen_graph): """Build a 2-layer graph with two connected SVs.""" - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - ], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, + edges=[("a0", "a1", 0.5)], ) - return graph @pytest.mark.timeout(30) def test_remove_edges_with_source_sink_ids(self, gen_graph): """Call remove_edges with source_ids/sink_ids (no atomic_edges) -- exercises shim.""" - graph = self._build_connected_graph(gen_graph) - sv0 = to_label(graph, 1, 0, 0, 0, 0) - sv1 = to_label(graph, 1, 0, 0, 0, 1) + cg, sv = self._build_connected_graph(gen_graph) + sv0 = sv["a0"] + sv1 = sv["a1"] # Verify they share a root before split - assert graph.get_root(sv0) == graph.get_root(sv1) + assert cg.get_root(sv0) == cg.get_root(sv1) # Use source_ids/sink_ids (the shim path) instead of atomic_edges - result = graph.remove_edges( + result = cg.remove_edges( "TestUser", source_ids=sv0, sink_ids=sv1, @@ -1521,17 +1181,17 @@ def test_remove_edges_with_source_sink_ids(self, gen_graph): assert len(result.new_root_ids) == 2 # After split, they should have different roots - assert graph.get_root(sv0) != graph.get_root(sv1) + assert cg.get_root(sv0) != cg.get_root(sv1) @pytest.mark.timeout(30) def test_remove_edges_shim_mismatched_lengths(self, gen_graph): """Shim path with mismatched source_ids/sink_ids lengths should raise.""" - graph = self._build_connected_graph(gen_graph) - sv0 = to_label(graph, 1, 0, 0, 0, 0) - sv1 = to_label(graph, 1, 0, 0, 0, 1) + cg, sv = self._build_connected_graph(gen_graph) + sv0 = sv["a0"] + sv1 = sv["a1"] with pytest.raises(PreconditionError, match="same number"): - graph.remove_edges( + cg.remove_edges( "TestUser", source_ids=[sv0, sv0], sink_ids=[sv1], @@ -1550,47 +1210,36 @@ class TestEarliestTimestamp: @pytest.mark.timeout(30) def test_get_earliest_timestamp_after_merge(self, gen_graph): """After creating a graph and performing a merge, get_earliest_timestamp should return a valid datetime.""" - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - ], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) # Perform a merge to generate operation logs - graph.add_edges( + cg.add_edges( "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], + [sv["a0"], sv["a1"]], affinities=[0.3], ) - ts = graph.get_earliest_timestamp() + ts = cg.get_earliest_timestamp() assert ts is not None assert isinstance(ts, datetime) @pytest.mark.timeout(30) def test_get_earliest_timestamp_stamped_by_root_build(self, gen_graph): """Building the root layer stamps earliest_ts strictly above every root's ts.""" - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV()}, ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) # no ops -> get_earliest_timestamp returns exactly the stamped boundary - boundary = graph.get_earliest_timestamp() - assert boundary == datetime.fromisoformat(graph.meta.custom_data["earliest_ts"]) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - [root_ts] = graph.get_node_timestamps(np.array([root]), return_numpy=False) + boundary = cg.get_earliest_timestamp() + assert boundary == datetime.fromisoformat(cg.meta.custom_data["earliest_ts"]) + root = cg.get_root(sv["a0"]) + [root_ts] = cg.get_node_timestamps(np.array([root]), return_numpy=False) assert root_ts < boundary def test_get_earliest_timestamp_falls_back_to_stamped(self, gen_graph): diff --git a/pychunkedgraph/tests/graph/test_lineage.py b/pychunkedgraph/tests/graph/test_lineage.py index 4a8dd3a01..26e7da41a 100644 --- a/pychunkedgraph/tests/graph/test_lineage.py +++ b/pychunkedgraph/tests/graph/test_lineage.py @@ -1,7 +1,6 @@ """Tests for pychunkedgraph.graph.lineage""" from datetime import datetime, timedelta, UTC -from math import inf import numpy as np import pytest @@ -18,35 +17,23 @@ ) from pychunkedgraph.graph import attributes -from ..helpers import create_chunk, to_label, fake_timestamp -from ...ingest.create.parent_layer import add_parent_chunk +from ..helpers import SV, build_graph class TestLineage: def _build_and_merge(self, gen_graph): """Build a graph with 2 isolated SVs, then merge them.""" - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - edges=[], - timestamp=fake_ts, - ) - - old_root_0 = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - old_root_1 = graph.get_root(to_label(graph, 1, 0, 0, 0, 1)) - - # Merge - result = graph.add_edges( - "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - affinities=[0.3], + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) + old_root_0 = cg.get_root(sv["a0"]) + old_root_1 = cg.get_root(sv["a1"]) + result = cg.add_edges("TestUser", [sv["a0"], sv["a1"]], affinities=[0.3]) new_root = result.new_root_ids[0] - return graph, old_root_0, old_root_1, new_root + return cg, old_root_0, old_root_1, new_root def test_get_latest_root_id_current(self, gen_graph): graph, _, _, new_root = self._build_and_merge(gen_graph) @@ -151,44 +138,20 @@ def _build_graph_with_two_merges(self, gen_graph): First merge SV0+SV1 -> root_A Then merge root_A+SV2 -> root_B """ - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - - fake_ts = fake_timestamp() - from ..helpers import create_chunk, to_label - - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - to_label(graph, 1, 0, 0, 0, 2), - ], - edges=[], - timestamp=fake_ts, - ) - - old_root_0 = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - old_root_1 = graph.get_root(to_label(graph, 1, 0, 0, 0, 1)) - old_root_2 = graph.get_root(to_label(graph, 1, 0, 0, 0, 2)) - - # First merge: SV0 + SV1 - result1 = graph.add_edges( - "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - affinities=[0.3], + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1), "a2": SV(seg=2)}, ) + old_root_0 = cg.get_root(sv["a0"]) + old_root_1 = cg.get_root(sv["a1"]) + old_root_2 = cg.get_root(sv["a2"]) + result1 = cg.add_edges("TestUser", [sv["a0"], sv["a1"]], affinities=[0.3]) mid_root = result1.new_root_ids[0] - - # Second merge: merged root + SV2 - result2 = graph.add_edges( - "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 2)], - affinities=[0.3], - ) + result2 = cg.add_edges("TestUser", [sv["a0"], sv["a2"]], affinities=[0.3]) final_root = result2.new_root_ids[0] - - return graph, old_root_0, old_root_1, old_root_2, mid_root, final_root + return cg, old_root_0, old_root_1, old_root_2, mid_root, final_root def test_future_root_ids_finds_chain(self, gen_graph): """get_future_root_ids from original root should find mid and final roots.""" @@ -222,29 +185,17 @@ class TestGetPastRootIdsTimestamps: def _build_and_merge(self, gen_graph): """Build a graph with 2 isolated SVs, then merge them.""" - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - - fake_ts = fake_timestamp() - from ..helpers import create_chunk, to_label - - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - edges=[], - timestamp=fake_ts, - ) - - old_root_0 = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - old_root_1 = graph.get_root(to_label(graph, 1, 0, 0, 0, 1)) - - result = graph.add_edges( - "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - affinities=[0.3], + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) + old_root_0 = cg.get_root(sv["a0"]) + old_root_1 = cg.get_root(sv["a1"]) + result = cg.add_edges("TestUser", [sv["a0"], sv["a1"]], affinities=[0.3]) new_root = result.new_root_ids[0] - return graph, old_root_0, old_root_1, new_root + return cg, old_root_0, old_root_1, new_root def test_past_root_ids_of_merged_root(self, gen_graph): """get_past_root_ids of the merged root should find old roots.""" @@ -274,29 +225,17 @@ class TestGetRootIdHistory: def _build_and_merge(self, gen_graph): """Build a graph with 2 isolated SVs, then merge them.""" - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - - fake_ts = fake_timestamp() - from ..helpers import create_chunk, to_label - - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - edges=[], - timestamp=fake_ts, - ) - - old_root_0 = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - old_root_1 = graph.get_root(to_label(graph, 1, 0, 0, 0, 1)) - - result = graph.add_edges( - "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - affinities=[0.3], + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) + old_root_0 = cg.get_root(sv["a0"]) + old_root_1 = cg.get_root(sv["a1"]) + result = cg.add_edges("TestUser", [sv["a0"], sv["a1"]], affinities=[0.3]) new_root = result.new_root_ids[0] - return graph, old_root_0, old_root_1, new_root + return cg, old_root_0, old_root_1, new_root def test_history_after_merge(self, gen_graph): """After merge, get_root_id_history should contain past and current root.""" @@ -343,43 +282,20 @@ def _build_graph_with_two_merges(self, gen_graph): First merge SV0+SV1 -> root_A Then merge root_A+SV2 -> root_B """ - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - - fake_ts = fake_timestamp() - - create_chunk( - graph, - vertices=[ - to_label(graph, 1, 0, 0, 0, 0), - to_label(graph, 1, 0, 0, 0, 1), - to_label(graph, 1, 0, 0, 0, 2), - ], - edges=[], - timestamp=fake_ts, - ) - - old_root_0 = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - old_root_1 = graph.get_root(to_label(graph, 1, 0, 0, 0, 1)) - old_root_2 = graph.get_root(to_label(graph, 1, 0, 0, 0, 2)) - - # First merge: SV0 + SV1 - result1 = graph.add_edges( - "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - affinities=[0.3], + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1), "a2": SV(seg=2)}, ) + old_root_0 = cg.get_root(sv["a0"]) + old_root_1 = cg.get_root(sv["a1"]) + old_root_2 = cg.get_root(sv["a2"]) + result1 = cg.add_edges("TestUser", [sv["a0"], sv["a1"]], affinities=[0.3]) mid_root = result1.new_root_ids[0] - - # Second merge: merged root + SV2 - result2 = graph.add_edges( - "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 2)], - affinities=[0.3], - ) + result2 = cg.add_edges("TestUser", [sv["a0"], sv["a2"]], affinities=[0.3]) final_root = result2.new_root_ids[0] - - return graph, old_root_0, old_root_1, old_root_2, mid_root, final_root + return cg, old_root_0, old_root_1, old_root_2, mid_root, final_root def test_history_contains_all_roots_from_old(self, gen_graph): """get_root_id_history from original root should contain all related roots.""" diff --git a/pychunkedgraph/tests/graph/test_operation.py b/pychunkedgraph/tests/graph/test_operation.py index 0a271d8d7..811f18b7a 100644 --- a/pychunkedgraph/tests/graph/test_operation.py +++ b/pychunkedgraph/tests/graph/test_operation.py @@ -10,7 +10,7 @@ import numpy as np import pytest -from ..helpers import create_chunk, to_label, fake_timestamp +from ..helpers import SV, build_graph from ...graph import attributes from ...graph.operation import ( GraphEditOperation, @@ -25,62 +25,38 @@ PostconditionError, SupervoxelSplitRequiredError, ) -from ...ingest.create.parent_layer import add_parent_chunk def _build_two_sv_disconnected(gen_graph): """2-layer graph, two disconnected SVs in the same chunk.""" - cg = gen_graph(n_layers=2, atomic_chunk_bounds=np.array([1, 1, 1])) - ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], - edges=[], - timestamp=ts, + return build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) - return cg, ts def _build_two_sv_connected(gen_graph): """2-layer graph, two connected SVs in the same chunk.""" - cg = gen_graph(n_layers=2, atomic_chunk_bounds=np.array([1, 1, 1])) - ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], - edges=[ - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5), - ], - timestamp=ts, + return build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, + edges=[("a0", "a1", 0.5)], ) - return cg, ts def _build_cross_chunk(gen_graph): """4-layer graph with cross-chunk edges suitable for MulticutOperation.""" - cg = gen_graph(n_layers=4) - ts = fake_timestamp() - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) - create_chunk( - cg, - vertices=[sv0, sv1], - edges=[ - (sv0, sv1, 0.5), - (sv0, to_label(cg, 1, 1, 0, 0, 0), inf), - ], - timestamp=ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, + edges=[("a0", "a1", 0.5), ("a0", "b", inf)], ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), sv0, inf)], - timestamp=ts, - ) - add_parent_chunk(cg, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(cg, 3, [1, 0, 0], n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], n_threads=1) - return cg, ts, sv0, sv1 + return cg, sv["a0"], sv["a1"] # =========================================================================== @@ -92,35 +68,22 @@ class TestOperationFromLogRecord: @pytest.fixture() def merged_graph(self, gen_graph): """Build a simple 2-chunk graph and perform a merge, returning (cg, operation_id).""" - cg = gen_graph(n_layers=3) - fake_ts = fake_timestamp() - - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a": SV(), "b": SV(x=1)}, + edges=[("a", "b", 0.5)], ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # Split first to get two separate roots split_result = cg.remove_edges( - "test_user", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 0), - mincut=False, + "test_user", source_ids=sv["a"], sink_ids=sv["b"], mincut=False ) # Now merge them back merge_result = cg.add_edges( "test_user", - atomic_edges=[[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0)]], + atomic_edges=[[sv["a"], sv["b"]]], source_coords=[0, 0, 0], sink_coords=[0, 0, 0], ) @@ -173,32 +136,19 @@ class TestOperationInversion: @pytest.fixture() def split_and_merge_ops(self, gen_graph): """Build graph, split, merge -- return (cg, merge_op_id, split_op_id).""" - cg = gen_graph(n_layers=3) - fake_ts = fake_timestamp() - - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a": SV(), "b": SV(x=1)}, + edges=[("a", "b", 0.5)], ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) split_result = cg.remove_edges( - "test_user", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 0), - mincut=False, + "test_user", source_ids=sv["a"], sink_ids=sv["b"], mincut=False ) merge_result = cg.add_edges( "test_user", - atomic_edges=[[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0)]], + atomic_edges=[[sv["a"], sv["b"]]], source_coords=[0, 0, 0], sink_coords=[0, 0, 0], ) @@ -231,29 +181,16 @@ class TestUndoRedoChainResolution: @pytest.fixture() def graph_with_undo(self, gen_graph): """Build graph, perform split, then undo -- return (cg, split_op_id, undo_result).""" - cg = gen_graph(n_layers=3) - fake_ts = fake_timestamp() - - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a": SV(), "b": SV(x=1)}, + edges=[("a", "b", 0.5)], ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # Split split_result = cg.remove_edges( - "test_user", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 0), - mincut=False, + "test_user", source_ids=sv["a"], sink_ids=sv["b"], mincut=False ) # Undo the split (= merge) undo_result = cg.undo_operation("test_user", split_result.operation_id) @@ -371,7 +308,7 @@ class TestFromLogRecordMulticutPath: def test_multicut_from_log_record(self, gen_graph): """A multicut operation's log, read back with multicut_as_split=False, should be reconstructed as MulticutOperation (lines 235-249).""" - cg, _, sv0, sv1 = _build_cross_chunk(gen_graph) + cg, sv0, sv1 = _build_cross_chunk(gen_graph) source_coords = [[0, 0, 0]] sink_coords = [[512, 0, 0]] try: @@ -404,9 +341,8 @@ class TestFromOperationId: @pytest.mark.timeout(30) def test_from_operation_id_merge(self, gen_graph): """from_operation_id should reconstruct a MergeOperation.""" - cg, _ = _build_two_sv_disconnected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) + cg, sv = _build_two_sv_disconnected(gen_graph) + sv0, sv1 = sv["a0"], sv["a1"] result = cg.add_edges("test_user", [sv0, sv1], affinities=[0.3]) op = GraphEditOperation.from_operation_id(cg, result.operation_id) assert isinstance(op, MergeOperation) @@ -416,9 +352,8 @@ def test_from_operation_id_merge(self, gen_graph): @pytest.mark.timeout(30) def test_from_operation_id_privileged(self, gen_graph): """from_operation_id with privileged_mode=True should propagate the flag (line 280).""" - cg, _ = _build_two_sv_disconnected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) + cg, sv = _build_two_sv_disconnected(gen_graph) + sv0, sv1 = sv["a0"], sv["a1"] result = cg.add_edges("test_user", [sv0, sv1], affinities=[0.3]) op = GraphEditOperation.from_operation_id( cg, result.operation_id, privileged_mode=True @@ -428,9 +363,8 @@ def test_from_operation_id_privileged(self, gen_graph): @pytest.mark.timeout(30) def test_from_operation_id_split(self, gen_graph): """from_operation_id should reconstruct a SplitOperation.""" - cg, _ = _build_two_sv_connected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) + cg, sv = _build_two_sv_connected(gen_graph) + sv0, sv1 = sv["a0"], sv["a1"] result = cg.remove_edges( "test_user", source_ids=sv0, sink_ids=sv1, mincut=False ) @@ -447,7 +381,7 @@ class TestMulticutInversion: @pytest.mark.timeout(30) def test_multicut_invert(self, gen_graph): """MulticutOperation.invert() -> MergeOperation with removed_edges as added_edges.""" - cg, _, sv0, sv1 = _build_cross_chunk(gen_graph) + cg, sv0, sv1 = _build_cross_chunk(gen_graph) mc_op = MulticutOperation( cg, user_id="test_user", @@ -472,8 +406,8 @@ class TestIDValidation: @pytest.mark.timeout(30) def test_merge_self_loop_raises(self, gen_graph): """added_edges where source == sink should raise PreconditionError (line 596).""" - cg, _ = _build_two_sv_disconnected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) + cg, sv = _build_two_sv_disconnected(gen_graph) + sv0 = sv["a0"] with pytest.raises(PreconditionError, match="self-loop"): MergeOperation( cg, @@ -486,8 +420,8 @@ def test_merge_self_loop_raises(self, gen_graph): @pytest.mark.timeout(30) def test_split_self_loop_raises(self, gen_graph): """removed_edges where source == sink should raise PreconditionError (line 733).""" - cg, _ = _build_two_sv_connected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) + cg, sv = _build_two_sv_connected(gen_graph) + sv0 = sv["a0"] with pytest.raises(PreconditionError, match="self-loop"): SplitOperation( cg, @@ -507,9 +441,8 @@ class TestEmptyCoordsAffinities: @pytest.mark.timeout(30) def test_empty_source_coords_becomes_none(self, gen_graph): """source_coords with size 0 should be stored as None (line 82).""" - cg, _ = _build_two_sv_disconnected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) + cg, sv = _build_two_sv_disconnected(gen_graph) + sv0, sv1 = sv["a0"], sv["a1"] op = MergeOperation( cg, user_id="test_user", @@ -523,9 +456,8 @@ def test_empty_source_coords_becomes_none(self, gen_graph): @pytest.mark.timeout(30) def test_empty_affinities_becomes_none(self, gen_graph): """affinities with size 0 should be stored as None (line 593).""" - cg, _ = _build_two_sv_disconnected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) + cg, sv = _build_two_sv_disconnected(gen_graph) + sv0, sv1 = sv["a0"], sv["a1"] op = MergeOperation( cg, user_id="test_user", @@ -546,18 +478,16 @@ class TestEditPreconditions: @pytest.mark.timeout(30) def test_merge_same_segment_raises(self, gen_graph): """Merging SVs already in the same root raises PreconditionError (line 618).""" - cg, _ = _build_two_sv_connected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) + cg, sv = _build_two_sv_connected(gen_graph) + sv0, sv1 = sv["a0"], sv["a1"] with pytest.raises(PreconditionError, match="different objects"): cg.add_edges("test_user", [sv0, sv1], affinities=[0.3]) @pytest.mark.timeout(30) def test_split_different_roots_raises(self, gen_graph): """Splitting SVs from different roots raises PreconditionError (line 765).""" - cg, _ = _build_two_sv_disconnected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) + cg, sv = _build_two_sv_disconnected(gen_graph) + sv0, sv1 = sv["a0"], sv["a1"] with pytest.raises(PreconditionError, match="same object"): cg.remove_edges("test_user", source_ids=sv0, sink_ids=sv1, mincut=False) @@ -570,24 +500,13 @@ class TestUndoRedoExecute: def _build_connected_cross_chunk(self, gen_graph): """Build a 3-layer graph with between-chunk edge -- suitable for split+undo.""" - cg = gen_graph(n_layers=3) - ts = fake_timestamp() - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 1, 0, 0, 0) - create_chunk( - cg, - vertices=[sv0], - edges=[(sv0, sv1, 0.5)], - timestamp=ts, - ) - create_chunk( - cg, - vertices=[sv1], - edges=[(sv1, sv0, 0.5)], - timestamp=ts, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a": SV(), "b": SV(x=1)}, + edges=[("a", "b", 0.5)], ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=ts, n_threads=1) - return cg, sv0, sv1 + return cg, sv["a"], sv["b"] @pytest.mark.timeout(60) def test_undo_split_restores_root(self, gen_graph): @@ -645,9 +564,8 @@ class TestUndoRedoInvert: @pytest.mark.timeout(60) def test_undo_invert_is_redo(self, gen_graph): """UndoOperation.invert() -> RedoOperation (line 1228).""" - cg, _ = _build_two_sv_disconnected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) + cg, sv = _build_two_sv_disconnected(gen_graph) + sv0, sv1 = sv["a0"], sv["a1"] merge_result = cg.add_edges("test_user", [sv0, sv1], affinities=[0.3]) undo_op = GraphEditOperation.undo_operation( @@ -661,9 +579,8 @@ def test_undo_invert_is_redo(self, gen_graph): @pytest.mark.timeout(60) def test_redo_invert_is_undo(self, gen_graph): """RedoOperation.invert() -> UndoOperation (line 1087).""" - cg, _ = _build_two_sv_disconnected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) + cg, sv = _build_two_sv_disconnected(gen_graph) + sv0, sv1 = sv["a0"], sv["a1"] merge_result = cg.add_edges("test_user", [sv0, sv1], affinities=[0.3]) redo_op = GraphEditOperation.redo_operation( @@ -684,9 +601,8 @@ class TestUndoRedoEdgeAttributes: @pytest.mark.timeout(60) def test_undo_merge_has_removed_edges(self, gen_graph): """Undoing a merge -> inverse is SplitOp -> undo should have removed_edges (line 1175).""" - cg, _ = _build_two_sv_disconnected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) + cg, sv = _build_two_sv_disconnected(gen_graph) + sv0, sv1 = sv["a0"], sv["a1"] merge_result = cg.add_edges("test_user", [sv0, sv1], affinities=[0.3]) undo_op = GraphEditOperation.undo_operation( @@ -698,9 +614,8 @@ def test_undo_merge_has_removed_edges(self, gen_graph): @pytest.mark.timeout(60) def test_undo_split_has_added_edges(self, gen_graph): """Undoing a split -> inverse is MergeOp -> undo should have added_edges (line 1173).""" - cg, _ = _build_two_sv_connected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) + cg, sv = _build_two_sv_connected(gen_graph) + sv0, sv1 = sv["a0"], sv["a1"] split_result = cg.remove_edges( "test_user", source_ids=sv0, sink_ids=sv1, mincut=False ) @@ -714,9 +629,8 @@ def test_undo_split_has_added_edges(self, gen_graph): @pytest.mark.timeout(60) def test_redo_merge_has_added_edges(self, gen_graph): """RedoOperation for a merge should have added_edges (line 1040-1041).""" - cg, _ = _build_two_sv_disconnected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) + cg, sv = _build_two_sv_disconnected(gen_graph) + sv0, sv1 = sv["a0"], sv["a1"] merge_result = cg.add_edges("test_user", [sv0, sv1], affinities=[0.3]) redo_op = GraphEditOperation.redo_operation( @@ -729,9 +643,8 @@ def test_redo_merge_has_added_edges(self, gen_graph): @pytest.mark.timeout(60) def test_redo_split_has_removed_edges(self, gen_graph): """RedoOperation for a split should have removed_edges (line 1042-1043).""" - cg, _ = _build_two_sv_connected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) + cg, sv = _build_two_sv_connected(gen_graph) + sv0, sv1 = sv["a0"], sv["a1"] split_result = cg.remove_edges( "test_user", source_ids=sv0, sink_ids=sv1, mincut=False ) @@ -752,23 +665,13 @@ class TestUndoRedoLogRecordTypes: def _build_and_split(self, gen_graph): """Build a cross-chunk graph and split it -- suitable for undo/redo.""" - cg = gen_graph(n_layers=3) - ts = fake_timestamp() - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 1, 0, 0, 0) - create_chunk( - cg, - vertices=[sv0], - edges=[(sv0, sv1, 0.5)], - timestamp=ts, - ) - create_chunk( - cg, - vertices=[sv1], - edges=[(sv1, sv0, 0.5)], - timestamp=ts, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a": SV(), "b": SV(x=1)}, + edges=[("a", "b", 0.5)], ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=ts, n_threads=1) + sv0, sv1 = sv["a"], sv["b"] split_result = cg.remove_edges( "test_user", source_ids=sv0, sink_ids=sv1, mincut=False ) @@ -806,9 +709,8 @@ class TestExecuteErrorHandling: @pytest.mark.timeout(30) def test_execute_precondition_error_clears_cache(self, gen_graph): """Trigger PreconditionError during merge (same-segment merge) and verify cache is cleared.""" - cg, _ = _build_two_sv_connected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) + cg, sv = _build_two_sv_connected(gen_graph) + sv0, sv1 = sv["a0"], sv["a1"] # Merging already-connected SVs raises PreconditionError with pytest.raises(PreconditionError, match="different objects"): @@ -822,9 +724,8 @@ def test_execute_postcondition_error_clears_cache(self, gen_graph): """PostconditionError during execute should also clear cache (lines 463-465).""" from unittest.mock import patch - cg, _ = _build_two_sv_disconnected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) + cg, sv = _build_two_sv_disconnected(gen_graph) + sv0, sv1 = sv["a0"], sv["a1"] # Mock _apply to raise PostconditionError with patch.object( @@ -843,9 +744,8 @@ def test_execute_assertion_error_clears_cache(self, gen_graph): """AssertionError/RuntimeError during execute should also clear cache (lines 466-468).""" from unittest.mock import patch - cg, _ = _build_two_sv_disconnected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) + cg, sv = _build_two_sv_disconnected(gen_graph) + sv0, sv1 = sv["a0"], sv["a1"] # Mock _apply to raise RuntimeError with patch.object( @@ -865,24 +765,13 @@ class TestUndoEdgeValidation: def _build_connected_cross_chunk(self, gen_graph): """Build a 3-layer graph with between-chunk edge suitable for split+undo.""" - cg = gen_graph(n_layers=3) - ts = fake_timestamp() - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 1, 0, 0, 0) - create_chunk( - cg, - vertices=[sv0], - edges=[(sv0, sv1, 0.5)], - timestamp=ts, - ) - create_chunk( - cg, - vertices=[sv1], - edges=[(sv1, sv0, 0.5)], - timestamp=ts, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a": SV(), "b": SV(x=1)}, + edges=[("a", "b", 0.5)], ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=ts, n_threads=1) - return cg, sv0, sv1 + return cg, sv["a"], sv["b"] @pytest.mark.timeout(60) def test_undo_split_restores_edges(self, gen_graph): @@ -908,9 +797,8 @@ def test_undo_split_restores_edges(self, gen_graph): @pytest.mark.timeout(60) def test_undo_merge_via_undo_operation_class(self, gen_graph): """UndoOperation on a merge constructs with inverse being SplitOperation.""" - cg, _ = _build_two_sv_disconnected(gen_graph) - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) + cg, sv = _build_two_sv_disconnected(gen_graph) + sv0, sv1 = sv["a0"], sv["a1"] # Merge merge_result = cg.add_edges("test_user", [sv0, sv1], affinities=[0.3]) diff --git a/pychunkedgraph/tests/graph/test_segmenthistory.py b/pychunkedgraph/tests/graph/test_segmenthistory.py index 6c2ac5340..aa2ca659d 100644 --- a/pychunkedgraph/tests/graph/test_segmenthistory.py +++ b/pychunkedgraph/tests/graph/test_segmenthistory.py @@ -14,30 +14,20 @@ from pychunkedgraph.graph import attributes -from ..helpers import create_chunk, to_label, fake_timestamp -from ...ingest.create.parent_layer import add_parent_chunk +from ..helpers import SV, build_graph class TestSegmentHistory: def _build_and_merge(self, gen_graph): - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - edges=[], - timestamp=fake_ts, - ) - - result = graph.add_edges( - "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - affinities=[0.3], + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) + result = cg.add_edges("TestUser", [sv["a0"], sv["a1"]], affinities=[0.3]) new_root = result.new_root_ids[0] - return graph, new_root + return cg, new_root def test_init(self, gen_graph): graph, new_root = self._build_and_merge(gen_graph) @@ -324,39 +314,27 @@ def test_removed_edges_on_merge_raises(self): class TestGetAllLogEntries: def test_empty_graph(self, gen_graph): """Create graph with no operations. get_all_log_entries should return empty list.""" - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - # Create a chunk with vertices but perform no edits - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV()}, ) - entries = get_all_log_entries(graph) + entries = get_all_log_entries(cg) assert isinstance(entries, list) assert len(entries) == 0 def test_basic(self, gen_graph): - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - edges=[], - timestamp=fake_ts, - ) - graph.add_edges( - "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - affinities=[0.3], + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) + cg.add_edges("TestUser", [sv["a0"], sv["a1"]], affinities=[0.3]) # get_all_log_entries iterates range(get_max_operation_id()) which # may not include the actual operation ID; verify it doesn't crash - entries = get_all_log_entries(graph) + entries = get_all_log_entries(cg) assert isinstance(entries, list) # If entries exist, verify LogEntry API works for entry in entries: @@ -370,24 +348,21 @@ class TestMergeLog: """Tests for SegmentHistory.merge_log() method (lines 245-268).""" def _build_and_merge(self, gen_graph): - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) - result = graph.add_edges( + result = cg.add_edges( "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], + [sv["a0"], sv["a1"]], affinities=[0.3], source_coords=[0, 0, 0], sink_coords=[1, 1, 1], ) new_root = result.new_root_ids[0] - return graph, new_root + return cg, new_root def test_merge_log_with_root(self, gen_graph): """merge_log(root_id=...) should return merge_edges and merge_edge_coords.""" @@ -424,22 +399,15 @@ class TestPastOperationIdsExtended: """Tests for SegmentHistory.past_operation_ids() (lines 270-292).""" def _build_and_merge(self, gen_graph): - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - edges=[], - timestamp=fake_ts, - ) - result = graph.add_edges( - "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - affinities=[0.3], + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) + result = cg.add_edges("TestUser", [sv["a0"], sv["a1"]], affinities=[0.3]) new_root = result.new_root_ids[0] - return graph, new_root + return cg, new_root def test_past_operation_ids_without_root(self, gen_graph): """past_operation_ids() without root_id iterates all root_ids.""" @@ -465,22 +433,15 @@ class TestPastFutureIdMappingExtended: """More thorough tests for past_future_id_mapping (lines 315-368).""" def _build_and_merge(self, gen_graph): - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - edges=[], - timestamp=fake_ts, - ) - result = graph.add_edges( - "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - affinities=[0.3], + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) + result = cg.add_edges("TestUser", [sv["a0"], sv["a1"]], affinities=[0.3]) new_root = result.new_root_ids[0] - return graph, new_root + return cg, new_root def test_past_future_id_mapping_without_root(self, gen_graph): """past_future_id_mapping() without root_id iterates all root_ids.""" @@ -509,19 +470,16 @@ class TestMergeSplitHistory: """Tests involving merge followed by split to cover more branches.""" def _build_merge_and_split(self, gen_graph): - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) # Merge - merge_result = graph.add_edges( + merge_result = cg.add_edges( "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], + [sv["a0"], sv["a1"]], affinities=[0.3], source_coords=[0, 0, 0], sink_coords=[1, 1, 1], @@ -529,14 +487,14 @@ def _build_merge_and_split(self, gen_graph): merge_root = merge_result.new_root_ids[0] # Split - split_result = graph.remove_edges( + split_result = cg.remove_edges( "TestUser", - source_ids=to_label(graph, 1, 0, 0, 0, 0), - sink_ids=to_label(graph, 1, 0, 0, 0, 1), + source_ids=sv["a0"], + sink_ids=sv["a1"], mincut=False, ) split_roots = split_result.new_root_ids - return graph, merge_root, split_roots + return cg, merge_root, split_roots def test_change_log_summary_with_split(self, gen_graph): """change_log_summary after merge+split should show both operations.""" @@ -593,17 +551,14 @@ def test_past_future_id_mapping_after_split(self, gen_graph): def test_collect_edited_sv_ids_no_edits(self, gen_graph): """collect_edited_sv_ids returns empty array when no edits exist for a root.""" - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV()}, ) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - sh = SegmentHistory(graph, root) + root = cg.get_root(sv["a0"]) + sh = SegmentHistory(cg, root) sv_ids = sh.collect_edited_sv_ids(root_id=root) assert isinstance(sv_ids, np.ndarray) assert sv_ids.dtype == np.uint64 @@ -611,17 +566,14 @@ def test_collect_edited_sv_ids_no_edits(self, gen_graph): def test_change_log_summary_no_operations(self, gen_graph): """change_log_summary with no operations should show zero splits/merges.""" - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV()}, ) - root = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - sh = SegmentHistory(graph, root) + root = cg.get_root(sv["a0"]) + sh = SegmentHistory(cg, root) summary = sh.change_log_summary(root_id=root) assert isinstance(summary, dict) assert summary["n_splits"] == 0 diff --git a/pychunkedgraph/tests/graph/test_sv_lookup_main.py b/pychunkedgraph/tests/graph/test_sv_lookup_main.py index e8be1d5af..1dd9becbe 100644 --- a/pychunkedgraph/tests/graph/test_sv_lookup_main.py +++ b/pychunkedgraph/tests/graph/test_sv_lookup_main.py @@ -17,8 +17,7 @@ from pychunkedgraph.graph.sv_lookup import resolve_supervoxels_at_coords from pychunkedgraph.graph.sv_lookup import main as sv_lookup_main -from ..helpers import create_chunk, to_label, fake_timestamp -from ...ingest.create.parent_layer import add_parent_chunk +from ..helpers import SV, build_graph UTC = timezone.utc @@ -61,39 +60,25 @@ def _build_two_sv_graph(gen_graph): and sv2 is in chunk (1,0,0). graph.meta._ws_cv is a _SliceableCV seeded with these SVs at coordinates (0,0,0), (1,0,0), (2,0,0) respectively. """ - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - - sv0 = to_label(graph, 1, 0, 0, 0, 0) - sv1 = to_label(graph, 1, 0, 0, 0, 1) - sv2 = to_label(graph, 1, 1, 0, 0, 0) - - create_chunk( - graph, - vertices=[sv0, sv1], - edges=[(sv0, sv1, 0.5), (sv0, sv2, inf)], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, + edges=[("a0", "a1", 0.5), ("a0", "b", inf)], ) - create_chunk( - graph, - vertices=[sv2], - edges=[(sv2, sv0, inf)], - timestamp=fake_ts, - ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) + sv0, sv1, sv2 = sv["a0"], sv["a1"], sv["b"] - root = graph.get_root(sv0) - assert graph.get_root(sv1) == root - assert graph.get_root(sv2) == root + root = cg.get_root(sv0) + assert cg.get_root(sv1) == root + assert cg.get_root(sv2) == root cv = _SliceableWS(shape=(8, 8, 8)) cv.set_voxel(0, 0, 0, sv0) cv.set_voxel(1, 0, 0, sv1) cv.set_voxel(2, 0, 0, sv2) - graph.meta.ws_ts_scale = lambda mip=0: cv + cg.meta.ws_ts_scale = lambda mip=0: cv - return graph, sv0, sv1, sv2, root + return cg, sv0, sv1, sv2, root def _build_two_root_graph(gen_graph): @@ -106,40 +91,26 @@ def _build_two_root_graph(gen_graph): Returns (graph, sv0, sv1, sv2, root_a, root_b) with sv0, sv2 under root_a and sv1 under root_b. graph.meta._ws_cv is seeded with the three SVs. """ - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - - sv0 = to_label(graph, 1, 0, 0, 0, 0) - sv1 = to_label(graph, 1, 0, 0, 0, 1) - sv2 = to_label(graph, 1, 1, 0, 0, 0) - - create_chunk( - graph, - vertices=[sv0, sv1], - edges=[(sv0, sv2, inf)], - timestamp=fake_ts, - ) - create_chunk( - graph, - vertices=[sv2], - edges=[(sv2, sv0, inf)], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, + edges=[("a0", "b", inf)], ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) + sv0, sv1, sv2 = sv["a0"], sv["a1"], sv["b"] - root_a = graph.get_root(sv0) - root_b = graph.get_root(sv1) + root_a = cg.get_root(sv0) + root_b = cg.get_root(sv1) assert root_a != root_b - assert graph.get_root(sv2) == root_a + assert cg.get_root(sv2) == root_a cv = _SliceableWS(shape=(8, 8, 8)) cv.set_voxel(0, 0, 0, sv0) cv.set_voxel(1, 0, 0, sv1) cv.set_voxel(2, 0, 0, sv2) - graph.meta.ws_ts_scale = lambda mip=0: cv + cg.meta.ws_ts_scale = lambda mip=0: cv - return graph, sv0, sv1, sv2, root_a, root_b + return cg, sv0, sv1, sv2, root_a, root_b class TestResolveSupervoxelsAtCoords: From 8d6d727489050a76e995899c586abea9a9d9518b Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Sun, 28 Jun 2026 23:20:08 +0000 Subject: [PATCH 08/12] test: eliminate hardcoded node coordinates from the suite MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reference every node through readable SV coordinates — build_graph for setup, the new label(cg, SV, layer) for construction/encoding tests — instead of positional to_label tuples. Adds an assert_graph_unchanged context manager for rejected-edit atomicity; drops the dead query fixture. Co-Authored-By: Claude --- pychunkedgraph/tests/conftest.py | 43 -- .../tests/graph/test_chunks_hierarchy.py | 4 +- .../tests/graph/test_chunks_utils.py | 22 +- .../tests/graph/test_edges_utils.py | 10 +- .../tests/graph/test_edits_extended.py | 25 +- .../tests/graph/test_graph_build.py | 202 +++---- .../tests/graph/test_graph_queries.py | 205 +++---- pychunkedgraph/tests/graph/test_history.py | 38 +- pychunkedgraph/tests/graph/test_locks.py | 194 ++----- pychunkedgraph/tests/graph/test_merge.py | 500 +++++------------- .../tests/graph/test_merge_split.py | 39 +- pychunkedgraph/tests/graph/test_mincut.py | 267 +++------- pychunkedgraph/tests/graph/test_misc.py | 129 ++--- .../tests/graph/test_node_conversion.py | 1 - pychunkedgraph/tests/graph/test_root_lock.py | 39 +- pychunkedgraph/tests/graph/test_split.py | 487 +++++++---------- .../tests/graph/test_stale_edges.py | 266 +++------- pychunkedgraph/tests/graph/test_undo_redo.py | 42 +- .../tests/graph/test_utils_id_helpers.py | 6 +- pychunkedgraph/tests/helpers.py | 17 + .../tests/ingest/test_ingest_cross_edges.py | 56 +- .../tests/ingest/test_ingest_parent_layer.py | 20 +- pychunkedgraph/tests/meshing/test_setup.py | 4 +- 23 files changed, 924 insertions(+), 1692 deletions(-) diff --git a/pychunkedgraph/tests/conftest.py b/pychunkedgraph/tests/conftest.py index 9d6dd6fdb..4f6e3e9b6 100644 --- a/pychunkedgraph/tests/conftest.py +++ b/pychunkedgraph/tests/conftest.py @@ -16,14 +16,11 @@ from ..ingest.utils import bootstrap from ..graph.edges import Edges from ..graph.chunkedgraph import ChunkedGraph -from ..ingest.create.parent_layer import add_parent_chunk from .helpers import ( CloudVolumeMock, TensorStoreMock, mock_ws_info, - create_chunk, - to_label, get_layer_chunk_bounds, ) from .hbase_mock_server import start_hbase_mock_server @@ -204,46 +201,6 @@ def fin(): return partial(_cgraph, request) -@pytest.fixture(scope="function") -def gen_graph_simplequerytest(request, gen_graph): - """ - ┌─────┬─────┬─────┐ - │ A¹ │ B¹ │ C¹ │ - │ 1 │ 3━2━┿━━4 │ - │ │ │ │ - └─────┴─────┴─────┘ - """ - from math import inf - - graph = gen_graph(n_layers=4) - - # Chunk A - create_chunk(graph, vertices=[to_label(graph, 1, 0, 0, 0, 0)], edges=[]) - - # Chunk B - create_chunk( - graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 1)], - edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 1), 0.5), - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 2, 0, 0, 0), inf), - ], - ) - - # Chunk C - create_chunk( - graph, - vertices=[to_label(graph, 1, 2, 0, 0, 0)], - edges=[(to_label(graph, 1, 2, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf)], - ) - - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 3, [1, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - - return graph - - @pytest.fixture(scope="session") def sv_data(): test_data_dir = "pychunkedgraph/tests/data" diff --git a/pychunkedgraph/tests/graph/test_chunks_hierarchy.py b/pychunkedgraph/tests/graph/test_chunks_hierarchy.py index 2b63bf84b..6e14dd7b2 100644 --- a/pychunkedgraph/tests/graph/test_chunks_hierarchy.py +++ b/pychunkedgraph/tests/graph/test_chunks_hierarchy.py @@ -5,7 +5,7 @@ from pychunkedgraph.graph.chunks import hierarchy from pychunkedgraph.graph.chunks import utils as chunk_utils -from ..helpers import to_label +from ..helpers import SV, label class TestGetChildrenChunkCoords: @@ -20,7 +20,7 @@ def test_basic(self, gen_graph): class TestGetChildrenChunkIds: def test_layer_1_returns_empty(self, gen_graph): graph = gen_graph(n_layers=4) - node_id = to_label(graph, 1, 0, 0, 0, 1) + node_id = label(graph, SV(seg=1)) result = hierarchy.get_children_chunk_ids(graph.meta, node_id) assert len(result) == 0 diff --git a/pychunkedgraph/tests/graph/test_chunks_utils.py b/pychunkedgraph/tests/graph/test_chunks_utils.py index 610c5b090..4e85f3537 100644 --- a/pychunkedgraph/tests/graph/test_chunks_utils.py +++ b/pychunkedgraph/tests/graph/test_chunks_utils.py @@ -5,13 +5,13 @@ from pychunkedgraph.graph.chunks import utils as chunk_utils +from ..helpers import SV, label + class TestGetChunkLayer: def test_basic(self, gen_graph): graph = gen_graph(n_layers=4) - from ..helpers import to_label - - node_id = to_label(graph, 1, 0, 0, 0, 1) + node_id = label(graph, SV(seg=1)) assert chunk_utils.get_chunk_layer(graph.meta, node_id) == 1 def test_higher_layer(self, gen_graph): @@ -23,11 +23,9 @@ def test_higher_layer(self, gen_graph): class TestGetChunkLayers: def test_multiple(self, gen_graph): graph = gen_graph(n_layers=4) - from ..helpers import to_label - ids = [ - to_label(graph, 1, 0, 0, 0, 1), - to_label(graph, 1, 1, 0, 0, 2), + label(graph, SV(seg=1)), + label(graph, SV(x=1, seg=2)), ] layers = chunk_utils.get_chunk_layers(graph.meta, ids) np.testing.assert_array_equal(layers, [1, 1]) @@ -61,9 +59,7 @@ def test_empty(self, gen_graph): class TestGetChunkId: def test_from_node_id(self, gen_graph): graph = gen_graph(n_layers=4) - from ..helpers import to_label - - node_id = to_label(graph, 1, 2, 3, 1, 5) + node_id = label(graph, SV(x=2, y=3, z=1, seg=5)) chunk_id = chunk_utils.get_chunk_id(graph.meta, node_id=node_id) coords = chunk_utils.get_chunk_coordinates(graph.meta, chunk_id) np.testing.assert_array_equal(coords, [2, 3, 1]) @@ -96,12 +92,10 @@ def test_basic(self, gen_graph): class TestGetChunkIdsFromNodeIds: def test_basic(self, gen_graph): graph = gen_graph(n_layers=4) - from ..helpers import to_label - ids = np.array( [ - to_label(graph, 1, 0, 0, 0, 1), - to_label(graph, 1, 1, 0, 0, 2), + label(graph, SV(seg=1)), + label(graph, SV(x=1, seg=2)), ], dtype=np.uint64, ) diff --git a/pychunkedgraph/tests/graph/test_edges_utils.py b/pychunkedgraph/tests/graph/test_edges_utils.py index 2e01898e6..998095bf3 100644 --- a/pychunkedgraph/tests/graph/test_edges_utils.py +++ b/pychunkedgraph/tests/graph/test_edges_utils.py @@ -11,7 +11,7 @@ ) from pychunkedgraph.graph import basetypes -from ..helpers import to_label +from ..helpers import SV, label class TestConcatenateChunkEdges: @@ -81,16 +81,16 @@ def test_empty(self, gen_graph): def test_same_chunk(self, gen_graph): graph = gen_graph(n_layers=4) - sv1 = to_label(graph, 1, 0, 0, 0, 1) - sv2 = to_label(graph, 1, 0, 0, 0, 2) + sv1 = label(graph, SV(seg=1)) + sv2 = label(graph, SV(seg=2)) edges = np.array([[sv1, sv2]], dtype=basetypes.NODE_ID) result = get_cross_chunk_edges_layer(graph.meta, edges) assert result[0] == 1 # same chunk -> layer 1 def test_adjacent_chunks(self, gen_graph): graph = gen_graph(n_layers=4) - sv1 = to_label(graph, 1, 0, 0, 0, 1) - sv2 = to_label(graph, 1, 1, 0, 0, 1) + sv1 = label(graph, SV(seg=1)) + sv2 = label(graph, SV(x=1, seg=1)) edges = np.array([[sv1, sv2]], dtype=basetypes.NODE_ID) result = get_cross_chunk_edges_layer(graph.meta, edges) assert result[0] >= 2 # different chunks -> higher layer diff --git a/pychunkedgraph/tests/graph/test_edits_extended.py b/pychunkedgraph/tests/graph/test_edits_extended.py index 576cb657f..5cfc997e7 100644 --- a/pychunkedgraph/tests/graph/test_edits_extended.py +++ b/pychunkedgraph/tests/graph/test_edits_extended.py @@ -8,8 +8,7 @@ from pychunkedgraph.graph.edits import flip_ids from pychunkedgraph.graph import basetypes -from ..helpers import create_chunk, to_label, fake_timestamp -from ...ingest.create.parent_layer import add_parent_chunk +from ..helpers import SV, build_graph class TestFlipIds: @@ -33,22 +32,14 @@ class TestInitOldHierarchy: def test_basic(self, gen_graph): from pychunkedgraph.graph.edits import _init_old_hierarchy - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - ], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1)}, + edges=[("a0", "a1", 0.5)], ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - sv = to_label(graph, 1, 0, 0, 0, 0) - l2_parent = graph.get_parent(sv) - result = _init_old_hierarchy(graph, np.array([l2_parent], dtype=np.uint64)) + l2_parent = cg.get_parent(sv["a0"]) + result = _init_old_hierarchy(cg, np.array([l2_parent], dtype=np.uint64)) assert l2_parent in result assert 2 in result[l2_parent] diff --git a/pychunkedgraph/tests/graph/test_graph_build.py b/pychunkedgraph/tests/graph/test_graph_build.py index f53cbede4..c052c6603 100644 --- a/pychunkedgraph/tests/graph/test_graph_build.py +++ b/pychunkedgraph/tests/graph/test_graph_build.py @@ -3,7 +3,7 @@ import numpy as np import pytest -from ..helpers import create_chunk, to_label, fake_timestamp +from ..helpers import SV, label, create_chunk, fake_timestamp from ...graph import attributes, basetypes, serializers from ...ingest.create.parent_layer import add_parent_chunk @@ -22,22 +22,22 @@ def test_build_single_node(self, gen_graph): cg = gen_graph(n_layers=2) # Add Chunk A - create_chunk(cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)]) + create_chunk(cg, vertices=[label(cg, SV())]) res = cg.client.read_all_rows() res.consume_all() - assert serializers.serialize_uint64(to_label(cg, 1, 0, 0, 0, 0)) in res.rows - parent = cg.get_parent(to_label(cg, 1, 0, 0, 0, 0)) - assert parent == to_label(cg, 2, 0, 0, 0, 1) + assert serializers.serialize_uint64(label(cg, SV())) in res.rows + parent = cg.get_parent(label(cg, SV())) + assert parent == label(cg, SV(seg=1), layer=2) # Check for the one Level 2 node that should have been created. - assert serializers.serialize_uint64(to_label(cg, 2, 0, 0, 0, 1)) in res.rows + assert serializers.serialize_uint64(label(cg, SV(seg=1), layer=2)) in res.rows atomic_cross_edge_d = cg.get_atomic_cross_edges( - np.array([to_label(cg, 2, 0, 0, 0, 1)], dtype=basetypes.NODE_ID) + np.array([label(cg, SV(seg=1), layer=2)], dtype=basetypes.NODE_ID) ) attr = attributes.Hierarchy.Child - row = res.rows[serializers.serialize_uint64(to_label(cg, 2, 0, 0, 0, 1))].cells[ + row = res.rows[serializers.serialize_uint64(label(cg, SV(seg=1), layer=2))].cells[ "0" ] children = attr.deserialize(row[attr.key][0].value) @@ -45,7 +45,7 @@ def test_build_single_node(self, gen_graph): for aces in atomic_cross_edge_d.values(): assert len(aces) == 0 - assert len(children) == 1 and children[0] == to_label(cg, 1, 0, 0, 0, 0) + assert len(children) == 1 and children[0] == label(cg, SV()) # Make sure there are not any more entries in the table # include counters, meta and version rows assert len(res.rows) == 1 + 1 + 2 + 1 + 1 @@ -66,29 +66,29 @@ def test_build_single_edge(self, gen_graph): # Add Chunk A create_chunk( cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5)], + vertices=[label(cg, SV()), label(cg, SV(seg=1))], + edges=[(label(cg, SV()), label(cg, SV(seg=1)), 0.5)], ) res = cg.client.read_all_rows() res.consume_all() - assert serializers.serialize_uint64(to_label(cg, 1, 0, 0, 0, 0)) in res.rows - parent = cg.get_parent(to_label(cg, 1, 0, 0, 0, 0)) - assert parent == to_label(cg, 2, 0, 0, 0, 1) + assert serializers.serialize_uint64(label(cg, SV())) in res.rows + parent = cg.get_parent(label(cg, SV())) + assert parent == label(cg, SV(seg=1), layer=2) - assert serializers.serialize_uint64(to_label(cg, 1, 0, 0, 0, 1)) in res.rows - parent = cg.get_parent(to_label(cg, 1, 0, 0, 0, 1)) - assert parent == to_label(cg, 2, 0, 0, 0, 1) + assert serializers.serialize_uint64(label(cg, SV(seg=1))) in res.rows + parent = cg.get_parent(label(cg, SV(seg=1))) + assert parent == label(cg, SV(seg=1), layer=2) # Check for the one Level 2 node that should have been created. - assert serializers.serialize_uint64(to_label(cg, 2, 0, 0, 0, 1)) in res.rows + assert serializers.serialize_uint64(label(cg, SV(seg=1), layer=2)) in res.rows atomic_cross_edge_d = cg.get_atomic_cross_edges( - np.array([to_label(cg, 2, 0, 0, 0, 1)], dtype=basetypes.NODE_ID) + np.array([label(cg, SV(seg=1), layer=2)], dtype=basetypes.NODE_ID) ) attr = attributes.Hierarchy.Child - row = res.rows[serializers.serialize_uint64(to_label(cg, 2, 0, 0, 0, 1))].cells[ + row = res.rows[serializers.serialize_uint64(label(cg, SV(seg=1), layer=2))].cells[ "0" ] children = attr.deserialize(row[attr.key][0].value) @@ -97,8 +97,8 @@ def test_build_single_edge(self, gen_graph): assert len(aces) == 0 assert ( len(children) == 2 - and to_label(cg, 1, 0, 0, 0, 0) in children - and to_label(cg, 1, 0, 0, 0, 1) in children + and label(cg, SV()) in children + and label(cg, SV(seg=1)) in children ) # Make sure there are not any more entries in the table @@ -122,91 +122,91 @@ def test_build_single_across_edge(self, gen_graph): # Chunk A create_chunk( cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), inf)], + vertices=[label(cg, SV())], + edges=[(label(cg, SV()), label(cg, SV(x=1)), inf)], ) # Chunk B create_chunk( cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), inf)], + vertices=[label(cg, SV(x=1))], + edges=[(label(cg, SV(x=1)), label(cg, SV()), inf)], ) add_parent_chunk(cg, 3, [0, 0, 0], n_threads=1) res = cg.client.read_all_rows() res.consume_all() - assert serializers.serialize_uint64(to_label(cg, 1, 0, 0, 0, 0)) in res.rows - parent = cg.get_parent(to_label(cg, 1, 0, 0, 0, 0)) - assert parent == to_label(cg, 2, 0, 0, 0, 1) + assert serializers.serialize_uint64(label(cg, SV())) in res.rows + parent = cg.get_parent(label(cg, SV())) + assert parent == label(cg, SV(seg=1), layer=2) - assert serializers.serialize_uint64(to_label(cg, 1, 1, 0, 0, 0)) in res.rows - parent = cg.get_parent(to_label(cg, 1, 1, 0, 0, 0)) - assert parent == to_label(cg, 2, 1, 0, 0, 1) + assert serializers.serialize_uint64(label(cg, SV(x=1))) in res.rows + parent = cg.get_parent(label(cg, SV(x=1))) + assert parent == label(cg, SV(x=1, seg=1), layer=2) # Check for the two Level 2 nodes that should have been created. Since Level 2 has the same # dimensions as Level 1, we also expect them to be in different chunks - # to_label(cg, 2, 0, 0, 0, 1) - assert serializers.serialize_uint64(to_label(cg, 2, 0, 0, 0, 1)) in res.rows + # label(cg, SV(seg=1), layer=2) + assert serializers.serialize_uint64(label(cg, SV(seg=1), layer=2)) in res.rows atomic_cross_edge_d = cg.get_atomic_cross_edges( - np.array([to_label(cg, 2, 0, 0, 0, 1)], dtype=basetypes.NODE_ID) + np.array([label(cg, SV(seg=1), layer=2)], dtype=basetypes.NODE_ID) ) atomic_cross_edge_d = atomic_cross_edge_d[ - np.uint64(to_label(cg, 2, 0, 0, 0, 1)) + np.uint64(label(cg, SV(seg=1), layer=2)) ] attr = attributes.Hierarchy.Child - row = res.rows[serializers.serialize_uint64(to_label(cg, 2, 0, 0, 0, 1))].cells[ + row = res.rows[serializers.serialize_uint64(label(cg, SV(seg=1), layer=2))].cells[ "0" ] children = attr.deserialize(row[attr.key][0].value) test_ace = np.array( - [to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0)], + [label(cg, SV()), label(cg, SV(x=1))], dtype=np.uint64, ) assert len(atomic_cross_edge_d[2]) == 1 assert test_ace in atomic_cross_edge_d[2] - assert len(children) == 1 and to_label(cg, 1, 0, 0, 0, 0) in children + assert len(children) == 1 and label(cg, SV()) in children - # to_label(cg, 2, 1, 0, 0, 1) - assert serializers.serialize_uint64(to_label(cg, 2, 1, 0, 0, 1)) in res.rows + # label(cg, SV(x=1, seg=1), layer=2) + assert serializers.serialize_uint64(label(cg, SV(x=1, seg=1), layer=2)) in res.rows atomic_cross_edge_d = cg.get_atomic_cross_edges( - np.array([to_label(cg, 2, 1, 0, 0, 1)], dtype=basetypes.NODE_ID) + np.array([label(cg, SV(x=1, seg=1), layer=2)], dtype=basetypes.NODE_ID) ) atomic_cross_edge_d = atomic_cross_edge_d[ - np.uint64(to_label(cg, 2, 1, 0, 0, 1)) + np.uint64(label(cg, SV(x=1, seg=1), layer=2)) ] attr = attributes.Hierarchy.Child - row = res.rows[serializers.serialize_uint64(to_label(cg, 2, 1, 0, 0, 1))].cells[ + row = res.rows[serializers.serialize_uint64(label(cg, SV(x=1, seg=1), layer=2))].cells[ "0" ] children = attr.deserialize(row[attr.key][0].value) test_ace = np.array( - [to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0)], + [label(cg, SV(x=1)), label(cg, SV())], dtype=np.uint64, ) assert len(atomic_cross_edge_d[2]) == 1 assert test_ace in atomic_cross_edge_d[2] - assert len(children) == 1 and to_label(cg, 1, 1, 0, 0, 0) in children + assert len(children) == 1 and label(cg, SV(x=1)) in children # Check for the one Level 3 node that should have been created. This one combines the two # connected components of Level 2 - # to_label(cg, 3, 0, 0, 0, 1) - assert serializers.serialize_uint64(to_label(cg, 3, 0, 0, 0, 1)) in res.rows + # label(cg, SV(seg=1), layer=3) + assert serializers.serialize_uint64(label(cg, SV(seg=1), layer=3)) in res.rows attr = attributes.Hierarchy.Child - row = res.rows[serializers.serialize_uint64(to_label(cg, 3, 0, 0, 0, 1))].cells[ + row = res.rows[serializers.serialize_uint64(label(cg, SV(seg=1), layer=3))].cells[ "0" ] children = attr.deserialize(row[attr.key][0].value) assert ( len(children) == 2 - and to_label(cg, 2, 0, 0, 0, 1) in children - and to_label(cg, 2, 1, 0, 0, 1) in children + and label(cg, SV(seg=1), layer=2) in children + and label(cg, SV(x=1, seg=1), layer=2) in children ) # Make sure there are not any more entries in the table @@ -230,88 +230,88 @@ def test_build_single_edge_and_single_across_edge(self, gen_graph): # Chunk A create_chunk( cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], + vertices=[label(cg, SV()), label(cg, SV(seg=1))], edges=[ - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5), - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), inf), + (label(cg, SV()), label(cg, SV(seg=1)), 0.5), + (label(cg, SV()), label(cg, SV(x=1)), inf), ], ) # Chunk B create_chunk( cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), inf)], + vertices=[label(cg, SV(x=1))], + edges=[(label(cg, SV(x=1)), label(cg, SV()), inf)], ) add_parent_chunk(cg, 3, np.array([0, 0, 0]), n_threads=1) res = cg.client.read_all_rows() res.consume_all() - assert serializers.serialize_uint64(to_label(cg, 1, 0, 0, 0, 0)) in res.rows - parent = cg.get_parent(to_label(cg, 1, 0, 0, 0, 0)) - assert parent == to_label(cg, 2, 0, 0, 0, 1) + assert serializers.serialize_uint64(label(cg, SV())) in res.rows + parent = cg.get_parent(label(cg, SV())) + assert parent == label(cg, SV(seg=1), layer=2) - # to_label(cg, 1, 0, 0, 0, 1) - assert serializers.serialize_uint64(to_label(cg, 1, 0, 0, 0, 1)) in res.rows - parent = cg.get_parent(to_label(cg, 1, 0, 0, 0, 1)) - assert parent == to_label(cg, 2, 0, 0, 0, 1) + # label(cg, SV(seg=1)) + assert serializers.serialize_uint64(label(cg, SV(seg=1))) in res.rows + parent = cg.get_parent(label(cg, SV(seg=1))) + assert parent == label(cg, SV(seg=1), layer=2) - # to_label(cg, 1, 1, 0, 0, 0) - assert serializers.serialize_uint64(to_label(cg, 1, 1, 0, 0, 0)) in res.rows - parent = cg.get_parent(to_label(cg, 1, 1, 0, 0, 0)) - assert parent == to_label(cg, 2, 1, 0, 0, 1) + # label(cg, SV(x=1)) + assert serializers.serialize_uint64(label(cg, SV(x=1))) in res.rows + parent = cg.get_parent(label(cg, SV(x=1))) + assert parent == label(cg, SV(x=1, seg=1), layer=2) # Check for the two Level 2 nodes that should have been created. Since Level 2 has the same # dimensions as Level 1, we also expect them to be in different chunks - # to_label(cg, 2, 0, 0, 0, 1) - assert serializers.serialize_uint64(to_label(cg, 2, 0, 0, 0, 1)) in res.rows - row = res.rows[serializers.serialize_uint64(to_label(cg, 2, 0, 0, 0, 1))].cells[ + # label(cg, SV(seg=1), layer=2) + assert serializers.serialize_uint64(label(cg, SV(seg=1), layer=2)) in res.rows + row = res.rows[serializers.serialize_uint64(label(cg, SV(seg=1), layer=2))].cells[ "0" ] - atomic_cross_edge_d = cg.get_atomic_cross_edges([to_label(cg, 2, 0, 0, 0, 1)]) + atomic_cross_edge_d = cg.get_atomic_cross_edges([label(cg, SV(seg=1), layer=2)]) atomic_cross_edge_d = atomic_cross_edge_d[ - np.uint64(to_label(cg, 2, 0, 0, 0, 1)) + np.uint64(label(cg, SV(seg=1), layer=2)) ] column = attributes.Hierarchy.Child children = column.deserialize(row[column.key][0].value) test_ace = np.array( - [to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0)], + [label(cg, SV()), label(cg, SV(x=1))], dtype=np.uint64, ) assert len(atomic_cross_edge_d[2]) == 1 assert test_ace in atomic_cross_edge_d[2] assert ( len(children) == 2 - and to_label(cg, 1, 0, 0, 0, 0) in children - and to_label(cg, 1, 0, 0, 0, 1) in children + and label(cg, SV()) in children + and label(cg, SV(seg=1)) in children ) - # to_label(cg, 2, 1, 0, 0, 1) - assert serializers.serialize_uint64(to_label(cg, 2, 1, 0, 0, 1)) in res.rows - row = res.rows[serializers.serialize_uint64(to_label(cg, 2, 1, 0, 0, 1))].cells[ + # label(cg, SV(x=1, seg=1), layer=2) + assert serializers.serialize_uint64(label(cg, SV(x=1, seg=1), layer=2)) in res.rows + row = res.rows[serializers.serialize_uint64(label(cg, SV(x=1, seg=1), layer=2))].cells[ "0" ] - atomic_cross_edge_d = cg.get_atomic_cross_edges([to_label(cg, 2, 1, 0, 0, 1)]) + atomic_cross_edge_d = cg.get_atomic_cross_edges([label(cg, SV(x=1, seg=1), layer=2)]) atomic_cross_edge_d = atomic_cross_edge_d[ - np.uint64(to_label(cg, 2, 1, 0, 0, 1)) + np.uint64(label(cg, SV(x=1, seg=1), layer=2)) ] children = column.deserialize(row[column.key][0].value) test_ace = np.array( - [to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0)], + [label(cg, SV(x=1)), label(cg, SV())], dtype=np.uint64, ) assert len(atomic_cross_edge_d[2]) == 1 assert test_ace in atomic_cross_edge_d[2] - assert len(children) == 1 and to_label(cg, 1, 1, 0, 0, 0) in children + assert len(children) == 1 and label(cg, SV(x=1)) in children # Check for the one Level 3 node that should have been created. This one combines the two # connected components of Level 2 - # to_label(cg, 3, 0, 0, 0, 1) - assert serializers.serialize_uint64(to_label(cg, 3, 0, 0, 0, 1)) in res.rows - row = res.rows[serializers.serialize_uint64(to_label(cg, 3, 0, 0, 0, 1))].cells[ + # label(cg, SV(seg=1), layer=3) + assert serializers.serialize_uint64(label(cg, SV(seg=1), layer=3)) in res.rows + row = res.rows[serializers.serialize_uint64(label(cg, SV(seg=1), layer=3))].cells[ "0" ] column = attributes.Hierarchy.Child @@ -319,8 +319,8 @@ def test_build_single_edge_and_single_across_edge(self, gen_graph): assert ( len(children) == 2 - and to_label(cg, 2, 0, 0, 0, 1) in children - and to_label(cg, 2, 1, 0, 0, 1) in children + and label(cg, SV(seg=1), layer=2) in children + and label(cg, SV(x=1, seg=1), layer=2) in children ) # Make sure there are not any more entries in the table @@ -342,10 +342,10 @@ def test_build_big_graph(self, gen_graph): cg = gen_graph(n_layers=5, atomic_chunk_bounds=atomic_chunk_bounds) # Preparation: Build Chunk A - create_chunk(cg, vertices=[to_label(cg, 1, 0, 0, 0, 0)], edges=[]) + create_chunk(cg, vertices=[label(cg, SV())], edges=[]) # Preparation: Build Chunk Z - create_chunk(cg, vertices=[to_label(cg, 1, 7, 7, 7, 0)], edges=[]) + create_chunk(cg, vertices=[label(cg, SV(x=7, y=7, z=7))], edges=[]) add_parent_chunk(cg, 3, [0, 0, 0], n_threads=1) add_parent_chunk(cg, 3, [3, 3, 3], n_threads=1) @@ -355,10 +355,10 @@ def test_build_big_graph(self, gen_graph): res = cg.client.read_all_rows() res.consume_all() - assert serializers.serialize_uint64(to_label(cg, 1, 0, 0, 0, 0)) in res.rows - assert serializers.serialize_uint64(to_label(cg, 1, 7, 7, 7, 0)) in res.rows - assert serializers.serialize_uint64(to_label(cg, 5, 0, 0, 0, 1)) in res.rows - assert serializers.serialize_uint64(to_label(cg, 5, 0, 0, 0, 2)) in res.rows + assert serializers.serialize_uint64(label(cg, SV())) in res.rows + assert serializers.serialize_uint64(label(cg, SV(x=7, y=7, z=7))) in res.rows + assert serializers.serialize_uint64(label(cg, SV(seg=1), layer=5)) in res.rows + assert serializers.serialize_uint64(label(cg, SV(seg=2), layer=5)) in res.rows @pytest.mark.timeout(30) def test_double_chunk_creation(self, gen_graph): @@ -378,7 +378,7 @@ def test_double_chunk_creation(self, gen_graph): fake_ts = fake_timestamp() create_chunk( cg, - vertices=[to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2)], + vertices=[label(cg, SV(seg=1)), label(cg, SV(seg=2))], edges=[], timestamp=fake_ts, ) @@ -386,7 +386,7 @@ def test_double_chunk_creation(self, gen_graph): # Preparation: Build Chunk B create_chunk( cg, - vertices=[to_label(cg, 1, 1, 0, 0, 1)], + vertices=[label(cg, SV(x=1, seg=1))], edges=[], timestamp=fake_ts, ) @@ -418,14 +418,14 @@ def test_double_chunk_creation(self, gen_graph): assert len(cg.range_read_chunk(cg.get_chunk_id(layer=3, x=0, y=0, z=0))) == 0 assert len(cg.range_read_chunk(cg.get_chunk_id(layer=4, x=0, y=0, z=0))) == 6 - assert cg.get_chunk_layer(cg.get_root(to_label(cg, 1, 0, 0, 0, 1))) == 4 - assert cg.get_chunk_layer(cg.get_root(to_label(cg, 1, 0, 0, 0, 2))) == 4 - assert cg.get_chunk_layer(cg.get_root(to_label(cg, 1, 1, 0, 0, 1))) == 4 + assert cg.get_chunk_layer(cg.get_root(label(cg, SV(seg=1)))) == 4 + assert cg.get_chunk_layer(cg.get_root(label(cg, SV(seg=2)))) == 4 + assert cg.get_chunk_layer(cg.get_root(label(cg, SV(x=1, seg=1)))) == 4 root_seg_ids = [ - cg.get_segment_id(cg.get_root(to_label(cg, 1, 0, 0, 0, 1))), - cg.get_segment_id(cg.get_root(to_label(cg, 1, 0, 0, 0, 2))), - cg.get_segment_id(cg.get_root(to_label(cg, 1, 1, 0, 0, 1))), + cg.get_segment_id(cg.get_root(label(cg, SV(seg=1)))), + cg.get_segment_id(cg.get_root(label(cg, SV(seg=2)))), + cg.get_segment_id(cg.get_root(label(cg, SV(x=1, seg=1)))), ] assert 4 in root_seg_ids diff --git a/pychunkedgraph/tests/graph/test_graph_queries.py b/pychunkedgraph/tests/graph/test_graph_queries.py index 1dccfb092..a9089ee76 100644 --- a/pychunkedgraph/tests/graph/test_graph_queries.py +++ b/pychunkedgraph/tests/graph/test_graph_queries.py @@ -3,7 +3,7 @@ import numpy as np import pytest -from ..helpers import create_chunk, to_label +from ..helpers import SV, build_graph class TestGraphSimpleQueries: @@ -15,39 +15,52 @@ class TestGraphSimpleQueries: └─────┴─────┴─────┘ 4: 1 2 0 0 0 ─── 2 2 0 0 1 ─── 3 1 0 0 1 ─┘ """ + def _build_graph(self, gen_graph): + return build_graph( + gen_graph, + n_layers=4, + supervoxels={ + "a0": SV(), + "b0": SV(x=1), + "b1": SV(x=1, seg=1), + "c0": SV(x=2), + }, + edges=[("b0", "b1", 0.5), ("b0", "c0", inf)], + ) + @pytest.mark.timeout(30) - def test_get_parent_and_children(self, gen_graph_simplequerytest): - cg = gen_graph_simplequerytest + def test_get_parent_and_children(self, gen_graph): + cg, sv = self._build_graph(gen_graph) - children10000 = cg.get_children(to_label(cg, 1, 0, 0, 0, 0)) - children11000 = cg.get_children(to_label(cg, 1, 1, 0, 0, 0)) - children11001 = cg.get_children(to_label(cg, 1, 1, 0, 0, 1)) - children12000 = cg.get_children(to_label(cg, 1, 2, 0, 0, 0)) + children10000 = cg.get_children(sv["a0"]) + children11000 = cg.get_children(sv["b0"]) + children11001 = cg.get_children(sv["b1"]) + children12000 = cg.get_children(sv["c0"]) - parent10000 = cg.get_parent(to_label(cg, 1, 0, 0, 0, 0)) - parent11000 = cg.get_parent(to_label(cg, 1, 1, 0, 0, 0)) - parent11001 = cg.get_parent(to_label(cg, 1, 1, 0, 0, 1)) - parent12000 = cg.get_parent(to_label(cg, 1, 2, 0, 0, 0)) + parent10000 = cg.get_parent(sv["a0"]) + parent11000 = cg.get_parent(sv["b0"]) + parent11001 = cg.get_parent(sv["b1"]) + parent12000 = cg.get_parent(sv["c0"]) - children20001 = cg.get_children(to_label(cg, 2, 0, 0, 0, 1)) - children21001 = cg.get_children(to_label(cg, 2, 1, 0, 0, 1)) - children22001 = cg.get_children(to_label(cg, 2, 2, 0, 0, 1)) + children20001 = cg.get_children(cg.get_node_id(np.uint64(1), layer=2, x=0, y=0, z=0)) + children21001 = cg.get_children(cg.get_node_id(np.uint64(1), layer=2, x=1, y=0, z=0)) + children22001 = cg.get_children(cg.get_node_id(np.uint64(1), layer=2, x=2, y=0, z=0)) - parent20001 = cg.get_parent(to_label(cg, 2, 0, 0, 0, 1)) - parent21001 = cg.get_parent(to_label(cg, 2, 1, 0, 0, 1)) - parent22001 = cg.get_parent(to_label(cg, 2, 2, 0, 0, 1)) + parent20001 = cg.get_parent(cg.get_node_id(np.uint64(1), layer=2, x=0, y=0, z=0)) + parent21001 = cg.get_parent(cg.get_node_id(np.uint64(1), layer=2, x=1, y=0, z=0)) + parent22001 = cg.get_parent(cg.get_node_id(np.uint64(1), layer=2, x=2, y=0, z=0)) - children30001 = cg.get_children(to_label(cg, 3, 0, 0, 0, 1)) - children31001 = cg.get_children(to_label(cg, 3, 1, 0, 0, 1)) + children30001 = cg.get_children(cg.get_node_id(np.uint64(1), layer=3, x=0, y=0, z=0)) + children31001 = cg.get_children(cg.get_node_id(np.uint64(1), layer=3, x=1, y=0, z=0)) - parent30001 = cg.get_parent(to_label(cg, 3, 0, 0, 0, 1)) - parent31001 = cg.get_parent(to_label(cg, 3, 1, 0, 0, 1)) + parent30001 = cg.get_parent(cg.get_node_id(np.uint64(1), layer=3, x=0, y=0, z=0)) + parent31001 = cg.get_parent(cg.get_node_id(np.uint64(1), layer=3, x=1, y=0, z=0)) - children40001 = cg.get_children(to_label(cg, 4, 0, 0, 0, 1)) - children40002 = cg.get_children(to_label(cg, 4, 0, 0, 0, 2)) + children40001 = cg.get_children(cg.get_node_id(np.uint64(1), layer=4, x=0, y=0, z=0)) + children40002 = cg.get_children(cg.get_node_id(np.uint64(2), layer=4, x=0, y=0, z=0)) - parent40001 = cg.get_parent(to_label(cg, 4, 0, 0, 0, 1)) - parent40002 = cg.get_parent(to_label(cg, 4, 0, 0, 0, 2)) + parent40001 = cg.get_parent(cg.get_node_id(np.uint64(1), layer=4, x=0, y=0, z=0)) + parent40002 = cg.get_parent(cg.get_node_id(np.uint64(2), layer=4, x=0, y=0, z=0)) # (non-existing) Children of L1 assert np.array_equal(children10000, []) is True @@ -56,38 +69,38 @@ def test_get_parent_and_children(self, gen_graph_simplequerytest): assert np.array_equal(children12000, []) is True # Parent of L1 - assert parent10000 == to_label(cg, 2, 0, 0, 0, 1) - assert parent11000 == to_label(cg, 2, 1, 0, 0, 1) - assert parent11001 == to_label(cg, 2, 1, 0, 0, 1) - assert parent12000 == to_label(cg, 2, 2, 0, 0, 1) + assert parent10000 == cg.get_node_id(np.uint64(1), layer=2, x=0, y=0, z=0) + assert parent11000 == cg.get_node_id(np.uint64(1), layer=2, x=1, y=0, z=0) + assert parent11001 == cg.get_node_id(np.uint64(1), layer=2, x=1, y=0, z=0) + assert parent12000 == cg.get_node_id(np.uint64(1), layer=2, x=2, y=0, z=0) # Children of L2 - assert len(children20001) == 1 and to_label(cg, 1, 0, 0, 0, 0) in children20001 + assert len(children20001) == 1 and sv["a0"] in children20001 assert ( len(children21001) == 2 - and to_label(cg, 1, 1, 0, 0, 0) in children21001 - and to_label(cg, 1, 1, 0, 0, 1) in children21001 + and sv["b0"] in children21001 + and sv["b1"] in children21001 ) - assert len(children22001) == 1 and to_label(cg, 1, 2, 0, 0, 0) in children22001 + assert len(children22001) == 1 and sv["c0"] in children22001 # Parent of L2 - assert parent20001 == to_label(cg, 4, 0, 0, 0, 1) - assert parent21001 == to_label(cg, 3, 0, 0, 0, 1) - assert parent22001 == to_label(cg, 3, 1, 0, 0, 1) + assert parent20001 == cg.get_node_id(np.uint64(1), layer=4, x=0, y=0, z=0) + assert parent21001 == cg.get_node_id(np.uint64(1), layer=3, x=0, y=0, z=0) + assert parent22001 == cg.get_node_id(np.uint64(1), layer=3, x=1, y=0, z=0) # Children of L3 assert len(children30001) == 1 and len(children31001) == 1 - assert to_label(cg, 2, 1, 0, 0, 1) in children30001 - assert to_label(cg, 2, 2, 0, 0, 1) in children31001 + assert cg.get_node_id(np.uint64(1), layer=2, x=1, y=0, z=0) in children30001 + assert cg.get_node_id(np.uint64(1), layer=2, x=2, y=0, z=0) in children31001 # Parent of L3 assert parent30001 == parent31001 assert ( - parent30001 == to_label(cg, 4, 0, 0, 0, 1) - and parent20001 == to_label(cg, 4, 0, 0, 0, 2) + parent30001 == cg.get_node_id(np.uint64(1), layer=4, x=0, y=0, z=0) + and parent20001 == cg.get_node_id(np.uint64(2), layer=4, x=0, y=0, z=0) ) or ( - parent30001 == to_label(cg, 4, 0, 0, 0, 2) - and parent20001 == to_label(cg, 4, 0, 0, 0, 1) + parent30001 == cg.get_node_id(np.uint64(2), layer=4, x=0, y=0, z=0) + and parent20001 == cg.get_node_id(np.uint64(1), layer=4, x=0, y=0, z=0) ) # Children of L4 @@ -100,27 +113,27 @@ def test_get_parent_and_children(self, gen_graph_simplequerytest): children2_separate = cg.get_children( [ - to_label(cg, 2, 0, 0, 0, 1), - to_label(cg, 2, 1, 0, 0, 1), - to_label(cg, 2, 2, 0, 0, 1), + cg.get_node_id(np.uint64(1), layer=2, x=0, y=0, z=0), + cg.get_node_id(np.uint64(1), layer=2, x=1, y=0, z=0), + cg.get_node_id(np.uint64(1), layer=2, x=2, y=0, z=0), ] ) assert len(children2_separate) == 3 - assert to_label(cg, 2, 0, 0, 0, 1) in children2_separate and np.all( - np.isin(children2_separate[to_label(cg, 2, 0, 0, 0, 1)], children20001) + assert cg.get_node_id(np.uint64(1), layer=2, x=0, y=0, z=0) in children2_separate and np.all( + np.isin(children2_separate[cg.get_node_id(np.uint64(1), layer=2, x=0, y=0, z=0)], children20001) ) - assert to_label(cg, 2, 1, 0, 0, 1) in children2_separate and np.all( - np.isin(children2_separate[to_label(cg, 2, 1, 0, 0, 1)], children21001) + assert cg.get_node_id(np.uint64(1), layer=2, x=1, y=0, z=0) in children2_separate and np.all( + np.isin(children2_separate[cg.get_node_id(np.uint64(1), layer=2, x=1, y=0, z=0)], children21001) ) - assert to_label(cg, 2, 2, 0, 0, 1) in children2_separate and np.all( - np.isin(children2_separate[to_label(cg, 2, 2, 0, 0, 1)], children22001) + assert cg.get_node_id(np.uint64(1), layer=2, x=2, y=0, z=0) in children2_separate and np.all( + np.isin(children2_separate[cg.get_node_id(np.uint64(1), layer=2, x=2, y=0, z=0)], children22001) ) children2_combined = cg.get_children( [ - to_label(cg, 2, 0, 0, 0, 1), - to_label(cg, 2, 1, 0, 0, 1), - to_label(cg, 2, 2, 0, 0, 1), + cg.get_node_id(np.uint64(1), layer=2, x=0, y=0, z=0), + cg.get_node_id(np.uint64(1), layer=2, x=1, y=0, z=0), + cg.get_node_id(np.uint64(1), layer=2, x=2, y=0, z=0), ], flatten=True, ) @@ -132,89 +145,89 @@ def test_get_parent_and_children(self, gen_graph_simplequerytest): ) @pytest.mark.timeout(30) - def test_get_root(self, gen_graph_simplequerytest): - cg = gen_graph_simplequerytest - root10000 = cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) - root11000 = cg.get_root(to_label(cg, 1, 1, 0, 0, 0)) - root11001 = cg.get_root(to_label(cg, 1, 1, 0, 0, 1)) - root12000 = cg.get_root(to_label(cg, 1, 2, 0, 0, 0)) + def test_get_root(self, gen_graph): + cg, sv = self._build_graph(gen_graph) + root10000 = cg.get_root(sv["a0"]) + root11000 = cg.get_root(sv["b0"]) + root11001 = cg.get_root(sv["b1"]) + root12000 = cg.get_root(sv["c0"]) with pytest.raises(Exception): cg.get_root(0) assert ( - root10000 == to_label(cg, 4, 0, 0, 0, 1) - and root11000 == root11001 == root12000 == to_label(cg, 4, 0, 0, 0, 2) + root10000 == cg.get_node_id(np.uint64(1), layer=4, x=0, y=0, z=0) + and root11000 == root11001 == root12000 == cg.get_node_id(np.uint64(2), layer=4, x=0, y=0, z=0) ) or ( - root10000 == to_label(cg, 4, 0, 0, 0, 2) - and root11000 == root11001 == root12000 == to_label(cg, 4, 0, 0, 0, 1) + root10000 == cg.get_node_id(np.uint64(2), layer=4, x=0, y=0, z=0) + and root11000 == root11001 == root12000 == cg.get_node_id(np.uint64(1), layer=4, x=0, y=0, z=0) ) @pytest.mark.timeout(30) - def test_get_subgraph_nodes(self, gen_graph_simplequerytest): - cg = gen_graph_simplequerytest - root1 = cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) - root2 = cg.get_root(to_label(cg, 1, 1, 0, 0, 0)) + def test_get_subgraph_nodes(self, gen_graph): + cg, sv = self._build_graph(gen_graph) + root1 = cg.get_root(sv["a0"]) + root2 = cg.get_root(sv["b0"]) lvl1_nodes_1 = cg.get_subgraph([root1], leaves_only=True) lvl1_nodes_2 = cg.get_subgraph([root2], leaves_only=True) assert len(lvl1_nodes_1) == 1 assert len(lvl1_nodes_2) == 3 - assert to_label(cg, 1, 0, 0, 0, 0) in lvl1_nodes_1 - assert to_label(cg, 1, 1, 0, 0, 0) in lvl1_nodes_2 - assert to_label(cg, 1, 1, 0, 0, 1) in lvl1_nodes_2 - assert to_label(cg, 1, 2, 0, 0, 0) in lvl1_nodes_2 + assert sv["a0"] in lvl1_nodes_1 + assert sv["b0"] in lvl1_nodes_2 + assert sv["b1"] in lvl1_nodes_2 + assert sv["c0"] in lvl1_nodes_2 - lvl2_parent = cg.get_parent(to_label(cg, 1, 1, 0, 0, 0)) + lvl2_parent = cg.get_parent(sv["b0"]) lvl1_nodes = cg.get_subgraph([lvl2_parent], leaves_only=True) assert len(lvl1_nodes) == 2 - assert to_label(cg, 1, 1, 0, 0, 0) in lvl1_nodes - assert to_label(cg, 1, 1, 0, 0, 1) in lvl1_nodes + assert sv["b0"] in lvl1_nodes + assert sv["b1"] in lvl1_nodes @pytest.mark.timeout(30) - def test_get_subgraph_edges(self, gen_graph_simplequerytest): - cg = gen_graph_simplequerytest - root1 = cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) - root2 = cg.get_root(to_label(cg, 1, 1, 0, 0, 0)) + def test_get_subgraph_edges(self, gen_graph): + cg, sv = self._build_graph(gen_graph) + root1 = cg.get_root(sv["a0"]) + root2 = cg.get_root(sv["b0"]) edges = cg.get_subgraph([root1], edges_only=True) assert len(edges) == 0 edges = cg.get_subgraph([root2], edges_only=True) - assert [to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 1)] in edges or [ - to_label(cg, 1, 1, 0, 0, 1), - to_label(cg, 1, 1, 0, 0, 0), + assert [sv["b0"], sv["b1"]] in edges or [ + sv["b1"], + sv["b0"], ] in edges - assert [to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 2, 0, 0, 0)] in edges or [ - to_label(cg, 1, 2, 0, 0, 0), - to_label(cg, 1, 1, 0, 0, 0), + assert [sv["b0"], sv["c0"]] in edges or [ + sv["c0"], + sv["b0"], ] in edges - lvl2_parent = cg.get_parent(to_label(cg, 1, 1, 0, 0, 0)) + lvl2_parent = cg.get_parent(sv["b0"]) edges = cg.get_subgraph([lvl2_parent], edges_only=True) - assert [to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 1)] in edges or [ - to_label(cg, 1, 1, 0, 0, 1), - to_label(cg, 1, 1, 0, 0, 0), + assert [sv["b0"], sv["b1"]] in edges or [ + sv["b1"], + sv["b0"], ] in edges - assert [to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 2, 0, 0, 0)] in edges or [ - to_label(cg, 1, 2, 0, 0, 0), - to_label(cg, 1, 1, 0, 0, 0), + assert [sv["b0"], sv["c0"]] in edges or [ + sv["c0"], + sv["b0"], ] in edges assert len(edges) == 1 @pytest.mark.timeout(30) - def test_get_subgraph_nodes_bb(self, gen_graph_simplequerytest): - cg = gen_graph_simplequerytest + def test_get_subgraph_nodes_bb(self, gen_graph): + cg, sv = self._build_graph(gen_graph) bb = np.array([[1, 0, 0], [2, 1, 1]], dtype=int) bb_coord = bb * cg.meta.graph_config.CHUNK_SIZE childs_1 = cg.get_subgraph( - [cg.get_root(to_label(cg, 1, 1, 0, 0, 1))], bbox=bb, leaves_only=True + [cg.get_root(sv["b1"])], bbox=bb, leaves_only=True ) childs_2 = cg.get_subgraph( - [cg.get_root(to_label(cg, 1, 1, 0, 0, 1))], + [cg.get_root(sv["b1"])], bbox=bb_coord, bbox_is_coordinate=True, leaves_only=True, diff --git a/pychunkedgraph/tests/graph/test_history.py b/pychunkedgraph/tests/graph/test_history.py index bb2573ce5..4c0d25d7e 100644 --- a/pychunkedgraph/tests/graph/test_history.py +++ b/pychunkedgraph/tests/graph/test_history.py @@ -3,11 +3,9 @@ import numpy as np import pytest -from ..helpers import create_chunk, to_label, fake_timestamp -from ...graph import ChunkedGraph +from ..helpers import SV, build_graph, fake_timestamp from ...graph.lineage import lineage_graph, get_root_id_history from ...graph.misc import get_delta_roots -from ...ingest.create.parent_layer import add_parent_chunk class TestGraphHistory: @@ -15,36 +13,22 @@ class TestGraphHistory: @pytest.mark.timeout(120) def test_cut_merge_history(self, gen_graph): - cg: ChunkedGraph = gen_graph(n_layers=3) fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "b": SV(x=1)}, + edges=[("a0", "b", 0.5)], timestamp=fake_ts, ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - - add_parent_chunk( - cg, - 3, - [0, 0, 0], - time_stamp=fake_ts, - n_threads=1, - ) - first_root = cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) - assert first_root == cg.get_root(to_label(cg, 1, 1, 0, 0, 0)) + first_root = cg.get_root(sv["a0"]) + assert first_root == cg.get_root(sv["b"]) timestamp_before_split = datetime.now(UTC) split_roots = cg.remove_edges( "Jane Doe", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 0), + source_ids=sv["a0"], + sink_ids=sv["b"], mincut=False, ).new_root_ids assert len(split_roots) == 2 @@ -56,7 +40,7 @@ def test_cut_merge_history(self, gen_graph): timestamp_after_split = datetime.now(UTC) merge_roots = cg.add_edges( "Jane Doe", - [to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0)], + [sv["a0"], sv["b"]], affinities=0.4, ).new_root_ids assert len(merge_roots) == 1 diff --git a/pychunkedgraph/tests/graph/test_locks.py b/pychunkedgraph/tests/graph/test_locks.py index 1fcec5912..a6e9eabc2 100644 --- a/pychunkedgraph/tests/graph/test_locks.py +++ b/pychunkedgraph/tests/graph/test_locks.py @@ -7,10 +7,9 @@ from ..helpers import ( RowKeyLockRegistry, - create_chunk, + SV, + build_graph, make_cg_with_row_key_lock_registry, - to_label, - fake_timestamp, ) from ...graph import attributes, exceptions from ...graph.locks import ( @@ -19,7 +18,6 @@ _l2_chunk_lock_row_key, ) from ...graph.lineage import get_future_root_ids -from ...ingest.create.parent_layer import add_parent_chunk class TestGraphLocks: @@ -39,35 +37,14 @@ def test_lock_unlock(self, gen_graph): (4) Try lock (opid = 2) """ - cg = gen_graph(n_layers=3) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2)], - edges=[], - timestamp=fake_ts, - ) - - # Preparation: Build Chunk B - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 1)], - edges=[], - timestamp=fake_ts, - ) - - add_parent_chunk( - cg, - 3, - [0, 0, 0], - time_stamp=fake_ts, - n_threads=1, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a1": SV(seg=1), "a2": SV(seg=2), "b1": SV(x=1, seg=1)}, ) operation_id_1 = cg.id_client.create_operation_id() - root_id = cg.get_root(to_label(cg, 1, 0, 0, 0, 1)) + root_id = cg.get_root(sv["a1"]) future_root_ids_d = {root_id: get_future_root_ids(cg, root_id)} assert cg.client.lock_roots( @@ -105,35 +82,14 @@ def test_lock_expiration(self, gen_graph): (2) Try lock (opid = 2) (3) Try lock (opid = 2) with retries """ - cg = gen_graph(n_layers=3) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2)], - edges=[], - timestamp=fake_ts, - ) - - # Preparation: Build Chunk B - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 1)], - edges=[], - timestamp=fake_ts, - ) - - add_parent_chunk( - cg, - 3, - [0, 0, 0], - time_stamp=fake_ts, - n_threads=1, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a1": SV(seg=1), "a2": SV(seg=2), "b1": SV(x=1, seg=1)}, ) operation_id_1 = cg.id_client.create_operation_id() - root_id = cg.get_root(to_label(cg, 1, 0, 0, 0, 1)) + root_id = cg.get_root(sv["a1"]) future_root_ids_d = {root_id: get_future_root_ids(cg, root_id)} assert cg.client.lock_roots( root_ids=[root_id], @@ -173,35 +129,14 @@ def test_lock_renew(self, gen_graph): (3) Try lock (opid = 2) with retries """ - cg = gen_graph(n_layers=3) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2)], - edges=[], - timestamp=fake_ts, - ) - - # Preparation: Build Chunk B - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 1)], - edges=[], - timestamp=fake_ts, - ) - - add_parent_chunk( - cg, - 3, - [0, 0, 0], - time_stamp=fake_ts, - n_threads=1, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a1": SV(seg=1), "a2": SV(seg=2), "b1": SV(x=1, seg=1)}, ) operation_id_1 = cg.id_client.create_operation_id() - root_id = cg.get_root(to_label(cg, 1, 0, 0, 0, 1)) + root_id = cg.get_root(sv["a1"]) future_root_ids_d = {root_id: get_future_root_ids(cg, root_id)} assert cg.client.lock_roots( root_ids=[root_id], @@ -225,38 +160,17 @@ def test_lock_merge_lock_old_id(self, gen_graph): (2) Try lock opid 2 --> should be successful and return new root id """ - cg = gen_graph(n_layers=3) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2)], - edges=[], - timestamp=fake_ts, - ) - - # Preparation: Build Chunk B - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 1)], - edges=[], - timestamp=fake_ts, - ) - - add_parent_chunk( - cg, - 3, - [0, 0, 0], - time_stamp=fake_ts, - n_threads=1, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a1": SV(seg=1), "a2": SV(seg=2), "b1": SV(x=1, seg=1)}, ) - root_id = cg.get_root(to_label(cg, 1, 0, 0, 0, 1)) + root_id = cg.get_root(sv["a1"]) new_root_ids = cg.add_edges( "Chuck Norris", - [to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2)], + [sv["a1"], sv["a2"]], affinities=1.0, ).new_root_ids @@ -291,35 +205,14 @@ def test_indefinite_lock(self, gen_graph): (4) Try lock (opid = 2), should get the normal lock """ - cg = gen_graph(n_layers=3) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2)], - edges=[], - timestamp=fake_ts, - ) - - # Preparation: Build Chunk B - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 1)], - edges=[], - timestamp=fake_ts, - ) - - add_parent_chunk( - cg, - 3, - [0, 0, 0], - time_stamp=fake_ts, - n_threads=1, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a1": SV(seg=1), "a2": SV(seg=2), "b1": SV(x=1, seg=1)}, ) operation_id_1 = cg.id_client.create_operation_id() - root_id = cg.get_root(to_label(cg, 1, 0, 0, 0, 1)) + root_id = cg.get_root(sv["a1"]) future_root_ids_d = {root_id: get_future_root_ids(cg, root_id)} assert cg.client.lock_roots_indefinitely( @@ -364,35 +257,14 @@ def test_indefinite_lock_with_normal_lock_expiration(self, gen_graph): """ # 1. TODO renew lock test when getting indefinite lock - cg = gen_graph(n_layers=3) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2)], - edges=[], - timestamp=fake_ts, - ) - - # Preparation: Build Chunk B - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 1)], - edges=[], - timestamp=fake_ts, - ) - - add_parent_chunk( - cg, - 3, - [0, 0, 0], - time_stamp=fake_ts, - n_threads=1, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a1": SV(seg=1), "a2": SV(seg=2), "b1": SV(x=1, seg=1)}, ) operation_id_1 = cg.id_client.create_operation_id() - root_id = cg.get_root(to_label(cg, 1, 0, 0, 0, 1)) + root_id = cg.get_root(sv["a1"]) future_root_ids_d = {root_id: get_future_root_ids(cg, root_id)} diff --git a/pychunkedgraph/tests/graph/test_merge.py b/pychunkedgraph/tests/graph/test_merge.py index d481e6184..32f16d305 100644 --- a/pychunkedgraph/tests/graph/test_merge.py +++ b/pychunkedgraph/tests/graph/test_merge.py @@ -4,10 +4,7 @@ import numpy as np import pytest -from ..helpers import create_chunk, to_label, fake_timestamp -from ...graph import ChunkedGraph -from ...graph import serializers -from ...ingest.create.parent_layer import add_parent_chunk +from ..helpers import SV, build_graph, assert_graph_unchanged class TestGraphMerge: @@ -23,22 +20,17 @@ def test_merge_pair_same_chunk(self, gen_graph): └─────┘ └─────┘ """ - atomic_chunk_bounds = np.array([1, 1, 1]) - cg = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) # Merge new_root_ids = cg.add_edges( "Jane Doe", - [to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 0)], + [sv["a1"], sv["a0"]], affinities=[0.3], ).new_root_ids @@ -46,12 +38,12 @@ def test_merge_pair_same_chunk(self, gen_graph): new_root_id = new_root_ids[0] # Check - assert cg.get_parent(to_label(cg, 1, 0, 0, 0, 0)) == new_root_id - assert cg.get_parent(to_label(cg, 1, 0, 0, 0, 1)) == new_root_id + assert cg.get_parent(sv["a0"]) == new_root_id + assert cg.get_parent(sv["a1"]) == new_root_id leaves = np.unique(cg.get_subgraph([new_root_id], leaves_only=True)) assert len(leaves) == 2 - assert to_label(cg, 1, 0, 0, 0, 0) in leaves - assert to_label(cg, 1, 0, 0, 0, 1) in leaves + assert sv["a0"] in leaves + assert sv["a1"] in leaves @pytest.mark.timeout(30) def test_merge_pair_neighboring_chunks(self, gen_graph): @@ -64,37 +56,16 @@ def test_merge_pair_neighboring_chunks(self, gen_graph): └─────┴─────┘ └─────┴─────┘ """ - cg = gen_graph(n_layers=3) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, - ) - - # Preparation: Build Chunk B - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[], - timestamp=fake_ts, - ) - - add_parent_chunk( - cg, - 3, - [0, 0, 0], - time_stamp=fake_ts, - n_threads=1, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "b": SV(x=1)}, ) # Merge new_root_ids = cg.add_edges( "Jane Doe", - [to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0)], + [sv["b"], sv["a0"]], affinities=0.3, ).new_root_ids @@ -102,12 +73,12 @@ def test_merge_pair_neighboring_chunks(self, gen_graph): new_root_id = new_root_ids[0] # Check - assert cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) == new_root_id - assert cg.get_root(to_label(cg, 1, 1, 0, 0, 0)) == new_root_id + assert cg.get_root(sv["a0"]) == new_root_id + assert cg.get_root(sv["b"]) == new_root_id leaves = np.unique(cg.get_subgraph([new_root_id], leaves_only=True)) assert len(leaves) == 2 - assert to_label(cg, 1, 0, 0, 0, 0) in leaves - assert to_label(cg, 1, 1, 0, 0, 0) in leaves + assert sv["a0"] in leaves + assert sv["b"] in leaves @pytest.mark.timeout(120) def test_merge_pair_disconnected_chunks(self, gen_graph): @@ -120,58 +91,16 @@ def test_merge_pair_disconnected_chunks(self, gen_graph): └─────┘ └─────┘ └─────┘ └─────┘ """ - cg = gen_graph(n_layers=5) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, - ) - - # Preparation: Build Chunk Z - create_chunk( - cg, - vertices=[to_label(cg, 1, 7, 7, 7, 0)], - edges=[], - timestamp=fake_ts, - ) - - add_parent_chunk( - cg, - 3, - [0, 0, 0], - time_stamp=fake_ts, - n_threads=1, - ) - add_parent_chunk( - cg, - 3, - [3, 3, 3], - time_stamp=fake_ts, - n_threads=1, - ) - add_parent_chunk( - cg, - 4, - [0, 0, 0], - time_stamp=fake_ts, - n_threads=1, - ) - add_parent_chunk( - cg, - 5, - [0, 0, 0], - time_stamp=fake_ts, - n_threads=1, + cg, sv = build_graph( + gen_graph, + n_layers=5, + supervoxels={"a0": SV(), "z": SV(x=7, y=7, z=7)}, ) # Merge result = cg.add_edges( "Jane Doe", - [to_label(cg, 1, 7, 7, 7, 0), to_label(cg, 1, 0, 0, 0, 0)], + [sv["z"], sv["a0"]], affinities=[0.3], ) new_root_ids, lvl2_node_ids = result.new_root_ids, result.new_lvl2_ids @@ -184,12 +113,12 @@ def test_merge_pair_disconnected_chunks(self, gen_graph): new_root_id = new_root_ids[0] # Check - assert cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) == new_root_id - assert cg.get_root(to_label(cg, 1, 7, 7, 7, 0)) == new_root_id + assert cg.get_root(sv["a0"]) == new_root_id + assert cg.get_root(sv["z"]) == new_root_id leaves = np.unique(cg.get_subgraph(new_root_id, leaves_only=True)) assert len(leaves) == 2 - assert to_label(cg, 1, 0, 0, 0, 0) in leaves - assert to_label(cg, 1, 7, 7, 7, 0) in leaves + assert sv["a0"] in leaves + assert sv["z"] in leaves @pytest.mark.timeout(30) def test_merge_pair_already_connected(self, gen_graph): @@ -203,15 +132,11 @@ def test_merge_pair_already_connected(self, gen_graph): └─────┘ └─────┘ """ - cg = gen_graph(n_layers=2) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5)], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + supervoxels={"a0": SV(), "a1": SV(seg=1)}, + edges=[("a0", "a1", 0.5)], ) res_old = cg.client.read_all_rows() @@ -221,7 +146,7 @@ def test_merge_pair_already_connected(self, gen_graph): with pytest.raises(Exception): cg.add_edges( "Jane Doe", - [to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 0)], + [sv["a1"], sv["a0"]], ) res_new = cg.client.read_all_rows() res_new.consume_all() @@ -246,29 +171,21 @@ def test_merge_triple_chain_to_full_circle_same_chunk(self, gen_graph): └─────┘ └─────┘ """ - cg = gen_graph(n_layers=2) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[ - to_label(cg, 1, 0, 0, 0, 0), - to_label(cg, 1, 0, 0, 0, 1), - to_label(cg, 1, 0, 0, 0, 2), - ], + cg, sv = build_graph( + gen_graph, + n_layers=2, + supervoxels={"a0": SV(), "a1": SV(seg=1), "a2": SV(seg=2)}, edges=[ - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 2), 0.5), - (to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2), 0.5), + ("a0", "a2", 0.5), + ("a1", "a2", 0.5), ], - timestamp=fake_ts, ) # Merge with pytest.raises(Exception): cg.add_edges( "Jane Doe", - [to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 0)], + [sv["a1"], sv["a0"]], affinities=0.3, ).new_root_ids @@ -283,41 +200,21 @@ def test_merge_triple_chain_to_full_circle_neighboring_chunks(self, gen_graph): └─────┴─────┘ └─────┴─────┘ """ - cg = gen_graph(n_layers=3) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, edges=[ - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5), - (to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 1, 0, 0, 0), inf), + ("a0", "a1", 0.5), + ("a1", "b", inf), ], - timestamp=fake_ts, - ) - - # Preparation: Build Chunk B - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), inf)], - timestamp=fake_ts, - ) - - add_parent_chunk( - cg, - 3, - [0, 0, 0], - time_stamp=fake_ts, - n_threads=1, ) # Merge with pytest.raises(Exception): cg.add_edges( "Jane Doe", - [to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0)], + [sv["b"], sv["a0"]], affinities=1.0, ).new_root_ids @@ -332,38 +229,20 @@ def test_merge_triple_chain_to_full_circle_disconnected_chunks(self, gen_graph): └─────┘ └─────┘ └─────┘ └─────┘ """ - cg = gen_graph(n_layers=5) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], + cg, sv = build_graph( + gen_graph, + n_layers=5, + supervoxels={"a0": SV(), "a1": SV(seg=1), "z": SV(x=7, y=7, z=7)}, edges=[ - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5), - (to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 7, 7, 7, 0), inf), + ("a0", "a1", 0.5), + ("a1", "z", inf), ], - timestamp=fake_ts, - ) - - # Preparation: Build Chunk B - create_chunk( - cg, - vertices=[to_label(cg, 1, 7, 7, 7, 0)], - edges=[(to_label(cg, 1, 7, 7, 7, 0), to_label(cg, 1, 0, 0, 0, 1), inf)], - timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 3, [3, 3, 3], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 4, [1, 1, 1], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 5, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - # Merge new_root_ids = cg.add_edges( "Jane Doe", - [to_label(cg, 1, 7, 7, 7, 0), to_label(cg, 1, 0, 0, 0, 0)], + [sv["z"], sv["a0"]], affinities=1.0, ).new_root_ids @@ -371,14 +250,14 @@ def test_merge_triple_chain_to_full_circle_disconnected_chunks(self, gen_graph): new_root_id = new_root_ids[0] # Check - assert cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) == new_root_id - assert cg.get_root(to_label(cg, 1, 0, 0, 0, 1)) == new_root_id - assert cg.get_root(to_label(cg, 1, 7, 7, 7, 0)) == new_root_id + assert cg.get_root(sv["a0"]) == new_root_id + assert cg.get_root(sv["a1"]) == new_root_id + assert cg.get_root(sv["z"]) == new_root_id leaves = np.unique(cg.get_subgraph(new_root_id, leaves_only=True)) assert len(leaves) == 3 - assert to_label(cg, 1, 0, 0, 0, 0) in leaves - assert to_label(cg, 1, 0, 0, 0, 1) in leaves - assert to_label(cg, 1, 7, 7, 7, 0) in leaves + assert sv["a0"] in leaves + assert sv["a1"] in leaves + assert sv["z"] in leaves @pytest.mark.timeout(30) def test_merge_same_node(self, gen_graph): @@ -391,31 +270,19 @@ def test_merge_same_node(self, gen_graph): └─────┘ """ - cg = gen_graph(n_layers=2) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + supervoxels={"a0": SV()}, ) - res_old = cg.client.read_all_rows() - res_old.consume_all() - # Merge - with pytest.raises(Exception): - cg.add_edges( - "Jane Doe", - [to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0)], - ) - - res_new = cg.client.read_all_rows() - res_new.consume_all() - - assert res_new.rows == res_old.rows + with assert_graph_unchanged(cg): + with pytest.raises(Exception): + cg.add_edges( + "Jane Doe", + [sv["a0"], sv["a0"]], + ) @pytest.mark.timeout(30) def test_merge_pair_abstract_nodes(self, gen_graph): @@ -424,41 +291,19 @@ def test_merge_pair_abstract_nodes(self, gen_graph): => Reject """ - cg = gen_graph(n_layers=3) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "b": SV(x=1)}, ) - # Preparation: Build Chunk B - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[], - timestamp=fake_ts, - ) - - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - - res_old = cg.client.read_all_rows() - res_old.consume_all() - # Merge - with pytest.raises(Exception): - cg.add_edges( - "Jane Doe", - [to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 2, 1, 0, 0, 1)], - ) - - res_new = cg.client.read_all_rows() - res_new.consume_all() - - assert res_new.rows == res_old.rows + with assert_graph_unchanged(cg): + with pytest.raises(Exception): + cg.add_edges( + "Jane Doe", + [sv["a0"], cg.get_node_id(np.uint64(1), layer=2, x=1, y=0, z=0)], + ) @pytest.mark.timeout(30) def test_diagonal_connections(self, gen_graph): @@ -474,44 +319,23 @@ def test_diagonal_connections(self, gen_graph): └─────┴─────┘ """ - cg = gen_graph(n_layers=3) - - # Chunk A - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], - edges=[ - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), inf), - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 1, 0, 0), inf), - ], - ) - - # Chunk B - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), inf)], - ) - - # Chunk C - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 1, 0, 0)], + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={ + "a0": SV(), + "a1": SV(seg=1), + "b": SV(x=1), + "c": SV(y=1), + "d": SV(x=1, y=1), + }, edges=[ - (to_label(cg, 1, 0, 1, 0, 0), to_label(cg, 1, 1, 1, 0, 0), inf), - (to_label(cg, 1, 0, 1, 0, 0), to_label(cg, 1, 0, 0, 0, 0), inf), + ("a0", "b", inf), + ("a0", "c", inf), + ("c", "d", inf), ], ) - # Chunk D - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 1, 0, 0)], - edges=[(to_label(cg, 1, 1, 1, 0, 0), to_label(cg, 1, 0, 1, 0, 0), inf)], - ) - - add_parent_chunk(cg, 3, [0, 0, 0], n_threads=1) - rr = cg.range_read_chunk(chunk_id=cg.get_chunk_id(layer=3, x=0, y=0, z=0)) root_ids_t0 = list(rr.keys()) @@ -523,7 +347,7 @@ def test_diagonal_connections(self, gen_graph): new_roots = cg.add_edges( "Jane Doe", - [to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], + [sv["a0"], sv["a1"]], affinities=[0.5], ).new_root_ids @@ -538,46 +362,26 @@ def test_diagonal_connections(self, gen_graph): @pytest.mark.timeout(240) def test_cross_edges(self, gen_graph): - cg = gen_graph(n_layers=5) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], - edges=[ - (to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 1, 0, 0, 0), inf), - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), inf), - ], - timestamp=fake_ts, - ) - - # Preparation: Build Chunk B - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 1)], + cg, sv = build_graph( + gen_graph, + n_layers=5, + supervoxels={ + "a0": SV(), + "a1": SV(seg=1), + "b0": SV(x=1), + "b1": SV(x=1, seg=1), + "c": SV(x=2), + }, edges=[ - (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), inf), - (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 1), inf), + ("a1", "b0", inf), + ("a0", "a1", inf), + ("b0", "b1", inf), ], - timestamp=fake_ts, ) - # Preparation: Build Chunk C - create_chunk( - cg, - vertices=[to_label(cg, 1, 2, 0, 0, 0)], - timestamp=fake_ts, - ) - - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 3, [1, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 5, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - new_roots = cg.add_edges( "Jane Doe", - [to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 2, 0, 0, 0)], + [sv["b0"], sv["c"]], affinities=0.9, ).new_root_ids @@ -601,30 +405,15 @@ def test_merge_creates_skip_connection(self, gen_graph): After merge: 1 and 2 are connected, hierarchy should skip intermediate empty layers. """ - cg = gen_graph(n_layers=5) - - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=5, + supervoxels={"a0": SV(), "z": SV(x=7, y=7, z=7)}, ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 7, 7, 7, 0)], - edges=[], - timestamp=fake_ts, - ) - - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 3, [3, 3, 3], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 5, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # Before merge: verify both nodes have root at layer 5 - root1_pre = cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) - root2_pre = cg.get_root(to_label(cg, 1, 7, 7, 7, 0)) + root1_pre = cg.get_root(sv["a0"]) + root2_pre = cg.get_root(sv["z"]) assert root1_pre != root2_pre assert cg.get_chunk_layer(root1_pre) == 5 assert cg.get_chunk_layer(root2_pre) == 5 @@ -632,7 +421,7 @@ def test_merge_creates_skip_connection(self, gen_graph): # Merge result = cg.add_edges( "Jane Doe", - [to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 7, 7, 7, 0)], + [sv["a0"], sv["z"]], affinities=[0.5], ) new_root_ids = result.new_root_ids @@ -640,8 +429,8 @@ def test_merge_creates_skip_connection(self, gen_graph): # After merge: single root, both supervoxels reachable new_root = new_root_ids[0] - assert cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) == new_root - assert cg.get_root(to_label(cg, 1, 7, 7, 7, 0)) == new_root + assert cg.get_root(sv["a0"]) == new_root + assert cg.get_root(sv["z"]) == new_root assert cg.get_chunk_layer(new_root) == 5 @pytest.mark.timeout(120) @@ -651,36 +440,21 @@ def test_merge_multi_layer_hierarchy_correctness(self, gen_graph): each supervoxel to root is valid — every node has a parent at a higher layer, and the root is reachable. """ - cg = gen_graph(n_layers=5) - - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=5, + supervoxels={"a0": SV(), "z": SV(x=7, y=7, z=7)}, ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 7, 7, 7, 0)], - edges=[], - timestamp=fake_ts, - ) - - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 3, [3, 3, 3], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 5, [0, 0, 0], time_stamp=fake_ts, n_threads=1) result = cg.add_edges( "Jane Doe", - [to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 7, 7, 7, 0)], + [sv["a0"], sv["z"]], affinities=[0.5], ) # Verify parent chain for both supervoxels - for sv in [to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 7, 7, 7, 0)]: - parents = cg.get_root(sv, get_all_parents=True) + for node in [sv["a0"], sv["z"]]: + parents = cg.get_root(node, get_all_parents=True) # Each parent should be at a strictly higher layer prev_layer = 1 for p in parents: @@ -704,28 +478,16 @@ def test_merge_no_skip_when_siblings_exist(self, gen_graph): │ 1 │ 2 │ └─────┴─────┘ """ - cg = gen_graph(n_layers=3) - - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "b": SV(x=1)}, ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[], - timestamp=fake_ts, - ) - - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # Merge result = cg.add_edges( "Jane Doe", - [to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0)], + [sv["a0"], sv["b"]], affinities=[0.5], ) diff --git a/pychunkedgraph/tests/graph/test_merge_split.py b/pychunkedgraph/tests/graph/test_merge_split.py index 2279b6c6a..7fdf3190c 100644 --- a/pychunkedgraph/tests/graph/test_merge_split.py +++ b/pychunkedgraph/tests/graph/test_merge_split.py @@ -1,17 +1,36 @@ -from datetime import datetime, timedelta, UTC from math import inf import numpy as np import pytest -from ..helpers import create_chunk, to_label +from ..helpers import SV, build_graph from ...graph import types class TestGraphMergeSplit: @pytest.mark.timeout(240) - def test_multiple_cuts_and_splits(self, gen_graph_simplequerytest): - cg = gen_graph_simplequerytest + def test_multiple_cuts_and_splits(self, gen_graph): + """ + ┌─────┬─────┬─────┐ + │ A¹ │ B¹ │ C¹ │ + │ 1 │ 3━2━┿━━4 │ + │ │ │ │ + └─────┴─────┴─────┘ + """ + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={ + "a0": SV(), + "b0": SV(x=1), + "b1": SV(x=1, seg=1), + "c0": SV(x=2), + }, + edges=[ + ("b0", "b1", 0.5), + ("b0", "c0", inf), + ], + ) rr = cg.range_read_chunk(chunk_id=cg.get_chunk_id(layer=4, x=0, y=0, z=0)) root_ids_t0 = list(rr.keys()) @@ -23,7 +42,7 @@ def test_multiple_cuts_and_splits(self, gen_graph_simplequerytest): for i in range(10): new_roots = cg.add_edges( "Jane Doe", - [to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 1)], + [sv["a0"], sv["b1"]], affinities=0.9, ).new_root_ids assert len(new_roots) == 1, new_roots @@ -35,8 +54,8 @@ def test_multiple_cuts_and_splits(self, gen_graph_simplequerytest): new_roots = cg.remove_edges( "John Doe", - source_ids=to_label(cg, 1, 1, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 1), + source_ids=sv["b0"], + sink_ids=sv["b1"], mincut=False, ).new_root_ids assert len(new_roots) == 2, new_roots @@ -52,8 +71,8 @@ def test_multiple_cuts_and_splits(self, gen_graph_simplequerytest): new_roots = cg.remove_edges( "Jane Doe", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 1), + source_ids=sv["a0"], + sink_ids=sv["b1"], mincut=False, ).new_root_ids assert len(new_roots) == 2, new_roots @@ -64,7 +83,7 @@ def test_multiple_cuts_and_splits(self, gen_graph_simplequerytest): new_roots = cg.add_edges( "Jane Doe", - [to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 1)], + [sv["b0"], sv["b1"]], affinities=0.9, ).new_root_ids assert len(new_roots) == 1, new_roots diff --git a/pychunkedgraph/tests/graph/test_mincut.py b/pychunkedgraph/tests/graph/test_mincut.py index 53d8355db..7879ccd78 100644 --- a/pychunkedgraph/tests/graph/test_mincut.py +++ b/pychunkedgraph/tests/graph/test_mincut.py @@ -3,9 +3,8 @@ import numpy as np import pytest -from ..helpers import create_chunk, to_label, fake_timestamp +from ..helpers import SV, build_graph, assert_graph_unchanged from ...graph import exceptions -from ...ingest.create.parent_layer import add_parent_chunk class TestGraphMinCut: @@ -21,39 +20,18 @@ def test_cut_regular_link(self, gen_graph): │ │ │ └─────┴─────┘ """ - - cg = gen_graph(n_layers=3) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - - # Preparation: Build Chunk B - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - - add_parent_chunk( - cg, - 3, - [0, 0, 0], - time_stamp=fake_ts, - n_threads=1, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "b": SV(x=1)}, + edges=[("a0", "b", 0.5)], ) # Mincut new_root_ids = cg.remove_edges( "Jane Doe", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 0), + source_ids=sv["a0"], + sink_ids=sv["b"], source_coords=[0, 0, 0], sink_coords=[ 2 * cg.meta.graph_config.CHUNK_SIZE[0], @@ -66,21 +44,15 @@ def test_cut_regular_link(self, gen_graph): # verify new state assert len(new_root_ids) == 2 - assert cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) != cg.get_root( - to_label(cg, 1, 1, 0, 0, 0) - ) + assert cg.get_root(sv["a0"]) != cg.get_root(sv["b"]) leaves = np.unique( - cg.get_subgraph( - [cg.get_root(to_label(cg, 1, 0, 0, 0, 0))], leaves_only=True - ) + cg.get_subgraph([cg.get_root(sv["a0"])], leaves_only=True) ) - assert len(leaves) == 1 and to_label(cg, 1, 0, 0, 0, 0) in leaves + assert len(leaves) == 1 and sv["a0"] in leaves leaves = np.unique( - cg.get_subgraph( - [cg.get_root(to_label(cg, 1, 1, 0, 0, 0))], leaves_only=True - ) + cg.get_subgraph([cg.get_root(sv["b"])], leaves_only=True) ) - assert len(leaves) == 1 and to_label(cg, 1, 1, 0, 0, 0) in leaves + assert len(leaves) == 1 and sv["b"] in leaves @pytest.mark.timeout(30) def test_cut_no_link(self, gen_graph): @@ -92,56 +64,27 @@ def test_cut_no_link(self, gen_graph): │ │ │ └─────┴─────┘ """ - - cg = gen_graph(n_layers=3) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, - ) - - # Preparation: Build Chunk B - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[], - timestamp=fake_ts, - ) - - add_parent_chunk( - cg, - 3, - [0, 0, 0], - time_stamp=fake_ts, - n_threads=1, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "b": SV(x=1)}, ) - res_old = cg.client.read_all_rows() - res_old.consume_all() - # Mincut - with pytest.raises(exceptions.PreconditionError): - cg.remove_edges( - "Jane Doe", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 0), - source_coords=[0, 0, 0], - sink_coords=[ - 2 * cg.meta.graph_config.CHUNK_SIZE[0], - 2 * cg.meta.graph_config.CHUNK_SIZE[1], - cg.meta.graph_config.CHUNK_SIZE[2], - ], - mincut=True, - ) - - res_new = cg.client.read_all_rows() - res_new.consume_all() - - assert res_new.rows == res_old.rows + with assert_graph_unchanged(cg): + with pytest.raises(exceptions.PreconditionError): + cg.remove_edges( + "Jane Doe", + source_ids=sv["a0"], + sink_ids=sv["b"], + source_coords=[0, 0, 0], + sink_coords=[ + 2 * cg.meta.graph_config.CHUNK_SIZE[0], + 2 * cg.meta.graph_config.CHUNK_SIZE[1], + cg.meta.graph_config.CHUNK_SIZE[2], + ], + mincut=True, + ) @pytest.mark.timeout(30) def test_cut_old_link(self, gen_graph): @@ -153,62 +96,34 @@ def test_cut_old_link(self, gen_graph): │ │ │ └─────┴─────┘ """ - - cg = gen_graph(n_layers=3) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - - # Preparation: Build Chunk B - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - - add_parent_chunk( - cg, - 3, - [0, 0, 0], - time_stamp=fake_ts, - n_threads=1, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "b": SV(x=1)}, + edges=[("a0", "b", 0.5)], ) cg.remove_edges( "John Doe", - source_ids=to_label(cg, 1, 1, 0, 0, 0), - sink_ids=to_label(cg, 1, 0, 0, 0, 0), + source_ids=sv["b"], + sink_ids=sv["a0"], mincut=False, ) - res_old = cg.client.read_all_rows() - res_old.consume_all() - # Mincut - with pytest.raises(exceptions.PreconditionError): - cg.remove_edges( - "Jane Doe", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 0), - source_coords=[0, 0, 0], - sink_coords=[ - 2 * cg.meta.graph_config.CHUNK_SIZE[0], - 2 * cg.meta.graph_config.CHUNK_SIZE[1], - cg.meta.graph_config.CHUNK_SIZE[2], - ], - mincut=True, - ) - - res_new = cg.client.read_all_rows() - res_new.consume_all() - - assert res_new.rows == res_old.rows + with assert_graph_unchanged(cg): + with pytest.raises(exceptions.PreconditionError): + cg.remove_edges( + "Jane Doe", + source_ids=sv["a0"], + sink_ids=sv["b"], + source_coords=[0, 0, 0], + sink_coords=[ + 2 * cg.meta.graph_config.CHUNK_SIZE[0], + 2 * cg.meta.graph_config.CHUNK_SIZE[1], + cg.meta.graph_config.CHUNK_SIZE[2], + ], + mincut=True, + ) @pytest.mark.timeout(30) def test_cut_indivisible_link(self, gen_graph): @@ -221,47 +136,22 @@ def test_cut_indivisible_link(self, gen_graph): │ │ │ └─────┴─────┘ """ - - cg = gen_graph(n_layers=3) - - # Preparation: Build Chunk A - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), inf)], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "b": SV(x=1)}, + edges=[("a0", "b", inf)], ) - # Preparation: Build Chunk B - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), inf)], - timestamp=fake_ts, - ) - - add_parent_chunk( - cg, - 3, - [0, 0, 0], - time_stamp=fake_ts, - n_threads=1, - ) - - original_parents_1 = cg.get_root( - to_label(cg, 1, 0, 0, 0, 0), get_all_parents=True - ) - original_parents_2 = cg.get_root( - to_label(cg, 1, 1, 0, 0, 0), get_all_parents=True - ) + original_parents_1 = cg.get_root(sv["a0"], get_all_parents=True) + original_parents_2 = cg.get_root(sv["b"], get_all_parents=True) # Mincut with pytest.raises(exceptions.PostconditionError): cg.remove_edges( "Jane Doe", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 0), + source_ids=sv["a0"], + sink_ids=sv["b"], source_coords=[0, 0, 0], sink_coords=[ 2 * cg.meta.graph_config.CHUNK_SIZE[0], @@ -271,8 +161,8 @@ def test_cut_indivisible_link(self, gen_graph): mincut=True, ) - new_parents_1 = cg.get_root(to_label(cg, 1, 0, 0, 0, 0), get_all_parents=True) - new_parents_2 = cg.get_root(to_label(cg, 1, 1, 0, 0, 0), get_all_parents=True) + new_parents_1 = cg.get_root(sv["a0"], get_all_parents=True) + new_parents_2 = cg.get_root(sv["b"], get_all_parents=True) assert np.all(np.array(original_parents_1) == np.array(new_parents_1)) assert np.all(np.array(original_parents_2) == np.array(new_parents_2)) @@ -285,31 +175,24 @@ def test_mincut_disrespects_sources_or_sinks(self, gen_graph): two sinks, this can happen when an edge along the only path between two sources or two sinks is cut. """ - cg = gen_graph(n_layers=2) - - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[ - to_label(cg, 1, 0, 0, 0, 0), - to_label(cg, 1, 0, 0, 0, 1), - to_label(cg, 1, 0, 0, 0, 2), - to_label(cg, 1, 0, 0, 0, 3), - ], - edges=[ - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 2), 2), - (to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2), 3), - (to_label(cg, 1, 0, 0, 0, 2), to_label(cg, 1, 0, 0, 0, 3), 10), - ], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + supervoxels={ + "a0": SV(), + "a1": SV(seg=1), + "a2": SV(seg=2), + "a3": SV(seg=3), + }, + edges=[("a0", "a2", 2), ("a1", "a2", 3), ("a2", "a3", 10)], ) # Mincut with pytest.raises(exceptions.PreconditionError): cg.remove_edges( "Jane Doe", - source_ids=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], - sink_ids=[to_label(cg, 1, 0, 0, 0, 3)], + source_ids=[sv["a0"], sv["a1"]], + sink_ids=[sv["a3"]], source_coords=[[0, 0, 0], [10, 0, 0]], sink_coords=[[5, 5, 0]], mincut=True, diff --git a/pychunkedgraph/tests/graph/test_misc.py b/pychunkedgraph/tests/graph/test_misc.py index 60786a84a..191d8110d 100644 --- a/pychunkedgraph/tests/graph/test_misc.py +++ b/pychunkedgraph/tests/graph/test_misc.py @@ -1,7 +1,6 @@ """Tests for pychunkedgraph.graph.misc""" from datetime import datetime, timedelta, UTC -from math import inf import numpy as np import pytest @@ -16,44 +15,32 @@ from pychunkedgraph.graph.edges import Edges from pychunkedgraph.graph.types import Agglomeration -from ..helpers import create_chunk, to_label, fake_timestamp -from ...ingest.create.parent_layer import add_parent_chunk +from ..helpers import SV, build_graph, fake_timestamp class TestGetLatestRoots: def test_basic(self, gen_graph): - graph = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), - ], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "a1": SV(seg=1)}, + edges=[("a0", "a1", 0.5)], ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - roots = get_latest_roots(graph) + roots = get_latest_roots(cg) assert len(roots) >= 1 def test_with_timestamp(self, gen_graph): - graph = gen_graph(n_layers=4) fake_ts = fake_timestamp() - - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV()}, timestamp=fake_ts, ) - add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) - add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) - roots_before = get_latest_roots(graph, fake_ts - timedelta(days=1)) - roots_after = get_latest_roots(graph) + roots_before = get_latest_roots(cg, fake_ts - timedelta(days=1)) + roots_after = get_latest_roots(cg) # Before creation, there should be no roots assert len(roots_before) == 0 assert len(roots_after) >= 1 @@ -61,65 +48,54 @@ def test_with_timestamp(self, gen_graph): class TestGetDeltaRoots: def test_basic(self, gen_graph): - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) before_merge = datetime.now(UTC) - graph.add_edges( + cg.add_edges( "TestUser", - [to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], + [sv["a0"], sv["a1"]], affinities=[0.3], ) - new_roots, expired_roots = get_delta_roots(graph, before_merge) + new_roots, expired_roots = get_delta_roots(cg, before_merge) assert len(new_roots) >= 1 class TestGetProofreadRootIds: def test_after_merge(self, gen_graph): """After a merge, get_proofread_root_ids should return old and new root IDs.""" - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - - fake_ts = fake_timestamp() - sv0 = to_label(graph, 1, 0, 0, 0, 0) - sv1 = to_label(graph, 1, 0, 0, 0, 1) - - create_chunk( - graph, - vertices=[sv0, sv1], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) before_merge = datetime.now(UTC) # Both SVs should have separate roots before merge - old_root0 = graph.get_root(sv0) - old_root1 = graph.get_root(sv1) + old_root0 = cg.get_root(sv["a0"]) + old_root1 = cg.get_root(sv["a1"]) assert old_root0 != old_root1 # Perform a merge - graph.add_edges( + cg.add_edges( "TestUser", - [sv0, sv1], + [sv["a0"], sv["a1"]], affinities=[0.3], ) # After merge, the two SVs share a new root - new_root = graph.get_root(sv0) - assert new_root == graph.get_root(sv1) + new_root = cg.get_root(sv["a0"]) + assert new_root == cg.get_root(sv["a1"]) - old_roots, new_roots = get_proofread_root_ids(graph, start_time=before_merge) + old_roots, new_roots = get_proofread_root_ids(cg, start_time=before_merge) # The new root from the merge should appear in new_roots assert new_root in new_roots @@ -129,20 +105,16 @@ def test_after_merge(self, gen_graph): def test_empty_when_no_operations(self, gen_graph): """When no operations occurred, both arrays should be empty.""" - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - - fake_ts = fake_timestamp() - create_chunk( - graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV()}, ) # Query a time range in the future where no operations exist future = datetime.now(UTC) + timedelta(days=1) - old_roots, new_roots = get_proofread_root_ids(graph, start_time=future) + old_roots, new_roots = get_proofread_root_ids(cg, start_time=future) assert len(old_roots) == 0 assert len(new_roots) == 0 @@ -268,26 +240,19 @@ class TestGetActivatedEdges: @pytest.mark.timeout(30) def test_returns_numpy_array_after_merge(self, gen_graph): """After merging two isolated SVs, get_activated_edges returns a numpy array.""" - atomic_chunk_bounds = np.array([1, 1, 1]) - graph = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - - fake_ts = fake_timestamp() - sv0 = to_label(graph, 1, 0, 0, 0, 0) - sv1 = to_label(graph, 1, 0, 0, 0, 1) - - create_chunk( - graph, - vertices=[sv0, sv1], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=np.array([1, 1, 1]), + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) # Merge the two isolated supervoxels - result = graph.add_edges( + result = cg.add_edges( "TestUser", - [sv0, sv1], + [sv["a0"], sv["a1"]], affinities=[0.3], ) - activated = get_activated_edges(graph, result.operation_id) + activated = get_activated_edges(cg, result.operation_id) assert isinstance(activated, np.ndarray) diff --git a/pychunkedgraph/tests/graph/test_node_conversion.py b/pychunkedgraph/tests/graph/test_node_conversion.py index 9181a8146..d2d27c60c 100644 --- a/pychunkedgraph/tests/graph/test_node_conversion.py +++ b/pychunkedgraph/tests/graph/test_node_conversion.py @@ -1,7 +1,6 @@ import numpy as np import pytest -from ..helpers import to_label from ...graph import serializers diff --git a/pychunkedgraph/tests/graph/test_root_lock.py b/pychunkedgraph/tests/graph/test_root_lock.py index b96a28eb6..9b46d3cd4 100644 --- a/pychunkedgraph/tests/graph/test_root_lock.py +++ b/pychunkedgraph/tests/graph/test_root_lock.py @@ -7,39 +7,28 @@ import numpy as np import pytest -from ..helpers import create_chunk, to_label, fake_timestamp +from ..helpers import SV, build_graph from ...graph import exceptions from ...graph.locks import RootLock -from ...ingest.create.parent_layer import add_parent_chunk class TestRootLock: @pytest.fixture() def simple_graph(self, gen_graph): """Build a 2-chunk graph with a single edge, return (cg, root_id).""" - cg = gen_graph(n_layers=3) - fake_ts = fake_timestamp() - - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "b": SV(x=1)}, + edges=[("a0", "b", 0.5)], ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - root_id = cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) - return cg, root_id + root_id = cg.get_root(sv["a0"]) + return cg, sv, root_id @pytest.mark.timeout(30) def test_successful_lock_and_release(self, simple_graph): """Lock acquired successfully inside context, released after exit.""" - cg, root_id = simple_graph + cg, sv, root_id = simple_graph with RootLock(cg, np.array([root_id])) as lock: assert lock.lock_acquired @@ -53,7 +42,7 @@ def test_successful_lock_and_release(self, simple_graph): @pytest.mark.timeout(30) def test_lock_released_on_exception(self, simple_graph): """Lock should be released even when an exception occurs inside the context.""" - cg, root_id = simple_graph + cg, sv, root_id = simple_graph with pytest.raises(exceptions.PreconditionError): with RootLock(cg, np.array([root_id])) as lock: @@ -67,18 +56,18 @@ def test_lock_released_on_exception(self, simple_graph): @pytest.mark.timeout(30) def test_operation_with_lock_succeeds(self, simple_graph): """A real graph operation (split) should succeed while holding the lock.""" - cg, root_id = simple_graph + cg, sv, root_id = simple_graph # Use the high-level API which acquires locks internally result = cg.remove_edges( "test_user", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 0), + source_ids=sv["a0"], + sink_ids=sv["b"], mincut=False, ) assert len(result.new_root_ids) == 2 # After operation, locks should be released — verify we can re-acquire - new_root = cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) + new_root = cg.get_root(sv["a0"]) with RootLock(cg, np.array([new_root])) as lock: assert lock.lock_acquired diff --git a/pychunkedgraph/tests/graph/test_split.py b/pychunkedgraph/tests/graph/test_split.py index c6a33dfb1..a0a3336cd 100644 --- a/pychunkedgraph/tests/graph/test_split.py +++ b/pychunkedgraph/tests/graph/test_split.py @@ -4,11 +4,9 @@ import numpy as np import pytest -from ..helpers import create_chunk, to_label, fake_timestamp -from ...graph import ChunkedGraph +from ..helpers import SV, build_graph, fake_timestamp, assert_graph_unchanged from ...graph import exceptions from ...graph.misc import get_latest_roots -from ...ingest.create.parent_layer import add_parent_chunk class TestGraphSplit: @@ -23,57 +21,55 @@ def test_split_pair_same_chunk(self, gen_graph): └─────┘ └─────┘ """ - cg: ChunkedGraph = gen_graph(n_layers=2) - - # Preparation: Build Chunk A fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5)], + cg, sv = build_graph( + gen_graph, + n_layers=2, + supervoxels={"a0": SV(), "a1": SV(seg=1)}, + edges=[("a0", "a1", 0.5)], timestamp=fake_ts, ) # Split new_root_ids = cg.remove_edges( "Jane Doe", - source_ids=to_label(cg, 1, 0, 0, 0, 1), - sink_ids=to_label(cg, 1, 0, 0, 0, 0), + source_ids=sv["a1"], + sink_ids=sv["a0"], mincut=False, ).new_root_ids # verify new state assert len(new_root_ids) == 2 - assert cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) != cg.get_root( - to_label(cg, 1, 0, 0, 0, 1) + assert cg.get_root(sv["a0"]) != cg.get_root( + sv["a1"] ) leaves = np.unique( cg.get_subgraph( - [cg.get_root(to_label(cg, 1, 0, 0, 0, 0))], leaves_only=True + [cg.get_root(sv["a0"])], leaves_only=True ) ) - assert len(leaves) == 1 and to_label(cg, 1, 0, 0, 0, 0) in leaves + assert len(leaves) == 1 and sv["a0"] in leaves leaves = np.unique( cg.get_subgraph( - [cg.get_root(to_label(cg, 1, 0, 0, 0, 1))], leaves_only=True + [cg.get_root(sv["a1"])], leaves_only=True ) ) - assert len(leaves) == 1 and to_label(cg, 1, 0, 0, 0, 1) in leaves + assert len(leaves) == 1 and sv["a1"] in leaves # verify old state cg.cache = None assert cg.get_root( - to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_ts - ) == cg.get_root(to_label(cg, 1, 0, 0, 0, 1), time_stamp=fake_ts) + sv["a0"], time_stamp=fake_ts + ) == cg.get_root(sv["a1"], time_stamp=fake_ts) leaves = np.unique( cg.get_subgraph( - [cg.get_root(to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_ts)], + [cg.get_root(sv["a0"], time_stamp=fake_ts)], leaves_only=True, ) ) assert len(leaves) == 2 - assert to_label(cg, 1, 0, 0, 0, 0) in leaves - assert to_label(cg, 1, 0, 0, 0, 1) in leaves + assert sv["a0"] in leaves + assert sv["a1"] in leaves assert len(get_latest_roots(cg)) == 2 assert len(get_latest_roots(cg, fake_ts)) == 1 @@ -87,21 +83,19 @@ def test_split_nonexisting_edge(self, gen_graph): │ 3 │ │ 3 │ └─────┘ └─────┘ """ - cg = gen_graph(n_layers=2) - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], + cg, sv = build_graph( + gen_graph, + n_layers=2, + supervoxels={"a0": SV(), "a1": SV(seg=1), "a2": SV(seg=2)}, edges=[ - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5), - (to_label(cg, 1, 0, 0, 0, 2), to_label(cg, 1, 0, 0, 0, 1), 0.5), + ("a0", "a1", 0.5), + ("a2", "a1", 0.5), ], - timestamp=fake_ts, ) new_root_ids = cg.remove_edges( "Jane Doe", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 0, 0, 0, 2), + source_ids=sv["a0"], + sink_ids=sv["a2"], mincut=False, ).new_root_ids assert len(new_root_ids) == 1 @@ -116,59 +110,52 @@ def test_split_pair_neighboring_chunks(self, gen_graph): │ │ │ │ │ │ └─────┴─────┘ └─────┴─────┘ """ - cg: ChunkedGraph = gen_graph(n_layers=3) fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 1.0)], - timestamp=fake_ts, - ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 1.0)], + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "b": SV(x=1)}, + edges=[("a0", "b", 1.0)], timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) new_root_ids = cg.remove_edges( "Jane Doe", - source_ids=to_label(cg, 1, 1, 0, 0, 0), - sink_ids=to_label(cg, 1, 0, 0, 0, 0), + source_ids=sv["b"], + sink_ids=sv["a0"], mincut=False, ).new_root_ids # verify new state assert len(new_root_ids) == 2 - assert cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) != cg.get_root( - to_label(cg, 1, 1, 0, 0, 0) + assert cg.get_root(sv["a0"]) != cg.get_root( + sv["b"] ) leaves = np.unique( cg.get_subgraph( - [cg.get_root(to_label(cg, 1, 0, 0, 0, 0))], leaves_only=True + [cg.get_root(sv["a0"])], leaves_only=True ) ) - assert len(leaves) == 1 and to_label(cg, 1, 0, 0, 0, 0) in leaves + assert len(leaves) == 1 and sv["a0"] in leaves leaves = np.unique( cg.get_subgraph( - [cg.get_root(to_label(cg, 1, 1, 0, 0, 0))], leaves_only=True + [cg.get_root(sv["b"])], leaves_only=True ) ) - assert len(leaves) == 1 and to_label(cg, 1, 1, 0, 0, 0) in leaves + assert len(leaves) == 1 and sv["b"] in leaves # verify old state assert cg.get_root( - to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_ts - ) == cg.get_root(to_label(cg, 1, 1, 0, 0, 0), time_stamp=fake_ts) + sv["a0"], time_stamp=fake_ts + ) == cg.get_root(sv["b"], time_stamp=fake_ts) leaves = np.unique( cg.get_subgraph( - [cg.get_root(to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_ts)], + [cg.get_root(sv["a0"], time_stamp=fake_ts)], leaves_only=True, ) ) assert len(leaves) == 2 - assert to_label(cg, 1, 0, 0, 0, 0) in leaves - assert to_label(cg, 1, 1, 0, 0, 0) in leaves + assert sv["a0"] in leaves + assert sv["b"] in leaves assert len(get_latest_roots(cg)) == 2 assert len(get_latest_roots(cg, fake_ts)) == 1 @@ -182,39 +169,29 @@ def test_split_verify_cross_chunk_edges(self, gen_graph): | │ 2 │ │ | │ 2 │ │ └─────┴─────┴─────┘ └─────┴─────┴─────┘ """ - cg: ChunkedGraph = gen_graph(n_layers=4) fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 1)], + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"b0": SV(x=1), "b1": SV(x=1, seg=1), "c": SV(x=2)}, edges=[ - (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 2, 0, 0, 0), inf), - (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 1), 0.5), + ("b0", "c", inf), + ("b0", "b1", 0.5), ], timestamp=fake_ts, ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 2, 0, 0, 0)], - edges=[(to_label(cg, 1, 2, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), inf)], - timestamp=fake_ts, - ) - - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 3, [1, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - assert cg.get_root(to_label(cg, 1, 1, 0, 0, 0)) == cg.get_root( - to_label(cg, 1, 1, 0, 0, 1) + assert cg.get_root(sv["b0"]) == cg.get_root( + sv["b1"] ) - assert cg.get_root(to_label(cg, 1, 1, 0, 0, 0)) == cg.get_root( - to_label(cg, 1, 2, 0, 0, 0) + assert cg.get_root(sv["b0"]) == cg.get_root( + sv["c"] ) new_root_ids = cg.remove_edges( "Jane Doe", - source_ids=to_label(cg, 1, 1, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 1), + source_ids=sv["b0"], + sink_ids=sv["b1"], mincut=False, ).new_root_ids @@ -229,11 +206,11 @@ def test_split_verify_cross_chunk_edges(self, gen_graph): # verify new state assert len(new_root_ids) == 2 - assert cg.get_root(to_label(cg, 1, 1, 0, 0, 0)) != cg.get_root( - to_label(cg, 1, 1, 0, 0, 1) + assert cg.get_root(sv["b0"]) != cg.get_root( + sv["b1"] ) - assert cg.get_root(to_label(cg, 1, 1, 0, 0, 0)) == cg.get_root( - to_label(cg, 1, 2, 0, 0, 0) + assert cg.get_root(sv["b0"]) == cg.get_root( + sv["c"] ) assert len(get_latest_roots(cg)) == 2 @@ -249,58 +226,47 @@ def test_split_verify_loop(self, gen_graph): | │ 3 2━━┿━━6 │ | │ 3 2━━┿━━6 │ └─────┴────────┴─────┘ └─────┴────────┴─────┘ """ - cg: ChunkedGraph = gen_graph(n_layers=4) fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[ - to_label(cg, 1, 1, 0, 0, 0), - to_label(cg, 1, 1, 0, 0, 1), - to_label(cg, 1, 1, 0, 0, 2), - to_label(cg, 1, 1, 0, 0, 3), - ], - edges=[ - (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 2, 0, 0, 0), inf), - (to_label(cg, 1, 1, 0, 0, 1), to_label(cg, 1, 2, 0, 0, 1), inf), - (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 2), 0.5), - (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 3), 0.5), - ], - timestamp=fake_ts, - ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 2, 0, 0, 0), to_label(cg, 1, 2, 0, 0, 1)], + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={ + "b0": SV(x=1), + "b1": SV(x=1, seg=1), + "b2": SV(x=1, seg=2), + "b3": SV(x=1, seg=3), + "c0": SV(x=2), + "c1": SV(x=2, seg=1), + }, edges=[ - (to_label(cg, 1, 2, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), inf), - (to_label(cg, 1, 2, 0, 0, 1), to_label(cg, 1, 1, 0, 0, 1), inf), - (to_label(cg, 1, 2, 0, 0, 1), to_label(cg, 1, 2, 0, 0, 0), 0.5), + ("b0", "c0", inf), + ("b1", "c1", inf), + ("b0", "b2", 0.5), + ("b0", "b3", 0.5), + ("c1", "c0", 0.5), ], timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 3, [1, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - - assert cg.get_root(to_label(cg, 1, 1, 0, 0, 0)) == cg.get_root( - to_label(cg, 1, 1, 0, 0, 1) + assert cg.get_root(sv["b0"]) == cg.get_root( + sv["b1"] ) - assert cg.get_root(to_label(cg, 1, 1, 0, 0, 0)) == cg.get_root( - to_label(cg, 1, 2, 0, 0, 0) + assert cg.get_root(sv["b0"]) == cg.get_root( + sv["c0"] ) new_root_ids = cg.remove_edges( "Jane Doe", - source_ids=to_label(cg, 1, 1, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 2), + source_ids=sv["b0"], + sink_ids=sv["b2"], mincut=False, ).new_root_ids assert len(new_root_ids) == 2 new_root_ids = cg.remove_edges( "Jane Doe", - source_ids=to_label(cg, 1, 1, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 3), + source_ids=sv["b0"], + sink_ids=sv["b3"], mincut=False, ).new_root_ids assert len(new_root_ids) == 2 @@ -318,13 +284,10 @@ def test_split_pair_already_disconnected(self, gen_graph): │ │ │ │ └─────┘ └─────┘ """ - cg: ChunkedGraph = gen_graph(n_layers=2) - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) res_old = cg.client.read_all_rows() res_old.consume_all() @@ -332,8 +295,8 @@ def test_split_pair_already_disconnected(self, gen_graph): with pytest.raises(exceptions.PreconditionError): cg.remove_edges( "Jane Doe", - source_ids=to_label(cg, 1, 0, 0, 0, 1), - sink_ids=to_label(cg, 1, 0, 0, 0, 0), + source_ids=sv["a1"], + sink_ids=sv["a0"], mincut=False, ) @@ -356,40 +319,40 @@ def test_split_full_circle_to_triple_chain_same_chunk(self, gen_graph): │ ┗3┛ │ │ ┗3┛ │ └─────┘ └─────┘ """ - cg: ChunkedGraph = gen_graph(n_layers=2) fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[ - to_label(cg, 1, 0, 0, 0, 0), - to_label(cg, 1, 0, 0, 0, 1), - to_label(cg, 1, 0, 0, 0, 2), - ], + cg, sv = build_graph( + gen_graph, + n_layers=2, + supervoxels={ + "a0": SV(), + "a1": SV(seg=1), + "a2": SV(seg=2), + }, edges=[ - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 2), 0.5), - (to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 2), 0.5), - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.3), + ("a0", "a2", 0.5), + ("a1", "a2", 0.5), + ("a0", "a1", 0.3), ], timestamp=fake_ts, ) new_root_ids = cg.remove_edges( "Jane Doe", - source_ids=to_label(cg, 1, 0, 0, 0, 1), - sink_ids=to_label(cg, 1, 0, 0, 0, 0), + source_ids=sv["a1"], + sink_ids=sv["a0"], mincut=False, ).new_root_ids # verify new state assert len(new_root_ids) == 1 - assert cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) == new_root_ids[0] - assert cg.get_root(to_label(cg, 1, 0, 0, 0, 1)) == new_root_ids[0] - assert cg.get_root(to_label(cg, 1, 0, 0, 0, 2)) == new_root_ids[0] + assert cg.get_root(sv["a0"]) == new_root_ids[0] + assert cg.get_root(sv["a1"]) == new_root_ids[0] + assert cg.get_root(sv["a2"]) == new_root_ids[0] leaves = np.unique(cg.get_subgraph([new_root_ids[0]], leaves_only=True)) assert len(leaves) == 3 # verify old state old_root_id = cg.get_root( - to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_ts + sv["a0"], time_stamp=fake_ts ) assert new_root_ids[0] != old_root_id assert len(get_latest_roots(cg)) == 1 @@ -405,47 +368,36 @@ def test_split_full_circle_to_triple_chain_neighboring_chunks(self, gen_graph): │ ┗3━┿━━┛ │ │ ┗3━┿━━┛ │ └─────┴─────┘ └─────┴─────┘ """ - cg: ChunkedGraph = gen_graph(n_layers=3) fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], - edges=[ - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5), - (to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 1, 0, 0, 0), 0.5), - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.3), - ], - timestamp=fake_ts, - ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "a1": SV(seg=1), "b": SV(x=1)}, edges=[ - (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5), - (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.3), + ("a0", "a1", 0.5), + ("a1", "b", 0.5), + ("a0", "b", 0.3), ], timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - new_root_ids = cg.remove_edges( "Jane Doe", - source_ids=to_label(cg, 1, 1, 0, 0, 0), - sink_ids=to_label(cg, 1, 0, 0, 0, 0), + source_ids=sv["b"], + sink_ids=sv["a0"], mincut=False, ).new_root_ids # verify new state assert len(new_root_ids) == 1 - assert cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) == new_root_ids[0] - assert cg.get_root(to_label(cg, 1, 0, 0, 0, 1)) == new_root_ids[0] - assert cg.get_root(to_label(cg, 1, 1, 0, 0, 0)) == new_root_ids[0] + assert cg.get_root(sv["a0"]) == new_root_ids[0] + assert cg.get_root(sv["a1"]) == new_root_ids[0] + assert cg.get_root(sv["b"]) == new_root_ids[0] leaves = np.unique(cg.get_subgraph([new_root_ids[0]], leaves_only=True)) assert len(leaves) == 3 # verify old state old_root_id = cg.get_root( - to_label(cg, 1, 0, 0, 0, 0), time_stamp=fake_ts + sv["a0"], time_stamp=fake_ts ) assert new_root_ids[0] != old_root_id assert len(get_latest_roots(cg)) == 1 @@ -461,28 +413,20 @@ def test_split_same_node(self, gen_graph): │ │ └─────┘ """ - cg: ChunkedGraph = gen_graph(n_layers=2) - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + supervoxels={"a0": SV()}, ) - res_old = cg.client.read_all_rows() - res_old.consume_all() - with pytest.raises(exceptions.PreconditionError): - cg.remove_edges( - "Jane Doe", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 0, 0, 0, 0), - mincut=False, - ) - - res_new = cg.client.read_all_rows() - res_new.consume_all() - assert res_new.rows == res_old.rows + with assert_graph_unchanged(cg): + with pytest.raises(exceptions.PreconditionError): + cg.remove_edges( + "Jane Doe", + source_ids=sv["a0"], + sink_ids=sv["a0"], + mincut=False, + ) @pytest.mark.timeout(30) def test_split_pair_abstract_nodes(self, gen_graph): @@ -491,35 +435,19 @@ def test_split_pair_abstract_nodes(self, gen_graph): => Reject """ - cg: ChunkedGraph = gen_graph(n_layers=3) - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "b": SV(x=1)}, ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[], - timestamp=fake_ts, - ) - - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - res_old = cg.client.read_all_rows() - res_old.consume_all() - with pytest.raises((exceptions.PreconditionError, AssertionError)): - cg.remove_edges( - "Jane Doe", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 2, 1, 0, 0, 1), - mincut=False, - ) - - res_new = cg.client.read_all_rows() - res_new.consume_all() - assert res_new.rows == res_old.rows + with assert_graph_unchanged(cg): + with pytest.raises((exceptions.PreconditionError, AssertionError)): + cg.remove_edges( + "Jane Doe", + source_ids=sv["a0"], + sink_ids=cg.get_node_id(np.uint64(1), layer=2, x=1, y=0, z=0), + mincut=False, + ) @pytest.mark.timeout(30) def test_diagonal_connections(self, gen_graph): @@ -534,35 +462,23 @@ def test_diagonal_connections(self, gen_graph): │ C¹ │ D¹ │ └─────┴─────┘ """ - cg: ChunkedGraph = gen_graph(n_layers=3) - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={ + "a0": SV(), + "a1": SV(seg=1), + "b": SV(x=1), + "c": SV(y=1), + "d": SV(x=1, y=1), + }, edges=[ - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5), - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), inf), - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 1, 0, 0), inf), + ("a0", "a1", 0.5), + ("a0", "b", inf), + ("a0", "c", inf), + ("c", "d", inf), ], ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), inf)], - ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 1, 0, 0)], - edges=[ - (to_label(cg, 1, 0, 1, 0, 0), to_label(cg, 1, 1, 1, 0, 0), inf), - (to_label(cg, 1, 0, 1, 0, 0), to_label(cg, 1, 0, 0, 0, 0), inf), - ], - ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 1, 0, 0)], - edges=[(to_label(cg, 1, 1, 1, 0, 0), to_label(cg, 1, 0, 1, 0, 0), inf)], - ) - add_parent_chunk(cg, 3, [0, 0, 0], n_threads=1) rr = cg.range_read_chunk(chunk_id=cg.get_chunk_id(layer=3, x=0, y=0, z=0)) root_ids_t0 = list(rr.keys()) @@ -570,17 +486,17 @@ def test_diagonal_connections(self, gen_graph): new_roots = cg.remove_edges( "Jane Doe", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 0, 0, 0, 1), + source_ids=sv["a0"], + sink_ids=sv["a1"], mincut=False, ).new_root_ids assert len(new_roots) == 2 - assert cg.get_root(to_label(cg, 1, 1, 1, 0, 0)) == cg.get_root( - to_label(cg, 1, 0, 1, 0, 0) + assert cg.get_root(sv["d"]) == cg.get_root( + sv["c"] ) - assert cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) == cg.get_root( - to_label(cg, 1, 0, 0, 0, 0) + assert cg.get_root(sv["a0"]) == cg.get_root( + sv["a0"] ) @@ -599,34 +515,24 @@ def test_split_multi_layer_hierarchy_correctness(self, gen_graph): │ │ │ └─────┴─────┘ """ - cg = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "b": SV(x=1)}, + edges=[("a0", "b", 0.5)], ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) result = cg.remove_edges( "Jane Doe", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 0), + source_ids=sv["a0"], + sink_ids=sv["b"], mincut=False, ) assert len(result.new_root_ids) == 2 # Verify parent chain for both supervoxels - for sv in [to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0)]: - parents = cg.get_root(sv, get_all_parents=True) + for sv_id in [sv["a0"], sv["b"]]: + parents = cg.get_root(sv_id, get_all_parents=True) prev_layer = 1 for p in parents: layer = cg.get_chunk_layer(p) @@ -648,51 +554,34 @@ def test_split_creates_isolated_components_with_skip_connections(self, gen_graph │ 1━━┿━━2━━┿━━3 │ => split 1-2 => 1 becomes isolated, 2-3 stay connected └─────┴─────┴─────┘ """ - cg = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "b": SV(x=1), "c": SV(x=2)}, edges=[ - (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5), - (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 2, 0, 0, 0), inf), + ("a0", "b", 0.5), + ("b", "c", inf), ], - timestamp=fake_ts, - ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 2, 0, 0, 0)], - edges=[(to_label(cg, 1, 2, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), inf)], - timestamp=fake_ts, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 3, [1, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # All three should share a root before split - root_pre = cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) - assert root_pre == cg.get_root(to_label(cg, 1, 1, 0, 0, 0)) - assert root_pre == cg.get_root(to_label(cg, 1, 2, 0, 0, 0)) + root_pre = cg.get_root(sv["a0"]) + assert root_pre == cg.get_root(sv["b"]) + assert root_pre == cg.get_root(sv["c"]) # Split 1 from 2 result = cg.remove_edges( "Jane Doe", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 0), + source_ids=sv["a0"], + sink_ids=sv["b"], mincut=False, ) assert len(result.new_root_ids) == 2 # Node 1 should be isolated, nodes 2 and 3 should share a root - root1 = cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) - root2 = cg.get_root(to_label(cg, 1, 1, 0, 0, 0)) - root3 = cg.get_root(to_label(cg, 1, 2, 0, 0, 0)) + root1 = cg.get_root(sv["a0"]) + root2 = cg.get_root(sv["b"]) + root3 = cg.get_root(sv["c"]) assert root1 != root2 assert root2 == root3 diff --git a/pychunkedgraph/tests/graph/test_stale_edges.py b/pychunkedgraph/tests/graph/test_stale_edges.py index 9ac1aa08b..2638e60a6 100644 --- a/pychunkedgraph/tests/graph/test_stale_edges.py +++ b/pychunkedgraph/tests/graph/test_stale_edges.py @@ -5,12 +5,13 @@ """ +from math import inf + import numpy as np import pytest -from ..helpers import create_chunk, to_label, fake_timestamp +from ..helpers import SV, build_graph from ...graph.edges.stale import get_stale_nodes, get_new_nodes -from ...ingest.create.parent_layer import add_parent_chunk class TestStaleEdges: @@ -26,31 +27,21 @@ def test_stale_nodes_detected_after_split(self, gen_graph): │ │ │ └─────┴─────┘ """ - cg = gen_graph(n_layers=3) - fake_ts = fake_timestamp() - - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "b": SV(x=1)}, + edges=[("a0", "b", inf)], ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # Get old parents before edit - old_root = cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) + old_root = cg.get_root(sv["a0"]) # Split cg.remove_edges( "test_user", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 0), + source_ids=sv["a0"], + sink_ids=sv["b"], mincut=False, ) @@ -69,34 +60,24 @@ def test_no_stale_nodes_for_current_ids(self, gen_graph): │ │ │ └─────┴─────┘ """ - cg = gen_graph(n_layers=3) - fake_ts = fake_timestamp() - - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "b": SV(x=1)}, + edges=[("a0", "b", inf)], ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # Split cg.remove_edges( "test_user", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 0), + source_ids=sv["a0"], + sink_ids=sv["b"], mincut=False, ) # Current roots should not be stale - new_root_1 = cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) - new_root_2 = cg.get_root(to_label(cg, 1, 1, 0, 0, 0)) + new_root_1 = cg.get_root(sv["a0"]) + new_root_2 = cg.get_root(sv["b"]) stale = get_stale_nodes(cg, [new_root_1, new_root_2]) assert new_root_1 not in stale assert new_root_2 not in stale @@ -113,32 +94,22 @@ def test_get_new_nodes_resolves_to_correct_layer(self, gen_graph): │ │ │ └─────┴─────┘ """ - cg = gen_graph(n_layers=3) - fake_ts = fake_timestamp() - - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "b": SV(x=1)}, + edges=[("a0", "b", inf)], ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) # Get L2 parent of SV 1 before edit - sv1 = to_label(cg, 1, 0, 0, 0, 0) + sv1 = sv["a0"] old_l2_parent = cg.get_parent(sv1) # Split cg.remove_edges( "test_user", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 0), + source_ids=sv["a0"], + sink_ids=sv["b"], mincut=False, ) @@ -158,41 +129,21 @@ def test_no_stale_nodes_in_unaffected_region(self, gen_graph): │ │ │ │ └─────┴─────┴─────┘ """ - cg = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - # Chunk C - isolated node, not connected to A or B - create_chunk( - cg, - vertices=[to_label(cg, 1, 2, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV(), "b": SV(x=1), "c": SV(x=2)}, + edges=[("a0", "b", inf)], ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 3, [1, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - # Get the isolated node's root before edit - isolated_root = cg.get_root(to_label(cg, 1, 2, 0, 0, 0)) + isolated_root = cg.get_root(sv["c"]) # Split nodes 1 and 2 cg.remove_edges( "test_user", - source_ids=to_label(cg, 1, 0, 0, 0, 0), - sink_ids=to_label(cg, 1, 1, 0, 0, 0), + source_ids=sv["a0"], + sink_ids=sv["b"], mincut=False, ) @@ -212,23 +163,17 @@ def test_get_new_nodes_returns_self_for_non_stale(self, gen_graph): │ │ └─────┘ """ - cg = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV()}, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - sv = to_label(cg, 1, 0, 0, 0, 0) - l2_parent = cg.get_parent(sv) + sv0 = sv["a0"] + l2_parent = cg.get_parent(sv0) # get_new_nodes at layer 2 should return the same L2 parent - result = get_new_nodes(cg, np.array([sv], dtype=np.uint64), layer=2) + result = get_new_nodes(cg, np.array([sv0], dtype=np.uint64), layer=2) assert result[0] == l2_parent @pytest.mark.timeout(30) @@ -242,26 +187,16 @@ def test_get_stale_nodes_empty_for_fresh_graph(self, gen_graph): │ │ │ └─────┴─────┘ """ - cg = gen_graph(n_layers=3) - fake_ts = fake_timestamp() - - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "b": SV(x=1)}, + edges=[("a0", "b", inf)], ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - root = cg.get_root(to_label(cg, 1, 0, 0, 0, 0)) - l2_0 = cg.get_parent(to_label(cg, 1, 0, 0, 0, 0)) - l2_1 = cg.get_parent(to_label(cg, 1, 1, 0, 0, 0)) + root = cg.get_root(sv["a0"]) + l2_0 = cg.get_parent(sv["a0"]) + l2_1 = cg.get_parent(sv["b"]) # No edits have been performed, so all nodes should be non-stale stale = get_stale_nodes(cg, [root, l2_0, l2_1]) @@ -279,25 +214,15 @@ def test_get_new_nodes_multiple_svs(self, gen_graph): │ │ │ └─────┴─────┘ """ - cg = gen_graph(n_layers=3) - fake_ts = fake_timestamp() - - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "b": SV(x=1)}, + edges=[("a0", "b", inf)], ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - sv1 = to_label(cg, 1, 0, 0, 0, 0) - sv2 = to_label(cg, 1, 1, 0, 0, 0) + sv1 = sv["a0"] + sv2 = sv["b"] svs = np.array([sv1, sv2], dtype=np.uint64) result = get_new_nodes(cg, svs, layer=2) @@ -318,24 +243,19 @@ def test_get_new_nodes_with_duplicate_svs(self, gen_graph): │ │ └─────┘ """ - cg = gen_graph(n_layers=3) - fake_ts = fake_timestamp() - - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV()}, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - sv = to_label(cg, 1, 0, 0, 0, 0) - svs = np.array([sv, sv, sv], dtype=np.uint64) + sv0 = sv["a0"] + svs = np.array([sv0, sv0, sv0], dtype=np.uint64) result = get_new_nodes(cg, svs, layer=2) assert result.shape == (3,) # All should map to the same L2 parent - expected = cg.get_parent(sv) + expected = cg.get_parent(sv0) assert np.all(result == expected) @pytest.mark.timeout(30) @@ -350,19 +270,16 @@ def test_get_stale_nodes_with_l2_ids_after_merge(self, gen_graph): └─────┘ """ atomic_chunk_bounds = np.array([1, 1, 1]) - cg = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = fake_timestamp() - - sv0 = to_label(cg, 1, 0, 0, 0, 0) - sv1 = to_label(cg, 1, 0, 0, 0, 1) - - create_chunk( - cg, - vertices=[sv0, sv1], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=atomic_chunk_bounds, + supervoxels={"a0": SV(), "a1": SV(seg=1)}, ) + sv0 = sv["a0"] + sv1 = sv["a1"] + # Get L2 parents before merge (each SV has its own L2 parent) old_l2_0 = cg.get_parent(sv0) old_l2_1 = cg.get_parent(sv1) @@ -391,17 +308,15 @@ def test_get_stale_nodes_returns_numpy_array(self, gen_graph): └─────┘ """ atomic_chunk_bounds = np.array([1, 1, 1]) - cg = gen_graph(n_layers=2, atomic_chunk_bounds=atomic_chunk_bounds) - fake_ts = fake_timestamp() - - sv0 = to_label(cg, 1, 0, 0, 0, 0) - create_chunk( - cg, - vertices=[sv0], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=2, + atomic_chunk_bounds=atomic_chunk_bounds, + supervoxels={"a0": SV()}, ) + sv0 = sv["a0"] + root = cg.get_root(sv0) stale = get_stale_nodes(cg, [root]) assert isinstance(stale, np.ndarray) @@ -417,22 +332,17 @@ def test_get_new_nodes_at_root_layer(self, gen_graph): │ │ └─────┘ """ - cg = gen_graph(n_layers=4) - fake_ts = fake_timestamp() - - sv = to_label(cg, 1, 0, 0, 0, 0) - create_chunk( - cg, - vertices=[sv], - edges=[], - timestamp=fake_ts, + cg, sv = build_graph( + gen_graph, + n_layers=4, + supervoxels={"a0": SV()}, ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - add_parent_chunk(cg, 4, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - root = cg.get_root(sv) + sv0 = sv["a0"] + + root = cg.get_root(sv0) root_layer = cg.get_chunk_layer(root) - result = get_new_nodes(cg, np.array([sv], dtype=np.uint64), layer=root_layer) + result = get_new_nodes(cg, np.array([sv0], dtype=np.uint64), layer=root_layer) assert result.shape == (1,) assert result[0] == root diff --git a/pychunkedgraph/tests/graph/test_undo_redo.py b/pychunkedgraph/tests/graph/test_undo_redo.py index 3858312f0..7f18f0d8c 100644 --- a/pychunkedgraph/tests/graph/test_undo_redo.py +++ b/pychunkedgraph/tests/graph/test_undo_redo.py @@ -8,8 +8,7 @@ import numpy as np import pytest -from ..helpers import create_chunk, to_label, fake_timestamp -from ...ingest.create.parent_layer import add_parent_chunk +from ..helpers import SV, build_graph class TestUndoRedo: @@ -23,30 +22,19 @@ def two_chunk_graph(self, gen_graph): │ │ │ └─────┴─────┘ """ - cg = gen_graph(n_layers=3) - fake_ts = fake_timestamp() - - create_chunk( - cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], - edges=[(to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), 0.5)], - timestamp=fake_ts, + return build_graph( + gen_graph, + n_layers=3, + supervoxels={"a0": SV(), "b": SV(x=1)}, + edges=[("a0", "b", 0.5)], ) - create_chunk( - cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], - edges=[(to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), 0.5)], - timestamp=fake_ts, - ) - add_parent_chunk(cg, 3, [0, 0, 0], time_stamp=fake_ts, n_threads=1) - return cg @pytest.mark.timeout(30) def test_undo_split_restores_merged_root(self, two_chunk_graph): """Split two nodes, undo — nodes should share a common root again.""" - cg = two_chunk_graph - sv1 = to_label(cg, 1, 0, 0, 0, 0) - sv2 = to_label(cg, 1, 1, 0, 0, 0) + cg, sv = two_chunk_graph + sv1 = sv["a0"] + sv2 = sv["b"] # Initially, both SVs share a root assert cg.get_root(sv1) == cg.get_root(sv2) @@ -67,9 +55,9 @@ def test_undo_split_restores_merged_root(self, two_chunk_graph): @pytest.mark.timeout(30) def test_redo_restores_operation_result(self, two_chunk_graph): """Split, undo, redo the original split — state should match the post-split state.""" - cg = two_chunk_graph - sv1 = to_label(cg, 1, 0, 0, 0, 0) - sv2 = to_label(cg, 1, 1, 0, 0, 0) + cg, sv = two_chunk_graph + sv1 = sv["a0"] + sv2 = sv["b"] # Split split_result = cg.remove_edges( @@ -90,9 +78,9 @@ def test_redo_restores_operation_result(self, two_chunk_graph): @pytest.mark.timeout(30) def test_undo_preserves_subgraph_leaves(self, two_chunk_graph): """After undo, subgraph leaves should match the pre-operation state.""" - cg = two_chunk_graph - sv1 = to_label(cg, 1, 0, 0, 0, 0) - sv2 = to_label(cg, 1, 1, 0, 0, 0) + cg, sv = two_chunk_graph + sv1 = sv["a0"] + sv2 = sv["b"] # Get initial leaf set initial_root = cg.get_root(sv1) diff --git a/pychunkedgraph/tests/graph/test_utils_id_helpers.py b/pychunkedgraph/tests/graph/test_utils_id_helpers.py index 00fb30442..ee9208499 100644 --- a/pychunkedgraph/tests/graph/test_utils_id_helpers.py +++ b/pychunkedgraph/tests/graph/test_utils_id_helpers.py @@ -5,13 +5,13 @@ from pychunkedgraph.graph.utils import id_helpers from pychunkedgraph.graph.chunks import utils as chunk_utils -from ..helpers import to_label +from ..helpers import SV, label class TestGetSegmentIdLimit: def test_basic(self, gen_graph): graph = gen_graph(n_layers=4) - node_id = to_label(graph, 1, 0, 0, 0, 1) + node_id = label(graph, SV(seg=1)) limit = id_helpers.get_segment_id_limit(graph.meta, node_id) assert limit > 0 assert isinstance(limit, np.uint64) @@ -20,7 +20,7 @@ def test_basic(self, gen_graph): class TestGetSegmentId: def test_basic(self, gen_graph): graph = gen_graph(n_layers=4) - node_id = to_label(graph, 1, 0, 0, 0, 42) + node_id = label(graph, SV(seg=42)) seg_id = id_helpers.get_segment_id(graph.meta, node_id) assert seg_id == 42 diff --git a/pychunkedgraph/tests/helpers.py b/pychunkedgraph/tests/helpers.py index 0ca36d535..e3f97ca05 100644 --- a/pychunkedgraph/tests/helpers.py +++ b/pychunkedgraph/tests/helpers.py @@ -1,5 +1,6 @@ import threading from collections import namedtuple +from contextlib import contextmanager from datetime import datetime, timedelta, UTC from functools import reduce from unittest.mock import MagicMock @@ -168,6 +169,11 @@ def get_layer_chunk_bounds( BuiltGraph = namedtuple("BuiltGraph", ["cg", "sv"]) +def label(cg, sv, layer=1): + """Node id for supervoxel coordinate ``sv`` at ``layer`` (layer 1 = the supervoxel).""" + return to_label(cg, layer, sv.x, sv.y, sv.z, sv.seg) + + def build_graph(gen_graph, n_layers, supervoxels, edges=(), *, timestamp=None, atomic_chunk_bounds=None): """Build a test graph from named supervoxels and edges; parents derived, at one ts. @@ -195,6 +201,17 @@ def build_graph(gen_graph, n_layers, supervoxels, edges=(), *, timestamp=None, a return BuiltGraph(cg, sv) +@contextmanager +def assert_graph_unchanged(cg): + """Assert the (rejected) edit run inside the block leaves every stored row untouched.""" + res_old = cg.client.read_all_rows() + res_old.consume_all() + yield + res_new = cg.client.read_all_rows() + res_new.consume_all() + assert res_new.rows == res_old.rows + + class RowKeyLockRegistry: """Thread-safe in-memory stand-in for kvdbclient's row-key lock API. diff --git a/pychunkedgraph/tests/ingest/test_ingest_cross_edges.py b/pychunkedgraph/tests/ingest/test_ingest_cross_edges.py index fef5c3a3f..f2923bea1 100644 --- a/pychunkedgraph/tests/ingest/test_ingest_cross_edges.py +++ b/pychunkedgraph/tests/ingest/test_ingest_cross_edges.py @@ -13,7 +13,7 @@ ) from pychunkedgraph.graph import basetypes -from ..helpers import create_chunk, to_label, fake_timestamp +from ..helpers import SV, label, create_chunk, fake_timestamp from ...ingest.create.parent_layer import add_parent_chunk @@ -72,7 +72,7 @@ def test_no_cross_edges(self, gen_graph): create_chunk( graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], + vertices=[label(graph, SV())], edges=[], timestamp=fake_ts, ) @@ -87,17 +87,17 @@ def test_with_cross_edges(self, gen_graph): create_chunk( graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], + vertices=[label(graph, SV())], edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), + (label(graph, SV()), label(graph, SV(x=1)), inf), ], timestamp=fake_ts, ) create_chunk( graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], + vertices=[label(graph, SV(x=1))], edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), + (label(graph, SV(x=1)), label(graph, SV()), inf), ], timestamp=fake_ts, ) @@ -114,7 +114,7 @@ def test_no_atomic_chunks_returns_empty(self, gen_graph): create_chunk( cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], + vertices=[label(cg, SV())], edges=[], timestamp=fake_ts, ) @@ -137,9 +137,9 @@ def test_basic_cross_edges(self, gen_graph): # Chunk A (0,0,0): sv 0 connected cross-chunk to chunk B create_chunk( cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], + vertices=[label(cg, SV())], edges=[ - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), inf), + (label(cg, SV()), label(cg, SV(x=1)), inf), ], timestamp=fake_ts, ) @@ -147,9 +147,9 @@ def test_basic_cross_edges(self, gen_graph): # Chunk B (1,0,0): sv 0 connected cross-chunk to chunk A create_chunk( cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], + vertices=[label(cg, SV(x=1))], edges=[ - (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), inf), + (label(cg, SV(x=1)), label(cg, SV()), inf), ], timestamp=fake_ts, ) @@ -179,20 +179,20 @@ def test_multiple_cross_edges(self, gen_graph): # Chunk A: two SVs, each cross-chunk connected create_chunk( cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1)], + vertices=[label(cg, SV()), label(cg, SV(seg=1))], edges=[ - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 1), 0.5), - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), inf), - (to_label(cg, 1, 0, 0, 0, 1), to_label(cg, 1, 1, 0, 0, 1), inf), + (label(cg, SV()), label(cg, SV(seg=1)), 0.5), + (label(cg, SV()), label(cg, SV(x=1)), inf), + (label(cg, SV(seg=1)), label(cg, SV(x=1, seg=1)), inf), ], timestamp=fake_ts, ) create_chunk( cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 1)], + vertices=[label(cg, SV(x=1)), label(cg, SV(x=1, seg=1))], edges=[ - (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), inf), - (to_label(cg, 1, 1, 0, 0, 1), to_label(cg, 1, 0, 0, 0, 1), inf), + (label(cg, SV(x=1)), label(cg, SV()), inf), + (label(cg, SV(x=1, seg=1)), label(cg, SV(seg=1)), inf), ], timestamp=fake_ts, ) @@ -217,18 +217,18 @@ def test_cross_edges_layer4(self, gen_graph): # SV at L1 [1,0,0] - on the right boundary of L3 [0,0,0] create_chunk( cg, - vertices=[to_label(cg, 1, 1, 0, 0, 0)], + vertices=[label(cg, SV(x=1))], edges=[ - (to_label(cg, 1, 1, 0, 0, 0), to_label(cg, 1, 2, 0, 0, 0), inf), + (label(cg, SV(x=1)), label(cg, SV(x=2)), inf), ], timestamp=fake_ts, ) # SV at L1 [2,0,0] - on the left boundary of L3 [1,0,0] create_chunk( cg, - vertices=[to_label(cg, 1, 2, 0, 0, 0)], + vertices=[label(cg, SV(x=2))], edges=[ - (to_label(cg, 1, 2, 0, 0, 0), to_label(cg, 1, 1, 0, 0, 0), inf), + (label(cg, SV(x=2)), label(cg, SV(x=1)), inf), ], timestamp=fake_ts, ) @@ -260,18 +260,18 @@ def test_no_threads_with_cross_edges(self, gen_graph): # SV at L1 [0,0,0] with cross edge to [2,0,0] (layer-3 cross edge) create_chunk( cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], + vertices=[label(cg, SV())], edges=[ - (to_label(cg, 1, 0, 0, 0, 0), to_label(cg, 1, 2, 0, 0, 0), inf), + (label(cg, SV()), label(cg, SV(x=2)), inf), ], timestamp=fake_ts, ) # SV at L1 [2,0,0] create_chunk( cg, - vertices=[to_label(cg, 1, 2, 0, 0, 0)], + vertices=[label(cg, SV(x=2))], edges=[ - (to_label(cg, 1, 2, 0, 0, 0), to_label(cg, 1, 0, 0, 0, 0), inf), + (label(cg, SV(x=2)), label(cg, SV()), inf), ], timestamp=fake_ts, ) @@ -294,7 +294,7 @@ def test_no_threads_empty_chunk(self, gen_graph): create_chunk( cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], + vertices=[label(cg, SV())], edges=[], timestamp=fake_ts, ) @@ -315,7 +315,7 @@ def test_no_cross_edges_returns_empty(self, gen_graph): create_chunk( cg, - vertices=[to_label(cg, 1, 0, 0, 0, 0)], + vertices=[label(cg, SV())], edges=[], timestamp=fake_ts, ) diff --git a/pychunkedgraph/tests/ingest/test_ingest_parent_layer.py b/pychunkedgraph/tests/ingest/test_ingest_parent_layer.py index 290f27250..f122204ce 100644 --- a/pychunkedgraph/tests/ingest/test_ingest_parent_layer.py +++ b/pychunkedgraph/tests/ingest/test_ingest_parent_layer.py @@ -5,7 +5,7 @@ import numpy as np import pytest -from ..helpers import create_chunk, to_label, fake_timestamp +from ..helpers import SV, label, create_chunk, fake_timestamp from ...ingest.create.parent_layer import add_parent_chunk @@ -16,9 +16,9 @@ def test_single_thread(self, gen_graph): create_chunk( graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1)], + vertices=[label(graph, SV()), label(graph, SV(seg=1))], edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 1), 0.5), + (label(graph, SV()), label(graph, SV(seg=1)), 0.5), ], timestamp=fake_ts, ) @@ -27,7 +27,7 @@ def test_single_thread(self, gen_graph): add_parent_chunk(graph, 3, [0, 0, 0], n_threads=1) # Verify parent was created - sv = to_label(graph, 1, 0, 0, 0, 0) + sv = label(graph, SV()) parent = graph.get_parent(sv) assert parent is not None assert graph.get_chunk_layer(parent) == 2 @@ -38,17 +38,17 @@ def test_multi_chunk(self, gen_graph): create_chunk( graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], + vertices=[label(graph, SV())], edges=[ - (to_label(graph, 1, 0, 0, 0, 0), to_label(graph, 1, 1, 0, 0, 0), inf), + (label(graph, SV()), label(graph, SV(x=1)), inf), ], timestamp=fake_ts, ) create_chunk( graph, - vertices=[to_label(graph, 1, 1, 0, 0, 0)], + vertices=[label(graph, SV(x=1))], edges=[ - (to_label(graph, 1, 1, 0, 0, 0), to_label(graph, 1, 0, 0, 0, 0), inf), + (label(graph, SV(x=1)), label(graph, SV()), inf), ], timestamp=fake_ts, ) @@ -57,6 +57,6 @@ def test_multi_chunk(self, gen_graph): add_parent_chunk(graph, 4, [0, 0, 0], n_threads=1) # Both SVs should share a root - root0 = graph.get_root(to_label(graph, 1, 0, 0, 0, 0)) - root1 = graph.get_root(to_label(graph, 1, 1, 0, 0, 0)) + root0 = graph.get_root(label(graph, SV())) + root1 = graph.get_root(label(graph, SV(x=1))) assert root0 == root1 diff --git a/pychunkedgraph/tests/meshing/test_setup.py b/pychunkedgraph/tests/meshing/test_setup.py index b31d54517..1f9c0d89d 100644 --- a/pychunkedgraph/tests/meshing/test_setup.py +++ b/pychunkedgraph/tests/meshing/test_setup.py @@ -2,7 +2,7 @@ from ...ingest.create.parent_layer import add_parent_chunk from ...meshing.setup import derive_initial_ts -from ..helpers import create_chunk, to_label, fake_timestamp +from ..helpers import SV, label, create_chunk, fake_timestamp def test_derive_initial_ts_uses_stamped_boundary(gen_graph): @@ -19,7 +19,7 @@ def test_derive_initial_ts_after_root_build(gen_graph): fake_ts = fake_timestamp() create_chunk( graph, - vertices=[to_label(graph, 1, 0, 0, 0, 0)], + vertices=[label(graph, SV())], edges=[], timestamp=fake_ts, ) From 37c4e81f56a7945d60721cbe0798a9636a9444a4 Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Mon, 29 Jun 2026 00:23:16 +0000 Subject: [PATCH 09/12] test: group edit-operation tests into an edits package Move the edit-operation tests into tests/graph/edits/ with a coverage README, and lift the duplicated split/merge/undo setup graphs into shared edits/conftest.py fixtures. Co-Authored-By: Claude --- pychunkedgraph/tests/graph/edits/README.md | 45 +++++++ pychunkedgraph/tests/graph/edits/__init__.py | 0 pychunkedgraph/tests/graph/edits/conftest.py | 49 ++++++++ .../graph/{ => edits}/test_edits_extended.py | 2 +- .../tests/graph/{ => edits}/test_edits_sv.py | 0 .../tests/graph/{ => edits}/test_merge.py | 2 +- .../graph/{ => edits}/test_merge_split.py | 4 +- .../tests/graph/{ => edits}/test_mincut.py | 4 +- .../tests/graph/{ => edits}/test_multicut.py | 4 +- .../tests/graph/{ => edits}/test_operation.py | 111 ++++-------------- .../tests/graph/{ => edits}/test_split.py | 6 +- .../tests/graph/{ => edits}/test_undo_redo.py | 37 ++---- 12 files changed, 138 insertions(+), 126 deletions(-) create mode 100644 pychunkedgraph/tests/graph/edits/README.md create mode 100644 pychunkedgraph/tests/graph/edits/__init__.py create mode 100644 pychunkedgraph/tests/graph/edits/conftest.py rename pychunkedgraph/tests/graph/{ => edits}/test_edits_extended.py (96%) rename pychunkedgraph/tests/graph/{ => edits}/test_edits_sv.py (100%) rename pychunkedgraph/tests/graph/{ => edits}/test_merge.py (99%) rename pychunkedgraph/tests/graph/{ => edits}/test_merge_split.py (97%) rename pychunkedgraph/tests/graph/{ => edits}/test_mincut.py (98%) rename pychunkedgraph/tests/graph/{ => edits}/test_multicut.py (96%) rename pychunkedgraph/tests/graph/{ => edits}/test_operation.py (90%) rename pychunkedgraph/tests/graph/{ => edits}/test_split.py (99%) rename pychunkedgraph/tests/graph/{ => edits}/test_undo_redo.py (73%) diff --git a/pychunkedgraph/tests/graph/edits/README.md b/pychunkedgraph/tests/graph/edits/README.md new file mode 100644 index 000000000..91a1be41b --- /dev/null +++ b/pychunkedgraph/tests/graph/edits/README.md @@ -0,0 +1,45 @@ +# Edit-operation tests + +The proofreading edit operations and their atomicity: + +- **merge** (`add_edges`) +- **split** / **mincut** / **multicut** (`remove_edges`) +- **undo / redo** +- **supervoxel split** + +Graphs are built with the `build_graph` factory (or explicit chunk construction +where the produced row structure is itself the subject). The core invariant under +test is that a **rejected** edit leaves every stored row untouched — expressed with +the `assert_graph_unchanged` context manager. + +## Coverage matrix + +| Operation | Topologies exercised | Accepted paths | Rejected paths (graph unchanged) | +|---|---|---|---| +| merge | same chunk, neighboring chunks, disconnected chunks, skip-connection across layers | new shared parent; cross-chunk edge re-mapping; multi-layer hierarchy correctness | already-connected pair; self-loop; abstract (non-L1) node; indirect cycle | +| split | same chunk, neighboring, disconnected; full-circle ↔ triple-chain; skip-connection components | component separation; pre-edit state still readable at the old timestamp | nonexistent edge; cross-layer / abstract-node edge | +| mincut | regular link, no link, previously-removed link, indivisible (`inf`) link | source/sink separation respecting affinities | disconnected source/sink; cut that isolates; cut disrespecting source/sink | +| multicut | path-augmented cut over real supervoxel data | multi-source / multi-sink cut | supervoxel-split-required is signalled | +| undo / redo | split ↔ merge inversion | state restoration; chain resolution; subgraph leaves preserved; edge re-validation on undo | — | +| supervoxel split | supervoxel-level split over real watershed data | split applied | — | + +Per-operation **atomicity** (graph unchanged after a rejected edit) is asserted for +the merge, split, and mincut rejection cases. + +## Gaps (→ edit-failure-handling follow-up) + +These are the known coverage holes; each is filled by the edit-failure-handling and +single-edit-rollback follow-ups, ticking off rows here as they land. + +- **Full-hierarchy invariant verifier after every edit** — layer monotonicity, + parent↔child bidirectionality, one root per connected component, no duplicate + children per layer, cross-edge validity. Today a rejected edit only compares the + stored rows, not these structural invariants. +- **Fail-before-persist proven for every bad input** — zero hierarchy mutation *and* + no `SUCCESS` operation-log row written. Two non-blocking rejection cases currently + `warn` instead of asserting the graph is unchanged. +- **Mid-write crash + recovery** — interrupt the operation mid-persist and prove + `repair_operation` restores a consistent state. +- **Concurrent-write isolation** — two edits on the same vs disjoint roots. +- **OCDBT-backed edit paths.** +- **The multicut scenario that currently skips itself** on a too-small graph. diff --git a/pychunkedgraph/tests/graph/edits/__init__.py b/pychunkedgraph/tests/graph/edits/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/pychunkedgraph/tests/graph/edits/conftest.py b/pychunkedgraph/tests/graph/edits/conftest.py new file mode 100644 index 000000000..38cef3539 --- /dev/null +++ b/pychunkedgraph/tests/graph/edits/conftest.py @@ -0,0 +1,49 @@ +"""Shared fixtures for the edit-operation tests.""" + +import pytest + +from ...helpers import SV, build_graph + + +@pytest.fixture() +def connected_pair(gen_graph): + """Two supervoxels joined by an edge across neighboring chunks (3 layers). + + ┌─────┬─────┐ + │ A¹ │ B¹ │ + │ a━━┿━━b │ + └─────┴─────┘ + """ + return build_graph( + gen_graph, + n_layers=3, + supervoxels={"a": SV(), "b": SV(x=1)}, + edges=[("a", "b", 0.5)], + ) + + +@pytest.fixture() +def split_then_merge_ops(connected_pair): + """`connected_pair` split apart then merged back; returns (cg, merge_op_id, split_op_id).""" + cg, sv = connected_pair + split_result = cg.remove_edges( + "test_user", source_ids=sv["a"], sink_ids=sv["b"], mincut=False + ) + merge_result = cg.add_edges( + "test_user", + atomic_edges=[[sv["a"], sv["b"]]], + source_coords=[0, 0, 0], + sink_coords=[0, 0, 0], + ) + return cg, merge_result.operation_id, split_result.operation_id + + +@pytest.fixture() +def split_then_undo(connected_pair): + """`connected_pair` split, then the split undone; returns (cg, split_op_id, undo_result).""" + cg, sv = connected_pair + split_result = cg.remove_edges( + "test_user", source_ids=sv["a"], sink_ids=sv["b"], mincut=False + ) + undo_result = cg.undo_operation("test_user", split_result.operation_id) + return cg, split_result.operation_id, undo_result diff --git a/pychunkedgraph/tests/graph/test_edits_extended.py b/pychunkedgraph/tests/graph/edits/test_edits_extended.py similarity index 96% rename from pychunkedgraph/tests/graph/test_edits_extended.py rename to pychunkedgraph/tests/graph/edits/test_edits_extended.py index 5cfc997e7..c8b798fc1 100644 --- a/pychunkedgraph/tests/graph/test_edits_extended.py +++ b/pychunkedgraph/tests/graph/edits/test_edits_extended.py @@ -8,7 +8,7 @@ from pychunkedgraph.graph.edits import flip_ids from pychunkedgraph.graph import basetypes -from ..helpers import SV, build_graph +from ...helpers import SV, build_graph class TestFlipIds: diff --git a/pychunkedgraph/tests/graph/test_edits_sv.py b/pychunkedgraph/tests/graph/edits/test_edits_sv.py similarity index 100% rename from pychunkedgraph/tests/graph/test_edits_sv.py rename to pychunkedgraph/tests/graph/edits/test_edits_sv.py diff --git a/pychunkedgraph/tests/graph/test_merge.py b/pychunkedgraph/tests/graph/edits/test_merge.py similarity index 99% rename from pychunkedgraph/tests/graph/test_merge.py rename to pychunkedgraph/tests/graph/edits/test_merge.py index 32f16d305..a8e29bcaa 100644 --- a/pychunkedgraph/tests/graph/test_merge.py +++ b/pychunkedgraph/tests/graph/edits/test_merge.py @@ -4,7 +4,7 @@ import numpy as np import pytest -from ..helpers import SV, build_graph, assert_graph_unchanged +from ...helpers import SV, build_graph, assert_graph_unchanged class TestGraphMerge: diff --git a/pychunkedgraph/tests/graph/test_merge_split.py b/pychunkedgraph/tests/graph/edits/test_merge_split.py similarity index 97% rename from pychunkedgraph/tests/graph/test_merge_split.py rename to pychunkedgraph/tests/graph/edits/test_merge_split.py index 7fdf3190c..cc4d0d96d 100644 --- a/pychunkedgraph/tests/graph/test_merge_split.py +++ b/pychunkedgraph/tests/graph/edits/test_merge_split.py @@ -3,8 +3,8 @@ import numpy as np import pytest -from ..helpers import SV, build_graph -from ...graph import types +from ...helpers import SV, build_graph +from ....graph import types class TestGraphMergeSplit: diff --git a/pychunkedgraph/tests/graph/test_mincut.py b/pychunkedgraph/tests/graph/edits/test_mincut.py similarity index 98% rename from pychunkedgraph/tests/graph/test_mincut.py rename to pychunkedgraph/tests/graph/edits/test_mincut.py index 7879ccd78..b6aba702d 100644 --- a/pychunkedgraph/tests/graph/test_mincut.py +++ b/pychunkedgraph/tests/graph/edits/test_mincut.py @@ -3,8 +3,8 @@ import numpy as np import pytest -from ..helpers import SV, build_graph, assert_graph_unchanged -from ...graph import exceptions +from ...helpers import SV, build_graph, assert_graph_unchanged +from ....graph import exceptions class TestGraphMinCut: diff --git a/pychunkedgraph/tests/graph/test_multicut.py b/pychunkedgraph/tests/graph/edits/test_multicut.py similarity index 96% rename from pychunkedgraph/tests/graph/test_multicut.py rename to pychunkedgraph/tests/graph/edits/test_multicut.py index 4edd5962f..57e815f79 100644 --- a/pychunkedgraph/tests/graph/test_multicut.py +++ b/pychunkedgraph/tests/graph/edits/test_multicut.py @@ -1,8 +1,8 @@ import numpy as np import pytest -from ...graph.edges import Edges -from ...graph.cutting import Cut, SvSplitRequired, run_multicut +from ....graph.edges import Edges +from ....graph.cutting import Cut, SvSplitRequired, run_multicut class TestGraphMultiCut: diff --git a/pychunkedgraph/tests/graph/test_operation.py b/pychunkedgraph/tests/graph/edits/test_operation.py similarity index 90% rename from pychunkedgraph/tests/graph/test_operation.py rename to pychunkedgraph/tests/graph/edits/test_operation.py index 811f18b7a..03f14a1e7 100644 --- a/pychunkedgraph/tests/graph/test_operation.py +++ b/pychunkedgraph/tests/graph/edits/test_operation.py @@ -10,9 +10,9 @@ import numpy as np import pytest -from ..helpers import SV, build_graph -from ...graph import attributes -from ...graph.operation import ( +from ...helpers import SV, build_graph +from ....graph import attributes +from ....graph.operation import ( GraphEditOperation, MergeOperation, MulticutOperation, @@ -20,7 +20,7 @@ RedoOperation, UndoOperation, ) -from ...graph.exceptions import ( +from ....graph.exceptions import ( PreconditionError, PostconditionError, SupervoxelSplitRequiredError, @@ -65,58 +65,34 @@ def _build_cross_chunk(gen_graph): class TestOperationFromLogRecord: """Test that GraphEditOperation.from_log_record correctly identifies operation types.""" - @pytest.fixture() - def merged_graph(self, gen_graph): - """Build a simple 2-chunk graph and perform a merge, returning (cg, operation_id).""" - cg, sv = build_graph( - gen_graph, - n_layers=3, - supervoxels={"a": SV(), "b": SV(x=1)}, - edges=[("a", "b", 0.5)], - ) - - # Split first to get two separate roots - split_result = cg.remove_edges( - "test_user", source_ids=sv["a"], sink_ids=sv["b"], mincut=False - ) - - # Now merge them back - merge_result = cg.add_edges( - "test_user", - atomic_edges=[[sv["a"], sv["b"]]], - source_coords=[0, 0, 0], - sink_coords=[0, 0, 0], - ) - return cg, merge_result.operation_id, split_result.operation_id - @pytest.mark.timeout(30) - def test_merge_log_record_type(self, merged_graph): + def test_merge_log_record_type(self, split_then_merge_ops): """MergeOperation should be correctly identified from a real merge log record.""" - cg, merge_op_id, _ = merged_graph + cg, merge_op_id, _ = split_then_merge_ops log_record, _ = cg.client.read_log_entry(merge_op_id) op_type = GraphEditOperation.get_log_record_type(log_record) assert op_type is MergeOperation @pytest.mark.timeout(30) - def test_split_log_record_type(self, merged_graph): + def test_split_log_record_type(self, split_then_merge_ops): """SplitOperation should be correctly identified from a real split log record.""" - cg, _, split_op_id = merged_graph + cg, _, split_op_id = split_then_merge_ops log_record, _ = cg.client.read_log_entry(split_op_id) op_type = GraphEditOperation.get_log_record_type(log_record) assert op_type is SplitOperation @pytest.mark.timeout(30) - def test_merge_from_log_record(self, merged_graph): + def test_merge_from_log_record(self, split_then_merge_ops): """from_log_record should return a MergeOperation for a real merge log.""" - cg, merge_op_id, _ = merged_graph + cg, merge_op_id, _ = split_then_merge_ops log_record, _ = cg.client.read_log_entry(merge_op_id) graph_op = GraphEditOperation.from_log_record(cg, log_record) assert isinstance(graph_op, MergeOperation) @pytest.mark.timeout(30) - def test_split_from_log_record(self, merged_graph): + def test_split_from_log_record(self, split_then_merge_ops): """from_log_record should return a SplitOperation for a real split log.""" - cg, _, split_op_id = merged_graph + cg, _, split_op_id = split_then_merge_ops log_record, _ = cg.client.read_log_entry(split_op_id) graph_op = GraphEditOperation.from_log_record(cg, log_record) assert isinstance(graph_op, SplitOperation) @@ -133,31 +109,10 @@ def test_unknown_log_record_fails(self, gen_graph): class TestOperationInversion: """Test that operation inversion produces the correct inverse type and edges.""" - @pytest.fixture() - def split_and_merge_ops(self, gen_graph): - """Build graph, split, merge -- return (cg, merge_op_id, split_op_id).""" - cg, sv = build_graph( - gen_graph, - n_layers=3, - supervoxels={"a": SV(), "b": SV(x=1)}, - edges=[("a", "b", 0.5)], - ) - - split_result = cg.remove_edges( - "test_user", source_ids=sv["a"], sink_ids=sv["b"], mincut=False - ) - merge_result = cg.add_edges( - "test_user", - atomic_edges=[[sv["a"], sv["b"]]], - source_coords=[0, 0, 0], - sink_coords=[0, 0, 0], - ) - return cg, merge_result.operation_id, split_result.operation_id - @pytest.mark.timeout(30) - def test_invert_merge_produces_split(self, split_and_merge_ops): + def test_invert_merge_produces_split(self, split_then_merge_ops): """Inverse of a MergeOperation should be a SplitOperation with matching edges.""" - cg, merge_op_id, _ = split_and_merge_ops + cg, merge_op_id, _ = split_then_merge_ops log_record, _ = cg.client.read_log_entry(merge_op_id) merge_op = GraphEditOperation.from_log_record(cg, log_record) inverted = merge_op.invert() @@ -165,9 +120,9 @@ def test_invert_merge_produces_split(self, split_and_merge_ops): assert np.all(np.equal(merge_op.added_edges, inverted.removed_edges)) @pytest.mark.timeout(30) - def test_invert_split_produces_merge(self, split_and_merge_ops): + def test_invert_split_produces_merge(self, split_then_merge_ops): """Inverse of a SplitOperation should be a MergeOperation with matching edges.""" - cg, _, split_op_id = split_and_merge_ops + cg, _, split_op_id = split_then_merge_ops log_record, _ = cg.client.read_log_entry(split_op_id) split_op = GraphEditOperation.from_log_record(cg, log_record) inverted = split_op.invert() @@ -178,45 +133,27 @@ def test_invert_split_produces_merge(self, split_and_merge_ops): class TestUndoRedoChainResolution: """Test undo/redo chain resolution through real graph operations.""" - @pytest.fixture() - def graph_with_undo(self, gen_graph): - """Build graph, perform split, then undo -- return (cg, split_op_id, undo_result).""" - cg, sv = build_graph( - gen_graph, - n_layers=3, - supervoxels={"a": SV(), "b": SV(x=1)}, - edges=[("a", "b", 0.5)], - ) - - # Split - split_result = cg.remove_edges( - "test_user", source_ids=sv["a"], sink_ids=sv["b"], mincut=False - ) - # Undo the split (= merge) - undo_result = cg.undo_operation("test_user", split_result.operation_id) - return cg, split_result.operation_id, undo_result - @pytest.mark.timeout(30) - def test_undo_log_record_type(self, graph_with_undo): + def test_undo_log_record_type(self, split_then_undo): """Undo operation log record should be identified as UndoOperation.""" - cg, _, undo_result = graph_with_undo + cg, _, undo_result = split_then_undo log_record, _ = cg.client.read_log_entry(undo_result.operation_id) op_type = GraphEditOperation.get_log_record_type(log_record) assert op_type is UndoOperation @pytest.mark.timeout(30) - def test_undo_from_log_resolves_correctly(self, graph_with_undo): + def test_undo_from_log_resolves_correctly(self, split_then_undo): """from_log_record on an undo record should resolve the chain to an UndoOperation.""" - cg, split_op_id, undo_result = graph_with_undo + cg, split_op_id, undo_result = split_then_undo log_record, _ = cg.client.read_log_entry(undo_result.operation_id) resolved_op = GraphEditOperation.from_log_record(cg, log_record) # Undo of a split -> UndoOperation whose inverse is a MergeOperation assert isinstance(resolved_op, UndoOperation) @pytest.mark.timeout(30) - def test_redo_after_undo(self, graph_with_undo): + def test_redo_after_undo(self, split_then_undo): """Redo of the original split (after undo) should produce a RedoOperation log.""" - cg, split_op_id, undo_result = graph_with_undo + cg, split_op_id, undo_result = split_then_undo # Redo the original split (which was undone) redo_result = cg.redo_operation("test_user", split_op_id) @@ -226,9 +163,9 @@ def test_redo_after_undo(self, graph_with_undo): assert isinstance(resolved_op, RedoOperation) @pytest.mark.timeout(30) - def test_undo_redo_chain_prevention(self, graph_with_undo): + def test_undo_redo_chain_prevention(self, split_then_undo): """Direct UndoOperation/RedoOperation on undo/redo targets should raise ValueError.""" - cg, _, undo_result = graph_with_undo + cg, _, undo_result = split_then_undo # Direct UndoOperation on an undo record should fail with pytest.raises(ValueError): diff --git a/pychunkedgraph/tests/graph/test_split.py b/pychunkedgraph/tests/graph/edits/test_split.py similarity index 99% rename from pychunkedgraph/tests/graph/test_split.py rename to pychunkedgraph/tests/graph/edits/test_split.py index a0a3336cd..40fb32994 100644 --- a/pychunkedgraph/tests/graph/test_split.py +++ b/pychunkedgraph/tests/graph/edits/test_split.py @@ -4,9 +4,9 @@ import numpy as np import pytest -from ..helpers import SV, build_graph, fake_timestamp, assert_graph_unchanged -from ...graph import exceptions -from ...graph.misc import get_latest_roots +from ...helpers import SV, build_graph, fake_timestamp, assert_graph_unchanged +from ....graph import exceptions +from ....graph.misc import get_latest_roots class TestGraphSplit: diff --git a/pychunkedgraph/tests/graph/test_undo_redo.py b/pychunkedgraph/tests/graph/edits/test_undo_redo.py similarity index 73% rename from pychunkedgraph/tests/graph/test_undo_redo.py rename to pychunkedgraph/tests/graph/edits/test_undo_redo.py index 7f18f0d8c..cc9136a8d 100644 --- a/pychunkedgraph/tests/graph/test_undo_redo.py +++ b/pychunkedgraph/tests/graph/edits/test_undo_redo.py @@ -8,32 +8,13 @@ import numpy as np import pytest -from ..helpers import SV, build_graph - class TestUndoRedo: - @pytest.fixture() - def two_chunk_graph(self, gen_graph): - """ - Build a 2-chunk graph with edge between SVs 1 and 2. - ┌─────┬─────┐ - │ A¹ │ B¹ │ - │ 1━━┿━━2 │ - │ │ │ - └─────┴─────┘ - """ - return build_graph( - gen_graph, - n_layers=3, - supervoxels={"a0": SV(), "b": SV(x=1)}, - edges=[("a0", "b", 0.5)], - ) - @pytest.mark.timeout(30) - def test_undo_split_restores_merged_root(self, two_chunk_graph): + def test_undo_split_restores_merged_root(self, connected_pair): """Split two nodes, undo — nodes should share a common root again.""" - cg, sv = two_chunk_graph - sv1 = sv["a0"] + cg, sv = connected_pair + sv1 = sv["a"] sv2 = sv["b"] # Initially, both SVs share a root @@ -53,10 +34,10 @@ def test_undo_split_restores_merged_root(self, two_chunk_graph): assert cg.get_root(sv1) == cg.get_root(sv2) @pytest.mark.timeout(30) - def test_redo_restores_operation_result(self, two_chunk_graph): + def test_redo_restores_operation_result(self, connected_pair): """Split, undo, redo the original split — state should match the post-split state.""" - cg, sv = two_chunk_graph - sv1 = sv["a0"] + cg, sv = connected_pair + sv1 = sv["a"] sv2 = sv["b"] # Split @@ -76,10 +57,10 @@ def test_redo_restores_operation_result(self, two_chunk_graph): assert cg.get_root(sv1) != cg.get_root(sv2) @pytest.mark.timeout(30) - def test_undo_preserves_subgraph_leaves(self, two_chunk_graph): + def test_undo_preserves_subgraph_leaves(self, connected_pair): """After undo, subgraph leaves should match the pre-operation state.""" - cg, sv = two_chunk_graph - sv1 = sv["a0"] + cg, sv = connected_pair + sv1 = sv["a"] sv2 = sv["b"] # Get initial leaf set From d56df7490e7c9f0d8c2d8639112ede4bfd24aea9 Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Mon, 29 Jun 2026 00:38:51 +0000 Subject: [PATCH 10/12] test: keep stale-edge tests on a between-chunk edge Restore the original finite affinity (a split test must exercise the same edge type); the migration had switched it to an inf cross-chunk edge. Co-Authored-By: Claude --- pychunkedgraph/tests/graph/test_stale_edges.py | 13 ++++++------- pychunkedgraph/tests/helpers.py | 2 +- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/pychunkedgraph/tests/graph/test_stale_edges.py b/pychunkedgraph/tests/graph/test_stale_edges.py index 2638e60a6..4b6b00070 100644 --- a/pychunkedgraph/tests/graph/test_stale_edges.py +++ b/pychunkedgraph/tests/graph/test_stale_edges.py @@ -5,7 +5,6 @@ """ -from math import inf import numpy as np import pytest @@ -31,7 +30,7 @@ def test_stale_nodes_detected_after_split(self, gen_graph): gen_graph, n_layers=3, supervoxels={"a0": SV(), "b": SV(x=1)}, - edges=[("a0", "b", inf)], + edges=[("a0", "b", 0.5)], ) # Get old parents before edit @@ -64,7 +63,7 @@ def test_no_stale_nodes_for_current_ids(self, gen_graph): gen_graph, n_layers=3, supervoxels={"a0": SV(), "b": SV(x=1)}, - edges=[("a0", "b", inf)], + edges=[("a0", "b", 0.5)], ) # Split @@ -98,7 +97,7 @@ def test_get_new_nodes_resolves_to_correct_layer(self, gen_graph): gen_graph, n_layers=3, supervoxels={"a0": SV(), "b": SV(x=1)}, - edges=[("a0", "b", inf)], + edges=[("a0", "b", 0.5)], ) # Get L2 parent of SV 1 before edit @@ -133,7 +132,7 @@ def test_no_stale_nodes_in_unaffected_region(self, gen_graph): gen_graph, n_layers=4, supervoxels={"a0": SV(), "b": SV(x=1), "c": SV(x=2)}, - edges=[("a0", "b", inf)], + edges=[("a0", "b", 0.5)], ) # Get the isolated node's root before edit @@ -191,7 +190,7 @@ def test_get_stale_nodes_empty_for_fresh_graph(self, gen_graph): gen_graph, n_layers=3, supervoxels={"a0": SV(), "b": SV(x=1)}, - edges=[("a0", "b", inf)], + edges=[("a0", "b", 0.5)], ) root = cg.get_root(sv["a0"]) @@ -218,7 +217,7 @@ def test_get_new_nodes_multiple_svs(self, gen_graph): gen_graph, n_layers=3, supervoxels={"a0": SV(), "b": SV(x=1)}, - edges=[("a0", "b", inf)], + edges=[("a0", "b", 0.5)], ) sv1 = sv["a0"] diff --git a/pychunkedgraph/tests/helpers.py b/pychunkedgraph/tests/helpers.py index e3f97ca05..c051559c0 100644 --- a/pychunkedgraph/tests/helpers.py +++ b/pychunkedgraph/tests/helpers.py @@ -178,7 +178,7 @@ def build_graph(gen_graph, n_layers, supervoxels, edges=(), *, timestamp=None, a """Build a test graph from named supervoxels and edges; parents derived, at one ts. supervoxels maps a name to its (x, y, z, seg) atomic coordinate; edges are - (name, name, affinity). Returns BuiltGraph(cg, sv, ts); sv maps each name to its node id. + (name, name, affinity). Returns BuiltGraph(cg, sv); sv maps each name to its node id. """ bounds = np.array([]) if atomic_chunk_bounds is None else atomic_chunk_bounds cg = gen_graph(n_layers=n_layers, atomic_chunk_bounds=bounds) From 82bbac4fba0cd37fcb4ced7a2a49fb6e2e444df7 Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Mon, 29 Jun 2026 02:35:02 +0000 Subject: [PATCH 11/12] test: cut real-time sleeps from the slowest lock/error tests Shrink the lock-acquire backoff where the failed acquire is only setup, lower the test lock-expiry, and skip the doomed best-effort error-artifact write that dominated the error-path test. Assertions unchanged. Co-Authored-By: Claude --- pychunkedgraph/tests/conftest.py | 2 +- pychunkedgraph/tests/graph/edits/test_operation.py | 5 +++-- pychunkedgraph/tests/graph/test_locks.py | 7 ++++++- 3 files changed, 10 insertions(+), 4 deletions(-) diff --git a/pychunkedgraph/tests/conftest.py b/pychunkedgraph/tests/conftest.py index 4f6e3e9b6..b76a30a5f 100644 --- a/pychunkedgraph/tests/conftest.py +++ b/pychunkedgraph/tests/conftest.py @@ -173,7 +173,7 @@ def _cgraph(request, n_layers=10, atomic_chunk_bounds: np.ndarray = np.array([]) "FANOUT": 2, "SPATIAL_BITS": 10, "ID_PREFIX": "", - "ROOT_LOCK_EXPIRY": timedelta(seconds=5), + "ROOT_LOCK_EXPIRY": timedelta(seconds=1), }, "backend_client": backend_client, "ingest_config": {}, diff --git a/pychunkedgraph/tests/graph/edits/test_operation.py b/pychunkedgraph/tests/graph/edits/test_operation.py index 03f14a1e7..ff7cacd21 100644 --- a/pychunkedgraph/tests/graph/edits/test_operation.py +++ b/pychunkedgraph/tests/graph/edits/test_operation.py @@ -684,10 +684,11 @@ def test_execute_assertion_error_clears_cache(self, gen_graph): cg, sv = _build_two_sv_disconnected(gen_graph) sv0, sv1 = sv["a0"], sv["a1"] - # Mock _apply to raise RuntimeError + # Mock _apply to raise RuntimeError; skip the best-effort error-artifact + # dump (a doomed GCS write to the read-only watershed bucket in tests). with patch.object( MergeOperation, "_apply", side_effect=RuntimeError("test runtime error") - ): + ), patch("pychunkedgraph.graph.err_dump.dump_err_artifact", return_value=None): with pytest.raises(RuntimeError, match="test runtime error"): cg.add_edges("test_user", [sv0, sv1], affinities=[0.3]) diff --git a/pychunkedgraph/tests/graph/test_locks.py b/pychunkedgraph/tests/graph/test_locks.py index a6e9eabc2..6ec1fb431 100644 --- a/pychunkedgraph/tests/graph/test_locks.py +++ b/pychunkedgraph/tests/graph/test_locks.py @@ -748,13 +748,18 @@ def test_partial_acquire_released_on_failure(self, monkeypatch): assert len(registry._held) == 1 assert next(iter(registry._held)) == _l2_chunk_lock_row_key(np.uint64(2)) - def test_privileged_mode_skips_acquire(self): + def test_privileged_mode_skips_acquire(self, monkeypatch): """Replay path: indefinite cells from the crashed op are still set, so a normal temporal acquire would refuse. Privileged mode bypasses the acquire entirely — the indefinite cells are the de-facto lock and the inner `IndefiniteL2ChunkLock(privileged=True)` releases them on exit. """ + # The contrast `normal` acquire below is expected to fail; shrink its + # backoff so the exhausted retry loop doesn't sleep the full schedule. + monkeypatch.setattr(L2ChunkLock, "_MAX_ACQUIRE_ATTEMPTS", 3) + monkeypatch.setattr(L2ChunkLock, "_ACQUIRE_BACKOFF_BASE_SEC", 0.01) + registry = RowKeyLockRegistry() # Crashed op's indefinite cells block a normal temporal acquire. crashed_op = np.uint64(42) From 2ca6fc9d386c1d00053522e28bbeb4729f9f77df Mon Sep 17 00:00:00 2001 From: Akhilesh Halageri Date: Mon, 29 Jun 2026 16:20:03 +0000 Subject: [PATCH 12/12] refactor: discover test backends from kvdbclient (>=0.8.0) Every test parametrizes over kvdbclient_testing.backends(); the per-backend branches, the bigtable emulator bootstrap, and the hbase mock are gone, now shipped by kvdbclient. bootstrap() resolves the config class via get_config_class(). Requires kvdbclient>=0.8.0. Co-Authored-By: Claude --- pychunkedgraph/ingest/utils.py | 7 +- pychunkedgraph/tests/conftest.py | 180 ++------ pychunkedgraph/tests/hbase_mock_server.py | 473 ---------------------- requirements.in | 2 +- requirements.txt | 2 +- 5 files changed, 35 insertions(+), 629 deletions(-) delete mode 100644 pychunkedgraph/tests/hbase_mock_server.py diff --git a/pychunkedgraph/ingest/utils.py b/pychunkedgraph/ingest/utils.py index fbd55dedc..87b048be4 100644 --- a/pychunkedgraph/ingest/utils.py +++ b/pychunkedgraph/ingest/utils.py @@ -8,7 +8,7 @@ from typing import Dict, Generator, Tuple import numpy as np -from kvdbclient import BigTableConfig, HBaseConfig +from kvdbclient import get_config_class from rich import box from rich.console import Group from rich.live import Live @@ -62,10 +62,7 @@ def bootstrap( TEST_RUN=test_run, ) backend_type = config["backend_client"].get("TYPE", "bigtable") - if backend_type == "hbase": - client_config = HBaseConfig(**config["backend_client"]["CONFIG"]) - else: - client_config = BigTableConfig(**config["backend_client"]["CONFIG"]) + client_config = get_config_class(backend_type)(**config["backend_client"]["CONFIG"]) client_info = BackendClientInfo(backend_type, client_config) graph_config = GraphConfig( diff --git a/pychunkedgraph/tests/conftest.py b/pychunkedgraph/tests/conftest.py index b76a30a5f..819072a97 100644 --- a/pychunkedgraph/tests/conftest.py +++ b/pychunkedgraph/tests/conftest.py @@ -1,167 +1,49 @@ -import atexit -import os -import signal -import subprocess -from functools import partial +from contextlib import ExitStack from datetime import timedelta +from functools import partial +import numpy as np import pytest -# Skip the old monolithic test file if it still exists (e.g., during branch transitions) -collect_ignore = ["test_uncategorized.py"] -import numpy as np -from google.auth import credentials -from google.cloud import bigtable +from kvdbclient_testing import backends as _backend_harnesses -from ..ingest.utils import bootstrap -from ..graph.edges import Edges from ..graph.chunkedgraph import ChunkedGraph - +from ..graph.edges import Edges +from ..ingest.utils import bootstrap from .helpers import ( CloudVolumeMock, TensorStoreMock, - mock_ws_info, get_layer_chunk_bounds, + mock_ws_info, ) -from .hbase_mock_server import start_hbase_mock_server - -_emulator_proc = None -_emulator_cleaned = False - - -def _delete_test_table(graph, backend="bigtable"): - """Test-only: delete the backing table for cleanup.""" - if backend == "bigtable": - graph.client._admin_table.delete() - else: - resp = graph.client._session.delete(graph.client._table_url("/schema")) - if resp.status_code not in (200, 404): - resp.raise_for_status() - - -def _cleanup_emulator(): - global _emulator_cleaned - if _emulator_cleaned or _emulator_proc is None: - return - _emulator_cleaned = True - try: - pgid = os.getpgid(_emulator_proc.pid) - os.killpg(pgid, signal.SIGTERM) - try: - _emulator_proc.wait(timeout=5) - except subprocess.TimeoutExpired: - os.killpg(pgid, signal.SIGKILL) - _emulator_proc.wait(timeout=5) - except (ProcessLookupError, OSError, ChildProcessError): - pass - # Hard kill cbtemulator in case it survived the process group signal - subprocess.run(["pkill", "-9", "cbtemulator"], stderr=subprocess.DEVNULL) - - -def setup_emulator_env(): - bt_env_init = subprocess.run( - ["gcloud", "beta", "emulators", "bigtable", "env-init"], stdout=subprocess.PIPE - ) - os.environ["BIGTABLE_EMULATOR_HOST"] = ( - bt_env_init.stdout.decode("utf-8").strip().split("=")[-1] - ) - - c = bigtable.Client( - project="IGNORE_ENVIRONMENT_PROJECT", - credentials=credentials.AnonymousCredentials(), - admin=True, - ) - t = c.instance("emulated_instance").table("emulated_table") - - try: - t.create() - return True - except Exception as err: - print("Bigtable Emulator not yet ready: %s" % err) - return False - - -@pytest.fixture(scope="session", autouse=True) -def bigtable_emulator(request): - global _emulator_proc, _emulator_cleaned - from time import sleep - - _emulator_cleaned = False - - # Kill any leftover emulator processes from previous runs - subprocess.run(["pkill", "-9", "cbtemulator"], stderr=subprocess.DEVNULL) - - # Start Emulator - _emulator_proc = subprocess.Popen( - [ - "gcloud", - "beta", - "emulators", - "bigtable", - "start", - "--host-port=localhost:8539", - ], - preexec_fn=os.setsid, - stdout=subprocess.PIPE, - ) - - # Register atexit handler as safety net for abnormal exits - atexit.register(_cleanup_emulator) - - # Wait for Emulator to start up - print("Waiting for BigTables Emulator to start up...", end="") - retries = 5 - while retries > 0: - if setup_emulator_env() is True: - break - else: - retries -= 1 - sleep(5) - - if retries == 0: - print( - "\nCouldn't start Bigtable Emulator. Make sure it is installed correctly." - ) - _cleanup_emulator() - exit(1) - - request.addfinalizer(_cleanup_emulator) +# Skip the old monolithic test file if it still exists (e.g., during branch transitions) +collect_ignore = ["test_uncategorized.py"] -@pytest.fixture(scope="session") -def hbase_emulator(): - """Start an in-process mock HBase REST server for the session.""" - _data, server, port = start_hbase_mock_server() - yield port - server.shutdown() +# Backends are discovered from KVDbClient; a new backend there is tested with no change here. +_HARNESSES = {h.name: h for h in _backend_harnesses()} -@pytest.fixture(scope="function", params=["bigtable", "hbase"]) -def gen_graph(request, bigtable_emulator, hbase_emulator): - backend = request.param +@pytest.fixture(scope="session", autouse=True) +def _backend_servers(): + """Start every available backend's local instance once for the session.""" + with ExitStack() as stack: + handles = {} + for name, harness in _HARNESSES.items(): + if harness.available(): + handles[name] = stack.enter_context(harness.server()) + yield handles + + +@pytest.fixture(scope="function", params=list(_HARNESSES)) +def gen_graph(request, _backend_servers): + name = request.param + if name not in _backend_servers: + pytest.skip(f"backend {name!r} unavailable in this environment") + harness = _HARNESSES[name] + handle = _backend_servers[name] def _cgraph(request, n_layers=10, atomic_chunk_bounds: np.ndarray = np.array([])): - if backend == "bigtable": - backend_client = { - "TYPE": "bigtable", - "CONFIG": { - "ADMIN": True, - "READ_ONLY": False, - "PROJECT": "IGNORE_ENVIRONMENT_PROJECT", - "INSTANCE": "emulated_instance", - "CREDENTIALS": credentials.AnonymousCredentials(), - "MAX_ROW_KEY_COUNT": 1000, - }, - } - else: - backend_client = { - "TYPE": "hbase", - "CONFIG": { - "BASE_URL": f"http://127.0.0.1:{hbase_emulator}", - "MAX_ROW_KEY_COUNT": 1000, - }, - } - config = { "data_source": { "EDGES": "gs://chunked-graph/minnie65_0/edges", @@ -175,7 +57,7 @@ def _cgraph(request, n_layers=10, atomic_chunk_bounds: np.ndarray = np.array([]) "ID_PREFIX": "", "ROOT_LOCK_EXPIRY": timedelta(seconds=1), }, - "backend_client": backend_client, + "backend_client": harness.backend_client(handle), "ingest_config": {}, } @@ -193,7 +75,7 @@ def _cgraph(request, n_layers=10, atomic_chunk_bounds: np.ndarray = np.array([]) graph.create() def fin(): - _delete_test_table(graph, backend) + harness.delete_table(graph) request.addfinalizer(fin) return graph diff --git a/pychunkedgraph/tests/hbase_mock_server.py b/pychunkedgraph/tests/hbase_mock_server.py deleted file mode 100644 index 0b9bb53b8..000000000 --- a/pychunkedgraph/tests/hbase_mock_server.py +++ /dev/null @@ -1,473 +0,0 @@ -"""In-process mock HBase REST (Stargate) server for testing. - -Implements the subset of the HBase REST API used by -``pychunkedgraph.graph.client.hbase.client.Client``. -Runs in a daemon thread on a random port using stdlib ``http.server``. -""" - -import base64 -import json -import struct -import threading -import time -import uuid -from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler -from urllib.parse import urlparse, parse_qs - -# --------------------------------------------------------------------------- -# Data structures -# --------------------------------------------------------------------------- - - -class ScannerState: - __slots__ = ("rows", "position", "batch_size") - - def __init__(self, rows, batch_size): - self.rows = rows # list of (row_key_bytes, cells_dict) - self.position = 0 - self.batch_size = batch_size - - -class HBaseMockData: - """Thread-safe shared state for the mock server.""" - - def __init__(self): - self.lock = threading.Lock() - # tables[table_name][row_key: bytes][col_spec: str] = [(value: bytes, ts_ms: int), ...] - self.tables: dict = {} - self.table_schemas: dict = {} - self.scanners: dict = {} - self._scanner_counter = 0 - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _b64enc(data: bytes) -> str: - return base64.b64encode(data).decode("ascii") - - -def _b64dec(s: str) -> bytes: - return base64.b64decode(s) - - -def _now_ms() -> int: - return int(time.time() * 1000) - - -def _insert_cell(cell_list: list, value: bytes, ts_ms: int): - """Insert (value, ts_ms) keeping descending ts order.""" - entry = (value, ts_ms) - for i, (_, t) in enumerate(cell_list): - if ts_ms >= t: - cell_list.insert(i, entry) - return - cell_list.append(entry) - - -# --------------------------------------------------------------------------- -# Request handler -# --------------------------------------------------------------------------- - - -def _make_handler_class(data: HBaseMockData): - """Create a handler class bound to the given shared data.""" - - class Handler(BaseHTTPRequestHandler): - - def log_message(self, format, *args): - pass # silence per-request logging - - # -- routing helpers ------------------------------------------------ - - def _parse(self): - parsed = urlparse(self.path) - parts = [p for p in parsed.path.split("/") if p] - query = parse_qs(parsed.query) - # flatten single-valued params - qflat = {k: v[0] if len(v) == 1 else v for k, v in query.items()} - return parts, qflat - - def _read_body(self) -> bytes: - length = int(self.headers.get("Content-Length", 0)) - return self.rfile.read(length) if length else b"" - - def _json_body(self): - raw = self._read_body() - return json.loads(raw) if raw else {} - - def _send_json(self, obj, code=200): - body = json.dumps(obj).encode() - self.send_response(code) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(body))) - self.end_headers() - self.wfile.write(body) - - def _send_empty(self, code=200): - self.send_response(code) - self.send_header("Content-Length", "0") - self.end_headers() - - def _send_bytes( - self, raw: bytes, code=200, content_type="application/octet-stream" - ): - self.send_response(code) - self.send_header("Content-Type", content_type) - self.send_header("Content-Length", str(len(raw))) - self.end_headers() - self.wfile.write(raw) - - # -- GET ------------------------------------------------------------ - - def do_GET(self): - parts, query = self._parse() - if len(parts) < 2: - return self._send_empty(404) - - table = parts[0] - - # GET /{table}/schema - if parts[1] == "schema": - with data.lock: - if table in data.table_schemas: - return self._send_json(data.table_schemas[table]) - return self._send_empty(404) - - # GET /{table}/scanner/{id} - if parts[1] == "scanner" and len(parts) >= 3: - scanner_id = parts[2] - with data.lock: - sc = data.scanners.get(scanner_id) - if sc is None: - return self._send_empty(404) - if sc.position >= len(sc.rows): - return self._send_empty(204) - end = sc.position + sc.batch_size - batch = sc.rows[sc.position : end] - sc.position = end - return self._send_json(self._rows_to_cellset(batch)) - - # GET /{table}/{key_b64}[/{col_spec}] - row_key = _b64dec(parts[1]) - col_specs = parts[2].split(",") if len(parts) >= 3 else None - max_versions = int(query.get("v", "1")) - ts_from = int(query["ts.from"]) if "ts.from" in query else None - ts_to = int(query["ts.to"]) if "ts.to" in query else None - - with data.lock: - tbl = data.tables.get(table) - if tbl is None or row_key not in tbl: - return self._send_empty(404) - row_data = tbl[row_key] - cells = self._filter_cells( - row_data, col_specs, ts_from, ts_to, max_versions - ) - if not cells: - return self._send_empty(404) - return self._send_json(self._rows_to_cellset([(row_key, cells)])) - - # -- PUT ------------------------------------------------------------ - - def do_PUT(self): - parts, query = self._parse() - if len(parts) < 2: - return self._send_empty(400) - - table = parts[0] - - # PUT /{table}/schema - if parts[1] == "schema": - body = self._json_body() - with data.lock: - data.table_schemas[table] = body - data.tables.setdefault(table, {}) - return self._send_empty(200) - - # PUT /{table}/scanner - if parts[1] == "scanner": - return self._handle_create_scanner(table) - - # check=put / check=delete - if "check" in query: - if query["check"] == "put": - return self._handle_check_and_put(table, parts, query) - if query["check"] == "delete": - return self._handle_check_and_delete(table, parts, query) - - # Batch write (PUT /{table}/{any_key}) - body = self._json_body() - with data.lock: - tbl = data.tables.setdefault(table, {}) - for row in body.get("Row", []): - rk = _b64dec(row["key"]) - row_dict = tbl.setdefault(rk, {}) - for cell in row.get("Cell", []): - col = _b64dec(cell["column"]).decode("utf-8") - val = _b64dec(cell["$"]) - ts = cell.get("timestamp", _now_ms()) - cell_list = row_dict.setdefault(col, []) - _insert_cell(cell_list, val, ts) - return self._send_empty(200) - - # -- POST (atomic increment) --------------------------------------- - - def do_POST(self): - parts, _ = self._parse() - if len(parts) < 3: - return self._send_empty(400) - table = parts[0] - row_key = _b64dec(parts[1]) - col_spec = parts[2] - - body = self._json_body() - # Extract increment value from CellSet body - inc_val = 0 - for row in body.get("Row", []): - for cell in row.get("Cell", []): - raw = _b64dec(cell["$"]) - inc_val = struct.unpack(">q", raw)[0] - break - break - - with data.lock: - tbl = data.tables.setdefault(table, {}) - row_dict = tbl.setdefault(row_key, {}) - cell_list = row_dict.get(col_spec, []) - current = 0 - if cell_list: - # latest value is first (newest) - try: - current = struct.unpack(">q", cell_list[0][0])[0] - except struct.error: - current = 0 - new_val = current + inc_val - new_bytes = struct.pack(">q", new_val) - ts = _now_ms() - new_list = [(new_bytes, ts)] - row_dict[col_spec] = new_list - - return self._send_bytes(new_bytes) - - # -- DELETE --------------------------------------------------------- - - def do_DELETE(self): - parts, _ = self._parse() - if len(parts) < 2: - return self._send_empty(400) - - table = parts[0] - - # DELETE /{table}/schema - if parts[1] == "schema": - with data.lock: - data.tables.pop(table, None) - data.table_schemas.pop(table, None) - return self._send_empty(200) - - # DELETE /{table}/scanner/{id} - if parts[1] == "scanner" and len(parts) >= 3: - with data.lock: - data.scanners.pop(parts[2], None) - return self._send_empty(200) - - row_key = _b64dec(parts[1]) - - with data.lock: - tbl = data.tables.get(table) - if tbl is None: - return self._send_empty(200) - - if len(parts) == 2: - # DELETE row - tbl.pop(row_key, None) - elif len(parts) == 3: - # DELETE column - col_spec = parts[2] - row_dict = tbl.get(row_key, {}) - row_dict.pop(col_spec, None) - if not row_dict: - tbl.pop(row_key, None) - elif len(parts) >= 4: - # DELETE cell version - col_spec = parts[2] - ts_ms = int(parts[3]) - row_dict = tbl.get(row_key, {}) - cell_list = row_dict.get(col_spec, []) - row_dict[col_spec] = [(v, t) for v, t in cell_list if t != ts_ms] - if not row_dict.get(col_spec): - row_dict.pop(col_spec, None) - if not row_dict: - tbl.pop(row_key, None) - - return self._send_empty(200) - - # -- Scanner -------------------------------------------------------- - - def _handle_create_scanner(self, table): - body = self._json_body() - start_row = _b64dec(body["startRow"]) if "startRow" in body else b"" - end_row = _b64dec(body["endRow"]) if "endRow" in body else None - batch_size = body.get("batch", 100) - col_filter = body.get("column") # list of col_spec strings or None - start_time = body.get("startTime") # ms, inclusive - end_time = body.get("endTime") # ms, exclusive - - with data.lock: - tbl = data.tables.get(table, {}) - filtered = [] - for rk in sorted(tbl.keys()): - if rk < start_row: - continue - if end_row is not None and rk >= end_row: - continue - cells = self._filter_cells( - tbl[rk], col_filter, start_time, end_time, max_versions=None - ) - if cells: - filtered.append((rk, cells)) - - data._scanner_counter += 1 - scanner_id = str(data._scanner_counter) - data.scanners[scanner_id] = ScannerState(filtered, batch_size) - - host, port = self.server.server_address - loc = f"http://{host}:{port}/{table}/scanner/{scanner_id}" - self.send_response(201) - self.send_header("Location", loc) - self.send_header("Content-Length", "0") - self.end_headers() - - # -- Check-and-put -------------------------------------------------- - - def _handle_check_and_put(self, table, parts, query): - body = self._json_body() - row_key = _b64dec(parts[1]) - check_col_spec = parts[2] # the column to check - - row_cells = body.get("Row", [{}])[0].get("Cell", []) - - if len(row_cells) >= 2: - # First cell is the check cell, second is the put cell - check_value = _b64dec(row_cells[0]["$"]) - put_cell = row_cells[1] - else: - # Single cell: check that column does NOT exist - check_value = None - put_cell = row_cells[0] if row_cells else None - - with data.lock: - tbl = data.tables.setdefault(table, {}) - row_dict = tbl.get(row_key, {}) - current_cells = row_dict.get(check_col_spec, []) - - if check_value is None: - # Column must not exist - if current_cells: - return self._send_empty(304) - else: - # Latest value must match - if not current_cells or current_cells[0][0] != check_value: - return self._send_empty(304) - - # Condition met - apply the put - if put_cell: - put_col = _b64dec(put_cell["column"]).decode("utf-8") - put_val = _b64dec(put_cell["$"]) - put_ts = put_cell.get("timestamp", _now_ms()) - row_dict = tbl.setdefault(row_key, {}) - cell_list = row_dict.setdefault(put_col, []) - _insert_cell(cell_list, put_val, put_ts) - - return self._send_empty(200) - - # -- Check-and-delete ----------------------------------------------- - - def _handle_check_and_delete(self, table, parts, query): - body = self._json_body() - row_key = _b64dec(parts[1]) - check_col_spec = parts[2] - - row_cells = body.get("Row", [{}])[0].get("Cell", []) - check_value = _b64dec(row_cells[0]["$"]) if row_cells else None - - with data.lock: - tbl = data.tables.get(table, {}) - row_dict = tbl.get(row_key, {}) - current_cells = row_dict.get(check_col_spec, []) - - if not current_cells or current_cells[0][0] != check_value: - return self._send_empty(304) - - # Match - delete the column - row_dict.pop(check_col_spec, None) - if not row_dict: - tbl.pop(row_key, None) - - return self._send_empty(200) - - # -- Utility -------------------------------------------------------- - - def _filter_cells(self, row_data, col_specs, ts_from, ts_to, max_versions): - """Filter a row's cells by column specs and time range.""" - result = {} - for col, cell_list in row_data.items(): - if col_specs and col not in col_specs: - continue - filtered = [] - for val, ts in cell_list: - if ts_from is not None and ts < ts_from: - continue - if ts_to is not None and ts >= ts_to: - continue - filtered.append((val, ts)) - if max_versions is not None: - filtered = filtered[:max_versions] - if filtered: - result[col] = filtered - return result - - def _rows_to_cellset(self, rows): - """Convert list of (row_key_bytes, cells_dict) to CellSet JSON.""" - out_rows = [] - for rk, cells in rows: - out_cells = [] - for col, cell_list in cells.items(): - for val, ts in cell_list: - out_cells.append( - { - "column": _b64enc(col.encode("utf-8")), - "$": _b64enc(val), - "timestamp": ts, - } - ) - out_rows.append( - { - "key": _b64enc(rk), - "Cell": out_cells, - } - ) - return {"Row": out_rows} - - return Handler - - -# --------------------------------------------------------------------------- -# Public API -# --------------------------------------------------------------------------- - - -def start_hbase_mock_server(host="127.0.0.1", port=0): - """Start mock HBase REST server in a daemon thread. - - Returns ``(data, server, port)`` where *port* is the actual bound port. - """ - mock_data = HBaseMockData() - handler_cls = _make_handler_class(mock_data) - server = ThreadingHTTPServer((host, port), handler_cls) - actual_port = server.server_address[1] - thread = threading.Thread(target=server.serve_forever, daemon=True) - thread.start() - return mock_data, server, actual_port diff --git a/requirements.in b/requirements.in index e85fadeef..f36cf492c 100644 --- a/requirements.in +++ b/requirements.in @@ -26,7 +26,7 @@ task-queue>=2.14.0 messagingclient>0.3.0 dracopy>=1.5.0 datastoreflex>=0.5.0 -kvdbclient>=0.7.1 +kvdbclient>=0.8.0 zstandard>=0.23.0 tinybrain>=1.7.0 pykdtree>=1.4.3 diff --git a/requirements.txt b/requirements.txt index d8c8a291a..5a2b3c6ce 100644 --- a/requirements.txt +++ b/requirements.txt @@ -198,7 +198,7 @@ jsonschema==4.26.0 # python-jsonschema-objects jsonschema-specifications==2025.9.1 # via jsonschema -kvdbclient==0.7.1 +kvdbclient==0.8.0 # via # -r requirements.in # cave-pipeline