FastAPI 용 factory_boy 만들기

만들게 된 계기

Django에서 테스트 코드를 작성할 때 테스트용 객체를 생성해 주는 factory_boy 라이브러리를 사용해 좀 더 손쉽게 테스트코드를 작성할 수 있었다. 마찬가지로 FastAPI에서도 사용하기 위해 factory_boy의 SQLAlchemyModelFactory를 사용하고자 했으나 다음과 같은 문제가 있었다.

• session 동적 할당이 불가능하다.
  • 클래스를 정의하는 시점에 Meta 클래스 내 sqlalchemy_session에 session을 정의해주어야 한다.
  • 이 방식은 테스트코드 런타임에 session을 동적 할당하는 것이 불가능해서 session 기반의 데이터들이 테스트 전반에 영향을 미치는 문제가 있다.
  • 아래는 sqlalchemy_session에 session을 할당하는 예시이다.

from sqlalchemy import Column, Integer, Unicode, create_engine
from sqlalchemy.orm import declarative_base, scoped_session, sessionmaker

engine = create_engine('sqlite://')
session = scoped_session(sessionmaker(bind=engine))
Base = declarative_base()

import factory

class UserFactory(factory.alchemy.SQLAlchemyModelFactory):
    class Meta:
        model = User
        sqlalchemy_session = session   # the SQLAlchemy session object

    id = factory.Sequence(lambda n: n)
    name = factory.Sequence(lambda n: 'User %d' % n)

 

사실 위의 문제는 22년에 issues에 제기되었고 해결 방법으로 factory boy 3.3.0에 sqlalchemy_session_factory 변수가 추가됐다. 그러나 또 문제가 있다. (참고)

• sqlalchemy_session_factory 기능이 제대로 동작하지 않는다.
  • 위의 문제를 해결하기 위해 sqlalchemy_session_factory 를 사용하면 다음과 같이 에러가 발생한다. (참고)

AttributeError: type object 'Meta' has no attribute 'sqlalchemy_session'

 

이 이슈 또한 main에 머지된 상태이지만 아직 릴리즈 되지 않았다. (참고)

 

해당 이슈가 릴리즈 되기를 기다려볼 수도 있지만, 도전 의식이 생겨 필요한 기능만 가볍게 구현해 보기로 했다.

내가 필요한 기능은 다음과 같다.

• 테스트용 객체를 생성할 수 있을 것.
• 테스트용 객체에 필드 값을 넘겨줄 수 있을 것.
• faker를 사용하여 샘플 데이터를 넣을 수 있을 것.
• 쿼리 조회나 relationship 조회 등에 문제가 없을 것.
• 객체의 특성에 따라 데이터를 설정해 볼 수 있을 것. 예를 들어 종료 예정인 특성을 설정하면 end_at이 now() + 1 되는 식이 가능할 것.

구현해 보기

기본적인 아이디어는 다음과 같다.

• session을 동적 할당할 수 있도록 한다.

class AbstractFactory:
	
	def __init__(self, session, *args, **params):
		self.session = session

• 추상 클래스를 상속받아 pydantic 모델 별 Factory를 정의한다.
  • factory_boy와 동일하게 Meta 클래스 내 model을 정의해서 해당 모델의 fields 정보를 가져올 수 있도록 한다.

class AbstractFactory:
	
	class Meta:
		abstract = True
		model = None
		
	def __init__(self, session, *args, **params):
		self.session = session
		self.model = self.Meta.model
		self.fields = vars(self.model)
		
class UserFactory(AbstractFactory):

	class Meta:
		model = User

 

여기서 UserFactory처럼 AbstractFactory를 상속받아 정의하는 클래스를 모델 팩토리 클래스라고 부르겠다.

• 모델 팩토리 클래스에는 클래스 변수로 필드값을 미리 선언할 수 있다.

class UserFactory(AbstractFactory):
	name = "jeanie"

	class Meta:
		model = User

• name을 “jeanie”로 고정하지 않고 Faker를 사용해 랜덤으로 생성하기 위해서 LazyAttribute를 추가했다.

class UserFactory(AbstractFactory):
	name = LazyAttribute(lambda obj: fake.name())

 

모델 팩토리 클래스를 사용해 다음과 같이 객체를 생성한다.

user = UserFactory(session=session, name="jeanie").generate()

 

이제 UserFactory를 생성했을 때, 처리되는 일련의 과정을 살펴보자.

• UserFactory 넘겨주는 session 인자값과 session으로, 속성 값을 params로 받는다.

class AbstractFactory:
	def __init__(self, session, *args, **params):
		self.session = session
		self.params = params
		self.model = self.Meta.model
		self.fields = vars(self.model)

 

다음으로 generate 메서드를 호출하여 User 객체를 생성한다.

• 클래스 속성 값(_attrs)을 가져온다.
  • Factory 클래스의 다중상속을 고려하여 cls.__mro__ 를 가져와 상속된 클래스별 속성(attr)을 가져온다.
  • 만약 가져온 속성(attr)이 콜러블하면, 속성값이 아닌 메서드이므로 pass 한다.
  • private 변수가 아니면서 Meta 클래스가 아니면 속성값을 append 하여 리턴한다.

