Add support for account level AIM disable. (#1175)

* Add support for account level AIM disable.

* Fix lint errors

tests/agent_unittests/test_agent_protocol.py:173:64: E261 at least two spaces before inline comment
tests/agent_unittests/test_agent_protocol.py:320:5: E303 too many blank lines (2)
tests/agent_unittests/test_connect_response_fields.py:164:8: W291 trailing whitespace
tests/agent_unittests/test_connect_response_fields.py:165:86: W291 trailing whitespace

* [Mega-Linter] Apply linters fixes

* Trigger tests

* Add debug log statement

* Fixup: unicode dropping

---------

Co-authored-by: Hannah Stepanek <[email protected]>
Co-authored-by: hmstepanek <[email protected]>

Author: umaannamalai

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"]

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.

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
 

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