From 3101a4d73b27beb9b70d923804aa77dd9dc06e07 Mon Sep 17 00:00:00 2001
From: Michael Brewer <michael.brewer@gyft.com>
Date: Fri, 21 Aug 2020 23:32:12 -0700
Subject: [PATCH] fix(capture_method): should yield inside with

Changes:
* capture_method should yield from within the "with" statement
* Add missing test cases

Closes #112
---
 aws_lambda_powertools/tracing/tracer.py |  3 +-
 tests/unit/test_tracing.py              | 91 +++++++++++++++++++++++++
 2 files changed, 92 insertions(+), 2 deletions(-)

diff --git a/aws_lambda_powertools/tracing/tracer.py b/aws_lambda_powertools/tracing/tracer.py
index 0ce55e60837..4c12be3fc26 100644
--- a/aws_lambda_powertools/tracing/tracer.py
+++ b/aws_lambda_powertools/tracing/tracer.py
@@ -487,14 +487,13 @@ def decorate(*args, **kwargs):
                     logger.debug(f"Calling method: {method_name}")
                     with method(*args, **kwargs) as return_val:
                         result = return_val
+                        yield result
                     self._add_response_as_metadata(function_name=method_name, data=result, subsegment=subsegment)
                 except Exception as err:
                     logger.exception(f"Exception received from '{method_name}' method")
                     self._add_full_exception_as_metadata(function_name=method_name, error=err, subsegment=subsegment)
                     raise
 
-                yield result
-
         return decorate
 
     def _decorate_sync_function(self, method: Callable = None):
diff --git a/tests/unit/test_tracing.py b/tests/unit/test_tracing.py
index 8f7d9a646dd..16c476ee0fc 100644
--- a/tests/unit/test_tracing.py
+++ b/tests/unit/test_tracing.py
@@ -383,6 +383,72 @@ def handler(event, context):
     assert "test result" in result
 
 
+def test_tracer_yield_from_context_manager_exception_metadata(mocker, provider_stub, in_subsegment_mock):
+    # GIVEN tracer is initialized
+    provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
+    tracer = Tracer(provider=provider, service="booking")
+
+    # WHEN capture_method decorator is used on a context manager
+    # and the method raises an exception
+    @tracer.capture_method
+    @contextlib.contextmanager
+    def yield_with_capture():
+        yield "partial"
+        raise ValueError("test")
+
+    with pytest.raises(ValueError):
+        with yield_with_capture() as partial_val:
+            assert partial_val == "partial"
+
+    # THEN we should add the exception using method name as key plus error
+    # and their service name as the namespace
+    put_metadata_mock_args = in_subsegment_mock.put_metadata.call_args[1]
+    assert put_metadata_mock_args["key"] == "yield_with_capture error"
+    assert isinstance(put_metadata_mock_args["value"], ValueError)
+    assert put_metadata_mock_args["namespace"] == "booking"
+
+
+def test_tracer_yield_from_nested_context_manager(mocker, provider_stub, in_subsegment_mock):
+    # GIVEN tracer is initialized
+    provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
+    tracer = Tracer(provider=provider, service="booking")
+
+    # WHEN capture_method decorator is used on a context manager nesting another context manager
+    class NestedContextManager(object):
+        def __enter__(self):
+            self._value = {"result": "test result"}
+            return self._value
+
+        def __exit__(self, exc_type, exc_val, exc_tb):
+            self._value["result"] = "exit was called before yielding"
+
+    @tracer.capture_method
+    @contextlib.contextmanager
+    def yield_with_capture():
+        with NestedContextManager() as nested_context:
+            yield nested_context
+
+    @tracer.capture_lambda_handler
+    def handler(event, context):
+        response = []
+        with yield_with_capture() as yielded_value:
+            response.append(yielded_value["result"])
+
+        return response
+
+    result = handler({}, {})
+
+    # THEN we should have a subsegment named after the method name
+    # and add its response as trace metadata
+    handler_trace, yield_function_trace = in_subsegment_mock.in_subsegment.call_args_list
+
+    assert "test result" in in_subsegment_mock.put_metadata.call_args[1]["value"]
+    assert in_subsegment_mock.in_subsegment.call_count == 2
+    assert handler_trace == mocker.call(name="## handler")
+    assert yield_function_trace == mocker.call(name="## yield_with_capture")
+    assert "test result" in result
+
+
 def test_tracer_yield_from_generator(mocker, provider_stub, in_subsegment_mock):
     # GIVEN tracer is initialized
     provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
@@ -411,3 +477,28 @@ def handler(event, context):
     assert handler_trace == mocker.call(name="## handler")
     assert generator_fn_trace == mocker.call(name="## generator_fn")
     assert "test result" in result
+
+
+def test_tracer_yield_from_generator_exception_metadata(mocker, provider_stub, in_subsegment_mock):
+    # GIVEN tracer is initialized
+    provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
+    tracer = Tracer(provider=provider, service="booking")
+
+    # WHEN capture_method decorator is used on a generator function
+    # and the method raises an exception
+    @tracer.capture_method
+    def generator_fn():
+        yield "partial"
+        raise ValueError("test")
+
+    with pytest.raises(ValueError):
+        gen = generator_fn()
+        list(gen)
+
+    # THEN we should add the exception using method name as key plus error
+    # and their service name as the namespace
+    put_metadata_mock_args = in_subsegment_mock.put_metadata.call_args[1]
+    assert put_metadata_mock_args["key"] == "generator_fn error"
+    assert put_metadata_mock_args["namespace"] == "booking"
+    assert isinstance(put_metadata_mock_args["value"], ValueError)
+    assert str(put_metadata_mock_args["value"]) == "test"