forked from DataDog/datadog-lambda-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_metric.py
More file actions
126 lines (99 loc) · 4.78 KB
/
test_metric.py
File metadata and controls
126 lines (99 loc) · 4.78 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
import os
import unittest
from unittest.mock import patch, call
from botocore.exceptions import ClientError as BotocoreClientError
from datadog.api.exceptions import ClientError
from datadog_lambda.metric import lambda_metric
from datadog_lambda.api import decrypt_kms_api_key, KMS_ENCRYPTION_CONTEXT_KEY
from datadog_lambda.thread_stats_writer import ThreadStatsWriter
from datadog_lambda.tags import _format_dd_lambda_layer_tag
class TestLambdaMetric(unittest.TestCase):
def setUp(self):
patcher = patch("datadog_lambda.metric.lambda_stats")
self.mock_metric_lambda_stats = patcher.start()
self.addCleanup(patcher.stop)
def test_lambda_metric_tagged_with_dd_lambda_layer(self):
lambda_metric("test", 1)
lambda_metric("test", 1, 123, [])
lambda_metric("test", 1, tags=["tag1:test"])
expected_tag = _format_dd_lambda_layer_tag()
self.mock_metric_lambda_stats.distribution.assert_has_calls(
[
call("test", 1, timestamp=None, tags=[expected_tag]),
call("test", 1, timestamp=123, tags=[expected_tag]),
call("test", 1, timestamp=None, tags=["tag1:test", expected_tag]),
]
)
# let's fake that the extension is present, this should override DD_FLUSH_TO_LOG
@patch("datadog_lambda.metric.should_use_extension", True)
def test_lambda_metric_flush_to_log_with_extension(self):
os.environ["DD_FLUSH_TO_LOG"] = "True"
lambda_metric("test", 1)
expected_tag = _format_dd_lambda_layer_tag()
self.mock_metric_lambda_stats.distribution.assert_has_calls(
[call("test", 1, timestamp=None, tags=[expected_tag])]
)
del os.environ["DD_FLUSH_TO_LOG"]
def test_lambda_metric_flush_to_log(self):
os.environ["DD_FLUSH_TO_LOG"] = "True"
lambda_metric("test", 1)
self.mock_metric_lambda_stats.distribution.assert_not_called()
del os.environ["DD_FLUSH_TO_LOG"]
class TestFlushThreadStats(unittest.TestCase):
def setUp(self):
patcher = patch(
"datadog.threadstats.reporters.HttpReporter.flush_distributions"
)
self.mock_threadstats_flush_distributions = patcher.start()
self.addCleanup(patcher.stop)
def test_retry_on_remote_disconnected(self):
# Raise the RemoteDisconnected error
lambda_stats = ThreadStatsWriter(True)
self.mock_threadstats_flush_distributions.side_effect = ClientError(
"POST",
"https://api.datadoghq.com/api/v1/distribution_points",
"RemoteDisconnected('Remote end closed connection without response')",
)
lambda_stats.flush()
self.assertEqual(self.mock_threadstats_flush_distributions.call_count, 2)
MOCK_FUNCTION_NAME = "myFunction"
# An API key encrypted with KMS and encoded as a base64 string
MOCK_ENCRYPTED_API_KEY_BASE64 = "MjIyMjIyMjIyMjIyMjIyMg=="
# The encrypted API key after it has been decoded from base64
MOCK_ENCRYPTED_API_KEY = "2222222222222222"
# The true value of the API key after decryption by KMS
EXPECTED_DECRYPTED_API_KEY = "1111111111111111"
class TestDecryptKMSApiKey(unittest.TestCase):
def test_key_encrypted_with_encryption_context(self):
os.environ["AWS_LAMBDA_FUNCTION_NAME"] = MOCK_FUNCTION_NAME
class MockKMSClient:
def decrypt(self, CiphertextBlob=None, EncryptionContext={}):
if (
EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY)
!= MOCK_FUNCTION_NAME
):
raise BotocoreClientError({}, "Decrypt")
if CiphertextBlob == MOCK_ENCRYPTED_API_KEY.encode("utf-8"):
return {
"Plaintext": EXPECTED_DECRYPTED_API_KEY.encode("utf-8"),
}
mock_kms_client = MockKMSClient()
decrypted_key = decrypt_kms_api_key(
mock_kms_client, MOCK_ENCRYPTED_API_KEY_BASE64
)
self.assertEqual(decrypted_key, EXPECTED_DECRYPTED_API_KEY)
del os.environ["AWS_LAMBDA_FUNCTION_NAME"]
def test_key_encrypted_without_encryption_context(self):
class MockKMSClient:
def decrypt(self, CiphertextBlob=None, EncryptionContext={}):
if EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != None:
raise BotocoreClientError({}, "Decrypt")
if CiphertextBlob == MOCK_ENCRYPTED_API_KEY.encode("utf-8"):
return {
"Plaintext": EXPECTED_DECRYPTED_API_KEY.encode("utf-8"),
}
mock_kms_client = MockKMSClient()
decrypted_key = decrypt_kms_api_key(
mock_kms_client, MOCK_ENCRYPTED_API_KEY_BASE64
)
self.assertEqual(decrypted_key, EXPECTED_DECRYPTED_API_KEY)