为Flask应用单元测试设置(模拟)请求头

28 投票
2 回答
19465 浏览
提问于 2025-04-17 18:18

有没有人知道在单元测试时,如何设置(模拟)FLask(Werkzeug)提供的请求对象的User-Agent?

目前,当我尝试获取像request.headers['User-Agent']这样的详细信息时,会出现KeyError错误,因为Flask的test_client()并没有设置这些内容。(下面是部分错误追踪信息)

在Flask项目的单元测试中尝试从请求对象获取User-Agent时,会出现KeyError错误。

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'

-- 更新 --

除了下面的(被接受的)解决方案外,environ_base的提示让我找到了另一个SO解决方案。这个解决方案的思路是为Flask应用创建一个包装类,并重写call方法,以自动设置环境变量。这样,所有调用时这些变量都会被设置。因此,我最终实现的解决方案是创建这个代理类:

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)

然后用这个代理包装WSGI容器:

app.wsgi_app = FlaskTestClientProxy(app.wsgi_app)
test_client = app.test_client()

2 个回答

8

虽然@chris-mckinnel写的那个方法可以用,但我不建议在这种情况下使用environ_base

你可以很简单地像下面这样做:

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

这样做可以达到目的,同时也让你的代码更清晰——清晰的代码比模糊的好。

如果你想知道可以传给test_request_context的内容,可以参考EnvironBuilder的定义;你可以在这里找到相关信息。

31

当你调用 get()post() 时,需要传入 environ_base。比如:

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

这样一来,你的 request.user_agent 就会是你传入的内容,你可以通过 request.headers['User-Agent'] 来访问它。

想了解更多信息,可以查看 http://werkzeug.pocoo.org/docs/test/#testing-api

撰写回答