add new on_after_test_span_finished event

lib/datadog/ci/test_retries/component.rb

@@ -71,9 +71,7 @@ def with_retries(&block)
           # @type var retry_strategy: Strategy::Base
           retry_strategy = nil
 
-          test_finished_callback = lambda do |tracer_span|
-            test_span = Datadog::CI::Test.new(tracer_span)
-
+          test_finished_callback = lambda do |test_span|
             if retry_strategy.nil?
               # we always run test at least once and after first pass create a correct retry strategy
               retry_strategy = build_strategy(test_span)

lib/datadog/ci/test_visibility/component.rb

@@ -57,13 +57,7 @@ def start_test_suite(test_suite_name, service: nil, tags: {})
         def trace_test(test_name, test_suite_name, service: nil, tags: {}, &block)
           if block
             @context.trace_test(test_name, test_suite_name, service: service, tags: tags) do |test|
-              # subscribe on test stop event
-              if test_finished_callback
-                events = test.tracer_span.send(:events)
-                events.after_stop.subscribe do |event|
-                  test_finished_callback&.call(event)
-                end
-              end
+              subscribe_to_after_stop_event(test.tracer_span)
 
               on_test_started(test)
               res = block.call(test)
@@ -72,6 +66,7 @@ def trace_test(test_name, test_suite_name, service: nil, tags: {}, &block)
             end
           else
             test = @context.trace_test(test_name, test_suite_name, service: service, tags: tags)
+            subscribe_to_after_stop_event(test.tracer_span)
             on_test_started(test)
             test
           end
@@ -216,13 +211,27 @@ def on_test_finished(test)
           test_optimisation.count_skipped_test(test)
 
           Telemetry.event_finished(test)
+
+          test_finished_callback&.call(test)
+        end
+
+        def on_after_test_span_finished(tracer_span)
+          # noop
         end
 
         # HELPERS
         def skip_tracing(block = nil)
           block&.call(nil)
         end
 
+        def subscribe_to_after_stop_event(tracer_span)
+          events = tracer_span.send(:events)
+
+          events.after_stop.subscribe do |span|
+            on_after_test_span_finished(span)
+          end
+        end
+
         def set_codeowners(test)
           source = test.source_file
           owners = @codeowners.list_owners(source) if source

sig/datadog/ci/test_visibility/component.rbs

@@ -41,6 +41,8 @@ module Datadog
 
         def deactivate_test_suite: (String test_suite_name) -> void
 
+        def test_finished_callback: () -> Proc?
+
         def set_test_finished_callback: (Proc callback) -> void
 
         def remove_test_finished_callback: () -> void
@@ -79,6 +81,10 @@ module Datadog
 
         def on_test_finished: (Datadog::CI::Test test) -> void
 
+        def on_after_test_span_finished: (Datadog::Tracing::SpanOperation span) -> void
+
+        def subscribe_to_after_stop_event: (Datadog::Tracing::SpanOperation span) -> void
+
         def test_optimisation: () -> Datadog::CI::TestOptimisation::Component
 
         def git_tree_upload_worker: () -> Datadog::CI::Worker