Fixes context propagation, adds context tests (#1369)

Author: terricain

aiobotocore/context.py

@@ -1,7 +1,6 @@
 from contextlib import asynccontextmanager
 from copy import deepcopy
 from functools import wraps
-from inspect import iscoroutinefunction
 
 from botocore.context import (
     ClientContext,
@@ -10,6 +9,8 @@
     set_context,
 )
 
+from ._helpers import resolve_awaitable
+
 
 @asynccontextmanager
 async def start_as_current_context(ctx=None):
@@ -31,10 +32,7 @@ def decorator(func):
         async def wrapper(*args, **kwargs):
             async with start_as_current_context():
                 if hook:
-                    if iscoroutinefunction(hook):
-                        await hook()
-                    else:
-                        hook()
+                    await resolve_awaitable(hook())
                 return await func(*args, **kwargs)
 
         return wrapper

aiobotocore/paginate.py

@@ -1,14 +1,23 @@
+from functools import partial
+
 import aioitertools
 import jmespath
 from botocore.exceptions import PaginationError
 from botocore.paginate import PageIterator, Paginator
+from botocore.useragent import register_feature_id
 from botocore.utils import merge_dicts, set_value_from_jmespath
 
+from .context import with_current_context
+
 
 class AioPageIterator(PageIterator):
     def __aiter__(self):
         return self.__anext__()
 
+    @with_current_context(partial(register_feature_id, 'PAGINATOR'))
+    async def _make_request(self, current_kwargs):
+        return await self._method(**current_kwargs)
+
     async def __anext__(self):
         current_kwargs = self._op_kwargs
         previous_next_token = None

aiobotocore/waiter.py

@@ -2,7 +2,6 @@
 from functools import partial
 
 # WaiterModel is required for client.py import
-from botocore.context import with_current_context
 from botocore.docs.docstring import WaiterDocstring
 from botocore.exceptions import ClientError
 from botocore.useragent import register_feature_id
@@ -19,6 +18,8 @@
     xform_name,
 )
 
