implement local context to keep track of current test being executed

lib/datadog/ci.rb

@@ -28,6 +28,19 @@ def trace(span_type, span_name, tags: {}, &block)
         recorder.trace(span_type, span_name, tags: tags, &block)
       end
 
+      def active_span(span_type)
+        span = recorder.active_span
+        span if span && span.span_type == span_type
+      end
+
+      def active_test
+        recorder.active_test
+      end
+
+      def deactivate_test(test)
+        recorder.deactivate_test(test)
+      end
+
       private
 
       def components

lib/datadog/ci/context/local.rb

@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+require "datadog/core/utils/sequence"
+
+module Datadog
+  module CI
+    module Context
+      class Local
+        def initialize
+          @key = "datadog_ci_active_test_#{Local.next_instance_id}"
+        end
+
+        def activate_test!(test)
+          raise "Nested tests are not supported. Currently active test: #{active_test}" unless active_test.nil?
+
+          if block_given?
+            begin
+              self.active_test = test
+              yield
+            ensure
+              self.active_test = nil
+            end
+          else
+            self.active_test = test
+          end
+        end
+
+        def deactivate_test!(test)
+          return if active_test.nil?
+
+          if active_test == test
+            self.active_test = nil
+          else
+            raise "Trying to deactivate test #{test}, but currently active test is #{active_test}"
+          end
+        end
+
+        def active_test
+          Thread.current[@key]
+        end
+
+        UNIQUE_INSTANCE_MUTEX = Mutex.new
+        UNIQUE_INSTANCE_GENERATOR = Datadog::Core::Utils::Sequence.new
+
+        private_constant :UNIQUE_INSTANCE_MUTEX, :UNIQUE_INSTANCE_GENERATOR
+
+        def self.next_instance_id
+          UNIQUE_INSTANCE_MUTEX.synchronize { UNIQUE_INSTANCE_GENERATOR.next }
+        end
+
+        private
+
+        def active_test=(test)
+          Thread.current[@key] = test
+        end
+      end
+    end
+  end
+end

lib/datadog/ci/contrib/cucumber/formatter.rb

@@ -27,7 +27,7 @@ def bind_events(config)
           end
 
           def on_test_case_started(event)
