From 04bd24b320deb15469010f088c3ccbe9701ee456 Mon Sep 17 00:00:00 2001
From: JoeStrout <joe@strout.net>
Date: Mon, 5 Aug 2024 15:36:56 +0000
Subject: [PATCH] Add test_custom_op.py to expose RecursionError bug.

---
 .../subchunkable/test_custom_op.py            | 60 +++++++++++++++++++
 1 file changed, 60 insertions(+)
 create mode 100644 tests/integration/subchunkable/test_custom_op.py

diff --git a/tests/integration/subchunkable/test_custom_op.py b/tests/integration/subchunkable/test_custom_op.py
new file mode 100644
index 000000000..f2bfced63
--- /dev/null
+++ b/tests/integration/subchunkable/test_custom_op.py
@@ -0,0 +1,60 @@
+# This file triggers a RecursionError inside id_generation.py.
+
+import attrs
+import zetta_utils
+from zetta_utils import mazepa
+from zetta_utils import builder
+import zetta_utils.mazepa_layer_processing.common
+from zetta_utils.layer.volumetric import VolumetricIndex, VolumetricLayer
+from zetta_utils.geometry import Vec3D
+import zetta_utils.tensor_ops.convert as convert
+
+@builder.register("CustomVolOp")
+@mazepa.taskable_operation_cls
+@attrs.frozen()
+class CustomVolOp:  # implements VolumetricOpProtocol
+    threshold: float = 0.1
+    connectivity: int = 6
+    ids_per_chunk: int = 10000
+
+    def get_input_resolution(self, dst_resolution: Vec3D) -> Vec3D:
+        # For simplicity, just return the destination resolution
+        return dst_resolution
+
+    def with_added_crop_pad(self, crop_pad: Vec3D = None) -> 'CustomVolOp':
+        return self # attrs.evolve(self, crop_pad=crop_pad)
+
+    def __call__(self,
+                 idx: VolumetricIndex,
+                 dst: VolumetricLayer,
+                 src: VolumetricLayer,
+                 *args, **kwargs):
+        # If the following line is commented out, the problem does not occur:
+        data_np = convert.to_np(src[idx][0])  # (use channel 0)
+
+        return None
+
+d = {'@type': 'build_subchunkable_apply_flow',
+ 'bbox': {'@type': 'BBox3D.from_coords',
+    'start_coord': [13996, 10633, 3062],
+    'end_coord': [14124, 10761, 3102],
+    'resolution': [16, 16, 42]
+ },
+ 'dst_resolution': [16, 16, 42],
+ 'processing_chunk_sizes': [[128, 128, 40]],
+ 'skip_intermediaries': True,
+ 'expand_bbox_processing': True,
+ 'op': {'@type': 'CustomVolOp'},
+ 'op_kwargs': {
+     'src': {'@type': 'build_cv_layer','path': 'gs://dkronauer-ant-001-synapse/test/syndet20240802'}
+ },
+ 'dst': {'@type': 'build_cv_layer',
+    'path': 'gs://dkronauer-ant-001-synapse/test/synseg20240802',
+    'info_reference_path': 'gs://dkronauer-ant-001-synapse/test/syndet20240802',
+    'on_info_exists': 'overwrite',
+    'info_field_overrides': {'num_channels': 1, 'data_type': 'int32'},
+    'info_add_scales': [[16, 16, 42]],
+    'info_add_scales_mode': 'replace'
+ }
+}
+flow = zetta_utils.builder.build(d)