如何用Python Unittest测试Token校验装饰器？

我需要测试一个JWT认证装饰器的功能，覆盖「Token缺失」和「无效Token」等场景。以下是相关代码：

装饰器代码（decorator.py） #

from functools import wraps
from flask import request, jsonify, make_response
import jwt

def token_required(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        data_header = request.headers['Authorization']
        token = str.replace(str(data_header), 'Bearer ', '')
        if not token:
            return jsonify({'Message':'Token is missing'}), 401
        try:    
            data = jwt.decode(token, "mysecret", algorithms=['HS256'])
        except jwt.InvalidTokenError:
            return make_response({'Message':'Invalid Token!'}, 403)
        return f(data, *args, **kwargs)

    return decorated

接口代码（get_user_uid.py） #

from flask import Blueprint
from your_module import token_required, service_profile

getuid_bp = Blueprint('getuid', __name__)

@getuid_bp.route('/users/<uid>', methods=['GET'])
@token_required
def profile_view(data, uid):
    profile_data = service_profile(data, uid)
    return profile_data

profile_view(data, uid)中的data是JWT解码后的值。我尝试了下面的测试代码，但不知道如何正确设置profile_view的data参数，也不确定mock装饰器里的request是否正确：

我尝试的测试代码 #

import mock
import jwt

def test_service_user_profile(self):
    data = {                                                              
        "address": "Ahmedabad",
        "dob": "2022-06-07T23:16:20",
        "email_address": "testuser58@example.com",
        "password": "sha256$WbLrMDWEUHKvLDSC$2a051f8689222fc9815cd00da2c3260e114c313ba20edebf8df17e6d409721ac",
        "uid": "75f9d2c5570f4c7bb5e28c5807f92c48",
        "user_name": "testuser58"
    }

    encoded_jwt = jwt.encode(data, "mysecret", algorithm="HS256")
    decoded_jwt = jwt.decode(encoded_jwt, "mysecret", algorithms="HS256")

    data_header = f'Bearer {encoded_jwt}'

    m1 = mock.MagicMock()
    m1.headers.return_value = data_header
    print("m", m1)
        
    response = profile_view(data,"75f9d2c5570f4c7bb5e28c5807f92c48")
    # mock_profile.assert_called_once()

测试Flask装饰器的核心是模拟请求上下文，因为装饰器依赖request对象。直接调用profile_view会绕过装饰器逻辑，所以推荐通过Flask测试客户端发起请求，或手动mockrequest对象后调用被装饰的函数。

方法一：使用Flask测试客户端（推荐，接近真实请求） #

这种方式无需手动mockrequest，直接模拟真实HTTP请求：

import unittest
from flask import Flask
from get_user_uid import getuid_bp
import jwt
from unittest.mock import patch

class TestTokenDecorator(unittest.TestCase):
    def setUp(self):
        self.app = Flask(__name__)
        self.app.register_blueprint(getuid_bp)
        self.client = self.app.test_client()
        self.valid_user_data = {                                                              
            "address": "Ahmedabad",
            "dob": "2022-06-07T23:16:20",
            "email_address": "testuser58@example.com",
            "password": "sha256$WbLrMDWEUHKvLDSC$2a051f8689222fc9815cd00da2c3260e114c313ba20edebf8df17e6d409721ac",
            "uid": "75f9d2c5570f4c7bb5e28c5807f92c48",
            "user_name": "testuser58"
        }
        self.valid_token = jwt.encode(self.valid_user_data, "mysecret", algorithm="HS256")

    # 测试Token缺失场景
    def test_token_missing(self):
        response = self.client.get('/users/75f9d2c5570f4c7bb5e28c5807f92c48')
        self.assertEqual(response.status_code, 401)
        self.assertEqual(response.json, {'Message':'Token is missing'})

    # 测试无效Token场景
    def test_invalid_token(self):
        invalid_token = jwt.encode(self.valid_user_data, "wrongsecret", algorithm="HS256")
        response = self.client.get(
            '/users/75f9d2c5570f4c7bb5e28c5807f92c48',
            headers={'Authorization': f'Bearer {invalid_token}'}
        )
        self.assertEqual(response.status_code, 403)
        self.assertEqual(response.json, {'Message':'Invalid Token!'})

    # 测试有效Token场景
    def test_valid_token(self):
        # Mock业务逻辑函数，避免依赖真实数据库/服务
        with patch('get_user_uid.service_profile') as mock_service:
            mock_service.return_value = {'status': 'success', 'data': self.valid_user_data}
            response = self.client.get(
                '/users/75f9d2c5570f4c7bb5e28c5807f92c48',
                headers={'Authorization': f'Bearer {self.valid_token}'}
            )
            self.assertEqual(response.status_code, 200)
            # 验证业务函数是否收到了解码后的JWT数据
            mock_service.assert_called_once_with(self.valid_user_data, "75f9d2c5570f4c7bb5e28c5807f92c48")

方法二：手动mock request对象（单独测试装饰器逻辑） #

如果只想测试装饰器本身的逻辑，可以直接mockrequest对象：

import unittest
from unittest.mock import patch
from decorator import token_required
import jwt

class TestTokenDecorator(unittest.TestCase):
    def setUp(self):
        self.valid_user_data = {                                                              
            "address": "Ahmedabad",
            "dob": "2022-06-07T23:16:20",
            "email_address": "testuser58@example.com",
            "password": "sha256$WbLrMDWEUHKvLDSC$2a051f8689222fc9815cd00da2c3260e114c313ba20edebf8df17e6d409721ac",
            "uid": "75f9d2c5570f4c7bb5e28c5807f92c48",
            "user_name": "testuser58"
        }
        self.valid_token = jwt.encode(self.valid_user_data, "mysecret", algorithm="HS256")

    # 测试Token缺失
    def test_token_missing(self):
        # 定义一个测试用的被装饰函数
        @token_required
        def test_func(data, uid):
            return data, uid

        # Mock request.headers，不提供Authorization
        with patch('flask.request') as mock_request:
            mock_request.headers = {}
            response, status_code = test_func(None, "test_uid")
            self.assertEqual(status_code, 401)
            self.assertEqual(response, {'Message':'Token is missing'})

    # 测试无效Token
    def test_invalid_token(self):
        @token_required
        def test_func(data, uid):
            return data, uid

        invalid_token = jwt.encode(self.valid_user_data, "wrongsecret", algorithm="HS256")
        with patch('flask.request') as mock_request:
            mock_request.headers = {'Authorization': f'Bearer {invalid_token}'}
            response, status_code = test_func(None, "test_uid")
            self.assertEqual(status_code, 403)
            self.assertEqual(response, {'Message':'Invalid Token!'})

    # 测试有效Token
    def test_valid_token(self):
        @token_required
        def test_func(data, uid):
            return data, uid

        with patch('flask.request') as mock_request:
            mock_request.headers = {'Authorization': f'Bearer {self.valid_token}'}
            data, uid = test_func(None, "test_uid")
            self.assertEqual(data, self.valid_user_data)
            self.assertEqual(uid, "test_uid")

内容的提问来源于stack exchange，提问作者Naveen Pandia