-            @current_feature_span = CI.start_test(
+            CI.start_test(
               event.test_case.name,
               tags: {
                 CI::Ext::Test::TAG_FRAMEWORK => Ext::FRAMEWORK,
@@ -41,35 +41,37 @@ def on_test_case_started(event)
           end
 
           def on_test_case_finished(event)
-            return if @current_feature_span.nil?
+            test_span = CI.active_test
+            return if test_span.nil?
 
             if event.result.skipped?
-              @current_feature_span.skipped!
+              test_span.skipped!
             elsif event.result.ok?
-              @current_feature_span.passed!
+              test_span.passed!
             elsif event.result.failed?
-              @current_feature_span.failed!
+              test_span.failed!
             end
 
-            @current_feature_span.finish
+            test_span.finish
           end
 
           def on_test_step_started(event)
-            @current_step_span = CI.trace(Ext::STEP_SPAN_TYPE, event.test_step.to_s)
+            CI.trace(Ext::STEP_SPAN_TYPE, event.test_step.to_s)
           end
 
           def on_test_step_finished(event)
-            return if @current_step_span.nil?
+            current_step_span = CI.active_span(Ext::STEP_SPAN_TYPE)
+            return if current_step_span.nil?
 
             if event.result.skipped?
-              @current_step_span.skipped!
+              current_step_span.skipped!
             elsif event.result.ok?
-              @current_step_span.passed!
+              current_step_span.passed!
             elsif event.result.failed?
-              @current_step_span.failed!(exception: event.result.exception)
+              current_step_span.failed!(exception: event.result.exception)
             end
 
-            @current_step_span.finish
+            current_step_span.finish
           end
 
           private

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

@@ -18,7 +18,7 @@ def before_setup
             path, = method(name).source_location
             test_suite = Pathname.new(path.to_s).relative_path_from(Pathname.pwd).to_s
 
-            test_span = CI.start_test(
+            CI.start_test(
               test_name,
               tags: {
                 CI::Ext::Test::TAG_FRAMEWORK => Ext::FRAMEWORK,
@@ -29,16 +29,12 @@ def before_setup
               service_name: configuration[:service_name],
               operation_name: configuration[:operation_name]
             )
-
-            Thread.current[:_datadog_test_span] = test_span
           end
 
           def after_teardown
-            test_span = Thread.current[:_datadog_test_span]
+            test_span = CI.active_test
             return super unless test_span
 
-            Thread.current[:_datadog_test_span] = nil
-
             case result_code
             when "."
               test_span.passed!

lib/datadog/ci/recorder.rb

@@ -8,7 +8,10 @@
 require_relative "ext/test"
 require_relative "ext/environment"
 
+require_relative "context/local"
+
 require_relative "span"
+require_relative "test"
 
 module Datadog
   module CI
@@ -18,6 +21,7 @@ class Recorder
 
       def initialize
         @environment_tags = Ext::Environment.tags(ENV)
+        @local_context = Context::Local.new
       end
 
       # Creates a new span for a CI test
@@ -31,7 +35,32 @@ def trace_test(test_name, service_name: nil, operation_name: "test", tags: {}, &
         tags[Ext::Test::TAG_NAME] = test_name
         tags.merge!(environment_tags)
 
-        create_datadog_span(operation_name, span_options: span_options, tags: tags, &block)
+        if block
+          Datadog::Tracing.trace(operation_name, **span_options) do |tracer_span, trace|
+            set_internal_tracing_context!(trace, tracer_span)
+
+            test = Test.new(tracer_span, tags)
+            # test.set_default_tags!
+            # test.set_tags!(environment_tags)
+            # test.set_tags!(tags)
+
+            @local_context.activate_test!(test) do
+              block.call(test)
+            end
+          end
+        else
+          tracer_span = Datadog::Tracing.trace(operation_name, **span_options)
+          trace = Datadog::Tracing.active_trace
+
+          set_internal_tracing_context!(trace, tracer_span)
+          test = Test.new(tracer_span, tags)
+          # test.set_default_tags!
+          # test.set_tags!(environment_tags)
+          # test.set_tags!(tags)
+
+          @local_context.activate_test!(test)
+          test
+        end
       end
 
       def trace(span_type, span_name, tags: {}, &block)
@@ -40,26 +69,33 @@ def trace(span_type, span_name, tags: {}, &block)
           span_type: span_type
         }
 
-        create_datadog_span(span_name, span_options: span_options, tags: tags, &block)
-      end
-
-      private
-
-      def create_datadog_span(span_name, span_options: {}, tags: {}, &block)
+        # create_datadog_span(span_name, span_options: span_options, tags: tags, &block)
         if block
           Datadog::Tracing.trace(span_name, **span_options) do |tracer_span, trace|
-            set_internal_tracing_context!(trace, tracer_span)
             block.call(Span.new(tracer_span, tags))
           end
         else
           tracer_span = Datadog::Tracing.trace(span_name, **span_options)
-          trace = Datadog::Tracing.active_trace
 
-          set_internal_tracing_context!(trace, tracer_span)
           Span.new(tracer_span, tags)
         end
       end
 
+      def active_test
+        @local_context.active_test
+      end
+
+      def deactivate_test(test)
+        @local_context.deactivate_test!(test)
+      end
+
+      def active_span
+        tracer_span = Datadog::Tracing.active_span
+        Span.new(tracer_span) if tracer_span
+      end
+
+      private
+
       def set_internal_tracing_context!(trace, span)
         # Sets trace's origin to ciapp-test
         trace.origin = Ext::Test::CONTEXT_ORIGIN if trace

lib/datadog/ci/span.rb

@@ -46,6 +46,10 @@ def finish
         tracer_span.finish
       end
 
+      def span_type
+        tracer_span.type
+      end
+
       private
 
       def set_tags!(tags)

lib/datadog/ci/test.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+require_relative "span"
+
+module Datadog
+  module CI
+    class Test < Span
+      def finish
+        super
+
+        CI.deactivate_test(self)
+      end
+    end
+  end
+end

sig/datadog/ci.rbs

@@ -1,11 +1,17 @@
 module Datadog
   module CI
-    def self.trace_test: (String span_name, ?service_name: String?, ?operation_name: String, ?tags: Hash[untyped, untyped]) ?{ (Datadog::CI::Span span) -> untyped } -> untyped
+    def self.trace_test: (String span_name, ?service_name: String?, ?operation_name: String, ?tags: Hash[untyped, untyped]) ?{ (Datadog::CI::Test test) -> untyped } -> untyped
 
-    def self.start_test: (String span_name, ?service_name: String?, ?operation_name: String, ?tags: Hash[untyped, untyped]) -> Datadog::CI::Span
+    def self.start_test: (String span_name, ?service_name: String?, ?operation_name: String, ?tags: Hash[untyped, untyped]) -> Datadog::CI::Test
 
     def self.trace: (String span_type, String span_name, ?tags: Hash[untyped, untyped]) ?{ (Datadog::CI::Span span) -> untyped } -> untyped
 
+    def self.active_test: () -> Datadog::CI::Test?
+
+    def self.active_span: (String span_type) -> Datadog::CI::Span?
+
+    def self.deactivate_test: (Datadog::CI::Test test) -> void
+
     def self.components: () -> Datadog::CI::Configuration::Components
 
     def self.recorder: () -> Datadog::CI::Recorder