Merging in latest boto.

tests/s3/test_resumable_uploads.py

 #!/usr/bin/env python
 #
 # Copyright 2010 Google Inc.
 #
 # Permission is hereby granted, free of charge, to any person obtaining a
 # copy of this software and associated documentation files (the
 # "Software"), to deal in the Software without restriction, including
 # without limitation the rights to use, copy, modify, merge, publish, dis-
 # tribute, sublicense, and/or sell copies of the Software, and to permit
 # persons to whom the Software is furnished to do so, subject to the fol-
 # lowing conditions:
 #
 # The above copyright notice and this permission notice shall be included
 # in all copies or substantial portions of the Software.
 #
 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 # OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABIL-
 # ITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
 # SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 # WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 # IN THE SOFTWARE.
 
 """
-Tests of resumable uploads.
+Tests of Google Cloud Storage resumable uploads.
 """
 
 import errno
 import getopt
 import os
 import random
 import re
 import shutil
 import socket
 import StringIO
 import sys
 import tempfile
 import time
 import unittest
 
 import boto
 from boto.exception import GSResponseError
 from boto.gs.resumable_upload_handler import ResumableUploadHandler
+from boto.exception import InvalidUriError
 from boto.exception import ResumableTransferDisposition
 from boto.exception import ResumableUploadException
 from boto.exception import StorageResponseError
-from boto.tests.cb_test_harnass import CallbackTestHarnass
+from cb_test_harnass import CallbackTestHarnass
+
+# We don't use the OAuth2 authentication plugin directly; importing it here
+# ensures that it's loaded and available by default.
+try:
+  from oauth2_plugin import oauth2_plugin
+except ImportError:
+  # Do nothing - if user doesn't have OAuth2 configured it doesn't matter;
+  # and if they do, the tests will fail (as they should in that case).
+  pass
 
 
 class ResumableUploadTests(unittest.TestCase):
     """
     Resumable upload test suite.
     """
 
     def get_suite_description(self):
         return 'Resumable upload test suite'
 
-    @classmethod
-    def setUp(cls):
+    def setUp(self):
         """
         Creates dst_key needed by all tests.
 
         This method's namingCase is required by the unittest framework.
         """
-        cls.dst_key = cls.dst_key_uri.new_key(validate=False)
+        self.dst_key = self.dst_key_uri.new_key(validate=False)
 
-    @classmethod
-    def tearDown(cls):
+    def tearDown(self):
         """
         Deletes any objects or files created by last test run.
 
         This method's namingCase is required by the unittest framework.
         """
         try:
-            cls.dst_key_uri.delete_key()
+            self.dst_key_uri.delete_key()
         except GSResponseError:
             # Ignore possible not-found error.
             pass
         # Recursively delete dst dir and then re-create it, so in effect we
         # remove all dirs and files under that directory.
-        shutil.rmtree(cls.tmp_dir)
-        os.mkdir(cls.tmp_dir)
+        shutil.rmtree(self.tmp_dir)
+        os.mkdir(self.tmp_dir)
 
     @staticmethod
     def build_test_input_file(size):
         buf = []
         # I manually construct the random data here instead of calling
         # os.urandom() because I want to constrain the range of data (in
         # this case to 0'..'9') so the test
         # code can easily overwrite part of the StringIO file with
         # known-to-be-different values.
         for i in range(size):
             buf.append(str(random.randint(0, 9)))
         file_as_string = ''.join(buf)
         return (file_as_string, StringIO.StringIO(file_as_string))
 
     @classmethod
