89 lines
2.9 KiB
Python
Executable File
89 lines
2.9 KiB
Python
Executable File
#!/usr/bin/env vpython3
|
|
# Copyright 2022 The Chromium Authors
|
|
# Use of this source code is governed by a BSD-style license that can be
|
|
# found in the LICENSE file.
|
|
|
|
import os
|
|
import subprocess
|
|
import tarfile
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
from gcs_download import DownloadAndUnpackFromCloudStorage
|
|
|
|
|
|
def _mock_task(status_code: int = 0, stderr: str = '') -> mock.Mock:
|
|
task_mock = mock.Mock()
|
|
attrs = {
|
|
'returncode': status_code,
|
|
'wait.return_value': status_code,
|
|
'communicate.return_value': (None, stderr.encode()),
|
|
}
|
|
task_mock.configure_mock(**attrs)
|
|
|
|
return task_mock
|
|
|
|
|
|
@mock.patch('tempfile.TemporaryDirectory')
|
|
@mock.patch('subprocess.run')
|
|
@mock.patch('tarfile.open')
|
|
@unittest.skipIf(os.name == 'nt', 'Fuchsia tests not supported on Windows')
|
|
class TestDownloadAndUnpackFromCloudStorage(unittest.TestCase):
|
|
def testHappyPath(self, mock_tarfile, mock_run, mock_tmp_dir):
|
|
mock_run.return_value = _mock_task()
|
|
|
|
tmp_dir = os.path.join('some', 'tmp', 'dir')
|
|
mock_tmp_dir.return_value.__enter__.return_value = tmp_dir
|
|
|
|
mock_seq = mock.Mock()
|
|
mock_seq.attach_mock(mock_run, 'Run')
|
|
mock_seq.attach_mock(mock_tarfile, 'Untar')
|
|
mock_seq.attach_mock(mock_tmp_dir, 'MkTmpD')
|
|
|
|
output_dir = os.path.join('output', 'dir')
|
|
DownloadAndUnpackFromCloudStorage('gs://some/url', output_dir)
|
|
|
|
image_tgz_path = os.path.join(tmp_dir, 'image.tgz')
|
|
mock_seq.assert_has_calls([
|
|
mock.call.MkTmpD(),
|
|
mock.call.MkTmpD().__enter__(),
|
|
mock.call.Run(mock.ANY,
|
|
stderr=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
check=True,
|
|
encoding='utf-8'),
|
|
mock.call.Untar(name=image_tgz_path, mode='r|gz'),
|
|
mock.call.Untar().extractall(path=output_dir),
|
|
mock.call.MkTmpD().__exit__(None, None, None)
|
|
],
|
|
any_order=False)
|
|
|
|
# Verify cmd.
|
|
cmd = ' '.join(mock_run.call_args[0][0])
|
|
self.assertRegex(
|
|
cmd, r'.*python3?\s.*gsutil.py\s+cp\s+gs://some/url\s+' + image_tgz_path)
|
|
|
|
def testFailedTarOpen(self, mock_tarfile, mock_run, mock_tmp_dir):
|
|
mock_run.return_value = _mock_task(stderr='some error')
|
|
mock_tarfile.side_effect = tarfile.ReadError()
|
|
|
|
with self.assertRaises(subprocess.CalledProcessError):
|
|
DownloadAndUnpackFromCloudStorage('', '')
|
|
mock_tmp_dir.assert_called_once()
|
|
mock_run.assert_called_once()
|
|
mock_tarfile.assert_called_once()
|
|
|
|
def testBadTaskStatusCode(self, mock_tarfile, mock_run, mock_tmp_dir):
|
|
mock_run.side_effect = subprocess.CalledProcessError(cmd='some/command',
|
|
returncode=1)
|
|
|
|
with self.assertRaises(subprocess.CalledProcessError):
|
|
DownloadAndUnpackFromCloudStorage('', '')
|
|
mock_run.assert_called_once()
|
|
mock_tarfile.assert_not_called()
|
|
mock_tmp_dir.assert_called_once()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|