Search code examples
pythonunit-testingflaskflask-principal

Unit-testing a flask-principal application


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.


Solution

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