Unit test case for file upload flask

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.