Setting HTTP Request Values For Flask Unit Testing

I recently had to do some unit testing for a Flask application code that looked at the user agent and IP address of the client. The first thing that you will realize when you try to get the User Agent from the request headers when you use the test_client() provided by Flask, is that the underlying Werkzeug library raises a KeyError saying that it can't find a value for the user agent.
For example, let's say that you have some code in your app that looks like this:

@app.route('/login', methods=['POST', 'OPTIONS'])
def login():
    print request.remote_addr
    print request.headers['User-Agent']

And then in your testing code you do something like:

client = app.test_client()
client.post('/login',
            data=json.dumps({
                'username': 'sheldon@cooper.com',
                'password': 'howimetyourmother'
            }), content_type='application/json')

Most likely, you will then get an error that looks something like:

File "/Users/leonard/app/rest/app.py", line 515, in login
    if request.headers['User-Agent']:
File "/Users/leonard/.virtualenvs/app/lib/python2.7/site-packages/werkzeug/datastructures.py", line 1229, in __getitem__
    return self.environ['HTTP_' + key]
    KeyError: 'HTTP_USER_AGENT'

As it turns out, there are two solutions to this problems that I found on StackOverflow (1, 2). This first option is to set the request environment variables on each call using the environ_base parameter provided by Werkzeug (more details on that here). For example, we can do something like:

client.post('/login',
            data=json.dumps({
                'username': 'sheldon@cooper.com',
                'password': 'howimetyourmother'
            }), content_type='application/json',
            environ_base={
                'HTTP_USER_AGENT': 'Chrome',
                'REMOTE_ADDR': '127.0.0.1'
            })

Although this method works perfectly well, it requires us to set the values on each call. If we have a really long unit testing script that makes many calls (as I do), then this becomes rather burdensome. The second solution is to create a proxy for the Flask app that overrides the __call__ method. So, the proxy looks something like:

class FlaskTestClientProxy(object):
    def __init__(self, app):
        self.app = app

    def __call__(self, environ, start_response):
        environ['REMOTE_ADDR'] = environ.get('REMOTE_ADDR', '127.0.0.1')
        environ['HTTP_USER_AGENT'] = environ.get('HTTP_USER_AGENT', 'Chrome')
        return self.app(environ, start_response)

Then when we create our test client, we simply make sure to wrap the WSGI app with our proxy, and voilà, it works.

app.wsgi_app = FlaskTestClientProxy(app.wsgi_app)
client = app.test_client()
client.post('/login',
            data=json.dumps({
                'username': 'sheldon@cooper.com',
                'password': 'howimetyourmother'
            }), content_type='application/json')