Advertisement
pierregm

Unit testing a flask-principal app

May 23rd, 2013
529
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 3.94 KB | None | 0 0
  1. import unittest
  2. from functools import wraps
  3.  
  4. import hmac
  5. from hashlib import sha1
  6.  
  7. import flask
  8.  
  9. from flask.ext.principal import (Principal, Permission, RoleNeed, Identity,
  10.                                  identity_changed, identity_loaded,
  11.                                  current_app)
  12.  
  13.  
  14. def roles_required(*roles):
  15.     """Decorator which specifies that a user must have all the specified roles.
  16.    Example::
  17.  
  18.        @app.route('/dashboard')
  19.        @roles_required('admin', 'editor')
  20.        def dashboard():
  21.            return 'Dashboard'
  22.  
  23.    The current user must have both the `admin` role and `editor` role in order
  24.    to view the page.
  25.  
  26.    :param args: The required roles.
  27.  
  28.    Source: https://github.com/mattupstate/flask-security/
  29.    """
  30.     def wrapper(fn):
  31.         @wraps(fn)
  32.         def decorated_view(*args, **kwargs):
  33.             perms = [Permission(RoleNeed(role)) for role in roles]
  34.             for perm in perms:
  35.                 if not perm.can():
  36.                     # return _get_unauthorized_view()
  37.                     flask.abort(403)
  38.             return fn(*args, **kwargs)
  39.         return decorated_view
  40.     return wrapper
  41.  
  42.  
  43.  
  44. def roles_accepted(*roles):
  45.     """Decorator which specifies that a user must have at least one of the
  46.    specified roles. Example::
  47.  
  48.        @app.route('/create_post')
  49.        @roles_accepted('editor', 'author')
  50.        def create_post():
  51.            return 'Create Post'
  52.  
  53.    The current user must have either the `editor` role or `author` role in
  54.    order to view the page.
  55.  
  56.    :param args: The possible roles.
  57.    """
  58.     def wrapper(fn):
  59.         print "roles_accepted: ", roles
  60.         @wraps(fn)
  61.         def decorated_view(*args, **kwargs):
  62.             perm = Permission(*[RoleNeed(role) for role in roles])
  63.             print "roles_accepted.permission:", perm, perm.can()
  64.             if perm.can():
  65.                 return fn(*args, **kwargs)
  66.             # return _get_unauthorized_view()
  67.             flask.abort(403)
  68.         return decorated_view
  69.     return wrapper
  70.  
  71.  
  72. def _on_principal_init(sender, identity):
  73.     if identity.id == 'admin':
  74.         identity.provides.add(RoleNeed('admin'))
  75.     identity.provides.add(RoleNeed('member'))
  76.  
  77.  
  78. def create_app():
  79.     app = flask.Flask(__name__)
  80.     app.debug = True
  81.     app.config.update(SECRET_KEY='secret',
  82.                       TESTING=True)
  83.     principal = Principal(app)
  84.     identity_loaded.connect(_on_principal_init)
  85.     #
  86.     @app.route('/')
  87.     def index():
  88.         return "OK"
  89.     #
  90.     @app.route('/member')
  91.     @roles_accepted('admin', 'member')
  92.     def role_needed(self):
  93.         return "OK"
  94.  
  95.     @app.route('/admin')
  96.     @roles_required('admin')
  97.     def connect_admin(self):
  98.         return "OK"
  99.  
  100.     @app.route('/admin_b')
  101.     @admin_permission.require()
  102.     def connect_admin_alt(self):
  103.         return "OK"
  104.  
  105.     return app
  106.  
  107.  
  108. admin_permission = Permission(RoleNeed('admin'))
  109.  
  110.  
  111.  
  112. class WorkshopTest(unittest.TestCase):
  113.     #
  114.     @classmethod
  115.     def setUpClass(cls):
  116.         app = create_app()
  117.         cls.app = app
  118.         cls.client = app.test_client()
  119.         cls.testing = app.test_request_context()
  120.     #
  121.  
  122.     def test_basic(self):
  123.         r = self.client.get('/')
  124.         print r.status_code
  125.         self.assertEqual(r.data, "OK")
  126.  
  127.     def test_member(self):
  128.         r = self.client.get('/member')
  129.         self.assertEqual(r.status_code, 403)
  130.         #
  131.         identity_changed.send(current_app, identity=Identity('admin'))
  132.         r = self.client.get('/member')
  133.         self.assertEqual(r.status_code, 200)
  134.         self.assertEqual(r.data, "OK")
  135.  
  136.     def test_admin_b(self):
  137.         with self.testing as c:
  138.             identity_changed.send(self.app, identity=Identity('admin'))
  139.             r = self.client.get('/admin_b')
  140.             self.assertEqual(r.status_code, 200)
  141.             self.assertEqual(r.data, "OK")
  142.  
  143.  
  144. if __name__ == '__main__':
  145.     unittest.main()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement