Merge pull request #2489 from danielgtaylor/glacier-hash-fix

Fix Glacier file object hash calculation. Fixes #2489, #2488.

2 files changed, 54 insertions, 3 deletions

diff --git a/tests/unit/glacier/test_utils.py b/tests/unit/glacier/test_utils.py
index a051b59e..bace2a38 100644
--- a/tests/unit/glacier/test_utils.py
+++ b/tests/unit/glacier/test_utils.py
@@ -19,13 +19,16 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 # IN THE SOFTWARE.
 #
-import time
 import logging
+import os
+import tempfile
+import time
 from hashlib import sha256
 from tests.unit import unittest
 
+from boto.compat import BytesIO, six, StringIO
 from boto.glacier.utils import minimum_part_size, chunk_hashes, tree_hash, \
-        bytes_to_hex
+        bytes_to_hex, compute_hashes_from_fileobj
 
 
 class TestPartSizeCalculations(unittest.TestCase):
@@ -114,3 +117,49 @@ class TestTreeHash(unittest.TestCase):
         self.assertEqual(
             self.calculate_tree_hash(''),
             b'e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855')
+
+
+class TestFileHash(unittest.TestCase):
+    def _gen_data(self):
+        # Generate some pseudo-random bytes of data. We include the
+        # hard-coded blob as an example that fails to decode via UTF-8.
+        return os.urandom(5000) + b'\xc2\x00'
+
+    def test_compute_hash_tempfile(self):
+        # Compute a hash from a file object. On Python 2 this uses a non-
+        # binary mode. On Python 3, however, binary mode is required for
+        # binary files. If not used, you will get UTF-8 code errors.
+        if six.PY2:
+            mode = "w+"
+        else:
+            mode = "wb+"
+
+        with tempfile.TemporaryFile(mode=mode) as f:
+            f.write(self._gen_data())
+            f.seek(0)
+
+            compute_hashes_from_fileobj(f, chunk_size=512)
+
+    @unittest.skipUnless(six.PY3, 'Python 3 requires reading binary!')
+    def test_compute_hash_tempfile_py3(self):
+        # Note the missing 'b' in the mode!
+        with tempfile.TemporaryFile(mode='w+') as f:
+            with self.assertRaises(ValueError):
+                compute_hashes_from_fileobj(f, chunk_size=512)
+
+        # What about file-like objects without a mode? If it has an
+        # encoding we use it, otherwise attempt UTF-8 encoding to
+        # bytes for hashing.
+        f = StringIO('test data' * 500)
+        compute_hashes_from_fileobj(f, chunk_size=512)
+
+    @unittest.skipUnless(six.PY2, 'Python 3 requires reading binary!')
+    def test_compute_hash_stringio(self):
+        # Python 2 binary data in StringIO example
+        f = StringIO(self._gen_data())
+        compute_hashes_from_fileobj(f, chunk_size=512)
+
+    def test_compute_hash_bytesio(self):
+        # Compute a hash from a file-like BytesIO object.
+        f = BytesIO(self._gen_data())
+        compute_hashes_from_fileobj(f, chunk_size=512)
diff --git a/tests/unit/glacier/test_vault.py b/tests/unit/glacier/test_vault.py
index b17a3eee..f532e3b9 100644
--- a/tests/unit/glacier/test_vault.py
+++ b/tests/unit/glacier/test_vault.py
@@ -44,7 +44,9 @@ class TestVault(unittest.TestCase):
     def tearDown(self):
         self.size_patch.stop()
 
-    def test_upload_archive_small_file(self):
+    @mock.patch('boto.glacier.vault.compute_hashes_from_fileobj',
+                return_value=[b'abc', b'123'])
+    def test_upload_archive_small_file(self, compute_hashes):
         self.getsize.return_value = 1
 
         self.api.upload_archive.return_value = {'ArchiveId': 'archive_id'}