From 73d116864a9805b0302445235ee8aa4e73f2e3b4 Mon Sep 17 00:00:00 2001
From: Marco Costa <marco.costa@datadoghq.com>
Date: Fri, 10 Jan 2025 15:55:29 -0800
Subject: [PATCH] Fix tests

---
 spec/datadog/tracing/integration_spec.rb      | 20 ++++++++++++-------
 spec/datadog/tracing/span_spec.rb             |  3 +--
 spec/datadog/tracing/transport/traces_spec.rb |  3 ++-
 3 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/spec/datadog/tracing/integration_spec.rb b/spec/datadog/tracing/integration_spec.rb
index f83022d9ee8..649a2f17505 100644
--- a/spec/datadog/tracing/integration_spec.rb
+++ b/spec/datadog/tracing/integration_spec.rb
@@ -867,7 +867,7 @@ def agent_receives_span_step3(previous_success)
     end
   end
 
-  shared_examples 'flushes traces with span events' do
+  shared_examples 'flushes traces with span events' do |native_span_events_support: true|
     context 'a trace with span events' do
       subject(:trace_with_event) do
         tracer.trace('parent_span') do |span|
@@ -876,18 +876,20 @@ def agent_receives_span_step3(previous_success)
             time_unix_nano: 123,
             attributes: { 'key' => 'value' }
           )
-        end.tap { try_wait_until(seconds: 2) { tracer.writer.stats[:traces_flushed] >= 1 } }
+        end
+
+        try_wait_until(seconds: 2) { tracer.writer.stats[:traces_flushed] >= 1 }
       end
 
       before do
         allow_any_instance_of(Datadog::Core::Remote::Transport::HTTP::Negotiation::Response)
-          .to receive(:span_events).and_return(native_span_events_support)
+          .to receive(:span_events).and_return(span_events_support)
 
         allow(encoder).to receive(:encode).and_wrap_original do |m, *args|
           encoded = m.call(*args)
           traces = encoder.decode(encoded)
 
-          traces = traces['traces'].first if traces.key?('traces')
+          traces = traces['traces'].first if traces.is_a?(Hash) # For Transport::IO
 
           expect(traces).to have(1).item
           @flushed_trace = traces.first
@@ -899,7 +901,11 @@ def agent_receives_span_step3(previous_success)
       let(:flushed_trace) { @flushed_trace }
 
       context 'with agent supporting native span events' do
-        let(:native_span_events_support) { true }
+        before do
+          skip 'Environment does not support native span events' unless native_span_events_support
+        end
+
+        let(:span_events_support) { true }
 
         it 'flushes events using the span_events field' do
           trace_with_event
@@ -918,7 +924,7 @@ def agent_receives_span_step3(previous_success)
       end
 
       context 'with agent not supporting native span events' do
-        let(:native_span_events_support) { false }
+        let(:span_events_support) { false }
 
         it 'flushes events using the span_events field' do
           trace_with_event
@@ -993,7 +999,7 @@ def agent_receives_span_step3(previous_success)
       expect(out).to have_received(:puts)
     end
 
-    it_behaves_like 'flushes traces with span events' do
+    it_behaves_like 'flushes traces with span events', native_span_events_support: false do
       let(:encoder) { Datadog.send(:components).tracer.writer.transport.encoder }
     end
   end
diff --git a/spec/datadog/tracing/span_spec.rb b/spec/datadog/tracing/span_spec.rb
index 1e79e3c2fd8..6a061114627 100644
--- a/spec/datadog/tracing/span_spec.rb
+++ b/spec/datadog/tracing/span_spec.rb
@@ -2,7 +2,6 @@
 
 require 'json'
 require 'msgpack'
-require 'pp'
 require 'time'
 
 require 'datadog/core'
@@ -268,7 +267,7 @@
     subject(:pretty_print) { PP.pp(span) }
 
     it 'output without errors' do
-      Datadog::Tracing.trace('x') { |s| pp s}
+      Datadog::Tracing.trace('x') { |s| pp s }
       # pretty_print
       # expect { pretty_print }.to output.to_stdout
     end
diff --git a/spec/datadog/tracing/transport/traces_spec.rb b/spec/datadog/tracing/transport/traces_spec.rb
index 788712ae259..5f9bd26aa59 100644
--- a/spec/datadog/tracing/transport/traces_spec.rb
+++ b/spec/datadog/tracing/transport/traces_spec.rb
@@ -187,7 +187,8 @@
 
       allow(Datadog::Tracing::Transport::Traces::Request).to receive(:new).and_return(request)
 
-      allow_any_instance_of(Datadog::Core::Remote::Transport::HTTP::Negotiation::Response).to receive(:span_events).and_return(native_events_supported)
+      allow_any_instance_of(Datadog::Core::Remote::Transport::HTTP::Negotiation::Response).to receive(:span_events)
+        .and_return(native_events_supported)
     end
 
     context 'which returns an OK response' do