在Flask測試中爲所有請求設置HTTP標頭

Question

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

Answer 1

您可以包裝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() 

Answer 2

的Client類取相同的參數EnvironBuilder類，其中之一是headers關鍵字參數。

因此，您可以簡單地使用client.get('/', headers={ ... })發送您的身份驗證。

現在，如果你想提供從客戶端頭的默認設置，你需要提供自己實現open它提供一個修改後的環境生成器（類似於make_test_environ_builder），並設置app.test_client_class指向您新班。

Answer 3

大廈@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") 

Answer 4

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

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

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

request.headers['Authorization']