flask-sqlalchemy、pytest 的单元测试和事务自动回滚

使用 flask-sqlalchemy 做数据库时,单元测试可以帮助发现一些可能意想不到的问题,像 delete-cascade 、数据长度、多对多关联等等。如果使用 alembic 管理数据库版本,还可以写些跟迁移相关的单元测试。在团队中实现规范的单元测试,再配合 flake8 / pep8 之类的代码规范工具,有助于提高代码的质量,让开发人员有意识去主动发现问题,在新功能进行回归测试时可重复使用单元测试的代码,避免中断已有功能。

在数据库的单元测试前,先要把数据库表结构创建起来。有两种方式来做这个事,一是用一个专门用于单元测试的数据库,重新进行 create_all 的建表操作;或者直接用开发用的数据库, upgrade 到最新的版本,然后在单元测试完成后再进行 downgrade 。两种方法都可以使用 pytest 的 fixture 实现,都可以使用。

做数据库的单元测试需要注意的是事务回滚。通常每个单元测试会针对某个单独的类,而为了避免单元测试所插入或删除的数据对其它单元测试的影响,需要把整个单元测试包含在一个事务中,单元测试结束时自动回滚。事务所针对的层次可以是一个模块,或者是单元测试函数。

假如项目的目录结构如下 ::

.
├── config.py
├── db.sqlite
├── manage.py
├── maze
├── migrations
├── requirements-dev.txt
├── requirements-test.txt
├── requirements.txt
├── tests
├── tox.ini

db.sqlite 为开发用的 sqlite 数据库,migrations 为数据库迁移文件,tests 中放单元测试。在配置 sqlite 路径时,需要注意使用绝对路径,否则会把 sqlite 文件写在 maze 目录下。

现在有 model 类的定义如下 ::

class User(db.Model, Base, UserMixin):
    """User model."""

    __tablename__ = 'users'
    username = db.Column(db.String(20), unique=True, nullable=False)
    email = db.Column(db.String(64), unique=True, nullable=False)
    password_hash = db.Column(db.String(128), nullable=False)
    
    @property
    def password(self):
        raise AttributeError('Not readable')
    
    @password.setter
    def password(self, value):
        self.password_hash = generate_password_hash(value)
    
    def save(self):
        db.session.add(self)
        db.session.flush()

为了测试这个类,需要先编写一些可行的测试框架代码,以在 pytest 中支持对 flask 的测试。

在 tests 中创建 conftest.py 文件,包含我们在整个测试中都需要使用的 fixture 。首先,需要创建一个 flask 的 fixture ,并创建一个测试上下文。

# -*- coding: utf-8 -*-

import pytest


@pytest.yield_fixture(scope='session',autouse=True)
def app():
    from maze import create_app
    app = create_app()
    ctx = app.test_request_context()
    with ctx:
        yield app

如果使用单独测试数据库的方式进行单元测试,那么加一个创建数据库的 fixture ,并在测试::

@pytest.yield_fixture(scope='session', autouse=True)
@pytest.mark.usefixtures('app')
def create_schema():
    from maze import db
    db.create_all()
    try:
        yield
    finally:
        db.drop_all()

那么在单元测试中,就已经可以通过 @pytest.mark.usefixtures('app') ,flask 的相关代码就能进行测试。为了支持自动回滚,还需要加一个 rollback 的 fixture ,如下 ::

@pytest.yield_fixture(autouse=True)
def rollback(app):
    from maze import db
    ctx = app.test_request_context()
    with ctx:
        db.session.begin(subtransactions=True)
        try:
            yield
        finally:
            db.session.rollback()
            db.session.remove()

rollback 可使用默认的 function scope ,如果需要在模块范围内才进行 rollback ,那么设置 scope='module' 。但要注意,原始代码中应用 db.session.flush() 的方式刷新 session 数据,而不是用 commitrollback 对本事务上次 commit 之后的数据才生效,因此如果单元测试中有 commit 的代码,部分数据会被写入数据库,从而达不到 rollback 的目的。

OK,现在在 test_user.py 的单元测试模块中写些测试代码 ::

@pytest.usefixtures('rollback')
def test_user():
    user = User(username='test', password='test123')
    user.save()
    assert user.id is not None
    
    user2 = User.query.filter_by(username='test').first()
    assert user2 is not None
    assert user == user2
    
    with pytest.raises(AttributeError):
        p = user.password   # noqa:F841

最后,执行 pytest 单元测试 ::

pytest tests

整个过程有几点要注意的:
(1) 为了支持自动回滚,原始代码不使用 commit 而要用 flush ,但要在 flask 中用 @app.teardown_request 注册一个处理函数,在函数中进行最终的 commit ;
(2) fixture 需要写在 conftest.py 文件中(或在 conftest.py 中进行 import),而单元测试模块不需要 import 这些 fixture ,否则会出现找不到 fixture (尤其在 pycharm IDE 直接调用 pytest 测试时非常容易出现);
(3) 注意 flask 的上下文处理,可使用 pytest-flask 这个插件处理一些问题。

pytest 的 fixture 编写方式和 unittest 、 nose 、 xUnit 等有点区别,它是以函数的方式出现,在一开始用的时候会不习惯,也没有用 setup 、 teardown 等常见的约定函数来注入依赖,而是用简单的函数参数的方式,用过一小段时间后,感觉也还挺好,简单明了处理的方式在编写代码时少了很多麻烦。