엔지니어 대한민국

제가 얼마전에 asyncio를 이용해서 동시다발적으로 task를 진행하는 코드를 짰는데요. Python에서 제공하는 unittest랑 mock 라이브러리는요, asyncio를 테스트할수 있는 특별한 기능이 없더라구요. 그래서 인터넷검색을 하던중에 마침 저랑 같은 고민을 하고 있는 사람이 있길래 그분이 소개한 내용을 여러분께 전달해드리려고 책상에 앉았습니다. 원본링크는 맨 아래에 걸어두었으니까 참고하세요.

동기적인 코드 테스트하기

비동기적인 코드를 알려드리기에 앞서서 차이점을 설명드리기위해서 일단 동기적인 코드를 unit test하는걸 먼저 알고갈게요. 예를들어, 우리가 테스트해야하는 함수가 receive.py에 들어있는 receive()라는 함수라고 가정하고 아래코드를 봐주세요.

def receive(packet_type, packet_data):
    """Receive packet from the client."""
    if packet_type == 'PING':
        send_to_client("PONG", packet_data)
    elif packet_type == 'MESSAGE':
        response = trigger_event('message', packet_data)
        send_to_client('MESSAGE', response)
    else:
        raise ValueError('Invalid packer type')

def send_to_client(packet_type, packet_data):
    """Implementation of this function not show."""
    pass

def tirgger_event(event_name, event_data):
    """Implementation of this function not show."""
    pass

코드를 대충보시면 아시겠지만, 현재 이 코드는 어떤 메세지를 받아서 client에게 응답하는 메세지서버에요. client가 서버에 요청을 하면 receive를 호출하고요, 호출받은 packet의 type에 따라서 각각 send_to_client()를 호출해서 응답을 합니다. receive()를 unit test하는건 그리 어려워 보이지않아요. 그런데 send_to_client()는 응답 데이타를 외부 사용자에게 어떤 메세지를 전달하는 기능을 하는 함수라서, 외부 서비스에 의존해야하는데 아시다시피 unit test는 functional test와 달리 외부서비스에 의존하지 않고 테스트를 할수 있어야하자나요. 마찬가지로, trigger_event()도 어떤 이벤트를 발생시켜서 다른 시스템이 그 데이타를 받아서 처리하도록 하게하고 심지어 이 함수는 결과값을 반환하기까지 합니다. 그걸 받아서 client에 전달해야하는거죠. 다행히 여러가지 테스트 케이스중에 적어도 한개는 좀 쉬워보이네요. 바로 packet_type이 ‘PING’, ‘MESSAGE’ 이 두개중에 하나가 아니면 ValueError를 raise하도록 만들어져있자나요. 그러면 일단 쉬운거부터 만들면서 진행해볼게요:

import unittest
from receive import receive

class TestReceive(unittest.TestCase):
    def test_invalid_packet(self):
        self.assertRaises(ValueError, receive, 'FOO', 'data')

다른 두개는 실제로 외부서비스를 호출하지 않게 하기 위해서 mock을 이용해서 가짜 함수로 대체할게요.

import uniitest
from unittest import mock
from receive import receive

class TestReceive(unittest.TestCase):
    ...
    @mock.patch('receive.send_to_client')
    def test_ping(self, send_to_client):
        receive('PING', 'data')
        send_to_clinet.assert_called_once_with('PONG', 'data')

mock으로 함수를 대체하면 실제로 그 함수를 호출하지 않고 mocking된 가짜함수를 호출하게 되는거에요. 그리고 mocking된 가짜함수가 반환할 값들을 다양하게 조작하고 경우의 수를 만들어서 모든 케이스를 다 커버하도록 조작 하는거죠. mocking을 하기 위해서 위의 코드에서는 mock.patch데코레이터를 사용했어요. 그리고 외부로 서비스되는 테스트는 functional테스트에서 하도록 할건데요. 그부분은 다음에 다루도록 할게요. 이번시간에는 unit test에 집중해서 설명을 하도록 하겠습니다. 이렇게 patch를 한 function은 MagicMock으로 mocking이 되는데요. MagicMock이 어떻게 생겼는지 아래 코드에서 보여드릴게요:

>>> from unittest import mock
>>> f = mock.MagicMock()
>>> f()
<MagicMock name='mock()' id='4373365928'>
>>> f('hello', 'world')
<MagicMock name='mock()' id='4373365928'>
>>> f.some_method('foo')
<MagicMock name='mock.some_method()' id='4373489200'>

위에서 보시다시피 patch말고도 MagicMock을 직접 이용해서 mocking을 할수 있는데요 사실 patch 데코레이터안에서 MagicMock을 사용하고 있기 때문에 patch로 mocking을 하나, MagicMock으로 코딩을하나 결과는 비슷하다고 보시면 되요. 마지막 라인에서 보시면요, MagixMock으로 정의된 오브젝트에 서브함수를 정의해 넣고 있습니다. 이렇게 mocking을 당한 함수들은요 처리가 되는 내내 감시를 받게 됩니다. 아래와 같이 몇번 호출되었는지, 함수에 호출된 인자값이 제대로 잘 갔는지 그런거를 비교할수 있어요.

>>> f('hello', 'world')
<MagicMock name='mock()' id='4373526568'>
>>> f.assert_called_once_with('hello', 'world')
>>> f.assert_called_once_with('bye', 'world')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File ".../unittest/mock.py", line 825, in assert_called_once_with
    return self.assert_called_with(*args, **kwargs)
  File ".../unittest/mock.py", line 814, in assert_called_with
    raise AssertionError(_error_message()) from cause
AssertionError: Expected call: mock('bye', 'world')
Actual call: mock('hello', 'world')

자 그러면 이제 packet_type이 ‘MESSAGE’인 경우의 테스트 코드를 한번 만들어 볼까요?

class TestReceive(unittest.TestCase):
    # ...
    @mock.patch('receive.trigger_event', return_value='my response')
    @mock.patch('receive.send_to_client')
    def test_message(self, send_to_client, trigger_event):
        receive('MESSAGE', 'data')
        trigger_event.assert_called_once_with('message', 'data')
        send_to_client.assert_called_once_with('MESSAGE', 'my response')

trigger_event랑 send_to_client를 mocking해야겠죠? 그리고 trigger_event에서 받은 값을 반환해야하니까 가짜 반환값을 ‘my_response’라는 문자열로 주었습니다. 그리고 receive를 호출했을때, 각 함수들이 정확한 함수의 인자들을 가지고 호출이 되었는지를 확인하는거죠.

AsyncIO 테스트

이제 앞서만든 동기적인 코드를 비동기적으로 변환해볼게요. async_receive.py라는 파일을 하나 만드시고요 아래 코드들 저장합니다.

async def receive(packet_type, packet_data):
    """Receive packet from the client."""
    if packet_type == 'PING':
        await send_to_client("PONG", packet_data)
    elif packet_type == 'MESSAGE':
        response = await trigger_event('message', packet_data)
        await send_to_client('MESSAGE', response)
    else:
        raise ValueError('Invalid packet type')

async def send_to_client(packet_type, packet_data):
    """Implementation of this function not shown."""
    pass

async def trigger_event(event_name, event_data):
    """Implementation of this function not shown."""
    pass

보시다시피, 3개 함수 모두다 async로 정의했구요, receive가 trigger_event랑 send_to_client를 호출할때는 await으로 기다리는 동안 다른 task들이 실행되도록 작성했어요.

자, 이제 unit test코드를 만들어 볼까요? 가장 우선적으로 receive함수를 테스트 해야겠죠? receive함수는 async로 선언이 되었기때문에 함수를 호출하면 함수가 실행되는게 아니구요, 대신에 coroutine이라는 object가 만들어집니다.

>>> from async_receive import receive
>>> receive('FOO', 'data')
<coroutine object receive at 0x10e5b32b0>

위의 코드를 보시면 나는 분명히 receive함수를 불렀는데 coroutine이라는 오브젝트를 반환했죠? 이것은 coroutine이라는 객체가 receive함수를 들고 event loop라는데 들어가서 CPU가 해당 coroutine의 작업을 처리해주기를 기다려야하기 때문이에요. FOO를 packet type으로 넘기면 ValueError가 나도록 만들었는데 에러는 안나고 이상한 객체만 보여주고, 왠지 coroutine에 등록을 하고 나면 뭔가 제어할수 있는 권한이 내 손을 떠나버린 느낌이 들죠? 그래서 함수를 실행하는걸 직접 제어하려면 event loop을 만들어서 그 안에서 실행되도록 해야해요.

