pytest-asyncio를 사용하여 비동기 ORM 서비스 테스트 코드 작성하기

sqlalchemy의 경우 async_engine이라 불리는 비동기 데이터베이스 엔진을 지원한다.

이 엔진은 asyncpg와 같은 비동기 데이터베이스 드라이버를 내부적으로 사용하여, 데이터베이스와의 상호작용을 비동기적으로 처리한다.

티켓 예매 시스템이라 가정하고 비동기 데이터베이스 엔진을 기반으로 아래와 같이 엔티티, 리포지토리, 서비스를 작성해보았다.

엔티티

User, Concert, Seat, Ticket, Reservation 의 여러 엔티티가 있다.

아래 코드는 엔티티에 대한 작성 예시다. 스압으로 사용자 엔티티 세부 정보에 대해서만 표시하였다.

# 사용자 엔티티
class User(Base):
    __tablename__ = "user"

    id = Column(Integer, primary_key=True, index=True)
    username = Column(String(100), nullable=False)
    email = Column(String(100), nullable=False, unique=True)
    created_at = Column(DateTime(timezone=True), default=func.now())

	# 사용자는 여러 예약을 가질 수 있다
    reservations = relationship("Reservation", back_populates="user")

# 콘서트 엔티티
class Concert(Base):
...
    seats = relationship("Seat", back_populates="concert")
    tickets = relationship("Ticket", back_populates="concert")
...

# 좌석 엔티티
class Seat(Base):
...
    concert = relationship("Concert", back_populates="seats")
    ticket = relationship("Ticket", back_populates="seat")
...

class Ticket(Base):
...
    concert = relationship("Concert", back_populates="tickets")
    seat = relationship("Seat", back_populates="ticket", uselist=False)
    reservation = relationship("Reservation", back_populates="ticket", uselist=False)
...

# 예약 엔티티
class Reservation(Base):
...
    user = relationship("User", back_populates="reservations")
    ticket = relationship("Ticket", back_populates="reservation")
...

리포지토리

데이터베이스와 상호작용하는 리포지토리를 아래와 같이 작성하였다.

데이터베이스와 상호작용하기 위한 비동기 세션 객체는 생성자를 통해 주입되도록 설정하였다.

class UserRepository:
    def __init__(self, session: AsyncSession):
        self.session = session

    async def get_user_by_id(self, user_id: int):
        result = await self.session.execute(select(User).filter(User.id == user_id))
        return result.scalars().first()

    async def create_user(self, username: str, email: str):
        user = User(username=username, email=email)
        self.session.add(user)
        await self.session.flush()
        await self.session.refresh(user)
        return user


class ConcertRepository:
    def __init__(self, session: AsyncSession):
        self.session = session

    async def get_concert_by_id(self, concert_id: int):
...

    async def create_concert(self, name: str, event_date):
...

    async def create_seats(self, concert_id: int, total_seats: int):
...


class TicketRepository:
    def __init__(self, session: AsyncSession):
        self.session = session

    async def get_ticket_by_id(self, ticket_id: int):
...

    async def create_ticket(self, concert_id: int, seat_id: int, price: float):
...

    async def update_ticket_status(self, ticket_id: int, status: TicketStatus):
...


class ReservationRepository:
    def __init__(self, session: AsyncSession):
        self.session = session

    async def create_reservation(self, user_id: int, ticket_id: int):
...

    async def update_reservation_status(self, reservation_id: int, status: ReservationStatus):
...

서비스

아래는 서비스 계층의 코드이다.

from sqlalchemy.ext.asyncio import AsyncSession
from app.repository import (
    UserRepository, ConcertRepository, TicketRepository, ReservationRepository,
    ReservationStatus, TicketStatus
)
from app.util import transactional


class UserService:
    def __init__(self, session: AsyncSession):
        self.session = session
        self.user_repo = UserRepository(session)

    @transactional
    async def register_user(self, username: str, email: str):
        return await self.user_repo.create_user(username, email)

    async def get_user(self, user_id: int):
        return await self.user_repo.get_user_by_id(user_id)