class AbstractFactory:
  ...
  
  @classmethod
  def _attrs(cls):
      attrs = []
      for parent in reversed(cls.__mro__):
          for v in vars(parent):
              attr = getattr(cls, v)
              if callable(attr) and not isinstance(attr, AbstractAttribute):
                  continue
              if not v.startswith("_") and v != "Meta":
                  attrs.append(v)
      return attrs

• _attrs와 params로 넘겨받은 속성 값을 업데이트한다.
  • 모델 필드가 클래스 변수로 선언되어 있다고 하더라도 인자로 넘겨받은 value를 반영해 주도록 한다.
  • 클래스 변수의 type이 LazyAttribute인 경우 호출하여, 그 값을 가져올 수 있도록 한다.

class AbstractFactory:
  ...
  
  def _build_params(self):
      build_params = {}

      for param, value in self.params.items():
          if param not in self.fields:
              raise Exception(f"Invalid params '{param}' on {self.model.__name__}")

          setattr(self, param, value)
          build_params[param] = value

      for attr in self._attrs():
          if attr not in self.fields:
              raise Exception(f"Invalid attrs '{attr}' on {self.model.__name__}")

          if attr in build_params:
              continue

          attr_v = getattr(self, attr)
          if isinstance(attr_v, LazyAttribute):
              build_params[attr] = attr_v(self)
              setattr(self, attr, build_params[attr])
          else:
              build_params[attr] = attr_v

      self._params = build_params

• User 객체를 생성한다.
  • 인자로 넘겨받은 session으로 커밋한다.
  • 만약 실패한다면 롤백하도록 한다.

class AbstractFactory:
	...
  def _generate(self):
    try:
        model_obj = self.model(**self._params)
        self.session.add(model_obj)
        self.session.commit()
        return model_obj
    except Exception as e:
        self.session.rollback()
        raise e

 

이로써 간단한 버전의 factory_boy를 구현했다. generate메서드의 리턴값이 model 객체이므로, model에 정의된 메서드나 relationship을 그대로 가져다 쓸 수 있다.

 

추가로 자주 쓰이는 특성을 하나의 코드로 묶고 싶을 수 있다. 이럴 때는 generate에 그 특성을 넘겨볼 수 있다.

예를 들어 탈퇴한 유저를 만든다고 하자.

class UserFactory(AbstractFactory):
	def generate(self, withdrew: bool = None):
		if withdrew:
			self._set_withdrew()
			
		user = super().generate()
		return user
	
	def _set_withdrew(self):
		self.params["status"] = "withdrew"
		self.params["deactivated_at"] = datetime.now()


withdrew_user = UserFactory(session=session).generate(withdrew=True)
assert withdrew_user.status == "withdrew"

 

마지막으로 pytest 함수 단위로 데이터를 생성하고 삭제하고 싶다. 다른 테스트 코드에서 생성된 데이터 때문에 테스트에 영향을 미치면 디버깅도 어렵고 테스트 코드끼리 서로 의존해 버리기 때문이다. 이를 위해 nested transaction을 사용한다. 아래 코드는 stackoverflow를 참고했다.

engine = create_engine(settings.database_url)
TestingSessionLocal = sessionmaker(autoflush=False, bind=engine)

@event.listens_for(engine, "connect")
def do_connect(dbapi_connection, connection_record):
    # disable pysqlite's emitting of the BEGIN statement entirely.
    # also stops it from emitting COMMIT before any DDL.
    dbapi_connection.isolation_level = None


@event.listens_for(engine, "begin")
def do_begin(conn):
    # emit our own BEGIN
    conn.exec_driver_sql("BEGIN")


@pytest.fixture()
def session():
    connection = engine.connect()
    transaction = connection.begin()
    session = TestingSessionLocal(bind=connection)

    nested = connection.begin_nested()

    @event.listens_for(session, "after_transaction_end")
    def end_savepoint(session, transaction):
        nonlocal nested
        if not nested.is_active:
            nested = connection.begin_nested()

    yield session
    session.close()
    transaction.rollback()
    connection.close()

Each time Session.begin_nested() is called, a new “BEGIN SAVEPOINT” command is emitted to the database within the scope of the current database transaction (starting one if not already in progress), and an object of type SessionTransaction is returned, which represents a handle to this SAVEPOINT. When the .commit() method on this object is called, “RELEASE SAVEPOINT” is emitted to the database, and if instead the 
.rollback() method is called, “ROLLBACK TO SAVEPOINT” is emitted. The enclosing database transaction remains in progress.

• begin_nested에 대한 설명을 가볍게 살펴보면, begin_nested를 호출하는 시점에 트랜잭션 안에 새로운 “BEGIN SAVEPOINT”가 시작되고 이를 제어할 수 있는 SessionTransaction 유형의 객체가 반환되어 가장 바깥은 트랜잭션 안에서 커밋과 롤백이 일어난다.
• 그러므로 가장 바깥 트랜잭션에서 커밋을 하지 않기 때문에 실제로 테이블에 데이터 쓰기 없이 test 함수 단위로 데이터를 생성하고 롤백할 수 있게 된다.

 

sqlalchemy_session_factory 버그 픽스 건이 배포되면 사용할 일은 없겠지만 factory_boy 동작 방식에 대해 좀 더 깊이 이해하고 하나의 공식 라이브러리를 작게나마 흉내 낼 수 있어 즐거웠다 :)

 

참고

https://factoryboy.readthedocs.io/en/stable/

https://docs.sqlalchemy.org/en/20/orm/session_transaction.html#using-savepoint