Flask Uploading files 一章就已經提供上傳檔案的範例程式,不過並未提及測試的部分,因此本文特別紀錄 Flask 檔案上傳與測試的範例。
本文環境
- Python 3.6.5
- macOS 10.14
- brew 2.1.7
- python-magic 0.4.15
- pytest 4.3.1
$ brew install libmagic
$ pip install python-magic pytest
範例
以下 app.py
包含首頁、檔案上傳、檔案上傳成功 3 個頁面,其中 upload()
本文最重要的部分,容後說明。
app.py
import os
import magic
from flask import Flask, request, redirect, url_for
from werkzeug.utils import secure_filename
ALLOWED_EXTENSIONS = {'jpg', 'jpeg'}
ALLOWED_MIME_TYPES = {'image/jpeg'}
app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024 # 16MB
@app.route('/', methods=['GET'])
def index():
return (
'<!doctype html>'
'<title>Upload File</title>'
'<h1>Upload File</h1>'
'<form method="post" enctype="multipart/form-data" action="/upload">'
'<input type="file" name="file">'
'<input type="submit" value="Upload">'
'</form>'
)
def is_allowed_file(file):
if '.' in file.filename:
ext = file.filename.rsplit('.', 1)[1].lower()
else:
return False
mime_type = magic.from_buffer(file.stream.read(), mime=True)
if (
mime_type in ALLOWED_MIME_TYPES and
ext in ALLOWED_EXTENSIONS
):
# move the cursor to the beginning
file.stream.seek(0,0)
return True
return False
@app.route('/upload', methods=['POST'])
def upload():
file = request.files['file']
if file and is_allowed_file(file):
filename = secure_filename(file.filename)
file.save(os.path.join('/tmp', filename))
return redirect(url_for('success'))
return redirect(url_for('index'))
@app.route('/success', methods=['GET'])
def success():
return (
'<!doctype html>'
'<title>Success</title>'
'<h1>Success</h1>'
)
上述範例可用以下指令執行:
$ env FLASK_APP=app.py flask run
首先, Flask 檔案上傳有 1 個重要設定 MAX_CONTENT_LENGTH
,該設定可限制檔案大小,如果超過檔案大小限制, Flask 將會拋出 RequestEntityTooLarge 的例外錯誤。
接著 upload() 函數中用 request.files['file']
取得使用者上傳的檔案, 鍵值(key) file
對應到的就是 index()
中的 name="file"
的 input 欄位。
為了確保上傳的檔案的安全性(例如使用者可能上傳病毒或其他奇怪檔案),正常來說,後端都應該檢查副檔名以及其 MIME type, 因此 is_allowed_file()
會檢查檔案的副檔名及 MIME type 。
通過 is_allowed_file()
檢查後,由於要將檔案存至檔案系統(file system)中,我們將檔名以 secure_filename()
函式進行過濾,避免檔名為類似 /../../../filename
的 Directory traversal attack 發生。
p.s. 如果不保留原始檔名,則可以忽略 secure_filename()
此步驟
最後才將檔案存至 /tmp
底下。
以上就是簡單的 Flask 檔案上傳範例。
測試
測試檔案上傳可使用 pytest 搭配 Flask test client 進行測試。
test_upload_file.py
import os
from app import app
def test_upload():
app.config['TESTING'] = True
with app.app_context():
client = app.test_client()
file = open('./test_file.jpg', 'rb')
resp = client.post(
'/upload',
content_type='multipart/form-data',
data={
'file': (file, 'test_file.jpg'),
},
)
assert resp.status_code >= 200
assert os.path.exists('/tmp/test_file.jpg') is True
上述執行測試指令為:
$ py.test test_upload_file.py
上述測試重點在於將 content_type
設定為 multipart/form-data
,並將要上傳的檔案指定為 (file, filename)
的 tuple 。
以上就是 Flask 檔案上傳及測試範例。
References
http://flask.palletsprojects.com/en/1.1.x/patterns/fileuploads/