如何为pytest设置current_user?

时间:2017-10-10 15:16:58

标签: python flask flask-login

我正在为在查询中使用当前登录用户的视图编写单元测试:

@app.route('/vendors/create', methods=['GET', 'POST'])
@login_required
def create_vendors():
    vendor_form = VendorForm(request.form)
    if vendor_form.validate_on_submit():
        vendor = db.session.query(Vendors).filter(Vendors.name == vendor_form.name.data,
                                                  Vendors.users.contains(g.user)).first()
        if vendor:
            flash('There is a vendor with this name already.')
            return redirect(url_for('create_vendors', vendor_form=vendor_form))
        vendor = Vendors(name=vendor_form.name.data, aws_access_key=vendor_form.aws_access_key.data,
                         merchant_key=vendor_form.merchant_key.data, secret_key=vendor_form.secret_key.data)
        vendor.users.append(current_user)
        db.session.add(vendor)
        db.session.commit()
        return redirect(url_for('vendors'))
    return render_template('edit_vendor.html', vendor_form=vendor_form)

和测试:

def test_create_vendor(self):
    response = TestViews.client.post('/vendors/create', follow_redirects=True,
                                     data=dict(name='Name', aws_access_key='aws_access_key',
                                               merchant_key='merchant_key',
                                               secret_key='secret_key'))
    assert response.status_code == 200
    vendors = db.session.query(Vendors).filter(Vendors.aws_access_key == 'aws_access_key').all()
    assert len(vendors) == 1
    assert vendors[0].name == 'Name'

并在我的设置中LOGIN_DISABLED = True

我的问题是,在运行测试current_user时从未设置过。

我尝试使用测试客户端登录并将我的测试包装在with TestLogin.client:中但current_user从未设置过。关于如何设置它的任何想法?

2 个答案:

答案 0 :(得分:7)

考虑我为此问题创建的以下示例应用

from flask import Flask, request,
from flask_login import LoginManager
from flask_login.utils import login_required, current_user, login_user

app = Flask(__name__)
app.secret_key = 'super secret key'
app.config['SESSION_TYPE'] = 'filesystem'

login_manager = LoginManager()
login_manager.init_app(app)

class MyUser():
    name = "tarun"

    def is_authenticated(self):
        return True

    def is_active(self):
        return True

    def is_anonymous(self):
        return False

    def get_id(self):
        return 10


@login_manager.user_loader
def load_user(user_id):
    return MyUser()

@app.route("/login", methods=["POST"])
def login():
    login_user(MyUser())
    return "Authenticated"

@app.route("/", methods=["GET", "POST"])
@login_required
def index():
    return "Authenticated " # + current_user


if __name__ == "__main__":
    app.run(debug=True)

测试登录流程有多种方法。

禁用登录和@login_required

在此,您不必在请求调用中使用current_user,只是保护流向经过身份验证的用户。在这种情况下,您可以使用以下方法

import flasktest as flaskr

class FlaskrTestCase(unittest.TestCase):

    def setUp(self):
        flaskr.app.testing = True
        flaskr.app.config['LOGIN_DISABLED'] = True
        flaskr.app.login_manager.init_app(flaskr.app)
        self.app = flaskr.app.test_client()

    def test_without_login(self):
        res = self.app.post("/")

        assert "Unauthorized" in res.data

在这种情况下,测试将失败,因为用户将允许访问经过身份验证的流程。并且current_user将设置为匿名用户。它与禁用@login_required一样好。

使用实际登录进行测试

要执行此操作,您需要确保可以使用POST登录到您实际登录的端点的应用

import unittest

import flasktest as flaskr

class FlaskrTestCase(unittest.TestCase):

    def setUp(self):
        flaskr.app.testing = True
        self.app = flaskr.app.test_client()

    def tearDown(self):
        pass

    def test_login(self):
        self.app.post("/login")

        res = self.app.get("/")

        assert "Authenticated" in res.data

    def test_without_login(self):
        res = self.app.post("/")

        assert "Unauthorized" in res.data

这甚至可以正确设置current_user。在我的情况下,我没有采用任何表单数据参数,在您的情况下,您将使用类似下面的内容

self.app.post('/login', data=dict(
        username=username,
        password=password
    ), follow_redirects=True)

使用会话Cookie

在这种情况下,您可以设置会话user_id_fresh

def test_login_session(self):
    with self.app.session_transaction() as sess:
        sess['user_id'] = '12'
        sess['_fresh'] = True  # https://flask-login.readthedocs.org/en/latest/#fresh-logins

    res = self.app.get("/")

    assert "Authenticated" in res.data

这将依次调用

@login_manager.user_loader
def load_user(user_id):
    return MyUser()

根据我们在会话user_id中提供的12获取实际用户

答案 1 :(得分:-1)

只需模拟flask_login.current_user并设置固定的返回值,如'testuser'