Merge pull request #307 from palazzem/aiobotocore-response

[aiobotocore] update async with context manager

Author: Emanuele Palazzetti

ddtrace/contrib/aiobotocore/patch.py

@@ -33,19 +33,19 @@ def unpatch():
 class WrappedClientResponseContentProxy(wrapt.ObjectProxy):
     def __init__(self, body, pin, parent_span):
         super(WrappedClientResponseContentProxy, self).__init__(body)
-        self.__pin = pin
-        self.__parent_span = parent_span
+        self._self_pin = pin
+        self._self_parent_span = parent_span
 
     @asyncio.coroutine
     def read(self, *args, **kwargs):
         # async read that must be child of the parent span operation
-        operation_name = '{}.read'.format(self.__parent_span.name)
+        operation_name = '{}.read'.format(self._self_parent_span.name)
 
-        with self.__pin.tracer.start_span(operation_name, child_of=self.__parent_span) as span:
+        with self._self_pin.tracer.start_span(operation_name, child_of=self._self_parent_span) as span:
             # inherit parent attributes
-            span.resource = self.__parent_span.resource
-            span.span_type = self.__parent_span.span_type
-            span.meta = dict(self.__parent_span.meta)
+            span.resource = self._self_parent_span.resource
+            span.span_type = self._self_parent_span.span_type
+            span.meta = dict(self._self_parent_span.meta)
 
             result = yield from self.__wrapped__.read(*args, **kwargs)  # noqa: E999
             span.set_tag('Length', len(result))
@@ -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