Skip to content

Commit 063de03

Browse files
committed
add tests
1 parent 94cc647 commit 063de03

File tree

1 file changed

+50
-6
lines changed

1 file changed

+50
-6
lines changed

tests/test_metric.py

Lines changed: 50 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -119,13 +119,57 @@ def test_retry_on_remote_disconnected(self):
119119

120120
def test_flush_stats_with_tags(self):
121121
lambda_stats = ThreadStatsWriter(True)
122+
original_constant_tags = lambda_stats.thread_stats.constant_tags.copy()
122123
tags = ["tag1:value1", "tag2:value2"]
123-
lambda_stats.flush(tags)
124-
self.mock_threadstats_flush_distributions.assert_called_once_with(
125-
lambda_stats.thread_stats._get_aggregate_metrics_and_dists(float("inf"))[1]
126-
)
127-
for tag in tags:
128-
self.assertTrue(tag in lambda_stats.thread_stats.constant_tags)
124+
125+
# Add a metric to be flushed
126+
lambda_stats.distribution("test.metric", 1, tags=["metric:tag"])
127+
128+
with patch.object(lambda_stats.thread_stats.reporter, 'flush_distributions') as mock_flush_distributions:
129+
lambda_stats.flush(tags)
130+
mock_flush_distributions.assert_called_once()
131+
# Verify that after flush, constant_tags is reset to original
132+
self.assertEqual(lambda_stats.thread_stats.constant_tags, original_constant_tags)
133+
134+
def test_flush_temp_constant_tags(self):
135+
lambda_stats = ThreadStatsWriter(flush_in_thread=True)
136+
lambda_stats.thread_stats.constant_tags = ["initial:tag"]
137+
original_constant_tags = lambda_stats.thread_stats.constant_tags.copy()
138+
139+
lambda_stats.distribution("test.metric", 1, tags=["metric:tag"])
140+
flush_tags = ["flush:tag1", "flush:tag2"]
141+
142+
with patch.object(lambda_stats.thread_stats.reporter, 'flush_distributions') as mock_flush_distributions:
143+
lambda_stats.flush(tags=flush_tags)
144+
mock_flush_distributions.assert_called_once()
145+
flushed_dists = mock_flush_distributions.call_args[0][0]
146+
147+
# Expected tags: original constant_tags + flush_tags + metric tags
148+
expected_tags = original_constant_tags + flush_tags + ["metric:tag"]
149+
150+
# Verify the tags on the metric
151+
self.assertEqual(len(flushed_dists), 1)
152+
metric = flushed_dists[0]
153+
self.assertEqual(sorted(metric['tags']), sorted(expected_tags))
154+
155+
# Verify that constant_tags is reset after flush
156+
self.assertEqual(lambda_stats.thread_stats.constant_tags, original_constant_tags)
157+
158+
# Repeat to ensure tags do not accumulate over multiple flushes
159+
new_flush_tags = ["flush:tag3"]
160+
lambda_stats.distribution("test.metric2", 2, tags=["metric2:tag"])
161+
162+
with patch.object(lambda_stats.thread_stats.reporter, 'flush_distributions') as mock_flush_distributions:
163+
lambda_stats.flush(tags=new_flush_tags)
164+
mock_flush_distributions.assert_called_once()
165+
flushed_dists = mock_flush_distributions.call_args[0][0]
166+
# Expected tags for the new metric
167+
expected_tags = original_constant_tags + new_flush_tags + ["metric2:tag"]
168+
169+
self.assertEqual(len(flushed_dists), 1)
170+
metric = flushed_dists[0]
171+
self.assertEqual(sorted(metric['tags']), sorted(expected_tags))
172+
self.assertEqual(lambda_stats.thread_stats.constant_tags, original_constant_tags)
129173

130174
def test_flush_stats_without_context(self):
131175
flush_stats(lambda_context=None)

0 commit comments

Comments
 (0)