diff --git a/tests/test_data.py b/tests/test_data.py
index ec43f08..d208b0f 100644
--- a/tests/test_data.py
+++ b/tests/test_data.py
@@ -34,12 +34,3 @@
         "voxel_size": 0.25,
     }
 )
-
-real_init_map_size = sum(
-    len(kitti.get_point_cloud(i).points)
-    for i in range(config.start_index, config.end_index)
-)
-
-real_image_count = config.end_index - (
-    config.start_index - config.start_image_index_offset
-)
diff --git a/tests/test_in_cube.py b/tests/test_in_cube.py
index 615cbc4..593799f 100644
--- a/tests/test_in_cube.py
+++ b/tests/test_in_cube.py
@@ -28,40 +28,48 @@
     "expected_points2instances",
     [
         (
-            np.array([
-                [0.1, 0.1, 5555],
-                [0.2, 0.2, 6666],
-                [0.3, 0.3, 7777],
-                [0.4, 0.4, 2],
-                [0.5, 0.5, 8888],
-                [0.6, 0.6, 9999],
-                [0.7, 0.7, 5],
-                [0.8, 0.8, 3],
-                [0.9, 0.9, 1],
-            ]),
-            np.array([
-                [1, 2, 3, 4, 5, 6],
-                [0, 2, 3, 4, 5, 6],
-                [1, 0, 3, 4, 5, 6],
-                [1, 2, 0, 4, 5, 6],
-                [1, 2, 3, 0, 5, 6],
-                [1, 2, 3, 4, 0, 6],
-                [1, 2, 3, 4, 5, 0],
-                [0, 0, 3, 4, 5, 6],
-                [1, 0, 0, 4, 5, 6],
-            ]),
-            np.array([
-                [0.4, 0.4, 2],
-                [0.7, 0.7, 5],
-                [0.8, 0.8, 3],
-                [0.9, 0.9, 1],
-            ]),
-            np.array([
-                [1, 2, 0, 4, 5, 6],
-                [1, 2, 3, 4, 5, 0],
-                [0, 0, 3, 4, 5, 6],
-                [1, 0, 0, 4, 5, 6],
-            ]),
+            np.array(
+                [
+                    [0.1, 0.1, 5555],
+                    [0.2, 0.2, 6666],
+                    [0.3, 0.3, 7777],
+                    [0.4, 0.4, 2],
+                    [0.5, 0.5, 8888],
+                    [0.6, 0.6, 9999],
+                    [0.7, 0.7, 5],
+                    [0.8, 0.8, 3],
+                    [0.9, 0.9, 1],
+                ]
+            ),
+            np.array(
+                [
+                    [1, 2, 3, 4, 5, 6],
+                    [0, 2, 3, 4, 5, 6],
+                    [1, 0, 3, 4, 5, 6],
+                    [1, 2, 0, 4, 5, 6],
+                    [1, 2, 3, 0, 5, 6],
+                    [1, 2, 3, 4, 0, 6],
+                    [1, 2, 3, 4, 5, 0],
+                    [0, 0, 3, 4, 5, 6],
+                    [1, 0, 0, 4, 5, 6],
+                ]
+            ),
+            np.array(
+                [
+                    [0.4, 0.4, 2],
+                    [0.7, 0.7, 5],
+                    [0.8, 0.8, 3],
+                    [0.9, 0.9, 1],
+                ]
+            ),
+            np.array(
+                [
+                    [1, 2, 0, 4, 5, 6],
+                    [1, 2, 3, 4, 5, 0],
+                    [0, 0, 3, 4, 5, 6],
+                    [1, 0, 0, 4, 5, 6],
+                ]
+            ),
         )
     ],
 )
diff --git a/tests/test_init_instances_matrix.py b/tests/test_init_instances_matrix.py
index dc51bb1..30389f5 100644
--- a/tests/test_init_instances_matrix.py
+++ b/tests/test_init_instances_matrix.py
@@ -18,7 +18,6 @@
 from services.preprocessing.init.instances_matrix import InitInstancesMatrixProcessor
 
 from tests.test_data import config
-from tests.test_data import real_image_count
 from tests.utils import generate_init_pcd
 
 
@@ -26,4 +25,8 @@
 def test_init_map(init_pcd: o3d.geometry.PointCloud):
     points2instances = InitInstancesMatrixProcessor().process(config, init_pcd)
 
+    real_image_count = config.end_index - (
+        config.start_index - config.start_image_index_offset
+    )
+
     assert points2instances.shape == (len(init_pcd.points), real_image_count)
diff --git a/tests/test_init_map.py b/tests/test_init_map.py
index 0867fff..77a358e 100644
--- a/tests/test_init_map.py
+++ b/tests/test_init_map.py
@@ -15,10 +15,14 @@
 from services.preprocessing.init.map import InitMapProcessor
 
 from tests.test_data import config
-from tests.test_data import real_init_map_size
 
 
 def test_init_map():
     init_pcd = InitMapProcessor().process(config)
 
+    real_init_map_size = sum(
+        len(config.dataset.get_point_cloud(i).points)
+        for i in range(config.start_index, config.end_index)
+    )
+
     assert len(init_pcd.points) == real_init_map_size