1
    2
    3
    4
    5
    6
    7
    8
    9
   10
   11
   12
   13
   14
   15
   16
   17
   18
   19
   20
   21
   22
   23
   24
   25
   26
   27
   28
   29
   30
   31
   32
   33
   34
   35
   36
   37
   38
   39
   40
   41
   42
   43
   44
   45
   46
   47
   48
   49
   50
   51
   52
   53
   54
   55
   56
   57
   58
   59
   60
   61
   62
   63
   64
   65
   66
   67
   68
   69
   70
   71
   72
   73
   74
   75
   76
   77
   78
   79
   80
   81
   82
   83
   84
   85


build / fuchsia / gcs_download_test.py [blame]

#!/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('gs_util_wrapper.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, check=True),
        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(tarfile.ReadError):
      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()