-
Notifications
You must be signed in to change notification settings - Fork 45
/
Copy pathtest_metric.py
177 lines (141 loc) · 7.05 KB
/
test_metric.py
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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
import os
import unittest
from unittest.mock import patch, call
from botocore.exceptions import ClientError as BotocoreClientError
from datadog.api.exceptions import ClientError
from datetime import datetime, timedelta
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 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"])
self.mock_metric_lambda_stats.distribution.assert_has_calls(
[
call("test", 1, timestamp=None, tags=[dd_lambda_layer_tag]),
call("test", 1, timestamp=123, tags=[dd_lambda_layer_tag]),
call(
"test", 1, timestamp=None, tags=["tag1:test", dd_lambda_layer_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)
self.mock_metric_lambda_stats.distribution.assert_has_calls(
[call("test", 1, timestamp=None, tags=[dd_lambda_layer_tag])]
)
del os.environ["DD_FLUSH_TO_LOG"]
@patch("datadog_lambda.metric.should_use_extension", True)
def test_lambda_metric_timestamp_with_extension(self):
patcher = patch("datadog_lambda.metric.extension_thread_stats")
self.mock_metric_extension_thread_stats = patcher.start()
self.addCleanup(patcher.stop)
delta = timedelta(minutes=1)
timestamp = int((datetime.now() - delta).timestamp())
lambda_metric("test_timestamp", 1, timestamp)
self.mock_metric_lambda_stats.distribution.assert_not_called()
self.mock_metric_extension_thread_stats.distribution.assert_called_with(
"test_timestamp", 1, timestamp=timestamp, tags=[dd_lambda_layer_tag]
)
@patch("datadog_lambda.metric.should_use_extension", True)
def test_lambda_metric_datetime_with_extension(self):
patcher = patch("datadog_lambda.metric.extension_thread_stats")
self.mock_metric_extension_thread_stats = patcher.start()
self.addCleanup(patcher.stop)
delta = timedelta(hours=5)
timestamp = datetime.now() - delta
lambda_metric("test_timestamp", 1, timestamp)
self.mock_metric_lambda_stats.distribution.assert_not_called()
self.mock_metric_extension_thread_stats.distribution.assert_not_called()
@patch("datadog_lambda.metric.should_use_extension", True)
def test_lambda_metric_invalid_timestamp_with_extension(self):
patcher = patch("datadog_lambda.metric.extension_thread_stats")
self.mock_metric_extension_thread_stats = patcher.start()
self.addCleanup(patcher.stop)
delta = timedelta(hours=5)
timestamp = int((datetime.now() - delta).timestamp())
lambda_metric("test_timestamp", 1, timestamp)
self.mock_metric_lambda_stats.distribution.assert_not_called()
self.mock_metric_extension_thread_stats.distribution.assert_not_called()
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)
def test_flush_stats_with_tags(self):
lambda_stats = ThreadStatsWriter(True)
tags = ["tag1:value1", "tag2:value2"]
lambda_stats.flush(tags)
self.mock_threadstats_flush_distributions.assert_called_once_with(
lambda_stats.thread_stats._get_aggregate_metrics_and_dists(float("inf"))[1]
)
for tag in tags:
self.assertTrue(tag in lambda_stats.thread_stats.constant_tags)
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)