Question

基于我之前的一个问题（How to unit test a Flask RESTful API），我试图在没有应用运行的情况下使用test_client测试Flask RESTful API，而不是使用requests应用程序正在运行。

作为一个简单示例，我有一个API（flaskapi2.py），其get函数使用登录装饰器：

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)

随着应用程序的运行，我在同一目录中运行这些单元测试（test_flaskapi2.py）：

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

所有测试都通过了。请注意，TestFlaskApiUsingRequests中的测试要求应用运行，而TestFlaskApi中的测试则不需要。

我的问题是，在使用requests时，我无法找到等效的test_client'Session object来“标准化”请求标头。这意味着如果我要编写更多测试，我必须单独将headers关键字参数传递给每个请求，这不是DRY。

如何为test_client制作“会话”？（看来这可以用Werkzeug的EnvironBuilder完成，但我无法快速弄清楚如何做到这一点。

Answer 1

为了在添加更多测试时保持代码DRY，而不是直接使用EnvironBuilder，我编写了一个装饰器authorized，它将所需的headers关键字参数添加到任何函数调用中。然后，在测试中，我拨打authorized(self.app.get)而不是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'})

测试全部按需通过。这个答案的灵感来自Python decorating functions before call，How can I pass a variable in a decorator to function's argument in a decorated function?和Flask and Werkzeug: Testing a post request with custom headers。

<强>更新

使用functools.partial可以使authorized包装器的定义更加简洁：

from functools import partial
def authorized(function):
    return partial(function, headers={'auth_token': AUTH_TOKEN})

使用test_client

1 个答案: