type checking

lib/datadog/ci.rb

@@ -12,13 +12,11 @@
 module Datadog
   # Public API for Datadog CI visibility
   module CI
-    module_function
-
-    def trace_test(test_name, service_name: nil, operation_name: nil, tags: {}, &block)
+    def self.trace_test(test_name, service_name: nil, operation_name: "test", tags: {}, &block)
       Recorder.trace_test(test_name, service_name: service_name, operation_name: operation_name, tags: tags, &block)
     end
 
-    def trace(span_type, span_name, &block)
+    def self.trace(span_type, span_name, &block)
       Recorder.trace(span_type, span_name, &block)
     end
   end

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

@@ -35,20 +35,19 @@ def before_setup
           end
 
           def after_teardown
-            return super unless @current_test_span
-
-            Thread.current[:_datadog_test_span] = nil
+            test_span = @current_test_span
+            return super unless test_span
 
             case result_code
             when "."
-              @current_test_span.passed!
+              test_span.passed!
             when "E", "F"
-              @current_test_span.failed!(failure)
+              test_span.failed!(failure)
             when "S"
-              @current_test_span.skipped!(nil, failure.message)
+              test_span.skipped!(nil, failure.message)
             end
 
-            @current_test_span.finish
+            test_span.finish
             @current_test_span = nil
 
             super

lib/datadog/ci/recorder.rb

@@ -14,7 +14,7 @@ module CI
     # Common behavior for CI tests
     module Recorder
       # Creates a new span for a CI test
-      def self.trace_test(test_name, service_name: nil, operation_name: nil, tags: {})
+      def self.trace_test(test_name, service_name: nil, operation_name: "test", tags: {}, &block)
         span_options = {
           resource: test_name,
           service: service_name,
@@ -23,10 +23,11 @@ def self.trace_test(test_name, service_name: nil, operation_name: nil, tags: {})
 
         tags[:test_name] = test_name
 
-        if block_given?
+        if block
           ::Datadog::Tracing.trace(operation_name, **span_options) do |span, trace|
             set_tags!(trace, span, tags)
-            yield(Test.new(span))
+
+            block.call(Test.new(span))
           end
         else
           span = ::Datadog::Tracing.trace(operation_name, **span_options)
@@ -36,15 +37,15 @@ def self.trace_test(test_name, service_name: nil, operation_name: nil, tags: {})
         end
       end
 
-      def self.trace(span_type, span_name)
+      def self.trace(span_type, span_name, &block)
         span_options = {
           resource: span_name,
           span_type: span_type
         }
 
-        if block_given?
+        if block
           ::Datadog::Tracing.trace(span_name, **span_options) do |tracer_span|
-            yield Span.new(tracer_span)
+            block.call(Span.new(tracer_span))
           end
         else
           tracer_span = Datadog::Tracing.trace(span_name, **span_options)

sig/datadog/ci.rbs

@@ -1,7 +1,7 @@
 module Datadog
   module CI
-    # See the writing guide of rbs: https://github.com/ruby/rbs#guides
-    class Error < StandardError
-    end
+    def self.trace_test: (String span_name, ?service_name: String?, ?operation_name: String, ?tags: Hash[untyped, untyped]) ?{ (Datadog::CI::Test span) -> untyped } -> untyped
+
+    def self.trace: (String span_type, String span_name) ?{ (Datadog::CI::Span span) -> untyped } -> untyped
   end
 end

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

@@ -5,6 +5,8 @@ module Datadog
         module Hooks : ::Minitest::Test
           include ::Minitest::Test::LifecycleHooks
 
+          @current_test_span: Datadog::CI::Test?
+
           def before_setup: () -> (nil | untyped)
 
           def after_teardown: () -> untyped

sig/datadog/ci/recorder.rbs

@@ -3,16 +3,13 @@ module Datadog
     module Recorder
       self.@environment_tags: Hash[String, String]
 
-      def self.trace: (untyped span_name, ?::Hash[untyped, untyped] options) ?{ (untyped, untyped) -> untyped } -> untyped
-      def self.set_tags!: (untyped trace, untyped span, ?::Hash[untyped, untyped] tags) -> untyped
+      def self.trace_test: (String span_name, ?service_name: String?, ?operation_name: String, ?tags: Hash[untyped, untyped]) ?{ (Datadog::CI::Test span) -> untyped } -> untyped
 
-      def self.passed!: (untyped span) -> untyped
+      def self.trace: (String span_type, String span_name) ?{ (Datadog::CI::Span span) -> untyped } -> untyped
 
-      def self.failed!: (untyped span, ?untyped? exception) -> untyped
+      def self.set_tags!: (Datadog::Tracing::TraceOperation trace, Datadog::Tracing::SpanOperation span, ?::Hash[untyped, untyped] tags)  -> untyped
 
-      def self.skipped!: (untyped span, ?untyped? exception) -> untyped
-
-      def self.set_environment_runtime_tags!: (untyped span) -> untyped
+      def self.set_environment_runtime_tags!: (Datadog::Tracing::SpanOperation span) -> untyped
     end
   end
 end

sig/datadog/ci/span.rbs

@@ -0,0 +1,19 @@
+module Datadog
+  module CI
+    class Span
+      @tracer_span: Datadog::Tracing::SpanOperation
+
+      attr_reader tracer_span: Datadog::Tracing::SpanOperation
+
+      def initialize: (Datadog::Tracing::SpanOperation tracer_span) -> void
+
+      def passed!: () -> void
+
+      def failed!: (?untyped? exception) -> void
+
+      def skipped!: (?untyped? exception, ?String? reason) -> void
+
+      def finish: () -> untyped
+    end
+  end
+end

sig/datadog/ci/test.rbs

@@ -0,0 +1,6 @@
+module Datadog
+  module CI
+    class Test < Span
+    end
+  end
+end

vendor/rbs/ddtrace/0/datadog/tracing.rbs

@@ -1,6 +1,6 @@
 module Datadog
   module Tracing
-    def self.active_trace: () -> Datadog::Tracing::TraceSegment
+    def self.active_trace: () -> Datadog::Tracing::TraceOperation
     def self.trace: (String span_name, Hash[untyped, untyped] options) ?{ (untyped span, untyped trace) -> untyped } -> Datadog::Tracing::SpanOperation
   end
 end