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

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 数据，而不是用 commit 。 rollback 对本事务上次 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 等常见的约定函数来注入依赖，而是用简单的函数参数的方式，用过一小段时间后，感觉也还挺好，简单明了处理的方式在编写代码时少了很多麻烦。