Authenticate Flask unit test client from another s

2019-04-15 08:01发布

Problem:

So my problem is I have a Flask microservice want to implement the unit tests to it so when I start writing my test cases I found that I need to authenticate the unit test client because of some endpoints need authorization and here comes the problem the whole authentication system in another service this service all can do about the authentication is to validate the JWT token and get user ID from it so here is one of the views.py

from flask_restful import Resource

from common.decorators import authorize


class PointsView(Resource):
    decorators = [authorize]

    def get(self, user):
        result = {"points": user.active_points}
        return result

and authorize decorator from decorators.py

import flask
import jwt
from jwt.exceptions import DecodeError, InvalidSignatureError
from functools import wraps
from flask import request
from flask import current_app as app

from app import db
from common.models import User
from common.utils import generate_error_response

def authorize(f):
    """This decorator for validate the logged in user """

    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'Authorization' not in request.headers:
            return "Unable to log in with provided credentials.", 403

        raw_token = request.headers.get('Authorization')
        if raw_token[0:3] != 'JWT':
            return generate_error_response("Unable to log in with provided credentials.", 403)
        token = str.replace(str(raw_token), 'JWT ', '')
        try:
            data = jwt_decode_handler(token)
        except (DecodeError, InvalidSignatureError):
            return generate_error_response("Unable to log in with provided credentials.", 403)

        user = User.query.filter_by(id=int(data['user_id'])).first()
        return f(user, *args, **kwargs)

    return decorated_function

and the test case from tests.py

import unittest

from app import create_app, db
from common.models import User


class TestMixin(object):
    """
    Methods to help all or most Test Cases
    """

    def __init__(self):
        self.user = None

    """ User Fixture for testing """

    def user_test_setup(self):
        self.user = User(
            username="user1",
            active_points=0
        )
        db.session.add(self.user)
        db.session.commit()

    def user_test_teardown(self):
        db.session.query(User).delete()
        db.session.commit()


class PointsTestCase(unittest.TestCase, TestMixin):
    """This class represents the points test case"""

    def setUp(self):
        """Define test variables and initialize app."""
        self.app = create_app("testing")
        self.client = self.app.test_client
        with self.app.app_context():
            self.user_test_setup()

    def test_get_points(self):
        """Test API can create a points (GET request)"""
        res = self.client().get('/user/points/')
        self.assertEqual(res.status_code, 200)
        self.assertEquals(res.data, {"active_points": 0})

    def tearDown(self):
        with self.app.app_context():
            self.user_test_teardown()


# Make the tests conveniently executable
if __name__ == "__main__":
    unittest.main()

My authentication system work as the following:

  1. Any service (include this one) request User service to get user JWT token
  2. Any service take the JWT token decoded and get the user ID from it
  3. Get the user object from the database by his ID

so I didn't know how to make the authentication flow in the test cases.

Here is 2 solutions i was thinking about

Solution 1 (not working):

I was thinking about using patchfrom unittest.mock but the problem have to write the patch and the new decorator which patch will replace it with the old before initializing the app but the problem is i need to use db to get user in the new decorator and to use db i need the app context and here the changes i made to tests.py:

import unittest

from app import create_app, db
from common.models import User


def fake_authorize(f):
    """This decorator for validate the logged in user """

    @wraps(f)
    def decorated_function(*args, **kwargs):
        user = db.session.query(User).first()
        return f(user, *args, **kwargs)

    return decorated_function


patch('common.decorators.authorize', fake_authorize).start()


class TestMixin(object):
    """
    Methods to help all or most Test Cases
    """

    def __init__(self):
        self.user = None

    """ User Fixture for testing """

    def user_test_setup(self):
        self.user = User(
            username="user1",
            active_points=0
        )
        db.session.add(self.user)
        db.session.commit()

    def user_test_teardown(self):
        db.session.query(User).delete()
        db.session.commit()


class PointsTestCase(unittest.TestCase, TestMixin):
    """This class represents the points test case"""

    def setUp(self):
        """Define test variables and initialize app."""
        self.app = create_app("testing")
        self.client = self.app.test_client
        with self.app.app_context():
            self.user_test_setup()

    def test_get_points(self):
        """Test API can create a points (GET request)"""
        res = self.client().get('/user/points/')
        self.assertEqual(res.status_code, 200)
        self.assertEquals(res.data, {"active_points": 0})

    def tearDown(self):
        with self.app.app_context():
            self.user_test_teardown()


# Make the tests conveniently executable
if __name__ == "__main__":
    unittest.main()

