text: #!/usr/bin/env python
  # Copyright (c) 2012 The Chromium Authors. All rights reserved.
  # Use of this source code is governed by a BSD-style license that can be
  # found in the LICENSE file.
  # pylint: disable=protected-access
  
  """Unit tests for download_from_google_storage.py."""
  
  import optparse
  import os
  import Queue
  import shutil
  import sys
  import tarfile
  import tempfile
  import threading
  import unittest
  
  
  sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
  
  import upload_to_google_storage
  import download_from_google_storage
  
  # ../third_party/gsutil/gsutil
  GSUTIL_DEFAULT_PATH = os.path.join(
      os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
      'gsutil.py')
  TEST_DIR = os.path.dirname(os.path.abspath(__file__))
  
  
  class GsutilMock(object):
    def __init__(self, path, boto_path, timeout=None):
      self.path = path
      self.timeout = timeout
      self.boto_path = boto_path
      self.expected = []
      self.history = []
      self.lock = threading.Lock()
  
    def add_expected(self, return_code, out, err, fn=None):
      self.expected.append((return_code, out, err, fn))
  
    def append_history(self, method, args):
      self.history.append((method, args))
  
    def call(self, *args):
      with self.lock:
        self.append_history('call', args)
        if self.expected:
          code, _out, _err, fn = self.expected.pop(0)
          if fn:
            fn()
          return code
        else:
          return 0
  
    def check_call(self, *args):
      with self.lock:
        self.append_history('check_call', args)
        if self.expected:
          code, out, err, fn = self.expected.pop(0)
          if fn:
            fn()
          return code, out, err
        else:
          return (0, '', '')
  
  
  class ChangedWorkingDirectory(object):
    def __init__(self, working_directory):
      self._old_cwd = ''
      self._working_directory = working_directory
  
    def __enter__(self):
      self._old_cwd = os.getcwd()
      print "Enter directory = ", self._working_directory
      os.chdir(self._working_directory)
  
    def __exit__(self, *_):
      print "Enter directory = ", self._old_cwd
      os.chdir(self._old_cwd)
  
  
  class GstoolsUnitTests(unittest.TestCase):
    def setUp(self):
      self.temp_dir = tempfile.mkdtemp(prefix='gstools_test')
      self.base_path = os.path.join(self.temp_dir, 'test_files')
      shutil.copytree(os.path.join(TEST_DIR, 'gstools'), self.base_path)
  
    def cleanUp(self):
      shutil.rmtree(self.temp_dir)
  
    def test_validate_tar_file(self):
      lorem_ipsum = os.path.join(self.base_path, 'lorem_ipsum.txt')
      with ChangedWorkingDirectory(self.base_path):
        # Sanity ok check.
        tar_dir = 'ok_dir'
        os.makedirs(os.path.join(self.base_path, tar_dir))
        tar = 'good.tar.gz'
        lorem_ipsum_copy = os.path.join(tar_dir, 'lorem_ipsum.txt')
        shutil.copyfile(lorem_ipsum, lorem_ipsum_copy)
        with tarfile.open(tar, 'w:gz') as tar:
          tar.add(lorem_ipsum_copy)
          self.assertTrue(
              download_from_google_storage._validate_tar_file(tar, tar_dir))
  
        # Test no links.
        tar_dir_link = 'for_tar_link'
        os.makedirs(tar_dir_link)
        link = os.path.join(tar_dir_link, 'link')
        os.symlink(lorem_ipsum, link)
        tar_with_links = 'with_links.tar.gz'
        with tarfile.open(tar_with_links, 'w:gz') as tar:
          tar.add(link)
          self.assertFalse(
              download_from_google_storage._validate_tar_file(tar, tar_dir_link))
  
        # Test not outside.
        tar_dir_outside = 'outside_tar'
        os.makedirs(tar_dir_outside)
        tar_with_outside = 'with_outside.tar.gz'
        with tarfile.open(tar_with_outside, 'w:gz') as tar:
          tar.add(lorem_ipsum)
          self.assertFalse(
              download_from_google_storage._validate_tar_file(tar,
                                                              tar_dir_outside))
        # Test no ..
        tar_with_dotdot = 'with_dotdot.tar.gz'
        dotdot_file = os.path.join(tar_dir, '..', tar_dir, 'lorem_ipsum.txt')
        with tarfile.open(tar_with_dotdot, 'w:gz') as tar:
          tar.add(dotdot_file)
          self.assertFalse(
              download_from_google_storage._validate_tar_file(tar,
                                                              tar_dir))
  
    def test_gsutil(self):
      gsutil = download_from_google_storage.Gsutil(GSUTIL_DEFAULT_PATH, None)
      self.assertEqual(gsutil.path, GSUTIL_DEFAULT_PATH)
      code, _, err = gsutil.check_call()
      self.assertEqual(code, 0)
      self.assertEqual(err, '')
  
    def test_get_sha1(self):
      lorem_ipsum = os.path.join(self.base_path, 'lorem_ipsum.txt')
      self.assertEqual(
          download_from_google_storage.get_sha1(lorem_ipsum),
          '7871c8e24da15bad8b0be2c36edc9dc77e37727f')
  
    def test_get_md5(self):
      lorem_ipsum = os.path.join(self.base_path, 'lorem_ipsum.txt')
      self.assertEqual(
          upload_to_google_storage.get_md5(lorem_ipsum),
          '634d7c1ed3545383837428f031840a1e')
  
    def test_get_md5_cached_read(self):
      lorem_ipsum = os.path.join(self.base_path, 'lorem_ipsum.txt')
      # Use a fake 'stale' MD5 sum.  Expected behavior is to return stale sum.
      self.assertEqual(
          upload_to_google_storage.get_md5_cached(lorem_ipsum),
          '734d7c1ed3545383837428f031840a1e')
  
    def test_get_md5_cached_write(self):
      lorem_ipsum2 = os.path.join(self.base_path, 'lorem_ipsum2.txt')
      lorem_ipsum2_md5 = os.path.join(self.base_path, 'lorem_ipsum2.txt.md5')
      if os.path.exists(lorem_ipsum2_md5):
        os.remove(lorem_ipsum2_md5)
      # Use a fake 'stale' MD5 sum.  Expected behavior is to return stale sum.
      self.assertEqual(
          upload_to_google_storage.get_md5_cached(lorem_ipsum2),
          '4c02d1eb455a0f22c575265d17b84b6d')
      self.assertTrue(os.path.exists(lorem_ipsum2_md5))
      self.assertEqual(
          open(lorem_ipsum2_md5, 'rb').read(),
          '4c02d1eb455a0f22c575265d17b84b6d')
      os.remove(lorem_ipsum2_md5)  # Clean up.
      self.assertFalse(os.path.exists(lorem_ipsum2_md5))
  
  
  class DownloadTests(unittest.TestCase):
    def setUp(self):
      self.gsutil = GsutilMock(GSUTIL_DEFAULT_PATH, None)
      self.temp_dir = tempfile.mkdtemp(prefix='gstools_test')
      self.checkout_test_files = os.path.join(
          TEST_DIR, 'gstools', 'download_test_data')
      self.base_path = os.path.join(
          self.temp_dir, 'download_test_data')
      shutil.copytree(self.checkout_test_files, self.base_path)
      self.base_url = 'gs://sometesturl'
      self.parser = optparse.OptionParser()
      self.queue = Queue.Queue()
      self.ret_codes = Queue.Queue()
      self.lorem_ipsum = os.path.join(self.base_path, 'lorem_ipsum.txt')
      self.lorem_ipsum_sha1 = '7871c8e24da15bad8b0be2c36edc9dc77e37727f'
      self.maxDiff = None
  
    def cleanUp(self):
      shutil.rmtree(self.temp_dir)
  
    def test_enumerate_files_non_recursive(self):
      queue_size = download_from_google_storage.enumerate_work_queue(
          self.base_path, self.queue, True, False, False, None, False, False)
      expected_queue = [
          ('e6c4fbd4fe7607f3e6ebf68b2ea4ef694da7b4fe',
              os.path.join(self.base_path, 'rootfolder_text.txt')),
         ('7871c8e24da15bad8b0be2c36edc9dc77e37727f',
              os.path.join(self.base_path, 'uploaded_lorem_ipsum.txt'))]
      self.assertEqual(sorted(expected_queue), sorted(self.queue.queue))
      self.assertEqual(queue_size, 2)
  
    def test_enumerate_files_recursive(self):
      queue_size = download_from_google_storage.enumerate_work_queue(
          self.base_path, self.queue, True, True, False, None, False, False)
      expected_queue = [
          ('e6c4fbd4fe7607f3e6ebf68b2ea4ef694da7b4fe',
              os.path.join(self.base_path, 'rootfolder_text.txt')),
          ('7871c8e24da15bad8b0be2c36edc9dc77e37727f',
              os.path.join(self.base_path, 'uploaded_lorem_ipsum.txt')),
          ('b5415aa0b64006a95c0c409182e628881d6d6463',
              os.path.join(self.base_path, 'subfolder', 'subfolder_text.txt'))]
      self.assertEqual(sorted(expected_queue), sorted(self.queue.queue))
      self.assertEqual(queue_size, 3)
  
    def test_download_worker_single_file(self):
      sha1_hash = '7871c8e24da15bad8b0be2c36edc9dc77e37727f'
      input_filename = '%s/%s' % (self.base_url, sha1_hash)
      output_filename = os.path.join(self.base_path, 'uploaded_lorem_ipsum.txt')
      self.queue.put((sha1_hash, output_filename))
      self.queue.put((None, None))
      stdout_queue = Queue.Queue()
      download_from_google_storage._downloader_worker_thread(
          0, self.queue, False, self.base_url, self.gsutil,
          stdout_queue, self.ret_codes, True, False)
      expected_calls = [
          ('check_call',
              ('ls', input_filename)),
          ('check_call',
              ('cp', input_filename, output_filename))]
      if sys.platform != 'win32':
        expected_calls.append(
            ('check_call',
             ('stat',
              'gs://sometesturl/7871c8e24da15bad8b0be2c36edc9dc77e37727f')))
      expected_output = [
          '0> Downloading %s...' % output_filename]
      expected_ret_codes = []
      self.assertEqual(list(stdout_queue.queue), expected_output)
      self.assertEqual(self.gsutil.history, expected_calls)
      self.assertEqual(list(self.ret_codes.queue), expected_ret_codes)
  
    def test_download_worker_skips_file(self):
      sha1_hash = 'e6c4fbd4fe7607f3e6ebf68b2ea4ef694da7b4fe'
      output_filename = os.path.join(self.base_path, 'rootfolder_text.txt')
      self.queue.put((sha1_hash, output_filename))
      self.queue.put((None, None))
      stdout_queue = Queue.Queue()
      download_from_google_storage._downloader_worker_thread(
          0, self.queue, False, self.base_url, self.gsutil,
          stdout_queue, self.ret_codes, True, False)
      expected_output = [
          '0> File %s exists and SHA1 matches. Skipping.' % output_filename
      ]
      self.assertEqual(list(stdout_queue.queue), expected_output)
      self.assertEqual(self.gsutil.history, [])
  
    def test_download_extract_archive(self):
      # Generate a gzipped tarfile
      output_filename = os.path.join(self.base_path, 'subfolder.tar.gz')
      output_dirname = os.path.join(self.base_path, 'subfolder')
      extracted_filename = os.path.join(output_dirname, 'subfolder_text.txt')
      with tarfile.open(output_filename, 'w:gz') as tar:
        tar.add(output_dirname, arcname='subfolder')
      shutil.rmtree(output_dirname)
      sha1_hash = download_from_google_storage.get_sha1(output_filename)
      input_filename = '%s/%s' % (self.base_url, sha1_hash)
      self.queue.put((sha1_hash, output_filename))
      self.queue.put((None, None))
      stdout_queue = Queue.Queue()
      download_from_google_storage._downloader_worker_thread(
          0, self.queue, True, self.base_url, self.gsutil,
          stdout_queue, self.ret_codes, True, True, delete=False)
      expected_calls = [
          ('check_call',
              ('ls', input_filename)),
          ('check_call',
              ('cp', input_filename, output_filename))]
      if sys.platform != 'win32':
        expected_calls.append(
            ('check_call',
             ('stat',
              'gs://sometesturl/%s' % sha1_hash)))
      expected_output = [
          '0> Downloading %s...' % output_filename]
      expected_output.extend([
          '0> Extracting 3 entries from %s to %s' % (output_filename,
                                                     output_dirname)])
      expected_ret_codes = []
      self.assertEqual(list(stdout_queue.queue), expected_output)
      self.assertEqual(self.gsutil.history, expected_calls)
      self.assertEqual(list(self.ret_codes.queue), expected_ret_codes)
      self.assertTrue(os.path.exists(output_dirname))
      self.assertTrue(os.path.exists(extracted_filename))
  
    def test_download_worker_skips_not_found_file(self):
      sha1_hash = '7871c8e24da15bad8b0be2c36edc9dc77e37727f'
      input_filename = '%s/%s' % (self.base_url, sha1_hash)
      output_filename = os.path.join(self.base_path, 'uploaded_lorem_ipsum.txt')
      self.queue.put((sha1_hash, output_filename))
      self.queue.put((None, None))
      stdout_queue = Queue.Queue()
      self.gsutil.add_expected(1, '', '')  # Return error when 'ls' is called.
      download_from_google_storage._downloader_worker_thread(
          0, self.queue, False, self.base_url, self.gsutil,
          stdout_queue, self.ret_codes, True, False)
      expected_output = [
          '0> Failed to fetch file %s for %s, skipping. [Err: ]' % (
              input_filename, output_filename),
      ]
      expected_calls = [
          ('check_call',
              ('ls', input_filename))
      ]
      expected_ret_codes = [
          (1, 'Failed to fetch file %s for %s. [Err: ]' % (
              input_filename, output_filename))
      ]
      self.assertEqual(list(stdout_queue.queue), expected_output)
      self.assertEqual(self.gsutil.history, expected_calls)
      self.assertEqual(list(self.ret_codes.queue), expected_ret_codes)
  
    def test_download_cp_fails(self):
      sha1_hash = '7871c8e24da15bad8b0be2c36edc9dc77e37727f'
      input_filename = '%s/%s' % (self.base_url, sha1_hash)
      output_filename = os.path.join(self.base_path, 'uploaded_lorem_ipsum.txt')
      self.gsutil.add_expected(0, '', '')
      self.gsutil.add_expected(101, '', 'Test error message.')
      code = download_from_google_storage.download_from_google_storage(
          input_filename=sha1_hash,
          base_url=self.base_url,
          gsutil=self.gsutil,
          num_threads=1,
          directory=False,
          recursive=False,
          force=True,
          output=output_filename,
          ignore_errors=False,
          sha1_file=False,
          verbose=True,
          auto_platform=False,
          extract=False)
      expected_calls = [
          ('check_call',
              ('ls', input_filename)),
          ('check_call',
              ('cp', input_filename, output_filename))
      ]
      if sys.platform != 'win32':
        expected_calls.append(
            ('check_call',
             ('stat',
              'gs://sometesturl/7871c8e24da15bad8b0be2c36edc9dc77e37727f')))
      self.assertEqual(self.gsutil.history, expected_calls)
      self.assertEqual(code, 101)
  
    def test_corrupt_download(self):
      q = Queue.Queue()
      out_q = Queue.Queue()
      ret_codes = Queue.Queue()
      tmp_dir = tempfile.mkdtemp()
      sha1_hash = '7871c8e24da15bad8b0be2c36edc9dc77e37727f'
      output_filename = os.path.join(tmp_dir, 'lorem_ipsum.txt')
      q.put(('7871c8e24da15bad8b0be2c36edc9dc77e37727f', output_filename))
      q.put((None, None))
      def _write_bad_file():
        with open(output_filename, 'w') as f:
          f.write('foobar')
      self.gsutil.add_expected(0, '', '')
      self.gsutil.add_expected(0, '', '', _write_bad_file)
      download_from_google_storage._downloader_worker_thread(
          1, q, True, self.base_url, self.gsutil, out_q, ret_codes, True, False)
      self.assertTrue(q.empty())
      msg = ('1> ERROR remote sha1 (%s) does not match expected sha1 (%s).' %
             ('8843d7f92416211de9ebb963ff4ce28125932878', sha1_hash))
      self.assertEquals(out_q.get(), '1> Downloading %s...' % output_filename)
      self.assertEquals(out_q.get(), msg)
      self.assertEquals(ret_codes.get(), (20, msg))
      self.assertTrue(out_q.empty())
      self.assertTrue(ret_codes.empty())
  
  
    def test_download_directory_no_recursive_non_force(self):
      sha1_hash = '7871c8e24da15bad8b0be2c36edc9dc77e37727f'
      input_filename = '%s/%s' % (self.base_url, sha1_hash)
      output_filename = os.path.join(self.base_path, 'uploaded_lorem_ipsum.txt')
      code = download_from_google_storage.download_from_google_storage(
          input_filename=self.base_path,
          base_url=self.base_url,
          gsutil=self.gsutil,
          num_threads=1,
          directory=True,
          recursive=False,
          force=False,
          output=None,
          ignore_errors=False,
          sha1_file=False,
          verbose=True,
          auto_platform=False,
          extract=False)
      expected_calls = [
          ('check_call',
              ('ls', input_filename)),
          ('check_call',
              ('cp', input_filename, output_filename))]
      if sys.platform != 'win32':
        expected_calls.append(
            ('check_call',
             ('stat',
              'gs://sometesturl/7871c8e24da15bad8b0be2c36edc9dc77e37727f')))
      self.assertEqual(self.gsutil.history, expected_calls)
      self.assertEqual(code, 0)
  
  
  if __name__ == '__main__':
    unittest.main()
  