tools/telemetry/third_party/gsutilz/gslib/tests/testcase/unit_testcase.py

   1 # -*- coding: utf-8 -*-
   2 # Copyright 2013 Google Inc. All Rights Reserved.
   3 #
   4 # Licensed under the Apache License, Version 2.0 (the "License");
   5 # you may not use this file except in compliance with the License.
   6 # You may obtain a copy of the License at
   7 #
   8 #     http://www.apache.org/licenses/LICENSE-2.0
   9 #
  10 # Unless required by applicable law or agreed to in writing, software
  11 # distributed under the License is distributed on an "AS IS" BASIS,
  12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 # See the License for the specific language governing permissions and
  14 # limitations under the License.
  15 """Contains gsutil base unit test case class."""
  16
  17 from __future__ import absolute_import
  18
  19 import logging
  20 import os
  21 import sys
  22 import tempfile
  23
  24 import boto
  25 from gslib import wildcard_iterator
  26 from gslib.boto_translation import BotoTranslation
  27 from gslib.cloud_api_delegator import CloudApiDelegator
  28 from gslib.command_runner import CommandRunner
  29 from gslib.cs_api_map import ApiMapConstants
  30 from gslib.cs_api_map import ApiSelector
  31 from gslib.tests.mock_logging_handler import MockLoggingHandler
  32 from gslib.tests.testcase import base
  33 import gslib.tests.util as util
  34 from gslib.tests.util import unittest
  35 from gslib.tests.util import WorkingDirectory
  36 from gslib.util import GsutilStreamHandler
  37
  38
  39 class GsutilApiUnitTestClassMapFactory(object):
  40   """Class map factory for use in unit tests.
  41
  42   BotoTranslation is used for all cases so that GSMockBucketStorageUri can
  43   be used to communicate with the mock XML service.
  44   """
  45
  46   @classmethod
  47   def GetClassMap(cls):
  48     """Returns a class map for use in unit tests."""
  49     gs_class_map = {
  50         ApiSelector.XML: BotoTranslation,
  51         ApiSelector.JSON: BotoTranslation
  52     }
  53     s3_class_map = {
  54         ApiSelector.XML: BotoTranslation
  55     }
  56     class_map = {
  57         'gs': gs_class_map,
  58         's3': s3_class_map
  59     }
  60     return class_map
  61
  62
  63 @unittest.skipUnless(util.RUN_UNIT_TESTS,
  64                      'Not running integration tests.')
  65 class GsUtilUnitTestCase(base.GsUtilTestCase):
  66   """Base class for gsutil unit tests."""
  67
  68   @classmethod
  69   def setUpClass(cls):
  70     base.GsUtilTestCase.setUpClass()
  71     cls.mock_bucket_storage_uri = util.GSMockBucketStorageUri
  72     cls.mock_gsutil_api_class_map_factory = GsutilApiUnitTestClassMapFactory
  73     cls.logger = logging.getLogger()
  74     cls.command_runner = CommandRunner(
  75         bucket_storage_uri_class=cls.mock_bucket_storage_uri,
  76         gsutil_api_class_map_factory=cls.mock_gsutil_api_class_map_factory)
  77
  78   def setUp(self):
  79     super(GsUtilUnitTestCase, self).setUp()
  80     self.bucket_uris = []
  81     self.stdout_save = sys.stdout
  82     self.stderr_save = sys.stderr
  83     fd, self.stdout_file = tempfile.mkstemp()
  84     sys.stdout = os.fdopen(fd, 'w+')
  85     fd, self.stderr_file = tempfile.mkstemp()
  86     sys.stderr = os.fdopen(fd, 'w+')
  87     self.accumulated_stdout = []
  88     self.accumulated_stderr = []
  89
  90     self.root_logger = logging.getLogger()
  91     self.is_debugging = self.root_logger.isEnabledFor(logging.DEBUG)
  92     self.log_handlers_save = self.root_logger.handlers
  93     fd, self.log_handler_file = tempfile.mkstemp()
  94     self.log_handler_stream = os.fdopen(fd, 'w+')
  95     self.temp_log_handler = GsutilStreamHandler(self.log_handler_stream)
  96     self.root_logger.handlers = [self.temp_log_handler]
  97
  98   def tearDown(self):
  99     super(GsUtilUnitTestCase, self).tearDown()
 100
 101     self.root_logger.handlers = self.log_handlers_save
 102     self.temp_log_handler.flush()
 103     self.temp_log_handler.close()
 104     self.log_handler_stream.seek(0)
 105     log_output = self.log_handler_stream.read()
 106     self.log_handler_stream.close()
 107     os.unlink(self.log_handler_file)
 108
 109     sys.stdout.seek(0)
 110     sys.stderr.seek(0)
 111     stdout = sys.stdout.read()
 112     stderr = sys.stderr.read()
 113     stdout += ''.join(self.accumulated_stdout)
 114     stderr += ''.join(self.accumulated_stderr)
 115     sys.stdout.close()
 116     sys.stderr.close()
 117     sys.stdout = self.stdout_save
 118     sys.stderr = self.stderr_save
 119     os.unlink(self.stdout_file)
 120     os.unlink(self.stderr_file)
 121
 122     if self.is_debugging and stdout:
 123       sys.stderr.write('==== stdout %s ====\n' % self.id())
 124       sys.stderr.write(stdout)
 125       sys.stderr.write('==== end stdout ====\n')
 126     if self.is_debugging and stderr:
 127       sys.stderr.write('==== stderr %s ====\n' % self.id())
 128       sys.stderr.write(stderr)
 129       sys.stderr.write('==== end stderr ====\n')
 130     if self.is_debugging and log_output:
 131       sys.stderr.write('==== log output %s ====\n' % self.id())
 132       sys.stderr.write(log_output)
 133       sys.stderr.write('==== end log output ====\n')
 134
 135   def RunCommand(self, command_name, args=None, headers=None, debug=0,
 136                  test_method=None, return_stdout=False, return_stderr=False,
 137                  return_log_handler=False, cwd=None):
 138     """Method for calling gslib.command_runner.CommandRunner.
 139
 140     Passes parallel_operations=False for all tests, optionally saving/returning
 141     stdout output. We run all tests multi-threaded, to exercise those more
 142     complicated code paths.
 143     TODO: Change to run with parallel_operations=True for all tests. At
 144     present when you do this it causes many test failures.
 145
 146     Args:
 147       command_name: The name of the command being run.
 148       args: Command-line args (arg0 = actual arg, not command name ala bash).
 149       headers: Dictionary containing optional HTTP headers to pass to boto.
 150       debug: Debug level to pass in to boto connection (range 0..3).
 151       test_method: Optional general purpose method for testing purposes.
 152                    Application and semantics of this method will vary by
 153                    command and test type.
 154       return_stdout: If True, will save and return stdout produced by command.
 155       return_stderr: If True, will save and return stderr produced by command.
 156       return_log_handler: If True, will return a MockLoggingHandler instance
 157            that was attached to the command's logger while running.
 158       cwd: The working directory that should be switched to before running the
 159            command. The working directory will be reset back to its original
 160            value after running the command. If not specified, the working
 161            directory is left unchanged.
 162
 163     Returns:
 164       One or a tuple of requested return values, depending on whether
 165       return_stdout, return_stderr, and/or return_log_handler were specified.
 166     """
 167     args = args or []
 168
 169     command_line = ' '.join([command_name] + args)
 170     if self.is_debugging:
 171       self.stderr_save.write('\nRunCommand of %s\n' % command_line)
 172
 173     # Save and truncate stdout and stderr for the lifetime of RunCommand. This
 174     # way, we can return just the stdout and stderr that was output during the
 175     # RunNamedCommand call below.
 176     sys.stdout.seek(0)
 177     sys.stderr.seek(0)
 178     stdout = sys.stdout.read()
 179     stderr = sys.stderr.read()
 180     if stdout:
 181       self.accumulated_stdout.append(stdout)
 182     if stderr:
 183       self.accumulated_stderr.append(stderr)
 184     sys.stdout.seek(0)
 185     sys.stderr.seek(0)
 186     sys.stdout.truncate()
 187     sys.stderr.truncate()
 188
 189     mock_log_handler = MockLoggingHandler()
 190     logging.getLogger(command_name).addHandler(mock_log_handler)
 191
 192     try:
 193       with WorkingDirectory(cwd):
 194         self.command_runner.RunNamedCommand(
 195             command_name, args=args, headers=headers, debug=debug,
 196             parallel_operations=False, test_method=test_method,
 197             do_shutdown=False)
 198     finally:
 199       sys.stdout.seek(0)
 200       stdout = sys.stdout.read()
 201       sys.stderr.seek(0)
 202       stderr = sys.stderr.read()
 203       logging.getLogger(command_name).removeHandler(mock_log_handler)
 204       mock_log_handler.close()
 205
 206       log_output = '\n'.join(
 207           '%s:\n  ' % level + '\n  '.join(records)
 208           for level, records in mock_log_handler.messages.iteritems()
 209           if records)
 210       if self.is_debugging and log_output:
 211         self.stderr_save.write(
 212             '==== logging RunCommand %s %s ====\n' % (self.id(), command_line))
 213         self.stderr_save.write(log_output)
 214         self.stderr_save.write('\n==== end logging ====\n')
 215       if self.is_debugging and stdout:
 216         self.stderr_save.write(
 217             '==== stdout RunCommand %s %s ====\n' % (self.id(), command_line))
 218         self.stderr_save.write(stdout)
 219         self.stderr_save.write('==== end stdout ====\n')
 220       if self.is_debugging and stderr:
 221         self.stderr_save.write(
 222             '==== stderr RunCommand %s %s ====\n' % (self.id(), command_line))
 223         self.stderr_save.write(stderr)
 224         self.stderr_save.write('==== end stderr ====\n')
 225
 226       # Reset stdout and stderr files, so that we won't print them out again
 227       # in tearDown if debugging is enabled.
 228       sys.stdout.seek(0)
 229       sys.stderr.seek(0)
 230       sys.stdout.truncate()
 231       sys.stderr.truncate()
 232
 233     to_return = []
 234     if return_stdout:
 235       to_return.append(stdout)
 236     if return_stderr:
 237       to_return.append(stderr)
 238     if return_log_handler:
 239       to_return.append(mock_log_handler)
 240     if len(to_return) == 1:
 241       return to_return[0]
 242     return tuple(to_return)
 243
 244   @classmethod
 245   def MakeGsUtilApi(cls, debug=0):
 246     gsutil_api_map = {
 247         ApiMapConstants.API_MAP: (
 248             cls.mock_gsutil_api_class_map_factory.GetClassMap()),
 249         ApiMapConstants.SUPPORT_MAP: {
 250             'gs': [ApiSelector.XML, ApiSelector.JSON],
 251             's3': [ApiSelector.XML]
 252         },
 253         ApiMapConstants.DEFAULT_MAP: {
 254             'gs': ApiSelector.JSON,
 255             's3': ApiSelector.XML
 256         }
 257     }
 258
 259     return CloudApiDelegator(
 260         cls.mock_bucket_storage_uri, gsutil_api_map, cls.logger, debug=debug)
 261
 262   @classmethod
 263   def _test_wildcard_iterator(cls, uri_or_str, debug=0):
 264     """Convenience method for instantiating a test instance of WildcardIterator.
 265
 266     This makes it unnecessary to specify all the params of that class
 267     (like bucket_storage_uri_class=mock_storage_service.MockBucketStorageUri).
 268     Also, naming the factory method this way makes it clearer in the test code
 269     that WildcardIterator needs to be set up for testing.
 270
 271     Args are same as for wildcard_iterator.wildcard_iterator(), except
 272     there are no class args for bucket_storage_uri_class or gsutil_api_class.
 273
 274     Args:
 275       uri_or_str: StorageUri or string representing the wildcard string.
 276       debug: debug level to pass to the underlying connection (0..3)
 277
 278     Returns:
 279       WildcardIterator, over which caller can iterate.
 280     """
 281     # TODO: Remove when tests no longer pass StorageUri arguments.
 282     uri_string = uri_or_str
 283     if hasattr(uri_or_str, 'uri'):
 284       uri_string = uri_or_str.uri
 285
 286     return wildcard_iterator.CreateWildcardIterator(
 287         uri_string, cls.MakeGsUtilApi())
 288
 289   @staticmethod
 290   def _test_storage_uri(uri_str, default_scheme='file', debug=0,
 291                         validate=True):
 292     """Convenience method for instantiating a testing instance of StorageUri.
 293
 294     This makes it unnecessary to specify
 295     bucket_storage_uri_class=mock_storage_service.MockBucketStorageUri.
 296     Also naming the factory method this way makes it clearer in the test
 297     code that StorageUri needs to be set up for testing.
 298
 299     Args, Returns, and Raises are same as for boto.storage_uri(), except there's
 300     no bucket_storage_uri_class arg.
 301
 302     Args:
 303       uri_str: Uri string to create StorageUri for.
 304       default_scheme: Default scheme for the StorageUri
 305       debug: debug level to pass to the underlying connection (0..3)
 306       validate: If True, validate the resource that the StorageUri refers to.
 307
 308     Returns:
 309       StorageUri based on the arguments.
 310     """
 311     return boto.storage_uri(uri_str, default_scheme, debug, validate,
 312                             util.GSMockBucketStorageUri)
 313
 314   def CreateBucket(self, bucket_name=None, test_objects=0, storage_class=None,
 315                    provider='gs'):
 316     """Creates a test bucket.
 317
 318     The bucket and all of its contents will be deleted after the test.
 319
 320     Args:
 321       bucket_name: Create the bucket with this name. If not provided, a
 322                    temporary test bucket name is constructed.
 323       test_objects: The number of objects that should be placed in the bucket or
 324                     a list of object names to place in the bucket. Defaults to
 325                     0.
 326       storage_class: storage class to use. If not provided we us standard.
 327       provider: string provider to use, default gs.
 328
 329     Returns:
 330       StorageUri for the created bucket.
 331     """
 332     bucket_name = bucket_name or self.MakeTempName('bucket')
 333     bucket_uri = boto.storage_uri(
 334         '%s://%s' % (provider, bucket_name.lower()),
 335         suppress_consec_slashes=False,
 336         bucket_storage_uri_class=util.GSMockBucketStorageUri)
 337     bucket_uri.create_bucket(storage_class=storage_class)
 338     self.bucket_uris.append(bucket_uri)
 339     try:
 340       iter(test_objects)
 341     except TypeError:
 342       test_objects = [self.MakeTempName('obj') for _ in range(test_objects)]
 343     for i, name in enumerate(test_objects):
 344       self.CreateObject(bucket_uri=bucket_uri, object_name=name,
 345                         contents='test %d' % i)
 346     return bucket_uri
 347
 348   def CreateObject(self, bucket_uri=None, object_name=None, contents=None):
 349     """Creates a test object.
 350
 351     Args:
 352       bucket_uri: The URI of the bucket to place the object in. If not
 353                   specified, a new temporary bucket is created.
 354       object_name: The name to use for the object. If not specified, a temporary
 355                    test object name is constructed.
 356       contents: The contents to write to the object. If not specified, the key
 357                 is not written to, which means that it isn't actually created
 358                 yet on the server.
 359
 360     Returns:
 361       A StorageUri for the created object.
 362     """
 363     bucket_uri = bucket_uri or self.CreateBucket()
 364     object_name = object_name or self.MakeTempName('obj')
 365     key_uri = bucket_uri.clone_replace_name(object_name)
 366     if contents is not None:
 367       key_uri.set_contents_from_string(contents)
 368     return key_uri