pax_global_header00006660000000000000000000000064136115410770014517gustar00rootroot0000000000000052 comment=e729a717e5d252f783f999b7318807133d666e2a Skyscanner-aiotask-context-9855ef4/000077500000000000000000000000001361154107700173255ustar00rootroot00000000000000Skyscanner-aiotask-context-9855ef4/.gitignore000066400000000000000000000020251361154107700213140ustar00rootroot00000000000000# Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] *$py.class # C extensions *.so # Distribution / packaging .Python env/ build/ develop-eggs/ dist/ downloads/ eggs/ .eggs/ lib/ lib64/ parts/ sdist/ var/ *.egg-info/ .installed.cfg *.egg # PyInstaller # Usually these files are written by a python script from a template # before PyInstaller builds the exe, so as to inject date/other infos into it. *.manifest *.spec # Installer logs pip-log.txt pip-delete-this-directory.txt # Unit test / coverage reports htmlcov/ .tox/ .coverage .coverage.* .cache nosetests.xml coverage.xml *,cover .hypothesis/ # Translations *.mo *.pot # Django stuff: *.log local_settings.py # Flask stuff: instance/ .webassets-cache # Scrapy stuff: .scrapy # Sphinx documentation docs/_build/ # PyBuilder target/ # IPython Notebook .ipynb_checkpoints # pyenv .python-version # celery beat schedule file celerybeat-schedule # dotenv .env # virtualenv venv/ ENV/ # Spyder project settings .spyderproject # Rope project settings .ropeproject Skyscanner-aiotask-context-9855ef4/.travis.yml000066400000000000000000000003051361154107700214340ustar00rootroot00000000000000language: python python: - "3.5" - "3.6" matrix: include: - python: 3.7 dist: xenial sudo: true install: pip install tox-travis flake8 script: tox - make syntax - tox Skyscanner-aiotask-context-9855ef4/LICENSE000066400000000000000000000020531361154107700203320ustar00rootroot00000000000000MIT License Copyright (c) 2016 Skyscanner Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. Skyscanner-aiotask-context-9855ef4/Makefile000066400000000000000000000002411361154107700207620ustar00rootroot00000000000000syntax: flake8 --exclude=examples/,.tox,.eggs cov: pytest --cov-report term-missing --cov=aiotask_context -sv tests/test_context.py test: pytest tests/ -sv Skyscanner-aiotask-context-9855ef4/README.md000066400000000000000000000266661361154107700206240ustar00rootroot00000000000000# Aiotask Context [![Build Status](https://travis-ci.org/Skyscanner/aiotask-context.svg?branch=master)](https://travis-ci.org/Skyscanner/aiotask-context) Store context information within the [asyncio.Task](https://docs.python.org/3/library/asyncio-task.html#task) object. For more information about why this package was developed, please read the blog post [From Flask to aiohttp](http://codevoyagers.com/2016/09/01/from-flask-to-aiohttp/). Supports both asyncio and uvloop loops. ## Installation ```bash pip install aiotask_context ``` ## Usage This package allows to store context information inside the [asyncio.Task](https://docs.python.org/3/library/asyncio-task.html#task) object. A typical use case for it is to pass information between coroutine calls without the need to do it explicitly using the called coroutine args. What this package is **NOT** for: - Don't fall into the bad pattern of storing everything your services need inside, this should only be used for objects or data that is needed by all or almost all the parts of your code where propagating it through args doesn't scale. - The context is a `dict` object so you can store any object you want inside. This opens the door to using it to change variables inside in the middle of an execution so other coroutines behave differently or other dirty usages. This is really **discouraged**. Now, a (simplified) example where you could apply this: In your application, to share the `request_id` between all the calls, you should do the following: ```python import asyncio async def my_coro_1(request_id): print(request_id) async def my_coro_2(request_id): await my_coro_1(request_id) async def my_coro_3(): request_id = "1234" await my_coro_2(request_id) if __name__ == '__main__': loop = asyncio.get_event_loop() loop.run_until_complete(my_coro_3()) ``` As you can see, this code smells a bit and feels like repeating yourself a lot (think about this example as if you had current API running in a framework and you needed the `request_id` everywhere to log it properly). With `aiotask_context` you can do: ```python import asyncio import aiotask_context as context async def my_coro_1(): print(context.get("request_id", default="Unknown")) async def my_coro_2(): print(context.get("request_id", default="Unknown")) await my_coro_1() async def my_coro_3(): context.set(key="request_id", value="1234") await my_coro_2() if __name__ == '__main__': loop = asyncio.get_event_loop() loop.set_task_factory(context.task_factory) loop.run_until_complete(my_coro_3()) ``` It also keeps the context between the calls like `ensure_future`, `wait_for`, `gather`, etc. That's why you have to change the [task factory](https://docs.python.org/3/library/asyncio-eventloop.html#asyncio.AbstractEventLoop.set_task_factory): ```python import asyncio import aiotask_context as context async def my_coro_0(): print("0: " + context.get("status")) async def my_coro_1(): context.set("status", "DONE") async def my_coro_2(): context.set("status", "RUNNING") print("2: " + context.get("status")) await asyncio.gather(asyncio.ensure_future(my_coro_1()), my_coro_0()) print("2: " + context.get("status")) if __name__ == '__main__': loop = asyncio.get_event_loop() loop.set_task_factory(context.task_factory) # This is the relevant line loop.run_until_complete(my_coro_2()) ``` You may also want to only keep a copy of the context between calls. For example, you have one task that spawns many others and do not want to reflect changes in one task's context into the other tasks. To do this you have two options: - `copying_task_factory` - uses a brand new copy of the context `dict` for *each* task - `chainmap_task_factory` - uses a `ChainMap` instead of a `dict` as context, which allows some of the values to be redefined while not creating a full copy of the context The following example yields the same results with both: ```python import asyncio import aiotask_context as context async def my_coro_0(): context.set("status", "FAILED") print("0: " + context.get("status")) async def my_coro_1(): context.set("status", "PENDING") print("1: " + context.get("status")) await my_coro_1_child() async def my_coro_1_child(): print("1 (child): " + context.get("status")) async def my_coro_spawner(): context.set("status", "RUNNING") print("2: " + context.get("status")) await asyncio.gather(asyncio.ensure_future(my_coro_1()), my_coro_0()) print("2: " + context.get("status")) if __name__ == '__main__': loop = asyncio.get_event_loop() for factory in (context.copying_task_factory, context.chainmap_task_factory): print('\nUsing', factory.__name__) loop.set_task_factory(factory) # This is the relevant line loop.run_until_complete(my_coro_spawner()) ``` ## Comparison of task factories / context types The difference between the task factories can be best illustrated by how they behave when inherited values are redefined or updated in child tasks. Consider: ```python import asyncio import aiotask_context as context async def my_coro_parent(loop): context.set("simple", "from parent") context.set("complex", ["from", "parent"]) print("parent before: simple={}, complex={}".format( context.get("simple"), context.get("complex"))) await loop.create_task(my_coro_child()) print("parent after: simple={}, complex={}".format( context.get("simple"), context.get("complex"))) async def my_coro_child(): context.set("simple", "from child") # redefine value completely context.get("complex")[1] = "child" # update existing object print("child: simple={}, complex={}".format( context.get("simple"), context.get("complex"))) if __name__ == '__main__': loop = asyncio.get_event_loop() for factory in (context.task_factory, context.copying_task_factory, context.chainmap_task_factory): print('\nUsing', factory.__name__) loop.set_task_factory(factory) loop.run_until_complete(my_coro_parent(loop)) ``` In this case the results are different for all three: ``` Using task_factory parent before: simple=from parent, complex=['from', 'parent'] child: simple=from child, complex=['from', 'child'] parent after: simple=from child, complex=['from', 'child'] Using copying_task_factory parent before: simple=from parent, complex=['from', 'parent'] child: simple=from child, complex=['from', 'child'] parent after: simple=from parent, complex=['from', 'parent'] Using chainmap_task_factory parent before: simple=from parent, complex=['from', 'parent'] child: simple=from child, complex=['from', 'child'] parent after: simple=from parent, complex=['from', 'child'] ``` ## Complete examples If you've reached this point it means you are interested. Here are a couple of complete examples with `aiohttp`: - Setting the `X-Request-ID` header and sharing it over your code: ```python """ POC to demonstrate the usage of the aiotask_context package for easily sharing the request_id in all your code. If you run this script, you can try to query with curl or the browser: $ curl http://127.0.0.1:8080/Manuel Hello, Manuel. Your request id is fdcde8e3-b2e0-4b9c-96ca-a7ce0c8749be. $ curl -H "X-Request-ID: myid" http://127.0.0.1:8080/Manuel Hello, Manuel. Your request id is myid. """ import uuid import asyncio import aiotask_context as context from aiohttp import web async def handle(request): name = request.match_info.get('name', "Anonymous") text = "Hello, {}. Your request id is {}.\n".format(name, context.get("X-Request-ID")) return web.Response(text=text) async def request_id_middleware(app, handler): async def middleware_handler(request): context.set("X-Request-ID", request.headers.get("X-Request-ID", str(uuid.uuid4()))) response = await handler(request) response.headers["X-Request-ID"] = context.get("X-Request-ID") return response return middleware_handler loop = asyncio.get_event_loop() loop.set_task_factory(context.task_factory) app = web.Application(middlewares=[request_id_middleware], loop=loop) app.router.add_route('GET', '/{name}', handle) web.run_app(app) ``` - Setting the request_id in all log calls: ```python """ POC to demonstrate the usage of the aiotask_context package for writing the request_id from aiohttp into every log call. If you run this script, you can try to query with curl or the browser: $ curl http://127.0.0.1:8080/Manuel Hello, Manuel. Your request id is fdcde8e3-b2e0-4b9c-96ca-a7ce0c8749be. $ curl -H "X-Request-ID: myid" http://127.0.0.1:8080/Manuel Hello, Manuel. Your request id is myid. In the terminal you should see something similar to: ======== Running on http://0.0.0.0:8080/ ======== (Press CTRL+C to quit) 2016-09-07 12:02:39,887 WARNING __main__:63 357ab21e-5f05-44eb-884b-0ce3ceebc1ce | First_call called 2016-09-07 12:02:39,887 ERROR __main__:67 357ab21e-5f05-44eb-884b-0ce3ceebc1ce | Second_call called 2016-09-07 12:02:39,887 INFO __main__:76 357ab21e-5f05-44eb-884b-0ce3ceebc1ce | Received new GET /Manuel call 2016-09-07 12:02:39,890 INFO aiohttp.access:405 357ab21e-5f05-44eb-884b-0ce3ceebc1ce | 127.0.0.1 - - [07/Sep/2016:10:02:39 +0000] "GET /Manuel HTTP/1.1" 200 70 "-" "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/52.0.2743.116 Safari/537.36" """ import asyncio import uuid import logging.config import aiotask_context as context from aiohttp import web class RequestIdFilter(logging.Filter): def filter(self, record): record.request_id = context.get("X-Request-ID") return True LOG_SETTINGS = { 'version': 1, 'disable_existing_loggers': False, 'handlers': { 'console': { 'class': 'logging.StreamHandler', 'level': 'INFO', 'formatter': 'default', 'filters': ['requestid'], }, }, 'filters': { 'requestid': { '()': RequestIdFilter, }, }, 'formatters': { 'default': { 'format': '%(asctime)s %(levelname)s %(name)s:%(lineno)d %(request_id)s | %(message)s', }, }, 'loggers': { '': { 'level': 'DEBUG', 'handlers': ['console'], 'propagate': True }, } } logging.config.dictConfig(LOG_SETTINGS) logger = logging.getLogger(__name__) logger.addFilter(RequestIdFilter()) async def first_call(): logger.warning("First_call called") async def second_call(): logger.error("Second_call called") async def handle(request): name = request.match_info.get('name') await asyncio.gather(first_call()) await second_call() logger.info("Received new GET /{} call".format(name)) text = "Hello, {}. Your request id is {}.\n".format(name, context.get("X-Request-ID")) return web.Response(text=text) async def request_id_middleware(app, handler): async def middleware_handler(request): context.set("X-Request-ID", request.headers.get("X-Request-ID", str(uuid.uuid4()))) response = await handler(request) response.headers["X-Request-ID"] = context.get("X-Request-ID") return response return middleware_handler if __name__ == "__main__": loop = asyncio.get_event_loop() loop.set_task_factory(context.task_factory) app = web.Application(middlewares=[request_id_middleware], loop=loop) app.router.add_route('GET', '/{name}', handle) web.run_app(app) ``` Skyscanner-aiotask-context-9855ef4/aiotask_context/000077500000000000000000000000001361154107700225245ustar00rootroot00000000000000Skyscanner-aiotask-context-9855ef4/aiotask_context/__init__.py000066400000000000000000000100041361154107700246300ustar00rootroot00000000000000import asyncio import logging import sys from collections import ChainMap from copy import deepcopy from functools import partial PY37 = sys.version_info >= (3, 7) if PY37: def asyncio_current_task(loop=None): """Return the current task or None.""" try: return asyncio.current_task(loop) except RuntimeError: # simulate old behaviour return None else: asyncio_current_task = asyncio.Task.current_task logger = logging.getLogger(__name__) NO_LOOP_EXCEPTION_MSG = "No event loop found, key {} couldn't be set" def dict_context_factory(parent_context=None, copy_context=False): """A traditional ``dict`` context to keep things simple""" if parent_context is None: # initial context return {} else: # inherit context new_context = parent_context if copy_context: new_context = deepcopy(new_context) return new_context def chainmap_context_factory(parent_context=None): """ A ``ChainMap`` context, to avoid copying any data and yet preserve strict one-way inheritance (just like with dict copying) """ if parent_context is None: # initial context return ChainMap() else: # inherit context if not isinstance(parent_context, ChainMap): # if a dict context was previously used, then convert # (without modifying the original dict) parent_context = ChainMap(parent_context) return parent_context.new_child() def task_factory(loop, coro, copy_context=False, context_factory=None): """ By default returns a task factory that uses a simple dict as the task context, but allows context creation and inheritance to be customized via ``context_factory``. """ context_factory = context_factory or partial( dict_context_factory, copy_context=copy_context) task = asyncio.tasks.Task(coro, loop=loop) if task._source_traceback: del task._source_traceback[-1] try: context = asyncio_current_task(loop=loop).context except AttributeError: context = None task.context = context_factory(context) return task def copying_task_factory(loop, coro): """ Returns a task factory that copies a task's context into new tasks instead of sharing it. :param loop: The active event loop :param coro: A coroutine object :return: A context-copying task factory """ return task_factory(loop, coro, copy_context=True) def chainmap_task_factory(loop, coro): """ Returns a task factory that uses ``ChainMap`` as the context. :param loop: The active event loop :param coro: A coroutine object :return: A ``ChainMap`` task factory """ return task_factory(loop, coro, context_factory=chainmap_context_factory) def get(key, default=None): """ Retrieves the value stored in key from the Task.context dict. If key does not exist, or there is no event loop running, default will be returned :param key: identifier for accessing the context dict. :param default: None by default, returned in case key is not found. :return: Value stored inside the dict[key]. """ current_task = asyncio_current_task() if not current_task: raise ValueError(NO_LOOP_EXCEPTION_MSG.format(key)) return current_task.context.get(key, default) def set(key, value): """ Sets the given value inside Task.context[key]. If the key does not exist it creates it. :param key: identifier for accessing the context dict. :param value: value to store inside context[key]. :raises """ current_task = asyncio_current_task() if not current_task: raise ValueError(NO_LOOP_EXCEPTION_MSG.format(key)) current_task.context[key] = value def clear(): """ Clear the Task.context. :raises ValueError: if no current task. """ current_task = asyncio_current_task() if not current_task: raise ValueError("No event loop found") current_task.context.clear() Skyscanner-aiotask-context-9855ef4/examples/000077500000000000000000000000001361154107700211435ustar00rootroot00000000000000Skyscanner-aiotask-context-9855ef4/examples/aiohttp_request_id.py000066400000000000000000000024351361154107700254150ustar00rootroot00000000000000""" POC to demonstrate the usage of the aiotask_context package for easily sharing the request_id in all your code. If you run this script, you can try to query with curl or the browser: $ curl http://127.0.0.1:8080/Manuel Hello, Manuel. Your request id is fdcde8e3-b2e0-4b9c-96ca-a7ce0c8749be. $ curl -H "X-Request-ID: myid" http://127.0.0.1:8080/Manuel Hello, Manuel. Your request id is myid. """ import uuid import asyncio from aiohttp import web import aiotask_context as context def handle(request): name = request.match_info.get('name') text = "Hello, {}. Your request id is {}.\n".format(name, context.get("X-Request-ID")) return web.Response(body=text.encode('utf-8')) async def request_id_middleware(app, handler): async def middleware_handler(request): context.set("X-Request-ID", request.headers.get("X-Request-ID", str(uuid.uuid4()))) response = await handler(request) response.headers["X-Request-ID"] = context.get("X-Request-ID") return response return middleware_handler if __name__ == "__main__": loop = asyncio.get_event_loop() loop.set_task_factory(context.task_factory) app = web.Application(middlewares=[request_id_middleware]) app.router.add_route('GET', '/{name}', handle) web.run_app(app) Skyscanner-aiotask-context-9855ef4/examples/asyncio_context_propagation.py000066400000000000000000000012501361154107700273270ustar00rootroot00000000000000import asyncio import uuid import aiotask_context as context async def my_coro(n): print(str(n) + ": " + context.get("request_id")) async def my_coro_2(): context.set("request_id", str(uuid.uuid4())) await asyncio.gather( asyncio.ensure_future(my_coro(0)), asyncio.wait_for(my_coro(1), 1), asyncio.shield(asyncio.wait_for(my_coro(2), 1)), my_coro(3)) if __name__ == '__main__': loop = asyncio.get_event_loop() # this ensure asyncio calls like ensure_future, wait_for, etc. # propagate the context to the new task that is created loop.set_task_factory(context.task_factory) loop.run_until_complete(my_coro_2()) Skyscanner-aiotask-context-9855ef4/examples/log_request_id.py000066400000000000000000000063151361154107700245270ustar00rootroot00000000000000""" POC to demonstrate the usage of the aiotask_context package for writing the request_id from aiohttp into every log call. If you run this script, you can try to query with curl or the browser: $ curl http://127.0.0.1:8080/Manuel Hello, Manuel. Your request id is fdcde8e3-b2e0-4b9c-96ca-a7ce0c8749be. $ curl -H "X-Request-ID: myid" http://127.0.0.1:8080/Manuel Hello, Manuel. Your request id is myid. In the terminal you should see something similar to: ======== Running on http://0.0.0.0:8080/ ======== (Press CTRL+C to quit) 2016-09-07 12:02:39,887 WARNING __main__:63 357ab21e-5f05-44eb-884b-0ce3ceebc1ce | First_call called 2016-09-07 12:02:39,887 ERROR __main__:67 357ab21e-5f05-44eb-884b-0ce3ceebc1ce | Second_call called 2016-09-07 12:02:39,887 INFO __main__:76 357ab21e-5f05-44eb-884b-0ce3ceebc1ce | Received new GET /Manuel call 2016-09-07 12:02:39,890 INFO aiohttp.access:405 357ab21e-5f05-44eb-884b-0ce3ceebc1ce | 127.0.0.1 - - [07/Sep/2016:10:02:39 +0000] "GET /Manuel HTTP/1.1" 200 70 "-" "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/52.0.2743.116 Safari/537.36" """ import uuid import logging.config import asyncio from aiohttp import web import aiotask_context as context class RequestIdFilter(logging.Filter): def filter(self, record): record.request_id = context.get("X-Request-ID") return True LOG_SETTINGS = { 'version': 1, 'disable_existing_loggers': False, 'handlers': { 'console': { 'class': 'logging.StreamHandler', 'level': 'INFO', 'formatter': 'default', 'filters': ['requestid'], }, }, 'filters': { 'requestid': { '()': RequestIdFilter, }, }, 'formatters': { 'default': { 'format': '%(asctime)s %(levelname)s %(name)s:%(lineno)d %(request_id)s | %(message)s', }, }, 'loggers': { '': { 'level': 'DEBUG', 'handlers': ['console'], 'propagate': True }, } } logging.config.dictConfig(LOG_SETTINGS) logger = logging.getLogger(__name__) logger.addFilter(RequestIdFilter()) async def first_call(): logger.warning("First_call called") async def second_call(): logger.error("Second_call called") async def handle(request): name = request.match_info.get('name') await first_call() await second_call() logger.info("Received new GET /{} call".format(name)) text = "Hello, {}. Your request id is {}.\n".format(name, context.get("X-Request-ID")) return web.Response(body=text.encode('utf-8')) async def request_id_middleware(app, handler): async def middleware_handler(request): context.set("X-Request-ID", request.headers.get("X-Request-ID", str(uuid.uuid4()))) response = await handler(request) response.headers["X-Request-ID"] = context.get("X-Request-ID") return response return middleware_handler if __name__ == "__main__": loop = asyncio.get_event_loop() loop.set_task_factory(context.task_factory) app = web.Application(middlewares=[request_id_middleware]) app.router.add_route('GET', '/{name}', handle) web.run_app(app) Skyscanner-aiotask-context-9855ef4/requirements-ci.txt000066400000000000000000000001521361154107700232000ustar00rootroot00000000000000pytest==3.8.0 pytest-randomly==1.2.3 pytest-asyncio==0.9.0 pytest-cov==2.6.0 flake8==3.5.0 uvloop==0.11.2 Skyscanner-aiotask-context-9855ef4/setup.cfg000066400000000000000000000001201361154107700211370ustar00rootroot00000000000000[aliases] test=pytest [pep8] max-line-length=100 [flake8] max-line-length=100 Skyscanner-aiotask-context-9855ef4/setup.py000066400000000000000000000011621361154107700210370ustar00rootroot00000000000000from setuptools import setup install_requires = [] tests_require = install_requires + ['pytest'] setup( name="aiotask_context", version="0.6.1", author="Manuel Miranda", author_email="manu.mirandad@gmail.com", description="Store context information inside the asyncio.Task object", classifiers=[ "Programming Language :: Python", "Programming Language :: Python :: 3.5", "Programming Language :: Python :: 3.6", "Programming Language :: Python :: 3.7", ], packages=['aiotask_context'], install_requires=install_requires, tests_require=tests_require, ) Skyscanner-aiotask-context-9855ef4/tests/000077500000000000000000000000001361154107700204675ustar00rootroot00000000000000Skyscanner-aiotask-context-9855ef4/tests/conftest.py000066400000000000000000000014571361154107700226750ustar00rootroot00000000000000import pytest import asyncio import uvloop import aiotask_context as context @pytest.fixture() def asyncio_loop(): asyncio.set_event_loop_policy(None) loop = asyncio.get_event_loop() yield loop loop.close() # restore the virgin state asyncio.set_event_loop_policy(None) @pytest.fixture() def uvloop_loop(): asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) loop = asyncio.get_event_loop() yield loop loop.close() # restore the virgin state asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) @pytest.fixture(params=[ 'asyncio_loop', 'uvloop_loop' ]) def event_loop(request): return request.getfixturevalue(request.param) @pytest.fixture(autouse=True) def context_loop(event_loop): event_loop.set_task_factory(context.task_factory) Skyscanner-aiotask-context-9855ef4/tests/test_acceptance.py000066400000000000000000000046011361154107700241670ustar00rootroot00000000000000import asyncio import random import pytest import uuid from collections import defaultdict import aiotask_context as context @asyncio.coroutine def dummy3(): yield from asyncio.sleep(random.uniform(0, 2)) return context.get("key") @asyncio.coroutine def dummy2(a, b): yield from asyncio.sleep(random.uniform(0, 2)) res = context.get("key") yield from asyncio.sleep(random.uniform(0, 2)) res1 = yield from dummy3() assert res == res1 return a, b, res @asyncio.coroutine def dummy1(n_tasks): context.set("key", str(uuid.uuid4())) tasks = [ asyncio.ensure_future( dummy2(id(context.asyncio_current_task()), n)) for n in range(n_tasks)] results = yield from asyncio.gather(*tasks) info = defaultdict(list) for taskid, n, key in results: info[key].append([taskid, n]) return info @pytest.mark.asyncio @asyncio.coroutine def test_ensure_future_concurrent(): n_tasks = 10 results = yield from asyncio.gather(*[dummy1(n_tasks=n_tasks) for x in range(1000)]) for r in results: assert len(r) == 1 for key, value in r.items(): assert len(value) == n_tasks @pytest.mark.asyncio @asyncio.coroutine def test_ensurefuture_context_propagation(): context.set("key", "value") @asyncio.coroutine def change_context(): assert context.get("key") == "value" context.set("key", "what") context.set("other", "data") yield from asyncio.ensure_future(change_context()) assert context.get("key") == "what" assert context.get("other") == "data" @pytest.mark.asyncio @asyncio.coroutine def test_waitfor_context_propagation(): context.set("key", "value") @asyncio.coroutine def change_context(): assert context.get("key") == "value" context.set("key", "what") context.set("other", "data") yield from asyncio.wait_for(change_context(), 1) assert context.get("key") == "what" assert context.get("other") == "data" @pytest.mark.asyncio @asyncio.coroutine def test_gather_context_propagation(): context.set("key", "value") @asyncio.coroutine def change_context(): assert context.get("key") == "value" context.set("key", "what") context.set("other", "data") yield from asyncio.gather(change_context()) assert context.get("key") == "what" assert context.get("other") == "data" Skyscanner-aiotask-context-9855ef4/tests/test_context.py000066400000000000000000000067241361154107700235750ustar00rootroot00000000000000import pytest import asyncio import traceback import aiotask_context as context @asyncio.coroutine def dummy(t=0): yield from asyncio.sleep(t) return True class TestSetGetClear: @pytest.mark.asyncio @asyncio.coroutine def test_get_ok(self): context.set("key", "value") assert context.get("key") == "value" @pytest.mark.asyncio @asyncio.coroutine def test_get_missing_key(self): context.set("random", "value") assert context.get("key", "default") == "default" assert context.get("key") is None @pytest.mark.asyncio @asyncio.coroutine def test_get_missing_context(self): assert context.get("key", "default") == "default" assert context.get("key") is None @pytest.mark.asyncio @asyncio.coroutine def test_set(self): context.set("key", "value") context.set("key", "updated_value") assert context.get("key") == "updated_value" def test_get_without_loop(self): with pytest.raises(ValueError): context.get("key") def test_set_without_loop(self): with pytest.raises(ValueError): context.set("random", "value") def test_loop_bug_aiohttp(self, event_loop): assert event_loop.run_until_complete(dummy()) is True asyncio.set_event_loop(None) assert event_loop.run_until_complete(dummy()) is True def test_closed_loop(self, event_loop): event_loop.close() with pytest.raises(RuntimeError): context.task_factory(event_loop, dummy()) @pytest.mark.asyncio @asyncio.coroutine def test_clear(self): context.set("key", "value") assert context.get("key") == "value" context.clear() assert context.get("key") is None def test_clear_without_loop(self): with pytest.raises(ValueError): context.clear() class TestTaskFactory: @pytest.mark.asyncio async def test_propagates_context(self, event_loop): context.set('key', 'value') task = context.task_factory(event_loop, dummy()) task.cancel() assert task.context == {'key': 'value'} @pytest.mark.asyncio async def test_sets_empty_context(self, event_loop): task = context.task_factory(event_loop, dummy()) task.cancel() assert task.context == {} @pytest.mark.asyncio async def test_sets_traceback(self, event_loop): event_loop.set_debug(True) task = context.task_factory(event_loop, dummy()) task.cancel() assert isinstance(task._source_traceback, traceback.StackSummary) @pytest.mark.asyncio async def test_propagates_copy_of_context(self, event_loop): @asyncio.coroutine def adds_to_context(): context.set('foo', 'bar') return True context.set('key', 'value') task = context.copying_task_factory(event_loop, adds_to_context()) await task assert task.context == {'key': 'value', 'foo': 'bar'} assert context.get('foo') is None @pytest.mark.asyncio async def test_propagates_chainmap_context(self, event_loop): @asyncio.coroutine def adds_to_context(): context.set('foo', 'bar') return True context.set('key', 'value') task = context.chainmap_task_factory(event_loop, adds_to_context()) await task assert task.context == {'key': 'value', 'foo': 'bar'} assert context.get('foo') is None Skyscanner-aiotask-context-9855ef4/tox.ini000066400000000000000000000010361361154107700206400ustar00rootroot00000000000000[tox] skipsdist = True envlist = setup py35 py36 py37 report [testenv:setup] deps = coverage setenv = COVERAGE_FILE = .coverage commands = coverage erase [testenv] setenv = COVERAGE_FILE = .coverage.{envname} deps = coverage pytest commands = pip install -e . pip install -r requirements-ci.txt coverage run --source aiotask_context -m pytest [testenv:report] deps = coverage setenv = COVERAGE_FILE = .coverage commands = coverage combine coverage report -m coverage xml