Preserve logger context across flushes

aws_embedded_metrics/logger/metrics_context.py

@@ -125,8 +125,17 @@ def create_copy_with_context(self) -> "MetricsContext":
         new_properties: Dict = {}
         new_properties.update(self.properties)
 
+        # dimensions added with put_dimension will not be copied.
+        # the reason for this is so that you can flush the same scope multiple
+        # times without stacking new dimensions. Example:
+        #
+        # @metric_scope
+        # def my_func(metrics):
+        #  metrics.put_dimensions(...)
+        #
+        # my_func()
+        # my_func()
         new_dimensions: List[Dict] = []
-        new_dimensions = new_dimensions + self.dimensions
 
         new_default_dimensions: Dict = {}
         new_default_dimensions.update(self.default_dimensions)

aws_embedded_metrics/logger/metrics_logger.py

@@ -41,7 +41,7 @@ async def flush(self) -> None:
 
         # accept and reset the context
         sink.accept(self.context)
-        self.context = MetricsContext.empty()
+        self.context = self.context.create_copy_with_context()
 
     def __configureContextForEnvironment(self, env: Environment) -> None:
         default_dimensions = {

tests/logger/test_metrics_context.py

@@ -197,7 +197,7 @@ def test_create_copy_with_context_copies_properties():
     assert context.properties is not new_context.properties
 
 
-def test_create_copy_with_context_copies_dimensions():
+def test_create_copy_with_context_does_not_copy_dimensions():
     # arrange
     context = MetricsContext()
     dimension_key = fake.word()
@@ -208,8 +208,7 @@ def test_create_copy_with_context_copies_dimensions():
     new_context = context.create_copy_with_context()
 
     # assert
-    assert context.dimensions == new_context.dimensions
-    assert context.dimensions is not new_context.dimensions
+    assert len(new_context.dimensions) == 0
 
 
 def test_create_copy_with_context_copies_default_dimensions():
@@ -225,6 +224,20 @@ def test_create_copy_with_context_copies_default_dimensions():
     assert context.default_dimensions is not new_context.default_dimensions
 
 
+def test_create_copy_with_context_does_not_copy_metrics():
+    # arrange
+    context = MetricsContext()
+    prop_key = fake.word()
+    prop_value = fake.word()
+    context.set_property(prop_key, prop_value)
+
+    # act
+    new_context = context.create_copy_with_context()
+
+    # assert
+    assert len(new_context.metrics) == 0
+
+
 def test_set_dimensions_overwrites_all_dimensions():
     # arrange
     context = MetricsContext()
@@ -238,3 +251,20 @@ def test_set_dimensions_overwrites_all_dimensions():
 
     # assert
     assert context.dimensions == expected_dimensions
+
+
+def test_create_copy_with_context_does_not_repeat_dimensions():
+    # arrange
+    context = MetricsContext()
+    expected_dimensions = {fake.word(): fake.word()}
+
+    custom = {fake.word(): fake.word()}
+    context.set_default_dimensions(expected_dimensions)
+    context.put_dimensions(custom)
+
+    new_context = context.create_copy_with_context()
+    new_context.set_default_dimensions(expected_dimensions)
+    new_context.put_dimensions(custom)
+
+    # assert
+    assert len(new_context.get_dimensions()) == 1

tests/logger/test_metrics_logger.py

@@ -180,6 +180,39 @@ async def test_can_set_namespace(mocker):
     assert context.namespace == expected_value
 
 
+@pytest.mark.asyncio
+async def test_context_is_preserved_across_flushes(mocker):
+    # arrange
+    expected_namespace = "Namespace"
+    metric_key = "Metric"
+    expected_dimension_key = "Dim"
+    expected_property_key = "Prop"
+    expected_value = "Value"
+
+    logger, sink, env = get_logger_and_sink(mocker)
+
+    logger.set_namespace(expected_namespace)
+    logger.set_property(expected_property_key, expected_value)
+    logger.set_dimensions({expected_dimension_key: expected_value})
+
+    # act
+    logger.put_metric(metric_key, 0)
+    await logger.flush()
+
+    context = sink.accept.call_args[0][0]
+    assert context.namespace == expected_namespace
+    assert context.properties[expected_property_key] == expected_value
+    assert context.metrics[metric_key].values == [0]
+
+    logger.put_metric(metric_key, 1)
+    await logger.flush()
+
+    context = sink.accept.call_args[0][0]
+    assert context.namespace == expected_namespace
+    assert context.properties[expected_property_key] == expected_value
+    assert context.metrics[metric_key].values == [1]
+
+
 # Test helper methods