Skip to content

Commit

Permalink
Merge branch 'main' into fix-kafka-tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@mergify
mergify[bot] authored Jul 14, 2023
2 parents ae0cece + 53fc51a commit c39d5be
Show file tree
Hide file tree
Showing 3 changed files with 153 additions and 76 deletions.
22 changes: 16 additions & 6 deletions newrelic/api/transaction.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1030,7 +1030,9 @@ def _create_distributed_trace_data(self):

settings = self._settings
account_id = settings.account_id
trusted_account_key = settings.trusted_account_key
trusted_account_key = settings.trusted_account_key or (
self._settings.serverless_mode.enabled and self._settings.account_id
)
application_id = settings.primary_application_id

if not (account_id and application_id and trusted_account_key and settings.distributed_tracing.enabled):
Expand Down Expand Up @@ -1121,7 +1123,10 @@ def _can_accept_distributed_trace_headers(self):
return False

settings = self._settings
if not (settings.distributed_tracing.enabled and settings.trusted_account_key):
trusted_account_key = settings.trusted_account_key or (
self._settings.serverless_mode.enabled and self._settings.account_id
)
if not (settings.distributed_tracing.enabled and trusted_account_key):
return False

if self._distributed_trace_state:
Expand Down Expand Up @@ -1167,10 +1172,13 @@ def _accept_distributed_trace_payload(self, payload, transport_type="HTTP"):

settings = self._settings
account_id = data.get("ac")
trusted_account_key = settings.trusted_account_key or (
self._settings.serverless_mode.enabled and self._settings.account_id
)

