instrument knapsack_pro runner in queue mode

Steepfile

@@ -25,4 +25,5 @@ target :lib do
   library "cucumber"
   library "msgpack"
   library "ci_queue"
+  library "knapsack_pro"
 end

lib/datadog/ci/contrib/rspec/knapsack_pro/extension.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require "knapsack_pro/extensions/rspec_extension"
+
+require_relative "runner"
+
+module Datadog
+  module CI
+    module Contrib
+      module RSpec
+        # Instrument RSpec::Core::Example
+        module KnapsackPro
+          module Extension
+            def self.included(base)
+              base.singleton_class.prepend(ClassMethods)
+            end
+
+            module ClassMethods
+              def setup!
+                super
+
+                ::RSpec::Core::Runner.include(Datadog::CI::Contrib::RSpec::KnapsackPro::Runner)
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/ci/contrib/rspec/knapsack_pro/runner.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require_relative "../../../ext/test"
+require_relative "../ext"
+
+module Datadog
+  module CI
+    module Contrib
+      module RSpec
+        module KnapsackPro
+          module Runner
+            def self.included(base)
+              base.prepend(InstanceMethods)
+            end
+
+            module InstanceMethods
+              def knapsack__run_specs(*)
+                return super unless datadog_configuration[:enabled]
+
+                test_session = CI.start_test_session(
+                  tags: {
+                    CI::Ext::Test::TAG_FRAMEWORK => Ext::FRAMEWORK,
+                    CI::Ext::Test::TAG_FRAMEWORK_VERSION => CI::Contrib::RSpec::Integration.version.to_s
+                  },
+                  service: datadog_configuration[:service_name]
+                )
+
+                test_module = CI.start_test_module(Ext::FRAMEWORK)
+
+                result = super
+                return result unless test_module && test_session
+
+                if result != 0
+                  test_module.failed!
+                  test_session.failed!
+                else
+                  test_module.passed!
+                  test_session.passed!
+                end
+                test_module.finish
+                test_session.finish
+
+                result
+              end
+
+              private
+
+              def datadog_configuration
+                Datadog.configuration.ci[:rspec]
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -27,6 +27,13 @@ def patch
               ::RSpec::Queue::Runner.include(Runner)
             end
 
+            # Knapsack Pro test runner instrumentation
+            # https://github.com/KnapsackPro/knapsack_pro-ruby
+            if knapsack_pro?
+              require_relative "knapsack_pro/extension"
+              ::KnapsackPro::Extensions::RSpecExtension.include(KnapsackPro::Extension)
+            end
+
             ::RSpec::Core::Runner.include(Runner)
             ::RSpec::Core::Example.include(Example)
             ::RSpec::Core::ExampleGroup.include(ExampleGroup)
@@ -35,6 +42,14 @@ def patch
           def ci_queue?
             defined?(::RSpec::Queue::Runner)
           end
+
+          def knapsack_pro?
+            knapsack_version = Gem.loaded_specs["knapsack_pro"]&.version
+
+            # additional instrumentation is needed for KnapsackPro version 7 and later
+            defined?(::KnapsackPro) &&
+              knapsack_version && knapsack_version >= Gem::Version.new("7")
+          end
         end
       end
     end

sig/datadog/ci/contrib/rspec/knapsack_pro/extension.rbs

@@ -0,0 +1,19 @@
+module Datadog
+  module CI
+    module Contrib
+      module RSpec
+        module KnapsackPro
+          module Extension
+            def self.included: (untyped base) -> untyped
+
+            module ClassMethods
+              include ::KnapsackPro::Extensions::RSpecExtension
+
+              def setup!: () -> void
+            end
+          end
+        end
+      end
+    end
+  end
+end

sig/datadog/ci/contrib/rspec/knapsack_pro/runner.rbs

@@ -0,0 +1,23 @@
+module Datadog
+  module CI
+    module Contrib
+      module RSpec
+        module KnapsackPro
+          module Runner
+            def self.included: (untyped base) -> untyped
+
+            module InstanceMethods
+              include ::KnapsackPro::Runners::Queue::RSpecRunner
+
+              def knapsack__run_specs: (untyped args) -> untyped
+
+              private
+
+              def datadog_configuration: () -> untyped
+            end
+          end
+        end
+      end
+    end
+  end
+end

sig/datadog/ci/contrib/rspec/patcher.rbs

@@ -10,6 +10,8 @@ module Datadog
           def self?.patch: () -> void
 
           def self?.ci_queue?: () -> bool
+
+          def self?.knapsack_pro?: () -> bool
         end
       end
     end

vendor/rbs/knapsack_pro/0/knapsack_pro.rbs

@@ -0,0 +1,17 @@
+module KnapsackPro
+  module Runners
+    module Queue
+      module RSpecRunner
+        def knapsack__run_specs: (untyped arg) -> untyped
+      end
+    end
+  end
+end
+
+module KnapsackPro
+  module Extensions
+    module RSpecExtension
+      def setup!: () -> void
+    end
+  end
+end