How to unit-testing Python Flask App
Lets say you have some production code which you want to test, you might wonder what is a good way to test it.
Here is some example how I have been doing it, hopefully it will be useful (ps. This is just one of many ways to do it)
Say you have some REST API routes
main.py
@app.route('/environment', method=['GET'])
def get_current_environment():
''' Get current application environment '''
return current_app.get('environment')
@app.route('/environment/<type>', method=['POST'])
def log_message(type):
''' Log message in the requested level
@param type: log level
'''
if type not in ('error', 'warn', 'info', 'debug', 'trace'):
return (False, 'Not a valid log type')
param = request.json()
message = param .get('message')
if not message:
return (False, 'Invalid payload given')
return (True, msg)
@app.route('/sum/', method=['PUT'])
def sum_data():
''' Load data from text file and print the sum of all number '''
# Assuming it is a valid file (to make the guidance short and clear)
file_data = request.form.get('file', None)
sum = 0
if not file_data:
return sum
# do some parsing of the file
return sum
Then to test the Flask routes you can either use:
- Using Flask
url_for
method to let it to decide which function to use. (Reason: Maybe routes will change but function name don’t change, which in my opinion is bad design because for API, route usually dont change a lot of times. - Using actual routes. (Reason: Most of the time function name will get renamed but routes stay the same)
test_main.py
import unittest
from flask import url_for
class TestMain(unittest.TestCase):
def test_get_current_environment(self):
resp = self.client.get(url_for('<module_name>.get_current_environment'))
# or you can do it this way
resp = self.client.get('/environment')
assert resp.data == <app env>, msg='<some debug msg, know why/where>'
def test_log_message(self):
payload = <some_payload_here>
resp = self.client.post(url_for('<module_name>.log_message', level='debug'),
data=json.dumps(payload), content_type='application/json')
# or you can do it this way too (its your option which one is better)
resp = self.client.post('/environment/debug', data=json.dumps(payload),
content_type='application/json')
self.assertEqual(resp.data, <valid resp data>
def test_sum_data(self):
with open('count_sum.txt', 'r') as file:
payload = {'file': file}
resp = self.client.put('/sum', data=payload)
# some assertion
You can do assertion however you want, I recommend using built in python unittest framework.
Note that to be able to use flask url_for
you have to activate the request context which basically you can set with self.app.test_request_context()
Lets go to another case when you wanted to modify flask session
to speed up testing.
Mocking session cookie data is one good way of doing it where the session cookie might store something like user_id
for authorization checking,
you can access and modify flask session as shown here
with app.test_client() as c:
with c.session_transaction() as sess:
sess['a_key'] = 'a value'
# once this is reached the session was stored
Other Guide that I find useful for doing unit-test using mock
What do you guys think? If you found any mistake feel free to comment below :)