Unit-testing a flask-principal application
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 hmacimport unittestfrom functools import wrapsfrom hashlib import sha1import flaskfrom flask.ext.principal import Principal, Permission, RoleNeed, Identity, \ identity_changed, identity_loaded current_appdef 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 wrapperdef 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 wrapperdef _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 appadmin_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()
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")
Along Matt's response, I've created a context manager to make the determine_identity a little cleaner:
@contextmanagerdef identity_setter(app, user): @app.before_request def determine_identity(): #see http://stackoverflow.com/questions/16712321/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