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?
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.
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()
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()
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
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']
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")
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.
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.