add fiber-local test_finished_callback to cleanup contribs code and e…

…ncapsulate all the retries logic inside TestRetries::Component

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

@@ -32,12 +32,11 @@ def run_one_method(klass, method_name)
               return super unless datadog_configuration[:enabled]
 
               result = nil
-              # retries here
-              test_retries_component.with_retries do |test_finished_callback|
-                Thread.current[:__dd_retry_callback] = test_finished_callback
 
+              test_retries_component.with_retries do
                 result = super
               end
+
               result
             end
 

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

@@ -50,7 +50,7 @@ def after_teardown
               test_span = test_visibility_component.active_test
               return super unless test_span
 
-              finish_with_result(test_span, result_code, Thread.current[:__dd_retry_callback])
+              finish_with_result(test_span, result_code)
               if Helpers.parallel?(self.class)
                 finish_with_result(test_span.test_suite, result_code)
               end
@@ -60,7 +60,7 @@ def after_teardown
 
             private
 
-            def finish_with_result(span, result_code, callback = nil)
+            def finish_with_result(span, result_code)
               return unless span
 
               case result_code
@@ -72,8 +72,6 @@ def finish_with_result(span, result_code, callback = nil)
                 span.skipped!(reason: failure.message)
               end
 
-              callback.call(span) if callback
-
               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/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

spec/datadog/ci/contrib/rspec/instrumentation_spec.rb

@@ -837,7 +837,7 @@ def rspec_skipped_session_run
       expect(test_spans).to have(6).items
 
       failed_spans, passed_spans = test_spans.partition { |span| span.get_tag("test.status") == "fail" }
-      expect(failed_spans).to have(4).items # see steps.rb
+      expect(failed_spans).to have(4).items
       expect(passed_spans).to have(2).items
 
       test_spans_by_test_name = test_spans.group_by { |span| span.get_tag("test.name") }