[SDTEST-437] Auto test retries for minitest

lib/datadog/ci/contrib/minitest/runner.rb

@@ -28,6 +28,18 @@ def init_plugins(*args)
               test_visibility_component.start_test_module(Ext::FRAMEWORK)
             end
 
+            def run_one_method(klass, method_name)
+              return super unless datadog_configuration[:enabled]
+
+              result = nil
+
+              test_retries_component.with_retries do
+                result = super
+              end
+
+              result
+            end
+
             private
 
             def datadog_configuration
@@ -37,6 +49,10 @@ def datadog_configuration
             def test_visibility_component
               Datadog.send(:components).test_visibility
             end
+
+            def test_retries_component
+              Datadog.send(:components).test_retries
+            end
           end
         end
       end

lib/datadog/ci/contrib/minitest/test.rb

@@ -71,6 +71,7 @@ def finish_with_result(span, result_code)
               when "S"
                 span.skipped!(reason: failure.message)
               end
+
               span.finish
             end
 

lib/datadog/ci/contrib/rspec/example.rb

@@ -40,7 +40,7 @@ def run(*args)
               # don't report test to RSpec::Core::Reporter until retries are done
               @skip_reporting = true
 
-              test_retries_component.with_retries do |retry_callback|
+              test_retries_component.with_retries do
                 test_visibility_component.trace_test(
                   test_name,
                   suite_name,
@@ -76,8 +76,6 @@ def run(*args)
                       exception: execution_result.pending_exception
                     )
                   end
-
-                  retry_callback.call(test_span)
                 end
               end
 

lib/datadog/ci/test_retries/component.rb

@@ -45,11 +45,15 @@ def with_retries(&block)
             end
           end
 
+          test_visibility_component.set_test_finished_callback(test_finished_callback)
+
           loop do
-            yield test_finished_callback
+            yield
 
             break unless retry_strategy&.should_retry?
           end
+        ensure
+          test_visibility_component.remove_test_finished_callback
         end
 
         def build_strategy(test_span)
@@ -70,6 +74,10 @@ def build_strategy(test_span)
         def should_retry_failed_test?(test_span)
           @retry_failed_tests_enabled && !!test_span&.failed? && @retry_failed_tests_count < @retry_failed_tests_total_limit
         end
+
+        def test_visibility_component
+          Datadog.send(:components).test_visibility
+        end
       end
     end
   end

lib/datadog/ci/test_visibility/component.rb

@@ -19,6 +19,8 @@ module TestVisibility
       class Component
         attr_reader :test_suite_level_visibility_enabled
 
+        FIBER_LOCAL_TEST_FINISHED_CALLBACK_KEY = :__dd_test_finished_callback
+
         def initialize(
           test_suite_level_visibility_enabled: false,
           codeowners: Codeowners::Parser.new(Git::LocalRepository.root).parse
@@ -125,6 +127,15 @@ def deactivate_test_suite(test_suite_name)
           @context.deactivate_test_suite(test_suite_name)
         end
 
+        # sets fiber-local callback to be called when test is finished
+        def set_test_finished_callback(callback)
+          Thread.current[FIBER_LOCAL_TEST_FINISHED_CALLBACK_KEY] = callback
+        end
+
+        def remove_test_finished_callback
+          Thread.current[FIBER_LOCAL_TEST_FINISHED_CALLBACK_KEY] = nil
+        end
+
         def itr_enabled?
           test_optimisation.enabled?
         end
@@ -192,6 +203,8 @@ def on_test_finished(test)
           test_optimisation.count_skipped_test(test)
 
           Telemetry.event_finished(test)
+
+          Thread.current[FIBER_LOCAL_TEST_FINISHED_CALLBACK_KEY]&.call(test)
         end
 
         # HELPERS

lib/datadog/ci/test_visibility/null_component.rb

@@ -47,6 +47,12 @@ def itr_enabled?
           false
         end
 
+        def set_test_finished_callback(_)
+        end
+
+        def remove_test_finished_callback
+        end
+
         private
 
         def skip_tracing(block = nil)

sig/datadog/ci/contrib/minitest/runner.rbs

@@ -10,11 +10,15 @@ module Datadog
 
             def init_plugins: (*untyped) -> (nil | untyped)
 
+            def run_one_method: (untyped klass, String method_name) -> untyped
+
             private
 
             def datadog_configuration: () -> untyped
 
             def test_visibility_component: () -> Datadog::CI::TestVisibility::Component
+
+            def test_retries_component: () -> Datadog::CI::TestRetries::Component
           end
         end
       end

sig/datadog/ci/contrib/minitest/test.rbs

@@ -13,6 +13,7 @@ module Datadog
 
           module InstanceMethods : ::Minitest::Test
             include ::Minitest::Test::LifecycleHooks
+
             extend ClassMethods
 
             def before_setup: () -> (nil | untyped)

sig/datadog/ci/test_retries/component.rbs

@@ -16,13 +16,15 @@ module Datadog
 
         def configure: (Datadog::CI::Remote::LibrarySettings library_settings) -> void
 
-        def with_retries: () { (untyped) -> void } -> void
+        def with_retries: () { () -> void } -> void
 
         def build_strategy: (Datadog::CI::Test test) -> Datadog::CI::TestRetries::Strategy::Base
 
         private
 
         def should_retry_failed_test?: (Datadog::CI::Test test) -> bool
+
+        def test_visibility_component: () -> Datadog::CI::TestVisibility::Component
       end
     end
   end

sig/datadog/ci/test_visibility/component.rbs

@@ -7,6 +7,8 @@ module Datadog
         @codeowners: Datadog::CI::Codeowners::Matcher
         @context: Datadog::CI::TestVisibility::Context
 
+        FIBER_LOCAL_TEST_FINISHED_CALLBACK_KEY: Symbol
+
         attr_reader test_suite_level_visibility_enabled: bool
 
         def initialize: (?test_suite_level_visibility_enabled: bool, ?codeowners: Datadog::CI::Codeowners::Matcher) -> void
@@ -39,6 +41,10 @@ module Datadog
 
         def deactivate_test_suite: (String test_suite_name) -> void
 
+        def set_test_finished_callback: (Proc callback) -> void
+
+        def remove_test_finished_callback: () -> void
+
         def itr_enabled?: () -> bool
 
         def shutdown!: () -> void

sig/datadog/ci/test_visibility/null_component.rbs

@@ -24,6 +24,10 @@ module Datadog
 
         def active_span: () -> nil
 
+        def set_test_finished_callback: (Proc callback) -> void
+
+        def remove_test_finished_callback: () -> void
+
         def shutdown!: () -> nil
 
         def itr_enabled?: () -> bool