Added MetricsLogger.add_stack_trace()

Author: benkehoe

aws_embedded_metrics/__init__.py

@@ -13,4 +13,6 @@
 
 name = "aws_embedded_metrics"
 
-from aws_embedded_metrics.metric_scope import metric_scope  # noqa: F401
+from aws_embedded_metrics.metric_scope import metric_scope  # noqa: F401 E402
+from aws_embedded_metrics.logger.metrics_logger import MetricsLogger  # noqa: F401 E402
+from aws_embedded_metrics.logger.metrics_context import MetricsContext  # noqa: F401 E402

aws_embedded_metrics/logger/metrics_logger.py

@@ -14,7 +14,9 @@
 from aws_embedded_metrics.environment import Environment
 from aws_embedded_metrics.logger.metrics_context import MetricsContext
 from aws_embedded_metrics.config import get_config
-from typing import Any, Awaitable, Callable, Dict
+from typing import Any, Awaitable, Callable, Dict, Tuple
+import sys
+import traceback
 
 Config = get_config()
 
@@ -73,6 +75,35 @@ def put_metric(self, key: str, value: float, unit: str = "None") -> "MetricsLogg
         self.context.put_metric(key, value, unit)
         return self
 
+    def add_stack_trace(self, key: str, details: Any = None, exc_info: Tuple = None) -> "MetricsLogger":
+        if not exc_info:
+            exc_info = sys.exc_info()
+
+        err_cls, err, tb = exc_info
+
+        if err_cls is None:
+            error_type = None
+            error_str = None
+            traceback_str = None
+        else:
+            if err_cls.__module__ == "builtins":
+                error_type = err_cls.__name__
+            else:
+                error_type = "{module}.{name}".format(module=err_cls.__module__, name=err_cls.__name__)
+            error_str = str(err)
+            traceback_str = ''.join(traceback.format_tb(tb))
+
+        trace_value = {}
+        if details:
+            trace_value["details"] = details
+        trace_value.update({
+            "error_type": error_type,
+            "error_str": error_str,
+            "traceback": traceback_str,
+        })
+        self.set_property(key, trace_value)
+        return self
+
     def new(self) -> "MetricsLogger":
         return MetricsLogger(
             self.resolve_environment, self.context.create_copy_with_context()

tests/logger/test_metrics_logger.py

@@ -7,6 +7,7 @@
 from asyncio import Future
 from importlib import reload
 import os
+import sys
 
 fake = Faker()
 
@@ -50,6 +51,108 @@ async def test_can_put_metric(mocker):
     assert context.metrics[expected_key].unit == "None"
 
 
+@pytest.mark.asyncio
+async def test_can_add_stack_trace(mocker):
+    # arrange
+    expected_key = fake.word()
+    expected_details = fake.word()
+    expected_error_str = fake.word()
+
+    logger, sink, env = get_logger_and_sink(mocker)
+
+    from configparser import Error  # Just some non-builtin exception
+
+    # act
+    try:
+        raise Error(expected_error_str)
+    except Error:
+        logger.add_stack_trace(expected_key, expected_details)
+    await logger.flush()
+
+    # assert
+    context = get_flushed_context(sink)
+    value = context.properties[expected_key]
+    assert isinstance(value, dict)
+    assert value["details"] == expected_details
+    assert value["error_type"] == "configparser.Error"
+    assert value["error_str"] == expected_error_str
+    assert value["traceback"].split("\n")[-2] == "    raise Error(expected_error_str)"
+
+
+@pytest.mark.asyncio
+async def test_can_add_stack_trace_for_builtin(mocker):
+    # arrange
+    expected_key = fake.word()
+    expected_details = fake.word()
+    expected_error_str = fake.word()
+
+    logger, sink, env = get_logger_and_sink(mocker)
+
+    # act
+    try:
+        raise ValueError(expected_error_str)
+    except ValueError:
+        logger.add_stack_trace(expected_key, expected_details)
+    await logger.flush()
+
+    # assert
+    context = get_flushed_context(sink)
+    value = context.properties[expected_key]
+    assert isinstance(value, dict)
+    assert value["details"] == expected_details
+    assert value["error_type"] == "ValueError"
+    assert value["error_str"] == expected_error_str
+    assert value["traceback"].split("\n")[-2] == "    raise ValueError(expected_error_str)"
+
+
+@pytest.mark.asyncio
+async def test_can_add_empty_stack_trace(mocker):
+    # arrange
+    expected_key = fake.word()
+
+    logger, sink, env = get_logger_and_sink(mocker)
+
+    # act
+    logger.add_stack_trace(expected_key)
+    await logger.flush()
+
+    # assert
+    context = get_flushed_context(sink)
+    value = context.properties[expected_key]
+    assert isinstance(value, dict)
+    assert "value" not in value
+    assert value["error_type"] is None
+    assert value["error_str"] is None
+    assert value["traceback"] is None
+
+
+@pytest.mark.asyncio
+async def test_can_add_stack_trace_manually(mocker):
+    # arrange
+    expected_key = fake.word()
+    expected_details = fake.word()
+    expected_error_str = fake.word()
+
+    logger, sink, env = get_logger_and_sink(mocker)
+
+    # act
+    try:
+        raise ValueError(expected_error_str)
+    except ValueError:
+        exc_info = sys.exc_info()
+    logger.add_stack_trace(expected_key, expected_details, exc_info=exc_info)
+    await logger.flush()
+
+    # assert
+    context = get_flushed_context(sink)
+    value = context.properties[expected_key]
+    assert isinstance(value, dict)
+    assert value["details"] == expected_details
+    assert value["error_type"] == "ValueError"
+    assert value["error_str"] == expected_error_str
+    assert value["traceback"].split("\n")[-2] == "    raise ValueError(expected_error_str)"
+
+
 @pytest.mark.asyncio
 async def test_put_metric_appends_values_to_array(mocker):
     # arrange