#!/usr/bin/env vpython3 # Copyright (c) 2024 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. """Unit tests for upload_to_google_storage_first_class.py.""" from io import StringIO import optparse import os import posixpath import shutil import sys import tarfile import tempfile import unittest sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) import upload_to_google_storage_first_class from download_from_google_storage_unittest import GsutilMock from download_from_google_storage_unittest import ChangedWorkingDirectory # ../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 UploadTests(unittest.TestCase): def setUp(self): self.gsutil = GsutilMock(GSUTIL_DEFAULT_PATH, None) self.temp_dir = tempfile.mkdtemp(prefix='gstools_test') self.base_path = os.path.join(self.temp_dir, 'gstools') shutil.copytree(os.path.join(TEST_DIR, 'gstools'), self.base_path) self.base_url = 'gs://sometesturl' self.parser = optparse.OptionParser() self.lorem_ipsum = os.path.join(self.base_path, 'lorem_ipsum.txt') def tearDown(self): shutil.rmtree(self.temp_dir) sys.stdin = sys.__stdin__ def test_upload_single_file_missing_generation(self): file = self.lorem_ipsum object_name = 'gs_object_name/version123' self.gsutil.add_expected(0, "", "") # ls call self.gsutil.add_expected(0, "", 'weird output') # cp call actual = upload_to_google_storage_first_class.upload_to_google_storage( file=file, base_url=self.base_url, object_name=object_name, gsutil=self.gsutil, force=True, gzip=False, dry_run=False) self.assertEqual( self.gsutil.history, [('check_call', ('ls', '%s/%s' % (self.base_url, object_name))), ('check_call', ('-h', 'Cache-Control:public, max-age=31536000', 'cp', '-v', file, '%s/%s' % (self.base_url, object_name)))]) self.assertEqual( actual, upload_to_google_storage_first_class.MISSING_GENERATION_MSG) def test_upload_single_file(self): file = self.lorem_ipsum object_name = 'gs_object_name/version123' self.gsutil.add_expected(0, "", "") # ls call expected_generation = '1712070862651948' expected_cp_status = ( f'Copying file://{file} [Content-Type=text/plain].\n' f'Created: {self.base_url}/{object_name}#{expected_generation}\n' 'Operation completed over 1 objects/8.0 B.') self.gsutil.add_expected(0, "", expected_cp_status) # cp call actual = upload_to_google_storage_first_class.upload_to_google_storage( file=file, base_url=self.base_url, object_name=object_name, gsutil=self.gsutil, force=True, gzip=False, dry_run=False) self.assertEqual( self.gsutil.history, [('check_call', ('ls', '%s/%s' % (self.base_url, object_name))), ('check_call', ('-h', 'Cache-Control:public, max-age=31536000', 'cp', '-v', file, '%s/%s' % (self.base_url, object_name)))]) self.assertEqual(actual, expected_generation) def test_upload_single_file_remote_exists(self): file = self.lorem_ipsum object_name = 'gs_object_name/version123' etag_string = 'ETag: 634d7c1ed3545383837428f031840a1e' self.gsutil.add_expected(0, b'', b'') self.gsutil.add_expected(0, etag_string, b'') with self.assertRaises(Exception): upload_to_google_storage_first_class.upload_to_google_storage( file=file, base_url=self.base_url, object_name=object_name, gsutil=self.gsutil, force=False, gzip=False, dry_run=False) self.assertEqual(self.gsutil.history, [('check_call', ('ls', '%s/%s' % (self.base_url, object_name))), ('check_call', ('ls', '-L', '%s/%s' % (self.base_url, object_name)))]) def test_create_archive(self): work_dir = os.path.join(self.base_path, 'download_test_data') with ChangedWorkingDirectory(work_dir): dirname = 'subfolder' dirs = [dirname] self.assertTrue( upload_to_google_storage_first_class.validate_archive_dirs( dirs)) tar_filename = upload_to_google_storage_first_class.create_archive( dirs) with tarfile.open(tar_filename, 'r:gz') as tar: content = map(lambda x: x.name, tar.getmembers()) self.assertIn(dirname, content) def test_create_archive_multiple_dirs(self): work_dir = os.path.join(self.base_path, 'download_test_data') with ChangedWorkingDirectory(work_dir): dirs = ['subfolder', 'subfolder2'] self.assertTrue( upload_to_google_storage_first_class.validate_archive_dirs( dirs)) tar_filename = upload_to_google_storage_first_class.create_archive( dirs) with tarfile.open(tar_filename, 'r:gz') as tar: content = map(lambda x: x.name, tar.getmembers()) for dirname in dirs: self.assertIn(dirname, content) @unittest.skipIf(sys.platform == 'win32', 'os.symlink does not exist on win') def test_validate_archive_dirs_fails(self): work_dir = os.path.join(self.base_path, 'download_test_data') with ChangedWorkingDirectory(work_dir): symlink = 'link' os.symlink(os.path.join(self.base_path, 'subfolder'), symlink) self.assertFalse( upload_to_google_storage_first_class.validate_archive_dirs( [symlink])) self.assertFalse( upload_to_google_storage_first_class.validate_archive_dirs( ['foobar'])) def test_dry_run(self): file = self.lorem_ipsum object_name = 'gs_object_name/version123' upload_to_google_storage_first_class.upload_to_google_storage( file=file, base_url=self.base_url, object_name=object_name, gsutil=self.gsutil, force=False, gzip=False, dry_run=True) self.assertEqual(self.gsutil.history, [('check_call', ('ls', '%s/%s' % (self.base_url, object_name))), ('check_call', ('ls', '-L', '%s/%s' % (self.base_url, object_name)))]) def test_get_targets_no_args(self): try: upload_to_google_storage_first_class.get_targets([], self.parser, False) self.fail() except SystemExit as e: self.assertEqual(e.code, 2) def test_get_targets_passthrough(self): result = upload_to_google_storage_first_class.get_targets( ['a', 'b', 'c', 'd', 'e'], self.parser, False) self.assertEqual(result, ['a', 'b', 'c', 'd', 'e']) def test_get_targets_multiple_stdin(self): inputs = ['a', 'b', 'c', 'd', 'e'] sys.stdin = StringIO(os.linesep.join(inputs)) result = upload_to_google_storage_first_class.get_targets(['-'], self.parser, False) self.assertEqual(result, inputs) def test_get_targets_multiple_stdin_null(self): inputs = ['a', 'b', 'c', 'd', 'e'] sys.stdin = StringIO('\0'.join(inputs)) result = upload_to_google_storage_first_class.get_targets(['-'], self.parser, True) self.assertEqual(result, inputs) if __name__ == '__main__': unittest.main()