如何使用 pytest 对 sqlalchemy orm 类进行单元测试

Question

我想编写一些 py.test 代码来测试 2 个基于 this Tutorial 创建的简单 sqlalchemy ORM 类。问题是，如何将 py.test 中的数据库设置为测试数据库并在测试完成后回滚所有更改？是否可以在不实际连接数据库的情况下模拟数据库并运行测试？

这是我的课程的代码：

from sqlalchemy import create_engine, ForeignKey
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import Column, Integer, String
from sqlalchemy.orm import sessionmaker, relationship

eng = create_engine('mssql+pymssql://user:pass@host/my_database')

Base = declarative_base(eng)
Session = sessionmaker(eng)
intern_session = Session()

class Author(Base):
    __tablename__ = "Authors"

    AuthorId = Column(Integer, primary_key=True)
    Name = Column(String)  
    Books = relationship("Book")

    def add_book(self, title):
        b = Book(Title=title, AuthorId=self.AuthorId)
        intern_session.add(b)
        intern_session.commit()

class Book(Base):
    __tablename__ = "Books"

    BookId = Column(Integer, primary_key=True)
    Title = Column(String)      
    AuthorId = Column(Integer, ForeignKey("Authors.AuthorId"))    

    Author = relationship("Author")                           

Answer 1

我通常这样做：

1. 我不使用模型声明来实例化引擎和会话，而是只声明一个没有绑定的 Base：

   Base = declarative_base()
   

   我只在需要时创建会话

   engine = create_engine('<the db url>')
   db_session = sessionmaker(bind=engine)
   

   您可以通过在add_book 方法中不使用intern_session 而是使用session 参数来执行相同的操作。

   def add_book(self, session, title):
       b = Book(Title=title, AuthorId=self.AuthorId)
       session.add(b)
       session.commit()
   

   它使您的代码更具可测试性，因为您现在可以在调用该方法时传递您选择的会话。 而且您不再被绑定到硬编码数据库 url 的会话所困。

2. 使用 pytest_addoption hook 向 pytest 添加自定义 --dburl 选项。

   只需将此添加到您的顶级conftest.py：

   def pytest_addoption(parser):
       parser.addoption('--dburl',
                        action='store',
                        default='<if needed, whatever your want>',
                        help='url of the database to use for tests')
   

   现在你可以运行pytest --dburl <url of the test database>

3. 然后您可以从request 夹具中检索dburl 选项

   • 来自自定义夹具：

     @pytest.fixture()
     def db_url(request):
         return request.config.getoption("--dburl")
         # ...
     

   • 测试内部：

     def test_something(request):
         db_url = request.config.getoption("--dburl")
         # ...
     

---

此时您可以：

• 在任何测试或夹具中获取测试db_url
• 用它来创建引擎
• 创建绑定到引擎的会话
• 将会话传递给经过测试的方法

在每个测试中都这样做很麻烦，因此您可以有效地使用 pytest 固定装置来简化过程。

以下是我使用的一些固定装置：

from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker


@pytest.fixture(scope='session')
def db_engine(request):
    """yields a SQLAlchemy engine which is suppressed after the test session"""
    db_url = request.config.getoption("--dburl")
    engine_ = create_engine(db_url, echo=True)

    yield engine_

    engine_.dispose()


@pytest.fixture(scope='session')
def db_session_factory(db_engine):
    """returns a SQLAlchemy scoped session factory"""
    return scoped_session(sessionmaker(bind=db_engine))


@pytest.fixture(scope='function')
def db_session(db_session_factory):
    """yields a SQLAlchemy connection which is rollbacked after the test"""
    session_ = db_session_factory()

    yield session_

    session_.rollback()
    session_.close()

使用db_session 夹具，您可以为每个单独的测试获得一个清新干净的db_session。 当测试结束时，db_session 被回滚，保持数据库干净。