diff --git a/newrelic/core/agent_protocol.py b/newrelic/core/agent_protocol.py index 9ac42e4050..b8ccff2b90 100644 --- a/newrelic/core/agent_protocol.py +++ b/newrelic/core/agent_protocol.py @@ -287,6 +287,7 @@ def _connect_payload(app_name, linked_applications, environment, settings): connect_settings = {} connect_settings["browser_monitoring.loader"] = settings["browser_monitoring.loader"] connect_settings["browser_monitoring.debug"] = settings["browser_monitoring.debug"] + connect_settings["ai_monitoring.enabled"] = settings["ai_monitoring.enabled"] security_settings = {} security_settings["capture_params"] = settings["capture_params"] diff --git a/newrelic/core/config.py b/newrelic/core/config.py index 5f6b772a36..9139b39344 100644 --- a/newrelic/core/config.py +++ b/newrelic/core/config.py @@ -1225,6 +1225,12 @@ def apply_server_side_settings(server_side_config=None, settings=_settings): settings_snapshot, "event_harvest_config.harvest_limits.span_event_data", span_event_harvest_limit ) + # Check to see if collect_ai appears in the connect response to handle account-level AIM toggling + collect_ai = server_side_config.get("collect_ai", None) + if collect_ai is not None: + apply_config_setting(settings_snapshot, "ai_monitoring.enabled", collect_ai) + _logger.debug("Setting ai_monitoring.enabled to value of collect_ai=%s", collect_ai) + # Since the server does not override this setting as it's an OTLP setting, # we must override it here manually by converting it into a per harvest cycle # value. diff --git a/tests/agent_unittests/test_agent_protocol.py b/tests/agent_unittests/test_agent_protocol.py index 7b865cb9d5..fde10ad294 100644 --- a/tests/agent_unittests/test_agent_protocol.py +++ b/tests/agent_unittests/test_agent_protocol.py @@ -170,7 +170,7 @@ def test_send(status_code): HttpClientRecorder.STATUS_CODE = status_code settings = finalize_application_settings( { - "request_headers_map": {"custom-header": u"value"}, # pylint: disable=W1406 + "request_headers_map": {"custom-header": "value"}, # pylint: disable=W1406 "agent_run_id": "RUN_TOKEN", } ) @@ -297,7 +297,13 @@ def connect_payload_asserts( with_kubernetes=True, ): payload_data = payload[0] + + # Turn off black formatting for this section of the code. While Python 2 has been + # EOL'd since 2020, New Relic still supports it and therefore this unicode assert + # needs the u"" still. + # fmt: off assert isinstance(payload_data["agent_version"], type(u"")) # pylint: disable=W1406 + # fmt: on assert payload_data["app_name"] == PAYLOAD_APP_NAME assert payload_data["display_host"] == DISPLAY_NAME assert payload_data["environment"] == ENVIRONMENT @@ -311,9 +317,10 @@ def connect_payload_asserts( assert len(payload_data["security_settings"]) == 2 assert payload_data["security_settings"]["capture_params"] == CAPTURE_PARAMS assert payload_data["security_settings"]["transaction_tracer"] == {"record_sql": RECORD_SQL} - assert len(payload_data["settings"]) == 2 + assert len(payload_data["settings"]) == 3 assert payload_data["settings"]["browser_monitoring.loader"] == (BROWSER_MONITORING_LOADER) assert payload_data["settings"]["browser_monitoring.debug"] == (BROWSER_MONITORING_DEBUG) + assert payload_data["settings"]["ai_monitoring.enabled"] is False utilization_len = 5 diff --git a/tests/agent_unittests/test_connect_response_fields.py b/tests/agent_unittests/test_connect_response_fields.py index 617cc7ce23..452be53f00 100644 --- a/tests/agent_unittests/test_connect_response_fields.py +++ b/tests/agent_unittests/test_connect_response_fields.py @@ -12,41 +12,41 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pytest import functools import time +import pytest from testing_support.fixtures import override_generic_settings -from newrelic.core.config import global_settings -from newrelic.core.agent_protocol import AgentProtocol + from newrelic.common.agent_http import DeveloperModeClient from newrelic.common.encoding_utils import json_encode - +from newrelic.core.agent_protocol import AgentProtocol +from newrelic.core.config import global_settings DEFAULT = object() LINKED_APPLICATIONS = [] ENVIRONMENT = [] NOW = time.time() EMPTY_SAMPLES = { - 'reservoir_size': 100, - 'events_seen': 0, + "reservoir_size": 100, + "events_seen": 0, } _all_endpoints = ( - ('send_metric_data', (NOW, NOW + 1, ())), - ('send_transaction_events', (EMPTY_SAMPLES, ())), - ('send_custom_events', (EMPTY_SAMPLES, ())), - ('send_error_events', (EMPTY_SAMPLES, ())), - ('send_transaction_traces', ([[]],)), - ('send_sql_traces', ([[]],)), - ('get_agent_commands', ()), - ('send_profile_data', ([[]],)), - ('send_errors', ([[]],)), - ('send_agent_command_results', ({0: {}},)), - ('agent_settings', ({},)), - ('send_span_events', (EMPTY_SAMPLES, ())), - ('shutdown_session', ()), + ("send_metric_data", (NOW, NOW + 1, ())), + ("send_transaction_events", (EMPTY_SAMPLES, ())), + ("send_custom_events", (EMPTY_SAMPLES, ())), + ("send_error_events", (EMPTY_SAMPLES, ())), + ("send_transaction_traces", ([[]],)), + ("send_sql_traces", ([[]],)), + ("get_agent_commands", ()), + ("send_profile_data", ([[]],)), + ("send_errors", ([[]],)), + ("send_agent_command_results", ({0: {}},)), + ("agent_settings", ({},)), + ("send_span_events", (EMPTY_SAMPLES, ())), + ("shutdown_session", ()), ) @@ -85,22 +85,17 @@ def send_request( ) -@pytest.mark.parametrize('headers_map_present', (True, False)) +@pytest.mark.parametrize("headers_map_present", (True, False)) def test_no_blob_behavior(headers_map_present): if headers_map_present: - connect_response_fields = {u"request_headers_map": None} - client_cls = functools.partial( - CustomTestClient, connect_response_fields=connect_response_fields) + connect_response_fields = {"request_headers_map": None} + client_cls = functools.partial(CustomTestClient, connect_response_fields=connect_response_fields) else: - client_cls = functools.partial( - CustomTestClient, connect_response_fields=DEFAULT) + client_cls = functools.partial(CustomTestClient, connect_response_fields=DEFAULT) protocol = AgentProtocol.connect( - 'app_name', - LINKED_APPLICATIONS, - ENVIRONMENT, - global_settings(), - client_cls=client_cls) + "app_name", LINKED_APPLICATIONS, ENVIRONMENT, global_settings(), client_cls=client_cls + ) protocol.send("shutdown") @@ -111,19 +106,14 @@ def test_no_blob_behavior(headers_map_present): def test_blob(): - request_headers_map = {u'X-Foo': u'Bar'} - connect_response_fields = {u"request_headers_map": request_headers_map} + request_headers_map = {"X-Foo": "Bar"} + connect_response_fields = {"request_headers_map": request_headers_map} - client_cls = functools.partial( - CustomTestClient, - connect_response_fields=connect_response_fields) + client_cls = functools.partial(CustomTestClient, connect_response_fields=connect_response_fields) protocol = AgentProtocol.connect( - 'app_name', - LINKED_APPLICATIONS, - ENVIRONMENT, - global_settings(), - client_cls=client_cls) + "app_name", LINKED_APPLICATIONS, ENVIRONMENT, global_settings(), client_cls=client_cls + ) protocol.send("shutdown") @@ -134,52 +124,71 @@ def test_blob(): } -@override_generic_settings(global_settings(), { - 'developer_mode': True, -}) +@override_generic_settings( + global_settings(), + { + "developer_mode": True, + }, +) def test_server_side_config_precedence(): connect_response_fields = { - u'agent_config': {u'span_events.enabled': True}, - u'span_events.enabled': False, + "agent_config": {"span_events.enabled": True}, + "span_events.enabled": False, } - client_cls = functools.partial( - CustomTestClient, - connect_response_fields=connect_response_fields) + client_cls = functools.partial(CustomTestClient, connect_response_fields=connect_response_fields) protocol = AgentProtocol.connect( - 'app_name', - LINKED_APPLICATIONS, - ENVIRONMENT, - global_settings(), - client_cls=client_cls) + "app_name", LINKED_APPLICATIONS, ENVIRONMENT, global_settings(), client_cls=client_cls + ) assert protocol.configuration.span_events.enabled is False -@override_generic_settings(global_settings(), { - 'developer_mode': True, -}) -@pytest.mark.parametrize("connect_response_fields", -( - {}, - {"span_event_harvest_config": {"report_period_ms": 60000, "harvest_limit": 123}}, - {"span_event_harvest_config": {}}) +@override_generic_settings( + global_settings(), + { + "developer_mode": True, + }, +) +@pytest.mark.parametrize( + "connect_response_fields", + ( + {}, + {"span_event_harvest_config": {"report_period_ms": 60000, "harvest_limit": 123}}, + {"span_event_harvest_config": {}}, + ), ) def test_span_event_harvest_config(connect_response_fields): - client_cls = functools.partial( - CustomTestClient, - connect_response_fields=connect_response_fields) + client_cls = functools.partial(CustomTestClient, connect_response_fields=connect_response_fields) protocol = AgentProtocol.connect( - 'app_name', - LINKED_APPLICATIONS, - ENVIRONMENT, - global_settings(), - client_cls=client_cls) + "app_name", LINKED_APPLICATIONS, ENVIRONMENT, global_settings(), client_cls=client_cls + ) if connect_response_fields and connect_response_fields["span_event_harvest_config"]: expected = 123 else: from newrelic.core.config import SPAN_EVENT_RESERVOIR_SIZE + expected = SPAN_EVENT_RESERVOIR_SIZE assert protocol.configuration.event_harvest_config.harvest_limits.span_event_data == expected + + +@override_generic_settings( + global_settings(), + { + "developer_mode": True, + }, +) +@pytest.mark.parametrize("connect_response_fields", ({}, {"collect_ai": True}, {"collect_ai": False})) +def test_account_level_aim(connect_response_fields): + client_cls = functools.partial(CustomTestClient, connect_response_fields=connect_response_fields) + + protocol = AgentProtocol.connect( + "app_name", LINKED_APPLICATIONS, ENVIRONMENT, global_settings(), client_cls=client_cls + ) + + if connect_response_fields and connect_response_fields["collect_ai"]: + assert protocol.configuration.ai_monitoring.enabled == connect_response_fields["collect_ai"] + else: + assert protocol.configuration.ai_monitoring.enabled is False