From b8a4c14a32cbafe36843ec5fd001b2ae061f377f Mon Sep 17 00:00:00 2001
From: Andrey Marchenko <andrey.marchenko@datadoghq.com>
Date: Wed, 22 Nov 2023 10:13:32 +0100
Subject: [PATCH] wip

---
 lib/datadog/ci/recorder.rb       | 10 ++++++
 spec/datadog/ci/recorder_spec.rb | 52 ++++++++++++++++----------------
 2 files changed, 36 insertions(+), 26 deletions(-)

diff --git a/lib/datadog/ci/recorder.rb b/lib/datadog/ci/recorder.rb
index 474b2744..5b9dde07 100644
--- a/lib/datadog/ci/recorder.rb
+++ b/lib/datadog/ci/recorder.rb
@@ -150,6 +150,16 @@ def build_span(tracer_span, tags)
 
         span
       end
+
+      def set_initial_tags(ci_span, tags)
+        tags.each do |key, value|
+          ci_span.set_default_tags
+          ci_span.set_environment_runtime_tags
+
+          ci_span.set_tags(tags)
+          ci_span.set_tags(environment_tags)
+        end
+      end
     end
   end
 end
diff --git a/spec/datadog/ci/recorder_spec.rb b/spec/datadog/ci/recorder_spec.rb
index 3b6b377a..473b50df 100644
--- a/spec/datadog/ci/recorder_spec.rb
+++ b/spec/datadog/ci/recorder_spec.rb
@@ -6,6 +6,10 @@
   let(:tags) { {} }
   let(:environment_tags) { Datadog::CI::Ext::Environment.tags(ENV) }
 
+  let(:ci_span) do
+    spy("CI object spy")
+  end
+
   subject(:recorder) { described_class.new }
 
   before do
@@ -21,15 +25,16 @@
     end
   end
 
-  describe "#trace_test" do
-    def expect_initialized_test
-      allow(Datadog::CI::Test).to receive(:new).with(span_op).and_return(ci_test)
-      expect(ci_test).to receive(:set_default_tags)
-      expect(ci_test).to receive(:set_environment_runtime_tags)
-      expect(ci_test).to receive(:set_tags).with(tags)
-      expect(ci_test).to receive(:set_tags).with(environment_tags)
+  shared_examples_for "initialize ci span with tags" do
+    it do
+      expect(ci_span).to have_received(:set_default_tags)
+      expect(ci_span).to have_received(:set_environment_runtime_tags)
+      expect(ci_span).to have_received(:set_tags).with(tags)
+      expect(ci_span).to have_received(:set_tags).with(environment_tags)
     end
+  end
 
+  describe "#trace_test" do
     context "when given a block" do
       subject(:trace) do
         recorder.trace_test(
@@ -42,7 +47,6 @@ def expect_initialized_test
       end
 
       let(:span_op) { Datadog::Tracing::SpanOperation.new(operation_name) }
-      let(:ci_test) { instance_double(Datadog::CI::Test) }
       let(:block) { proc { |s| block_spy.call(s) } }
       let(:block_result) { double("result") }
       let(:block_spy) { spy("block") }
@@ -61,7 +65,7 @@ def expect_initialized_test
               }
             )
 
-            expect_initialized_test
+            allow(Datadog::CI::Test).to receive(:new).with(span_op).and_return(ci_span)
 
             trace_block.call(span_op, trace_op)
           end
@@ -70,7 +74,9 @@ def expect_initialized_test
       end
 
       it_behaves_like "internal tracing context"
-      it { expect(block_spy).to have_received(:call).with(ci_test) }
+      it_behaves_like "initialize ci span with tags"
+
+      it { expect(block_spy).to have_received(:call).with(ci_span) }
       it { is_expected.to be(block_result) }
     end
 
@@ -84,7 +90,6 @@ def expect_initialized_test
         )
       end
       let(:span_op) { Datadog::Tracing::SpanOperation.new(operation_name) }
-      let(:ci_test) { instance_double(Datadog::CI::Test) }
 
       before do
         allow(Datadog::Tracing)
@@ -99,24 +104,18 @@ def expect_initialized_test
           )
           .and_return(span_op)
 
-        expect_initialized_test
+        allow(Datadog::CI::Test).to receive(:new).with(span_op).and_return(ci_span)
 
         trace
       end
 
       it_behaves_like "internal tracing context"
-      it { is_expected.to be(ci_test) }
+      it_behaves_like "initialize ci span with tags"
+      it { is_expected.to be(ci_span) }
     end
   end
 
   describe "#trace" do
-    def expect_initialized_span
-      allow(Datadog::CI::Span).to receive(:new).with(span_op).and_return(ci_span)
-      expect(ci_span).to receive(:set_default_tags)
-      expect(ci_span).to receive(:set_environment_runtime_tags)
-      expect(ci_span).to receive(:set_tags).with(tags)
-    end
-
     let(:tags) { {"my_tag" => "my_value"} }
     let(:span_type) { "step" }
     let(:span_name) { "span name" }
@@ -134,7 +133,6 @@ def expect_initialized_span
       end
 
       let(:span_op) { Datadog::Tracing::SpanOperation.new(span_name) }
-      let(:ci_span) { instance_double(Datadog::CI::Span) }
       let(:block) { proc { |s| block_spy.call(s) } }
       let(:block_result) { double("result") }
       let(:block_spy) { spy("block") }
@@ -154,11 +152,12 @@ def expect_initialized_span
             trace_block.call(span_op, trace_op)
           end
 
-        expect_initialized_span
+        allow(Datadog::CI::Span).to receive(:new).with(span_op).and_return(ci_span)
 
         trace
       end
 
+      it_behaves_like "initialize ci span with tags"
       it { expect(block_spy).to have_received(:call).with(ci_span) }
       it { is_expected.to be(block_result) }
     end
@@ -187,11 +186,12 @@ def expect_initialized_span
           )
           .and_return(span_op)
 
-        expect_initialized_span
+        allow(Datadog::CI::Span).to receive(:new).with(span_op).and_return(ci_span)
 
         trace
       end
 
+      it_behaves_like "initialize ci span with tags"
       it { is_expected.to be(ci_span) }
     end
   end
@@ -271,10 +271,10 @@ def expect_initialized_span
           }
         )
         .and_return(span_op)
-
-      allow(Datadog::CI::TestSession).to receive(:new).with(span_op).and_return(ci_test_session)
     end
 
-    it { is_expected.to be(ci_test_session) }
+    it do
+      expect(start_test_session.tracer_span).to be(span_op)
+    end
   end
 end