TDD Excel File Download API with Flask

Test-driven development (TDD) is a very useful development technique. However, some features are particularly difficult to test. One example is to download files.

Recently, I have just “TDDed” a file download feature in python Flask. During development, there was not a complete guide on how to TDD it. So to help save effort and time for other engineer friends, this blog will give a simplified example of using TDD to drive an excel file download API.

Flask for web and openpyxl for excel writer are used in this blog, but the logic should be similar for different frameworks and libraries.

There are four major steps in this blog:

• Setup Development Environment
• Write Test First
• Fix Failing Test
• Refactor

You are welcome to follow these steps or just view the full example at this Github Repo.

Setup Development Environment

• Create a directory

mkdir tdd_excel_download
cd tdd_excel_download

• Create a virtual environment

virtual environments help to isolate python dependencies. virtualenv is used in this blog, but you are free to create a virtual environment in any tool you like.

# create a virtual environment at destination venv
virtualenv venv
# activate virtual environment
. venv/bin/activate

• Install Dependencies

pip install Flask openpyxl

• Create a Starter Flask Project

Flask providerstest_client to facilitate the testing of the Flask app. So let’s create an empty Flask app first. Create app.py with the following content:

from flask import Flask

app = Flask(__name__)


if __name__ == '__main__':
    app.run(debug=False, host='0.0.0.0')

And we have a development environment ready. Let’s move on to write our first test.

Write Test First

Let’s create the test file tests/test_excel_file_download.py with the following content:

import unittest
from io import BytesIO

from openpyxl import load_workbook

from app import app


class TestExcelFileDownload(unittest.TestCase):
    def test_download_excel_file(self):
        # Setup / Arrange
        client = app.test_client()

        # Execute / Act
        response = client.get('/excel/download')

        # Assert
        # Assert file name is correct
        self.assertEqual(200, response.status_code)
        self.assertEqual('attachment; filename=tdd-excel.xlsx', response.headers['Content-Disposition'])

        # Assert file content is correct
        expected_file_content = [
            ['TDD', 'is', 'AWESOME!'],
            ['Except', 'when', 'it', 'is', 'particularly', 'hard'],
            ['But', 'we', 'can', 'handle', 'it'],
        ]
        wb = load_workbook(filename=BytesIO(response.data))
        ws = wb.worksheets[0]
        actual_file_content = [[cell.value for cell in row if cell.value is not None] for row in ws.rows]
        self.assertEqual(expected_file_content, actual_file_content)

Most tests consist of three parts: Setup / Arrange, Execute / Act, and Assert. In this test, the setup is to get the test_client provided by Flask app. The execute part is to make an HTTP call to the endpoint /excel/download. Finally, the assertions are to assert the HTTP response status code is 200 OK with the header attachment; filename=tdd-excel.xlsx. And the file content is the same as expected.

Run the test with the following command:

python -m unittest tests.test_excel_file_download

We should have a failing test at the first assertion: 404 != 200. This is because we don’t have the endpoint yet. So let’s fix the failing test together.

Fix Failing Test

To fix the failing test, furnish the app.py with the following content:

from io import BytesIO

from flask import Flask, send_file
from openpyxl import Workbook

app = Flask(__name__)


@app.route('/excel/download')
def download_excel():
    wb = Workbook()
    ws = wb.active
    ws.append(['TDD', 'is', 'AWESOME!'])
    ws.append(['Except', 'when', 'it', 'is', 'particularly', 'hard'])
    ws.append(['But', 'we', 'can', 'handle', 'it'])

    file_stream = BytesIO()
    wb.save(file_stream)
    file_stream.seek(0)

    return send_file(file_stream, attachment_filename="tdd-excel.xlsx", as_attachment=True)


if __name__ == '__main__':
    app.run(debug=False, host='0.0.0.0')

In more details, we have added a GET request endpoint at /excel/download. When a GET request initiates, a new Excel Workbook is created with content and saved to a byte stream. Then the byte stream is sent as a file with an attachment file name using the Flask send_file method.

Last but not least, let’s refactor a bit.

Refactor

Currently, the download_excel method is doing two things: prepare an Excel Workbook and send the file in response. Let’s extract the preparation of the Excel Workbook to a separate method.

from io import BytesIO

from flask import Flask, send_file
from openpyxl import Workbook

app = Flask(__name__)


def get_tdd_excel_workbook():
    wb = Workbook()
    ws = wb.active
    ws.append(['TDD', 'is', 'AWESOME!'])
    ws.append(['Except', 'when', 'it', 'is', 'particularly', 'hard'])
    ws.append(['But', 'we', 'can', 'handle', 'it'])
    return wb


@app.route('/excel/download')
def download_excel():
    wb = get_tdd_excel_workbook()

    file_stream = BytesIO()
    wb.save(file_stream)
    file_stream.seek(0)

    return send_file(file_stream, attachment_filename="tdd-excel.xlsx", as_attachment=True)


if __name__ == '__main__':
    app.run(debug=False, host='0.0.0.0')

Happy Coding! Maybe with a cup of coffee~ ☕☕☕