Merge pull request #2534 from newrelic/net-http-openai-tests

Net::HTTP OpenAI tests

test/new_relic/agent/llm_test.rb

@@ -0,0 +1,129 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+require_relative '../../test_helper'
+
+module NewRelic
+  module Agent
+    class LLMTest < Minitest::Test
+      def setup
+        NewRelic::Agent::LLM.remove_instance_variable(:@openai) if NewRelic::Agent::LLM.instance_variable_defined?(:@openai)
+      end
+
+      def test_openai_true_when_ruby_openai_prepend_and_ai_monitoring_enabled
+        # with_config doesn't work because the value for
+        # instrumentation.ruby_openai will be overriden during
+        # dependency detection.
+        NewRelic::Agent.stub(:config, {:'instrumentation.ruby_openai' => :prepend, :'ai_monitoring.enabled' => true}) do
+          assert_truthy NewRelic::Agent::LLM.openai?
+        end
+      end
+
+      def test_openai_true_when_ruby_openai_chain_and_ai_monitoring_enabled
+        NewRelic::Agent.stub(:config, {:'instrumentation.ruby_openai' => :chain, :'ai_monitoring.enabled' => true}) do
+          assert_truthy NewRelic::Agent::LLM.openai?
+        end
+      end
+
+      def test_openai_false_when_ruby_openai_auto
+        NewRelic::Agent.stub(:config, {:'instrumentation.ruby_openai' => :auto, :'ai_monitoring.enabled' => true}) do
+          refute_predicate(NewRelic::Agent::LLM, :openai?)
+        end
+      end
+
+      def test_openai_false_when_ruby_openai_disabled
+        NewRelic::Agent.stub(:config, {:'instrumentation.ruby_openai' => :disabled, :'ai_monitoring.enabled' => true}) do
+          refute_predicate(NewRelic::Agent::LLM, :openai?)
+        end
+      end
+
+      def test_openai_false_when_ruby_openai_unsatisfied
+        NewRelic::Agent.stub(:config, {:'instrumentation.ruby_openai' => :unsatisfied, :'ai_monitoring.enabled' => true}) do
+          refute_predicate(NewRelic::Agent::LLM, :openai?)
+        end
+      end
+
+      def test_openai_false_when_ai_monitoring_disabled
+        NewRelic::Agent.stub(:config, {:'instrumentation.ruby_openai' => :prepend, :'ai_monitoring.enabled' => false}) do
+          refute_predicate(NewRelic::Agent::LLM, :openai?)
+        end
+      end
+
+      def test_openai_value_memoized # could probs be a better test
+        NewRelic::Agent.stub(:config, {:'instrumentation.ruby_openai' => :prepend, :'ai_monitoring.enabled' => true}) do
+          refute NewRelic::Agent::LLM.instance_variable_defined?(:@openai)
+          assert_predicate(NewRelic::Agent::LLM, :openai?)
+          assert NewRelic::Agent::LLM.instance_variable_defined?(:@openai)
+        end
+      end
+
+      def test_openai_parent_with_all_the_things
+        NewRelic::Agent.stub(:config, {:'instrumentation.ruby_openai' => :prepend, :'ai_monitoring.enabled' => true}) do
+          parent_segment = NewRelic::Agent::Transaction::AbstractSegment.new('Llm/embedding/OpenAI/embeddings')
+          segment = NewRelic::Agent::Transaction::AbstractSegment.new('Anything')
+          segment.parent = parent_segment
+
+          assert_truthy NewRelic::Agent::LLM.openai_parent?(segment)
+        end
+      end
+
+      def test_openai_parent_without_segment
+        NewRelic::Agent.stub(:config, {:'instrumentation.ruby_openai' => :auto, :'ai_monitoring.enabled' => true}) do
+          refute NewRelic::Agent::LLM.openai_parent?(nil)
+        end
+      end
+
+      def test_opneai_parent_without_parent
+        NewRelic::Agent.stub(:config, {:'instrumentation.ruby_openai' => :auto, :'ai_monitoring.enabled' => true}) do
+          segment = NewRelic::Agent::Transaction::AbstractSegment.new('Anything')
+
+          refute NewRelic::Agent::LLM.openai_parent?(segment)
+        end
+      end
+
+      def test_opneai_parent_without_name
+        NewRelic::Agent.stub(:config, {:'instrumentation.ruby_openai' => :auto, :'ai_monitoring.enabled' => true}) do
+          parent_segment = NewRelic::Agent::Transaction::AbstractSegment.new(nil)
+          segment = NewRelic::Agent::Transaction::AbstractSegment.new('Anything')
+          segment.parent = parent_segment
+
+          refute NewRelic::Agent::LLM.openai_parent?(segment)
+        end
+      end
+
+      def test_opneai_parent_without_match
+        NewRelic::Agent.stub(:config, {:'instrumentation.ruby_openai' => :auto, :'ai_monitoring.enabled' => true}) do
+          parent_segment = NewRelic::Agent::Transaction::AbstractSegment.new('doesnt_match')
+          segment = NewRelic::Agent::Transaction::AbstractSegment.new('Anything')
+          segment.parent = parent_segment
+
+          refute NewRelic::Agent::LLM.openai_parent?(segment)
+        end
+      end
+
+      def test_populate_openai_response_headers_without_llm_event
+        NewRelic::Agent.stub(:config, {:'instrumentation.ruby_openai' => :auto, :'ai_monitoring.enabled' => true}) do
+          response = {}
+          parent = NewRelic::Agent::Transaction::AbstractSegment.new('Llm/embedding/OpenAI/embeddings')
+
+          refute parent.instance_variable_defined?(:@llm_event)
+          assert_nil NewRelic::Agent::LLM.populate_openai_response_headers(response, parent)
+        end
+      end
+
+      def test_populate_openai_response_headers_with_llm_event_calls_llm_method
+        NewRelic::Agent.stub(:config, {:'instrumentation.ruby_openai' => :auto, :'ai_monitoring.enabled' => true}) do
+          request_id = 'PIR8SH1P'
+          response = {'x-request-id' => [request_id]}
+          parent = NewRelic::Agent::Transaction::AbstractSegment.new('Llm/embedding/OpenAI/embeddings')
+          mock_llm_event = NewRelic::Agent::Llm::Embedding.new
+          parent.llm_event = mock_llm_event
+          NewRelic::Agent::LLM.populate_openai_response_headers(response, parent)
+
+          assert_equal request_id, mock_llm_event.request_id
+        end
+      end
+    end
+  end
+end

test/new_relic/net_http_test_cases.rb

@@ -137,4 +137,28 @@ def test_ipv6_host_get_request_records_metric
       'External/[::1]/Net::HTTP/GET' => {:call_count => 1}
     )
   end
+
+  def test_does_not_attempt_to_populate_response_headers_without_openai
+    populate_was_called = false
+    NewRelic::Agent::LLM.stub(:openai_parent?, false) do
+      NewRelic::Agent::LLM.stub(:populate_openai_response_headers, proc { |_h| populate_was_called = true }) do
+        in_transaction do
+          get_response
+        end
+      end
+    end
+
+    refute populate_was_called
+  end
+
+  def test_attempts_to_populate_response_headers_with_openai
+    populate_was_called = false
+    NewRelic::Agent::LLM.stub(:openai_parent?, true) do
+      NewRelic::Agent::LLM.stub(:populate_openai_response_headers, proc { |_h| populate_was_called = true }) do
+        in_transaction { get_response }
+      end
+    end
+
+    assert populate_was_called
+  end
 end