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