Unit-testing a flask-principal application

Question

All, I'm writing a flask application that depends on flask-principal for managing user roles. I'd like to write some simple unit tests to check which views can be accessed by which user. An example of code is posted on pastebin to avoid cluttering this post. In short, I define a few routes, decorating some so that they can be accessed only by users with the proper role, then try to access them in a test.

In the code pasted, the test_member and test_admin_b both fail, complaining about a PermissionDenied. Obviously, I'm failing to declare the user properly; at least, the info about the user roles is not in the right context.

Any help or insight about the complexities of context processing will be deeply appreciated.

Elisa

Answer 1

Flask-Principal does not store information for you between requests. It's up to you to do this however you like. Keep that in mind and think about your tests for a moment. You call the test_request_context method in the setUpClass method. This creates a new request context. You are also making test client calls with self.client.get(..) in your tests. These calls create additional request contexts that are not shared between each other. Thus, your calls to identity_changed.send(..) do not happen with the context of the requests that are checking for permissions. I've gone ahead and edited your code to make the tests pass in hopes that it will help you understand. Pay special attention to the before_request filter I added in the create_app method.

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()

Pharmacology

That's what I feared: I'm getting lost in contexts. AFAIU, your determine_identity will be called before the request is processed, using the same context, right? So, I need to declare an identity somewhere in that context, or retrieve it from some global context, or create it on the fly from some extra arguments passed to the request (eg, the query_string)... I'll try to post some solutions in another answer, I'd be quite grateful if you could let me know what you thought. – Elisa

Correct. But I'm not sure why you fear this. And yes, the determine_identity function will be called on every request and share the same context with your view methods. Determining the identity all depends on how you plan to authenticate users. For example, if you want a session based authentication mechanism you should pair Flask-Principal with Flask-Login. If you're building an API that is stateless, you should pass auth parameters in headers or use basic http auth and determine the user in the determine_identity from those values. – Pharmacology

One thing I failed to mention is that Flask-Principal, by default, saves the identity in the session, so the first time you call the identity_changed.send method it will store the identity in the session and load it for every request except for static endpoints. – Pharmacology

To put things back in context (pun intended): I'm taking over an application that works as expected but has a dire lack of unittests (bad, bad kitty). This app uses flask.login for identification and flask.principal for role management. I wanted to add the missing tests without modifying the current logic. Once a user is logged in, her roles are set and her Identity saved in the session, as you pointed out. In the unit tests, I need to switch identities and that's where your determine_identy comes to play... – Elisa

@MattW How to deal with the roles and permissions stored in a db? In my case, permissions are constant one and roles are created dynamically. For ex, an user can have any number of roles. For each role, there should be some set of permissions. In this case, How I assign more than one roles to the identity? – Oe

Answer 2

As Matt explained, it's only a matter of context. Thanks to his explanations, I came with two different ways to switch identities during unit tests.

Before all, let's modify a bit the application creation:

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

---

A first possibility is to create a function that loads an identity before each request in our test. The easiest is to declare it in the setUpClass of the test suite after the app is created, using the app.before_request decorator:

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))

Then, the tests become:

    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")

(Incidentally, the very last test shows that Matt's decorator are far easier to use....)

---

A second approach uses the test_request_context function with a with ... to create a temporary context. No need to define a function decorated by @app.before_request, just pass the route to test as argument of test_request_context, send the identity_changed signal in the context and use the .full_dispatch_request method

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")

Elisa

Answer 3

Along Matt's response, I've created a context manager to make the determine_identity a little cleaner:

@contextmanager
def identity_setter(app, user):
    @app.before_request
    def determine_identity():
        #see https://mcmap.net/q/1695268/-unit-testing-a-flask-principal-application 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)]}

So when I run my test it looks like:

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')

I hope this helps, and I would welcome any feedback :)

~Victor

Ochlocracy