diff --git a/tests/advanced_streaming/test_advanced_streaming_noslip.py b/tests/advanced_streaming/test_advanced_streaming_noslip.py
index 547d14ac4669197233f5ef22e17483917388f275..d6acc6d6047af5fed5b71795d6238104b2e349d2 100644
--- a/tests/advanced_streaming/test_advanced_streaming_noslip.py
+++ b/tests/advanced_streaming/test_advanced_streaming_noslip.py
@@ -43,7 +43,7 @@ def test_advanced_streaming_noslip_single_cell(stencil, streaming_pattern, prev_
     index_struct_dtype = numpy_data_type_for_boundary_object(noslip, stencil.D)
 
     index_field = Field('indexVector', FieldType.INDEXED, index_struct_dtype, layout=[0],
-                        shape=(TypedSymbol("indexVectorSize", np.int64), 1), strides=(1, 1))
+                        shape=(TypedSymbol("indexVectorSize", np.int32), 1), strides=(1, 1))
     index_vector = np.array([pos + (d,) for d in range(stencil.Q)], dtype=index_struct_dtype)
 
     ast = create_lattice_boltzmann_boundary_kernel(pdf_field,
diff --git a/tests/benchmark/test_benchmark.py b/tests/benchmark/test_benchmark.py
index a59d0510d190cac2a34067c4085ce34e04695b0a..1680a547d1921958afa36cb2c5eff02cfc696447 100644
--- a/tests/benchmark/test_benchmark.py
+++ b/tests/benchmark/test_benchmark.py
@@ -4,7 +4,13 @@ import numpy as np
 import sympy as sp
 
 from lbmpy.scenarios import create_lid_driven_cavity
-from pystencils.cpu.cpujit import add_or_change_compiler_flags
+from lbmpy._compat import IS_PYSTENCILS_2
+
+if IS_PYSTENCILS_2:
+    from pystencils.backend.jit.legacy_cpu import add_or_change_compiler_flags
+else:
+    from pystencils.cpu.cpujit import add_or_change_compiler_flags
+
 from pystencils.runhelper import ParameterStudy