>>> import asyncio
>>> asyncio.get_event_loop().run_until_complete(receive('FOO', 'data'))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File ".../asyncio/base_events.py", line 466, in run_until_complete
    return future.result()
  File ".../async_receive.py", line 19, in receive
    raise ValueError('Invalid packet type')
ValueError: Invalid packet type

asyncio에서 제공하는 get_event_loop을 초기화해서 run_until_complete이라는 옵션으로 receive를 실행합니다. 그러면 이제 ValueError가 반환되는게 보이죠? 이걸 좀 간편하게 갖다 쓰기 위해서 저는 _run이라는 함수를 만들었어요.

import asyncio

def _run(coro):
    return asyncio.get_event_loop().run_until_complete(coro)

이 함수를 사용하면 테스트 코드가 좀더 깔끔해지죠

import unittest
from unittest import mock
from async_receive import receive

class TestReceive(unittest.TestCase):
    def test_invalid_packet(self):
        self.assertRaises(ValueError, _run, receive('FOO', 'data'))

ValueError를 반환하는 경우는 이렇게 처리가 가능하지만, 나머지 두가지 경우는 이거보다 좀더 복잡한데요. ‘PING’의 경우에는 send_to_client()를 mocking 해야되요. 위에 동기식 테스트에서 했던것처럼 말이에요. 그런데 여기서 비동기에서의 문제점이 하나 있는데요, receive함수에서 send_to_client를 호출할때 await으로 호출하고 있죠? 그런데 mock 객체는 await으로 실행될수가 없어요.

mock 객체를 사용할수 없다면, 도대체 어떻게해야만 send_to_client를 mocking할수 있을까요? send_to_client함수는 async함수니까 그냥 호출이 되면 아까 보여드린대로 coroutine 객체를 반환할거에요. 당연히 unit test중에 coroutine이 반환되는건 아무도 원치 않을거에요. 어차피 그 안에 함수는 우리가 제어할수도 없으니까요. 그렇다면 우리가 원하는건 바로 coroutine을 MagicMock을 이용해서 mocking하는 겁니다. 그러면 바로 우리가 원하는 대로 함수를 호출할수가 있어요.

갑자기 두통이 오시나요? 사실 저도 이거 고민하면서 두통이 생겼었어요. 결국 제가 찾은 방법은 또다른 helper함수를 통해서 coroutine을 mocking하자는 거였죠.

def AsyncMock(*args, **kwargs):
    m = mock.MagicMock(*args, **kwargs)

    async def mock_coro(*args, **kwargs):
        return m(*args, **kwargs)

    mock_coro.mock = m
    return mock_coro

위의 함수를 차근차근 한번 살펴볼게요. 함수 중간에 보시면 비동기로 선언되어있는 mock_coro라는 내부함수가 있죠? 얘가 AsyncMock에 넘겨진 함수의 인자를 고스란히 받게 됩니다. 맨 마지막 줄에 보시면 AsyncMock함수는 이 내부함수 mock_coro를 반환하게 되는데요. 제가 방금 위에서 말씀드린대로 실제로 돌아가는 모양새에 맞게 비동기함수를 mocking해야하는데요, 바로 이 mock_coro()함수가 그 역할을 해주는거에요. 이게 비동기로 선언이 되어있기때문에 가능한일입니다.

이 coroutine이 실행될때, 우리는 MagicMock객체를 만들고 싶은거자나요. 위의 AsyncMock()이 호출되면 가장먼저 m이 초기화되면서, 이 mocking된 object를 만드는거죠. MagicMock이 함수인자를 다 받아줄수 있어서 AsyncMock으로 받은 모든 인자를 m에 죄다 넘겨줄수가 있어요. 이때 넘겨주는 인자에는 return_value같은거도 다 넘어간답니다. 이렇게 선언된 m은 어떻게 활용되느냐, 바로 mock_coro함수에서 이걸 객체로 구현하는거에요. 바로 이때 모든 함수의 인자들이 실제 호출한 함수로 전달이 되는 순간인거죠.

