From 9e388f13331ef6cf962fb1f7e6d5274b408d3b10 Mon Sep 17 00:00:00 2001
From: Oleg Pudeyev <code@olegp.name>
Date: Thu, 18 Jul 2024 12:38:13 -0400
Subject: [PATCH] di code tracker test

---
 lib/datadog/di/code_tracker.rb               | 14 ++++---
 spec/datadog/di/code_tracker_spec.rb         | 41 ++++++++++++++++++++
 spec/datadog/di/code_tracker_test_class_1.rb |  2 +
 3 files changed, 51 insertions(+), 6 deletions(-)
 create mode 100644 spec/datadog/di/code_tracker_spec.rb
 create mode 100644 spec/datadog/di/code_tracker_test_class_1.rb

diff --git a/lib/datadog/di/code_tracker.rb b/lib/datadog/di/code_tracker.rb
index d041da251aa..9c1987424d5 100644
--- a/lib/datadog/di/code_tracker.rb
+++ b/lib/datadog/di/code_tracker.rb
@@ -13,7 +13,7 @@ module DI
     # recreated when the DI component is created.
     class CodeTracker
       def initialize
-        @file_registry = Concurrent::Map.new
+        @registry = Concurrent::Map.new
       end
 
       def start
@@ -37,10 +37,10 @@ def start
           #
           # For now just map the path to the instruction sequence.
           path = tp.instruction_sequence.path
-          file_registry[path] = tp.instruction_sequence
+          registry[path] = tp.instruction_sequence
 
           # TODO fix this to properly deal with paths
-          file_registry[File.basename(path)] = tp.instruction_sequence
+          registry[File.basename(path)] = tp.instruction_sequence
 
           DI.component&.hook_manager&.install_pending_line_hooks(path)
         end
@@ -56,7 +56,7 @@ def active?
       # Returns the RubVM::InstructionSequence (i.e. the compiled code)
       # for the provided path.
       def [](path)
-        file_registry[path]
+        registry[path]
       end
 
       def stop
@@ -65,12 +65,14 @@ def stop
         # Clear the instance variable so that the trace point may be
         # reinstated in the future.
         @compiled_trace_point = nil
-        file_registry.clear
+        registry.clear
       end
 
       private
 
-      attr_reader :file_registry
+      # Mapping from paths of loaded files to RubyVM::InstructionSequence
+      # objects representing compiled code of those files.
+      attr_reader :registry
     end
   end
 end
diff --git a/spec/datadog/di/code_tracker_spec.rb b/spec/datadog/di/code_tracker_spec.rb
new file mode 100644
index 00000000000..ff8ff9013e1
--- /dev/null
+++ b/spec/datadog/di/code_tracker_spec.rb
@@ -0,0 +1,41 @@
+require 'datadog/di/code_tracker'
+
+RSpec.describe Datadog::DI::CodeTracker do
+  let(:tracker) do
+    described_class.new
+  end
+
+  describe '.new' do
+    it 'creates an instance' do
+      expect(tracker).to be_a(described_class)
+    end
+  end
+
+  describe '#start' do
+    after do
+      tracker.stop
+    end
+
+    it 'tracks loaded files' do
+      # The expectations appear to be lazy-loaded, therefore
+      # we need to invoke the same expectation before starting
+      # code tracking as we'll be using later in the test.
+      expect(tracker.send(:registry)).to be_empty
+      tracker.start
+      # Should still be empty here.
+      expect(tracker.send(:registry)).to be_empty
+      require_relative 'code_tracker_test_class_1'
+      # TODO due to a hack we currently have 2 entries for every file,
+      # one with full path and one with basename only.
+      expect(tracker.send(:registry).each.to_a.length).to eq(2)
+
+      path = tracker.send(:registry).each.to_a.first.first
+      # The full path is dependent on the environment/system
+      # running the tests, but we can assert on the basename
+      # which will be the same.
+      expect(File.basename(path)).to eq('code_tracker_test_class_1.rb')
+      # And, we should in fact have a full path.
+      expect(path).to start_with('/')
+    end
+  end
+end
diff --git a/spec/datadog/di/code_tracker_test_class_1.rb b/spec/datadog/di/code_tracker_test_class_1.rb
new file mode 100644
index 00000000000..60f4b30f80c
--- /dev/null
+++ b/spec/datadog/di/code_tracker_test_class_1.rb
@@ -0,0 +1,2 @@
+class CodeTrackerTestClass1
+end