[aiobotocore] update context manager so it works using async with

Author: Emanuele Palazzetti

ddtrace/contrib/aiobotocore/patch.py

@@ -56,11 +56,14 @@ def read(self, *args, **kwargs):
     if PYTHON_VERSION >= (3, 5, 0):
         @asyncio.coroutine
         def __aenter__(self):
-            return self.__wrapped__.__aenter__()
+            # call the wrapped method but return the object proxy
+            yield from self.__wrapped__.__aenter__()
+            return self
 
         @asyncio.coroutine
         def __aexit__(self, *args, **kwargs):
-            return self.__wrapped__.__aexit__(*args, **kwargs)
+            response = yield from self.__wrapped__.__aexit__(*args, **kwargs)
+            return response
 
 
 def truncate_arg_value(value, max_len=1024):

tests/contrib/aiobotocore/test_35.py

@@ -0,0 +1,56 @@
+from nose.tools import eq_, ok_, assert_raises
+from botocore.errorfactory import ClientError
+
+from ddtrace.contrib.aiobotocore.patch import patch, unpatch
+
+from .utils import aiobotocore_client
+from ..asyncio.utils import AsyncioTestCase, mark_asyncio
+from ...test_tracer import get_dummy_tracer
+
+
+class AIOBotocoreTest(AsyncioTestCase):
+    """Botocore integration testsuite"""
+    def setUp(self):
+        super(AIOBotocoreTest, self).setUp()
+        patch()
+        self.tracer = get_dummy_tracer()
+
+    def tearDown(self):
+        super(AIOBotocoreTest, self).tearDown()
+        unpatch()
+        self.tracer = None
+
+    @mark_asyncio
+    async def test_response_context_manager(self):
+        # the client should call the wrapped __aenter__ and return the
+        # object proxy
+        with aiobotocore_client('s3', self.tracer) as s3:
+            # prepare S3 and flush traces if any
+            await s3.create_bucket(Bucket='tracing')
+            await s3.put_object(Bucket='tracing', Key='apm', Body=b'')
+            self.tracer.writer.pop_traces()
+            # `async with` under test
+            response = await s3.get_object(Bucket='tracing', Key='apm')
+            async with response['Body'] as stream:
+                await stream.read()
+
+        traces = self.tracer.writer.pop_traces()
+        eq_(len(traces), 2)
+        eq_(len(traces[0]), 1)
+        eq_(len(traces[1]), 1)
+
+        span = traces[0][0]
+        eq_(span.get_tag('aws.operation'), 'GetObject')
+        eq_(span.get_tag('http.status_code'), '200')
+        eq_(span.service, 'aws.s3')
+        eq_(span.resource, 's3.getobject')
+
+        read_span = traces[1][0]
+        eq_(read_span.get_tag('aws.operation'), 'GetObject')
+        eq_(read_span.get_tag('http.status_code'), '200')
+        eq_(read_span.service, 'aws.s3')
+        eq_(read_span.resource, 's3.getobject')
+        eq_(read_span.name, 's3.command.read')
+        # enforce parenting
+        eq_(read_span.parent_id, span.span_id)
+        eq_(read_span.trace_id, span.trace_id)

tox.ini

@@ -210,7 +210,8 @@ commands =
     {py27}-pylons: nosetests {posargs} tests/contrib/pylons
     {py27,py34}-boto: nosetests {posargs} tests/contrib/boto
     {py27,py34}-botocore: nosetests {posargs} tests/contrib/botocore
-    aiobotocore{02,03,04}: nosetests {posargs} tests/contrib/aiobotocore
+    py{34}-aiobotocore{02,03,04}: nosetests {posargs} --exclude=".*(test_35).*" tests/contrib/aiobotocore
+    py{35,36}-aiobotocore{02,03,04}: nosetests {posargs} tests/contrib/aiobotocore
     bottle{12}: nosetests {posargs} tests/contrib/bottle/test.py
     bottle-autopatch{12}: ddtrace-run nosetests {posargs} tests/contrib/bottle/test_autopatch.py
     cassandra{35,36,37,38}: nosetests {posargs} tests/contrib/cassandra