Set HTTP headers for all requests in a Flask test

Question

I'm using Flask and have endpoints which require authorization (and occasionally other app-specific headers). In my tests use the test_client function to create a client and then do the various get, put, delete calls. All of these calls will require authorization, and other headers to be added. How can I setup the test client to put such headers on all of the requests?

Answer 1

The Client class takes the same arguments as the EnvironBuilder class, among which is the headers keyword argument.

So you can simply use client.get( '/', headers={ ... } ) to send in your authentication.

Now if you'd like to provide a default set of headers from the client, you'd need to provide your own implementation of open which supplies a modified environment builder (akin to make_test_environ_builder) and set app.test_client_class to point to your new class.

Answer 2

Furthering the suggestion from @soulseekah, it's not too difficult to extend the test client and point your app at it. I did this recently to have a default api key in my test headers. The example given is using a py.test fixture but can easily be adapted to 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

You can wrap the WSGI app and inject headers there:

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

Thanks ArturM

using factory-boy and HTTP_AUTHORIZATION as auth method for API, fixture will looks like:

@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

You can set header inside test client.

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

Then you can use header from request:

request.headers['Authorization']

Answer 6

Building on @DazWorrall answer, and looking into the Werkzeug source code, I ended up with the following wrapper for passing default Headers that I needed for authentication:

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)

You can then use it like:

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

I needed to add an authorization header to all requests in a test, with a value depending on the test (admin user, simple user).

I didn't find how to parametrize the header (the credentials) by parametrizing the fixture creating the app, because this fixture is already parametrized to set the config class.

I did it using context variables (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

It seems a bit too much for the job, and it adds a level of indentation, but the good thing about it is that I don't have to invoke more pytest parametrization trickery to get the fixture to do what I need, and I can even change the header value in the middle of a test.