the problem here is this line user = db.session.query(User).first() need the app context and patch must be called before the app initializing to work

Solution 2:

I was thinking about editing my original decorator to check if the environment was testing then get the first user and return it to the view and here is the decorator after i add this part to it :

def authorize(f):
    """This decorator for validate the logged in user """

    @wraps(f)
    def decorated_function(*args, **kwargs):
        testing = app.config.get("TESTING")
        if testing:
            user = db.session.query(User).first()
            return f(user, *args, **kwargs)

        print("real authorize working !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        if 'Authorization' not in request.headers:
            return "Unable to log in with provided credentials.", 403

        raw_token = request.headers.get('Authorization')
        if raw_token[0:3] != 'JWT':
            return generate_error_response("Unable to log in with provided credentials.", 403)
        token = str.replace(str(raw_token), 'JWT ', '')
        try:
            data = jwt_decode_handler(token)
        except (DecodeError, InvalidSignatureError):
            return generate_error_response("Unable to log in with provided credentials.", 403)

        user = User.query.filter_by(id=int(data['user_id'])).first()
        return f(user, *args, **kwargs)

    return decorated_function

my concern about this solution is i add 2 more steps in the original decorator in other words every user will pass throw 2 unnecessary lines

2条回答
一夜七次
2楼-- · 2019-04-15 08:44

Could you create some mock tokens in your unit testing framework (that your decorator can actually decode like in a real request) and send them in with your test client? An example of how that might look can be seen here: https://github.com/vimalloc/flask-jwt-extended/blob/master/tests/test_view_decorators.py#L321

孤傲高冷的网名
3楼-- · 2019-04-15 08:52

Here is just an example. I skipped some little things such as create_app, jwt.decode(token) etc. I'm sure you can understand the main approach. Structure:

src
├── __init__.py # empty
├── app.py
└── auth_example.py

app.py:

from flask import Flask

from src.auth_example import current_identity, authorize

app = Flask(__name__)


@app.route('/')
@authorize()
def main():
    """
    You can use flask_restful - doesn't matter
    Do here all what you need:
        user = User.query.filter_by(id=int(current_identity['user_id'])).first()
        etc..
    just demo - return current user_id
    """

    return current_identity['user_id']

auth_example.py:

from flask import request, _request_ctx_stack
from functools import wraps
from werkzeug.local import LocalProxy

current_identity = LocalProxy(lambda: getattr(_request_ctx_stack.top, 'current_identity', None))


def jwt_decode_handler(token):
    """
    just do here all what you need. Should return current user data
    :param str token:
    :return: dict
    """
    # return jwt.decode(token), but now - just demo
    raise Exception('just demo')


def authorize():
    def _authorize(f):

        @wraps(f)
        def __authorize(*args, **kwargs):
            if 'Authorization' not in request.headers:
                return "Unable to log in with provided credentials.", 403

            raw_token = request.headers.get('Authorization')
            if raw_token[0:3] != 'JWT':
                return "Unable to log in with provided credentials.", 403
            token = str.replace(str(raw_token), 'JWT ', '')
            try:
                # I don't know do you use Flask-JWT or not
                # this is doesn't matter - all what you need is just to mock jwt_decode_handler result 
                _request_ctx_stack.top.current_identity = jwt_decode_handler(token)
            except Exception:
                return "Unable to log in with provided credentials.", 403

            return f(*args, **kwargs)

        return __authorize
    return _authorize

Our test:

import unittest

from mock import patch

from src.app import app

app.app_context().push()


class TestExample(unittest.TestCase):

    def test_main_403(self):
        # just a demo that @authorize works fine
        result = app.test_client().get('/')
        self.assertEqual(result.status_code, 403)

    def test_main_ok(self):
        expected = '1'
        # we say that jwt_decode_handler will return {'user_id': '1'}
        patcher = patch('src.auth_example.jwt_decode_handler', return_value={'user_id': expected})
        patcher.start()
        result = app.test_client().get(
            '/',
            # send a header to skip errors in the __authorize
            headers={
                'Authorization': 'JWT=blabla',
            },
        )
        # as you can see current_identity['user_id'] is '1' (so, it was mocked in view)
        self.assertEqual(result.data, expected)
        patcher.stop()

So, in your case you need just mock jwt_decode_handler. Also I recommend do not add any additional arguments inside a decorators. It will be hard to debugging when you have more than two decorators with a different arguments, recursion, hard processing etc.

Hope this helps.

查看更多
登录 后发表回答