@@ -119,13 +119,57 @@ def test_retry_on_remote_disconnected(self):
119
119
120
120
def test_flush_stats_with_tags (self ):
121
121
lambda_stats = ThreadStatsWriter (True )
122
+ original_constant_tags = lambda_stats .thread_stats .constant_tags .copy ()
122
123
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 )
129
173
130
174
def test_flush_stats_without_context (self ):
131
175
flush_stats (lambda_context = None )
0 commit comments