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('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, 51 stderr=subprocess.PIPE, 52 stdout=subprocess.PIPE, 53 check=True, 54 encoding='utf-8'), 55 mock.call.Untar(name=image_tgz_path, mode='r|gz'), 56 mock.call.Untar().extractall(path=output_dir), 57 mock.call.MkTmpD().__exit__(None, None, None) 58 ], 59 any_order=False) 60 61 # Verify cmd. 62 cmd = ' '.join(mock_run.call_args[0][0]) 63 self.assertRegex( 64 cmd, r'.*python3?\s.*gsutil.py\s+cp\s+gs://some/url\s+' + image_tgz_path) 65 66 def testFailedTarOpen(self, mock_tarfile, mock_run, mock_tmp_dir): 67 mock_run.return_value = _mock_task(stderr='some error') 68 mock_tarfile.side_effect = tarfile.ReadError() 69 70 with self.assertRaises(subprocess.CalledProcessError): 71 DownloadAndUnpackFromCloudStorage('', '') 72 mock_tmp_dir.assert_called_once() 73 mock_run.assert_called_once() 74 mock_tarfile.assert_called_once() 75 76 def testBadTaskStatusCode(self, mock_tarfile, mock_run, mock_tmp_dir): 77 mock_run.side_effect = subprocess.CalledProcessError(cmd='some/command', 78 returncode=1) 79 80 with self.assertRaises(subprocess.CalledProcessError): 81 DownloadAndUnpackFromCloudStorage('', '') 82 mock_run.assert_called_once() 83 mock_tarfile.assert_not_called() 84 mock_tmp_dir.assert_called_once() 85 86 87if __name__ == '__main__': 88 unittest.main() 89