정리를 좀 하자면요,
1) 일반 함수를 하나 선언하고
2) 내부변수 m에 MagicMock을 생성합니다. 이때, 일반 함수에서 받은 모든 인자를 MagicMock에 넘겨주어 mocking함수를 m에 할당합니다.
3) 내부에 비동기 함수를 하나 만듭니다. 그 함수는 호출될때 부모함수가 받은 모든 인자를 받게 됩니다. 이 내부함수는 방금 MagicMock으로 선언한 m을 호출한 결과를 반환합니다. 이때 일반 함수에서 받은 모든 인자를 그래도 m에 넘겨주도록합니다.
4) 내부함수를 바깥에서 호출할수 있게 함수자체를 return하려고 하는데 그전에 방금 만든 MigicMock을 할당한 m을 내부함수에 mock이라는 객체를 하나 붙여서 내부비동기함수.mock()이라고 호출하면 호출되도록 callable함수를 함수.mock()에 저장합니다.

이제 우리에게 마지막으로 필요한게 한가지 남았는데 바로 테스트코드에서 이 m객체를 사용하도록 만드는 방법이 필요해요. 이 객체가 함수 바깥에서도 접근해서 사용하도록 만들려면, 이걸 넘겨주는 뭔가가 있어야겠죠? 그래서 만든게 mock_coro함수인거죠. 지금 제가 하는게 이상하다고 생각하시는 분들이 계실까봐 잠시 설명을 드리자면요, Python에서 함수들은 전부 Object입니다. 그래서 여러분이 내부함수를 선언하시면요 그게 바로 해당 Object의 custom attribute이 되는거에요. 클래스로 치면 내부 메쏘드 정도라고 생각하시면 될것 같아요.

백문이 불여일견이라고요, 이렇게 우왕좌왕 설명을 드리는것보다 코드로 정리해서 직접 눈으로 확인시켜드리는게 훨씬 효과적일거 같다는 생각이 드네요. 아래 코드를 봐주세요.

>>> import asyncio
>>> from test_async_receive import AsyncMock
>>> f = AsyncMock(return_value='hello!')
>>> f('foo', 'bar')
<coroutine object AsyncMock.<locals>.mock_coro at 0x10ef84ca8>
>>> asyncio.get_event_loop().run_until_complete(f('foo', 'bar'))
'hello!'
>>> f.mock.assert_called_once_with('foo', 'bar')
>>> f.mock.assert_called_once_with('foo')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File ".../unittest/mock.py", line 825, in assert_called_once_with
    return self.assert_called_with(*args, **kwargs)
  File ".../unittest/mock.py", line 814, in assert_called_with
    raise AssertionError(_error_message()) from cause
AssertionError: Expected call: mock('foo')
Actual call: mock('foo', 'bar')

위의 예제에서 보시다시피요, async로 선언된 mock_coro를 반환하는 AsyncMock함수(=Object)를 invoke해서 f에 assign했죠? 만약에 제가 함수 f()를 invoke한다면, coroutine을 하나 갖게 되겠죠? 지금 우리가 만든게 실제 async함수가 동작하는 원리랑 똑같이 작동하도록 만든거에요.

f(‘foo’, ‘bar’)가 loop를 돌고 있다면, 기존에 mock에 설정해둔 return value를 받아오게 되는거죠. 아까 AsyncMock선언한 코드를 보시면요, 내부함수를 반환하기 전에 MagicMock으로 만든 m을 mock라는 변수에다 할당합니다. 그러니까 그 coroutine이 실제 holding하고 있는 task 즉, 그안에 MagicMock에 접근을 하려면 f가 아니라 f.mock으로 접근을 해야겠죠. 코드는 엄청 짧은데 정말 복잡한 코드에요. 근데 이게 동작은 정말 잘해요.

아래에 ‘PING’의 경우를 unit test해본 거에요.

class TestReceive(unittest.TestCase):
    # ...
    @mock.patch('async_receive.send_to_client', new=AsyncMock())
    def test_ping(self):
        _run(receive('PING', 'data'))
        from async_receive import send_to_client
        send_to_client.mock.assert_called_once_with('PONG', 'data')

