build/fuchsia/gcs_download_test.py

   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
   6 import os
   7 import subprocess
   8 import tarfile
   9 import unittest
  10 from unittest import mock
  11
  12 from gcs_download import DownloadAndUnpackFromCloudStorage
  13
  14
  15 def _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')
  31 class 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
  87 if __name__ == '__main__':
  88   unittest.main()