# If trust key doesn't exist in the payload, use account_id
received_trust_key = data.get("tk", account_id)
if settings.trusted_account_key != received_trust_key:
if trusted_account_key != received_trust_key:
self._record_supportability("Supportability/DistributedTrace/AcceptPayload/Ignored/UntrustedAccount")
if settings.debug.log_untrusted_distributed_trace_keys:
_logger.debug(
Expand Down Expand Up @@ -1279,8 +1287,10 @@ def accept_distributed_trace_headers(self, headers, transport_type="HTTP"):
tracestate = ensure_str(tracestate)
try:
vendors = W3CTraceState.decode(tracestate)
tk = self._settings.trusted_account_key
payload = vendors.pop(tk + "@nr", "")
trusted_account_key = self._settings.trusted_account_key or (
self._settings.serverless_mode.enabled and self._settings.account_id
)
payload = vendors.pop(trusted_account_key + "@nr", "")
self.tracing_vendors = ",".join(vendors.keys())
self.tracestate = vendors.text(limit=31)
except:
Expand All @@ -1289,7 +1299,7 @@ def accept_distributed_trace_headers(self, headers, transport_type="HTTP"):
# Remove trusted new relic header if available and parse
if payload:
try:
tracestate_data = NrTraceState.decode(payload, tk)
tracestate_data = NrTraceState.decode(payload, trusted_account_key)
except:
tracestate_data = None
if tracestate_data:
Expand Down
63 changes: 63 additions & 0 deletions tests/agent_features/test_distributed_tracing.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@

from newrelic.api.application import application_instance
from newrelic.api.background_task import BackgroundTask, background_task
from newrelic.api.external_trace import ExternalTrace
from newrelic.api.time_trace import current_trace
from newrelic.api.transaction import (
accept_distributed_trace_headers,
Expand Down Expand Up @@ -361,3 +362,65 @@ def test_current_span_id_inside_transaction():
def test_current_span_id_outside_transaction():
span_id = current_span_id()
assert span_id is None


@pytest.mark.parametrize("trusted_account_key", ("1", None), ids=("tk_set", "tk_unset"))
def test_outbound_dt_payload_generation(trusted_account_key):
@override_application_settings(
{
"distributed_tracing.enabled": True,
"account_id": "1",
"trusted_account_key": trusted_account_key,
"primary_application_id": "1",
}
)
@background_task(name="test_outbound_dt_payload_generation")
def _test_outbound_dt_payload_generation():
transaction = current_transaction()
payload = ExternalTrace.generate_request_headers(transaction)
if trusted_account_key:
assert payload
# Ensure trusted account key present as vendor
assert dict(payload)["tracestate"].startswith("1@nr=")
else:
assert not payload

_test_outbound_dt_payload_generation()


@pytest.mark.parametrize("trusted_account_key", ("1", None), ids=("tk_set", "tk_unset"))
def test_inbound_dt_payload_acceptance(trusted_account_key):
@override_application_settings(
{
"distributed_tracing.enabled": True,
"account_id": "1",
"trusted_account_key": trusted_account_key,
"primary_application_id": "1",
}
)
@background_task(name="_test_inbound_dt_payload_acceptance")
def _test_inbound_dt_payload_acceptance():
transaction = current_transaction()

payload = {
"v": [0, 1],
"d": {
"ty": "Mobile",
"ac": "1",
"tk": "1",
"ap": "2827902",
"pa": "5e5733a911cfbc73",
"id": "7d3efb1b173fecfa",
"tr": "d6b4ba0c3a712ca",
"ti": 1518469636035,
"tx": "8703ff3d88eefe9d",
},
}

result = transaction.accept_distributed_trace_payload(payload)
if trusted_account_key:
assert result
else:
assert not result

_test_inbound_dt_payload_acceptance()
144 changes: 74 additions & 70 deletions tests/agent_features/test_serverless_mode.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,16 @@
# limitations under the License.

import json

import pytest
from testing_support.fixtures import override_generic_settings
from testing_support.validators.validate_serverless_data import validate_serverless_data
from testing_support.validators.validate_serverless_metadata import (
validate_serverless_metadata,
)
from testing_support.validators.validate_serverless_payload import (
validate_serverless_payload,
)

from newrelic.api.application import application_instance
from newrelic.api.background_task import background_task
Expand All @@ -22,23 +31,14 @@
from newrelic.api.transaction import current_transaction
from newrelic.core.config import global_settings

from testing_support.fixtures import override_generic_settings
from testing_support.validators.validate_serverless_data import (
validate_serverless_data)
from testing_support.validators.validate_serverless_payload import (
validate_serverless_payload)
from testing_support.validators.validate_serverless_metadata import (
validate_serverless_metadata)


@pytest.fixture(scope='function')
@pytest.fixture(scope="function")
def serverless_application(request):
settings = global_settings()
orig = settings.serverless_mode.enabled
settings.serverless_mode.enabled = True

application_name = 'Python Agent Test (test_serverless_mode:%s)' % (
request.node.name)
application_name = "Python Agent Test (test_serverless_mode:%s)" % (request.node.name)
application = application_instance(application_name)
application.activate()

Expand All @@ -48,17 +48,18 @@ def serverless_application(request):


def test_serverless_payload(capsys, serverless_application):

@override_generic_settings(serverless_application.settings, {
'distributed_tracing.enabled': True,
})
@override_generic_settings(
serverless_application.settings,
{
"distributed_tracing.enabled": True,
},
)
@validate_serverless_data(
expected_methods=('metric_data', 'analytic_event_data'),
forgone_methods=('preconnect', 'connect', 'get_agent_commands'))
expected_methods=("metric_data", "analytic_event_data"),
forgone_methods=("preconnect", "connect", "get_agent_commands"),
)
@validate_serverless_payload()
@background_task(
application=serverless_application,
name='test_serverless_payload')
@background_task(application=serverless_application, name="test_serverless_payload")
def _test():
transaction = current_transaction()
assert transaction.settings.serverless_mode.enabled
Expand All @@ -75,17 +76,15 @@ def _test():


def test_no_cat_headers(serverless_application):
@background_task(
application=serverless_application,
name='test_cat_headers')
@background_task(application=serverless_application, name="test_cat_headers")
def _test_cat_headers():
transaction = current_transaction()

payload = ExternalTrace.generate_request_headers(transaction)
assert not payload

trace = ExternalTrace('testlib', 'http://example.com')
response_headers = [('X-NewRelic-App-Data', 'Cookies')]
trace = ExternalTrace("testlib", "http://example.com")
response_headers = [("X-NewRelic-App-Data", "Cookies")]
with trace:
trace.process_response_headers(response_headers)

Expand All @@ -94,61 +93,66 @@ def _test_cat_headers():
_test_cat_headers()


def test_dt_outbound(serverless_application):
@override_generic_settings(serverless_application.settings, {
'distributed_tracing.enabled': True,
'account_id': '1',
'trusted_account_key': '1',
'primary_application_id': '1',
})
@background_task(
application=serverless_application,
name='test_dt_outbound')
def _test_dt_outbound():
@pytest.mark.parametrize("trusted_account_key", ("1", None), ids=("tk_set", "tk_unset"))
def test_outbound_dt_payload_generation(serverless_application, trusted_account_key):
@override_generic_settings(
serverless_application.settings,
{
"distributed_tracing.enabled": True,
"account_id": "1",
"trusted_account_key": trusted_account_key,
"primary_application_id": "1",
},
)
@background_task(application=serverless_application, name="test_outbound_dt_payload_generation")
def _test_outbound_dt_payload_generation():
transaction = current_transaction()
payload = ExternalTrace.generate_request_headers(transaction)
assert payload

_test_dt_outbound()


def test_dt_inbound(serverless_application):
@override_generic_settings(serverless_application.settings, {
'distributed_tracing.enabled': True,
'account_id': '1',
'trusted_account_key': '1',
'primary_application_id': '1',
})
@background_task(
application=serverless_application,
name='test_dt_inbound')
def _test_dt_inbound():
# Ensure trusted account key or account ID present as vendor
assert dict(payload)["tracestate"].startswith("1@nr=")

_test_outbound_dt_payload_generation()


@pytest.mark.parametrize("trusted_account_key", ("1", None), ids=("tk_set", "tk_unset"))
def test_inbound_dt_payload_acceptance(serverless_application, trusted_account_key):
@override_generic_settings(
serverless_application.settings,
{
"distributed_tracing.enabled": True,
"account_id": "1",
"trusted_account_key": trusted_account_key,
"primary_application_id": "1",
},
)
@background_task(application=serverless_application, name="test_inbound_dt_payload_acceptance")
def _test_inbound_dt_payload_acceptance():
transaction = current_transaction()

payload = {
'v': [0, 1],
'd': {
'ty': 'Mobile',
'ac': '1',
'tk': '1',
'ap': '2827902',
'pa': '5e5733a911cfbc73',
'id': '7d3efb1b173fecfa',
'tr': 'd6b4ba0c3a712ca',
'ti': 1518469636035,
'tx': '8703ff3d88eefe9d',
}
"v": [0, 1],
"d": {
"ty": "Mobile",
"ac": "1",
"tk": "1",
"ap": "2827902",
"pa": "5e5733a911cfbc73",
"id": "7d3efb1b173fecfa",
"tr": "d6b4ba0c3a712ca",
"ti": 1518469636035,
"tx": "8703ff3d88eefe9d",
},
}

result = transaction.accept_distributed_trace_payload(payload)
assert result

_test_dt_inbound()
_test_inbound_dt_payload_acceptance()


@pytest.mark.parametrize('arn_set', (True, False))
@pytest.mark.parametrize("arn_set", (True, False))
def test_payload_metadata_arn(serverless_application, arn_set):

# If the session object gathers the arn from the settings object before the
# lambda handler records it there, then this test will fail.

Expand All @@ -157,17 +161,17 @@ def test_payload_metadata_arn(serverless_application, arn_set):

arn = None
if arn_set:
arn = 'arrrrrrrrrrRrrrrrrrn'
arn = "arrrrrrrrrrRrrrrrrrn"

settings.aws_lambda_metadata.update({'arn': arn, 'function_version': '$LATEST'})
settings.aws_lambda_metadata.update({"arn": arn, "function_version": "$LATEST"})

class Context(object):
invoked_function_arn = arn

@validate_serverless_metadata(exact_metadata={'arn': arn})
@validate_serverless_metadata(exact_metadata={"arn": arn})
@lambda_handler(application=serverless_application)
def handler(event, context):
assert settings.aws_lambda_metadata['arn'] == arn
assert settings.aws_lambda_metadata["arn"] == arn
return {}

try:
Expand Down

0 comments on commit c39d5be

Please sign in to comment.