1#!/usr/bin/env vpython3 2# Copyright 2022 The Chromium Authors 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5 6import os 7import subprocess 8import tarfile 9import unittest 10from unittest import mock 11 12from gcs_download import DownloadAndUnpackFromCloudStorage 13 14 15def _mock_task(status_code: int = 0, stderr: str = '') -> mock.Mock: 16 task_mock = mock.Mock() 17 attrs = { 18 'returncode': status_code, 19 'wait.return_value': status_code, 20 'communicate.return_value': (None, stderr.encode()), 21 } 22 task_mock.configure_mock(**attrs) 23 24 return task_mock 25 26 27@mock.patch('tempfile.TemporaryDirectory') 28@mock.patch('gs_util_wrapper.subprocess.run') 29@mock.patch('tarfile.open') 30@unittest.skipIf(os.name == 'nt', 'Fuchsia tests not supported on Windows') 31class TestDownloadAndUnpackFromCloudStorage(unittest.TestCase): 32 def testHappyPath(self, mock_tarfile, mock_run, mock_tmp_dir): 33 mock_run.return_value = _mock_task() 34 35 tmp_dir = os.path.join('some', 'tmp', 'dir') 36 mock_tmp_dir.return_value.__enter__.return_value = tmp_dir 37 38 mock_seq = mock.Mock() 39 mock_seq.attach_mock(mock_run, 'Run') 40 mock_seq.attach_mock(mock_tarfile, 'Untar') 41 mock_seq.attach_mock(mock_tmp_dir, 'MkTmpD') 42 43 output_dir = os.path.join('output', 'dir') 44 DownloadAndUnpackFromCloudStorage('gs://some/url', output_dir) 45 46 image_tgz_path = os.path.join(tmp_dir, 'image.tgz') 47 mock_seq.assert_has_calls([ 48 mock.call.MkTmpD(), 49 mock.call.MkTmpD().__enter__(), 50 mock.call.Run(mock.ANY, check=True), 51 mock.call.Untar(name=image_tgz_path, mode='r|gz'), 52 mock.call.Untar().extractall(path=output_dir), 53 mock.call.MkTmpD().__exit__(None, None, None) 54 ], 55 any_order=False) 56 57 # Verify cmd. 58 cmd = ' '.join(mock_run.call_args[0][0]) 59 self.assertRegex( 60 cmd, 61 r'.*python3?\s.*gsutil.py\s+cp\s+gs://some/url\s+' + image_tgz_path) 62 63 def testFailedTarOpen(self, mock_tarfile, mock_run, mock_tmp_dir): 64 mock_run.return_value = _mock_task(stderr='some error') 65 mock_tarfile.side_effect = tarfile.ReadError() 66 67 with self.assertRaises(tarfile.ReadError): 68 DownloadAndUnpackFromCloudStorage('', '') 69 mock_tmp_dir.assert_called_once() 70 mock_run.assert_called_once() 71 mock_tarfile.assert_called_once() 72 73 def testBadTaskStatusCode(self, mock_tarfile, mock_run, mock_tmp_dir): 74 mock_run.side_effect = subprocess.CalledProcessError(cmd='some/command', 75 returncode=1) 76 77 with self.assertRaises(subprocess.CalledProcessError): 78 DownloadAndUnpackFromCloudStorage('', '') 79 mock_run.assert_called_once() 80 mock_tarfile.assert_not_called() 81 mock_tmp_dir.assert_called_once() 82 83 84if __name__ == '__main__': 85 unittest.main() 86