在 Flask 测试中为所有请求设置 HTTP 标头

Question

我正在使用 Flask 并且有需要授权的端点（有时还有其他特定于应用程序的标头）。在我的测试中，使用test_client 函数创建一个客户端，然后执行各种get、put、delete 调用。所有这些调用都需要授权，并且需要添加其他标头。如何设置测试客户端以将此类标头放在所有请求上？

Answer 1

Client 类与EnvironBuilder 类采用相同的参数，其中headers 关键字参数。

所以您可以简单地使用client.get( '/', headers={ ... } ) 发送您的身份验证。

现在，如果您想从客户端提供一组默认标头，您需要提供自己的 open 实现，它提供修改后的环境构建器（类似于 make_test_environ_builder）并设置 @987654326 @ 指向你的新类。

Answer 2

根据@soulseekah 的建议，扩展测试客户端并将您的应用程序指向它并不难。我最近这样做是为了在我的测试标头中有一个默认的 api 密钥。给出的示例使用 py.test 夹具，但可以轻松适应 unittest/nosetests。

from flask import testing
from werkzeug.datastructures import Headers


class TestClient(testing.FlaskClient):
    def open(self, *args, **kwargs):
        api_key_headers = Headers({
            'x-api-key': 'TEST-API-KEY'
        })
        headers = kwargs.pop('headers', Headers())
        headers.extend(api_key_headers)
        kwargs['headers'] = headers
        return super().open(*args, **kwargs)


@pytest.fixture(scope='session')
def test_client(app):
    app.test_client_class = TestClient
    return app.test_client()

Answer 3

您可以包装 WSGI 应用并在其中注入标头：

from flask import Flask, request
import unittest

def create_app():
    app = Flask(__name__)

    @app.route('/')
    def index():
        return request.headers.get('Custom', '')

    return app

class TestAppWrapper(object):

    def __init__(self, app):
        self.app = app

    def __call__(self, environ, start_response):
        environ['HTTP_CUSTOM'] = 'Foo'
        return self.app(environ, start_response)


class Test(unittest.TestCase):

    def setUp(self):
        self.app = create_app()
        self.app.wsgi_app = TestAppWrapper(self.app.wsgi_app)
        self.client = self.app.test_client()

    def test_header(self):
        resp = self.client.get('/')
        self.assertEqual('Foo', resp.data)


if __name__ == '__main__':
    unittest.main()

Answer 4

感谢阿图姆

使用 factory-boy 和 HTTP_AUTHORIZATION 作为 API 的身份验证方法，fixture 将如下所示：

@pytest.fixture(scope='function')
def test_client(flask_app):
    def get_user():
        user = UserDataFactory()
        db.session.commit()
        return user

    token = get_user().get_auth_token()
    client = app.test_client()
    client.environ_base['HTTP_AUTHORIZATION'] = 'Bearer ' + token
    return client

Answer 5

您可以在测试客户端中设置标头。

client = app.test_client()
client.environ_base['HTTP_AUTHORIZATION'] = 'Bearer your_token'

然后你可以使用请求头：

request.headers['Authorization']

Answer 6

基于@DazWorrall 的回答，并查看 Werkzeug 源代码，我最终得到了以下包装器，用于传递身份验证所需的默认标头：

class TestAppWrapper:
    """ This lets the user define custom defaults for the test client.
    """

    def build_header_dict(self):
        """ Inspired from : https://github.com/pallets/werkzeug/blob/master/werkzeug/test.py#L591 """
        header_dict = {}
        for key, value in self._default_headers.items():
            new_key = 'HTTP_%s' % key.upper().replace('-', '_')
            header_dict[new_key] = value
        return header_dict

    def __init__(self, app, default_headers={}):
        self.app = app
        self._default_headers = default_headers

    def __call__(self, environ, start_response):
        new_environ = self.build_header_dict()
        new_environ.update(environ)
        return self.app(new_environ, start_response)

然后你可以像这样使用它：

class BaseControllerTest(unittest.TestCase):

    def setUp(self):
        _, headers = self.get_user_and_auth_headers() # Something like: {'Authorization': 'Bearer eyJhbGciOiJ...'}
        app.wsgi_app = TestAppWrapper(app.wsgi_app, headers)
        self.app = app.test_client()

    def test_some_request(self):
        response = self.app.get("/some_endpoint_that_needs_authentication_header")

Answer 7

我需要为测试中的所有请求添加一个授权标头，其值取决于测试（管理员用户、简单用户）。

我没有找到如何通过参数化创建应用程序的夹具来参数化标头（凭据），因为此夹具已经参数化以设置配置类。

我使用 上下文变量 (Python 3.7+) 做到了。

tests/__init__.py

# Empty. Needed to import common.py.

tests/common.py

from contextvars import ContextVar
from contextlib import AbstractContextManager

from my_application.settings import Config


# Unrelated part creating config classes
class TestConfig(Config):
    TESTING = True
    AUTH_ENABLED = False


class AuthTestConfig(TestConfig):
    AUTH_ENABLED = True


# "Interesting" part creating the context variable...
AUTH_HEADER = ContextVar("auth_header", default=None)


# ... and the context manager to use it
class AuthHeader(AbstractContextManager):
    def __init__(self, creds):
        self.creds = creds

    def __enter__(self):
        self.token = AUTH_HEADER.set('Basic ' + self.creds)

    def __exit__(self, *args, **kwargs):
        AUTH_HEADER.reset(self.token)

conftest.py

import flask.testing

from my_application import create_app

from tests.common import TestConfig, AUTH_HEADER


class TestClient(flask.testing.FlaskClient):
    def open(self, *args, **kwargs):
        auth_header = AUTH_HEADER.get()
        if auth_header:
            (
                kwargs
                .setdefault("headers", {})
                .setdefault("Authorization", auth_header)
            )
        return super().open(*args, **kwargs)


@pytest.fixture(params=(TestConfig, ))
def app(request, database):
    application = create_app(request.param)
    application.test_client_class = TestClient
    yield application

test_users.py

import pytest

from tests.common import AuthTestConfig, AuthHeader


class TestUsersApi:

    # app fixture parametrization is used to set the config class
    @pytest.mark.parametrize("app", (AuthTestConfig, ), indirect=True)
    def test_users_as_admin_api(self, app):

        client = app.test_client()

        # Calling the context manager to specify the credentials for the auth header
        creds = ...  # Define credentials here
        with AuthHeader(creds):

            ret = client.get(/users/)
            assert ret.status_code == 200

这项工作似乎有点太多了，而且它增加了一定程度的缩进，但它的好处是我不必调用更多的 pytest 参数化技巧来让夹具完成我需要的工作，我什至可以在测试过程中更改标头值。