Flask- Pytest mock file write for test file

Viewed 25

New to unit testing in python. I am writing some Pytests for my file upload endpoint. While my test works, I do not want my test to write on my disk (which is what it does currently). What is the best way to mock the file save function and just check if it was called or not. Here is my code:

from io import BytesIO


def test_upload_file_post_success(client):
    data = {
        'fields': 'somefield',
        'file': (BytesIO(b'FILE CONTENT'), 'test.csv')
    }
    response = client.post('/file-upload',
                           content_type='multipart/form-data',
                           data=data)
    assert response.status_code == 200
    

What I have tried so far:

  
from werkzeug.datastructures import FileStorage
from unittest.mock import MagicMock

@patch('werkzeug.datastructures.FileStorage')
def test_upload_files_success(mocked_fs, client):
    mocked_fs.save.return_value = None
    file = FileStorage(
        stream=BytesIO(b'FILE CONTENT'),
        filename='test_page_one_input.csv',
        content_type='text/csv'
    )
    file.save = MagicMock(return_value=None)
    
    response = client.post(
        '/file-upload',
        content_type='multipart/form-data',
        data={"file": file}
    )
    mocked_fs.save.assert_called_once()

Server code:


class UploadFiles(Resource):
    def post(self):
        # check if the post request has the file part
        if 'file' not in request.files:
            resp = jsonify({'message': 'No file in the request'})
            resp.status_code = 400
            return resp
        file = request.files['file']
        filename, extension = get_file_name_and_extension(file.filename)
        if file and extension in PARSERS.keys():
            filename = secure_filename(file.filename)
            # parse file and validate
            filepath = os.path.join('api/upload/datasets', filename)
            file.save(filepath)
            try:
                validate_file_content(filepath, PARSERS[extension])
                upload_to_storage(filename, file_version)
                resp = jsonify({'message': 'File successfully uploaded'})
                resp.status_code = 200
                return resp
            except (InvalidFileError, EmptyFileError, SaveFileError) as e:
                os.remove(filepath)
                resp = jsonify({'message': 'File validation Failed',
                                'details': str(e)})
                resp.status_code = 500
                return resp
        # check if valid file extension
        if extension not in PARSERS.keys():
            resp = jsonify(
                {'message': 'Invalid file format. Allowed file types are csv'})
            resp.status_code = 400
            return resp
        else:
            raise OSError(
                "Error occured while uploading file. Please try again later")


0 Answers
Related