You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
296 lines
11 KiB
Python
296 lines
11 KiB
Python
#!/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=W0212
|
|
|
|
"""Unit tests for download_from_google_storage.py."""
|
|
|
|
import optparse
|
|
import os
|
|
import Queue
|
|
import shutil
|
|
import sys
|
|
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__))),
|
|
'third_party', 'gsutil', 'gsutil')
|
|
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):
|
|
self.expected.append((return_code, out, err))
|
|
|
|
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:
|
|
return self.expected.pop(0)[0]
|
|
else:
|
|
return 0
|
|
|
|
def check_call(self, *args):
|
|
with self.lock:
|
|
self.append_history('check_call', args)
|
|
if self.expected:
|
|
return self.expected.pop(0)
|
|
else:
|
|
return (0, '', '')
|
|
|
|
|
|
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_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)
|
|
expected_calls = [
|
|
('check_call',
|
|
('ls', input_filename)),
|
|
('check_call',
|
|
('cp', '-q', input_filename, output_filename))]
|
|
if sys.platform != 'win32':
|
|
expected_calls.append(
|
|
('check_call',
|
|
('ls',
|
|
'-L',
|
|
'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)
|
|
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_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)
|
|
expected_output = [
|
|
'0> File %s for %s does not exist, skipping.' % (
|
|
input_filename, output_filename),
|
|
]
|
|
expected_calls = [
|
|
('check_call',
|
|
('ls', input_filename))
|
|
]
|
|
expected_ret_codes = [
|
|
(1, 'File %s for %s does not exist.' % (
|
|
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)
|
|
expected_calls = [
|
|
('check_call',
|
|
('ls', input_filename)),
|
|
('check_call',
|
|
('cp', '-q', input_filename, output_filename))
|
|
]
|
|
if sys.platform != 'win32':
|
|
expected_calls.append(
|
|
('check_call',
|
|
('ls',
|
|
'-L',
|
|
'gs://sometesturl/7871c8e24da15bad8b0be2c36edc9dc77e37727f')))
|
|
self.assertEqual(self.gsutil.history, expected_calls)
|
|
self.assertEqual(code, 101)
|
|
|
|
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)
|
|
expected_calls = [
|
|
('check_call',
|
|
('ls', input_filename)),
|
|
('check_call',
|
|
('cp', '-q', input_filename, output_filename))]
|
|
if sys.platform != 'win32':
|
|
expected_calls.append(
|
|
('check_call',
|
|
('ls',
|
|
'-L',
|
|
'gs://sometesturl/7871c8e24da15bad8b0be2c36edc9dc77e37727f')))
|
|
self.assertEqual(self.gsutil.history, expected_calls)
|
|
self.assertEqual(code, 0)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|