2013-05-07 124 views
5

我正在使用Flask並且有需要授權的端點(偶爾還有其他應用程序特定的標頭)。在我的測試中,使用test_client函數創建一個客戶端,然後執行各種獲取,放入,刪除調用。所有這些呼叫都需要授權,並添加其他頭文件。我如何設置測試客戶端將所有請求放在這樣的頭文件中?在Flask測試中爲所有請求設置HTTP標頭

回答

11

您可以包裝WSGI應用程序,並有注入頭:

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() 
2

大廈@DazWorrall答案,並尋找到WERKZEUG源代碼,我結束了以下包裝傳遞,我需要進行身份驗證默認頭:

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) 

然後,您可以使用它像:

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") 
0

您可以在測試客戶端內設置標題。

client = app.test_client() 
client.environ_base['HTTP_AUTHORIZATION'] = 'Bearer your_token' 

然後你可以使用頭從請求:

request.headers['Authorization']