class ConcertService:
    def __init__(self, session: AsyncSession):
        self.session = session
        self.concert_repo = ConcertRepository(session)

    async def get_seats(self, concert_id: int):
        return await self.concert_repo.get_seats_by_concert_id(concert_id=concert_id)

    @transactional
    async def create_concert(self, name: str, event_date):
        return await self.concert_repo.create_concert(name, event_date)

    @transactional
    async def create_concert_seats(self, concert_id: int, total_seats: int):
        return await self.concert_repo.create_seats(concert_id, total_seats)


class TicketService:
    def __init__(self, session: AsyncSession):
        self.session = session
        self.ticket_repo = TicketRepository(session)

    async def get_ticket_by_id(self, ticket_id: int):
        return await self.ticket_repo.get_ticket_by_id(ticket_id)

    @transactional
    async def create_ticket(self, concert_id: int, seat_id: int, price: float):
        return await self.ticket_repo.create_ticket(concert_id, seat_id, price)

    @transactional
    async def update_ticket_status(self, ticket_id: int, status):
        return await self.ticket_repo.update_ticket_status(ticket_id, status)


class ReservationService:
    def __init__(self, session: AsyncSession):
        self.session = session
        self.reservation_repo = ReservationRepository(session)
        self.ticket_service = TicketService(session)

    @transactional
    async def reserve_ticket(self, user_id: int, ticket_id: int):
        # TicketService를 통해 티켓 상태 확인 및 상태 업데이트
        ticket = await self.ticket_service.get_ticket_by_id(ticket_id)
        if ticket and ticket.status == TicketStatus.available:
            # 예약 생성
            reservation = await self.reservation_repo.create_reservation(user_id, ticket_id)
            # 티켓을 'reserved' 상태로 업데이트
            await self.ticket_service.update_ticket_status(ticket_id, TicketStatus.reserved)
            return reservation
        return None

    @transactional
    async def confirm_reservation(self, reservation_id: int):
        # 예약 상태를 'confirmed'로 업데이트
        reservation = await self.reservation_repo.update_reservation_status(reservation_id, ReservationStatus.confirmed)
        if reservation:
            # 티켓을 'sold' 상태로 업데이트
            await self.ticket_service.update_ticket_status(reservation.ticket_id, TicketStatus.sold)
        return reservation

transactional

위 코드를 보면 transactional 데코레이터를 사용하여 트랜잭션을 적용하고자 하는 구간을 정의하고 제어한다.

해당 데코레이터는 트랜잭션을 수동으로 처리하는 반복적인 코드를 피하고, 기존 함수의 수정 없이 트랜잭션 처리 로직을 확장하기 위해 작성하였다.

이를 통해, 함수 호출 시 자동으로 트랜잭션이 시작되고, 성공적으로 완료되면 커밋하며, 예외가 발생하면 롤백한다.

데코레이터가 중첩되는 경우도 고려하여 in_transaction 메서드를 통해 체크하고 중첩된 경우 기존 트랜잭션을 사용하도록 하였다.

코드는 아래와 같다.

def transactional(func):
    async def wrapper(self, *args, **kwargs):
        session = getattr(self, 'session', None)

        if session is None:
            raise ValueError("Session not provided")

        if session.in_transaction():
            return await func(self, *args, **kwargs)

        async with session.begin():
            try:
                return await func(self, *args, **kwargs)
            except Exception:
                await session.rollback()
                raise

    return wrapper

서비스 레이어까지의 작성은 모두 마무리 되었다.
그렇다면 비동기 엔진을 기반으로 데이터베이스와 상호작용하는 시스템의 서비스 레이어는 어떻게 테스트 해야할까 ?

테스트 코드 작성

우선 나는 여러 테스트 프레임워크 중에서도 가장 대중적이고 간결하게 테스트 작성이 가능한 pytest를 사용하였다.

비동기 함수를 테스트 하기전 사전 설치해야할 패키지들이 있다.

