Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Set HTTP headers for all requests in a Flask test

Tags:

python

flask

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?

like image 589
edA-qa mort-ora-y Avatar asked May 07 '13 09:05

edA-qa mort-ora-y


5 Answers

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.

like image 132
soulseekah Avatar answered Nov 08 '22 21:11

soulseekah


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()
like image 26
Coxy Avatar answered Nov 08 '22 22:11

Coxy


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()
like image 13
DazWorrall Avatar answered Nov 08 '22 21:11

DazWorrall


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
like image 11
mexekanez Avatar answered Nov 08 '22 23:11

mexekanez


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']
like image 7
ArturM Avatar answered Nov 08 '22 21:11

ArturM