Update based on feedback.

aws · Feb 27, 2016 · 6083393 · 6083393
1 parent 23b71f0
commit 6083393
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 23 deletions.
diff --git a/awscli/customizations/s3/fileinfo.py b/awscli/customizations/s3/fileinfo.py
@@ -61,7 +61,7 @@ def save_file(filename, response_data, last_update, is_stream=False):
         with open(filename, 'wb') as out_file:
             write_to_file(out_file, etag, file_chunks, md5)
 
-    if md5 and etag != md5.hexdigest():
+    if md5 is not None and etag != md5.hexdigest():
         if not is_stream:
             os.remove(filename)
         raise MD5Error(filename)
@@ -93,7 +93,7 @@ def write_to_file(out_file, etag, file_chunks, md5=None, is_stream=False):
     """
     body = b''
     for chunk in file_chunks:
-        if md5 and not _is_multipart_etag(etag):
+        if md5 is not None and not _is_multipart_etag(etag):
             md5.update(chunk)
         if is_stream:
             body += chunk

diff --git a/tests/unit/customizations/s3/test_fileinfo.py b/tests/unit/customizations/s3/test_fileinfo.py
@@ -25,18 +25,30 @@
 from awscli.customizations.s3.fileinfo import TaskInfo
 
 
-class BaseMD5Test(unittest.TestCase):
+class TestSaveFile(unittest.TestCase):
     def setUp(self):
+        self.tempdir = tempfile.mkdtemp()
+        self.filename = os.path.join(self.tempdir, 'dir1', 'dir2', 'foo.txt')
+        etag = '3858f62230ac3c915f300c664312c63f'
+        self.response_data = {
+            'Body': six.BytesIO(b'foobar'),
+            'ETag': '"%s"' % etag,
+        }
+        self.last_update = datetime.now()
+
+        # Setup MD5 patches
         self.md5_object = mock.Mock()
-        self.md5_object.hexdigest = mock.Mock(return_value=b'foo')
+        self.md5_object.hexdigest = mock.Mock(return_value=etag)
         md5_builder = mock.Mock(return_value=self.md5_object)
         self.md5_patch = mock.patch('hashlib.md5', md5_builder)
         self.md5_patch.start()
         self._md5_available_patch = None
         self.set_md5_available()
 
     def tearDown(self):
-        super(BaseMD5Test, self).tearDown()
+        shutil.rmtree(self.tempdir)
+
+        # Tear down MD5 patches
         self.md5_patch.stop()
         if self._md5_available_patch:
             self._md5_available_patch.stop()
@@ -49,24 +61,6 @@ def set_md5_available(self, is_available=True):
             'awscli.customizations.s3.fileinfo.MD5_AVAILABLE', is_available)
         self._md5_available_patch.start()
 
-
-class TestSaveFile(BaseMD5Test):
-    def setUp(self):
-        super(TestSaveFile, self).setUp()
-        self.tempdir = tempfile.mkdtemp()
-        self.filename = os.path.join(self.tempdir, 'dir1', 'dir2', 'foo.txt')
-        etag = '3858f62230ac3c915f300c664312c63f'
-        self.md5_object.hexdigest = mock.Mock(return_value=etag)
-        self.response_data = {
-            'Body': six.BytesIO(b'foobar'),
-            'ETag': '"%s"' % etag,
-        }
-        self.last_update = datetime.now()
-
-    def tearDown(self):
-        super(TestSaveFile, self).tearDown()
-        shutil.rmtree(self.tempdir)
-
     def test_save_file(self):
         fileinfo.save_file(self.filename, self.response_data, self.last_update)
         self.assertTrue(os.path.isfile(self.filename))