From df6d3ca0acb102e7b55e2c0b89ed32c4a46aabdc Mon Sep 17 00:00:00 2001
From: Tanna McClure <tmcclure@newrelic.com>
Date: Thu, 20 Jun 2024 10:45:10 -0500
Subject: [PATCH] refactor

---
 .../aws_sqs/instrumentation.rb                | 48 +++++--------------
 .../awssqs/awssqs_instrumentation_test.rb     | 38 +--------------
 2 files changed, 14 insertions(+), 72 deletions(-)

diff --git a/lib/new_relic/agent/instrumentation/aws_sqs/instrumentation.rb b/lib/new_relic/agent/instrumentation/aws_sqs/instrumentation.rb
index f18d94d577..e78bfcb9d1 100644
--- a/lib/new_relic/agent/instrumentation/aws_sqs/instrumentation.rb
+++ b/lib/new_relic/agent/instrumentation/aws_sqs/instrumentation.rb
@@ -7,60 +7,36 @@ module AwsSqs
     MESSAGING_LIBRARY = 'SQS'
 
     def send_message_with_new_relic(*args)
-      segment = nil
-      begin
-        info = get_url_info(args[0])
-        segment = NewRelic::Agent::Tracer.start_message_broker_segment(
-          action: :produce,
-          library: MESSAGING_LIBRARY,
-          destination_type: :queue,
-          destination_name: info[:queue_name]
-        )
-        add_aws_attributes(segment, info)
-      rescue => e
-        NewRelic::Agent.logger.error('Error starting message broker segment in Aws::SQS::Client#send_message ', e)
-      end
-      NewRelic::Agent::Tracer.capture_segment_error(segment) do
+      with_tracing(:produce, args) do
         yield
       end
-    ensure
-      segment&.finish
     end
 
     def send_message_batch_with_new_relic(*args)
-      segment = nil
-      begin
-        info = get_url_info(args[0])
-        segment = NewRelic::Agent::Tracer.start_message_broker_segment(
-          action: :produce,
-          library: MESSAGING_LIBRARY,
-          destination_type: :queue,
-          destination_name: info[:queue_name]
-        )
-        add_aws_attributes(segment, info)
-      rescue => e
-        NewRelic::Agent.logger.error('Error starting message broker segment in Aws::SQS::Client#send_message_batch ', e)
-      end
-      NewRelic::Agent::Tracer.capture_segment_error(segment) do
+      with_tracing(:produce, args) do
         yield
       end
-    ensure
-      segment&.finish
     end
 
     def receive_message_with_new_relic(*args)
+      with_tracing(:consume, args) do
+        yield
+      end
+    end
+
+    def with_tracing(action, params)
       segment = nil
       begin
-        info = get_url_info(args[0])
+        info = get_url_info(params[0])
         segment = NewRelic::Agent::Tracer.start_message_broker_segment(
-          action: :consume,
+          action: action,
           library: MESSAGING_LIBRARY,
           destination_type: :queue,
           destination_name: info[:queue_name]
         )
         add_aws_attributes(segment, info)
       rescue => e
-        NewRelic::Agent.logger.error('Error starting message broker segment in Aws::SQS::Client#receive_message ', e)
+        NewRelic::Agent.logger.error('Error starting message broker segment in Aws::SQS::Client', e)
       end
       NewRelic::Agent::Tracer.capture_segment_error(segment) do
         yield
@@ -69,6 +45,8 @@ def receive_message_with_new_relic(*args)
       segment&.finish
     end
 
+    private
+
     def add_aws_attributes(segment, info)
       return unless segment
 
diff --git a/test/multiverse/suites/awssqs/awssqs_instrumentation_test.rb b/test/multiverse/suites/awssqs/awssqs_instrumentation_test.rb
index 0455c44730..03e30ea5be 100644
--- a/test/multiverse/suites/awssqs/awssqs_instrumentation_test.rb
+++ b/test/multiverse/suites/awssqs/awssqs_instrumentation_test.rb
@@ -102,42 +102,6 @@ def test_error_send_message
       end
     end
 
-    assert_log_contains(log, 'Error starting message broker segment in Aws::SQS::Client#send_message')
-  end
-
-  def test_error_send_message_batch
-    client = create_client
-
-    log = with_array_logger(:info) do
-      in_transaction do |txn|
-        begin
-          client.send_message_batch({
-            queue_url: 42
-          })
-        rescue
-          # will cause an error in the instrumentation, but also will make the sdk raise an error
-        end
-      end
-    end
-
-    assert_log_contains(log, 'Error starting message broker segment in Aws::SQS::Client#send_message_batch')
-  end
-
-  def test_error_receive_message
-    client = create_client
-
-    log = with_array_logger(:info) do
-      in_transaction do |txn|
-        begin
-          client.receive_message({
-            queue_url: 42
-          })
-        rescue
-          # will cause an error in the instrumentation, but also will make the sdk raise an error
-        end
-      end
-    end
-
-    assert_log_contains(log, 'Error starting message broker segment in Aws::SQS::Client#receive_message')
+    assert_log_contains(log, 'Error starting message broker segment in Aws::SQS::Client')
   end
 end