우선 비동기 함수 테스트를 위해 pytest-asyncio 패키지는 필수로 설치해야한다.

pip install pytest-asyncio

추가로 pytest를 수행할때 pytest.ini에 별도의 환경변수를 적용하기위해 pytest-env도 설치하였다.

pip install pytest-env

설치 후 아래와 같이 서비스 레이어를 테스트 하기위한 테스트 코드를 작성하였다.
비동기 함수를 테스트하기 위해서는 pytest.mark.asyncio 데코레이터를 사용해야한다.
그리고 테스트 함수에서 사용되는 픽스처를 정의할때는 pytest_asyncio.fixture로 정의되어야한다.

아래는 작성한 테스트 코드의 일부이다.

async def get_session():
    async with session_factory() as session:
        yield session
        
        
@pytest_asyncio.fixture
async def session():
    async for s in get_session():
        await s.begin()
        try:
            yield s
        except Exception as e:
            await s.rollback()
            raise e
        finally:
            await s.rollback()


# UserService 테스트
@pytest.mark.asyncio
async def test_사용자_등록(session: AsyncSession):
    user_service = UserService(session)
    user = await user_service.register_user(username="john_doe", email="john@example.com")

    assert user.username == "john_doe"
    assert user.email == "john@example.com"


@pytest.mark.asyncio
async def test_사용자정보_얻기(session: AsyncSession):
    user_service = UserService(session)
    created_user = await user_service.register_user(username="john_doe", email="john@example.com")
    user = await user_service.get_user(created_user.id)

    assert user.id == created_user.id
    assert user.username == "john_doe"
    assert user.email == "john@example.com"
    
...

위 테스트 코드를 실행해보았다.
그런데 첫번째 테스트 통과이후 나머지 테스트 케이스들이 전부 실패하면서 아래와 같은 에러가 발생했다.

에러.. ?

RuntimeError: Task <Task pending name='Task-6' coro=<test_사용자정보_얻기() running at /Users/soon/ticket_booking_system/tests/app/service_test.py:36> cb=[_run_until_complete_cb() at /opt/homebrew/Cellar/python@3.12/3.12.7/Frameworks/Python.framework/Versions/3.12/lib/python3.12/asyncio/base_events.py:182]> got Future <Future pending cb=[Protocol._on_waiter_completed()]> attached to a different loop

내 질문의 수준이 낮아서 그런건지 몰라도 GPT 형님한테 물어봐도 잘 모르더라..

왜 에러가 발생했는지 도무지 감이 잡히지 않아 pytest가 비동기 함수를 테스트 할때 어떻게 동작하는지에 대한 파악이 필요하다고 생각했다.

pytest는 기본적으로 테스트가 순차적으로 실행되는 특징을 가졌다.
즉, 하나의 테스트가 완료되어야 다른 테스트가 실행된다.

이는 pytest.mark.asyncio 데코레이터를 사용하여 테스트할 때도 동일하다.

pytest.mark.asyncio 데코레이터를 통해서 scope를 설정할 수 있는데 scope로 지정한 범위에 따라서 이벤트 루프 생명주기가 달라진다.

예를 들면 scope를 function으로 지정할 경우 테스트 함수마다 새로운 이벤트 루프를 사용하고 class로 지정하면 클래스 마다 새로운 이벤트 루프를 사용하는 식이다.

# 테스트 함수 실행 시 새로운 이벤트 루프를 생성하여 실행
@pytest.mark.asyncio(scope="function")
async def test_function_scope():
    await asyncio.sleep(1)

이벤트 루프

이벤트 루프는 비동기 작업을 스케줄링하고 관리하여, 태스크들이 대기 중일 때 다른 작업을 실행할 수 있도록 하는 비동기 프로그래밍의 메커니즘이다.

Task는 비동기 함수(비동기 코루틴)을 관리하는 단위이다.

