Handle flaky tests

test/CMakeLists.txt

@@ -5,11 +5,12 @@ if (NOT WIN32)
     target_link_libraries(run_test PUBLIC OpenMP::OpenMP_CXX)
 endif()
 target_link_libraries(run_test PUBLIC perf_lib layers_lib)
-target_link_libraries(run_test PUBLIC gtest_main)
+target_link_libraries(run_test PUBLIC gtest)
 target_link_libraries(run_test PUBLIC ReadLib)
 target_link_libraries(run_test PUBLIC reader_lib)
 
 
+target_include_directories(run_test PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}")
 target_include_directories(run_test PRIVATE "${CMAKE_SOURCE_DIR}/app/ReaderImage")
 
 if (NOT CMAKE_BUILD_TYPE)

test/inference/test_inference.cpp

@@ -12,6 +12,7 @@
 #include "layers/OutputLayer.hpp"
 #include "layers/PoolingLayer.hpp"
 #include "layers/SplitLayer.hpp"
+#include "utils/flaky_test_runner.hpp"
 
 using namespace it_lab_ai;
 
@@ -487,7 +488,7 @@ TEST(bfs, check_struct_layer_added) {
   ASSERT_EQ(tmp, res);
 }
 
-TEST(bfs, check_struct_graph_split) {
+FLAKY_TEST(bfs, check_struct_graph_split) {
   std::vector<std::vector<std::pair<int, int>>> split = {
       {{12, 0}, {13, 0}, {14, 0}}};
   Graph graph(151, split);
@@ -565,4 +566,5 @@ TEST(bfs, check_struct_graph_split) {
   std::vector<int> tmp = *output.as<int>();
   std::vector<int> res(36, 81);
   ASSERT_EQ(tmp, res);
-}
+}
+FLAKY_END_TEST

test/main.cpp

@@ -0,0 +1,25 @@
+#include <gtest/gtest.h>
+
+#include <cstdlib>
+#include <cstring>
+#include <iostream>
+
+int main(int argc, char** argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+
+  const char* flaky_disabled = std::getenv("DISABLE_FLAKY_TESTS");
+  bool flaky_enabled =
+      (flaky_disabled == nullptr) || (std::strcmp(flaky_disabled, "1") != 0 &&
+                                      std::strcmp(flaky_disabled, "true") != 0);
+
+  if (flaky_enabled) {
+    const char* retries = std::getenv("FLAKY_RETRIES");
+    int retry_count = retries ? std::atoi(retries) : 10;
+    std::cout << "Flaky test support enabled. Max retries: " << retry_count
+              << std::endl;
+    std::cout << "Use FLAKY_TEST/FLAKY_TEST_F macros to create flaky tests."
+              << std::endl;
+  }
+
+  return RUN_ALL_TESTS();
+}

test/utils/flaky_test_runner.hpp

@@ -0,0 +1,84 @@
+#pragma once
+
+#include <gtest/gtest.h>
+
+#include <cstdlib>
+#include <cstring>
+#include <iostream>
+
+namespace test_utils {
+
+static int getFlakyRetries() {
+  const char* retry_env = std::getenv("FLAKY_RETRIES");
+  if (retry_env != nullptr) {
+    try {
+      int retries = std::stoi(retry_env);
+      return std::max(1, std::min(retries, 100));
+    } catch (const std::exception&) {
+      return 10;
+    }
+  }
+  return 10;
+}
+
+static bool isFlakyTestsEnabled() {
+  const char* disabled = std::getenv("DISABLE_FLAKY_TESTS");
+  if (disabled == nullptr) {
+    return true;
+  }
+  return std::strcmp(disabled, "1") != 0 && std::strcmp(disabled, "true") != 0;
+}
+
+template <typename TestFunc>
+void runFlakyTest(const char* test_name, TestFunc test_func) {
+  if (!isFlakyTestsEnabled()) {
+    test_func();
+    return;
+  }
+
+  int max_retries = getFlakyRetries();
+
+  for (int attempt = 1; attempt <= max_retries; ++attempt) {
+    try {
+      if (attempt > 1) {
+        std::cout << "[FLAKY RETRY " << attempt << "/" << max_retries << "] "
+                  << test_name << std::endl;
+      }
+
+      test_func();
+
+      if (attempt > 1) {
+        std::cout << "[FLAKY SUCCESS] " << test_name << " passed on attempt "
+                  << attempt << std::endl;
+      }
+      return;
+
+    } catch (...) {
+      if (attempt == max_retries) {
+        std::cout << "[FLAKY EXHAUSTED] " << test_name << " failed after "
+                  << max_retries << " attempts" << std::endl;
+        throw;
+      } else if (attempt == 1) {
+        std::cout << "[FLAKY FAILED] " << test_name << " failed on attempt "
+                  << attempt << ", retrying..." << std::endl;
+      }
+    }
+  }
+}
+
+#define FLAKY_TEST(test_case_name, test_name) \
+  TEST(test_case_name, test_name) {           \
+    auto flaky_test_body = []()
+
+#define FLAKY_TEST_F(test_fixture, test_name) \
+  TEST_F(test_fixture, test_name) {           \
+    auto flaky_test_body = [this]()
+
+#define FLAKY_END_TEST                                               \
+  ;                                                                  \
+  test_utils::runFlakyTest(                                          \
+      testing::UnitTest::GetInstance()->current_test_info()->name(), \
+      flaky_test_body);                                              \
+  }
+
+}