diff --git a/tests/agent_features/test_custom_events.py b/tests/agent_features/test_custom_events.py index 4fcb8c290b..1951a291f5 100644 --- a/tests/agent_features/test_custom_events.py +++ b/tests/agent_features/test_custom_events.py @@ -19,6 +19,8 @@ function_not_called, override_application_settings, reset_core_stats_engine, +) +from testing_support.validators.validate_custom_event import ( validate_custom_event_count, validate_custom_event_in_application_stats_engine, ) diff --git a/tests/agent_features/test_high_security_mode.py b/tests/agent_features/test_high_security_mode.py index 86e8ed2b5e..ae64be2ef9 100644 --- a/tests/agent_features/test_high_security_mode.py +++ b/tests/agent_features/test_high_security_mode.py @@ -22,9 +22,11 @@ override_generic_settings, reset_core_stats_engine, validate_attributes_complete, + validate_request_params_omitted, +) +from testing_support.validators.validate_custom_event import ( validate_custom_event_count, validate_custom_event_in_application_stats_engine, - validate_request_params_omitted, ) from testing_support.validators.validate_custom_parameters import ( validate_custom_parameters, diff --git a/tests/agent_features/test_record_llm_feedback_event.py b/tests/agent_features/test_record_llm_feedback_event.py index e61c7b530d..aee570fadb 100644 --- a/tests/agent_features/test_record_llm_feedback_event.py +++ b/tests/agent_features/test_record_llm_feedback_event.py @@ -12,10 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from testing_support.fixtures import ( - reset_core_stats_engine, - validate_custom_event_count, -) +from testing_support.fixtures import reset_core_stats_engine +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from newrelic.api.background_task import background_task diff --git a/tests/external_botocore/test_bedrock_chat_completion.py b/tests/external_botocore/test_bedrock_chat_completion.py index fc25e15adb..fe183c27dc 100644 --- a/tests/external_botocore/test_bedrock_chat_completion.py +++ b/tests/external_botocore/test_bedrock_chat_completion.py @@ -37,7 +37,6 @@ override_llm_token_callback_settings, reset_core_stats_engine, validate_attributes, - validate_custom_event_count, ) from testing_support.ml_testing_utils import ( # noqa: F401 add_token_count_to_events, @@ -49,6 +48,7 @@ llm_token_count_callback, set_trace_info, ) +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_error_trace_attributes import ( validate_error_trace_attributes, diff --git a/tests/external_botocore/test_bedrock_chat_completion_via_langchain.py b/tests/external_botocore/test_bedrock_chat_completion_via_langchain.py index e144868e37..00be00e17a 100644 --- a/tests/external_botocore/test_bedrock_chat_completion_via_langchain.py +++ b/tests/external_botocore/test_bedrock_chat_completion_via_langchain.py @@ -18,12 +18,9 @@ chat_completion_langchain_expected_streaming_events, ) from conftest import BOTOCORE_VERSION # pylint: disable=E0611 -from testing_support.fixtures import ( - reset_core_stats_engine, - validate_attributes, - validate_custom_event_count, -) +from testing_support.fixtures import reset_core_stats_engine, validate_attributes from testing_support.ml_testing_utils import set_trace_info # noqa: F401 +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_transaction_metrics import ( validate_transaction_metrics, diff --git a/tests/external_botocore/test_bedrock_embeddings.py b/tests/external_botocore/test_bedrock_embeddings.py index 946858e8dd..0c9e45244a 100644 --- a/tests/external_botocore/test_bedrock_embeddings.py +++ b/tests/external_botocore/test_bedrock_embeddings.py @@ -30,7 +30,6 @@ override_llm_token_callback_settings, reset_core_stats_engine, validate_attributes, - validate_custom_event_count, ) from testing_support.ml_testing_utils import ( # noqa: F401 add_token_count_to_events, @@ -41,6 +40,7 @@ llm_token_count_callback, set_trace_info, ) +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_error_trace_attributes import ( validate_error_trace_attributes, diff --git a/tests/mlmodel_langchain/test_chain.py b/tests/mlmodel_langchain/test_chain.py index f6f877d63c..69c54b9509 100644 --- a/tests/mlmodel_langchain/test_chain.py +++ b/tests/mlmodel_langchain/test_chain.py @@ -25,17 +25,14 @@ from langchain.prompts import ChatPromptTemplate from langchain.schema import BaseOutputParser from mock import patch -from testing_support.fixtures import ( - reset_core_stats_engine, - validate_attributes, - validate_custom_event_count, -) +from testing_support.fixtures import reset_core_stats_engine, validate_attributes from testing_support.ml_testing_utils import ( # noqa: F401 disabled_ai_monitoring_record_content_settings, disabled_ai_monitoring_settings, events_sans_content, set_trace_info, ) +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_error_trace_attributes import ( validate_error_trace_attributes, diff --git a/tests/mlmodel_langchain/test_tool.py b/tests/mlmodel_langchain/test_tool.py index 84ba7cf369..dac87c421e 100644 --- a/tests/mlmodel_langchain/test_tool.py +++ b/tests/mlmodel_langchain/test_tool.py @@ -21,16 +21,13 @@ import pytest from langchain.tools import tool from mock import patch -from testing_support.fixtures import ( - reset_core_stats_engine, - validate_attributes, - validate_custom_event_count, -) +from testing_support.fixtures import reset_core_stats_engine, validate_attributes from testing_support.ml_testing_utils import ( # noqa: F401 disabled_ai_monitoring_record_content_settings, disabled_ai_monitoring_settings, set_trace_info, ) +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_error_trace_attributes import ( validate_error_trace_attributes, diff --git a/tests/mlmodel_langchain/test_vectorstore.py b/tests/mlmodel_langchain/test_vectorstore.py index e090b1892f..f587e93b52 100644 --- a/tests/mlmodel_langchain/test_vectorstore.py +++ b/tests/mlmodel_langchain/test_vectorstore.py @@ -19,16 +19,13 @@ import pytest from langchain_community.document_loaders import PyPDFLoader from langchain_community.vectorstores.faiss import FAISS -from testing_support.fixtures import ( - reset_core_stats_engine, - validate_attributes, - validate_custom_event_count, -) +from testing_support.fixtures import reset_core_stats_engine, validate_attributes from testing_support.ml_testing_utils import ( # noqa: F401 disabled_ai_monitoring_record_content_settings, disabled_ai_monitoring_settings, set_trace_info, ) +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_error_trace_attributes import ( validate_error_trace_attributes, diff --git a/tests/mlmodel_openai/test_chat_completion.py b/tests/mlmodel_openai/test_chat_completion.py index 5d7cf11d81..229ec44272 100644 --- a/tests/mlmodel_openai/test_chat_completion.py +++ b/tests/mlmodel_openai/test_chat_completion.py @@ -17,7 +17,6 @@ override_llm_token_callback_settings, reset_core_stats_engine, validate_attributes, - validate_custom_event_count, ) from testing_support.ml_testing_utils import ( # noqa: F401 add_token_count_to_events, @@ -29,6 +28,7 @@ llm_token_count_callback, set_trace_info, ) +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_transaction_metrics import ( validate_transaction_metrics, diff --git a/tests/mlmodel_openai/test_chat_completion_error.py b/tests/mlmodel_openai/test_chat_completion_error.py index e1fe993292..d3ed79bea9 100644 --- a/tests/mlmodel_openai/test_chat_completion_error.py +++ b/tests/mlmodel_openai/test_chat_completion_error.py @@ -19,7 +19,6 @@ dt_enabled, override_llm_token_callback_settings, reset_core_stats_engine, - validate_custom_event_count, ) from testing_support.ml_testing_utils import ( # noqa: F401 add_token_count_to_events, @@ -28,6 +27,7 @@ llm_token_count_callback, set_trace_info, ) +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_error_trace_attributes import ( validate_error_trace_attributes, diff --git a/tests/mlmodel_openai/test_chat_completion_error_v1.py b/tests/mlmodel_openai/test_chat_completion_error_v1.py index 735664cac2..dc1cdcfb78 100644 --- a/tests/mlmodel_openai/test_chat_completion_error_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_error_v1.py @@ -18,7 +18,6 @@ dt_enabled, override_llm_token_callback_settings, reset_core_stats_engine, - validate_custom_event_count, ) from testing_support.ml_testing_utils import ( # noqa: F401 add_token_count_to_events, @@ -27,6 +26,7 @@ llm_token_count_callback, set_trace_info, ) +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_error_trace_attributes import ( validate_error_trace_attributes, diff --git a/tests/mlmodel_openai/test_chat_completion_stream.py b/tests/mlmodel_openai/test_chat_completion_stream.py index 18cee59ab8..3c32dd9f05 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream.py +++ b/tests/mlmodel_openai/test_chat_completion_stream.py @@ -17,7 +17,6 @@ override_llm_token_callback_settings, reset_core_stats_engine, validate_attributes, - validate_custom_event_count, ) from testing_support.ml_testing_utils import ( # noqa: F401 add_token_count_to_events, @@ -29,6 +28,7 @@ llm_token_count_callback, set_trace_info, ) +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_transaction_metrics import ( validate_transaction_metrics, diff --git a/tests/mlmodel_openai/test_chat_completion_stream_error.py b/tests/mlmodel_openai/test_chat_completion_stream_error.py index 01270491c5..2dc0400b7e 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_error.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_error.py @@ -19,7 +19,6 @@ dt_enabled, override_llm_token_callback_settings, reset_core_stats_engine, - validate_custom_event_count, ) from testing_support.ml_testing_utils import ( # noqa: F401 add_token_count_to_events, @@ -28,6 +27,7 @@ llm_token_count_callback, set_trace_info, ) +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_error_trace_attributes import ( validate_error_trace_attributes, diff --git a/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py b/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py index e9bba66b48..33bda04f78 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_error_v1.py @@ -19,7 +19,6 @@ dt_enabled, override_llm_token_callback_settings, reset_core_stats_engine, - validate_custom_event_count, ) from testing_support.ml_testing_utils import ( # noqa: F401 add_token_count_to_events, @@ -28,6 +27,7 @@ llm_token_count_callback, set_trace_info, ) +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_error_trace_attributes import ( validate_error_trace_attributes, diff --git a/tests/mlmodel_openai/test_chat_completion_stream_v1.py b/tests/mlmodel_openai/test_chat_completion_stream_v1.py index 01b2a09e66..4ee6da1809 100644 --- a/tests/mlmodel_openai/test_chat_completion_stream_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_stream_v1.py @@ -17,7 +17,6 @@ override_llm_token_callback_settings, reset_core_stats_engine, validate_attributes, - validate_custom_event_count, ) from testing_support.ml_testing_utils import ( # noqa: F401 add_token_count_to_events, @@ -29,6 +28,7 @@ llm_token_count_callback, set_trace_info, ) +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_transaction_metrics import ( validate_transaction_metrics, diff --git a/tests/mlmodel_openai/test_chat_completion_v1.py b/tests/mlmodel_openai/test_chat_completion_v1.py index 1ad95c274a..5b7c294b08 100644 --- a/tests/mlmodel_openai/test_chat_completion_v1.py +++ b/tests/mlmodel_openai/test_chat_completion_v1.py @@ -17,7 +17,6 @@ override_llm_token_callback_settings, reset_core_stats_engine, validate_attributes, - validate_custom_event_count, ) from testing_support.ml_testing_utils import ( # noqa: F401 add_token_count_to_events, @@ -29,6 +28,7 @@ llm_token_count_callback, set_trace_info, ) +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_transaction_metrics import ( validate_transaction_metrics, diff --git a/tests/mlmodel_openai/test_embeddings.py b/tests/mlmodel_openai/test_embeddings.py index 5532d475de..488f60d2e1 100644 --- a/tests/mlmodel_openai/test_embeddings.py +++ b/tests/mlmodel_openai/test_embeddings.py @@ -17,7 +17,6 @@ override_llm_token_callback_settings, reset_core_stats_engine, validate_attributes, - validate_custom_event_count, ) from testing_support.ml_testing_utils import ( # noqa: F401 add_token_count_to_events, @@ -27,6 +26,7 @@ llm_token_count_callback, set_trace_info, ) +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_transaction_metrics import ( validate_transaction_metrics, diff --git a/tests/mlmodel_openai/test_embeddings_error.py b/tests/mlmodel_openai/test_embeddings_error.py index 6dd465bc8a..020648d0c2 100644 --- a/tests/mlmodel_openai/test_embeddings_error.py +++ b/tests/mlmodel_openai/test_embeddings_error.py @@ -18,7 +18,6 @@ dt_enabled, override_llm_token_callback_settings, reset_core_stats_engine, - validate_custom_event_count, ) from testing_support.ml_testing_utils import ( # noqa: F401 add_token_count_to_events, @@ -27,6 +26,7 @@ llm_token_count_callback, set_trace_info, ) +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_error_trace_attributes import ( validate_error_trace_attributes, diff --git a/tests/mlmodel_openai/test_embeddings_error_v1.py b/tests/mlmodel_openai/test_embeddings_error_v1.py index 06a38cf460..2464ebfe20 100644 --- a/tests/mlmodel_openai/test_embeddings_error_v1.py +++ b/tests/mlmodel_openai/test_embeddings_error_v1.py @@ -20,7 +20,6 @@ dt_enabled, override_llm_token_callback_settings, reset_core_stats_engine, - validate_custom_event_count, ) from testing_support.ml_testing_utils import ( # noqa: F401 add_token_count_to_events, @@ -29,6 +28,7 @@ llm_token_count_callback, set_trace_info, ) +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_error_trace_attributes import ( validate_error_trace_attributes, diff --git a/tests/mlmodel_openai/test_embeddings_stream_v1.py b/tests/mlmodel_openai/test_embeddings_stream_v1.py index 3ea8be4aa3..0517801b39 100644 --- a/tests/mlmodel_openai/test_embeddings_stream_v1.py +++ b/tests/mlmodel_openai/test_embeddings_stream_v1.py @@ -14,11 +14,9 @@ import pytest from conftest import get_openai_version # pylint: disable=E0611 -from testing_support.fixtures import ( - reset_core_stats_engine, - validate_custom_event_count, -) +from testing_support.fixtures import reset_core_stats_engine from testing_support.ml_testing_utils import set_trace_info # noqa: F401 +from testing_support.validators.validate_custom_event import validate_custom_event_count from newrelic.api.background_task import background_task diff --git a/tests/mlmodel_openai/test_embeddings_v1.py b/tests/mlmodel_openai/test_embeddings_v1.py index 38777e2c56..7bafa53181 100644 --- a/tests/mlmodel_openai/test_embeddings_v1.py +++ b/tests/mlmodel_openai/test_embeddings_v1.py @@ -17,7 +17,6 @@ override_llm_token_callback_settings, reset_core_stats_engine, validate_attributes, - validate_custom_event_count, ) from testing_support.ml_testing_utils import ( # noqa: F401 add_token_count_to_events, @@ -27,6 +26,7 @@ llm_token_count_callback, set_trace_info, ) +from testing_support.validators.validate_custom_event import validate_custom_event_count from testing_support.validators.validate_custom_events import validate_custom_events from testing_support.validators.validate_transaction_metrics import ( validate_transaction_metrics, diff --git a/tests/testing_support/fixtures.py b/tests/testing_support/fixtures.py index 26ad53829b..b5565a2706 100644 --- a/tests/testing_support/fixtures.py +++ b/tests/testing_support/fixtures.py @@ -33,13 +33,13 @@ user_attributes_added, ) -from newrelic.api.ml_model import set_llm_token_count_callback from newrelic.admin.record_deploy import record_deploy from newrelic.api.application import ( application_instance, application_settings, register_application, ) +from newrelic.api.ml_model import set_llm_token_count_callback from newrelic.common.agent_http import DeveloperModeClient from newrelic.common.encoding_utils import json_encode, obfuscate from newrelic.common.object_names import callable_name @@ -895,56 +895,6 @@ def _validate_application_exception_message(wrapped, instance, args, kwargs): return _validate_application_exception_message -def _validate_custom_event(recorded_event, required_event): - assert len(recorded_event) == 2 # [intrinsic, user attributes] - - intrinsics = recorded_event[0] - - assert intrinsics["type"] == required_event[0]["type"] - - now = time.time() - assert isinstance(intrinsics["timestamp"], int) - assert intrinsics["timestamp"] <= 1000.0 * now - assert intrinsics["timestamp"] >= 1000.0 * required_event[0]["timestamp"] - - assert recorded_event[1].items() == required_event[1].items() - - -def validate_custom_event_in_application_stats_engine(required_event): - @function_wrapper - def _validate_custom_event_in_application_stats_engine(wrapped, instance, args, kwargs): - try: - result = wrapped(*args, **kwargs) - except: - raise - else: - stats = core_application_stats_engine(None) - assert stats.custom_events.num_samples == 1 - - custom_event = next(iter(stats.custom_events)) - _validate_custom_event(custom_event, required_event) - - return result - - return _validate_custom_event_in_application_stats_engine - - -def validate_custom_event_count(count): - @function_wrapper - def _validate_custom_event_count(wrapped, instance, args, kwargs): - try: - result = wrapped(*args, **kwargs) - except: - raise - else: - stats = core_application_stats_engine(None) - assert stats.custom_events.num_samples == count, "Expected %d, got %d" % (count, stats.custom_events.num_samples) - - return result - - return _validate_custom_event_count - - def _validate_node_parenting(node, expected_node): assert node.exclusive >= 0, "node.exclusive = %s" % node.exclusive diff --git a/tests/testing_support/validators/validate_custom_event.py b/tests/testing_support/validators/validate_custom_event.py new file mode 100644 index 0000000000..7456cb156e --- /dev/null +++ b/tests/testing_support/validators/validate_custom_event.py @@ -0,0 +1,69 @@ +# Copyright 2010 New Relic, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import time + +from testing_support.fixtures import core_application_stats_engine + +from newrelic.common.object_wrapper import function_wrapper + + +def _validate_custom_event(recorded_event, required_event): + assert len(recorded_event) == 2 # [intrinsic, user attributes] + + intrinsics = recorded_event[0] + + assert intrinsics["type"] == required_event[0]["type"] + + now = time.time() + assert isinstance(intrinsics["timestamp"], int) + assert intrinsics["timestamp"] <= 1000.0 * now + assert intrinsics["timestamp"] >= 1000.0 * required_event[0]["timestamp"] + + assert recorded_event[1].items() == required_event[1].items() + + +def validate_custom_event_in_application_stats_engine(required_event): + @function_wrapper + def _validate_custom_event_in_application_stats_engine(wrapped, instance, args, kwargs): + try: + result = wrapped(*args, **kwargs) + except: + raise + else: + stats = core_application_stats_engine(None) + assert stats.custom_events.num_samples == 1 + + custom_event = next(iter(stats.custom_events)) + _validate_custom_event(custom_event, required_event) + + return result + + return _validate_custom_event_in_application_stats_engine + + +def validate_custom_event_count(count): + @function_wrapper + def _validate_custom_event_count(wrapped, instance, args, kwargs): + try: + result = wrapped(*args, **kwargs) + except: + raise + else: + stats = core_application_stats_engine(None) + assert stats.custom_events.num_samples == count + + return result + + return _validate_custom_event_count