При подготовке материала для курса, нам периодически попадаются интересные статьи, которыми хотелось бы поделиться с вами!

Автор Stefan Scherfke “Testing (asyncio) coroutines with pytest”



PyTest — отличный пакет для тестирования на Python, и с давних пор один из моих любимых пакетов в целом. Он значительно облегчает написание тестов и обладает широкими возможностями по составлению отчетов о непройденных тестах.

Тем не менее, на момент версии 2.7, он менее эффективен в тестировании (asyncio) подпрограмм. Поэтому не стоит пытаться их тестировать таким способом:

# tests/test_coros.py

import asyncio

def test_coro():
    loop = asyncio.get_event_loop()

    @asyncio.coroutine
    def do_test():
        yield from asyncio.sleep(0.1, loop=loop)
        assert 0  # onoes!

    loop.run_until_complete(do_test())

В таком методе много недостатков и излишеств. Единственные интересные строки — те, что содержат операторы yield from и assert.

Каждый тесткейс должен иметь свой событийный цикл, который завершается корректно вне зависимости от успешности прохождения теста.

Применить yield вышеописанным способом не получится, pytest решит что наш тест возвращает новые тесткейсы.

Поэтому необходимо создать отдельную подпрограмму, в которой содержится фактический тест, а для его выполнения запускается событийный цикл.
Тесты будут чище, если сделать вот так:

# tests/test_coros.py

@asyncio.coroutine
def test_coro(loop):
    yield from asyncio.sleep(0.1, loop=loop)
    assert 0

Pytest обладает гибкой системой плагинов, благодаря чему имплементация такого поведения возможна. Но, к сожалению, большая часть нужных хуков задокументирована плохо или совсем никак, поэтому узнавать способы их выполнения проблематично.

Локальный попапочный плагин создаётся просто потому что это чуть проще, чем создавать “настоящий” внешний плагин. Pytest находит в каждой тестовой директории файл с названием conftest.py и применяет фикстуры и хуки, имплементированные в нём, ко всем тестам этой директории.

Начнём с написания фикстуры, которая создаёт новый событийный цикл для каждого тесткейса и закрывает его корректно по окончании теста:

# tests/conftest.py

import asyncio
import pytest


@pytest.yield_fixture
def loop():
    # Настройка
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    yield loop

    # Очистка
    loop.close()


# tests/test_coros.py

def test_coro(loop):
    @asyncio.coroutine
    def do_test():
        yield from asyncio.sleep(0.1, loop=loop)
        assert 0  # onoes!

    loop.run_until_complete(do_test())
 

Перед каждым тестом pytest выполняет фикстуру loop до первого оператора yield. То, что возвращает yield, передаётся в качестве аргумента loop (то есть цикла) нашего тесткейса. Когда тест окончен (успешно или нет), pytest завершает выполнение фикстуры цикла, закрывая его корректно. Таким же образом можно написать тестовую фикстуру, которая создаёт сокет и закрывает его после каждого теста (фикстура сокета может зависеть от фикстуры цикла точно так же, как в нашем примере. Круто, не правда ли?)

Но до конца ещё далеко. Нужно научить pytest выполнять наши тестовые подпрограммы. Для этого нужно изменить, как asyncio подпрограммы собираются (они должны собираться как обычные тестовые функции, а не как тестовые генераторы) и как они выполняются (с помощью loop.run_until_complete()):

# tests/conftest.py

