fix(frontend-python): ensure shift strategies coverage

zama-ai · Jan 5, 2024 · 2cf1508 · 2cf1508
1 parent 69efcee
commit 2cf1508
Showing 1 changed file with 13 additions and 10 deletions.
diff --git a/frontends/concrete-python/tests/execution/test_shift.py b/frontends/concrete-python/tests/execution/test_shift.py
@@ -33,7 +33,18 @@
             ]
         ]
 
-for _ in range(8):
+NB_RANDOM_TESTS = 8
+
+strategies = [
+    fhe.BitwiseStrategy.ONE_TLU_PROMOTED,
+    fhe.BitwiseStrategy.THREE_TLU_CASTED,
+    fhe.BitwiseStrategy.TWO_TLU_BIGGER_PROMOTED_SMALLER_CASTED,
+    fhe.BitwiseStrategy.TWO_TLU_BIGGER_CASTED_SMALLER_PROMOTED,
+    fhe.BitwiseStrategy.CHUNKED,
+]
+random_tests_stategies = strategies * (NB_RANDOM_TESTS // len(strategies) + 1)
+
+for _ in range(NB_RANDOM_TESTS):
     cases.append(
         [
             # operation
@@ -49,15 +60,7 @@
             random.choice([(), (2,), (3, 2)]),
             random.choice([(), (2,), (3, 2)]),
             # strategy
-            random.choice(
-                [
-                    fhe.BitwiseStrategy.ONE_TLU_PROMOTED,
-                    fhe.BitwiseStrategy.THREE_TLU_CASTED,
-                    fhe.BitwiseStrategy.TWO_TLU_BIGGER_PROMOTED_SMALLER_CASTED,
-                    fhe.BitwiseStrategy.TWO_TLU_BIGGER_CASTED_SMALLER_PROMOTED,
-                    fhe.BitwiseStrategy.CHUNKED,
-                ]
-            ),
+            random_tests_stategies.pop(),
         ]
     )