여기에서 send_to_client를 AsyncMock으로 mocking합니다. 그게 비동기니까 send_to_client는 반환값이 아닌 coroutine을 반환할거니까 그 모양새를 비슷하게 만들어 놓은 AsyncMock를 구현해서 바꿔치기 하는거죠. 그리고 receive를 할때는 에러든 뭐든 일단 볼수 있게, 직접 호출하지 않고, Mock에서 제공하는 event_loop을 초기화해서 사용할수 있도록 _run()함수를 호출합니다.

아래는 ‘MESSAGE’의 경우를 테스트해본 결과인데요, 어딘가 뭔가 조금더 복잡해 진듯 하네요.

class TestReceive(unittest.TestCase):
    # ...
    @mock.patch('async_receive.send_to_client', new=AsyncMock())
    @mock.patch('async_receive.trigger_event', new=AsyncMock(return_value='my response'))
    def test_message(self):
        _run(receive('MESSAGE', 'data'))
        from async_receive import send_to_client, trigger_event
        trigger_event.mock.assert_called_once_with('message', 'data')
        send_to_client.mock.assert_called_once_with('MESSAGE', 'my response')

여기서는 message에 trigger_event가 반환한 값을 넘겨주어야하니까 AsyncMock을 호출할때 return_value도 함께 넘겨줍니다. 그러면 내부에 비동기로 선언한 함수를 호출할때도, 마찬가지로 해당 인자를 넘겨받아서 실행하게 됩니다.

아래는 필요한 함수와 코드들을 한번에 모아봤어요. 코드 전체다 보시고 싶으시면 https://github.com/miguelgrinberg/asyncio-testing.

async_receive.py

async def receive(packet_type, packet_data):
    """Receive packet from the client."""
    if packet_type == 'PING':
        await send_to_client("PONG", packet_data)
    elif packet_type == 'MESSAGE':
        response = await trigger_event('message', packet_data)
        await send_to_client('MESSAGE', response)
    else:
        raise ValueError('Invalid packet type')

async def send_to_client(packet_type, packet_data):
    """Implementation of this function not shown."""
    pass

async def trigger_event(event_name, event_data):
    """Implementation of this function not shown."""
    pass

test_async_receive.py

import unittest
from unittest import mock
from async_receive import receive
import asyncio

def _run(coro):
    return asyncio.get_event_loop().run_until_complete(coro)

def AsyncMock(*args, **kwargs):
    m = mock.MagicMock(*args, **kwargs)

    async def mock_coro(*args, **kwargs):
        return m(*args, **kwargs)

    mock_coro.mock = m
    return mock_coro

class TestReceive(unittest.TestCase):
    def test_invalid_packet(self):
        self.assertRaises(ValueError, _run, receive('FOO', 'data'))

    @mock.patch('async_receive.send_to_client', new=AsyncMock())
    def test_ping(self):
        _run(receive('PING', 'data'))
        from async_receive import send_to_client
        send_to_client.mock.assert_called_once_with('PONG', 'data')

    @mock.patch('async_receive.send_to_client', new=AsyncMock())
    @mock.patch('async_receive.trigger_event', new=AsyncMock(return_value='my response'))
    def test_message(self):
        _run(receive('MESSAGE', 'data'))
        from async_receive import send_to_client, trigger_event
        trigger_event.mock.assert_called_once_with('message', 'data')
        send_to_client.mock.assert_called_once_with('MESSAGE', 'my response')

테스트 결과

> pip install pytest
> pytest test_async_receive.py
======================================================================== test session starts ========================================================================
platform darwin -- Python 3.7.3, pytest-5.3.5, py-1.8.1, pluggy-0.13.1
rootdir: /Users/damazzang/Desktop
collected 3 items

test_async_receive.py ...                                                                                                                                     [100%]

========================================================================= 3 passed in 0.06s =========================================================================

부디 제 글이 조금이나마 비동기 테스트를 하시는데 도움이 되셨길 바랍니다. 혹시 번역이 부족하면 아래 original article을 참고하세요.

Source: https://blog.miguelgrinberg.com/post/unit-testing-asyncio-code