diff --git a/tests/unit_tests/test_mesh_to_vtk.py b/tests/unit_tests/test_mesh_to_vtk.py
index e487b3668b..b57235edda 100644
--- a/tests/unit_tests/test_mesh_to_vtk.py
+++ b/tests/unit_tests/test_mesh_to_vtk.py
@@ -32,7 +32,7 @@ def test_write_data_to_vtk(mesh, tmpdir):
     # BUILD
     filename = tmpdir / "out.vtk"
 
-    data = np.random.random(mesh.dimension[0]*mesh.dimension[1]*mesh.dimension[2])
+    data = np.random.random(mesh.num_mesh_cells)
 
     # RUN
     mesh.write_data_to_vtk(filename=filename, datasets={"label1": data, "label2": data})