Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
Browse files

Adding a better way of building auth decorators, then converting the

login_required to use the new stuff... plus an "admin_required"
decorator
  • Loading branch information...
commit 253eaa2aa9477ab9b8d070d71baf8992500e45e8 1 parent 56d7ee6
@coleifer authored
View
27 flask_peewee/auth.py
@@ -99,17 +99,24 @@ def get_urls(self):
def get_login_form(self):
return LoginForm
- def login_required(self, func):
- @functools.wraps(func)
- def inner(*args, **kwargs):
- user = self.get_logged_in_user()
-
- if not user:
- login_url = url_for('%s.login' % self.blueprint.name, next=get_next())
- return redirect(login_url)
+ def test_user(self, test_fn):
+ def decorator(fn):
+ @functools.wraps(fn)
+ def inner(*args, **kwargs):
+ user = self.get_logged_in_user()
+
+ if not user or not test_fn(user):
+ login_url = url_for('%s.login' % self.blueprint.name, next=get_next())
+ return redirect(login_url)
+ return fn(*args, **kwargs)
+ return inner
+ return decorator
- return func(*args, **kwargs)
- return inner
+ def login_required(self, func):
+ return self.test_user(lambda u: True)(func)
+
+ def admin_required(self, func):
+ return self.test_user(lambda u: u.admin)(func)
def authenticate(self, username, password):
active = self.User.select().where(active=True)
View
21 flask_peewee/tests/auth.py
@@ -162,3 +162,24 @@ def test_login_required(self):
self.assertEqual(resp.status_code, 200)
self.assertEqual(auth.get_logged_in_user(), self.admin)
+
+ def test_admin_required(self):
+ self.create_users()
+
+ with self.flask_app.test_client() as c:
+ resp = c.get('/secret/')
+ self.assertEqual(resp.status_code, 302)
+ self.assertTrue(resp.headers['location'].endswith('/accounts/login/?next=%2Fsecret%2F'))
+
+ self.login('normal', 'normal', c)
+
+ resp = c.get('/secret/')
+ self.assertEqual(resp.status_code, 302)
+ self.assertTrue(resp.headers['location'].endswith('/accounts/login/?next=%2Fsecret%2F'))
+ self.assertEqual(auth.get_logged_in_user(), self.normal)
+
+ self.login('admin', 'admin', c)
+ resp = c.get('/secret/')
+ self.assertEqual(resp.status_code, 200)
+
+ self.assertEqual(auth.get_logged_in_user(), self.admin)
View
5 flask_peewee/tests/test_app.py
@@ -211,6 +211,11 @@ def homepage():
def private_timeline():
return Response()
+@app.route('/secret/')
+@auth.admin_required
+def secret_area():
+ return Response()
+
admin.setup()
api.setup()
Please sign in to comment.
Something went wrong with that request. Please try again.