+    def get_dst_bucket_uri(cls, debug):
+        """A unique bucket to test."""
+        hostname = socket.gethostname().split('.')[0]
+        uri_base_str = 'gs://res-upload-test-%s-%s-%s' % (
+            hostname, os.getpid(), int(time.time()))
+        return boto.storage_uri('%s-dst' % uri_base_str, debug=debug)
+
+    @classmethod
+    def get_dst_key_uri(cls):
+        """A key to test."""
+        return cls.dst_bucket_uri.clone_replace_name('obj')
+
+    @classmethod
+    def get_staged_host(cls):
+        """URL of an existing bucket."""
+        return 'pub.commondatastorage.googleapis.com'
+
+    @classmethod
+    def get_invalid_upload_id(cls):
+        return (
+            'http://%s/?upload_id='
+            'AyzB2Uo74W4EYxyi5dp_-r68jz8rtbvshsv4TX7srJVkJ57CxTY5Dw2' % (
+                cls.get_staged_host()))
+
+    @classmethod
     def set_up_class(cls, debug):
         """
         Initializes test suite.
         """
 
         # Use a designated tmpdir prefix to make it easy to find the end of
         # the tmp path.
         cls.tmpdir_prefix = 'tmp_resumable_upload_test'
 
         # Create test source file data.
         cls.empty_src_file_size = 0
         (cls.empty_src_file_as_string, cls.empty_src_file) = (
             cls.build_test_input_file(cls.empty_src_file_size))
         cls.small_src_file_size = 2 * 1024  # 2 KB.
         (cls.small_src_file_as_string, cls.small_src_file) = (
             cls.build_test_input_file(cls.small_src_file_size))
         cls.larger_src_file_size = 500 * 1024  # 500 KB.
         (cls.larger_src_file_as_string, cls.larger_src_file) = (
             cls.build_test_input_file(cls.larger_src_file_size))
         cls.largest_src_file_size = 1024 * 1024  # 1 MB.
         (cls.largest_src_file_as_string, cls.largest_src_file) = (
             cls.build_test_input_file(cls.largest_src_file_size))
 
         # Create temp dir.
         cls.tmp_dir = tempfile.mkdtemp(prefix=cls.tmpdir_prefix)
 
         # Create the test bucket.
-        hostname = socket.gethostname().split('.')[0]
-        cls.uri_base_str = 'gs://res_upload_test_%s_%s_%s' % (
-            hostname, os.getpid(), int(time.time()))
-        cls.dst_bucket_uri = boto.storage_uri('%s_dst' %
-                                              cls.uri_base_str, debug=debug)
+        cls.dst_bucket_uri = cls.get_dst_bucket_uri(debug)
         cls.dst_bucket_uri.create_bucket()
-        cls.dst_key_uri = cls.dst_bucket_uri.clone_replace_name('obj')
+        cls.dst_key_uri = cls.get_dst_key_uri()
 
         cls.tracker_file_name = '%s%suri_tracker' % (cls.tmp_dir, os.sep)
 
         cls.syntactically_invalid_tracker_file_name = (
             '%s%ssynt_invalid_uri_tracker' % (cls.tmp_dir, os.sep))
         f = open(cls.syntactically_invalid_tracker_file_name, 'w')
         f.write('ftp://example.com')
         f.close()
 
-        cls.invalid_upload_id = (
-            'http://pub.commondatastorage.googleapis.com/?upload_id='
-            'AyzB2Uo74W4EYxyi5dp_-r68jz8rtbvshsv4TX7srJVkJ57CxTY5Dw2')
+        cls.invalid_upload_id = cls.get_invalid_upload_id()
         cls.invalid_upload_id_tracker_file_name = (
             '%s%sinvalid_upload_id_tracker' % (cls.tmp_dir, os.sep))
         f = open(cls.invalid_upload_id_tracker_file_name, 'w')
         f.write(cls.invalid_upload_id)
         f.close()
 
         cls.created_test_data = True
 
     @classmethod
     def tear_down_class(cls):
         """
         Deletes bucket and tmp dir created by set_up_class.
         """
         if not hasattr(cls, 'created_test_data'):
             return
