Unit test case for file upload flask - python-3.x

Question

I have created a flask application, where I am uploading a file and then predicting the type of the file. I want to write unit test case for the same. I am new to unit test in python and therefore very confused!. There are 2 parts to my code, the first is the Main function, which then calls the classification method.
main.py - here the file is being uploaded and then we call the func_predict function which returns the output
upload_parser = api.parser()
upload_parser.add_argument('file', location='files',
type=FileStorage, required=True)
#api.route('/classification')
#api.expect(upload_parser)
class classification(Resource):
def post(self):
"""
predict the document
"""
args = upload_parser.parse_args()
uploaded_file = args['file']
filename = uploaded_file.filename
prediction,confidence = func_predict(uploaded_file)
return {'file_name':filename,'prediction': prediction,'confidence':confidence}, 201
predict.py : this file contains the func_predict function which does the actual prediction work. It takes the uploaded file as an input
def func_predict(file):
filename = file.filename #filename
extension = os.path.splitext(filename)[1][1:].lower() #file_extension
path = os.path.join(UPLOAD_FOLDER, filename) #store the temporary path of the file
output = {}
try:
# Does some processing.... some lines which are not relevant and then returns the two values
return (''.join(y_pred),max_prob)
Now my confusion is, How do i mock the uploaded file, the uploaded file is of FileStorage type. Also which method should i perform the testing for, should it be the '/classification' or the func_predict.
I have tried the below method, though I did not get any success in this.
I created a test.py file and imported the classification method from main.py and then passed a filename to the data
from flask import Flask, Request
import io
import unittest
from main import classification
class TestFileFail(unittest.TestCase):
def test_1(self):
app = Flask(__name__)
app.debug = True
app.request_class = MyRequest
client = app.test_client()
resp = client.post(
'/classification',
data = {
'file': 'C:\\Users\\aswathi.nambiar\\Desktop\\Desktop docs\\W8_ECI_1.pdf'
}, content_type='multipart/form-data'
)
print(resp.data)
self.assertEqual(
'ok',
resp.data,
)
if __name__ == '__main__':
unittest.main()
I am completely lost! I know there have been earlier questions, but I am not able to figure out .

Answer 1

I have finally stumbled upon how to test it, in case anybody was looking out for something similar.
from predict_main_restplus import func_predict
from werkzeug.datastructures import FileStorage
file = None
def test_classification_correct():
with open('W8-EXP_1.pdf', 'rb') as fp:
file = FileStorage(fp)
a , b = func_predict(file)
assert (a, b) == ('W-8EXP',90.15652760121652)
So, here we are testing the prediction function in predict.py, it returns two values, prediction result and the confidence of the prediction. We can mock the upload using open(file) and then wrapping it with FileStorage. This worked for me.