이벤트 루프의 흐름
1. Task가 생성되면 Task 큐에 적재되고 Ready 상태가 된다.
2. Event loop는 Task 큐에서 Task 1을 꺼내 실행하고 await 구문을 만나게 되면 네트워크 I/O, 파일 I/O와 같은 비동기 작업을 os 또는 스레드 풀로 요청한다.
3. 비동기 작업 요청이후 Task 1는 Block상태가 되며 비동기 작업이 완료될때까지 대기한다.
3-1. 그리고 바로 Task 2가 실행된다.
4. Task 1에서 실행된 비동기 작업이 끝나면 이벤트 루프로 신호가 오고 이벤트 루프는 신호를 받으면 Task 1을 Ready상태로 전환하고 Task 큐에 다시 적재한다.

위의 흐름으로 비동기 로직이 돌아간다.

이해를 위해 아래 예제를 한번보자.

async def task1():
    print("Task 1 실행")
    await asyncio.sleep(3)  # 3초 대기
    print("Task 1 완료")

async def task2():
    print("Task 2 실행")
    await asyncio.sleep(2)  # 2초 대기
    print("Task 2 완료")

async def task3():
    print("Task 3 실행")
    await asyncio.sleep(1)  # 1초 대기
    print("Task 3 완료")

async def main():
    tasks = [
        asyncio.create_task(task1()),
        asyncio.create_task(task2()),
        asyncio.create_task(task3())
    ]
    await asyncio.gather(*tasks)

asyncio.run(main())

asyncio.run(main())의 경우 이벤트 루프를 생성하고 실행하는 코드이다.
이벤트 루프가 생성되고 task1, task2, task3 코루틴이 Task로 래핑되어 이벤트 루프에 등록된다.
그럼 비동기 작업의 시간에 따라서 task3, task2, task1 순으로 신호를 받고 완료되는 것을 볼 수 있다.

실행 결과:

Task 1 실행
Task 2 실행
Task 3 실행
Task 3 완료
Task 2 완료
Task 1 완료

테스트 코드 분석

자 이제 다시 pytest로 돌아와서 비동기 메커니즘을 바탕으로 아래 테스트 코드의 동작 흐름을 분석해보자.

@pytest_asyncio.fixture
async def session():
    async for s in get_session():
        await s.begin()
        try:
            yield s
        except Exception as e:
            await s.rollback()
            raise e
        finally:
            await s.rollback()


# UserService 테스트
@pytest.mark.asyncio
async def test_사용자_등록(session: AsyncSession):
    user_service = UserService(session)
    user = await user_service.register_user(username="john_doe", email="john@example.com")

    assert user.username == "john_doe"
    assert user.email == "john@example.com"
    
@pytest.mark.asyncio
async def test_사용자정보_얻기(session: AsyncSession):
    user_service = UserService(session)
    created_user = await user_service.register_user(username="john_doe", email="john@example.com")
    user = await user_service.get_user(created_user.id)

    assert user.id == created_user.id
    assert user.username == "john_doe"
    assert user.email == "john@example.com"

pytest.mark.asyncio 데코레이터의 경우 파라미터로 아래와 같이 scope을 제공할 수 있다.

@pytest.mark.asyncio(scope="function")
async def test_func_1(module_fixture):
    print(f"테스트 함수 1 실행: {module_fixture}")
    await asyncio.sleep(1)

설정하지 않을 경우 기본 값은 function인데, 이는 테스트 함수가 실행될 때마다 새로운 이벤트 루프를 생성하여 테스트 함수를 실행시키겠다는 의미다.

session으로 값을 지정할 경우 pytest 세션이 시작될때 이벤트 루프를 생성하고 그것을 여러 테스트 함수를 실행할때마다 재사용하겠다는 의미이다.

테스트의 경우 독립적인 상태를 유지하고 각 테스트 케이스마다 순수한 상태에서 실행하는 게 적합하다고 판단했기에 현재 function scope을 적용해도 문제가 없다고 판단했다.

"그래 여기까지는 일단 문제가 없는 것 같다.."

자 그렇다면 이제 pytest 비동기 픽스처를 한번 보자.

pytest 비동기 픽스처

