• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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
16from __future__ import absolute_import
17
18from contextlib import contextmanager
19import functools
20import os
21import pkgutil
22import posixpath
23import re
24import tempfile
25import unittest
26import urlparse
27
28import boto
29import crcmod
30import gslib.tests as gslib_tests
31from gslib.util import UsingCrcmodExtension
32
33if not hasattr(unittest.TestCase, 'assertIsNone'):
34  # external dependency unittest2 required for Python <= 2.6
35  import unittest2 as unittest  # pylint: disable=g-import-not-at-top
36
37# Flags for running different types of tests.
38RUN_INTEGRATION_TESTS = True
39RUN_UNIT_TESTS = True
40RUN_S3_TESTS = False
41
42PARALLEL_COMPOSITE_UPLOAD_TEST_CONFIG = '/tmp/.boto.parallel_upload_test_config'
43
44
45def _HasS3Credentials():
46  return (boto.config.get('Credentials', 'aws_access_key_id', None) and
47          boto.config.get('Credentials', 'aws_secret_access_key', None))
48
49HAS_S3_CREDS = _HasS3Credentials()
50
51
52def _HasGSHost():
53  return boto.config.get('Credentials', 'gs_host', None) is not None
54
55HAS_GS_HOST = _HasGSHost()
56
57
58def _UsingJSONApi():
59  return boto.config.get('GSUtil', 'prefer_api', 'json').upper() != 'XML'
60
61USING_JSON_API = _UsingJSONApi()
62
63
64def _ArgcompleteAvailable():
65  argcomplete = None
66  try:
67    # pylint: disable=g-import-not-at-top
68    import argcomplete
69  except ImportError:
70    pass
71  return argcomplete is not None
72
73ARGCOMPLETE_AVAILABLE = _ArgcompleteAvailable()
74
75
76def _NormalizeURI(uri):
77  """Normalizes the path component of a URI.
78
79  Args:
80    uri: URI to normalize.
81
82  Returns:
83    Normalized URI.
84
85  Examples:
86    gs://foo//bar -> gs://foo/bar
87    gs://foo/./bar -> gs://foo/bar
88  """
89  # Note: we have to do this dance of changing gs:// to file:// because on
90  # Windows, the urlparse function won't work with URL schemes that are not
91  # known. urlparse('gs://foo/bar') on Windows turns into:
92  #     scheme='gs', netloc='', path='//foo/bar'
93  # while on non-Windows platforms, it turns into:
94  #     scheme='gs', netloc='foo', path='/bar'
95  uri = uri.replace('gs://', 'file://')
96  parsed = list(urlparse.urlparse(uri))
97  parsed[2] = posixpath.normpath(parsed[2])
98  if parsed[2].startswith('//'):
99    # The normpath function doesn't change '//foo' -> '/foo' by design.
100    parsed[2] = parsed[2][1:]
101  unparsed = urlparse.urlunparse(parsed)
102  unparsed = unparsed.replace('file://', 'gs://')
103  return unparsed
104
105
106def GenerationFromURI(uri):
107  """Returns a the generation for a StorageUri.
108
109  Args:
110    uri: boto.storage_uri.StorageURI object to get the URI from.
111
112  Returns:
113    Generation string for the URI.
114  """
115  if not (uri.generation or uri.version_id):
116    if uri.scheme == 's3': return 'null'
117  return uri.generation or uri.version_id
118
119
120def ObjectToURI(obj, *suffixes):
121  """Returns the storage URI string for a given StorageUri or file object.
122
123  Args:
124    obj: The object to get the URI from. Can be a file object, a subclass of
125         boto.storage_uri.StorageURI, or a string. If a string, it is assumed to
126         be a local on-disk path.
127    *suffixes: Suffixes to append. For example, ObjectToUri(bucketuri, 'foo')
128               would return the URI for a key name 'foo' inside the given
129               bucket.
130
131  Returns:
132    Storage URI string.
133  """
134  if isinstance(obj, file):
135    return 'file://%s' % os.path.abspath(os.path.join(obj.name, *suffixes))
136  if isinstance(obj, basestring):
137    return 'file://%s' % os.path.join(obj, *suffixes)
138  uri = obj.uri
139  if suffixes:
140    uri = _NormalizeURI('/'.join([uri] + list(suffixes)))
141
142  # Storage URIs shouldn't contain a trailing slash.
143  if uri.endswith('/'):
144    uri = uri[:-1]
145  return uri
146
147# The mock storage service comes from the Boto library, but it is not
148# distributed with Boto when installed as a package. To get around this, we
149# copy the file to gslib/tests/mock_storage_service.py when building the gsutil
150# package. Try and import from both places here.
151# pylint: disable=g-import-not-at-top
152try:
153  from gslib.tests import mock_storage_service
154except ImportError:
155  try:
156    from boto.tests.integration.s3 import mock_storage_service
157  except ImportError:
158    try:
159      from tests.integration.s3 import mock_storage_service
160    except ImportError:
161      import mock_storage_service
162
163
164class GSMockConnection(mock_storage_service.MockConnection):
165
166  def __init__(self, *args, **kwargs):
167    kwargs['provider'] = 'gs'
168    self.debug = 0
169    super(GSMockConnection, self).__init__(*args, **kwargs)
170
171mock_connection = GSMockConnection()
172
173
174class GSMockBucketStorageUri(mock_storage_service.MockBucketStorageUri):
175
176  def connect(self, access_key_id=None, secret_access_key=None):
177    return mock_connection
178
179  def compose(self, components, headers=None):
180    """Dummy implementation to allow parallel uploads with tests."""
181    return self.new_key()
182
183
184TEST_BOTO_REMOVE_SECTION = 'TestRemoveSection'
185
186
187def _SetBotoConfig(section, name, value, revert_list):
188  """Sets boto configuration temporarily for testing.
189
190  SetBotoConfigForTest and SetBotoConfigFileForTest should be called by tests
191  instead of this function. Those functions will ensure that the configuration
192  is reverted to its original setting using _RevertBotoConfig.
193
194  Args:
195    section: Boto config section to set
196    name: Boto config name to set
197    value: Value to set
198    revert_list: List for tracking configs to revert.
199  """
200  prev_value = boto.config.get(section, name, None)
201  if not boto.config.has_section(section):
202    revert_list.append((section, TEST_BOTO_REMOVE_SECTION, None))
203    boto.config.add_section(section)
204  revert_list.append((section, name, prev_value))
205  if value is None:
206    boto.config.remove_option(section, name)
207  else:
208    boto.config.set(section, name, value)
209
210
211def _RevertBotoConfig(revert_list):
212  """Reverts boto config modifications made by _SetBotoConfig.
213
214  Args:
215    revert_list: List of boto config modifications created by calls to
216                 _SetBotoConfig.
217  """
218  sections_to_remove = []
219  for section, name, value in revert_list:
220    if value is None:
221      if name == TEST_BOTO_REMOVE_SECTION:
222        sections_to_remove.append(section)
223      else:
224        boto.config.remove_option(section, name)
225    else:
226      boto.config.set(section, name, value)
227  for section in sections_to_remove:
228    boto.config.remove_section(section)
229
230
231def SequentialAndParallelTransfer(func):
232  """Decorator for tests that perform file to object transfers, or vice versa.
233
234  This forces the test to run once normally, and again with special boto
235  config settings that will ensure that the test follows the parallel composite
236  upload and/or sliced object download code paths.
237
238  Args:
239    func: Function to wrap.
240
241  Returns:
242    Wrapped function.
243  """
244  @functools.wraps(func)
245  def Wrapper(*args, **kwargs):
246    # Run the test normally once.
247    func(*args, **kwargs)
248
249    if not RUN_S3_TESTS and UsingCrcmodExtension(crcmod):
250      # Try again, forcing parallel upload and sliced download.
251      with SetBotoConfigForTest([
252          ('GSUtil', 'parallel_composite_upload_threshold', '1'),
253          ('GSUtil', 'sliced_object_download_threshold', '1'),
254          ('GSUtil', 'sliced_object_download_max_components', '3'),
255          ('GSUtil', 'check_hashes', 'always')]):
256        func(*args, **kwargs)
257
258  return Wrapper
259
260
261@contextmanager
262def SetBotoConfigForTest(boto_config_list):
263  """Sets the input list of boto configs for the duration of a 'with' clause.
264
265  Args:
266    boto_config_list: list of tuples of:
267      (boto config section to set, boto config name to set, value to set)
268
269  Yields:
270    Once after config is set.
271  """
272  revert_configs = []
273  tmp_filename = None
274  try:
275    tmp_fd, tmp_filename = tempfile.mkstemp(prefix='gsutil-temp-cfg')
276    os.close(tmp_fd)
277    for boto_config in boto_config_list:
278      _SetBotoConfig(boto_config[0], boto_config[1], boto_config[2],
279                     revert_configs)
280    with open(tmp_filename, 'w') as tmp_file:
281      boto.config.write(tmp_file)
282
283    with SetBotoConfigFileForTest(tmp_filename):
284      yield
285  finally:
286    _RevertBotoConfig(revert_configs)
287    if tmp_filename:
288      try:
289        os.remove(tmp_filename)
290      except OSError:
291        pass
292
293
294@contextmanager
295def SetEnvironmentForTest(env_variable_dict):
296  """Sets OS environment variables for a single test."""
297
298  def _ApplyDictToEnvironment(dict_to_apply):
299    for k, v in dict_to_apply.iteritems():
300      old_values[k] = os.environ.get(k)
301      if v is not None:
302        os.environ[k] = v
303      elif k in os.environ:
304        del os.environ[k]
305
306  old_values = {}
307  for k in env_variable_dict:
308    old_values[k] = os.environ.get(k)
309
310  try:
311    _ApplyDictToEnvironment(env_variable_dict)
312    yield
313  finally:
314    _ApplyDictToEnvironment(old_values)
315
316
317@contextmanager
318def SetBotoConfigFileForTest(boto_config_path):
319  """Sets a given file as the boto config file for a single test."""
320  # Setup for entering "with" block.
321  try:
322    old_boto_config_env_variable = os.environ['BOTO_CONFIG']
323    boto_config_was_set = True
324  except KeyError:
325    boto_config_was_set = False
326  os.environ['BOTO_CONFIG'] = boto_config_path
327
328  try:
329    yield
330  finally:
331    # Teardown for exiting "with" block.
332    if boto_config_was_set:
333      os.environ['BOTO_CONFIG'] = old_boto_config_env_variable
334    else:
335      os.environ.pop('BOTO_CONFIG', None)
336
337
338def GetTestNames():
339  """Returns a list of the names of the test modules in gslib.tests."""
340  matcher = re.compile(r'^test_(?P<name>.*)$')
341  names = []
342  for _, modname, _ in pkgutil.iter_modules(gslib_tests.__path__):
343    m = matcher.match(modname)
344    if m:
345      names.append(m.group('name'))
346  return names
347
348
349@contextmanager
350def WorkingDirectory(new_working_directory):
351  """Changes the working directory for the duration of a 'with' call.
352
353  Args:
354    new_working_directory: The directory to switch to before executing wrapped
355      code. A None value indicates that no switching is necessary.
356
357  Yields:
358    Once after working directory has been changed.
359  """
360  prev_working_directory = None
361  try:
362    prev_working_directory = os.getcwd()
363  except OSError:
364    # This can happen if the current working directory no longer exists.
365    pass
366
367  if new_working_directory:
368    os.chdir(new_working_directory)
369
370  try:
371    yield
372  finally:
373    if new_working_directory and prev_working_directory:
374      os.chdir(prev_working_directory)
375