lib/opentelemetry/test_helpers.rb in opentelemetry-test-helpers-0.3.0 vs lib/opentelemetry/test_helpers.rb in opentelemetry-test-helpers-0.4.0

@@ -53,14 +53,12 @@
     ensure
       env_to_reset.each_pair { |k, v| ENV[k] = v }
       keys_to_delete.each { |k| ENV.delete(k) }
     end
 
-    def with_ids(trace_id, span_id)
+    def with_ids(trace_id, span_id, &block)
       OpenTelemetry::Trace.stub(:generate_trace_id, trace_id) do
-        OpenTelemetry::Trace.stub(:generate_span_id, span_id) do
-          yield
-        end
+        OpenTelemetry::Trace.stub(:generate_span_id, span_id, &block)
       end
     end
 
     def create_span_data(name: '', kind: nil, status: nil, parent_span_id: OpenTelemetry::Trace::INVALID_SPAN_ID,
                          total_recorded_attributes: 0, total_recorded_events: 0, total_recorded_links: 0, start_timestamp: OpenTelemetry::TestHelpers.exportable_timestamp,