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

时间:2013-05-07 09:46:53

标签: python flask

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

7 个答案:

答案 0 :(得分:12)

您可以包装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()

答案 1 :(得分:12)

Client类使用与EnvironBuilder类相同的参数,其中包含headers关键字参数。

因此,您只需使用client.get( '/', headers={ ... } )发送身份验证即可。

现在,如果您想从客户端提供一组默认标头,则需要提供自己的open实现,该实现提供修改后的环境构建器(类似于make_test_environ_builder)并设置app.test_client_class以指向您的新班级。

答案 2 :(得分:3)

在@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")

答案 3 :(得分:3)

从@soulseekah继续提出建议,扩展测试客户端并将应用程序指向它并不困难。我最近这样做是为了在我的测试标题中有一个默认的api密钥。给出的示例是使用py.test fixture,但可以很容易地适应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()

答案 4 :(得分:1)

您可以在测试客户端内设置标题。

localhost/registration.html

然后你可以使用来自请求的标题:

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

答案 5 :(得分:1)

感谢ArturM

使用 factory-boy HTTP_AUTHORIZATION 作为API的身份验证方法,固定装置将如下所示:

@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

答案 6 :(得分:0)

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

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

我使用上下文变量(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 参数化技巧来让夹具做我需要的事情,我什至可以在测试过程中更改标题值。