Skip to content

Commit 92a35b8

Browse files
jcstorms1jcstorms1
committed
update tag collection and tests
1 parent 92a6ed4 commit 92a35b8

File tree

3 files changed

+155
-30
lines changed

3 files changed

+155
-30
lines changed

‎datadog_lambda/tags.py‎

Lines changed: 39 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -25,30 +25,48 @@ def tag_dd_lambda_layer(tags):
2525
return [dd_lambda_layer_tag]
2626

2727

28-
def parse_lambda_tags_from_arn(arn):
28+
def parse_lambda_tags_from_arn(lambda_context):
2929
"""Generate the list of lambda tags based on the data in the arn
3030
Args:
31-
arn (str): Lambda ARN.
32-
ex: arn:aws:lambda:us-east-1:123597598159:function:my-lambda[:optional-version]
31+
lambda_context: Aws lambda context object
32+
ex: lambda_context.arn = arn:aws:lambda:us-east-1:123597598159:function:my-lambda:1
3333
"""
34-
# Cap the number of times to split
35-
split_arn = arn.split(":")
34+
# Set up flag for extra testing to distinguish between a version or alias
35+
hasAlias = False
36+
# Cap the number of times to spli
37+
split_arn = lambda_context.invoked_function_arn.split(":")
3638

3739
if len(split_arn) > 7:
40+
hasAlias = True
3841
_, _, _, region, account_id, _, function_name, alias = split_arn
39-
resource = function_name + ":" + alias
40-
4142
else:
4243
_, _, _, region, account_id, _, function_name = split_arn
43-
resource = function_name
4444

