Repository: awslabs/aws-embedded-metrics-python

Issue:
Preserve logger context across flushes
See https://github.com/awslabs/aws-embedded-metrics-node/issues/15

https://github.com/awslabs/aws-embedded-metrics-python/blob/6cf02e6a5b1849be7aaa5300f2ec9425c5fbc398/aws_embedded_metrics/logger/metrics_logger.py#L44

Discussion:

Patch:
diff --git a/aws_embedded_metrics/logger/metrics_context.py b/aws_embedded_metrics/logger/metrics_context.py
index 5d762e3..283f369 100644
--- a/aws_embedded_metrics/logger/metrics_context.py
+++ b/aws_embedded_metrics/logger/metrics_context.py
@@ -125,8 +125,17 @@ class MetricsContext(object):
         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)
diff --git a/aws_embedded_metrics/logger/metrics_logger.py b/aws_embedded_metrics/logger/metrics_logger.py
index 60843bc..69e1b62 100644
--- a/aws_embedded_metrics/logger/metrics_logger.py
+++ b/aws_embedded_metrics/logger/metrics_logger.py
@@ -41,7 +41,7 @@ class MetricsLogger:
 
         # 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 = {


Test patch:
diff --git a/tests/logger/test_metrics_context.py b/tests/logger/test_metrics_context.py
index e188314..550440d 100644
--- a/tests/logger/test_metrics_context.py
+++ b/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
diff --git a/tests/logger/test_metrics_logger.py b/tests/logger/test_metrics_logger.py
index 23a33a5..895b02b 100644
--- a/tests/logger/test_metrics_logger.py
+++ b/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