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.
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()
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:
#contextmanager
def 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
Related
I'm developing Flask application and I have a decorator called login_required which checks if the user is logged in, and sends the current user to the next function.
def login_required(function):
#wraps(function)
def decorator(*args, **kwargs):
token = None
if 'x-access-tokens' in request.headers:
token = request.headers['x-access-tokens']
if not token:
return jsonify({'message': 'a valid token is missing'}), 401
try:
data = jwt.decode(token, app.secret_key)
current_user = User.query.filter_by(username=data['username']).first()
except:
return jsonify({'message': 'token is invalid'}), 401
return function(current_user, *args, **kwargs)
return decorator
So, the callback function is declared like this.
#blueprint.route('/settings', methods=['GET'])
#login_required
def settings(current_user):
# here I check if the current_user is admin or not
...
# function body
Now I want to implement an admin_required decorator which depends on login_required decorator, to be easy using it within functions like this.
#blueprint.route('/settings', methods=['GET'])
#admin_required
def settings(current_user):
...
# function body
How can I achieve this?
So you can create your functionality like this
def decorator1(function):
def fun(*args, **kwargs):
user = function(*args, **kwargs)
print("here")
return user
return fun
def decorator2(function):
def fun(*args, **kwargs):
# firslty i will decorate function with decorator1
decor1 = decorator1(function)
decor1_result = decor1(*args, **kwargs)
# do operation with above result in your case now check if user is admin
print("here in decor2")
return decor1_result + 5
return fun
#decorator2
def test_function(a,b):
return a+b
# Now you want to access a+b and then add c to it
I have included comments for better understanding.
In your case decorator1 will be login_required decorator
and decorator2 will be admin_check decorator. So while creating admin check decorator you can access login_required decorator inside it.
I created a simple Server Interceptor that retrieves the user based on the JWT token.
But now I would like to make it available to all the methods of my services.
At the moment im using decorators. But I would like to avoid having to decorate all the methods. In case, decorate only those that do not need the user.
Some one can give me a clue ?
here is my code:
class AuthInterceptor(grpc.ServerInterceptor):
"""Authorization Interceptor"""
def __init__(self, loader):
self._loader = loader
def intercept_service(self, continuation, handler_call_details):
# Authenticate if we not requesting token.
if not handler_call_details.method.endswith('GetToken'):
# My Authentication class.
auth = EosJWTAuth()
# Authenticate using the headers tokens to get the user.
user = auth.authenticate(
dict(handler_call_details.invocation_metadata))[0]
# Do something here to pass the authenticated user to the functions.
cont = continuation(handler_call_details)
return cont
And I'd like my methods can to access the user in a way like this.
class UserService(BaseService, users_pb2_grpc.UserServicer):
"""User service."""
def get_age(self, request, context):
"""Get user's age"""
user = context.get_user()
# or user = context.user
# or user = self.user
# os user = request.get_user()
return pb.UserInfo(name=user.name, age=user.age)
This is a common need for web servers, and it is a good idea to add decorators to the handlers to explicitly set requirement for authentication/authorization. It helps readability, and reduces the overall complexity.
However, here is a workaround to solve your question. It uses Python metaclass to automatically decorate every servicer method.
import grpc
import functools
import six
def auth_decorator(func):
#functools.wraps(func)
def wrapper(request, context):
if not func.__name__.endswith('GetToken'):
auth = FooAuthClass()
try:
user = auth.authenticate(
dict(context.invocation_metadata)
)[0]
request.user = user
except UserNotFound:
context.abort(
grpc.StatusCode.UNAUTHENTICATED,
'Permission denied.',
)
return func(request, context)
return wrapper
class AuthMeta:
def __new__(self, class_name, bases, namespace):
for key, value in list(namespace.items()):
if callable(value):
namespace[key] = auth_decorator(value)
return type.__new__(self, class_name, bases, namespace)
class BusinessServer(FooServicer, six.with_metaclass(AuthMeta)):
def LogicA(self, request, context):
# request.user accessible
...
def LogicGetToken(self, request, context):
# request.user not accessible
...
Building on a previous question of mine (How to unit test a Flask RESTful API), I'm trying to test a Flask RESTful API using a test_client without the app running, rather than using requests while the app is running.
As a simple example, I have an API (flaskapi2.py) with a get function which uses a login decorator:
import flask
import flask_restful
from functools import wraps
app = flask.Flask(__name__)
api = flask_restful.Api(app)
AUTH_TOKEN = "foobar"
def login_required(f):
#wraps(f)
def decorated_function(*args, **kwargs):
if flask.request.headers.get("auth_token") == AUTH_TOKEN:
return f(*args, **kwargs)
else:
return flask.abort(401) # Return HTTP status code for 'Unauthorized'
return decorated_function
class HelloWorld(flask_restful.Resource):
#login_required
def get(self):
return {'hello': 'world'}
api.add_resource(HelloWorld, '/')
if __name__ == "__main__":
app.run(debug=True)
With the app running, I run these unit tests (test_flaskapi2.py in the same directory):
import unittest
import flaskapi2
import requests
import json
AUTH_TOKEN = "foobar"
class TestFlaskApiUsingRequests(unittest.TestCase):
def setUp(self):
self.session = requests.Session()
self.session.headers.update({'auth_token': AUTH_TOKEN})
def test_hello_world(self):
response = self.session.get('http://localhost:5000')
self.assertEqual(response.json(), {'hello': 'world'})
def test_hello_world_does_not_work_without_login(self):
response = requests.get('http://localhost:5000') # Make an unauthorized GET request
self.assertEqual(response.status_code, 401) # The HTTP status code received should be 401 'Unauthorized'
class TestFlaskApi(unittest.TestCase):
def setUp(self):
self.app = flaskapi2.app.test_client()
def test_hello_world(self):
response = self.app.get('/', headers={'auth_token': AUTH_TOKEN})
self.assertEqual(json.loads(response.get_data()), {'hello': 'world'})
if __name__ == "__main__":
unittest.main()
All the tests pass. Note that the tests in TestFlaskApiUsingRequests require the app to be running, whereas those in TestFlaskApi don't.
My problem is that I haven't been able to find the equivalent of requests' Session object to 'standardize' the request headers when using the test_client. This means that if I were to write more tests, I would have to pass the headers keyword argument to each request individually, which is not DRY.
How can I make a 'session' for the test_client? (It seems like this can be done with Werkzeug's EnvironBuilder but I wasn't able to quickly figure out how to do this).
In order to keep the code DRY when adding more tests, instead of using EnvironBuilder directly I wrote a decorator authorized which adds the required headers keyword argument to any function call. Then, in the test I call authorized(self.app.get) instead of self.app.get:
def authorized(function):
def wrap_function(*args, **kwargs):
kwargs['headers'] = {'auth_token': AUTH_TOKEN}
return function(*args, **kwargs)
return wrap_function
class TestFlaskApi(unittest.TestCase):
def setUp(self):
self.app = flaskapi2.app.test_client()
def test_hello_world(self):
response = self.app.get('/', headers={'auth_token': AUTH_TOKEN})
self.assertEqual(json.loads(response.get_data()), {'hello': 'world'})
def test_hello_world_authorized(self): # Same as the previous test but using a decorator
response = authorized(self.app.get)('/')
self.assertEqual(json.loads(response.get_data()), {'hello': 'world'})
The tests all pass as desired. This answer was inspired by Python decorating functions before call, How can I pass a variable in a decorator to function's argument in a decorated function?, and Flask and Werkzeug: Testing a post request with custom headers.
Update
The definition of the authorized wrapper can be made more succinct using functools.partial:
from functools import partial
def authorized(function):
return partial(function, headers={'auth_token': AUTH_TOKEN})
For every test I write that uses my apps models, I seem to have to use the current apps context:
SomeTestCase(unittest2.TestCase):
setUp(self):
self.app = Flask(__name__)
...
test_something(self):
with self.app.app_context():
# Do something
Is there a way to tell all of my tests to run using the current apps context to save me having this line in all of my tests?
I found the answer I was looking for by looking at the way the Flask-Testing extensions TestCase sets itself up, i.e., pushing the testing context to the _ctx stack inside a function that's called from within it's __call__ method.
class BaseTestCase(unittest2.TestCase):
def __call__(self, result=None):
try:
self._pre_setup()
super(BaseTestCase, self).__call__(result)
finally:
self._post_teardown()
def _pre_setup(self):
self.app = create_app()
self.client = self.app.test_client()
self._ctx = self.app.test_request_context()
self._ctx.push()
def _post_teardown(self):
if getattr(self, '_ctx') and self._ctx is not None:
self._ctx.pop()
del self._ctx
And my test:
class SomeTestCase(BaseTestCase):
test_something(self):
# Test something - we're using the right app context here
You can try something like below.
DISCLAIMER: I just came up with the idea and didn't test this solution thoroughly although it seems to work. It is also IMHO rather ugly.
from functools import wraps
def with_context(test):
#wraps(test)
def _with_context(self):
with self.app.app_context():
test(self)
return _with_context
SomeTestCase(unittest2.TestCase):
setUp(self):
self.app = Flask(__name__)
...
#with_context
test_something(self):
# Do something
Depending on how you test you may be able to use the test client. Example:
SomeTestCase(unittest2.TestCase):
setUp(self):
self.app = Flask(__name__)
self.client = self.app.text_client()
test_something(self):
response = self.client.get('/something')
# check response
I'm writing some unit tests for a Django project, and I was wondering if its possible (or necessary?) to test some of the decorators that I wrote for it.
Here is an example of a decorator that I wrote:
class login_required(object):
def __init__(self, f):
self.f = f
def __call__(self, *args):
request = args[0]
if request.user and request.user.is_authenticated():
return self.f(*args)
return redirect('/login')
Simply:
from nose.tools import assert_equal
from mock import Mock
class TestLoginRequired(object):
def test_no_user(self):
func = Mock()
decorated_func = login_required(func)
request = prepare_request_without_user()
response = decorated_func(request)
assert not func.called
# assert response is redirect
def test_bad_user(self):
func = Mock()
decorated_func = login_required(func)
request = prepare_request_with_non_authenticated_user()
response = decorated_func(request)
assert not func.called
# assert response is redirect
def test_ok(self):
func = Mock(return_value='my response')
decorated_func = login_required(func)
request = prepare_request_with_ok_user()
response = decorated_func(request)
func.assert_called_with(request)
assert_equal(response, 'my response')
The mock library helps here.
A decorator like this might be tested simply thanks to duck-typing. Just supply a mock object to the call function, that seems to hold and act as a request, and see if you get the expected behaviour.
When it is necessary to use unit tests is quite individual i'd say. The example you give contain such basic code that one might say that it isn't necessary. But then again, the cost of testing a class like this is equally low.
Example for Django's UnitTest
class TestCaseExample(TestCase):
def test_decorator(self):
request = HttpRequest()
# Set the required properties of your request
function = lambda x: x
decorator = login_required(function)
response = decorator(request)
self.assertRedirects(response)
In general, the approach I've utilized is the following:
Set up your request.
Create a dummy function to allow the decorator magic to happen (lambda). This is where you can control the number of arguments that will eventually be passed into the decorator.
Conduct an assertion based on your decorator's response.
For those looking for a django type decorator test, this is how I ran tests on my custom django decorator:
common/decorators.py
from functools import wraps
from django.http import Http404
def condition_passes_test(test_func, fail_msg=''):
"""
Decorator for views that checks that a condition passes the given test,
raising a 404 if condition fails
"""
def decorator(view_func):
#wraps(view_func)
def _wrapped_view(request, *args, **kwargs):
if test_func():
return view_func(request, *args, **kwargs)
else:
raise Http404(fail_msg)
return _wrapped_view
return decorator
The test:
import django
from django.test import TestCase
from django.http import Http404
from django.http import HttpResponse
from django.test import RequestFactory
from common import decorators
class TestCommonDecorators(TestCase):
def shortDescription(self):
return None
def test_condition_passes_test1(self):
"""
Test case where we raise a 404 b/c test function returns False
"""
def func():
return False
#decorators.condition_passes_test(func)
def a_view(request):
return HttpResponse('a response for request {}'.format(request))
request = RequestFactory().get('/a/url')
with self.assertRaises(django.http.response.Http404):
a_view(request)
def test_condition_passes_test2(self):
"""
Test case where we get 200 b/c test function returns True
"""
def func():
return True
#decorators.condition_passes_test(func)
def a_view(request):
return HttpResponse('a response for request {}'.format(request))
request = RequestFactory().get('/app_health/z')
request = a_view(request)
self.assertEquals(request.status_code, 200)
NOTE: I Am using it as such on my view where my_test_function returns True or False:
#method_decorator(condition_passes_test(my_test_function, fail_msg="Container view disabled"), name='dispatch')