Skip to content

Commit 83d7914

Browse files
maxdaymaxday
committed
flake and lint
1 parent 2cd0228 commit 83d7914

File tree

6 files changed

+71
-60
lines changed

6 files changed

+71
-60
lines changed

‎datadog_lambda/constants.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,8 @@ class TraceContextSource(object):
3535
EVENT = "event"
3636
DDTRACE = "ddtrace"
3737

38+
3839
# X-Ray deamon
3940
class XrayDeamon(object):
4041
HOST = "127.0.0.1"
4142
PORT = 2000
42-
43-

‎datadog_lambda/tracing.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,12 @@
1010
from datadog_lambda.constants import (
1111
SamplingPriority,
1212
TraceHeader,
13-
XraySubsegment,
1413
TraceContextSource,
1514
)
1615
from datadog_lambda.xray import (
1716
send_segment,
1817
parse_xray_header,
19-
XRAY_TRACE_ID_HEADER_NAME
18+
XRAY_TRACE_ID_HEADER_NAME,
2019
)
2120
from ddtrace import tracer, patch
2221
from ddtrace import __version__ as ddtrace_version
@@ -32,6 +31,7 @@
3231

3332
FUNCTION_NAME_HEADER_NAME = "AWS_LAMBDA_FUNCTION_NAME"
3433

34+
3535
def _convert_xray_trace_id(xray_trace_id):
3636
"""
3737
Convert X-Ray trace id (hex)'s last 63 bits to a Datadog trace id (int).
@@ -72,7 +72,11 @@ def _get_xray_trace_context():
7272
logger.debug(
7373
"Converted trace context %s from X-Ray segment %s",
7474
trace_context,
75-
(xray_trace_entity["trace_id"], xray_trace_entity["parent_id"], xray_trace_entity["sampled"]),
75+
(
76+
xray_trace_entity["trace_id"],
77+
xray_trace_entity["parent_id"],
78+
xray_trace_entity["sampled"],
79+
),
7680
)
7781
return trace_context
7882

‎datadog_lambda/xray.py

Lines changed: 33 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,17 @@
11
import os
22
import logging
3-
import json
3+
import json
44
import binascii
55
import time
6+
import socket
67

78
from datadog_lambda.constants import XraySubsegment, TraceContextSource
89

910
XRAY_TRACE_ID_HEADER_NAME = "_X_AMZN_TRACE_ID"
1011
XRAY_DAEMON_ADDRESS = "AWS_XRAY_DAEMON_ADDRESS"
1112
logger = logging.getLogger(__name__)
1213

14+
1315
def get_xray_host_port(adress):
1416
if adress == "":
1517
logger.debug("X-Ray daemon env var not set, not sending sub-segment")
@@ -22,6 +24,7 @@ def get_xray_host_port(adress):
2224
host = parts[0]
2325
return (host, port)
2426

27+
2528
def send(host_port_tuple, payload):
2629
sock = None
2730
try:
@@ -36,18 +39,17 @@ def send(host_port_tuple, payload):
3639
except Exception as e_close:
3740
logger.error("Error while closing the socket: %s", str(e_close))
3841

42+
3943
def build_segment_payload(payload):
4044
if payload is None:
4145
return None
42-
header = json.dumps({
43-
"format": "json",
44-
"version": 1
45-
})
46+
header = json.dumps({"format": "json", "version": 1})
4647
return header + "\n" + payload
4748

49+
4850
def parse_xray_header(raw_trace_id):
4951
# Example: Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1
50-
logger.debug("Reading trace context from env var %s", raw_trace_id);
52+
logger.debug("Reading trace context from env var %s", raw_trace_id)
5153
if len(raw_trace_id) == 0:
5254
return None
5355
parts = raw_trace_id.split(";")
@@ -66,39 +68,45 @@ def parse_xray_header(raw_trace_id):
6668
"parent_id": parent,
6769
"trace_id": root,
6870
"sampled": sampled,
69-
"source": TraceContextSource.XRAY
71+
"source": TraceContextSource.XRAY,
7072
}
7173

74+
7275
def generate_random_id():
73-
return binascii.b2a_hex(os.urandom(8)).decode('utf-8')
76+
return binascii.b2a_hex(os.urandom(8)).decode("utf-8")
77+
7478

7579
def build_segment(context, key, metadata):
7680

77-
segment = json.dumps({
78-
"id": generate_random_id(),
79-
"trace_id": context["trace_id"],
80-
"parent_id": context["parent_id"],
81-
"name": XraySubsegment.NAME,
82-
"start_time": time.time(),
83-
"end_time": time.time(),
84-
"type": "subsegment",
85-
"metadata": {
86-
XraySubsegment.NAMESPACE: {
87-
key: metadata,
88-
}
81+
segment = json.dumps(
82+
{
83+
"id": generate_random_id(),
84+
"trace_id": context["trace_id"],
85+
"parent_id": context["parent_id"],
86+
"name": XraySubsegment.NAME,
87+
"start_time": time.time(),
88+
"end_time": time.time(),
89+
"type": "subsegment",
90+
"metadata": {
91+
XraySubsegment.NAMESPACE: {
92+
key: metadata,
93+
}
94+
},
8995
}
90-
})
96+
)
9197
return segment
9298

99+
93100
def send_segment(key, metadata):
94101
host_port_tuple = get_xray_host_port(os.environ.get(XRAY_DAEMON_ADDRESS, ""))
95102
if host_port_tuple is None:
96103
return None
97-
context = parse_xray_header(os.environ.get(XRAY_TRACE_ID_HEADER_NAME,""))
104+
context = parse_xray_header(os.environ.get(XRAY_TRACE_ID_HEADER_NAME, ""))
98105
if context is None:
99-
logger.debug("Failed to create the segment since it was not possible to get the trace context for the header")
106+
logger.debug(
107+
"Failed to create segment since it was not possible to get trace context from header"
108+
)
100109
return None
101110
segment = build_segment(context, key, metadata)
102111
segment_payload = build_segment_payload(segment)
103-
send(host_port_tuple, seegment_payload)
104-
112+
send(host_port_tuple, segment_payload)

‎tests/test_tracing.py

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,11 @@ def test_without_datadog_trace_headers(self):
7979
self.assertEqual(source, "xray")
8080
self.assertDictEqual(
8181
ctx,
82-
{"trace-id": fake_xray_header_value_root_decimal, "parent-id": fake_xray_header_value_parent_decimal, "sampling-priority": "2"},
82+
{
83+
"trace-id": fake_xray_header_value_root_decimal,
84+
"parent-id": fake_xray_header_value_parent_decimal,
85+
"sampling-priority": "2",
86+
},
8387
)
8488
self.assertDictEqual(
8589
get_dd_trace_context(),
@@ -100,7 +104,11 @@ def test_with_incomplete_datadog_trace_headers(self):
100104
self.assertEqual(source, "xray")
101105
self.assertDictEqual(
102106
ctx,
103-
{"trace-id": fake_xray_header_value_root_decimal, "parent-id": fake_xray_header_value_parent_decimal, "sampling-priority": "2"},
107+
{
108+
"trace-id": fake_xray_header_value_root_decimal,
109+
"parent-id": fake_xray_header_value_parent_decimal,
110+
"sampling-priority": "2",
111+
},
104112
)
105113
self.assertDictEqual(
106114
get_dd_trace_context(),
@@ -305,7 +313,7 @@ def test_with_legacy_client_context_datadog_trace_data(self):
305313
self.mock_send_segment.assert_called()
306314
self.mock_send_segment.assert_called_with(
307315
XraySubsegment.TRACE_KEY,
308-
{"trace-id": "666", "parent-id": "777", "sampling-priority": "1"}
316+
{"trace-id": "666", "parent-id": "777", "sampling-priority": "1"},
309317
)
310318

311319
def test_with_new_client_context_datadog_trace_data(self):
@@ -338,7 +346,7 @@ def test_with_new_client_context_datadog_trace_data(self):
338346
self.mock_send_segment.assert_called()
339347
self.mock_send_segment.assert_called_with(
340348
XraySubsegment.TRACE_KEY,
341-
{"trace-id": "666", "parent-id": "777", "sampling-priority": "1"}
349+
{"trace-id": "666", "parent-id": "777", "sampling-priority": "1"},
342350
)
343351

344352
def test_with_complete_datadog_trace_headers_with_mixed_casing(self):
@@ -378,7 +386,7 @@ def test_with_complete_datadog_trace_headers_with_trigger_tags(self):
378386
{
379387
"function_trigger.event_source": "sqs",
380388
"function_trigger.event_source_arn": "arn:aws:sqs:us-east-1:123456789012:MyQueue",
381-
}
389+
},
382390
),
383391
]
384392
)

‎tests/test_wrapper.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -132,8 +132,6 @@ def lambda_handler(event, context):
132132

133133
del os.environ["DD_FLUSH_TO_LOG"]
134134

135-
136-
137135
def test_errors_metric(self):
138136
@datadog_lambda_wrapper
139137
def lambda_handler(event, context):

‎tests/test_xray.py

Lines changed: 17 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,10 @@
44

55
from unittest import mock
66

7-
from datadog_lambda.xray import (
8-
get_xray_host_port,
9-
build_segment_payload,
10-
build_segment
11-
)
7+
from datadog_lambda.xray import get_xray_host_port, build_segment_payload, build_segment
8+
129

1310
class TestXRay(unittest.TestCase):
14-
1511
def test_get_xray_host_port_empty_(self):
1612
result = get_xray_host_port("")
1713
self.assertIsNone(result)
@@ -26,35 +22,33 @@ def test_get_xray_host_port_success(self):
2622
self.assertEqual(1000, result[1])
2723

2824
def test_build_segment_payload_ok(self):
29-
exected_text = "{\"format\": \"json\", \"version\": 1}\nmyPayload"
25+
exected_text = '{"format": "json", "version": 1}\nmyPayload'
3026
self.assertEqual(exected_text, build_segment_payload("myPayload"))
3127

3228
def test_build_segment_payload_no_payload(self):
3329
self.assertIsNone(build_segment_payload(None))
3430

35-
@mock.patch('time.time', mock.MagicMock(return_value=1111))
36-
@mock.patch('datadog_lambda.xray.generate_random_id', mock.MagicMock(return_value="1234abcd"))
31+
@mock.patch("time.time", mock.MagicMock(return_value=1111))
32+
@mock.patch(
33+
"datadog_lambda.xray.generate_random_id",
34+
mock.MagicMock(return_value="1234abcd"),
35+
)
3736
def test_build_segment(self):
3837
context = {
3938
"trace_id": 111000111,
4039
"parent_id": 222000222,
4140
}
4241

43-
value = json.dumps({
44-
"a": "aaa",
45-
"b": "bbb"
46-
})
42+
value = json.dumps({"a": "aaa", "b": "bbb"})
4743
result = build_segment(context, "myKey", "myValue")
4844
jsonResult = json.loads(result)
4945
metadataJson = jsonResult["metadata"]
5046

51-
self.assertEqual("1234abcd",jsonResult["id"])
52-
self.assertEqual(1111,jsonResult["start_time"])
53-
self.assertEqual(1111,jsonResult["end_time"])
54-
self.assertEqual(111000111,jsonResult["trace_id"])
55-
self.assertEqual(222000222,jsonResult["parent_id"])
56-
self.assertEqual("datadog-metadata",jsonResult["name"])
57-
self.assertEqual("subsegment",jsonResult["type"])
58-
self.assertEqual("myValue",metadataJson["datadog"]["myKey"])
59-
60-
47+
self.assertEqual("1234abcd", jsonResult["id"])
48+
self.assertEqual(1111, jsonResult["start_time"])
49+
self.assertEqual(1111, jsonResult["end_time"])
50+
self.assertEqual(111000111, jsonResult["trace_id"])
51+
self.assertEqual(222000222, jsonResult["parent_id"])
52+
self.assertEqual("datadog-metadata", jsonResult["name"])
53+
self.assertEqual("subsegment", jsonResult["type"])
54+
self.assertEqual("myValue", metadataJson["datadog"]["myKey"])

0 commit comments

Comments
 (0)