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() 函式進行過濾,避免檔名為類似 /../../../filenameDirectory 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/