45-
return [
45+
# Add the standard tags to a list
46+
tags = [
4647
"region:{}".format(region),
4748
"account_id:{}".format(account_id),
4849
"functionname:{}".format(function_name),
49-
"resource:{}".format(resource),
5050
]
5151

52+
# Check if we have a version or alias
53+
if hasAlias:
54+
# If $Latest, drop the $ for datadog tag convention. A lambda alias can't start with $
55+
if alias.startswith("$"):
56+
alias = alias[1:]
57+
# Versions are numeric. Aliases need the executed version tag
58+
elif not check_if_number(alias):
59+
tags.append("executedversion:{}".format(lambda_context.function_version))
60+
# create resource tag with function name and alias/version
61+
resource = "resource:{}:{}".format(function_name, alias)
62+
else:
63+
# Resource is only the function name otherwise
64+
resource = "resource:{}".format(function_name)
65+
66+
tags.append(resource)
67+
68+
return tags
69+
5270

5371
def get_runtime_tag():
5472
"""Get the runtime tag from the current Python version
@@ -63,10 +81,19 @@ def get_runtime_tag():
6381
def get_enhanced_metrics_tags(lambda_context):
6482
"""Get the list of tags to apply to enhanced metrics
6583
"""
66-
return parse_lambda_tags_from_arn(lambda_context.invoked_function_arn) + [
84+
return parse_lambda_tags_from_arn(lambda_context) + [
6785
get_cold_start_tag(),
6886
"memorysize:{}".format(lambda_context.memory_limit_in_mb),
69-
"executedversion:{}".format(lambda_context.function_version),
7087
get_runtime_tag(),
7188
_format_dd_lambda_layer_tag(),
7289
]
90+
91+
92+
def check_if_number(alias):
93+
""" Check if the alias is a version or number. Python 2 has no easy way to test this like Python 3
94+
"""
95+
try:
96+
float(alias)
97+
return True
98+
except ValueError:
99+
return False

‎tests/test_tags.py‎

Lines changed: 46 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,29 +1,68 @@
11
import unittest
22

33
try:
4-
from unittest.mock import patch
4+
from unittest.mock import patch, MagicMock
55
except ImportError:
6-
from mock import patch
6+
from mock import patch, MagicMock
77

88
from datadog_lambda.tags import parse_lambda_tags_from_arn, get_runtime_tag
99

1010

11+
def get_mock_context(
12+
invoked_function_arn="arn:aws:lambda:us-east-1:1234597598159:function:swf-hello-test:$Latest",
13+
function_version="1",
14+
):
15+
lambda_context = MagicMock()
16+
lambda_context.invoked_function_arn = invoked_function_arn
17+
lambda_context.function_version = function_version
18+
return lambda_context
19+
20+
1121
class TestMetricTags(unittest.TestCase):
1222
def setUp(self):
1323
patcher = patch("datadog_lambda.tags.python_version_tuple")
1424
self.mock_python_version_tuple = patcher.start()
1525
self.addCleanup(patcher.stop)
1626

17-
def test_parse_lambda_tags_from_arn(self):
27+
def test_parse_lambda_tags_from_arn_latest(self):
28+
self.assertListEqual(
29+
parse_lambda_tags_from_arn(get_mock_context()),
30+
[
31+
"region:us-east-1",
32+
"account_id:1234597598159",
33+
"functionname:swf-hello-test",
34+
"resource:swf-hello-test:Latest",
35+
],
36+
)
37+
38+
def test_parse_lambda_tags_from_arn_version(self):
39+
lambda_context = get_mock_context()
40+
lambda_context.invoked_function_arn = (
41+
"arn:aws:lambda:us-east-1:1234597598159:function:swf-hello-test:3"
42+
)
43+
self.assertListEqual(
44+
parse_lambda_tags_from_arn(lambda_context),
45+
[
46+
"region:us-east-1",
47+
"account_id:1234597598159",
48+
"functionname:swf-hello-test",
49+
"resource:swf-hello-test:3",
50+
],
51+
)
52+
53+
def test_parse_lambda_tags_from_arn_alias(self):
54+
lambda_context = get_mock_context()
55+
lambda_context.invoked_function_arn = (
56+
"arn:aws:lambda:us-east-1:1234597598159:function:swf-hello-test:my_alias-1"
57+
)
1858
self.assertListEqual(
19-
parse_lambda_tags_from_arn(
20-
"arn:aws:lambda:us-east-1:1234597598159:function:swf-hello-test:function-alias"
21-
),
59+
parse_lambda_tags_from_arn(lambda_context),
2260
[
2361
"region:us-east-1",
2462
"account_id:1234597598159",
2563
"functionname:swf-hello-test",
26-
"resource:swf-hello-test:function-alias",
64+
"executedversion:1",
65+
"resource:swf-hello-test:my_alias-1",
2766
],
2867
)
2968

‎tests/test_wrapper.py‎

Lines changed: 70 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
def get_mock_context(
1414
aws_request_id="request-id-1",
1515
memory_limit_in_mb="256",
16-
invoked_function_arn="arn:aws:lambda:us-west-1:123457598159:function:python-layer-test:foobar",
16+
invoked_function_arn="arn:aws:lambda:us-west-1:123457598159:function:python-layer-test:1",
1717
function_version="1",
1818
):
1919
lambda_context = MagicMock()
@@ -146,10 +146,9 @@ def lambda_handler(event, context):
146146
"region:us-west-1",
147147
"account_id:123457598159",
148148
"functionname:python-layer-test",
149-
"resource:python-layer-test:foobar",
149+
"resource:python-layer-test:1",
150150
"cold_start:true",
151151
"memorysize:256",
152-
"executedversion:1",
153152
"runtime:python2.7",
154153
"dd_lambda_layer:datadog-python27_0.1.0",
155154
],
@@ -176,10 +175,9 @@ def lambda_handler(event, context):
176175
"region:us-west-1",
177176
"account_id:123457598159",
178177
"functionname:python-layer-test",
179-
"resource:python-layer-test:foobar",
178+
"resource:python-layer-test:1",
180179
"cold_start:true",
181180
"memorysize:256",
182-
"executedversion:1",
183181
"runtime:python2.7",
184182
"dd_lambda_layer:datadog-python27_0.1.0",
185183
],
@@ -191,10 +189,9 @@ def lambda_handler(event, context):
191189
"region:us-west-1",
192190
"account_id:123457598159",
193191
"functionname:python-layer-test",
194-
"resource:python-layer-test:foobar",
192+
"resource:python-layer-test:1",
195193
"cold_start:true",
196194
"memorysize:256",
197-
"executedversion:1",
198195
"runtime:python2.7",
199196
"dd_lambda_layer:datadog-python27_0.1.0",
200197
],
@@ -226,10 +223,9 @@ def lambda_handler(event, context):
226223
"region:us-west-1",
227224
"account_id:123457598159",
228225
"functionname:python-layer-test",
229-
"resource:python-layer-test:foobar",
226+
"resource:python-layer-test:1",
230227
"cold_start:true",
231228
"memorysize:256",
232-
"executedversion:1",
233229
"runtime:python2.7",
234230
"dd_lambda_layer:datadog-python27_0.1.0",
235231
],
@@ -241,17 +237,80 @@ def lambda_handler(event, context):
241237
"region:us-west-1",
242238
"account_id:123457598159",
243239
"functionname:python-layer-test",
244-
"resource:python-layer-test:foobar",
240+
"resource:python-layer-test:1",
245241
"cold_start:false",
246242
"memorysize:256",
247-
"executedversion:1",
248243
"runtime:python2.7",
249244
"dd_lambda_layer:datadog-python27_0.1.0",
250245
],
251246
),
252247
]
253248
)
254249

250+
def test_enhanced_metrics_latest(self):
251+
@datadog_lambda_wrapper
252+
def lambda_handler(event, context):
253+
lambda_metric("test.metric", 100)
254+
255+
lambda_event = {}
256+
lambda_context = get_mock_context()
257+
258+
lambda_context.invoked_function_arn = (
259+
"arn:aws:lambda:us-west-1:123457598159:function:python-layer-test:$Latest"
260+
)
261+
lambda_handler(lambda_event, lambda_context)
262+
263+
self.mock_write_metric_point_to_stdout.assert_has_calls(
264+
[
265+
call(
266+
"aws.lambda.enhanced.invocations",
267+
1,
268+
tags=[
269+
"region:us-west-1",
270+
"account_id:123457598159",
271+
"functionname:python-layer-test",
272+
"resource:python-layer-test:Latest",
273+
"cold_start:true",
274+
"memorysize:256",
275+
"runtime:python2.7",
276+
"dd_lambda_layer:datadog-python27_0.1.0",
277+
],
278+
)
279+
]
280+
)
281+
282+
def test_enhanced_metrics_alias(self):
283+
@datadog_lambda_wrapper
284+
def lambda_handler(event, context):
285+
lambda_metric("test.metric", 100)
286+
287+
lambda_event = {}
288+
lambda_context = get_mock_context()
289+
# tests wouldn't run because line was too long
290+
alias_arn = "arn:aws:lambda:us-west-1:123457598159:function:python-layer-test:My_alias-1"
291+
lambda_context.invoked_function_arn = alias_arn
292+
lambda_handler(lambda_event, lambda_context)
293+
294+
self.mock_write_metric_point_to_stdout.assert_has_calls(
295+
[
296+
call(
297+
"aws.lambda.enhanced.invocations",
298+
1,
299+
tags=[
300+
"region:us-west-1",
301+
"account_id:123457598159",
302+
"functionname:python-layer-test",
303+
"executedversion:1",
304+
"resource:python-layer-test:My_alias-1",
305+
"cold_start:true",
306+
"memorysize:256",
307+
"runtime:python2.7",
308+
"dd_lambda_layer:datadog-python27_0.1.0",
309+
],
310+
)
311+
]
312+
)
313+
255314
def test_no_enhanced_metrics_without_env_var(self):
256315
os.environ["DD_ENHANCED_METRICS"] = "false"
257316

0 commit comments

Comments
 (0)