def pytest_pycollect_makeitem(collector, name, obj):
    """Собери asyncio подпрограммы как обычные функции, а не генераторы.”
    if collector.funcnamefilter(name) and asyncio.iscoroutinefunction(obj):
        # Мы возвращаем список объектов тестовой функции 
        # В зависимости от фикстур, одна функция тестирования может вернуть несколько тестовых объектов
        # например, когда оно декорировано "pytest.mark.parametrize()".
        return list(collector._genfunctions(name, obj))
    # else:
    #    Мы возвращаем None, и дефолтное поведение pytest’а применяется к “obj”


def pytest_pyfunc_call(pyfuncitem):
    """Если ``pyfuncitem.obj`` - функция asyncio подпрограммы, выполнить её с помощью событийного цикла вместо прямого обращения"""
    testfunction = pyfuncitem.obj

    if not asyncio.iscoroutinefunction(testfunction):
        # Вернуть None, если это не подпрограмма. Pytest будет обращаться
        # с тестовым объектом обычным способом
        return

    # Вернуть нужные аргументы из всех доступных фикстур:
    funcargs = pyfuncitem.funcargs  # словарь со всеми фикстурами
    argnames = pyfuncitem._fixtureinfo.argnames  # аргументы для этого теста
    testargs = {arg: funcargs[arg] for arg in argnames}

    # Создать объект-генератор для этого теста (тест пока не выполнен!)
    coro = testfunction(**testargs)

    # Запустить подпрограмму в событийном цикле и выполнить тест
    loop = testargs['loop'] if loop in testargs else asyncio.get_event_loop()
    loop.run_until_complete(coro)

    return True  # Сообщить pytest’у, что мы выполнили тест
 

Этот плагин работает в pytest версии 2.4 и выше. Я проверял его работоспособность с версиями 2.6 и 2.7.

И всё бы было хорошо, но вскоре после опубликования данного решения в Stack Overflow, появился плагин PyTest-Asyncio, но Стефан нисколько не расстроился, а сделал подробный разбор этого плагина

Продвинутое тестирование асинхронного кода

В своей первой статье я показал, как pytest помогает писать качественные тесты. Фикстуры позволяют создавать чистый событийный цикл для каждого тесткейса, а благодаря системе плагинов можно писать тестовые функции, которые на самом деле являются сопрограммами asyncio.

Но пока велась работа над этим материалом, Тин Твртковиц создал плагин pytest-asyncio.

Если коротко, он позволяет делать так:

import asyncio
import time

import pytest

@pytest.mark.asyncio
async def test_coro(event_loop):
    before = time.monotonic()
    await asyncio.sleep(0.1, loop=event_loop)
    after = time.monotonic()
    assert after - before >= 0.1
 

Вместо:

import asyncio
import time

def test_coro():
    loop = asyncio.new_event_loop()
    try:
        asyncio.set_event_loop(loop)

        before = time.monotonic()
        loop.run_until_complete(asyncio.sleep(0.1, loop=loop))
        after = time.monotonic()
        assert after - before >= 0.1
    finally:
        loop.close()
 

Использование pytest-asyncio наглядно улучшает тест (и это не предел возможностей плагина!).

Когда я работал над aiomas, столкнулся с дополнительными требованиями, выполнить которые было не так просто.

Немного о том, что представляет собой сам aiomas. Он добавляет три слоя абстракции asyncio транспортам:

• Слой channel позволяет отправлять сообщения в формате JSON или MsgPack по схеме “запрос-ответ”. Он использует кастомный протокол, работающий с разными типами транспортов: TCP сокетами, сокетами домена Unix и кастомным транспортом — local queue.
• Слой RPC создаёт систему удалённого вызова процедур (Remote Procedure Call) над channel слоем.
• Слой agent (для мультиагентных систем) содержит ещё больше деталей, связанных с сетью, и позволяет писать классы, которые вызывают методы других классов через сетевое подключение.

Простейший пример того, как работает слой channel:

import aiomas


async def handle_client(channel):
    """Выполнить подключение"""
    req = await channel.recv()
    print(req.content)
    await req.reply('cya')
    await channel.close()


async def client():
    """Клиентская сопрограмма: Отправить приветствие серверу и дожидаться ответа"""
    channel = await aiomas.channel.open_connection(('localhost', 5555))
    rep = await channel.send('ohai')
    print(rep)
    await channel.close()


server = aiomas.run(aiomas.channel.start_server(
    ('localhost', 5555), handle_client))
aiomas.run(client())

server.close()
aiomas.run(server.wait_closed())
 

Требования к тестам

Нужен чистый событийный цикл для каждого теста.

Это можно сделать с помощью фикстуры event_loop, которая есть в pytest-asyncio.
Каждый тест должен быть запущен с каждым возможным транспортом (TCP сокет, сокет домена Unix, ...).

Теоретически это можно решить с помощью декоратора pytest.mark.parametrize() (но не в нашем случае, как дальше это станет ясно).

Каждому тесту нужна клиентская сопрограмма. В идеале, сам тест.

Декоратор pytest.mark.asyncio в pytest-asyncio справляется с этой задачей.

Каждому тесту необходим сервер с кастомным колбеком для клиентских соединений. По завершении теста сервера должны быть выключены вне зависимости от результата тестирования.

Кажется, что сопрограмма может решить эту задачу, но каждому серверу необходим конкретный колбек для управления клиентскими соединениями. Что усложняет решение проблемы. Не хочется получать ошибки “Address already in use”, если один из тестов не выполнится. Фикстура unused_tcp_port в pytest-asyncio спешит на помощь.

Не хочется постоянно использовать loop.run_until_complete().

И декоратор pytest.mark.asyncio решает задачу.

Обобщим, что осталось решить: каждому тесту нужно две фикстуры (одна для событийного цикла и еще одна для типа адреса), но я хочу объединить их в одну. Нужно создать фикстуру для настройки сервера, но как это сделать?

Первый подход

Можно обернуть цикл и тип адреса в фикстуру. Назовём её ctx (сокращенно от test context). Благодаря параметрам фикстуры легко создать отдельную на каждый тип адреса.

import tempfile
import py
import pytest


class Context:
    def __init__(self, loop, addr):
        self.loop = loop
        self.addr = addr


@pytest.fixture(params=['tcp', 'unix'])
def ctx(request, event_loop, unused_tcp_port, short_tmpdir):
    """Сгенерировать тесты с TCP сокетами и сокетами домена Unix."""
    addr_type = request.param
    if addr_type == 'tcp':
        addr = ('127.0.0.1', unused_tcp_port)
    elif addr_type == 'unix':
        addr = short_tmpdir.join('sock').strpath
    else:
        raise RuntimeError('Unknown addr type: %s' % addr_type)

    ctx = Context(event_loop, addr)
    return ctx


@pytest.yield_fixture()
def short_tmpdir():
    """Сгенерировать временную директорию для сокетов домена Unix. Пути, предоставленные фикстурой pytest’а tmpdir, слишком длинные на некоторых платформах"""
    with tempfile.TemporaryDirectory() as tdir:
        yield py.path.local(tdir)
 

Это позволяет написать тесты так:

import aiomas

@pytest.mark.asyncio
async def test_channel(ctx):
    results = []

    async def handle_client(channel):
        req = await channel.recv()
        results.append(req.content)
        await req.reply('cya')
        await channel.close()


    server = await aiomas.channel.start_server(ctx.addr, handle_client)
    try:
        channel = await aiomas.channel.open_connection(ctx.addr)
        rep = await channel.send('ohai')
        results.append(rep)
        await channel.close()

    finally:
        server.close()
        await server.wait_closed()

    assert results == ['ohai', 'cya']
 

Работает уже неплохо, и каждый тест, использующий фикстуру ctx, запускается единожды для каждого типа адреса.

Тем не менее, остаётся две проблемы:

1. Фикстура всегда требует неиспользуемый TCP порт + временную директорию — несмотря на то что нам понадобится только одно из двух.
2. Настройка сервера (и его закрытие) включает в себя некоторое количество кода, который будет идентичен для всех тестов, и поэтому должен быть включён в фикстуру. Но она не будет работать напрямую, потому что каждому серверу нужен тест-зависимый колбек (это видно в строке, где мы создаём сервер server = await ...). Но без серверной фикстуры нет и teardown’а для него…

Посмотрим, как можно решить эти проблемы.

Второй подход

Первая проблема решается использованием метода getfuncargvalue(), принадлежащего объекту request, который получает наша фикстура. Этим методом можно вручную вызвать её функцию:

@pytest.fixture(params=['tcp', 'unix'])
def ctx(request, event_loop):
    """Сгенерировать тесты с TCP сокетами и сокетами домена Unix."""
    addr_type = request.param
    if addr_type == 'tcp':
        port = request.getfuncargvalue('unused_tcp_port')
        addr = ('127.0.0.1', port)
    elif addr_type == 'unix':
        tmpdir = request.getfuncargvalue('short_tmpdir')
        addr = tmpdir.join('sock').strpath
    else:
        raise RuntimeError('Unknown addr type: %s' % addr_type)

    ctx = Context(event_loop, addr)
    return ctx

Для решения второй проблемы можно расширить класс Context, который передаётся в каждый тест. Добавляем метод Context.start_server(client_handler), который можно вызывать прямо из тестов. А также добавляем завершающий teardown в нашу ctx фикстуру, которая закроет сервер после окончания. Помимо всего прочего, можно создать несколько функций для шорткатов.

import asyncio
import tempfile
import py
import pytest


class Context:
    def __init__(self, loop, addr):
        self.loop = loop
        self.addr = addr
        self.server = None

    async def connect(self, **kwargs):
        """Создать и вернуть соединение "self.addr"."""
        return (await aiomas.channel.open_connection(
            self.addr, loop=self.loop, **kwargs))

    async def start_server(self, handle_client, **kwargs):
        """Запустить сервер с колбеком *handle_client*, слушающим
        "self.addr"."""
        self.server = await aiomas.channel.start_server(
            self.addr, handle_client, loop=self.loop, **kwargs)

    async def start_server_and_connect(self, handle_client,
                                       server_kwargs=None,
                                       client_kwargs=None):
        """Шорткат для::

            await ctx.start_server(...)
            channel = await ctx.connect()"

        """
        if server_kwargs is None:
            server_kwargs = {}

        if client_kwargs is None:
            client_kwargs = {}

        await self.start_server(handle_client, **server_kwargs)
        return (await self.connect(**client_kwargs))

    async def close_server(self):
        """Закрыть сервер."""
        if self.server is not None:
            server, self.server = self.server, None
            server.close()
            await server.wait_closed()


@pytest.yield_fixture(params=['tcp', 'unix'])
def ctx(request, event_loop):
    """Сгененировать тесты с TCP сокетами и сокетами домена Unix."""
    addr_type = request.param
    if addr_type == 'tcp':
        port = request.getfuncargvalue('unused_tcp_port')
        addr = ('127.0.0.1', port)
    elif addr_type == 'unix':
        tmpdir = request.getfuncargvalue('short_tmpdir')
        addr = tmpdir.join('sock').strpath
    else:
        raise RuntimeError('Unknown addr type: %s' % addr_type)

    ctx = Context(event_loop, addr)

    yield ctx

    # Выключить сервер и ожидать завершения всех нерешенных задач:
    aiomas.run(ctx.close_server())
    aiomas.run(asyncio.gather(*asyncio.Task.all_tasks(event_loop),
                              return_exceptions=True))

Тесткейс становится ощутимо короче, проще для восприятия и надёжней благодаря такому дополнительному функционалу:

import aiomas

@pytest.mark.asyncio
async def test_channel(ctx):
    results = []

    async def handle_client(channel):
        req = await channel.recv()
        results.append(req.content)
        await req.reply('cya')
        await channel.close()


    channel = await ctx.start_server_and_connect(handle_client)
    rep = await channel.send('ohai')
    results.append(rep)
    await channel.close()

    assert results == ['ohai', 'cya']

Фикстура ctx (и класс Context), конечно, не самая короткая, что я когда-либо писал, но она помогла избавить мои тесты примерно от 200 строчек шаблонного кода.

The end

Изящных решений и надёжного кода!