Search code examples
pythoncommand-line-interfacepytestpython-asynciopython-click

How can I test async click commands from an async pytest function?


I'm trying to test a click command that is async from pytest, but I am hitting the limits of my knowledge of asyncio (or approaching the problem with a wrong architecture)

On one side, I have a click command line, that creates a grpclib channel to hit a grpc api.

import asyncio
from grpclib import Channel
from functools import wraps

def async_cmd(func):
  @wraps(func)
  def wrapper(*args, **kwargs):
    return asyncio.run(func(*args, **kwargs))
  return wrapper

@click.command
@async_cmd
async def main():
  async with Channel('127.0.0.1', 1234) as channel:
    blah = await something(channel)
    do_stuff_with(blah)
  return 0

Now I'm trying to test things using pytest and pytest-asyncio:

from click.testing import CliRunner
from cli import main
from grpclib.testing import ChannelFor
import pytest

@pytest.mark.asyncio
async def test_main()
  async with ChannelFor([Service()]) as test_channel:
    # Plan is to eventually mock grpclib.Channel with test_channel here.
    runner = CliRunner()
    runner.invoke(main)

My issue is that the async_cmd around main expects to call asyncio.run. But by the time the test_main method is called, a loop is already running (launched by pytest).

What should I do?

  • Should I modify my wrapper to join an existing loop (and how so?).
  • Should I mock something somewhere?
  • Should I just change my code do have my main just responsible for parsing the arguments and calling another function?

Solution

  • You are running your own event loop in the async_cmd decorator with this:

    asyncio.run(func(*args, **kwargs))
    

    Therefore, it is not apparent that you need to use @pytest.mark.asyncio, I suggest trying your testing without it.

    If you need an async context manager for a Mock, you can init the context manager in a hook called via the mock as shown below in test_hook().

    Test Code (for the test code)

    import asyncio
    import click
    import functools as ft
    import pytest
    import time
    from unittest import mock
    from click.testing import CliRunner
    
    
    class AsyncContext():
    
        def __init__(self, delay):
            self.delay = delay
    
        async def __aenter__(self):
            await asyncio.sleep(self.delay)
            return self.delay
    
        async def __aexit__(self, exc_type, exc, tb):
            await asyncio.sleep(self.delay)
    
    TestAsyncContext = AsyncContext
    
    def async_cmd(func):
        @ft.wraps(func)
        def wrapper(*args, **kwargs):
            return asyncio.run(func(*args, **kwargs))
        return wrapper
    
    
    @click.command()
    @async_cmd
    async def cli():
        async with TestAsyncContext(0.5) as delay:
            await asyncio.sleep(delay)
        print('hello')
    
    
    @pytest.mark.parametrize('use_mock, min_time, max_time',
                             ((True, 2.5, 3.5), (False, 1.0, 2.0)))
    def test_async_cli(use_mock, min_time, max_time):
        def test_hook(delay):
            return AsyncContext(delay + 0.5)
    
        runner = CliRunner()
        start = time.time()
        if use_mock:
            with mock.patch('test_code.TestAsyncContext', test_hook):
                result = runner.invoke(cli)
        else:
            result = runner.invoke(cli)
        stop = time.time()
        assert result.exit_code == 0
        assert result.stdout == 'hello\n'
        assert min_time < stop - start < max_time
    

    Test Results

    ============================= test session starts =============================
    collecting ... collected 2 items
    
    test_code.py::test_async_cli[True-2.5-3.5] 
    test_code.py::test_async_cli[False-1.0-2.0] 
    
    ============================== 2 passed in 4.57s ==============================