+from .context import with_current_context
+
 
 def create_waiter_with_client(waiter_name, waiter_model, client):
     """

tests/botocore_tests/__init__.py

@@ -24,6 +24,7 @@
 from botocore.compat import parse_qs, urlparse
 
 import aiobotocore.session
+from aiobotocore.awsrequest import AioAWSResponse
 
 _LOADER = botocore.loaders.Loader()
 
@@ -93,3 +94,72 @@ def assert_url_equal(url1, url2):
     assert parts1.hostname == parts2.hostname
     assert parts1.port == parts2.port
     assert parse_qs(parts1.query) == parse_qs(parts2.query)
+
+
+class HTTPStubberException(Exception):
+    pass
+
+
+class BaseHTTPStubber:
+    class AsyncFileWrapper:
+        def __init__(self, body: bytes):
+            self._body = body
+
+        async def read(self):
+            return self._body
+
+    def __init__(self, obj_with_event_emitter, strict=True):
+        self.reset()
+        self._strict = strict
+        self._obj_with_event_emitter = obj_with_event_emitter
+
+    def reset(self):
+        self.requests = []
+        self.responses = []
+
+    def add_response(
+        self, url='https://example.com', status=200, headers=None, body=b''
+    ):
+        if headers is None:
+            headers = {}
+
+        response = AioAWSResponse(
+            url, status, headers, self.AsyncFileWrapper(body)
+        )
+        self.responses.append(response)
+
+    @property
+    def _events(self):
+        raise NotImplementedError('_events')
+
+    def start(self):
+        self._events.register('before-send', self)
+
+    def stop(self):
+        self._events.unregister('before-send', self)
+
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.stop()
+
+    def __call__(self, request, **kwargs):
+        self.requests.append(request)
+        if self.responses:
+            response = self.responses.pop(0)
+            if isinstance(response, Exception):
+                raise response
+            else:
+                return response
+        elif self._strict:
+            raise HTTPStubberException('Insufficient responses')
+        else:
+            return None
+
+
+class ClientHTTPStubber(BaseHTTPStubber):
+    @property
+    def _events(self):
+        return self._obj_with_event_emitter.meta.events

tests/botocore_tests/functional/__init__.py

@@ -0,0 +1,123 @@
+import asyncio
+
+from aiobotocore.session import AioSession
+
+from ...mock_server import AIOServer
+from .. import ClientHTTPStubber
+
+
+def get_captured_ua_strings(stubber):
+    """Get captured request-level user agent strings from stubber.
+    :type stubber: tests.BaseHTTPStubber
+    """
+    return [req.headers['User-Agent'].decode() for req in stubber.requests]
+
+
+def parse_registered_feature_ids(ua_string):
+    """Parse registered feature ids in user agent string.
+    :type ua_string: str
+    :rtype: list[str]
+    """
+    ua_fields = ua_string.split(' ')
+    feature_field = [field for field in ua_fields if field.startswith('m/')][0]
+    return feature_field[2:].split(',')
+
+
+async def test_user_agent_has_registered_feature_id():
+    session = AioSession()
+
+    async with (
+        AIOServer() as server,
+        session.create_client(
+            's3',
+            endpoint_url=server.endpoint_url,
+            aws_secret_access_key='xxx',
+            aws_access_key_id='xxx',
+        ) as s3_client,
+    ):
+        with ClientHTTPStubber(s3_client) as stub_client:
+            stub_client.add_response()
+            paginator = s3_client.get_paginator('list_buckets')
+            # The `paginate()` method registers `'PAGINATOR': 'C'`
+            async for _ in paginator.paginate():
+                pass
+
+        ua_string = get_captured_ua_strings(stub_client)[0]
+        feature_list = parse_registered_feature_ids(ua_string)
+        assert 'C' in feature_list
+
+
+async def test_registered_feature_ids_dont_bleed_between_requests():
+    session = AioSession()
+
+    async with (
+        AIOServer() as server,
+        session.create_client(
+            's3',
+            endpoint_url=server.endpoint_url,
+            aws_secret_access_key='xxx',
+            aws_access_key_id='xxx',
+        ) as s3_client,
+    ):
+        with ClientHTTPStubber(s3_client) as stub_client:
+            stub_client.add_response()
+            waiter = s3_client.get_waiter('bucket_exists')
+            # The `wait()` method registers `'WAITER': 'B'`
+            await waiter.wait(Bucket='mybucket')
+
+            stub_client.add_response()
+            paginator = s3_client.get_paginator('list_buckets')
+            # The `paginate()` method registers `'PAGINATOR': 'C'`
+            async for _ in paginator.paginate():
+                pass
+
+        ua_strings = get_captured_ua_strings(stub_client)
+        waiter_feature_list = parse_registered_feature_ids(ua_strings[0])
+        assert 'B' in waiter_feature_list
+
+        paginator_feature_list = parse_registered_feature_ids(ua_strings[1])
+        assert 'C' in paginator_feature_list
+        assert 'B' not in paginator_feature_list
+
+
+# This tests context's bleeding across tasks instead
+async def test_registered_feature_ids_dont_bleed_across_threads():
+    session = AioSession()
+
+    async with (
+        AIOServer() as server,
+        session.create_client(
+            's3',
+            endpoint_url=server.endpoint_url,
+            aws_secret_access_key='xxx',
+            aws_access_key_id='xxx',
+        ) as s3_client,
+    ):
+
+        async def wait():
+            with ClientHTTPStubber(s3_client) as stub_client:
+                stub_client.add_response()
+                waiter = s3_client.get_waiter('bucket_exists')
+                # The `wait()` method registers `'WAITER': 'B'`
+                await waiter.wait(Bucket='mybucket')
+            ua_string = get_captured_ua_strings(stub_client)[0]
+            return parse_registered_feature_ids(ua_string)
+
+        async def paginate():
+            with ClientHTTPStubber(s3_client) as stub_client:
+                stub_client.add_response()
+                paginator = s3_client.get_paginator('list_buckets')
+                # The `paginate()` method registers `'PAGINATOR': 'C'`
+                async for _ in paginator.paginate():
+                    pass
+            ua_string = get_captured_ua_strings(stub_client)[0]
+            return parse_registered_feature_ids(ua_string)
+
+        waiter_features, paginator_features = await asyncio.gather(
+            wait(), paginate()
+        )
+
+        assert 'B' in waiter_features
+        assert 'C' not in waiter_features
+        assert 'C' in paginator_features
+        assert 'B' not in paginator_features

tests/botocore_tests/functional/test_useragent.py

@@ -786,6 +786,12 @@ def test_protocol_parsers():
                 'a7e83728338e61ff2ca0a26c6f03c67cbabffc32',
             },
         ),
+        (
+            PageIterator._make_request,
+            {
+                'e926671018897ac5851a3add5d2bc15a2d6142df',
+            },
+        ),
         (
             PageIterator.result_key_iters,
             {