使用单元测试对Flask-Principal应用进行测试

Question

大家好，我正在编写一个flask应用程序，该应用程序依赖于flask-principal来管理用户角色。我想编写一些简单的单元测试来检查哪些视图可以被哪些用户访问。为了避免混乱，我在pastebin上发布了代码示例。简而言之，我定义了一些路由，对其中的一些进行了修饰，以便只能被具有适当角色的用户访问，然后在一个测试中尝试访问它们。

在粘贴的代码中，test_member和test_admin_b都失败了，并抱怨PermissionDenied。显然，我没有正确声明用户；至少，关于用户角色的信息不在正确的上下文中。

非常感谢任何有关上下文处理复杂性的帮助或见解。

- Pierre GM

Answer 1

Flask-Principal不会在请求之间为您存储信息。您需要自行决定如何存储信息。请记住这一点，并花一些时间考虑一下您的测试。您在setUpClass方法中调用test_request_context方法。这会创建一个新的请求上下文。您还在测试中使用self.client.get(..)进行测试客户端调用。这些调用会创建额外的请求上下文，彼此之间不共享。因此，您对identity_changed.send(..)的调用不会发生在检查权限的请求上下文中。我已经编辑了您的代码，使测试通过，希望这能帮助您理解。请特别注意我在create_app方法中添加的before_request过滤器。

import hmac
import unittest

from functools import wraps
from hashlib import sha1

import flask

from flask.ext.principal import Principal, Permission, RoleNeed, Identity, \
    identity_changed, identity_loaded current_app


def roles_required(*roles):
    """Decorator which specifies that a user must have all the specified roles.
    Example::

        @app.route('/dashboard')
        @roles_required('admin', 'editor')
        def dashboard():
            return 'Dashboard'

    The current user must have both the `admin` role and `editor` role in order
    to view the page.

    :param args: The required roles.

    Source: https://github.com/mattupstate/flask-security/
    """
    def wrapper(fn):
        @wraps(fn)
        def decorated_view(*args, **kwargs):
            perms = [Permission(RoleNeed(role)) for role in roles]
            for perm in perms:
                if not perm.can():
                    # return _get_unauthorized_view()
                    flask.abort(403)
            return fn(*args, **kwargs)
        return decorated_view
    return wrapper



def roles_accepted(*roles):
    """Decorator which specifies that a user must have at least one of the
    specified roles. Example::

        @app.route('/create_post')
        @roles_accepted('editor', 'author')
        def create_post():
            return 'Create Post'

    The current user must have either the `editor` role or `author` role in
    order to view the page.

    :param args: The possible roles.
    """
    def wrapper(fn):
        @wraps(fn)
        def decorated_view(*args, **kwargs):
            perm = Permission(*[RoleNeed(role) for role in roles])
            if perm.can():
                return fn(*args, **kwargs)
            flask.abort(403)
        return decorated_view
    return wrapper


def _on_principal_init(sender, identity):
    if identity.id == 'admin':
        identity.provides.add(RoleNeed('admin'))
    identity.provides.add(RoleNeed('member'))


def create_app():
    app = flask.Flask(__name__)
    app.debug = True
    app.config.update(SECRET_KEY='secret', TESTING=True)
    principal = Principal(app)
    identity_loaded.connect(_on_principal_init)

    @app.before_request
    def determine_identity():
        # This is where you get your user authentication information. This can
        # be done many ways. For instance, you can store user information in the
        # session from previous login mechanism, or look for authentication
        # details in HTTP headers, the querystring, etc...
        identity_changed.send(current_app._get_current_object(), identity=Identity('admin'))

    @app.route('/')
    def index():
        return "OK"

    @app.route('/member')
    @roles_accepted('admin', 'member')
    def role_needed():
        return "OK"

    @app.route('/admin')
    @roles_required('admin')
    def connect_admin():
        return "OK"

    @app.route('/admin_b')
    @admin_permission.require()
    def connect_admin_alt():
        return "OK"

    return app


admin_permission = Permission(RoleNeed('admin'))


class WorkshopTest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        app = create_app()
        cls.app = app
        cls.client = app.test_client()

    def test_basic(self):
        r = self.client.get('/')
        self.assertEqual(r.data, "OK")

    def test_member(self):
        r = self.client.get('/member')
        self.assertEqual(r.status_code, 200)
        self.assertEqual(r.data, "OK")

    def test_admin_b(self):
        r = self.client.get('/admin_b')
        self.assertEqual(r.status_code, 200)
        self.assertEqual(r.data, "OK")


if __name__ == '__main__':
    unittest.main()

- Matt W

---

这正是我担心的：我在上下文中迷失了方向。据我所知，您的 determine_identity 将在请求处理之前使用相同的上下文进行调用，对吗？因此，我需要在该上下文中声明一个标识，或从某个全局上下文中检索它，或者根据传递给请求的一些额外参数（例如 query_string）即时创建它... 我将尝试在另一个答案中发布一些解决方案，如果您能让我知道您的想法，我将不胜感激。 - Pierre GM

正确。但我不确定你为什么害怕这个。是的，determine_identity函数将在每个请求上被调用，并与您的视图方法共享相同的上下文。确定身份完全取决于您计划如何验证用户。例如，如果您想要一个基于会话的身份验证机制，您应该将Flask-Principal与Flask-Login配对使用。如果您正在构建一个无状态的API，则应在标头中传递auth参数或使用基本的http auth，并从这些值中在determine_identity中确定用户。 - Matt W

我忘了提到的是，Flask-Principal默认将身份信息保存在会话中，因此第一次调用“identity_changed.send”方法时，它将在会话中存储身份信息，并为每个请求加载它，但静态端点除外。 - Matt W