pytest 비동기 픽스처 사용시 asyncio_default_fixture_loop_scope 설정을 하도록 되어 있는데 이 역시 나는 function을 적용한 상태였다.

이때 비동기 픽스처의 scope에 따라서 테스트마다 픽스처를 재생성할건지, 아니면 세션에서 한번만 픽스처를 생성할건지 scope에 따라 픽스처의 생명주기가 결정된다.

참고로 픽스처와 테스트 함수는 동일한 이벤트 루프를 사용한다.

이 역시도 그다지 문제가 되진 않는 듯하다.

에러의 이유가 뭘까

픽스처도 테스트 함수도 그다지 문제가 되는 부분은 보이지 않는다.
그래서 에러를 다시 한번 자세히 읽어보았다.

RuntimeError: Task <Task pending name='Task-6' coro=<test_사용자정보_얻기() running at /Users/soon/ticket_booking_system/tests/app/service_test.py:36> cb=[_run_until_complete_cb() at /opt/homebrew/Cellar/python@3.12/3.12.7/Frameworks/Python.framework/Versions/3.12/lib/python3.12/asyncio/base_events.py:182]> got Future <Future pending cb=[Protocol._on_waiter_completed()]> attached to a different loop

에러 내용을 그대로 읽어보면 아래와 같다.

"Task-6이 이벤트 루프에서 실행되는 동안 다른 이벤트 루프에 연결된 Future 객체를 받았다."

테스트 함수마다 이벤트 루프가 별도 실행되고 픽스처로 새로 생성되는 데 도대체 어떻게 다른 이벤트 루프로 부터 객체를 받을 수가 있단 말인가..

"그렇다면 혹시 픽스처는 새로 생성되더라도 픽스처에서 동일한 자원을 계속 참조하는 부분이 있지않을까"

비동기 데이터베이스 엔진과 관련하여 검색을 하다보니 하나의 단서를 발견하게 되었다.

"비동기 데이터베이스 엔진은 커넥션 풀을 기본으로 활성화하고 각 커넥션을 재사용한다. 그리고 쿼리 전송이나 I/O 작업 같은 비동기 작업은 Future 객체로 관리되며 이는 커넥션과 연결되어 이러한 작업들은 비동기 데이터베이스 엔진에서 관리한다."

이를 통해서 "아 이전 이벤트 루프에서 사용된 커넥션이 새로 생성된 이벤트 루프에서 다시 재사용되면서 이전 커넥션과 연결되어 이전 이벤트 루프에서 생성된 Future 객체가 현재 이벤트 루프와 충돌하면서 에러가 발생했구나"라는 것을 알 수 있었다.

(Future 객체는 특정 이벤트 루프에 의존하며 다른 이벤트 루프와는 호환되지 않는다.)

참고: 이벤트 루프가 종료되었음에도 아직도 Future객체가 메모리에 남아 있는 이유는 커넥션 객체가 참조하고 있었기 때문이다.

이제 원인을 알았으니 해결방법을 아래와 같이 쉽게 정리하였다.

"매 테스트마다 커넥션 풀에 남아있는 커넥션들을 다 정리하면 되겠다."

아래는 수정한 코드이다.

@pytest_asyncio.fixture
async def session():
    async for s in get_session():
        await s.begin()
        try:
            yield s
        except Exception as e:
            await s.rollback()
            raise e
        finally:
            await s.rollback()

    await async_engine.dispose()

async_engine.dispose()를 호출하여 픽스처가 종료될때마다 커넥션 풀을 정리해주었다.

이제 테스트가 매우 잘 통과한다 !

정리

pytest를 비동기 함수를 테스트하고 분석하는 과정에서 비동기 메커니즘에 대한 이해도가 많이 올라갔다.
단순히 문제를 해결하기 위한 과정이었음에도 상당히 많은 인사이트와 배움이 있었던 것 같다.
비동기 패러다임에 대해서는 사실 감각적으로만 알고 있었는데 정리하는 시간을 통해 자세히 알게되어 매우 유익한 시간이었다.