Equivalent of a requests Session object when unit testing a Flask RESTful API using a test_client

Question

Building on a previous question of mine (How to unit test a Flask RESTful API), I'm trying to test a Flask RESTful API using a test_client without the app running, rather than using requests while the app is running.

As a simple example, I have an API (flaskapi2.py) with a get function which uses a login decorator:

import flask
import flask_restful
from functools import wraps

app = flask.Flask(__name__)
api = flask_restful.Api(app)

AUTH_TOKEN = "foobar"

def login_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if flask.request.headers.get("auth_token") == AUTH_TOKEN:
            return f(*args, **kwargs)
        else:
            return flask.abort(401)     # Return HTTP status code for 'Unauthorized'
    return decorated_function

class HelloWorld(flask_restful.Resource):
    @login_required
    def get(self):
        return {'hello': 'world'}

api.add_resource(HelloWorld, '/')

if __name__ == "__main__":
    app.run(debug=True)

With the app running, I run these unit tests (test_flaskapi2.py in the same directory):

import unittest
import flaskapi2
import requests
import json

AUTH_TOKEN = "foobar"

class TestFlaskApiUsingRequests(unittest.TestCase):
    def setUp(self):
        self.session = requests.Session()
        self.session.headers.update({'auth_token': AUTH_TOKEN})

    def test_hello_world(self):
        response = self.session.get('http://localhost:5000')
        self.assertEqual(response.json(), {'hello': 'world'})

    def test_hello_world_does_not_work_without_login(self):
        response = requests.get('http://localhost:5000')        # Make an unauthorized GET request
        self.assertEqual(response.status_code, 401)             # The HTTP status code received should be 401 'Unauthorized'


class TestFlaskApi(unittest.TestCase):
    def setUp(self):
        self.app = flaskapi2.app.test_client()

    def test_hello_world(self):
        response = self.app.get('/', headers={'auth_token': AUTH_TOKEN})
        self.assertEqual(json.loads(response.get_data()), {'hello': 'world'})


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

All the tests pass. Note that the tests in TestFlaskApiUsingRequests require the app to be running, whereas those in TestFlaskApi don't.

My problem is that I haven't been able to find the equivalent of requests' Session object to 'standardize' the request headers when using the test_client. This means that if I were to write more tests, I would have to pass the headers keyword argument to each request individually, which is not DRY.

How can I make a 'session' for the test_client? (It seems like this can be done with Werkzeug's EnvironBuilder but I wasn't able to quickly figure out how to do this).

Answer 1

In order to keep the code DRY when adding more tests, instead of using EnvironBuilder directly I wrote a decorator authorized which adds the required headers keyword argument to any function call. Then, in the test I call authorized(self.app.get) instead of self.app.get:

def authorized(function):
    def wrap_function(*args, **kwargs):
        kwargs['headers'] = {'auth_token': AUTH_TOKEN}
        return function(*args, **kwargs)
    return wrap_function

class TestFlaskApi(unittest.TestCase):
    def setUp(self):
        self.app = flaskapi2.app.test_client()

    def test_hello_world(self):
        response = self.app.get('/', headers={'auth_token': AUTH_TOKEN})
        self.assertEqual(json.loads(response.get_data()), {'hello': 'world'})

    def test_hello_world_authorized(self):          # Same as the previous test but using a decorator
        response = authorized(self.app.get)('/')
        self.assertEqual(json.loads(response.get_data()), {'hello': 'world'})

The tests all pass as desired. This answer was inspired by Python decorating functions before call, How can I pass a variable in a decorator to function's argument in a decorated function?, and Flask and Werkzeug: Testing a post request with custom headers.

Update

The definition of the authorized wrapper can be made more succinct using functools.partial:

from functools import partial
def authorized(function):
    return partial(function, headers={'auth_token': AUTH_TOKEN})