一个使用Flask-Login登录后的Pytest测试用例的坑

前言

最近写一个项目，用到了 Flask-Login 实现用户登录状态。项目最后需要编写测试用例，但是测试代码卡在确认用户登录状态的部分就写不下去了，研究了好久才找到原因，分享一下 pytest 测试 Flask 应用上的一个坑儿.

一个最基础的例子

为了演示这个问题，我写了一个小型应用，全部代码可以从 dongweiming/mp 获取。

这里只列出部分核心代码。首先是测试部分代码:

from app import db, User


def test_user_info(client):
    response = client.get('/api/user/info')
    assert response.json == {}


def test_login(client):
    db.session.add(User(name='user1'))
    db.session.commit()
    response = client.post('/api/login', json={
        'id': '1'})
    assert response.json['id'] == 1


def test_after_login(client):
    response = client.post('/api/login', json={
        'id': 1})
    response = client.get('/api/user/info')
    assert response.json == {}

其中test_login里面给数据库里添加了一个用户记录，这是因为前面没有注册相关的代码。这么写是因为本来是一个专门用例测试注册的，里面有添加用户的逻辑，为了精简代码就在这里先添加用户再测试这个用户的登录了。当然，这么写也是为了重现问题，最后总结时会提到。

这部分测试代码，第一个函数中因为没有登录，所以获取不到用户信息；第二个函数中测试了登录功能，当请求结束后会返回用户信息（包含 id=1）；最后一部分也就是出问题的部分，本来是测试登陆后的逻辑，按照设想，第一步登录，然后在请求用户信息应该可以拿到对应用户信息了。

但事实上，response.json的返回结果还是{}。也就是说，前面先请求登录 API/api/login的那部分就没生效。

初步排查

因为 Flask-Login 的逻辑主要是在登录后给浏览器设置 httponly 为 True 的 Cookie，这样下次请求时候带着这个正确的 Cookie 后端就认为此时已经登录。所以我一开始认为是 client 请求登录 API 后并没有设置对应的 Cookie，但事实上:

第一个 response，response.headers ["Set-Cookie"] 里表示确实已经设置过 Cookie
第二个 response，response.request.headers 里确实有这个 Cookie，那么请求是正确的，但是 API 没有返回正确的用户信息

所以先确认问题不在这个 client 实例上。

深入排查

接着看一下conftest.py的代码:

import os

import pytest

from app import app as _app, db


@pytest.fixture
def app():
    _app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///test.db'

    with _app.app_context():
        db.create_all()
    yield _app

    os.remove('test.db')


@pytest.fixture
def client(app):
    return app.test_client()

conftest.py用来定义被测试用例们共享的 fixture、钩子等设置。例如上面测试用例函数中的参数client就是一个 fixture，它就是这里定义的。

fixture是一个函数，可以用作初始化或者返回了某种数据，所有测试都可以访问它，这样就达到了代码重用和逻辑隔离。例如上面的 2 个 fixture:

app。包含了初始化 (重建数据库表) 设置的 Flask 实例，通过 yield 迭代出来，不过最后还会做清理工作 (删除 SQLite 文件)`
client。使用了 app 这个 fixture，返回了测试客户端，这是 Flask 自带的类。

一眼看去这部分代码并没有问题。既然是测试，当然不能用假数据也不要做monkeypatch，就应该真实的创建数据库，操作数据。

我在 IPython 里面试验，是符合预期的:

In [1]: from flask.testing import FlaskClient

In [2]: client = FlaskClient(app)

In [3]: response = client.post(
   ...:     "/api/login", json={'id': '1'}
   ...: )

In [4]: response = client.get("/api/user/info")

In [5]: response.json
Out[5]:
{'id': 1}

这不会是一个 BUG 吧，哈哈！不过等我阅读了对应源码，确认过程虽然和执行 pytest 的方式下有一些逻辑差别，但是并不影响预期结果。

找到灵感

困扰了很久，网上一顿搜索，最终在 pytest-flask 的一个 issue 下找到相关讨论，不过按照我的理解其实可以过滤大部分回答。有些人说自己的解决方案可以工作我也不是很能理解，可能是代码逻辑有区别。

不过这个讨论中有 2 个人说的内容给我了启发，首先是这个:

@vitalk: the client fixture pushes a new request context, so the 1st example doesn’t work because the current_user is anonymous.

我开始觉得这个问题是 Flask 设计的请求上下文或者应用上下文作怪了。这个作者的意思是 client 每次都会使用一个新的请求上下文，啥也不说了，先试验一下，加个 print

@pytest.fixture
def client(app):
    print('Hit')
    return app.test_client()

看看它在测试过程中被调用的频率:

❯ pytest -s  # pytest默认会捕捉各种输出，除非测试用例失败否则都过滤掉了，通过-s可以关闭捕捉(等于--capture=no)
...
collected 3 items

tests/test_api.py Hit
.Hit
.Hit
.

...

这里省略了大部分无关输出。可以看到 Hit 出现了三次，也就是每个测试用例出现了 1 次，但是这也说明在test_after_login只出现了一次，所以不是它的问题。那么我开始重点怀疑第一个 fixture 有什么问题了，不过以我当时的知识还不知道是什么，各种尝试去 hack 代码也没有让测试符合预期

后来又回来看这个 issue，看到之前被一直忽略的最后一个回复:

@jab: To work around this, I ended up creating my own client fixture with a longer scope (e.g. "module" or "session" both work) rather than using pytest-flask, which doesn't currently allow customization of its client fixture's scope. In case this helps anyone else!

一开始我并没懂他的意思。仔细看他说这不是pytest-flask的问题，而是修改 fixture 的应用范围到module或者session就可以了。按着关键词搜了一下，果然可以 (不过我是在 app 上使用了更大范围):

@pytest.fixture(scope='session')
def app():
    _app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///test.db'

    with _app.app_context():
        db.create_all()
    yield _app

    os.remove('test.db')

也就是给pytest.fixture添加参数 scope。这样测试就通过了，开心😄。

思考

最后阅读了pytest.fixture的源码，才理解为什么。分享一下，pytest.fixture的 scope 参数支持 5 个选项，表示不同的范围:

function。默认值，每个测试功能运行一次，函数结束后清理。←注意这句哈
class。每个测试类运行一次。
package。每个包运行一次。
module。每个模块运行一次。
session。每个会话运行一次。

所以问题就是 app 这个 fixture 在每个测试函数结束后就会被调用:

@pytest.fixture
def app():
    print('Hit')
    _app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///test.db'

    with _app.app_context():
        db.create_all()
    yield _app

    os.remove('test.db')

这次打印依然可以看到 3 次Hit:

❯ pytest -s
...
collected 3 items

tests/test_api.py Hit
.Hit
.Hit
.

我恍然大悟，这就是问题，和 Flask 上下文无关，而是和在 fixture 里做数据清理有关。注意看代码，每次都会在结束后删掉数据库文件，下次又会重新创建数据库，对于测试用例来说，在test_login添加了新用户，但是等到test_after_login里重新创建数据库早没数据了。所以找不到这个用户。这样改也是可以的:

def test_after_login(client):
    db.session.add(User(name='user1'))
    db.session.commit()
    response = client.post('/api/login', json={
        'id': 1})
    response = client.get('/api/user/info')
    assert response.json['id'] == 1