Improve bit mask test

pystencils_tests/test_bit_masks.py

+import pytest
 import numpy as np
+
 import pystencils as ps
 from pystencils import Field, Assignment, create_kernel
 from pystencils.bit_masks import flag_cond
 
 
-def test_flag_condition():
+@pytest.mark.parametrize('mask_type', [np.uint8, np.uint16, np.uint32, np.uint64])
+def test_flag_condition(mask_type):
     f_arr = np.zeros((2, 2, 2), dtype=np.float64)
-    # TODO different uints
-    mask_arr = np.zeros((2, 2), dtype=np.uint64)
+    mask_arr = np.zeros((2, 2), dtype=mask_type)
 
     mask_arr[0, 1] = (1 << 3)
     mask_arr[1, 0] = (1 << 5)
@@ -18,7 +20,7 @@ def test_flag_condition():
 
     v1 = 42.3
     v2 = 39.7
-    v3 = 119.87
+    v3 = 119
 
     assignments = [
         Assignment(f(0), flag_cond(3, mask(0), v1)),
@@ -28,6 +30,7 @@ def test_flag_condition():
     kernel = create_kernel(assignments).compile()
     kernel(f=f_arr, mask=mask_arr)
     code = ps.get_code_str(kernel)
+    assert '119.0' in code
 
     reference = np.zeros((2, 2, 2), dtype=np.float64)
     reference[0, 1, 0] = v1