Search code examples
unit-testingdjango-rest-framework

Write test for large file size. Django rest framwork


In my Django application there is an endpoint that handles file upload. In the serializer I define the file size validation like below

extra_kwargs = {
    "image_file": {
        "validators": [
            get_file_size_validator(1e8),
        ],
               
},

The validator function looks like this

def get_file_size_validator(file_size):
    def validator_func(file):
        if file_size < file.size:
            raise ValidationError(f"Maximum file size is {file_size}")

    return validator_func

I would like to write a test where I want to test that user cannot upload large files.

What I have tried so far.

  1. SimpleUploadedFile("foo.pdf", b"aaaaa" * (9**10), content_type="image/jpeg")
  2. InMemoryUploadedFile(BytesIO(b"large_content"), None, "foo.jpg", "image/jpeg", lambda: 1e9, None, None)
  3. Used mock.patch.object(InMemoryUploadedFile, "size", return_value=1e9):
  4. Mocked os.path.getsize as recommended here https://stackoverflow.com/a/35246275/1847898

The first solution gives me MemoryError and for all the others file size is returned as 13 in validator_func

Any ideas how I can achieve this?


Solution

  • You should mock the MemoryFileUploadHandler.file_complete method instead of SimpleUploadedFile or InMemoryUploadedFile.

    from unittest.mock import patch
    
    from django.core.files.uploadedfile import InMemoryUploadedFile, SimpleUploadedFile
    from rest_framework.test import APITestCase
    
    
    class TestUploadAPI(APITestCase):
        url = "/api/to/the/end-point/"
    
        @patch("django.core.files.uploadhandler.MemoryFileUploadHandler.file_complete")
        def test_validate_size(self, mock_file_complete):
            file_name = "any-file-name"
            file_content_type = "image/png"
            max_size = 1e8
    
            mock_file_complete.return_value = InMemoryUploadedFile(
                file=b"",
                field_name=None,
                name=file_name,
                content_type=file_content_type,
                size=max_size + 1,
                charset=None,
            )
    
            file = SimpleUploadedFile(
                name=file_name,
                content=b"",
                content_type=file_content_type,
            )
            response = self.client.post(self.url, {"file": file})
    
            self.assertEqual(response.status_code, 400)
            self.assertEqual(
                response.json(),
                {
                    "file": ["Maximum file size is 100000000.0"],
                },
            )
            mock_file_complete.assert_called_once()