Skip to content

Commit

Permalink
Add support for account level AIM disable. (#1175)
Browse files Browse the repository at this point in the history 
* 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]>
  • Loading branch information
@umaannamalai @hmstepanek
3 people authored Jul 11, 2024
1 parent 9b6b2e5 commit 5469c4a
Show file tree
Hide file tree
Showing 4 changed files with 94 additions and 71 deletions.
1 change: 1 addition & 0 deletions newrelic/core/agent_protocol.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down
6 changes: 6 additions & 0 deletions newrelic/core/config.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
11 changes: 9 additions & 2 deletions tests/agent_unittests/test_agent_protocol.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
}
)
Expand Down Expand Up @@ -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
Expand All @@ -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

Expand Down
147 changes: 78 additions & 69 deletions tests/agent_unittests/test_connect_response_fields.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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", ()),
)


Expand Down Expand Up @@ -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")

Expand All @@ -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")

Expand All @@ -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

0 comments on commit 5469c4a

Please sign in to comment.