-        # Call cls.tearDown() in case the tests got interrupted, to ensure
-        # dst objects get deleted.
-        cls.tearDown()
 
         # Retry (for up to 2 minutes) the bucket gets deleted (it may not
         # the first time round, due to eventual consistency of bucket delete
         # operations).
         for i in range(60):
             try:
                 cls.dst_bucket_uri.delete_bucket()
                 break
             except StorageResponseError:
                 print 'Test bucket (%s) not yet deleted, still trying' % (
@@ skipping 31 matching lines @@
             tracker_file_name=self.tracker_file_name, num_retries=0)
         try:
             self.dst_key.set_contents_from_file(
                 self.small_src_file, cb=harnass.call,
                 res_upload_handler=res_upload_handler)
             self.fail('Did not get expected ResumableUploadException')
         except ResumableUploadException, e:
             # We'll get a ResumableUploadException at this point because
             # of CallbackTestHarnass (above). Check that the tracker file was
             # created correctly.
-            self.assertEqual(e.disposition, ResumableTransferDisposition.ABORT)
+            self.assertEqual(e.disposition,
+                             ResumableTransferDisposition.ABORT_CUR_PROCESS)
             self.assertTrue(os.path.exists(self.tracker_file_name))
             f = open(self.tracker_file_name)
             uri_from_file = f.readline().strip()
             f.close()
             self.assertEqual(uri_from_file,
                              res_upload_handler.get_tracker_uri())
 
     def test_retryable_exception_recovery(self):
         """
         Tests handling of a retryable exception
         """
         # Test one of the RETRYABLE_EXCEPTIONS.
         exception = ResumableUploadHandler.RETRYABLE_EXCEPTIONS[0]
         harnass = CallbackTestHarnass(exception=exception)
         res_upload_handler = ResumableUploadHandler(num_retries=1)
         self.dst_key.set_contents_from_file(
             self.small_src_file, cb=harnass.call,
             res_upload_handler=res_upload_handler)
         # Ensure uploaded object has correct content.
         self.assertEqual(self.small_src_file_size, self.dst_key.size)
         self.assertEqual(self.small_src_file_as_string,
                          self.dst_key.get_contents_as_string())
 
