From 2aad9d5d4873471290697fc32a42a76c9623561e Mon Sep 17 00:00:00 2001 From: Kayla Reopelle Date: Wed, 27 Mar 2024 17:49:31 -0700 Subject: [PATCH 1/5] Move LLM class method tests into llm_test file These tests do not require Net::HTTP instrumentation. Now, the location of the tests more closely align with where they are defined. --- .../suites/net_http/net_http_openai_test.rb | 128 ----------------- test/new_relic/agent/llm_test.rb | 132 ++++++++++++++++++ 2 files changed, 132 insertions(+), 128 deletions(-) delete mode 100644 test/multiverse/suites/net_http/net_http_openai_test.rb create mode 100644 test/new_relic/agent/llm_test.rb diff --git a/test/multiverse/suites/net_http/net_http_openai_test.rb b/test/multiverse/suites/net_http/net_http_openai_test.rb deleted file mode 100644 index 249b138de8..0000000000 --- a/test/multiverse/suites/net_http/net_http_openai_test.rb +++ /dev/null @@ -1,128 +0,0 @@ -# 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 - -class NetHttpOpenAITest < Minitest::Test - include NewRelic::Agent::Instrumentation::NetHTTP - - def setup - NewRelic::Agent::LLM.remove_instance_variable(:@openai) if NewRelic::Agent::LLM.instance_variable_defined?(:@openai) - end - - # Mocha used in test_populate_openai_response_headers_with_llm_event_calls_llm_method - def teardown - mocha_teardown - 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 - response = {} - parent = NewRelic::Agent::Transaction::AbstractSegment.new('Llm/embedding/OpenAI/embeddings') - mock_llm_event = NewRelic::Agent::Llm::Embedding.new - mock_llm_event.expects(:populate_openai_response_headers).with(response) - parent.llm_event = mock_llm_event - NewRelic::Agent::LLM.populate_openai_response_headers(response, parent) - end - end -end diff --git a/test/new_relic/agent/llm_test.rb b/test/new_relic/agent/llm_test.rb new file mode 100644 index 0000000000..c8b7fbb82c --- /dev/null +++ b/test/new_relic/agent/llm_test.rb @@ -0,0 +1,132 @@ +# 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 + + # Mocha used in test_populate_openai_response_headers_with_llm_event_calls_llm_method + def teardown + mocha_teardown + 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 + response = {} + parent = NewRelic::Agent::Transaction::AbstractSegment.new('Llm/embedding/OpenAI/embeddings') + mock_llm_event = NewRelic::Agent::Llm::Embedding.new + mock_llm_event.expects(:populate_openai_response_headers).with(response) + parent.llm_event = mock_llm_event + NewRelic::Agent::LLM.populate_openai_response_headers(response, parent) + end + end + end + end +end From 582eab130999d3413987ce869af8b0cf7f12b5e9 Mon Sep 17 00:00:00 2001 From: Kayla Reopelle Date: Wed, 27 Mar 2024 17:53:41 -0700 Subject: [PATCH 2/5] Add tests for OpenAI behavior within Net::HTTP --- test/new_relic/net_http_test_cases.rb | 35 +++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/test/new_relic/net_http_test_cases.rb b/test/new_relic/net_http_test_cases.rb index 0f98bb5371..3cd27e8bd0 100644 --- a/test/new_relic/net_http_test_cases.rb +++ b/test/new_relic/net_http_test_cases.rb @@ -137,4 +137,39 @@ def test_ipv6_host_get_request_records_metric 'External/[::1]/Net::HTTP/GET' => {:call_count => 1} ) end + + class OpenAITestError < StandardError; end + + def test_does_not_attempt_to_populate_response_headers_without_openai + segment = Minitest::Mock.new + + NewRelic::Agent::LLM.stub(:openai_parent?, false, segment) do + # raise if populate_openai_response_headers is called + NewRelic::Agent::LLM.stub(:populate_openai_response_headers, -> { raise OpenAITestError.new }) do + in_transaction do + get_response + end + end + end + end + + def test_attempts_to_populate_response_headers_with_openai + segment = Minitest::Mock.new + wrapped_response = Minitest::Mock.new + mock_proc = proc { |*_args| raise OpenAITestError } + result = nil + expected_result = 'expected_result' + + NewRelic::Agent::LLM.stub(:openai_parent?, true, segment) do + result = NewRelic::Agent::LLM.stub(:populate_openai_response_headers, mock_proc) do + in_transaction do + get_response + rescue OpenAITestError + expected_result + end + end + end + + assert_equal expected_result, result + end end From ec4269bc62c36136da398f2e236b0aa2d256facc Mon Sep 17 00:00:00 2001 From: Kayla Reopelle Date: Thu, 28 Mar 2024 14:56:12 -0700 Subject: [PATCH 3/5] Add begin to block with rescue --- test/new_relic/net_http_test_cases.rb | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/test/new_relic/net_http_test_cases.rb b/test/new_relic/net_http_test_cases.rb index 3cd27e8bd0..f45f7b699b 100644 --- a/test/new_relic/net_http_test_cases.rb +++ b/test/new_relic/net_http_test_cases.rb @@ -162,10 +162,12 @@ def test_attempts_to_populate_response_headers_with_openai NewRelic::Agent::LLM.stub(:openai_parent?, true, segment) do result = NewRelic::Agent::LLM.stub(:populate_openai_response_headers, mock_proc) do - in_transaction do - get_response + begin + in_transaction do + get_response + end rescue OpenAITestError - expected_result + return expected_result end end end From 77e852227aa9836eb1c411edc96b7d2c25eb3437 Mon Sep 17 00:00:00 2001 From: fallwith Date: Thu, 28 Mar 2024 16:48:55 -0700 Subject: [PATCH 4/5] simplify tests for AI header setting / not setting use a variable in combination with a proc to confirm the stubbed method was indeed invoked as expected or not invoked as expected. this gives us an assertion in the non-invoked path, which is nice. i believe this fixes (switching from `-> {}` to `proc { |_h| }` an issue to where if the `raise` actually was attempted in the non-invoked path, it would fail. removed unused variables / mocks --- test/new_relic/net_http_test_cases.rb | 33 ++++++++------------------- 1 file changed, 10 insertions(+), 23 deletions(-) diff --git a/test/new_relic/net_http_test_cases.rb b/test/new_relic/net_http_test_cases.rb index f45f7b699b..8fd9f18874 100644 --- a/test/new_relic/net_http_test_cases.rb +++ b/test/new_relic/net_http_test_cases.rb @@ -138,40 +138,27 @@ def test_ipv6_host_get_request_records_metric ) end - class OpenAITestError < StandardError; end - def test_does_not_attempt_to_populate_response_headers_without_openai - segment = Minitest::Mock.new - - NewRelic::Agent::LLM.stub(:openai_parent?, false, segment) do - # raise if populate_openai_response_headers is called - NewRelic::Agent::LLM.stub(:populate_openai_response_headers, -> { raise OpenAITestError.new }) do + 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 - segment = Minitest::Mock.new - wrapped_response = Minitest::Mock.new - mock_proc = proc { |*_args| raise OpenAITestError } - result = nil - expected_result = 'expected_result' - - NewRelic::Agent::LLM.stub(:openai_parent?, true, segment) do - result = NewRelic::Agent::LLM.stub(:populate_openai_response_headers, mock_proc) do - begin - in_transaction do - get_response - end - rescue OpenAITestError - return expected_result - end + 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_equal expected_result, result + assert populate_was_called end end From 7bf893aa65ef17d151db76d8e9aaf75957b34c78 Mon Sep 17 00:00:00 2001 From: Kayla Reopelle Date: Fri, 29 Mar 2024 12:31:51 -0700 Subject: [PATCH 5/5] Refactor test to remove Mocha calls This test will also look to see that a request ID was correctly applied to the LLM event. --- test/new_relic/agent/llm_test.rb | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/test/new_relic/agent/llm_test.rb b/test/new_relic/agent/llm_test.rb index c8b7fbb82c..5b22a396bf 100644 --- a/test/new_relic/agent/llm_test.rb +++ b/test/new_relic/agent/llm_test.rb @@ -11,11 +11,6 @@ def setup NewRelic::Agent::LLM.remove_instance_variable(:@openai) if NewRelic::Agent::LLM.instance_variable_defined?(:@openai) end - # Mocha used in test_populate_openai_response_headers_with_llm_event_calls_llm_method - def teardown - mocha_teardown - 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 @@ -119,12 +114,14 @@ def test_populate_openai_response_headers_without_llm_event 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 - response = {} + 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 - mock_llm_event.expects(:populate_openai_response_headers).with(response) 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