diff --git a/model_server/extensions/chaeo/tests/test_zstack.py b/model_server/extensions/chaeo/tests/test_zstack.py
index 7dd2c05a3e5da0e2cf3a62c092b2b29dad07da72..cc700b49aba513f3d793fda87406fcdd211cb575 100644
--- a/model_server/extensions/chaeo/tests/test_zstack.py
+++ b/model_server/extensions/chaeo/tests/test_zstack.py
@@ -8,7 +8,7 @@ from model_server.extensions.chaeo.conf.testing import multichannel_zstack, pixe
 from extensions.chaeo.params import RoiSetExportParams, RoiSetMetaParams
 from model_server.extensions.chaeo.products import export_patches_from_zstack, export_multichannel_patches_from_zstack, export_patch_masks_from_zstack
 from model_server.extensions.chaeo.workflows import infer_object_map_from_zstack
-from model_server.extensions.chaeo.zmask import build_zmask_from_object_mask, get_label_ids, RoiSet
+from model_server.extensions.chaeo.zmask import get_label_ids, RoiSet
 from model_server.base.accessors import generate_file_accessor, InMemoryDataAccessor, write_accessor_data_to_file
 from model_server.extensions.ilastik.models import IlastikPixelClassifierModel
 from model_server.base.models import DummyInstanceSegmentationModel
@@ -42,11 +42,16 @@ class TestZStackDerivedDataProducts(unittest.TestCase):
 
     def test_zmask_makes_correct_boxes(self, mask_type='boxes', **kwargs):
         id_map = get_label_ids(self.seg_mask)
-        zmask, meta, _, interm = build_zmask_from_object_mask(
+        roiset = RoiSet(
             id_map,
             self.stack_ch_pa,
-            params=RoiSetMetaParams(mask_type=mask_type, filters=kwargs.get('filters')),
+            params=RoiSetMetaParams(
+                mask_type=mask_type, filters=kwargs.get('filters')
+            )
         )
+        zmask = roiset.get_zmask()
+        meta = roiset.zmask_meta
+        interm = roiset.interm
         zmask_acc = InMemoryDataAccessor(zmask)
         self.assertTrue(zmask_acc.is_mask())
 
@@ -118,11 +123,9 @@ class TestZStackDerivedDataProducts(unittest.TestCase):
 
     def test_flatten_image(self):
         id_map = get_label_ids(self.seg_mask)
-        _, _, df, _ = build_zmask_from_object_mask(
-            id_map,
-            self.stack_ch_pa,
-            params=RoiSetMetaParams(mask_type='boxes')
-        )
+
+        roiset = RoiSet(id_map, self.stack_ch_pa, params=RoiSetMetaParams(mask_type='boxes'))
+        df = roiset.df
 
         from model_server.extensions.chaeo.zmask import project_stack_from_focal_points