+    def test_broken_pipe_recovery(self):
+        """
+        Tests handling of a Broken Pipe (which interacts with an httplib bug)
+        """
+        exception = IOError(errno.EPIPE, "Broken pipe")
+        harnass = CallbackTestHarnass(exception=exception)
+        res_upload_handler = ResumableUploadHandler(num_retries=1)
+        self.dst_key.set_contents_from_file(
+            self.small_src_file, cb=harnass.call,
+            res_upload_handler=res_upload_handler)
+        # Ensure uploaded object has correct content.
+        self.assertEqual(self.small_src_file_size, self.dst_key.size)
+        self.assertEqual(self.small_src_file_as_string,
+                         self.dst_key.get_contents_as_string())
+
     def test_non_retryable_exception_handling(self):
         """
         Tests a resumable upload that fails with a non-retryable exception
         """
         harnass = CallbackTestHarnass(
             exception=OSError(errno.EACCES, 'Permission denied'))
         res_upload_handler = ResumableUploadHandler(num_retries=1)
         try:
             self.dst_key.set_contents_from_file(
                 self.small_src_file, cb=harnass.call,
@@ skipping 44 matching lines @@
         harnass = CallbackTestHarnass(
             fail_after_n_bytes=self.larger_src_file_size/2, num_times_to_fail=2)
         res_upload_handler = ResumableUploadHandler(
             tracker_file_name=self.tracker_file_name, num_retries=1)
         try:
             self.dst_key.set_contents_from_file(
                 self.larger_src_file, cb=harnass.call,
                 res_upload_handler=res_upload_handler)
             self.fail('Did not get expected ResumableUploadException')
         except ResumableUploadException, e:
-            self.assertEqual(e.disposition, ResumableTransferDisposition.ABORT)
+            self.assertEqual(e.disposition,
+                             ResumableTransferDisposition.ABORT_CUR_PROCESS)
             # Ensure a tracker file survived.
             self.assertTrue(os.path.exists(self.tracker_file_name))
         # Try it one more time; this time should succeed.
         self.dst_key.set_contents_from_file(
             self.larger_src_file, cb=harnass.call,
             res_upload_handler=res_upload_handler)
         self.assertEqual(self.larger_src_file_size, self.dst_key.size)
         self.assertEqual(self.larger_src_file_as_string,
                          self.dst_key.get_contents_as_string())
         self.assertFalse(os.path.exists(self.tracker_file_name))
@@ skipping 69 matching lines @@
         # retries (initial upload request will establish expected size to
         # upload server).
         res_upload_handler = ResumableUploadHandler(
             tracker_file_name=self.tracker_file_name, num_retries=0)
         try:
             self.dst_key.set_contents_from_file(
                 self.larger_src_file, cb=harnass.call,
                 res_upload_handler=res_upload_handler)
             self.fail('Did not get expected ResumableUploadException')
         except ResumableUploadException, e:
-            self.assertEqual(e.disposition, ResumableTransferDisposition.ABORT)
+            # First abort (from harnass-forced failure) should be
+            # ABORT_CUR_PROCESS.
+            self.assertEqual(e.disposition, ResumableTransferDisposition.ABORT_CUR_PROCESS)
             # Ensure a tracker file survived.
             self.assertTrue(os.path.exists(self.tracker_file_name))
         # Try it again, this time with different size source file.
         # Wait 1 second between retry attempts, to give upload server a
         # chance to save state so it can respond to changed file size with
         # 500 response in the next attempt.
         time.sleep(1)
         try:
             self.dst_key.set_contents_from_file(
                 self.largest_src_file, res_upload_handler=res_upload_handler)
             self.fail('Did not get expected ResumableUploadException')
         except ResumableUploadException, e:
+            # This abort should be a hard abort (file size changing during
+            # transfer).
             self.assertEqual(e.disposition, ResumableTransferDisposition.ABORT)
-            self.assertNotEqual(
-                e.message.find('attempt to upload a different size file'), -1)
+            self.assertNotEqual(e.message.find('file size changed'), -1, e.message) 
 
     def test_upload_with_file_size_change_during_upload(self):
         """
         Tests resumable upload on a file that changes sizes while upload
         in progress
         """
         # Create a file we can change during the upload.
         test_file_size = 500 * 1024  # 500 KB.
         test_file = self.build_test_input_file(test_file_size)[1]
         harnass = CallbackTestHarnass(fp_to_change=test_file,
@@ skipping 29 matching lines @@
                 res_upload_handler=res_upload_handler)
             self.fail('Did not get expected ResumableUploadException')
         except ResumableUploadException, e:
             self.assertEqual(e.disposition, ResumableTransferDisposition.ABORT)
             # Ensure the file size didn't change.
             test_file.seek(0, os.SEEK_END)
             self.assertEqual(test_file_size, test_file.tell())
             self.assertNotEqual(
                 e.message.find('md5 signature doesn\'t match etag'), -1)
             # Ensure the bad data wasn't left around.
-            all_keys = self.dst_key_uri.get_all_keys()
-            self.assertEqual(0, len(all_keys))
+            try:
+              self.dst_key_uri.get_key()
+              self.fail('Did not get expected InvalidUriError')
+            except InvalidUriError, e:
+              pass
 
     def test_upload_with_content_length_header_set(self):
         """
         Tests resumable upload on a file when the user supplies a
         Content-Length header. This is used by gsutil, for example,
         to set the content length when gzipping a file.
         """
         res_upload_handler = ResumableUploadHandler()
         try:
             self.dst_key.set_contents_from_file(
@@ skipping 74 matching lines @@
     # that do it, with camelCase versions of these names).
     try:
         print 'Setting up %s...' % test_class.get_suite_description()
         test_class.set_up_class(debug)
         print 'Running %s...' % test_class.get_suite_description()
         unittest.TextTestRunner(verbosity=2).run(suite)
     finally:
         print 'Cleaning up after %s...' % test_class.get_suite_description()
         test_class.tear_down_class()
         print ''