为了让事情回到上下文（双关语意），我接手了一个功能正常但缺乏单元测试的应用程序（糟糕，糟糕的小猫）。该应用程序使用 flask.login 进行身份验证和 flask.principal 进行角色管理。我想在不修改当前逻辑的情况下添加缺失的测试。一旦用户登录，她的角色就被设置并且她的 Identity 保存在会话中，正如你所指出的那样。在单元测试中，我需要切换身份，这就是你的 determine_identity 起作用的地方... - Pierre GM

@MattW 如何处理存储在数据库中的角色和权限？在我的情况下，权限是恒定的，而角色是动态创建的。例如，一个用户可以拥有任意数量的角色。对于每个角色，都应该有一些权限集。在这种情况下，我如何将多个角色分配给身份？ - Avinash Raj

Answer 2

正如 Matt解释的那样，这只是一个上下文问题。在他的解释下，我想出了两种不同的方法来在单元测试期间切换标识。

首先，让我们稍微修改一下应用程序创建：

def _on_principal_init(sender, identity):
    "Sets the roles for the 'admin' and 'member' identities"
    if identity.id:
        if identity.id == 'admin':
            identity.provides.add(RoleNeed('admin'))
        identity.provides.add(RoleNeed('member'))

def create_app():
    app = flask.Flask(__name__)
    app.debug = True
    app.config.update(SECRET_KEY='secret',
                      TESTING=True)
    principal = Principal(app)
    identity_loaded.connect(_on_principal_init)
    #
    @app.route('/')
    def index():
        return "OK"
    #
    @app.route('/member')
    @roles_accepted('admin', 'member')
    def role_needed():
        return "OK"
    #
    @app.route('/admin')
    @roles_required('admin')
    def connect_admin():
        return "OK"

    # Using `flask.ext.principal` `Permission.require`...
    # ... instead of Matt's decorators
    @app.route('/admin_alt')
    @admin_permission.require()
    def connect_admin_alt():
        return "OK"

    return app

---

一种可能性是创建一个函数，在我们的测试中在每个请求之前加载一个身份。最简单的方法是在应用程序创建后，在测试套件的setUpClass中声明它，并使用app.before_request装饰器：

class WorkshopTestOne(unittest.TestCase):
    #
    @classmethod
    def setUpClass(cls):
        app = create_app()
        cls.app = app
        cls.client = app.test_client()

        @app.before_request
        def get_identity():
            idname = flask.request.args.get('idname', '') or None
            print "Notifying that we're using '%s'" % idname
            identity_changed.send(current_app._get_current_object(),
                                  identity=Identity(idname))

然后，测试变成了：

    def test_admin(self):
        r = self.client.get('/admin')
        self.assertEqual(r.status_code, 403)
        #
        r = self.client.get('/admin', query_string={'idname': "member"})
        self.assertEqual(r.status_code, 403)
        #
        r = self.client.get('/admin', query_string={'idname': "admin"})
        self.assertEqual(r.status_code, 200)
        self.assertEqual(r.data, "OK")
    #
    def test_admin_alt(self):
        try:
            r = self.client.get('/admin_alt')
        except flask.ext.principal.PermissionDenied:
            pass
        #
        try:
            r = self.client.get('/admin_alt', query_string={'idname': "member"})
        except flask.ext.principal.PermissionDenied:
            pass
        #
        try:
            r = self.client.get('/admin_alt', query_string={'idname': "admin"})
        except flask.ext.principal.PermissionDenied:
            raise
        self.assertEqual(r.data, "OK")

（顺便说一下，最后一个测试表明Matt的装饰器要容易使用得多...）

---

第二种方法使用test_request_context函数和with ...创建临时上下文。不需要定义一个由@app.before_request装饰的函数，只需将要测试的路由作为test_request_context的参数传递，在上下文中发送identity_changed信号，并使用.full_dispatch_request方法。

class WorkshopTestTwo(unittest.TestCase):
    #
    @classmethod
    def setUpClass(cls):
        app = create_app()
        cls.app = app
        cls.client = app.test_client()
        cls.testing = app.test_request_context


    def test_admin(self):
        with self.testing("/admin") as c:
            r = c.app.full_dispatch_request()
            self.assertEqual(r.status_code, 403)
        #
        with self.testing("/admin") as c:
            identity_changed.send(c.app, identity=Identity("member"))
            r = c.app.full_dispatch_request()
            self.assertEqual(r.status_code, 403)
        #
        with self.testing("/admin") as c:
            identity_changed.send(c.app, identity=Identity("admin"))
            r = c.app.full_dispatch_request()
            self.assertEqual(r.status_code, 200)
            self.assertEqual(r.data, "OK")

- Pierre GM

Answer 3

根据Matt的回答，我创建了一个上下文管理器，使determine_identity函数更加简洁：

@contextmanager
def identity_setter(app, user):
    @app.before_request
    def determine_identity():
        #see https://dev59.com/LXLYa4cB1Zd3GeqPX3Wx for details
        identity_changed.send(current_app._get_current_object(), identity=Identity(user.id))
    determine_identity.remove_after_identity_test = True
    try: 
        yield
    finally:
        #if there are errors in the code under trest I need this to be run or the addition of the decorator could affect other tests
        app.before_request_funcs = {None: [e for e in app.before_request_funcs[None] if not getattr(e,'remove_after_identity_test', False)]}

所以当我运行我的测试时，它看起来像：

with identity_setter(self.app,user):
           with user_set(self.app, user):
                with self.app.test_client() as c:
                    response = c.get('/orders/' + order.public_key + '/review')

希望这能帮到您，欢迎提出任何反馈意见 :)

~维克多

- VFR292