如何使用pytest对SQLAlchemy Orm类进行单元测试

Question

我想编写一些py.test代码来测试2个基于此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)
   

   您可以通过不使用做同样的intern_session在你的add_book方法，而是用一个session参数。

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

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

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

   只需将此添加到您的顶级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. 然后您可以dburl从request夹具中检索选项

   • 从自定义装置：

     @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回滚，保持数据库干净。