• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 unittest
7from unittest import mock
8
9from parameterized import parameterized
10
11from update_images import _GetImageOverrideInfo
12from update_images import GetImageLocationInfo
13
14
15@mock.patch('update_images.GetSDKOverrideGCSPath')
16class TestGetImageOverrideInfo(unittest.TestCase):
17  def testLocationIsNone(self, mock_sdk_loc):
18    mock_sdk_loc.return_value = None
19
20    actual = _GetImageOverrideInfo()
21    self.assertIsNone(actual)
22
23  def testBadLocationStr(self, mock_sdk_loc):
24    mock_sdk_loc.return_value = 'bad-format-string'
25
26    with self.assertRaises(Exception):
27      _GetImageOverrideInfo()
28
29  @parameterized.expand([
30      ('gs://my-bucket/development/my-hash/sdk', {
31          'bucket': 'my-bucket',
32          'image_hash': 'my-hash'
33      }),
34      ('gs://my-bucket/development/my-hash', {
35          'bucket': 'my-bucket',
36          'image_hash': 'my-hash'
37      }),
38      ('gs://my-bucket/development/my-hash/', {
39          'bucket': 'my-bucket',
40          'image_hash': 'my-hash'
41      }),
42  ])
43  def testValidLocation(self, mock_sdk_loc, in_path, expected):
44    mock_sdk_loc.return_value = in_path
45
46    actual = _GetImageOverrideInfo()
47    self.assertEqual(actual, expected)
48
49
50@mock.patch('update_images.GetImageHash')
51@mock.patch('update_images.GetOverrideCloudStorageBucket')
52@mock.patch('update_images._GetImageOverrideInfo')
53class TestGetImageLocationInfo(unittest.TestCase):
54  def testNoOverride(self, mock_image_override, mock_override_bucket,
55                     mock_image_hash):
56    mock_image_override.return_value = None
57    mock_override_bucket.return_value = None
58    mock_image_hash.return_value = 'image-hash'
59
60    actual = GetImageLocationInfo('my-bucket')
61    self.assertEqual(actual, {
62        'bucket': 'my-bucket',
63        'image_hash': 'image-hash',
64    })
65
66  def testOverride(self, mock_image_override, mock_override_bucket,
67                   mock_image_hash):
68    override_info = {
69        'bucket': 'override-bucket',
70        'image_hash': 'override-hash',
71    }
72    mock_image_override.return_value = override_info
73    mock_override_bucket.return_value = None
74    mock_image_hash.return_value = 'image-hash'
75
76    actual = GetImageLocationInfo('my-bucket')
77    self.assertEqual(actual, override_info)
78
79  def testNoAllowOverride(self, mock_image_override, mock_override_bucket,
80                          mock_image_hash):
81    override_info = {
82        'bucket': 'override-bucket',
83        'image_hash': 'override-hash',
84    }
85    mock_image_override.return_value = override_info
86    mock_override_bucket.return_value = None
87    mock_image_hash.return_value = 'image-hash'
88
89    actual = GetImageLocationInfo('my-bucket', allow_override=False)
90    self.assertEqual(actual, {
91        'bucket': 'my-bucket',
92        'image_hash': 'image-hash',
93    })
94
95
96if __name__ == '__main__':
97  unittest.main()
98