Skip to main content
 首页 » 编程设计

python之从另一个服务(微服务架构)验证 Flask 单元测试客户端

2024年09月03日20小虾米

问题:

所以我的问题是我有一个 Flask 微服务想要对其实现单元测试,所以当我开始编写测试用例时,我发现我需要对单元测试客户端进行身份验证,因为某些端点需要授权,问题就来了另一个服务中的整个身份验证系统该服务可以做的关于身份验证的所有事情都是验证 JWT token 并从中获取用户 ID 所以这里是 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 
 

并从 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 

和来自 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() 
 

我的身份验证系统工作如下:

  1. 任何服务(包括此服务)请求用户服务以获取用户 JWT token
  2. 任何服务都会解码 JWT token 并获取用户 ID 从中
  3. 通过ID从数据库中获取用户对象

所以我不知道如何在测试用例中进行身份验证。

请您参考如下方法:

这只是一个例子。 create_app, jwt.decode(token) 等一些小东西我跳过了。我相信你能理解主要的方法。结构:

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

应用程序.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 

我们的测试:

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() 

因此,在您的情况下,您只需要模拟 jwt_decode_handler。另外我建议不要在装饰器中添加任何额外的参数。当你有两个以上具有不同参数、递归、硬处理等的装饰器时,将很难调试。

希望这对您有所帮助。