From ced48cee5fbb923a91e4757d4f15375345cc8805 Mon Sep 17 00:00:00 2001
From: Andrey <andrey.marchenko@datadoghq.com>
Date: Wed, 20 Nov 2024 14:23:45 +0100
Subject: [PATCH] extract patch_integration method in Contrib::Instrumentation

---
 lib/datadog/ci/contrib/instrumentation.rb  | 53 ++++++++++++----------
 sig/datadog/ci/contrib/instrumentation.rbs |  5 +-
 sig/datadog/ci/contrib/integration.rbs     |  2 +
 3 files changed, 34 insertions(+), 26 deletions(-)

diff --git a/lib/datadog/ci/contrib/instrumentation.rb b/lib/datadog/ci/contrib/instrumentation.rb
index a283fe6c..9430980a 100644
--- a/lib/datadog/ci/contrib/instrumentation.rb
+++ b/lib/datadog/ci/contrib/instrumentation.rb
@@ -21,29 +21,12 @@ def self.register_integration(integration_class)
         # This method is called when user has `c.ci.instrument :integration_name` in their code.
         def self.instrument(integration_name, options = {}, &block)
           integration = fetch_integration(integration_name)
+          # when manually instrumented, it might be configured via code
           integration.configure(options, &block)
 
           return unless integration.enabled
 
-          patch_results = integration.patch
-          if patch_results[:ok]
-            # try to patch dependant integrations (for example knapsack that depends on rspec)
-            dependants = integration.dependants
-              .map { |name| fetch_integration(name) }
-              .filter { |integration| integration.patchable? }
-
-            Datadog.logger.debug("Found dependent integrations for #{integration_name}: #{dependants}")
-
-            dependants.each do |dependent_integration|
-              dependent_integration.patch
-            end
-          else
-            error_message = <<-ERROR
-                  Available?: #{patch_results[:available]}, Loaded?: #{patch_results[:loaded]},
-                  Compatible?: #{patch_results[:compatible]}, Patchable?: #{patch_results[:patchable]}"
-            ERROR
-            Datadog.logger.warn("Unable to patch #{integration_name} (#{error_message})")
-          end
+          patch_integration(integration, with_dependencies: true)
         end
 
         # This method instruments all additional test libraries (ex: selenium-webdriver) that need to be instrumented
@@ -58,15 +41,11 @@ def self.instrument_on_session_start
 
           @registry.each do |name, integration|
             next unless integration.late_instrument?
+            next unless integration.enabled
 
             Datadog.logger.debug "#{name} is allowed to be late instrumented"
 
-            patch_results = integration.patch
-            if patch_results[:ok]
-              Datadog.logger.debug("#{name} is patched")
-            else
-              Datadog.logger.debug("#{name} is not patched (#{patch_results})")
-            end
+            patch_integration(integration)
           end
         end
 
@@ -82,6 +61,30 @@ def self.integration_name(subclass)
           raise "Integration name could not be derived for #{subclass}" if result.nil?
           result
         end
+
+        def self.patch_integration(integration, with_dependencies: false)
+          patch_results = integration.patch
+
+          if patch_results[:ok]
+            Datadog.logger.debug("#{integration.class} is patched")
+
+            return unless with_dependencies
+
+            # try to patch dependant integrations (for example knapsack that depends on rspec)
+            dependants = integration.dependants
+              .map { |name| fetch_integration(name) }
+              .filter { |integration| integration.patchable? }
+
+            Datadog.logger.debug("Found dependent integrations for #{integration.class}: #{dependants}")
+
+            dependants.each do |dependent_integration|
+              patch_integration(dependent_integration, with_dependencies: true)
+            end
+
+          else
+            Datadog.logger.debug("Attention: #{integration.class} is not patched (#{patch_results})")
+          end
+        end
       end
     end
   end
diff --git a/sig/datadog/ci/contrib/instrumentation.rbs b/sig/datadog/ci/contrib/instrumentation.rbs
index a10222fd..1f511280 100644
--- a/sig/datadog/ci/contrib/instrumentation.rbs
+++ b/sig/datadog/ci/contrib/instrumentation.rbs
@@ -11,13 +11,16 @@ module Datadog
 
         def self.instrument: (Symbol integration_name, ?::Hash[untyped, untyped] options) { (?) -> untyped } -> void
 
+        def self.instrument_on_session_start: () -> void
+
         def self.fetch_integration: (Symbol name) -> untyped
 
         def self.integration_name: (Class) -> Symbol
 
         def self.register_integration: (Class integration) -> void
 
-        def self.instrument_on_session_start: () -> void
+        def self.patch_integration: (Contrib::Integration integration, ?with_dependencies: bool) -> void
+
       end
     end
   end
diff --git a/sig/datadog/ci/contrib/integration.rbs b/sig/datadog/ci/contrib/integration.rbs
index 96be033b..a5f6fbac 100644
--- a/sig/datadog/ci/contrib/integration.rbs
+++ b/sig/datadog/ci/contrib/integration.rbs
@@ -24,6 +24,8 @@ module Datadog
 
         def patch: () -> Hash[Symbol, bool]
 
+        def dependants: () -> Array[Symbol]
+
         def late_instrument?: () -> bool
 
         def new_configuration: () -> Datadog::CI::Contrib::Settings