1*6777b538SAndroid Build Coastguard Worker#!/usr/bin/env vpython3 2*6777b538SAndroid Build Coastguard Worker# Copyright 2022 The Chromium Authors 3*6777b538SAndroid Build Coastguard Worker# Use of this source code is governed by a BSD-style license that can be 4*6777b538SAndroid Build Coastguard Worker# found in the LICENSE file. 5*6777b538SAndroid Build Coastguard Worker 6*6777b538SAndroid Build Coastguard Worker# TODO(crbug.com/1496426): Remove this file. 7*6777b538SAndroid Build Coastguard Worker 8*6777b538SAndroid Build Coastguard Workerimport unittest 9*6777b538SAndroid Build Coastguard Workerfrom unittest import mock 10*6777b538SAndroid Build Coastguard Worker 11*6777b538SAndroid Build Coastguard Workerfrom parameterized import parameterized 12*6777b538SAndroid Build Coastguard Worker 13*6777b538SAndroid Build Coastguard Workerfrom update_images import _GetImageOverrideInfo 14*6777b538SAndroid Build Coastguard Workerfrom update_images import GetImageLocationInfo 15*6777b538SAndroid Build Coastguard Worker 16*6777b538SAndroid Build Coastguard Worker 17*6777b538SAndroid Build Coastguard Worker@mock.patch('update_images.GetSDKOverrideGCSPath') 18*6777b538SAndroid Build Coastguard Workerclass TestGetImageOverrideInfo(unittest.TestCase): 19*6777b538SAndroid Build Coastguard Worker def testLocationIsNone(self, mock_sdk_loc): 20*6777b538SAndroid Build Coastguard Worker mock_sdk_loc.return_value = None 21*6777b538SAndroid Build Coastguard Worker 22*6777b538SAndroid Build Coastguard Worker actual = _GetImageOverrideInfo() 23*6777b538SAndroid Build Coastguard Worker self.assertIsNone(actual) 24*6777b538SAndroid Build Coastguard Worker 25*6777b538SAndroid Build Coastguard Worker def testBadLocationStr(self, mock_sdk_loc): 26*6777b538SAndroid Build Coastguard Worker mock_sdk_loc.return_value = 'bad-format-string' 27*6777b538SAndroid Build Coastguard Worker 28*6777b538SAndroid Build Coastguard Worker with self.assertRaises(Exception): 29*6777b538SAndroid Build Coastguard Worker _GetImageOverrideInfo() 30*6777b538SAndroid Build Coastguard Worker 31*6777b538SAndroid Build Coastguard Worker @parameterized.expand([ 32*6777b538SAndroid Build Coastguard Worker ('gs://my-bucket/development/my-hash/sdk', { 33*6777b538SAndroid Build Coastguard Worker 'bucket': 'my-bucket', 34*6777b538SAndroid Build Coastguard Worker 'image_hash': 'my-hash' 35*6777b538SAndroid Build Coastguard Worker }), 36*6777b538SAndroid Build Coastguard Worker ('gs://my-bucket/development/my-hash', { 37*6777b538SAndroid Build Coastguard Worker 'bucket': 'my-bucket', 38*6777b538SAndroid Build Coastguard Worker 'image_hash': 'my-hash' 39*6777b538SAndroid Build Coastguard Worker }), 40*6777b538SAndroid Build Coastguard Worker ('gs://my-bucket/development/my-hash/', { 41*6777b538SAndroid Build Coastguard Worker 'bucket': 'my-bucket', 42*6777b538SAndroid Build Coastguard Worker 'image_hash': 'my-hash' 43*6777b538SAndroid Build Coastguard Worker }), 44*6777b538SAndroid Build Coastguard Worker ]) 45*6777b538SAndroid Build Coastguard Worker def testValidLocation(self, mock_sdk_loc, in_path, expected): 46*6777b538SAndroid Build Coastguard Worker mock_sdk_loc.return_value = in_path 47*6777b538SAndroid Build Coastguard Worker 48*6777b538SAndroid Build Coastguard Worker actual = _GetImageOverrideInfo() 49*6777b538SAndroid Build Coastguard Worker self.assertEqual(actual, expected) 50*6777b538SAndroid Build Coastguard Worker 51*6777b538SAndroid Build Coastguard Worker 52*6777b538SAndroid Build Coastguard Worker@mock.patch('update_images.GetImageHash') 53*6777b538SAndroid Build Coastguard Worker@mock.patch('update_images.GetOverrideCloudStorageBucket') 54*6777b538SAndroid Build Coastguard Worker@mock.patch('update_images._GetImageOverrideInfo') 55*6777b538SAndroid Build Coastguard Workerclass TestGetImageLocationInfo(unittest.TestCase): 56*6777b538SAndroid Build Coastguard Worker def testNoOverride(self, mock_image_override, mock_override_bucket, 57*6777b538SAndroid Build Coastguard Worker mock_image_hash): 58*6777b538SAndroid Build Coastguard Worker mock_image_override.return_value = None 59*6777b538SAndroid Build Coastguard Worker mock_override_bucket.return_value = None 60*6777b538SAndroid Build Coastguard Worker mock_image_hash.return_value = 'image-hash' 61*6777b538SAndroid Build Coastguard Worker 62*6777b538SAndroid Build Coastguard Worker actual = GetImageLocationInfo('my-bucket') 63*6777b538SAndroid Build Coastguard Worker self.assertEqual(actual, { 64*6777b538SAndroid Build Coastguard Worker 'bucket': 'my-bucket', 65*6777b538SAndroid Build Coastguard Worker 'image_hash': 'image-hash', 66*6777b538SAndroid Build Coastguard Worker }) 67*6777b538SAndroid Build Coastguard Worker 68*6777b538SAndroid Build Coastguard Worker def testOverride(self, mock_image_override, mock_override_bucket, 69*6777b538SAndroid Build Coastguard Worker mock_image_hash): 70*6777b538SAndroid Build Coastguard Worker override_info = { 71*6777b538SAndroid Build Coastguard Worker 'bucket': 'override-bucket', 72*6777b538SAndroid Build Coastguard Worker 'image_hash': 'override-hash', 73*6777b538SAndroid Build Coastguard Worker } 74*6777b538SAndroid Build Coastguard Worker mock_image_override.return_value = override_info 75*6777b538SAndroid Build Coastguard Worker mock_override_bucket.return_value = None 76*6777b538SAndroid Build Coastguard Worker mock_image_hash.return_value = 'image-hash' 77*6777b538SAndroid Build Coastguard Worker 78*6777b538SAndroid Build Coastguard Worker actual = GetImageLocationInfo('my-bucket') 79*6777b538SAndroid Build Coastguard Worker self.assertEqual(actual, override_info) 80*6777b538SAndroid Build Coastguard Worker 81*6777b538SAndroid Build Coastguard Worker def testNoAllowOverride(self, mock_image_override, mock_override_bucket, 82*6777b538SAndroid Build Coastguard Worker mock_image_hash): 83*6777b538SAndroid Build Coastguard Worker override_info = { 84*6777b538SAndroid Build Coastguard Worker 'bucket': 'override-bucket', 85*6777b538SAndroid Build Coastguard Worker 'image_hash': 'override-hash', 86*6777b538SAndroid Build Coastguard Worker } 87*6777b538SAndroid Build Coastguard Worker mock_image_override.return_value = override_info 88*6777b538SAndroid Build Coastguard Worker mock_override_bucket.return_value = None 89*6777b538SAndroid Build Coastguard Worker mock_image_hash.return_value = 'image-hash' 90*6777b538SAndroid Build Coastguard Worker 91*6777b538SAndroid Build Coastguard Worker actual = GetImageLocationInfo('my-bucket', allow_override=False) 92*6777b538SAndroid Build Coastguard Worker self.assertEqual(actual, { 93*6777b538SAndroid Build Coastguard Worker 'bucket': 'my-bucket', 94*6777b538SAndroid Build Coastguard Worker 'image_hash': 'image-hash', 95*6777b538SAndroid Build Coastguard Worker }) 96*6777b538SAndroid Build Coastguard Worker 97*6777b538SAndroid Build Coastguard Worker 98*6777b538SAndroid Build Coastguard Workerif __name__ == '__main__': 99*6777b538SAndroid Build Coastguard Worker unittest.main() 100