Repository: spotify/luigi
File: test/contrib/s3_test.py (Python)

Code before:
        original_size = s3_client.get_key(original).size
        copy_size = s3_client.get_key(copy).size
        self.assertEqual(original_size, copy_size)

    @mock_s3
    def _run_multipart_test(self, part_size, file_size, **kwargs):
        create_bucket()
        file_contents = b"a" * file_size

        s3_path = 's3://mybucket/putMe'
        tmp_file = tempfile.NamedTemporaryFile(mode='wb', delete=True)
        tmp_file_path = tmp_file.name
        tmp_file.write(file_contents)
        tmp_file.flush()

        s3_client = S3Client(AWS_ACCESS_KEY, AWS_SECRET_KEY)

        s3_client.put_multipart(tmp_file_path, s3_path,
                                part_size=part_size, **kwargs)
        self.assertTrue(s3_client.exists(s3_path))
        file_size = os.path.getsize(tmp_file.name)
        key_size = s3_client.get_key(s3_path).size
        self.assertEqual(file_size, key_size)
        tmp_file.close()

    def _run_copy_response_test(self, response):
        num, size = response
        self.assertIsInstance(response, tuple)

        # only check >= minimum possible values
        self.assertGreaterEqual(num, 1)
        self.assertGreaterEqual(size, 0)


Review comment (suggestion):
I'd find the way to check that the size is strictly equal to the sum instead of using GreaterEqual. Same thing to the number of items.

Code after:
        # As all the multipart put methods use `self._run_multipart_test`
        # we can just use this key
        original = 's3://mybucket/putMe'
        copy = 's3://mybucket/putMe_copy'

        # Copy the file from old location to new
        s3_client = S3Client(AWS_ACCESS_KEY, AWS_SECRET_KEY)
        if is_multipart:
            # 5MB is minimum part size, use it here so we don't have to generate huge files to test
            # the multipart upload in moto
            part_size = (1024 ** 2) * 5
            response = s3_client.copy(original, copy, part_size=part_size, threads=4)
        else:
            response = s3_client.copy(original, copy, threads=4)

        self._run_copy_response_test(response, expected_num=expected_num, expected_size=expected_size)

        # We can't use etags to compare between multipart and normal keys,
        # so we fall back to using the file size
        original_size = s3_client.get_key(original).size
        copy_size = s3_client.get_key(copy).size
        self.assertEqual(original_size, copy_size)

    @mock_s3
    def _run_multipart_test(self, part_size, file_size, **kwargs):
        create_bucket()
        file_contents = b"a" * file_size

        s3_path = 's3://mybucket/putMe'
        tmp_file = tempfile.NamedTemporaryFile(mode='wb', delete=True)
        tmp_file_path = tmp_file.name
        tmp_file.write(file_contents)
        tmp_file.flush()

        s3_client = S3Client(AWS_ACCESS_KEY, AWS_SECRET_KEY)

        s3_client.put_multipart(tmp_file_path, s3_path,
                                part_size=part_size, **kwargs)
        self.assertTrue(s3_client.exists(s3_path))
        file_size = os.path.getsize(tmp_file.name)
        key_size = s3_client.get_key(s3_path).size
        self.assertEqual(file_size, key_size)
        tmp_file.close()

        return 1, key_size

    def _run_copy_response_test(self, response, expected_num=None, expected_size=None):
        num, size = response
        self.assertIsInstance(response, tuple)

        # only check >= minimum possible value if not provided expected value