28

Does anyone know of a way to set (mock) the User-Agent of the request object provided by FLask (Werkzeug) during unit testing?

As it currently stands, when I attempt to obtain details such as the request.headers['User-Agent'] a KeyError is raised as the Flask test_client() doesn't set these up. (See partial stack trace below)

When attempting to get the User-Agent from the request object in a Flask project during unit testing, a KeyError is raised.

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

-- UPDATE --

Along with the (accepted) solution below, the environ_base hint lead me to this other SO solution. The premise of this solution is to create a wrapper class for the Flask app and override the call method to automatically set the environment variables. This way, the variables are set for all calls. So, the solution I ended up implementing is creating this proxy class:

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)

And then wrapping the WSGI container with that proxy:

app.wsgi_app = FlaskTestClientProxy(app.wsgi_app)
test_client = app.test_client()
Community
  • 1
  • 1
prschmid
  • 498
  • 1
  • 6
  • 9
  • Your solution does not work for me, while Chris' does. I am using a pytest fixture for testing so that may have something to do with it: app = create_app(config_class=config.TestConfig); testing_client = app.test_client(); ctx = app.test_request_context() ; ctx.push(); yield testing_client; ctx.pop() – marvin Dec 16 '19 at 15:27

2 Answers2

31

You need to pass in environ_base when you call get() or post(). E.g.,

client = app.test_client()
response = client.get('/your/url/', 
                      environ_base={'HTTP_USER_AGENT': 'Chrome, etc'})

Then your request.user_agent should be whatever you pass in, and you can access it via request.headers['User-Agent'].

See http://werkzeug.pocoo.org/docs/test/#testing-api for more info.

Chris McKinnel
  • 14,694
  • 6
  • 64
  • 67
  • Thanks! Although I ended up using the solution I included in my original question above, this would have worked perfectly fine as well. – prschmid Mar 07 '13 at 21:43
  • Nice one, I think that solution is far more elegant if you need to set those environment variables on all your tests... In fact I think you've convinced me to change my Flask apps to do the same thing! – Chris McKinnel Mar 07 '13 at 21:48
  • This also works in `app.test_request_context` (maybe useful if you're trying to write unit tests instead of integration tests) – gatoatigrado Oct 07 '13 at 02:33
  • 1
    How do we override 'Host' i.e. ```request.host``` using this method? – 4Z4T4R Sep 02 '14 at 21:00
  • 3
    RE: @toszter -- Pass in ```headers``` containing a list of tuples e.g. ```headers=[('Host','unit-tester'),]``` -- ```environ_override``` doesn't work, nor did ```environ_base```. :( – 4Z4T4R Sep 02 '14 at 21:14
  • There are some inconsistencies around the expected format of these keys. AFAICT `HTTP_USER_AGENT` is a Flask special case that uses all underscores, and everything else uses either e.g. `HTTP_CONTENT-TYPE` or `CONTENT_TYPE`. Because of the way the keys are updated, a different tuple is created for each key format, so you may need to one or the other depending on context. e.g. to override `Content-Length` you must use `CONTENT_LENGTH`, not `HTTP_CONTENT-LENGTH`. See https://github.com/pallets/werkzeug/blob/810cd659153a7c9923e2557a00c587c8bc77d74a/src/werkzeug/test.py#L763-L817 – wst Dec 13 '22 at 23:49
8

Even though what @chris-mckinnel wrote works, I wouldn't opt to use environ_base in this case.

You can easily do something as shown below:

with app.test_request_context('url', headers={'User-Agent': 'Your value'}):
    pass

This should do the job and at the same time it keeps your code explicit - explicit is better than implicit.

If you want to know what you can pass to the test_request_context reference the EnvironBuilder definition; which can be found here.

Julian Camilleri
  • 2,975
  • 1
  • 25
  • 34