From 738b88250d48f3c69455961bf7b35961e7a850aa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexandre=20P=C3=A9r=C3=A9?= <alexandre.pere@zama.ai>
Date: Mon, 11 Dec 2023 14:32:17 +0100
Subject: [PATCH] chore(compiler): fix broken retry for flaky tests

---
 .../workflows/compiler_build_and_test_cpu.yml |  7 +++++
 compilers/concrete-compiler/compiler/Makefile |  4 +--
 .../tests/end_to_end_tests/end_to_end_test.cc | 28 +++++++++++++++----
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/compiler_build_and_test_cpu.yml b/.github/workflows/compiler_build_and_test_cpu.yml
index 920724e1e8..e0962ae0a2 100644
--- a/.github/workflows/compiler_build_and_test_cpu.yml
+++ b/.github/workflows/compiler_build_and_test_cpu.yml
@@ -36,6 +36,7 @@ jobs:
     runs-on: ${{ github.event.inputs.runner_name }}
     if: ${{ !cancelled() }}
     steps:
+  
       - name: Instance configuration used
         run: |
           echo "IDs: ${{ inputs.instance_id }}"
@@ -142,6 +143,12 @@ jobs:
             mkdir -p /tmp/concrete_compiler/gpu_tests/
             make MINIMAL_TESTS=${{ env.MINIMAL_TESTS }} DATAFLOW_EXECUTION_ENABLED=ON CCACHE=ON Python3_EXECUTABLE=$PYTHON_EXEC BUILD_DIR=/build run-tests run-end-to-end-dataflow-tests
             chmod -R ugo+rwx /tmp/KeySetCache
+            
+      - name: Analyze logs
+        run: |
+            cd build/gtest-parallel-logs/passed
+            echo "/!\ WARNING RETRY TEST: Different results ..." > testing.log
+            ls -1 | xargs grep -H "WARNING RETRY" | sed -e "s/.log.*//g" | uniq -c | sed -re "s/ *([0-9]*) (.*)/::warning ::Test \2 retried \1 times/g" | cat
 
       # - name: Archive python package
       #   uses: actions/upload-artifact@v3
diff --git a/compilers/concrete-compiler/compiler/Makefile b/compilers/concrete-compiler/compiler/Makefile
index 305187203f..dc0ed3173f 100644
--- a/compilers/concrete-compiler/compiler/Makefile
+++ b/compilers/concrete-compiler/compiler/Makefile
@@ -215,7 +215,7 @@ GTEST_PARALLEL_CMD=
 GTEST_PARALLEL_SEPARATOR=
 else
 GTEST_PARALLEL_PY=./gtest-parallel/gtest-parallel
-GTEST_PARALLEL_CMD=$(Python3_EXECUTABLE) $(GTEST_PARALLEL_PY)
+GTEST_PARALLEL_CMD=$(Python3_EXECUTABLE) $(GTEST_PARALLEL_PY) -d $(BUILD_DIR)
 GTEST_PARALLEL_SEPARATOR=--
 $(GTEST_PARALLEL_PY):
 	git clone https://github.com/google/gtest-parallel.git
@@ -299,7 +299,7 @@ run-end-to-end-tests: $(GTEST_PARALLEL_PY) build-end-to-end-tests generate-cpu-t
 	$(foreach optimizer_strategy,$(OPTIMIZATION_STRATEGY_TO_TEST), $(foreach security,$(SECURITY_TO_TEST), \
 	  $(GTEST_PARALLEL_CMD) $(BUILD_DIR)/tools/concretelang/tests/end_to_end_tests/end_to_end_test \
 	  $(GTEST_PARALLEL_SEPARATOR) --backend=cpu --security-level=$(security) \
-	  --optimizer-strategy=$(optimizer_strategy) --retry-failing-tests=1 $(FIXTURE_CPU_DIR)/*.yaml || exit $$?;))
+	  --optimizer-strategy=$(optimizer_strategy) --retry-failing-tests=5 $(FIXTURE_CPU_DIR)/end_to_end_fhe.yaml || exit $$?;))
 
 ### end-to-end-tests GPU
 
diff --git a/compilers/concrete-compiler/compiler/tests/end_to_end_tests/end_to_end_test.cc b/compilers/concrete-compiler/compiler/tests/end_to_end_tests/end_to_end_test.cc
index 0674ddff59..49d3929493 100644
--- a/compilers/concrete-compiler/compiler/tests/end_to_end_tests/end_to_end_test.cc
+++ b/compilers/concrete-compiler/compiler/tests/end_to_end_tests/end_to_end_test.cc
@@ -1,6 +1,7 @@
 #include <cstdint>
 #include <filesystem>
 #include <gtest/gtest.h>
+#include <llvm/Support/raw_ostream.h>
 #include <regex>
 #include <type_traits>
 
@@ -55,25 +56,40 @@ class EndToEndTest : public ::testing::Test {
   }
 
   void testOnce() {
+
+    // We try for a a number of times
     for (auto tests_rep = 0; tests_rep <= retryFailingTests; tests_rep++) {
       // We execute the circuit.
       auto maybeRes = testCircuit->call(args);
       ASSERT_OUTCOME_HAS_VALUE(maybeRes);
       auto result = maybeRes.value();
 
-      /* Check result */
+      /* Check results */
+      bool allgood = true;
       for (size_t i = 0; i < desc.outputs.size(); i++) {
         auto maybeErr = checkResult(desc.outputs[i], result[i]);
-        if (maybeErr && tests_rep < 2) {
+        if (maybeErr) {
+          allgood = false;
+          llvm::errs() << "/!\\ WARNING RETRY TEST: " << maybeErr << "\n";
           llvm::consumeError(std::move(maybeErr));
-          ASSERT_OUTCOME_HAS_VALUE(
-              testCircuit->generateKeyset(tests_rep + 1, tests_rep + 1));
           break;
-        } else {
-          ASSERT_LLVM_ERROR(std::move(maybeErr));
         }
       }
+
+      // If OK we return
+      if (allgood) {
+        return;
+      }
+
+      // Otherwise we reset the keyset
+      llvm::errs() << "Regenerating keyset with seed: ";
+      llvm::errs() << tests_rep + 1;
+      llvm::errs() << "\n";
+      ASSERT_OUTCOME_HAS_VALUE(testCircuit->generateKeyset(tests_rep + 1, tests_rep + 1));
     }
+
+    // If all attempts fail, we return an error.
+    ASSERT_TRUE(false) << "Test failed after multiple attempts";
   }
 
   void testErrorRate() {