From 6ceefdc64b1625b37be30b54ffb2487587869a67 Mon Sep 17 00:00:00 2001
From: Oleg Pudeyev <code@olegp.name>
Date: Mon, 15 Jul 2024 12:38:15 -0400
Subject: [PATCH] di delayed

---
 lib/datadog/di/hook_manager.rb       | 8 ++++++++
 spec/datadog/di/hook_manager_spec.rb | 7 +++++++
 2 files changed, 15 insertions(+)

diff --git a/lib/datadog/di/hook_manager.rb b/lib/datadog/di/hook_manager.rb
index 88d3aa987a2..5dc2e13b7cc 100644
--- a/lib/datadog/di/hook_manager.rb
+++ b/lib/datadog/di/hook_manager.rb
@@ -145,6 +145,14 @@ def hook_line(file, line_no, &block)
       # defined when the hook request is received.
       attr_reader :definition_trace_point
 
+      def pending_methods
+        PENDING_METHODS
+      end
+
+      def instrumented_methods
+        INSTRUMENTED_METHODS
+      end
+
       PENDING_METHODS = Concurrent::Map.new
       INSTRUMENTED_METHODS = Concurrent::Map.new
       INSTRUMENTED_LINES = Concurrent::Map.new
diff --git a/spec/datadog/di/hook_manager_spec.rb b/spec/datadog/di/hook_manager_spec.rb
index c3ee4254830..9408d19c4be 100644
--- a/spec/datadog/di/hook_manager_spec.rb
+++ b/spec/datadog/di/hook_manager_spec.rb
@@ -127,12 +127,19 @@ def hook_test_method_with_kwarg(kwarg:)
           invoked = true
         end).to be false
 
+        expect(manager.send(:pending_methods)[[:HookManagerTestLateDefinition, :test_method]]).to be_a(Proc)
+        expect(manager.send(:instrumented_methods)[[:HookManagerTestLateDefinition, :test_method]]).to be nil
+
         class HookManagerTestLateDefinition
           def test_method
             42
           end
         end
 
+        # Method should now be hooked, and no longer pending
+        expect(manager.send(:pending_methods)[[:HookManagerTestLateDefinition, :test_method]]).to be nil
+        expect(manager.send(:instrumented_methods)[[:HookManagerTestLateDefinition, :test_method]]).to be_a(Integer)
+
         expect(HookManagerTestLateDefinition.new.test_method).to eq 42
 
         expect(invoked).to be true