xref: /aosp_15_r20/tools/acloud/internal/lib/gcompute_client_test.py (revision 800a58d989c669b8eb8a71d8df53b1ba3d411444)
1*800a58d9SAndroid Build Coastguard Worker#!/usr/bin/env python
2*800a58d9SAndroid Build Coastguard Worker#
3*800a58d9SAndroid Build Coastguard Worker# Copyright 2016 - The Android Open Source Project
4*800a58d9SAndroid Build Coastguard Worker#
5*800a58d9SAndroid Build Coastguard Worker# Licensed under the Apache License, Version 2.0 (the "License");
6*800a58d9SAndroid Build Coastguard Worker# you may not use this file except in compliance with the License.
7*800a58d9SAndroid Build Coastguard Worker# You may obtain a copy of the License at
8*800a58d9SAndroid Build Coastguard Worker#
9*800a58d9SAndroid Build Coastguard Worker#     http://www.apache.org/licenses/LICENSE-2.0
10*800a58d9SAndroid Build Coastguard Worker#
11*800a58d9SAndroid Build Coastguard Worker# Unless required by applicable law or agreed to in writing, software
12*800a58d9SAndroid Build Coastguard Worker# distributed under the License is distributed on an "AS IS" BASIS,
13*800a58d9SAndroid Build Coastguard Worker# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14*800a58d9SAndroid Build Coastguard Worker# See the License for the specific language governing permissions and
15*800a58d9SAndroid Build Coastguard Worker# limitations under the License.
16*800a58d9SAndroid Build Coastguard Worker"""Tests for acloud.internal.lib.gcompute_client."""
17*800a58d9SAndroid Build Coastguard Worker# pylint: disable=too-many-lines
18*800a58d9SAndroid Build Coastguard Worker
19*800a58d9SAndroid Build Coastguard Workerimport copy
20*800a58d9SAndroid Build Coastguard Workerimport os
21*800a58d9SAndroid Build Coastguard Worker
22*800a58d9SAndroid Build Coastguard Workerimport unittest
23*800a58d9SAndroid Build Coastguard Worker
24*800a58d9SAndroid Build Coastguard Workerfrom unittest import mock
25*800a58d9SAndroid Build Coastguard Worker
26*800a58d9SAndroid Build Coastguard Worker# pylint: disable=import-error
27*800a58d9SAndroid Build Coastguard Workerfrom acloud import errors
28*800a58d9SAndroid Build Coastguard Workerfrom acloud.internal import constants
29*800a58d9SAndroid Build Coastguard Workerfrom acloud.internal.lib import driver_test_lib
30*800a58d9SAndroid Build Coastguard Workerfrom acloud.internal.lib import gcompute_client
31*800a58d9SAndroid Build Coastguard Workerfrom acloud.internal.lib import utils
32*800a58d9SAndroid Build Coastguard Worker
33*800a58d9SAndroid Build Coastguard Worker
34*800a58d9SAndroid Build Coastguard WorkerGS_IMAGE_SOURCE_URI = "https://storage.googleapis.com/fake-bucket/fake.tar.gz"
35*800a58d9SAndroid Build Coastguard WorkerGS_IMAGE_SOURCE_DISK = (
36*800a58d9SAndroid Build Coastguard Worker    "https://www.googleapis.com/compute/v1/projects/fake-project/zones/"
37*800a58d9SAndroid Build Coastguard Worker    "us-east1-d/disks/fake-disk")
38*800a58d9SAndroid Build Coastguard WorkerPROJECT = "fake-project"
39*800a58d9SAndroid Build Coastguard Worker
40*800a58d9SAndroid Build Coastguard Worker
41*800a58d9SAndroid Build Coastguard Worker# pylint: disable=protected-access, too-many-public-methods
42*800a58d9SAndroid Build Coastguard Workerclass ComputeClientTest(driver_test_lib.BaseDriverTest):
43*800a58d9SAndroid Build Coastguard Worker    """Test ComputeClient."""
44*800a58d9SAndroid Build Coastguard Worker
45*800a58d9SAndroid Build Coastguard Worker    PROJECT_OTHER = "fake-project-other"
46*800a58d9SAndroid Build Coastguard Worker    INSTANCE = "fake-instance"
47*800a58d9SAndroid Build Coastguard Worker    IMAGE = "fake-image"
48*800a58d9SAndroid Build Coastguard Worker    IMAGE_URL = "http://fake-image-url"
49*800a58d9SAndroid Build Coastguard Worker    IMAGE_OTHER = "fake-image-other"
50*800a58d9SAndroid Build Coastguard Worker    DISK = "fake-disk"
51*800a58d9SAndroid Build Coastguard Worker    MACHINE_TYPE = "fake-machine-type"
52*800a58d9SAndroid Build Coastguard Worker    MACHINE_TYPE_URL = "http://fake-machine-type-url"
53*800a58d9SAndroid Build Coastguard Worker    METADATA = ("metadata_key", "metadata_value")
54*800a58d9SAndroid Build Coastguard Worker    ACCELERATOR_URL = "http://speedy-gpu"
55*800a58d9SAndroid Build Coastguard Worker    NETWORK = "fake-network"
56*800a58d9SAndroid Build Coastguard Worker    NETWORK_URL = "http://fake-network-url"
57*800a58d9SAndroid Build Coastguard Worker    SUBNETWORK_URL = "http://fake-subnetwork-url"
58*800a58d9SAndroid Build Coastguard Worker    ZONE = "fake-zone"
59*800a58d9SAndroid Build Coastguard Worker    REGION = "fake-region"
60*800a58d9SAndroid Build Coastguard Worker    OPERATION_NAME = "fake-op"
61*800a58d9SAndroid Build Coastguard Worker    IMAGE_FINGERPRINT = "L_NWHuz7wTY="
62*800a58d9SAndroid Build Coastguard Worker    GPU = "fancy-graphics"
63*800a58d9SAndroid Build Coastguard Worker    SSHKEY = (
64*800a58d9SAndroid Build Coastguard Worker        "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDBkTOTRze9v2VOqkkf7RG"
65*800a58d9SAndroid Build Coastguard Worker        "jSkg6Z2kb9Q9UHsDGatvend3fmjIw1Tugg0O7nnjlPkskmlgyd4a/j99WOeLL"
66*800a58d9SAndroid Build Coastguard Worker        "CPk6xPyoVjrPUVBU/pAk09ORTC4Zqk6YjlW7LOfzvqmXhmIZfYu6Q4Yt50pZzhl"
67*800a58d9SAndroid Build Coastguard Worker        "lllfu26nYjY7Tg12D019nJi/kqPX5+NKgt0LGXTu8T1r2Gav/q4V7QRWQrB8Eiu"
68*800a58d9SAndroid Build Coastguard Worker        "pxXR7I2YhynqovkEt/OXG4qWgvLEXGsWtSQs0CtCzqEVxz0Y9ECr7er4VdjSQxV"
69*800a58d9SAndroid Build Coastguard Worker        "AaeLAsQsK9ROae8hMBFZ3//8zLVapBwpuffCu+fUoql9qeV9xagZcc9zj8XOUOW"
70*800a58d9SAndroid Build Coastguard Worker        "ApiihqNL1111 [email protected]")
71*800a58d9SAndroid Build Coastguard Worker    EXTRA_SCOPES = ["scope1"]
72*800a58d9SAndroid Build Coastguard Worker
73*800a58d9SAndroid Build Coastguard Worker    def setUp(self):
74*800a58d9SAndroid Build Coastguard Worker        """Set up test."""
75*800a58d9SAndroid Build Coastguard Worker        super().setUp()
76*800a58d9SAndroid Build Coastguard Worker        self.Patch(gcompute_client.ComputeClient, "InitResourceHandle")
77*800a58d9SAndroid Build Coastguard Worker        fake_cfg = mock.MagicMock()
78*800a58d9SAndroid Build Coastguard Worker        fake_cfg.project = PROJECT
79*800a58d9SAndroid Build Coastguard Worker        fake_cfg.extra_scopes = self.EXTRA_SCOPES
80*800a58d9SAndroid Build Coastguard Worker        self.compute_client = gcompute_client.ComputeClient(
81*800a58d9SAndroid Build Coastguard Worker            fake_cfg, mock.MagicMock())
82*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service = mock.MagicMock()
83*800a58d9SAndroid Build Coastguard Worker
84*800a58d9SAndroid Build Coastguard Worker        self._disk_args = copy.deepcopy(gcompute_client.BASE_DISK_ARGS)
85*800a58d9SAndroid Build Coastguard Worker        self._disk_args["initializeParams"] = {"diskName": self.INSTANCE,
86*800a58d9SAndroid Build Coastguard Worker                                               "sourceImage": self.IMAGE_URL}
87*800a58d9SAndroid Build Coastguard Worker
88*800a58d9SAndroid Build Coastguard Worker    # pylint: disable=invalid-name
89*800a58d9SAndroid Build Coastguard Worker    def _SetupMocksForGetOperationStatus(self, mock_result, operation_scope):
90*800a58d9SAndroid Build Coastguard Worker        """A helper class for setting up mocks for testGetOperationStatus*.
91*800a58d9SAndroid Build Coastguard Worker
92*800a58d9SAndroid Build Coastguard Worker        Args:
93*800a58d9SAndroid Build Coastguard Worker            mock_result: The result to return by _GetOperationStatus.
94*800a58d9SAndroid Build Coastguard Worker            operation_scope: A value of OperationScope.
95*800a58d9SAndroid Build Coastguard Worker
96*800a58d9SAndroid Build Coastguard Worker        Returns:
97*800a58d9SAndroid Build Coastguard Worker            A mock for Resource object.
98*800a58d9SAndroid Build Coastguard Worker        """
99*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
100*800a58d9SAndroid Build Coastguard Worker        mock_api = mock.MagicMock()
101*800a58d9SAndroid Build Coastguard Worker        if operation_scope == gcompute_client.OperationScope.GLOBAL:
102*800a58d9SAndroid Build Coastguard Worker            self.compute_client._service.globalOperations = mock.MagicMock(
103*800a58d9SAndroid Build Coastguard Worker                return_value=resource_mock)
104*800a58d9SAndroid Build Coastguard Worker        elif operation_scope == gcompute_client.OperationScope.ZONE:
105*800a58d9SAndroid Build Coastguard Worker            self.compute_client._service.zoneOperations = mock.MagicMock(
106*800a58d9SAndroid Build Coastguard Worker                return_value=resource_mock)
107*800a58d9SAndroid Build Coastguard Worker        elif operation_scope == gcompute_client.OperationScope.REGION:
108*800a58d9SAndroid Build Coastguard Worker            self.compute_client._service.regionOperations = mock.MagicMock(
109*800a58d9SAndroid Build Coastguard Worker                return_value=resource_mock)
110*800a58d9SAndroid Build Coastguard Worker        resource_mock.get = mock.MagicMock(return_value=mock_api)
111*800a58d9SAndroid Build Coastguard Worker        mock_api.execute = mock.MagicMock(return_value=mock_result)
112*800a58d9SAndroid Build Coastguard Worker        return resource_mock
113*800a58d9SAndroid Build Coastguard Worker
114*800a58d9SAndroid Build Coastguard Worker    def testGetOperationStatusGlobal(self):
115*800a58d9SAndroid Build Coastguard Worker        """Test _GetOperationStatus for global."""
116*800a58d9SAndroid Build Coastguard Worker        resource_mock = self._SetupMocksForGetOperationStatus(
117*800a58d9SAndroid Build Coastguard Worker            {"status": "GOOD"}, gcompute_client.OperationScope.GLOBAL)
118*800a58d9SAndroid Build Coastguard Worker        status = self.compute_client._GetOperationStatus(
119*800a58d9SAndroid Build Coastguard Worker            {"name": self.OPERATION_NAME},
120*800a58d9SAndroid Build Coastguard Worker            gcompute_client.OperationScope.GLOBAL)
121*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(status, "GOOD")
122*800a58d9SAndroid Build Coastguard Worker        resource_mock.get.assert_called_with(
123*800a58d9SAndroid Build Coastguard Worker            project=PROJECT, operation=self.OPERATION_NAME)
124*800a58d9SAndroid Build Coastguard Worker
125*800a58d9SAndroid Build Coastguard Worker    def testGetOperationStatusZone(self):
126*800a58d9SAndroid Build Coastguard Worker        """Test _GetOperationStatus for zone."""
127*800a58d9SAndroid Build Coastguard Worker        resource_mock = self._SetupMocksForGetOperationStatus(
128*800a58d9SAndroid Build Coastguard Worker            {"status": "GOOD"}, gcompute_client.OperationScope.ZONE)
129*800a58d9SAndroid Build Coastguard Worker        status = self.compute_client._GetOperationStatus(
130*800a58d9SAndroid Build Coastguard Worker            {"name": self.OPERATION_NAME}, gcompute_client.OperationScope.ZONE,
131*800a58d9SAndroid Build Coastguard Worker            self.ZONE)
132*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(status, "GOOD")
133*800a58d9SAndroid Build Coastguard Worker        resource_mock.get.assert_called_with(
134*800a58d9SAndroid Build Coastguard Worker            project=PROJECT,
135*800a58d9SAndroid Build Coastguard Worker            operation=self.OPERATION_NAME,
136*800a58d9SAndroid Build Coastguard Worker            zone=self.ZONE)
137*800a58d9SAndroid Build Coastguard Worker
138*800a58d9SAndroid Build Coastguard Worker    def testGetOperationStatusRegion(self):
139*800a58d9SAndroid Build Coastguard Worker        """Test _GetOperationStatus for region."""
140*800a58d9SAndroid Build Coastguard Worker        resource_mock = self._SetupMocksForGetOperationStatus(
141*800a58d9SAndroid Build Coastguard Worker            {"status": "GOOD"}, gcompute_client.OperationScope.REGION)
142*800a58d9SAndroid Build Coastguard Worker        self.compute_client._GetOperationStatus(
143*800a58d9SAndroid Build Coastguard Worker            {"name": self.OPERATION_NAME},
144*800a58d9SAndroid Build Coastguard Worker            gcompute_client.OperationScope.REGION, self.REGION)
145*800a58d9SAndroid Build Coastguard Worker        resource_mock.get.assert_called_with(
146*800a58d9SAndroid Build Coastguard Worker            project=PROJECT, operation=self.OPERATION_NAME, region=self.REGION)
147*800a58d9SAndroid Build Coastguard Worker
148*800a58d9SAndroid Build Coastguard Worker    def testGetOperationStatusError(self):
149*800a58d9SAndroid Build Coastguard Worker        """Test _GetOperationStatus failed."""
150*800a58d9SAndroid Build Coastguard Worker        self._SetupMocksForGetOperationStatus(
151*800a58d9SAndroid Build Coastguard Worker            {"error": {"errors": ["error1", "error2"]}},
152*800a58d9SAndroid Build Coastguard Worker            gcompute_client.OperationScope.GLOBAL)
153*800a58d9SAndroid Build Coastguard Worker        self.assertRaisesRegex(
154*800a58d9SAndroid Build Coastguard Worker                              errors.DriverError,
155*800a58d9SAndroid Build Coastguard Worker                              "Get operation state failed.*error1.*error2",
156*800a58d9SAndroid Build Coastguard Worker                              self.compute_client._GetOperationStatus,
157*800a58d9SAndroid Build Coastguard Worker                              {"name": self.OPERATION_NAME},
158*800a58d9SAndroid Build Coastguard Worker                              gcompute_client.OperationScope.GLOBAL)
159*800a58d9SAndroid Build Coastguard Worker
160*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(errors, "GceOperationTimeoutError")
161*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(utils, "PollAndWait")
162*800a58d9SAndroid Build Coastguard Worker    def testWaitOnOperation(self, mock_poll, mock_gce_operation_timeout_error):
163*800a58d9SAndroid Build Coastguard Worker        """Test WaitOnOperation."""
164*800a58d9SAndroid Build Coastguard Worker        mock_error = mock.MagicMock()
165*800a58d9SAndroid Build Coastguard Worker        mock_gce_operation_timeout_error.return_value = mock_error
166*800a58d9SAndroid Build Coastguard Worker        self.compute_client.WaitOnOperation(
167*800a58d9SAndroid Build Coastguard Worker            operation={"name": self.OPERATION_NAME},
168*800a58d9SAndroid Build Coastguard Worker            operation_scope=gcompute_client.OperationScope.REGION,
169*800a58d9SAndroid Build Coastguard Worker            scope_name=self.REGION)
170*800a58d9SAndroid Build Coastguard Worker        mock_poll.assert_called_with(
171*800a58d9SAndroid Build Coastguard Worker            func=self.compute_client._GetOperationStatus,
172*800a58d9SAndroid Build Coastguard Worker            expected_return="DONE",
173*800a58d9SAndroid Build Coastguard Worker            timeout_exception=mock_error,
174*800a58d9SAndroid Build Coastguard Worker            timeout_secs=self.compute_client.OPERATION_TIMEOUT_SECS,
175*800a58d9SAndroid Build Coastguard Worker            sleep_interval_secs=self.compute_client.OPERATION_POLL_INTERVAL_SECS,
176*800a58d9SAndroid Build Coastguard Worker            operation={"name": self.OPERATION_NAME},
177*800a58d9SAndroid Build Coastguard Worker            operation_scope=gcompute_client.OperationScope.REGION,
178*800a58d9SAndroid Build Coastguard Worker            scope_name=self.REGION)
179*800a58d9SAndroid Build Coastguard Worker
180*800a58d9SAndroid Build Coastguard Worker    def testGetImage(self):
181*800a58d9SAndroid Build Coastguard Worker        """Test GetImage."""
182*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
183*800a58d9SAndroid Build Coastguard Worker        mock_api = mock.MagicMock()
184*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.images = mock.MagicMock(
185*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
186*800a58d9SAndroid Build Coastguard Worker        resource_mock.get = mock.MagicMock(return_value=mock_api)
187*800a58d9SAndroid Build Coastguard Worker        mock_api.execute = mock.MagicMock(return_value={"name": self.IMAGE})
188*800a58d9SAndroid Build Coastguard Worker        result = self.compute_client.GetImage(self.IMAGE)
189*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(result, {"name": self.IMAGE})
190*800a58d9SAndroid Build Coastguard Worker        resource_mock.get.assert_called_with(project=PROJECT, image=self.IMAGE)
191*800a58d9SAndroid Build Coastguard Worker
192*800a58d9SAndroid Build Coastguard Worker    def testGetImageOther(self):
193*800a58d9SAndroid Build Coastguard Worker        """Test GetImage with other project."""
194*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
195*800a58d9SAndroid Build Coastguard Worker        mock_api = mock.MagicMock()
196*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.images = mock.MagicMock(
197*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
198*800a58d9SAndroid Build Coastguard Worker        resource_mock.get = mock.MagicMock(return_value=mock_api)
199*800a58d9SAndroid Build Coastguard Worker        mock_api.execute = mock.MagicMock(return_value={"name": self.IMAGE_OTHER})
200*800a58d9SAndroid Build Coastguard Worker        result = self.compute_client.GetImage(
201*800a58d9SAndroid Build Coastguard Worker            image_name=self.IMAGE_OTHER,
202*800a58d9SAndroid Build Coastguard Worker            image_project=self.PROJECT_OTHER)
203*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(result, {"name": self.IMAGE_OTHER})
204*800a58d9SAndroid Build Coastguard Worker        resource_mock.get.assert_called_with(
205*800a58d9SAndroid Build Coastguard Worker            project=self.PROJECT_OTHER, image=self.IMAGE_OTHER)
206*800a58d9SAndroid Build Coastguard Worker
207*800a58d9SAndroid Build Coastguard Worker    def testCreateImageWithSourceURI(self):
208*800a58d9SAndroid Build Coastguard Worker        """Test CreateImage with src uri."""
209*800a58d9SAndroid Build Coastguard Worker        source_uri = GS_IMAGE_SOURCE_URI
210*800a58d9SAndroid Build Coastguard Worker        source_disk = None
211*800a58d9SAndroid Build Coastguard Worker        labels = None
212*800a58d9SAndroid Build Coastguard Worker        expected_body = {"name": self.IMAGE,
213*800a58d9SAndroid Build Coastguard Worker                         "rawDisk": {"source": GS_IMAGE_SOURCE_URI}}
214*800a58d9SAndroid Build Coastguard Worker        mock_check = self.Patch(gcompute_client.ComputeClient,
215*800a58d9SAndroid Build Coastguard Worker                                "CheckImageExists",
216*800a58d9SAndroid Build Coastguard Worker                                return_value=False)
217*800a58d9SAndroid Build Coastguard Worker        mock_wait = self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
218*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
219*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.images = mock.MagicMock(
220*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
221*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert = mock.MagicMock()
222*800a58d9SAndroid Build Coastguard Worker        self.compute_client.CreateImage(
223*800a58d9SAndroid Build Coastguard Worker            image_name=self.IMAGE, source_uri=source_uri,
224*800a58d9SAndroid Build Coastguard Worker            source_disk=source_disk, labels=labels)
225*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert.assert_called_with(
226*800a58d9SAndroid Build Coastguard Worker            project=PROJECT, body=expected_body)
227*800a58d9SAndroid Build Coastguard Worker        mock_wait.assert_called_with(
228*800a58d9SAndroid Build Coastguard Worker            operation=mock.ANY,
229*800a58d9SAndroid Build Coastguard Worker            operation_scope=gcompute_client.OperationScope.GLOBAL)
230*800a58d9SAndroid Build Coastguard Worker        mock_check.assert_called_with(self.IMAGE)
231*800a58d9SAndroid Build Coastguard Worker
232*800a58d9SAndroid Build Coastguard Worker    def testCreateImageWithSourceDisk(self):
233*800a58d9SAndroid Build Coastguard Worker        """Test CreateImage with src disk."""
234*800a58d9SAndroid Build Coastguard Worker        source_uri = None
235*800a58d9SAndroid Build Coastguard Worker        source_disk = GS_IMAGE_SOURCE_DISK
236*800a58d9SAndroid Build Coastguard Worker        labels = None
237*800a58d9SAndroid Build Coastguard Worker        expected_body = {"name": self.IMAGE,
238*800a58d9SAndroid Build Coastguard Worker                         "sourceDisk": GS_IMAGE_SOURCE_DISK}
239*800a58d9SAndroid Build Coastguard Worker        mock_check = self.Patch(gcompute_client.ComputeClient,
240*800a58d9SAndroid Build Coastguard Worker                                "CheckImageExists",
241*800a58d9SAndroid Build Coastguard Worker                                return_value=False)
242*800a58d9SAndroid Build Coastguard Worker        mock_wait = self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
243*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
244*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.images = mock.MagicMock(
245*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
246*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert = mock.MagicMock()
247*800a58d9SAndroid Build Coastguard Worker        self.compute_client.CreateImage(
248*800a58d9SAndroid Build Coastguard Worker            image_name=self.IMAGE, source_uri=source_uri,
249*800a58d9SAndroid Build Coastguard Worker            source_disk=source_disk, labels=labels)
250*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert.assert_called_with(
251*800a58d9SAndroid Build Coastguard Worker            project=PROJECT, body=expected_body)
252*800a58d9SAndroid Build Coastguard Worker        mock_wait.assert_called_with(
253*800a58d9SAndroid Build Coastguard Worker            operation=mock.ANY,
254*800a58d9SAndroid Build Coastguard Worker            operation_scope=gcompute_client.OperationScope.GLOBAL)
255*800a58d9SAndroid Build Coastguard Worker        mock_check.assert_called_with(self.IMAGE)
256*800a58d9SAndroid Build Coastguard Worker
257*800a58d9SAndroid Build Coastguard Worker    def testCreateImageWithSourceDiskAndLabel(self):
258*800a58d9SAndroid Build Coastguard Worker        """Test CreateImage with src disk and label."""
259*800a58d9SAndroid Build Coastguard Worker        source_uri = None
260*800a58d9SAndroid Build Coastguard Worker        source_disk = GS_IMAGE_SOURCE_DISK
261*800a58d9SAndroid Build Coastguard Worker        labels = {"label1": "xxx"}
262*800a58d9SAndroid Build Coastguard Worker        expected_body = {"name": self.IMAGE,
263*800a58d9SAndroid Build Coastguard Worker                         "sourceDisk": GS_IMAGE_SOURCE_DISK,
264*800a58d9SAndroid Build Coastguard Worker                         "labels": {"label1": "xxx"}}
265*800a58d9SAndroid Build Coastguard Worker        mock_check = self.Patch(gcompute_client.ComputeClient,
266*800a58d9SAndroid Build Coastguard Worker                                "CheckImageExists",
267*800a58d9SAndroid Build Coastguard Worker                                return_value=False)
268*800a58d9SAndroid Build Coastguard Worker        mock_wait = self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
269*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
270*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.images = mock.MagicMock(
271*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
272*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert = mock.MagicMock()
273*800a58d9SAndroid Build Coastguard Worker        self.compute_client.CreateImage(
274*800a58d9SAndroid Build Coastguard Worker            image_name=self.IMAGE, source_uri=source_uri,
275*800a58d9SAndroid Build Coastguard Worker            source_disk=source_disk, labels=labels)
276*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert.assert_called_with(
277*800a58d9SAndroid Build Coastguard Worker            project=PROJECT, body=expected_body)
278*800a58d9SAndroid Build Coastguard Worker        mock_wait.assert_called_with(
279*800a58d9SAndroid Build Coastguard Worker            operation=mock.ANY,
280*800a58d9SAndroid Build Coastguard Worker            operation_scope=gcompute_client.OperationScope.GLOBAL)
281*800a58d9SAndroid Build Coastguard Worker        mock_check.assert_called_with(self.IMAGE)
282*800a58d9SAndroid Build Coastguard Worker
283*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "GetImage")
284*800a58d9SAndroid Build Coastguard Worker    def testSetImageLabel(self, mock_get_image):
285*800a58d9SAndroid Build Coastguard Worker        """Test SetImageLabel."""
286*800a58d9SAndroid Build Coastguard Worker        with mock.patch.object(self.compute_client._service, "images",
287*800a58d9SAndroid Build Coastguard Worker                               return_value=mock.MagicMock(
288*800a58d9SAndroid Build Coastguard Worker                                   setLabels=mock.MagicMock())) as _:
289*800a58d9SAndroid Build Coastguard Worker            image = {"name": self.IMAGE,
290*800a58d9SAndroid Build Coastguard Worker                     "sourceDisk": GS_IMAGE_SOURCE_DISK,
291*800a58d9SAndroid Build Coastguard Worker                     "labelFingerprint": self.IMAGE_FINGERPRINT,
292*800a58d9SAndroid Build Coastguard Worker                     "labels": {"a": "aaa", "b": "bbb"}}
293*800a58d9SAndroid Build Coastguard Worker            mock_get_image.return_value = image
294*800a58d9SAndroid Build Coastguard Worker            new_labels = {"a": "xxx", "c": "ccc"}
295*800a58d9SAndroid Build Coastguard Worker            # Test
296*800a58d9SAndroid Build Coastguard Worker            self.compute_client.SetImageLabels(
297*800a58d9SAndroid Build Coastguard Worker                self.IMAGE, new_labels)
298*800a58d9SAndroid Build Coastguard Worker            # Check result
299*800a58d9SAndroid Build Coastguard Worker            expected_labels = {"a": "xxx", "b": "bbb", "c": "ccc"}
300*800a58d9SAndroid Build Coastguard Worker            self.compute_client._service.images().setLabels.assert_called_with(
301*800a58d9SAndroid Build Coastguard Worker                project=PROJECT,
302*800a58d9SAndroid Build Coastguard Worker                resource=self.IMAGE,
303*800a58d9SAndroid Build Coastguard Worker                body={
304*800a58d9SAndroid Build Coastguard Worker                    "labels": expected_labels,
305*800a58d9SAndroid Build Coastguard Worker                    "labelFingerprint": self.IMAGE_FINGERPRINT
306*800a58d9SAndroid Build Coastguard Worker                })
307*800a58d9SAndroid Build Coastguard Worker
308*800a58d9SAndroid Build Coastguard Worker    def testCreateImageRaiseDriverErrorWithValidInput(self):
309*800a58d9SAndroid Build Coastguard Worker        """Test CreateImage with valid input."""
310*800a58d9SAndroid Build Coastguard Worker        source_uri = GS_IMAGE_SOURCE_URI
311*800a58d9SAndroid Build Coastguard Worker        source_disk = GS_IMAGE_SOURCE_DISK
312*800a58d9SAndroid Build Coastguard Worker        self.Patch(gcompute_client.ComputeClient, "CheckImageExists", return_value=False)
313*800a58d9SAndroid Build Coastguard Worker        self.assertRaises(errors.DriverError, self.compute_client.CreateImage,
314*800a58d9SAndroid Build Coastguard Worker                          image_name=self.IMAGE, source_uri=source_uri,
315*800a58d9SAndroid Build Coastguard Worker                          source_disk=source_disk)
316*800a58d9SAndroid Build Coastguard Worker
317*800a58d9SAndroid Build Coastguard Worker    def testCreateImageRaiseDriverErrorWithInvalidInput(self):
318*800a58d9SAndroid Build Coastguard Worker        """Test CreateImage with valid input."""
319*800a58d9SAndroid Build Coastguard Worker        source_uri = None
320*800a58d9SAndroid Build Coastguard Worker        source_disk = None
321*800a58d9SAndroid Build Coastguard Worker        self.Patch(gcompute_client.ComputeClient, "CheckImageExists", return_value=False)
322*800a58d9SAndroid Build Coastguard Worker        self.assertRaises(errors.DriverError, self.compute_client.CreateImage,
323*800a58d9SAndroid Build Coastguard Worker                          image_name=self.IMAGE, source_uri=source_uri,
324*800a58d9SAndroid Build Coastguard Worker                          source_disk=source_disk)
325*800a58d9SAndroid Build Coastguard Worker
326*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "DeleteImage")
327*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "CheckImageExists",
328*800a58d9SAndroid Build Coastguard Worker                       side_effect=[False, True])
329*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation",
330*800a58d9SAndroid Build Coastguard Worker                       side_effect=errors.DriverError("Expected fake error"))
331*800a58d9SAndroid Build Coastguard Worker    def testCreateImageFail(self, mock_wait, mock_check, mock_delete):
332*800a58d9SAndroid Build Coastguard Worker        """Test CreateImage fails."""
333*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
334*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.images = mock.MagicMock(
335*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
336*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert = mock.MagicMock()
337*800a58d9SAndroid Build Coastguard Worker
338*800a58d9SAndroid Build Coastguard Worker        expected_body = {
339*800a58d9SAndroid Build Coastguard Worker            "name": self.IMAGE,
340*800a58d9SAndroid Build Coastguard Worker            "rawDisk": {
341*800a58d9SAndroid Build Coastguard Worker                "source": GS_IMAGE_SOURCE_URI,
342*800a58d9SAndroid Build Coastguard Worker            },
343*800a58d9SAndroid Build Coastguard Worker        }
344*800a58d9SAndroid Build Coastguard Worker        self.assertRaisesRegex(
345*800a58d9SAndroid Build Coastguard Worker            errors.DriverError,
346*800a58d9SAndroid Build Coastguard Worker            "Expected fake error",
347*800a58d9SAndroid Build Coastguard Worker            self.compute_client.CreateImage,
348*800a58d9SAndroid Build Coastguard Worker            image_name=self.IMAGE,
349*800a58d9SAndroid Build Coastguard Worker            source_uri=GS_IMAGE_SOURCE_URI)
350*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert.assert_called_with(
351*800a58d9SAndroid Build Coastguard Worker            project=PROJECT, body=expected_body)
352*800a58d9SAndroid Build Coastguard Worker        mock_wait.assert_called_with(
353*800a58d9SAndroid Build Coastguard Worker            operation=mock.ANY,
354*800a58d9SAndroid Build Coastguard Worker            operation_scope=gcompute_client.OperationScope.GLOBAL)
355*800a58d9SAndroid Build Coastguard Worker        mock_check.assert_called_with(self.IMAGE)
356*800a58d9SAndroid Build Coastguard Worker        mock_delete.assert_called_with(self.IMAGE)
357*800a58d9SAndroid Build Coastguard Worker
358*800a58d9SAndroid Build Coastguard Worker    def testCheckImageExistsTrue(self):
359*800a58d9SAndroid Build Coastguard Worker        """Test CheckImageExists return True."""
360*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
361*800a58d9SAndroid Build Coastguard Worker        mock_api = mock.MagicMock()
362*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.images = mock.MagicMock(
363*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
364*800a58d9SAndroid Build Coastguard Worker        resource_mock.get = mock.MagicMock(return_value=mock_api)
365*800a58d9SAndroid Build Coastguard Worker        mock_api.execute = mock.MagicMock(return_value={"name": self.IMAGE})
366*800a58d9SAndroid Build Coastguard Worker        self.assertTrue(self.compute_client.CheckImageExists(self.IMAGE))
367*800a58d9SAndroid Build Coastguard Worker
368*800a58d9SAndroid Build Coastguard Worker    def testCheckImageExistsFalse(self):
369*800a58d9SAndroid Build Coastguard Worker        """Test CheckImageExists return False."""
370*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
371*800a58d9SAndroid Build Coastguard Worker        mock_api = mock.MagicMock()
372*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.images = mock.MagicMock(
373*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
374*800a58d9SAndroid Build Coastguard Worker        resource_mock.get = mock.MagicMock(return_value=mock_api)
375*800a58d9SAndroid Build Coastguard Worker        mock_api.execute = mock.MagicMock(
376*800a58d9SAndroid Build Coastguard Worker            side_effect=errors.ResourceNotFoundError(404, "no image"))
377*800a58d9SAndroid Build Coastguard Worker        self.assertFalse(self.compute_client.CheckImageExists(self.IMAGE))
378*800a58d9SAndroid Build Coastguard Worker
379*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation")
380*800a58d9SAndroid Build Coastguard Worker    def testDeleteImage(self, mock_wait):
381*800a58d9SAndroid Build Coastguard Worker        """Test DeleteImage."""
382*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
383*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.images = mock.MagicMock(
384*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
385*800a58d9SAndroid Build Coastguard Worker        resource_mock.delete = mock.MagicMock()
386*800a58d9SAndroid Build Coastguard Worker        self.compute_client.DeleteImage(self.IMAGE)
387*800a58d9SAndroid Build Coastguard Worker        resource_mock.delete.assert_called_with(
388*800a58d9SAndroid Build Coastguard Worker            project=PROJECT, image=self.IMAGE)
389*800a58d9SAndroid Build Coastguard Worker        self.assertTrue(mock_wait.called)
390*800a58d9SAndroid Build Coastguard Worker
391*800a58d9SAndroid Build Coastguard Worker    def _SetupBatchHttpRequestMock(self):
392*800a58d9SAndroid Build Coastguard Worker        """Setup BatchHttpRequest mock."""
393*800a58d9SAndroid Build Coastguard Worker        requests = {}
394*800a58d9SAndroid Build Coastguard Worker
395*800a58d9SAndroid Build Coastguard Worker        def _Add(request, callback, request_id):
396*800a58d9SAndroid Build Coastguard Worker            requests[request_id] = (request, callback)
397*800a58d9SAndroid Build Coastguard Worker
398*800a58d9SAndroid Build Coastguard Worker        def _Execute():
399*800a58d9SAndroid Build Coastguard Worker            for rid in requests:
400*800a58d9SAndroid Build Coastguard Worker                _, callback = requests[rid]
401*800a58d9SAndroid Build Coastguard Worker                callback(
402*800a58d9SAndroid Build Coastguard Worker                    request_id=rid, response=mock.MagicMock(), exception=None)
403*800a58d9SAndroid Build Coastguard Worker        mock_batch = mock.MagicMock()
404*800a58d9SAndroid Build Coastguard Worker        mock_batch.add = _Add
405*800a58d9SAndroid Build Coastguard Worker        mock_batch.execute = _Execute
406*800a58d9SAndroid Build Coastguard Worker        self.Patch(self.compute_client._service,
407*800a58d9SAndroid Build Coastguard Worker                   "new_batch_http_request",
408*800a58d9SAndroid Build Coastguard Worker                   return_value=mock_batch)
409*800a58d9SAndroid Build Coastguard Worker
410*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation")
411*800a58d9SAndroid Build Coastguard Worker    def testDeleteImages(self, mock_wait):
412*800a58d9SAndroid Build Coastguard Worker        """Test DeleteImages."""
413*800a58d9SAndroid Build Coastguard Worker        self._SetupBatchHttpRequestMock()
414*800a58d9SAndroid Build Coastguard Worker        fake_images = ["fake_image_1", "fake_image_2"]
415*800a58d9SAndroid Build Coastguard Worker        mock_api = mock.MagicMock()
416*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
417*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.images = mock.MagicMock(
418*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
419*800a58d9SAndroid Build Coastguard Worker        resource_mock.delete = mock.MagicMock(return_value=mock_api)
420*800a58d9SAndroid Build Coastguard Worker        # Call the API.
421*800a58d9SAndroid Build Coastguard Worker        deleted, failed, error_msgs = self.compute_client.DeleteImages(
422*800a58d9SAndroid Build Coastguard Worker            fake_images)
423*800a58d9SAndroid Build Coastguard Worker        # Verify
424*800a58d9SAndroid Build Coastguard Worker        calls = [
425*800a58d9SAndroid Build Coastguard Worker            mock.call(project=PROJECT, image="fake_image_1"),
426*800a58d9SAndroid Build Coastguard Worker            mock.call(project=PROJECT, image="fake_image_2")
427*800a58d9SAndroid Build Coastguard Worker        ]
428*800a58d9SAndroid Build Coastguard Worker        resource_mock.delete.assert_has_calls(calls, any_order=True)
429*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(mock_wait.call_count, 2)
430*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(error_msgs, [])
431*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(failed, [])
432*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(set(deleted), set(fake_images))
433*800a58d9SAndroid Build Coastguard Worker
434*800a58d9SAndroid Build Coastguard Worker    def testListImages(self):
435*800a58d9SAndroid Build Coastguard Worker        """Test ListImages."""
436*800a58d9SAndroid Build Coastguard Worker        fake_token = "fake_next_page_token"
437*800a58d9SAndroid Build Coastguard Worker        image_1 = "image_1"
438*800a58d9SAndroid Build Coastguard Worker        image_2 = "image_2"
439*800a58d9SAndroid Build Coastguard Worker        response_1 = {"items": [image_1], "nextPageToken": fake_token}
440*800a58d9SAndroid Build Coastguard Worker        response_2 = {"items": [image_2]}
441*800a58d9SAndroid Build Coastguard Worker        self.Patch(
442*800a58d9SAndroid Build Coastguard Worker            gcompute_client.ComputeClient,
443*800a58d9SAndroid Build Coastguard Worker            "Execute",
444*800a58d9SAndroid Build Coastguard Worker            side_effect=[response_1, response_2])
445*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
446*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.images = mock.MagicMock(
447*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
448*800a58d9SAndroid Build Coastguard Worker        resource_mock.list = mock.MagicMock()
449*800a58d9SAndroid Build Coastguard Worker        images = self.compute_client.ListImages()
450*800a58d9SAndroid Build Coastguard Worker        calls = [
451*800a58d9SAndroid Build Coastguard Worker            mock.call(project=PROJECT, filter=None, pageToken=None),
452*800a58d9SAndroid Build Coastguard Worker            mock.call(project=PROJECT, filter=None, pageToken=fake_token)
453*800a58d9SAndroid Build Coastguard Worker        ]
454*800a58d9SAndroid Build Coastguard Worker        resource_mock.list.assert_has_calls(calls)
455*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(images, [image_1, image_2])
456*800a58d9SAndroid Build Coastguard Worker
457*800a58d9SAndroid Build Coastguard Worker    def testListImagesFromExternalProject(self):
458*800a58d9SAndroid Build Coastguard Worker        """Test ListImages which accepts different project."""
459*800a58d9SAndroid Build Coastguard Worker        image = "image_1"
460*800a58d9SAndroid Build Coastguard Worker        response = {"items": [image]}
461*800a58d9SAndroid Build Coastguard Worker        self.Patch(gcompute_client.ComputeClient, "Execute", side_effect=[response])
462*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
463*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.images = mock.MagicMock(
464*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
465*800a58d9SAndroid Build Coastguard Worker        resource_mock.list = mock.MagicMock()
466*800a58d9SAndroid Build Coastguard Worker        images = self.compute_client.ListImages(
467*800a58d9SAndroid Build Coastguard Worker            image_project="fake-project-2")
468*800a58d9SAndroid Build Coastguard Worker        calls = [
469*800a58d9SAndroid Build Coastguard Worker            mock.call(project="fake-project-2", filter=None, pageToken=None)]
470*800a58d9SAndroid Build Coastguard Worker        resource_mock.list.assert_has_calls(calls)
471*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(images, [image])
472*800a58d9SAndroid Build Coastguard Worker
473*800a58d9SAndroid Build Coastguard Worker    def testGetInstance(self):
474*800a58d9SAndroid Build Coastguard Worker        """Test GetInstance."""
475*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
476*800a58d9SAndroid Build Coastguard Worker        mock_api = mock.MagicMock()
477*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.instances = mock.MagicMock(
478*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
479*800a58d9SAndroid Build Coastguard Worker        resource_mock.get = mock.MagicMock(return_value=mock_api)
480*800a58d9SAndroid Build Coastguard Worker        mock_api.execute = mock.MagicMock(return_value={"name": self.INSTANCE})
481*800a58d9SAndroid Build Coastguard Worker        result = self.compute_client.GetInstance(self.INSTANCE, self.ZONE)
482*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(result, {"name": self.INSTANCE})
483*800a58d9SAndroid Build Coastguard Worker        resource_mock.get.assert_called_with(
484*800a58d9SAndroid Build Coastguard Worker            project=PROJECT, zone=self.ZONE, instance=self.INSTANCE)
485*800a58d9SAndroid Build Coastguard Worker
486*800a58d9SAndroid Build Coastguard Worker    def testListInstances(self):
487*800a58d9SAndroid Build Coastguard Worker        """Test ListInstances."""
488*800a58d9SAndroid Build Coastguard Worker        instance_1 = "instance_1"
489*800a58d9SAndroid Build Coastguard Worker        instance_2 = "instance_2"
490*800a58d9SAndroid Build Coastguard Worker        response = {"items": {'zones/fake_zone': {"instances": [instance_1, instance_2]}}}
491*800a58d9SAndroid Build Coastguard Worker        self.Patch(
492*800a58d9SAndroid Build Coastguard Worker            gcompute_client.ComputeClient,
493*800a58d9SAndroid Build Coastguard Worker            "Execute",
494*800a58d9SAndroid Build Coastguard Worker            side_effect=[response])
495*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
496*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.instances = mock.MagicMock(
497*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
498*800a58d9SAndroid Build Coastguard Worker        resource_mock.aggregatedList = mock.MagicMock()
499*800a58d9SAndroid Build Coastguard Worker        instances = self.compute_client.ListInstances()
500*800a58d9SAndroid Build Coastguard Worker        calls = [
501*800a58d9SAndroid Build Coastguard Worker            mock.call(
502*800a58d9SAndroid Build Coastguard Worker                project=PROJECT,
503*800a58d9SAndroid Build Coastguard Worker                filter=None,
504*800a58d9SAndroid Build Coastguard Worker                pageToken=None),
505*800a58d9SAndroid Build Coastguard Worker        ]
506*800a58d9SAndroid Build Coastguard Worker        resource_mock.aggregatedList.assert_has_calls(calls)
507*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(instances, [instance_1, instance_2])
508*800a58d9SAndroid Build Coastguard Worker
509*800a58d9SAndroid Build Coastguard Worker    def testGetZoneByInstance(self):
510*800a58d9SAndroid Build Coastguard Worker        """Test GetZoneByInstance."""
511*800a58d9SAndroid Build Coastguard Worker        instance_1 = "instance_1"
512*800a58d9SAndroid Build Coastguard Worker        response = {"items": {'zones/fake_zone': {"instances": [instance_1]}}}
513*800a58d9SAndroid Build Coastguard Worker        self.Patch(
514*800a58d9SAndroid Build Coastguard Worker            gcompute_client.ComputeClient,
515*800a58d9SAndroid Build Coastguard Worker            "Execute",
516*800a58d9SAndroid Build Coastguard Worker            side_effect=[response])
517*800a58d9SAndroid Build Coastguard Worker        expected_zone = "fake_zone"
518*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(self.compute_client.GetZoneByInstance(instance_1),
519*800a58d9SAndroid Build Coastguard Worker                         expected_zone)
520*800a58d9SAndroid Build Coastguard Worker
521*800a58d9SAndroid Build Coastguard Worker        # Test unable to find 'zone' from instance name.
522*800a58d9SAndroid Build Coastguard Worker        response = {"items": {'zones/fake_zone': {"warning": "No instances."}}}
523*800a58d9SAndroid Build Coastguard Worker        self.Patch(
524*800a58d9SAndroid Build Coastguard Worker            gcompute_client.ComputeClient,
525*800a58d9SAndroid Build Coastguard Worker            "Execute",
526*800a58d9SAndroid Build Coastguard Worker            side_effect=[response])
527*800a58d9SAndroid Build Coastguard Worker        with self.assertRaises(errors.GetGceZoneError):
528*800a58d9SAndroid Build Coastguard Worker            self.compute_client.GetZoneByInstance(instance_1)
529*800a58d9SAndroid Build Coastguard Worker
530*800a58d9SAndroid Build Coastguard Worker    def testGetZonesByInstances(self):
531*800a58d9SAndroid Build Coastguard Worker        """Test GetZonesByInstances."""
532*800a58d9SAndroid Build Coastguard Worker        instances = ["instance_1", "instance_2"]
533*800a58d9SAndroid Build Coastguard Worker        # Test instances in the same zone.
534*800a58d9SAndroid Build Coastguard Worker        self.Patch(
535*800a58d9SAndroid Build Coastguard Worker            gcompute_client.ComputeClient,
536*800a58d9SAndroid Build Coastguard Worker            "GetZoneByInstance",
537*800a58d9SAndroid Build Coastguard Worker            side_effect=["zone_1", "zone_1"])
538*800a58d9SAndroid Build Coastguard Worker        expected_result = {"zone_1": ["instance_1", "instance_2"]}
539*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(self.compute_client.GetZonesByInstances(instances),
540*800a58d9SAndroid Build Coastguard Worker                         expected_result)
541*800a58d9SAndroid Build Coastguard Worker
542*800a58d9SAndroid Build Coastguard Worker        # Test instances in different zones.
543*800a58d9SAndroid Build Coastguard Worker        self.Patch(
544*800a58d9SAndroid Build Coastguard Worker            gcompute_client.ComputeClient,
545*800a58d9SAndroid Build Coastguard Worker            "GetZoneByInstance",
546*800a58d9SAndroid Build Coastguard Worker            side_effect=["zone_1", "zone_2"])
547*800a58d9SAndroid Build Coastguard Worker        expected_result = {"zone_1": ["instance_1"],
548*800a58d9SAndroid Build Coastguard Worker                           "zone_2": ["instance_2"]}
549*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(self.compute_client.GetZonesByInstances(instances),
550*800a58d9SAndroid Build Coastguard Worker                         expected_result)
551*800a58d9SAndroid Build Coastguard Worker
552*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "GetImage")
553*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "GetNetworkUrl")
554*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "GetSubnetworkUrl")
555*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "GetMachineType")
556*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation")
557*800a58d9SAndroid Build Coastguard Worker    @mock.patch("getpass.getuser", return_value="fake_user")
558*800a58d9SAndroid Build Coastguard Worker    def testCreateInstance(self, _get_user, mock_wait, mock_get_mach_type,
559*800a58d9SAndroid Build Coastguard Worker                           mock_get_subnetwork_url, mock_get_network_url,
560*800a58d9SAndroid Build Coastguard Worker                           mock_get_image):
561*800a58d9SAndroid Build Coastguard Worker        """Test CreateInstance."""
562*800a58d9SAndroid Build Coastguard Worker        mock_get_mach_type.return_value = {"selfLink": self.MACHINE_TYPE_URL}
563*800a58d9SAndroid Build Coastguard Worker        mock_get_network_url.return_value = self.NETWORK_URL
564*800a58d9SAndroid Build Coastguard Worker        mock_get_subnetwork_url.return_value = self.SUBNETWORK_URL
565*800a58d9SAndroid Build Coastguard Worker        mock_get_image.return_value = {"selfLink": self.IMAGE_URL}
566*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
567*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.instances = mock.MagicMock(
568*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
569*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert = mock.MagicMock()
570*800a58d9SAndroid Build Coastguard Worker        self.Patch(
571*800a58d9SAndroid Build Coastguard Worker            self.compute_client,
572*800a58d9SAndroid Build Coastguard Worker            "_GetExtraDiskArgs",
573*800a58d9SAndroid Build Coastguard Worker            return_value=[{"fake_extra_arg": "fake_extra_value"}])
574*800a58d9SAndroid Build Coastguard Worker        extra_disk_name = "gce-x86-userdebug-2345-abcd-data"
575*800a58d9SAndroid Build Coastguard Worker        expected_disk_args = [self._disk_args]
576*800a58d9SAndroid Build Coastguard Worker        expected_disk_args.extend([{"fake_extra_arg": "fake_extra_value"}])
577*800a58d9SAndroid Build Coastguard Worker        expected_scope = []
578*800a58d9SAndroid Build Coastguard Worker        expected_scope.extend(self.compute_client.DEFAULT_INSTANCE_SCOPE)
579*800a58d9SAndroid Build Coastguard Worker        expected_scope.extend(self.EXTRA_SCOPES)
580*800a58d9SAndroid Build Coastguard Worker
581*800a58d9SAndroid Build Coastguard Worker        expected_body = {
582*800a58d9SAndroid Build Coastguard Worker            "machineType": self.MACHINE_TYPE_URL,
583*800a58d9SAndroid Build Coastguard Worker            "name": self.INSTANCE,
584*800a58d9SAndroid Build Coastguard Worker            "networkInterfaces": [
585*800a58d9SAndroid Build Coastguard Worker                {
586*800a58d9SAndroid Build Coastguard Worker                    "network": self.NETWORK_URL,
587*800a58d9SAndroid Build Coastguard Worker                    "subnetwork": self.SUBNETWORK_URL,
588*800a58d9SAndroid Build Coastguard Worker                    "accessConfigs": [
589*800a58d9SAndroid Build Coastguard Worker                        {"name": "External NAT",
590*800a58d9SAndroid Build Coastguard Worker                         "type": "ONE_TO_ONE_NAT"}
591*800a58d9SAndroid Build Coastguard Worker                    ],
592*800a58d9SAndroid Build Coastguard Worker                }
593*800a58d9SAndroid Build Coastguard Worker            ],
594*800a58d9SAndroid Build Coastguard Worker            "disks": expected_disk_args,
595*800a58d9SAndroid Build Coastguard Worker            "serviceAccounts": [
596*800a58d9SAndroid Build Coastguard Worker                {"email": "default",
597*800a58d9SAndroid Build Coastguard Worker                 "scopes": expected_scope}
598*800a58d9SAndroid Build Coastguard Worker            ],
599*800a58d9SAndroid Build Coastguard Worker            "metadata": {
600*800a58d9SAndroid Build Coastguard Worker                "items": [{"key": self.METADATA[0],
601*800a58d9SAndroid Build Coastguard Worker                           "value": self.METADATA[1]}],
602*800a58d9SAndroid Build Coastguard Worker            },
603*800a58d9SAndroid Build Coastguard Worker            "labels":{constants.LABEL_CREATE_BY: "fake_user"},
604*800a58d9SAndroid Build Coastguard Worker            "enableVtpm": True,
605*800a58d9SAndroid Build Coastguard Worker        }
606*800a58d9SAndroid Build Coastguard Worker
607*800a58d9SAndroid Build Coastguard Worker        self.compute_client.CreateInstance(
608*800a58d9SAndroid Build Coastguard Worker            instance=self.INSTANCE,
609*800a58d9SAndroid Build Coastguard Worker            image_name=self.IMAGE,
610*800a58d9SAndroid Build Coastguard Worker            machine_type=self.MACHINE_TYPE,
611*800a58d9SAndroid Build Coastguard Worker            metadata={self.METADATA[0]: self.METADATA[1]},
612*800a58d9SAndroid Build Coastguard Worker            network=self.NETWORK,
613*800a58d9SAndroid Build Coastguard Worker            zone=self.ZONE,
614*800a58d9SAndroid Build Coastguard Worker            extra_disk_name=extra_disk_name,
615*800a58d9SAndroid Build Coastguard Worker            extra_scopes=self.EXTRA_SCOPES)
616*800a58d9SAndroid Build Coastguard Worker
617*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert.assert_called_with(
618*800a58d9SAndroid Build Coastguard Worker            project=PROJECT, zone=self.ZONE, body=expected_body)
619*800a58d9SAndroid Build Coastguard Worker        mock_wait.assert_called_with(
620*800a58d9SAndroid Build Coastguard Worker            mock.ANY,
621*800a58d9SAndroid Build Coastguard Worker            operation_scope=gcompute_client.OperationScope.ZONE,
622*800a58d9SAndroid Build Coastguard Worker            scope_name=self.ZONE)
623*800a58d9SAndroid Build Coastguard Worker
624*800a58d9SAndroid Build Coastguard Worker
625*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "GetImage")
626*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "GetNetworkUrl")
627*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "GetSubnetworkUrl")
628*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "GetMachineType")
629*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation")
630*800a58d9SAndroid Build Coastguard Worker    @mock.patch("getpass.getuser", return_value="fake_user")
631*800a58d9SAndroid Build Coastguard Worker    def testCreateInstanceWithTags(self,
632*800a58d9SAndroid Build Coastguard Worker                                   _get_user,
633*800a58d9SAndroid Build Coastguard Worker                                   mock_wait,
634*800a58d9SAndroid Build Coastguard Worker                                   mock_get_mach_type,
635*800a58d9SAndroid Build Coastguard Worker                                   mock_get_subnetwork_url,
636*800a58d9SAndroid Build Coastguard Worker                                   mock_get_network_url,
637*800a58d9SAndroid Build Coastguard Worker                                   mock_get_image):
638*800a58d9SAndroid Build Coastguard Worker        """Test CreateInstance."""
639*800a58d9SAndroid Build Coastguard Worker        mock_get_mach_type.return_value = {"selfLink": self.MACHINE_TYPE_URL}
640*800a58d9SAndroid Build Coastguard Worker        mock_get_network_url.return_value = self.NETWORK_URL
641*800a58d9SAndroid Build Coastguard Worker        mock_get_subnetwork_url.return_value = self.SUBNETWORK_URL
642*800a58d9SAndroid Build Coastguard Worker        mock_get_image.return_value = {"selfLink": self.IMAGE_URL}
643*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
644*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.instances = mock.MagicMock(
645*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
646*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert = mock.MagicMock()
647*800a58d9SAndroid Build Coastguard Worker        self.Patch(
648*800a58d9SAndroid Build Coastguard Worker            self.compute_client,
649*800a58d9SAndroid Build Coastguard Worker            "_GetExtraDiskArgs",
650*800a58d9SAndroid Build Coastguard Worker            return_value=[{"fake_extra_arg": "fake_extra_value"}])
651*800a58d9SAndroid Build Coastguard Worker        extra_disk_name = "gce-x86-userdebug-2345-abcd-data"
652*800a58d9SAndroid Build Coastguard Worker        expected_disk_args = [self._disk_args]
653*800a58d9SAndroid Build Coastguard Worker        expected_disk_args.extend([{"fake_extra_arg": "fake_extra_value"}])
654*800a58d9SAndroid Build Coastguard Worker        expected_scope = []
655*800a58d9SAndroid Build Coastguard Worker        expected_scope.extend(self.compute_client.DEFAULT_INSTANCE_SCOPE)
656*800a58d9SAndroid Build Coastguard Worker        expected_scope.extend(self.EXTRA_SCOPES)
657*800a58d9SAndroid Build Coastguard Worker
658*800a58d9SAndroid Build Coastguard Worker        expected_body = {
659*800a58d9SAndroid Build Coastguard Worker            "machineType": self.MACHINE_TYPE_URL,
660*800a58d9SAndroid Build Coastguard Worker            "name": self.INSTANCE,
661*800a58d9SAndroid Build Coastguard Worker            "networkInterfaces": [
662*800a58d9SAndroid Build Coastguard Worker                {
663*800a58d9SAndroid Build Coastguard Worker                    "network": self.NETWORK_URL,
664*800a58d9SAndroid Build Coastguard Worker                    "subnetwork": self.SUBNETWORK_URL,
665*800a58d9SAndroid Build Coastguard Worker                    "accessConfigs": [
666*800a58d9SAndroid Build Coastguard Worker                        {"name": "External NAT",
667*800a58d9SAndroid Build Coastguard Worker                         "type": "ONE_TO_ONE_NAT"}
668*800a58d9SAndroid Build Coastguard Worker                    ],
669*800a58d9SAndroid Build Coastguard Worker                }
670*800a58d9SAndroid Build Coastguard Worker            ],
671*800a58d9SAndroid Build Coastguard Worker            'tags': {'items': ['https-server']},
672*800a58d9SAndroid Build Coastguard Worker            "disks": expected_disk_args,
673*800a58d9SAndroid Build Coastguard Worker            "serviceAccounts": [
674*800a58d9SAndroid Build Coastguard Worker                {"email": "default",
675*800a58d9SAndroid Build Coastguard Worker                 "scopes": expected_scope}
676*800a58d9SAndroid Build Coastguard Worker            ],
677*800a58d9SAndroid Build Coastguard Worker            "metadata": {
678*800a58d9SAndroid Build Coastguard Worker                "items": [{"key": self.METADATA[0],
679*800a58d9SAndroid Build Coastguard Worker                           "value": self.METADATA[1]}],
680*800a58d9SAndroid Build Coastguard Worker            },
681*800a58d9SAndroid Build Coastguard Worker            "labels":{'created_by': "fake_user"},
682*800a58d9SAndroid Build Coastguard Worker            "enableVtpm": True,
683*800a58d9SAndroid Build Coastguard Worker        }
684*800a58d9SAndroid Build Coastguard Worker
685*800a58d9SAndroid Build Coastguard Worker        self.compute_client.CreateInstance(
686*800a58d9SAndroid Build Coastguard Worker            instance=self.INSTANCE,
687*800a58d9SAndroid Build Coastguard Worker            image_name=self.IMAGE,
688*800a58d9SAndroid Build Coastguard Worker            machine_type=self.MACHINE_TYPE,
689*800a58d9SAndroid Build Coastguard Worker            metadata={self.METADATA[0]: self.METADATA[1]},
690*800a58d9SAndroid Build Coastguard Worker            network=self.NETWORK,
691*800a58d9SAndroid Build Coastguard Worker            zone=self.ZONE,
692*800a58d9SAndroid Build Coastguard Worker            extra_disk_name=extra_disk_name,
693*800a58d9SAndroid Build Coastguard Worker            tags=["https-server"],
694*800a58d9SAndroid Build Coastguard Worker            extra_scopes=self.EXTRA_SCOPES)
695*800a58d9SAndroid Build Coastguard Worker
696*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert.assert_called_with(
697*800a58d9SAndroid Build Coastguard Worker            project=PROJECT, zone=self.ZONE, body=expected_body)
698*800a58d9SAndroid Build Coastguard Worker        mock_wait.assert_called_with(
699*800a58d9SAndroid Build Coastguard Worker            mock.ANY,
700*800a58d9SAndroid Build Coastguard Worker            operation_scope=gcompute_client.OperationScope.ZONE,
701*800a58d9SAndroid Build Coastguard Worker            scope_name=self.ZONE)
702*800a58d9SAndroid Build Coastguard Worker
703*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "GetAcceleratorUrl")
704*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "GetImage")
705*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "GetNetworkUrl")
706*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "GetSubnetworkUrl")
707*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "GetMachineType")
708*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation")
709*800a58d9SAndroid Build Coastguard Worker    @mock.patch("getpass.getuser", return_value="fake_user")
710*800a58d9SAndroid Build Coastguard Worker    def testCreateInstanceWithGpu(self, _get_user, mock_wait, mock_get_mach,
711*800a58d9SAndroid Build Coastguard Worker                                  mock_get_subnetwork, mock_get_network,
712*800a58d9SAndroid Build Coastguard Worker                                  mock_get_image, mock_get_accel):
713*800a58d9SAndroid Build Coastguard Worker        """Test CreateInstance with a GPU parameter not set to None."""
714*800a58d9SAndroid Build Coastguard Worker        mock_get_mach.return_value = {"selfLink": self.MACHINE_TYPE_URL}
715*800a58d9SAndroid Build Coastguard Worker        mock_get_network.return_value = self.NETWORK_URL
716*800a58d9SAndroid Build Coastguard Worker        mock_get_subnetwork.return_value = self.SUBNETWORK_URL
717*800a58d9SAndroid Build Coastguard Worker        mock_get_accel.return_value = self.ACCELERATOR_URL
718*800a58d9SAndroid Build Coastguard Worker        mock_get_image.return_value = {"selfLink": self.IMAGE_URL}
719*800a58d9SAndroid Build Coastguard Worker
720*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
721*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.instances = mock.MagicMock(
722*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
723*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert = mock.MagicMock()
724*800a58d9SAndroid Build Coastguard Worker
725*800a58d9SAndroid Build Coastguard Worker        expected_body = {
726*800a58d9SAndroid Build Coastguard Worker            "machineType":
727*800a58d9SAndroid Build Coastguard Worker                self.MACHINE_TYPE_URL,
728*800a58d9SAndroid Build Coastguard Worker            "name":
729*800a58d9SAndroid Build Coastguard Worker                self.INSTANCE,
730*800a58d9SAndroid Build Coastguard Worker            "networkInterfaces": [{
731*800a58d9SAndroid Build Coastguard Worker                "network": self.NETWORK_URL,
732*800a58d9SAndroid Build Coastguard Worker                "subnetwork": self.SUBNETWORK_URL,
733*800a58d9SAndroid Build Coastguard Worker                "accessConfigs": [{
734*800a58d9SAndroid Build Coastguard Worker                    "name": "External NAT",
735*800a58d9SAndroid Build Coastguard Worker                    "type": "ONE_TO_ONE_NAT"
736*800a58d9SAndroid Build Coastguard Worker                }],
737*800a58d9SAndroid Build Coastguard Worker            }],
738*800a58d9SAndroid Build Coastguard Worker            "disks": [self._disk_args],
739*800a58d9SAndroid Build Coastguard Worker            "serviceAccounts": [{
740*800a58d9SAndroid Build Coastguard Worker                "email": "default",
741*800a58d9SAndroid Build Coastguard Worker                "scopes": self.compute_client.DEFAULT_INSTANCE_SCOPE
742*800a58d9SAndroid Build Coastguard Worker            }],
743*800a58d9SAndroid Build Coastguard Worker            "scheduling": {
744*800a58d9SAndroid Build Coastguard Worker                "onHostMaintenance": "terminate"
745*800a58d9SAndroid Build Coastguard Worker            },
746*800a58d9SAndroid Build Coastguard Worker            "guestAccelerators": [{
747*800a58d9SAndroid Build Coastguard Worker                "acceleratorCount": 1,
748*800a58d9SAndroid Build Coastguard Worker                "acceleratorType": "http://speedy-gpu"
749*800a58d9SAndroid Build Coastguard Worker            }],
750*800a58d9SAndroid Build Coastguard Worker            "metadata": {
751*800a58d9SAndroid Build Coastguard Worker                "items": [{
752*800a58d9SAndroid Build Coastguard Worker                    "key": self.METADATA[0],
753*800a58d9SAndroid Build Coastguard Worker                    "value": self.METADATA[1]
754*800a58d9SAndroid Build Coastguard Worker                }],
755*800a58d9SAndroid Build Coastguard Worker            },
756*800a58d9SAndroid Build Coastguard Worker            "labels":{'created_by': "fake_user"},
757*800a58d9SAndroid Build Coastguard Worker            "enableVtpm": True,
758*800a58d9SAndroid Build Coastguard Worker        }
759*800a58d9SAndroid Build Coastguard Worker
760*800a58d9SAndroid Build Coastguard Worker        self.compute_client.CreateInstance(
761*800a58d9SAndroid Build Coastguard Worker            instance=self.INSTANCE,
762*800a58d9SAndroid Build Coastguard Worker            image_name=self.IMAGE,
763*800a58d9SAndroid Build Coastguard Worker            machine_type=self.MACHINE_TYPE,
764*800a58d9SAndroid Build Coastguard Worker            metadata={self.METADATA[0]: self.METADATA[1]},
765*800a58d9SAndroid Build Coastguard Worker            network=self.NETWORK,
766*800a58d9SAndroid Build Coastguard Worker            zone=self.ZONE,
767*800a58d9SAndroid Build Coastguard Worker            gpu=self.GPU,
768*800a58d9SAndroid Build Coastguard Worker            extra_scopes=None)
769*800a58d9SAndroid Build Coastguard Worker
770*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert.assert_called_with(
771*800a58d9SAndroid Build Coastguard Worker            project=PROJECT, zone=self.ZONE, body=expected_body)
772*800a58d9SAndroid Build Coastguard Worker        mock_wait.assert_called_with(
773*800a58d9SAndroid Build Coastguard Worker            mock.ANY, operation_scope=gcompute_client.OperationScope.ZONE,
774*800a58d9SAndroid Build Coastguard Worker            scope_name=self.ZONE)
775*800a58d9SAndroid Build Coastguard Worker
776*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation")
777*800a58d9SAndroid Build Coastguard Worker    def testDeleteInstance(self, mock_wait):
778*800a58d9SAndroid Build Coastguard Worker        """Test DeleteInstance."""
779*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
780*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.instances = mock.MagicMock(
781*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
782*800a58d9SAndroid Build Coastguard Worker        resource_mock.delete = mock.MagicMock()
783*800a58d9SAndroid Build Coastguard Worker        self.compute_client.DeleteInstance(
784*800a58d9SAndroid Build Coastguard Worker            instance=self.INSTANCE, zone=self.ZONE)
785*800a58d9SAndroid Build Coastguard Worker        resource_mock.delete.assert_called_with(
786*800a58d9SAndroid Build Coastguard Worker            project=PROJECT, zone=self.ZONE, instance=self.INSTANCE)
787*800a58d9SAndroid Build Coastguard Worker        mock_wait.assert_called_with(
788*800a58d9SAndroid Build Coastguard Worker            mock.ANY,
789*800a58d9SAndroid Build Coastguard Worker            operation_scope=gcompute_client.OperationScope.ZONE,
790*800a58d9SAndroid Build Coastguard Worker            scope_name=self.ZONE)
791*800a58d9SAndroid Build Coastguard Worker
792*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation")
793*800a58d9SAndroid Build Coastguard Worker    def testDeleteInstances(self, mock_wait):
794*800a58d9SAndroid Build Coastguard Worker        """Test DeleteInstances."""
795*800a58d9SAndroid Build Coastguard Worker        self._SetupBatchHttpRequestMock()
796*800a58d9SAndroid Build Coastguard Worker        fake_instances = ["fake_instance_1", "fake_instance_2"]
797*800a58d9SAndroid Build Coastguard Worker        mock_api = mock.MagicMock()
798*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
799*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.instances = mock.MagicMock(
800*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
801*800a58d9SAndroid Build Coastguard Worker        resource_mock.delete = mock.MagicMock(return_value=mock_api)
802*800a58d9SAndroid Build Coastguard Worker        deleted, failed, error_msgs = self.compute_client.DeleteInstances(
803*800a58d9SAndroid Build Coastguard Worker            fake_instances, self.ZONE)
804*800a58d9SAndroid Build Coastguard Worker        calls = [
805*800a58d9SAndroid Build Coastguard Worker            mock.call(
806*800a58d9SAndroid Build Coastguard Worker                project=PROJECT,
807*800a58d9SAndroid Build Coastguard Worker                instance="fake_instance_1",
808*800a58d9SAndroid Build Coastguard Worker                zone=self.ZONE),
809*800a58d9SAndroid Build Coastguard Worker            mock.call(
810*800a58d9SAndroid Build Coastguard Worker                project=PROJECT,
811*800a58d9SAndroid Build Coastguard Worker                instance="fake_instance_2",
812*800a58d9SAndroid Build Coastguard Worker                zone=self.ZONE)
813*800a58d9SAndroid Build Coastguard Worker        ]
814*800a58d9SAndroid Build Coastguard Worker        resource_mock.delete.assert_has_calls(calls, any_order=True)
815*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(mock_wait.call_count, 2)
816*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(error_msgs, [])
817*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(failed, [])
818*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(set(deleted), set(fake_instances))
819*800a58d9SAndroid Build Coastguard Worker
820*800a58d9SAndroid Build Coastguard Worker    def testCreateDiskWithProject(self):
821*800a58d9SAndroid Build Coastguard Worker        """Test CreateDisk with images using a set project."""
822*800a58d9SAndroid Build Coastguard Worker        source_project = "fake-image-project"
823*800a58d9SAndroid Build Coastguard Worker        expected_project_to_use = "fake-image-project"
824*800a58d9SAndroid Build Coastguard Worker        mock_wait = self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
825*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
826*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.disks = mock.MagicMock(
827*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
828*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert = mock.MagicMock()
829*800a58d9SAndroid Build Coastguard Worker        self.compute_client.CreateDisk(
830*800a58d9SAndroid Build Coastguard Worker            "fake_disk", "fake_image", 10, self.ZONE, source_project=source_project)
831*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert.assert_called_with(
832*800a58d9SAndroid Build Coastguard Worker            project=PROJECT,
833*800a58d9SAndroid Build Coastguard Worker            zone=self.ZONE,
834*800a58d9SAndroid Build Coastguard Worker            sourceImage="projects/%s/global/images/fake_image" %
835*800a58d9SAndroid Build Coastguard Worker            expected_project_to_use,
836*800a58d9SAndroid Build Coastguard Worker            body={
837*800a58d9SAndroid Build Coastguard Worker                "name":
838*800a58d9SAndroid Build Coastguard Worker                    "fake_disk",
839*800a58d9SAndroid Build Coastguard Worker                "sizeGb":
840*800a58d9SAndroid Build Coastguard Worker                    10,
841*800a58d9SAndroid Build Coastguard Worker                "type":
842*800a58d9SAndroid Build Coastguard Worker                    "projects/%s/zones/%s/diskTypes/pd-standard" % (PROJECT,
843*800a58d9SAndroid Build Coastguard Worker                                                                    self.ZONE)
844*800a58d9SAndroid Build Coastguard Worker            })
845*800a58d9SAndroid Build Coastguard Worker        self.assertTrue(mock_wait.called)
846*800a58d9SAndroid Build Coastguard Worker
847*800a58d9SAndroid Build Coastguard Worker    def testCreateDiskWithNoSourceProject(self):
848*800a58d9SAndroid Build Coastguard Worker        """Test CreateDisk with images with no set project."""
849*800a58d9SAndroid Build Coastguard Worker        source_project = None
850*800a58d9SAndroid Build Coastguard Worker        expected_project_to_use = PROJECT
851*800a58d9SAndroid Build Coastguard Worker        mock_wait = self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
852*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
853*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.disks = mock.MagicMock(
854*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
855*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert = mock.MagicMock()
856*800a58d9SAndroid Build Coastguard Worker        self.compute_client.CreateDisk(
857*800a58d9SAndroid Build Coastguard Worker            "fake_disk", "fake_image", 10, self.ZONE, source_project=source_project)
858*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert.assert_called_with(
859*800a58d9SAndroid Build Coastguard Worker            project=PROJECT,
860*800a58d9SAndroid Build Coastguard Worker            zone=self.ZONE,
861*800a58d9SAndroid Build Coastguard Worker            sourceImage="projects/%s/global/images/fake_image" %
862*800a58d9SAndroid Build Coastguard Worker            expected_project_to_use,
863*800a58d9SAndroid Build Coastguard Worker            body={
864*800a58d9SAndroid Build Coastguard Worker                "name":
865*800a58d9SAndroid Build Coastguard Worker                    "fake_disk",
866*800a58d9SAndroid Build Coastguard Worker                "sizeGb":
867*800a58d9SAndroid Build Coastguard Worker                    10,
868*800a58d9SAndroid Build Coastguard Worker                "type":
869*800a58d9SAndroid Build Coastguard Worker                    "projects/%s/zones/%s/diskTypes/pd-standard" % (PROJECT,
870*800a58d9SAndroid Build Coastguard Worker                                                                    self.ZONE)
871*800a58d9SAndroid Build Coastguard Worker            })
872*800a58d9SAndroid Build Coastguard Worker        self.assertTrue(mock_wait.called)
873*800a58d9SAndroid Build Coastguard Worker
874*800a58d9SAndroid Build Coastguard Worker    def testCreateDiskWithTypeStandard(self):
875*800a58d9SAndroid Build Coastguard Worker        """Test CreateDisk with images using standard."""
876*800a58d9SAndroid Build Coastguard Worker        disk_type = gcompute_client.PersistentDiskType.STANDARD
877*800a58d9SAndroid Build Coastguard Worker        expected_disk_type_string = "pd-standard"
878*800a58d9SAndroid Build Coastguard Worker        mock_wait = self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
879*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
880*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.disks = mock.MagicMock(
881*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
882*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert = mock.MagicMock()
883*800a58d9SAndroid Build Coastguard Worker        self.compute_client.CreateDisk(
884*800a58d9SAndroid Build Coastguard Worker            "fake_disk",
885*800a58d9SAndroid Build Coastguard Worker            "fake_image",
886*800a58d9SAndroid Build Coastguard Worker            10,
887*800a58d9SAndroid Build Coastguard Worker            self.ZONE,
888*800a58d9SAndroid Build Coastguard Worker            source_project="fake-project",
889*800a58d9SAndroid Build Coastguard Worker            disk_type=disk_type)
890*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert.assert_called_with(
891*800a58d9SAndroid Build Coastguard Worker            project=PROJECT,
892*800a58d9SAndroid Build Coastguard Worker            zone=self.ZONE,
893*800a58d9SAndroid Build Coastguard Worker            sourceImage="projects/%s/global/images/fake_image" % "fake-project",
894*800a58d9SAndroid Build Coastguard Worker            body={
895*800a58d9SAndroid Build Coastguard Worker                "name":
896*800a58d9SAndroid Build Coastguard Worker                    "fake_disk",
897*800a58d9SAndroid Build Coastguard Worker                "sizeGb":
898*800a58d9SAndroid Build Coastguard Worker                    10,
899*800a58d9SAndroid Build Coastguard Worker                "type":
900*800a58d9SAndroid Build Coastguard Worker                    "projects/%s/zones/%s/diskTypes/%s" %
901*800a58d9SAndroid Build Coastguard Worker                    (PROJECT, self.ZONE, expected_disk_type_string)
902*800a58d9SAndroid Build Coastguard Worker            })
903*800a58d9SAndroid Build Coastguard Worker        self.assertTrue(mock_wait.called)
904*800a58d9SAndroid Build Coastguard Worker
905*800a58d9SAndroid Build Coastguard Worker    def testCreateDiskWithTypeSSD(self):
906*800a58d9SAndroid Build Coastguard Worker        """Test CreateDisk with images using standard."""
907*800a58d9SAndroid Build Coastguard Worker        disk_type = gcompute_client.PersistentDiskType.SSD
908*800a58d9SAndroid Build Coastguard Worker        expected_disk_type_string = "pd-ssd"
909*800a58d9SAndroid Build Coastguard Worker        mock_wait = self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
910*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
911*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.disks = mock.MagicMock(
912*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
913*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert = mock.MagicMock()
914*800a58d9SAndroid Build Coastguard Worker        self.compute_client.CreateDisk(
915*800a58d9SAndroid Build Coastguard Worker            "fake_disk",
916*800a58d9SAndroid Build Coastguard Worker            "fake_image",
917*800a58d9SAndroid Build Coastguard Worker            10,
918*800a58d9SAndroid Build Coastguard Worker            self.ZONE,
919*800a58d9SAndroid Build Coastguard Worker            source_project="fake-project",
920*800a58d9SAndroid Build Coastguard Worker            disk_type=disk_type)
921*800a58d9SAndroid Build Coastguard Worker        resource_mock.insert.assert_called_with(
922*800a58d9SAndroid Build Coastguard Worker            project=PROJECT,
923*800a58d9SAndroid Build Coastguard Worker            zone=self.ZONE,
924*800a58d9SAndroid Build Coastguard Worker            sourceImage="projects/%s/global/images/fake_image" % "fake-project",
925*800a58d9SAndroid Build Coastguard Worker            body={
926*800a58d9SAndroid Build Coastguard Worker                "name":
927*800a58d9SAndroid Build Coastguard Worker                    "fake_disk",
928*800a58d9SAndroid Build Coastguard Worker                "sizeGb":
929*800a58d9SAndroid Build Coastguard Worker                    10,
930*800a58d9SAndroid Build Coastguard Worker                "type":
931*800a58d9SAndroid Build Coastguard Worker                    "projects/%s/zones/%s/diskTypes/%s" %
932*800a58d9SAndroid Build Coastguard Worker                    (PROJECT, self.ZONE, expected_disk_type_string)
933*800a58d9SAndroid Build Coastguard Worker            })
934*800a58d9SAndroid Build Coastguard Worker        self.assertTrue(mock_wait.called)
935*800a58d9SAndroid Build Coastguard Worker
936*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation")
937*800a58d9SAndroid Build Coastguard Worker    def testAttachDisk(self, mock_wait):
938*800a58d9SAndroid Build Coastguard Worker        """Test AttachDisk."""
939*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
940*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.instances = mock.MagicMock(
941*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
942*800a58d9SAndroid Build Coastguard Worker        resource_mock.attachDisk = mock.MagicMock()
943*800a58d9SAndroid Build Coastguard Worker        self.compute_client.AttachDisk(
944*800a58d9SAndroid Build Coastguard Worker            "fake_instance_1", self.ZONE, deviceName="fake_disk",
945*800a58d9SAndroid Build Coastguard Worker            source="fake-selfLink")
946*800a58d9SAndroid Build Coastguard Worker        resource_mock.attachDisk.assert_called_with(
947*800a58d9SAndroid Build Coastguard Worker            project=PROJECT,
948*800a58d9SAndroid Build Coastguard Worker            zone=self.ZONE,
949*800a58d9SAndroid Build Coastguard Worker            instance="fake_instance_1",
950*800a58d9SAndroid Build Coastguard Worker            body={
951*800a58d9SAndroid Build Coastguard Worker                "deviceName": "fake_disk",
952*800a58d9SAndroid Build Coastguard Worker                "source": "fake-selfLink"
953*800a58d9SAndroid Build Coastguard Worker            })
954*800a58d9SAndroid Build Coastguard Worker        self.assertTrue(mock_wait.called)
955*800a58d9SAndroid Build Coastguard Worker
956*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation")
957*800a58d9SAndroid Build Coastguard Worker    def testDetachDisk(self, mock_wait):
958*800a58d9SAndroid Build Coastguard Worker        """Test DetachDisk."""
959*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
960*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.instances = mock.MagicMock(
961*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
962*800a58d9SAndroid Build Coastguard Worker        resource_mock.detachDisk = mock.MagicMock()
963*800a58d9SAndroid Build Coastguard Worker        self.compute_client.DetachDisk("fake_instance_1", self.ZONE, "fake_disk")
964*800a58d9SAndroid Build Coastguard Worker        resource_mock.detachDisk.assert_called_with(
965*800a58d9SAndroid Build Coastguard Worker            project=PROJECT,
966*800a58d9SAndroid Build Coastguard Worker            zone=self.ZONE,
967*800a58d9SAndroid Build Coastguard Worker            instance="fake_instance_1",
968*800a58d9SAndroid Build Coastguard Worker            deviceName="fake_disk")
969*800a58d9SAndroid Build Coastguard Worker        self.assertTrue(mock_wait.called)
970*800a58d9SAndroid Build Coastguard Worker
971*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "GetAcceleratorUrl")
972*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation")
973*800a58d9SAndroid Build Coastguard Worker    def testAttachAccelerator(self, mock_wait, mock_get_accel):
974*800a58d9SAndroid Build Coastguard Worker        """Test AttachAccelerator."""
975*800a58d9SAndroid Build Coastguard Worker        mock_get_accel.return_value = self.ACCELERATOR_URL
976*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
977*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.instances = mock.MagicMock(
978*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
979*800a58d9SAndroid Build Coastguard Worker        resource_mock.attachAccelerator = mock.MagicMock()
980*800a58d9SAndroid Build Coastguard Worker        self.compute_client.AttachAccelerator("fake_instance_1", self.ZONE, 1,
981*800a58d9SAndroid Build Coastguard Worker                                              "nvidia-tesla-k80")
982*800a58d9SAndroid Build Coastguard Worker        resource_mock.setMachineResources.assert_called_with(
983*800a58d9SAndroid Build Coastguard Worker            project=PROJECT,
984*800a58d9SAndroid Build Coastguard Worker            zone=self.ZONE,
985*800a58d9SAndroid Build Coastguard Worker            instance="fake_instance_1",
986*800a58d9SAndroid Build Coastguard Worker            body={
987*800a58d9SAndroid Build Coastguard Worker                "guestAccelerators": [{
988*800a58d9SAndroid Build Coastguard Worker                    "acceleratorType": self.ACCELERATOR_URL,
989*800a58d9SAndroid Build Coastguard Worker                    "acceleratorCount": 1
990*800a58d9SAndroid Build Coastguard Worker                }]
991*800a58d9SAndroid Build Coastguard Worker            })
992*800a58d9SAndroid Build Coastguard Worker        self.assertTrue(mock_wait.called)
993*800a58d9SAndroid Build Coastguard Worker
994*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation")
995*800a58d9SAndroid Build Coastguard Worker    def testBatchExecuteOnInstances(self, mock_wait):
996*800a58d9SAndroid Build Coastguard Worker        """Test BatchExecuteOnInstances."""
997*800a58d9SAndroid Build Coastguard Worker        self._SetupBatchHttpRequestMock()
998*800a58d9SAndroid Build Coastguard Worker        action = mock.MagicMock(return_value=mock.MagicMock())
999*800a58d9SAndroid Build Coastguard Worker        fake_instances = ["fake_instance_1", "fake_instance_2"]
1000*800a58d9SAndroid Build Coastguard Worker        done, failed, error_msgs = self.compute_client._BatchExecuteOnInstances(
1001*800a58d9SAndroid Build Coastguard Worker            fake_instances, self.ZONE, action)
1002*800a58d9SAndroid Build Coastguard Worker        calls = [mock.call(instance="fake_instance_1"),
1003*800a58d9SAndroid Build Coastguard Worker                 mock.call(instance="fake_instance_2")]
1004*800a58d9SAndroid Build Coastguard Worker        action.assert_has_calls(calls, any_order=True)
1005*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(mock_wait.call_count, 2)
1006*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(set(done), set(fake_instances))
1007*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(error_msgs, [])
1008*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(failed, [])
1009*800a58d9SAndroid Build Coastguard Worker
1010*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation")
1011*800a58d9SAndroid Build Coastguard Worker    def testResetInstance(self, mock_wait):
1012*800a58d9SAndroid Build Coastguard Worker        """Test ResetInstance."""
1013*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
1014*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.instances = mock.MagicMock(
1015*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
1016*800a58d9SAndroid Build Coastguard Worker        resource_mock.reset = mock.MagicMock()
1017*800a58d9SAndroid Build Coastguard Worker        self.compute_client.ResetInstance(
1018*800a58d9SAndroid Build Coastguard Worker            instance=self.INSTANCE, zone=self.ZONE)
1019*800a58d9SAndroid Build Coastguard Worker        resource_mock.reset.assert_called_with(
1020*800a58d9SAndroid Build Coastguard Worker            project=PROJECT, zone=self.ZONE, instance=self.INSTANCE)
1021*800a58d9SAndroid Build Coastguard Worker        mock_wait.assert_called_with(
1022*800a58d9SAndroid Build Coastguard Worker            mock.ANY,
1023*800a58d9SAndroid Build Coastguard Worker            operation_scope=gcompute_client.OperationScope.ZONE,
1024*800a58d9SAndroid Build Coastguard Worker            scope_name=self.ZONE)
1025*800a58d9SAndroid Build Coastguard Worker
1026*800a58d9SAndroid Build Coastguard Worker    def _CompareMachineSizeTestHelper(self,
1027*800a58d9SAndroid Build Coastguard Worker                                      machine_info_1,
1028*800a58d9SAndroid Build Coastguard Worker                                      machine_info_2,
1029*800a58d9SAndroid Build Coastguard Worker                                      expected_result=None,
1030*800a58d9SAndroid Build Coastguard Worker                                      expected_error_type=None):
1031*800a58d9SAndroid Build Coastguard Worker        """Helper class for testing CompareMachineSize.
1032*800a58d9SAndroid Build Coastguard Worker
1033*800a58d9SAndroid Build Coastguard Worker        Args:
1034*800a58d9SAndroid Build Coastguard Worker            machine_info_1: A dictionary representing the first machine size.
1035*800a58d9SAndroid Build Coastguard Worker            machine_info_2: A dictionary representing the second machine size.
1036*800a58d9SAndroid Build Coastguard Worker            expected_result: An integer, 0, 1 or -1, or None if not set.
1037*800a58d9SAndroid Build Coastguard Worker            expected_error_type: An exception type, if set will check for exception.
1038*800a58d9SAndroid Build Coastguard Worker        """
1039*800a58d9SAndroid Build Coastguard Worker        mock_get_mach_type = self.Patch(
1040*800a58d9SAndroid Build Coastguard Worker            gcompute_client.ComputeClient,
1041*800a58d9SAndroid Build Coastguard Worker            "GetMachineType",
1042*800a58d9SAndroid Build Coastguard Worker            side_effect=[machine_info_1, machine_info_2])
1043*800a58d9SAndroid Build Coastguard Worker        if expected_error_type:
1044*800a58d9SAndroid Build Coastguard Worker            self.assertRaises(expected_error_type,
1045*800a58d9SAndroid Build Coastguard Worker                              self.compute_client.CompareMachineSize, "name1",
1046*800a58d9SAndroid Build Coastguard Worker                              "name2", self.ZONE)
1047*800a58d9SAndroid Build Coastguard Worker        else:
1048*800a58d9SAndroid Build Coastguard Worker            result = self.compute_client.CompareMachineSize("name1", "name2",
1049*800a58d9SAndroid Build Coastguard Worker                                                            self.ZONE)
1050*800a58d9SAndroid Build Coastguard Worker            self.assertEqual(result, expected_result)
1051*800a58d9SAndroid Build Coastguard Worker
1052*800a58d9SAndroid Build Coastguard Worker        mock_get_mach_type.assert_has_calls(
1053*800a58d9SAndroid Build Coastguard Worker            [mock.call("name1", self.ZONE), mock.call("name2", self.ZONE)])
1054*800a58d9SAndroid Build Coastguard Worker
1055*800a58d9SAndroid Build Coastguard Worker    def testCompareMachineSizeSmall(self):
1056*800a58d9SAndroid Build Coastguard Worker        """Test CompareMachineSize where the first one is smaller."""
1057*800a58d9SAndroid Build Coastguard Worker        machine_info_1 = {"guestCpus": 10, "memoryMb": 100}
1058*800a58d9SAndroid Build Coastguard Worker        machine_info_2 = {"guestCpus": 10, "memoryMb": 200}
1059*800a58d9SAndroid Build Coastguard Worker        self._CompareMachineSizeTestHelper(machine_info_1, machine_info_2, -1)
1060*800a58d9SAndroid Build Coastguard Worker
1061*800a58d9SAndroid Build Coastguard Worker    def testCompareMachineSizeSmallSmallerOnSecond(self):
1062*800a58d9SAndroid Build Coastguard Worker        """Test CompareMachineSize where the first one is smaller."""
1063*800a58d9SAndroid Build Coastguard Worker        machine_info_1 = {"guestCpus": 11, "memoryMb": 100}
1064*800a58d9SAndroid Build Coastguard Worker        machine_info_2 = {"guestCpus": 10, "memoryMb": 200}
1065*800a58d9SAndroid Build Coastguard Worker        self._CompareMachineSizeTestHelper(machine_info_1, machine_info_2, -1)
1066*800a58d9SAndroid Build Coastguard Worker
1067*800a58d9SAndroid Build Coastguard Worker    def testCompareMachineSizeLarge(self):
1068*800a58d9SAndroid Build Coastguard Worker        """Test CompareMachineSize where the first one is larger."""
1069*800a58d9SAndroid Build Coastguard Worker        machine_info_1 = {"guestCpus": 11, "memoryMb": 200}
1070*800a58d9SAndroid Build Coastguard Worker        machine_info_2 = {"guestCpus": 10, "memoryMb": 100}
1071*800a58d9SAndroid Build Coastguard Worker        self._CompareMachineSizeTestHelper(machine_info_1, machine_info_2, 1)
1072*800a58d9SAndroid Build Coastguard Worker
1073*800a58d9SAndroid Build Coastguard Worker    def testCompareMachineSizeLargeWithEqualElement(self):
1074*800a58d9SAndroid Build Coastguard Worker        """Test CompareMachineSize where the first one is larger."""
1075*800a58d9SAndroid Build Coastguard Worker        machine_info_1 = {"guestCpus": 10, "memoryMb": 200}
1076*800a58d9SAndroid Build Coastguard Worker        machine_info_2 = {"guestCpus": 10, "memoryMb": 100}
1077*800a58d9SAndroid Build Coastguard Worker        self._CompareMachineSizeTestHelper(machine_info_1, machine_info_2, 1)
1078*800a58d9SAndroid Build Coastguard Worker
1079*800a58d9SAndroid Build Coastguard Worker    def testCompareMachineSizeEqual(self):
1080*800a58d9SAndroid Build Coastguard Worker        """Test CompareMachineSize where two machine sizes are equal."""
1081*800a58d9SAndroid Build Coastguard Worker        machine_info = {"guestCpus": 10, "memoryMb": 100}
1082*800a58d9SAndroid Build Coastguard Worker        self._CompareMachineSizeTestHelper(machine_info, machine_info, 0)
1083*800a58d9SAndroid Build Coastguard Worker
1084*800a58d9SAndroid Build Coastguard Worker    def testCompareMachineSizeBadMetric(self):
1085*800a58d9SAndroid Build Coastguard Worker        """Test CompareMachineSize with bad metric."""
1086*800a58d9SAndroid Build Coastguard Worker        machine_info = {"unknown_metric": 10, "memoryMb": 100}
1087*800a58d9SAndroid Build Coastguard Worker        self._CompareMachineSizeTestHelper(
1088*800a58d9SAndroid Build Coastguard Worker            machine_info, machine_info, expected_error_type=errors.DriverError)
1089*800a58d9SAndroid Build Coastguard Worker
1090*800a58d9SAndroid Build Coastguard Worker    def testGetMachineType(self):
1091*800a58d9SAndroid Build Coastguard Worker        """Test GetMachineType."""
1092*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
1093*800a58d9SAndroid Build Coastguard Worker        mock_api = mock.MagicMock()
1094*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.machineTypes = mock.MagicMock(
1095*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
1096*800a58d9SAndroid Build Coastguard Worker        resource_mock.get = mock.MagicMock(return_value=mock_api)
1097*800a58d9SAndroid Build Coastguard Worker        mock_api.execute = mock.MagicMock(
1098*800a58d9SAndroid Build Coastguard Worker            return_value={"name": self.MACHINE_TYPE})
1099*800a58d9SAndroid Build Coastguard Worker        result = self.compute_client.GetMachineType(self.MACHINE_TYPE,
1100*800a58d9SAndroid Build Coastguard Worker                                                    self.ZONE)
1101*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(result, {"name": self.MACHINE_TYPE})
1102*800a58d9SAndroid Build Coastguard Worker        resource_mock.get.assert_called_with(
1103*800a58d9SAndroid Build Coastguard Worker            project=PROJECT,
1104*800a58d9SAndroid Build Coastguard Worker            zone=self.ZONE,
1105*800a58d9SAndroid Build Coastguard Worker            machineType=self.MACHINE_TYPE)
1106*800a58d9SAndroid Build Coastguard Worker
1107*800a58d9SAndroid Build Coastguard Worker    def _GetSerialPortOutputTestHelper(self, response):
1108*800a58d9SAndroid Build Coastguard Worker        """Helper function for testing GetSerialPortOutput.
1109*800a58d9SAndroid Build Coastguard Worker
1110*800a58d9SAndroid Build Coastguard Worker        Args:
1111*800a58d9SAndroid Build Coastguard Worker            response: A dictionary representing a fake response.
1112*800a58d9SAndroid Build Coastguard Worker        """
1113*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
1114*800a58d9SAndroid Build Coastguard Worker        mock_api = mock.MagicMock()
1115*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.instances = mock.MagicMock(
1116*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
1117*800a58d9SAndroid Build Coastguard Worker        resource_mock.getSerialPortOutput = mock.MagicMock(
1118*800a58d9SAndroid Build Coastguard Worker            return_value=mock_api)
1119*800a58d9SAndroid Build Coastguard Worker        mock_api.execute = mock.MagicMock(return_value=response)
1120*800a58d9SAndroid Build Coastguard Worker
1121*800a58d9SAndroid Build Coastguard Worker        if "contents" in response:
1122*800a58d9SAndroid Build Coastguard Worker            result = self.compute_client.GetSerialPortOutput(
1123*800a58d9SAndroid Build Coastguard Worker                instance=self.INSTANCE, zone=self.ZONE)
1124*800a58d9SAndroid Build Coastguard Worker            self.assertEqual(result, "fake contents")
1125*800a58d9SAndroid Build Coastguard Worker        else:
1126*800a58d9SAndroid Build Coastguard Worker            self.assertRaisesRegex(
1127*800a58d9SAndroid Build Coastguard Worker                errors.DriverError,
1128*800a58d9SAndroid Build Coastguard Worker                "Malformed response.*",
1129*800a58d9SAndroid Build Coastguard Worker                self.compute_client.GetSerialPortOutput,
1130*800a58d9SAndroid Build Coastguard Worker                instance=self.INSTANCE,
1131*800a58d9SAndroid Build Coastguard Worker                zone=self.ZONE)
1132*800a58d9SAndroid Build Coastguard Worker        resource_mock.getSerialPortOutput.assert_called_with(
1133*800a58d9SAndroid Build Coastguard Worker            project=PROJECT,
1134*800a58d9SAndroid Build Coastguard Worker            zone=self.ZONE,
1135*800a58d9SAndroid Build Coastguard Worker            instance=self.INSTANCE,
1136*800a58d9SAndroid Build Coastguard Worker            port=1)
1137*800a58d9SAndroid Build Coastguard Worker
1138*800a58d9SAndroid Build Coastguard Worker    def testGetSerialPortOutput(self):
1139*800a58d9SAndroid Build Coastguard Worker        """Test GetSerialPortOutput."""
1140*800a58d9SAndroid Build Coastguard Worker        response = {"contents": "fake contents"}
1141*800a58d9SAndroid Build Coastguard Worker        self._GetSerialPortOutputTestHelper(response)
1142*800a58d9SAndroid Build Coastguard Worker
1143*800a58d9SAndroid Build Coastguard Worker    def testGetSerialPortOutputFail(self):
1144*800a58d9SAndroid Build Coastguard Worker        """Test GetSerialPortOutputFail."""
1145*800a58d9SAndroid Build Coastguard Worker        response = {"malformed": "fake contents"}
1146*800a58d9SAndroid Build Coastguard Worker        self._GetSerialPortOutputTestHelper(response)
1147*800a58d9SAndroid Build Coastguard Worker
1148*800a58d9SAndroid Build Coastguard Worker    def testGetInstanceNamesByIPs(self):
1149*800a58d9SAndroid Build Coastguard Worker        """Test GetInstanceNamesByIPs."""
1150*800a58d9SAndroid Build Coastguard Worker        good_instance = {
1151*800a58d9SAndroid Build Coastguard Worker            "name": "instance_1",
1152*800a58d9SAndroid Build Coastguard Worker            "networkInterfaces": [
1153*800a58d9SAndroid Build Coastguard Worker                {
1154*800a58d9SAndroid Build Coastguard Worker                    "accessConfigs": [
1155*800a58d9SAndroid Build Coastguard Worker                        {"natIP": "172.22.22.22"},
1156*800a58d9SAndroid Build Coastguard Worker                    ],
1157*800a58d9SAndroid Build Coastguard Worker                },
1158*800a58d9SAndroid Build Coastguard Worker            ],
1159*800a58d9SAndroid Build Coastguard Worker        }
1160*800a58d9SAndroid Build Coastguard Worker        bad_instance = {"name": "instance_2"}
1161*800a58d9SAndroid Build Coastguard Worker        self.Patch(
1162*800a58d9SAndroid Build Coastguard Worker            gcompute_client.ComputeClient,
1163*800a58d9SAndroid Build Coastguard Worker            "ListInstances",
1164*800a58d9SAndroid Build Coastguard Worker            return_value=[good_instance, bad_instance])
1165*800a58d9SAndroid Build Coastguard Worker        ip_name_map = self.compute_client.GetInstanceNamesByIPs(
1166*800a58d9SAndroid Build Coastguard Worker            ips=["172.22.22.22", "172.22.22.23"])
1167*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(ip_name_map, {"172.22.22.22": "instance_1",
1168*800a58d9SAndroid Build Coastguard Worker                                       "172.22.22.23": None})
1169*800a58d9SAndroid Build Coastguard Worker
1170*800a58d9SAndroid Build Coastguard Worker    def testRsaNotInMetadata(self):
1171*800a58d9SAndroid Build Coastguard Worker        """Test rsa not in metadata."""
1172*800a58d9SAndroid Build Coastguard Worker        fake_user = "fake_user"
1173*800a58d9SAndroid Build Coastguard Worker        fake_ssh_key = "fake_ssh"
1174*800a58d9SAndroid Build Coastguard Worker        metadata = {
1175*800a58d9SAndroid Build Coastguard Worker            "kind": "compute#metadata",
1176*800a58d9SAndroid Build Coastguard Worker            "fingerprint": "a-23icsyx4E=",
1177*800a58d9SAndroid Build Coastguard Worker            "items": [
1178*800a58d9SAndroid Build Coastguard Worker                {
1179*800a58d9SAndroid Build Coastguard Worker                    "key": "sshKeys",
1180*800a58d9SAndroid Build Coastguard Worker                    "value": "%s:%s" % (fake_user, self.SSHKEY)
1181*800a58d9SAndroid Build Coastguard Worker                }
1182*800a58d9SAndroid Build Coastguard Worker            ]
1183*800a58d9SAndroid Build Coastguard Worker        }
1184*800a58d9SAndroid Build Coastguard Worker        # Test rsa doesn't exist in metadata.
1185*800a58d9SAndroid Build Coastguard Worker        new_entry = "%s:%s" % (fake_user, fake_ssh_key)
1186*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(True, gcompute_client.RsaNotInMetadata(metadata, new_entry))
1187*800a58d9SAndroid Build Coastguard Worker
1188*800a58d9SAndroid Build Coastguard Worker        # Test rsa exists in metadata.
1189*800a58d9SAndroid Build Coastguard Worker        exist_entry = "%s:%s" %(fake_user, self.SSHKEY)
1190*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(False, gcompute_client.RsaNotInMetadata(metadata, exist_entry))
1191*800a58d9SAndroid Build Coastguard Worker
1192*800a58d9SAndroid Build Coastguard Worker    def testGetSshKeyFromMetadata(self):
1193*800a58d9SAndroid Build Coastguard Worker        """Test get ssh key from metadata."""
1194*800a58d9SAndroid Build Coastguard Worker        fake_user = "fake_user"
1195*800a58d9SAndroid Build Coastguard Worker        metadata_key_exist_value_is_empty = {
1196*800a58d9SAndroid Build Coastguard Worker            "kind": "compute#metadata",
1197*800a58d9SAndroid Build Coastguard Worker            "fingerprint": "a-23icsyx4E=",
1198*800a58d9SAndroid Build Coastguard Worker            "items": [
1199*800a58d9SAndroid Build Coastguard Worker                {
1200*800a58d9SAndroid Build Coastguard Worker                    "key": "sshKeys",
1201*800a58d9SAndroid Build Coastguard Worker                    "value": ""
1202*800a58d9SAndroid Build Coastguard Worker                }
1203*800a58d9SAndroid Build Coastguard Worker            ]
1204*800a58d9SAndroid Build Coastguard Worker        }
1205*800a58d9SAndroid Build Coastguard Worker        metadata_key_exist = {
1206*800a58d9SAndroid Build Coastguard Worker            "kind": "compute#metadata",
1207*800a58d9SAndroid Build Coastguard Worker            "fingerprint": "a-23icsyx4E=",
1208*800a58d9SAndroid Build Coastguard Worker            "items": [
1209*800a58d9SAndroid Build Coastguard Worker                {
1210*800a58d9SAndroid Build Coastguard Worker                    "key": "sshKeys",
1211*800a58d9SAndroid Build Coastguard Worker                    "value": "%s:%s" % (fake_user, self.SSHKEY)
1212*800a58d9SAndroid Build Coastguard Worker                }
1213*800a58d9SAndroid Build Coastguard Worker            ]
1214*800a58d9SAndroid Build Coastguard Worker        }
1215*800a58d9SAndroid Build Coastguard Worker        metadata_key_not_exist = {
1216*800a58d9SAndroid Build Coastguard Worker            "kind": "compute#metadata",
1217*800a58d9SAndroid Build Coastguard Worker            "fingerprint": "a-23icsyx4E=",
1218*800a58d9SAndroid Build Coastguard Worker            "items": [
1219*800a58d9SAndroid Build Coastguard Worker                {
1220*800a58d9SAndroid Build Coastguard Worker                }
1221*800a58d9SAndroid Build Coastguard Worker            ]
1222*800a58d9SAndroid Build Coastguard Worker        }
1223*800a58d9SAndroid Build Coastguard Worker        expected_key_exist_value_is_empty = {
1224*800a58d9SAndroid Build Coastguard Worker            "key": "sshKeys",
1225*800a58d9SAndroid Build Coastguard Worker            "value": ""
1226*800a58d9SAndroid Build Coastguard Worker        }
1227*800a58d9SAndroid Build Coastguard Worker        expected_key_exist = {
1228*800a58d9SAndroid Build Coastguard Worker            "key": "sshKeys",
1229*800a58d9SAndroid Build Coastguard Worker            "value": "%s:%s" % (fake_user, self.SSHKEY)
1230*800a58d9SAndroid Build Coastguard Worker        }
1231*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(expected_key_exist_value_is_empty,
1232*800a58d9SAndroid Build Coastguard Worker                         gcompute_client.GetSshKeyFromMetadata(metadata_key_exist_value_is_empty))
1233*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(expected_key_exist,
1234*800a58d9SAndroid Build Coastguard Worker                         gcompute_client.GetSshKeyFromMetadata(metadata_key_exist))
1235*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(None,
1236*800a58d9SAndroid Build Coastguard Worker                         gcompute_client.GetSshKeyFromMetadata(metadata_key_not_exist))
1237*800a58d9SAndroid Build Coastguard Worker
1238*800a58d9SAndroid Build Coastguard Worker
1239*800a58d9SAndroid Build Coastguard Worker    def testGetRsaKeyPathExistsFalse(self):
1240*800a58d9SAndroid Build Coastguard Worker        """Test the rsa key path not exists."""
1241*800a58d9SAndroid Build Coastguard Worker        fake_ssh_rsa_path = "/path/to/test_rsa.pub"
1242*800a58d9SAndroid Build Coastguard Worker        self.Patch(os.path, "exists", return_value=False)
1243*800a58d9SAndroid Build Coastguard Worker        self.assertRaisesRegex(errors.DriverError,
1244*800a58d9SAndroid Build Coastguard Worker                               "RSA file %s does not exist." % fake_ssh_rsa_path,
1245*800a58d9SAndroid Build Coastguard Worker                               gcompute_client.GetRsaKey,
1246*800a58d9SAndroid Build Coastguard Worker                               ssh_rsa_path=fake_ssh_rsa_path)
1247*800a58d9SAndroid Build Coastguard Worker
1248*800a58d9SAndroid Build Coastguard Worker    def testGetRsaKey(self):
1249*800a58d9SAndroid Build Coastguard Worker        """Test get the rsa key."""
1250*800a58d9SAndroid Build Coastguard Worker        fake_ssh_rsa_path = "/path/to/test_rsa.pub"
1251*800a58d9SAndroid Build Coastguard Worker        self.Patch(os.path, "exists", return_value=True)
1252*800a58d9SAndroid Build Coastguard Worker        m = mock.mock_open(read_data=self.SSHKEY)
1253*800a58d9SAndroid Build Coastguard Worker        with mock.patch("builtins.open", m):
1254*800a58d9SAndroid Build Coastguard Worker            result = gcompute_client.GetRsaKey(fake_ssh_rsa_path)
1255*800a58d9SAndroid Build Coastguard Worker            self.assertEqual(self.SSHKEY, result)
1256*800a58d9SAndroid Build Coastguard Worker
1257*800a58d9SAndroid Build Coastguard Worker    def testUpdateRsaInMetadata(self):
1258*800a58d9SAndroid Build Coastguard Worker        """Test update rsa in metadata."""
1259*800a58d9SAndroid Build Coastguard Worker        fake_ssh_key = "fake_ssh"
1260*800a58d9SAndroid Build Coastguard Worker        fake_metadata_sshkeys_not_exist = {
1261*800a58d9SAndroid Build Coastguard Worker            "kind": "compute#metadata",
1262*800a58d9SAndroid Build Coastguard Worker            "fingerprint": "a-23icsyx4E=",
1263*800a58d9SAndroid Build Coastguard Worker            "items": [
1264*800a58d9SAndroid Build Coastguard Worker                {
1265*800a58d9SAndroid Build Coastguard Worker                    "key": "not_sshKeys",
1266*800a58d9SAndroid Build Coastguard Worker                    "value": ""
1267*800a58d9SAndroid Build Coastguard Worker                }
1268*800a58d9SAndroid Build Coastguard Worker            ]
1269*800a58d9SAndroid Build Coastguard Worker        }
1270*800a58d9SAndroid Build Coastguard Worker        new_entry = "new_user:%s" % fake_ssh_key
1271*800a58d9SAndroid Build Coastguard Worker        expected = {
1272*800a58d9SAndroid Build Coastguard Worker            "kind": "compute#metadata",
1273*800a58d9SAndroid Build Coastguard Worker            "fingerprint": "a-23icsyx4E=",
1274*800a58d9SAndroid Build Coastguard Worker            "items": [
1275*800a58d9SAndroid Build Coastguard Worker                {
1276*800a58d9SAndroid Build Coastguard Worker                    "key": "not_sshKeys",
1277*800a58d9SAndroid Build Coastguard Worker                    "value": ""
1278*800a58d9SAndroid Build Coastguard Worker                },
1279*800a58d9SAndroid Build Coastguard Worker                {
1280*800a58d9SAndroid Build Coastguard Worker                    "key": "sshKeys",
1281*800a58d9SAndroid Build Coastguard Worker                    "value": new_entry
1282*800a58d9SAndroid Build Coastguard Worker                }
1283*800a58d9SAndroid Build Coastguard Worker            ]
1284*800a58d9SAndroid Build Coastguard Worker        }
1285*800a58d9SAndroid Build Coastguard Worker        self.Patch(os.path, "exists", return_value=True)
1286*800a58d9SAndroid Build Coastguard Worker        self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
1287*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
1288*800a58d9SAndroid Build Coastguard Worker        self.compute_client.SetInstanceMetadata = mock.MagicMock(
1289*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
1290*800a58d9SAndroid Build Coastguard Worker        # Test the key item not exists in the metadata.
1291*800a58d9SAndroid Build Coastguard Worker        self.compute_client.UpdateRsaInMetadata(
1292*800a58d9SAndroid Build Coastguard Worker            "fake_zone",
1293*800a58d9SAndroid Build Coastguard Worker            "fake_instance",
1294*800a58d9SAndroid Build Coastguard Worker            fake_metadata_sshkeys_not_exist,
1295*800a58d9SAndroid Build Coastguard Worker            new_entry)
1296*800a58d9SAndroid Build Coastguard Worker        self.compute_client.SetInstanceMetadata.assert_called_with(
1297*800a58d9SAndroid Build Coastguard Worker            "fake_zone",
1298*800a58d9SAndroid Build Coastguard Worker            "fake_instance",
1299*800a58d9SAndroid Build Coastguard Worker            expected)
1300*800a58d9SAndroid Build Coastguard Worker
1301*800a58d9SAndroid Build Coastguard Worker        # Test the key item exists in the metadata.
1302*800a58d9SAndroid Build Coastguard Worker        fake_metadata_ssh_keys_exists = {
1303*800a58d9SAndroid Build Coastguard Worker            "kind": "compute#metadata",
1304*800a58d9SAndroid Build Coastguard Worker            "fingerprint": "a-23icsyx4E=",
1305*800a58d9SAndroid Build Coastguard Worker            "items": [
1306*800a58d9SAndroid Build Coastguard Worker                {
1307*800a58d9SAndroid Build Coastguard Worker                    "key": "sshKeys",
1308*800a58d9SAndroid Build Coastguard Worker                    "value": "old_user:%s" % self.SSHKEY
1309*800a58d9SAndroid Build Coastguard Worker                }
1310*800a58d9SAndroid Build Coastguard Worker            ]
1311*800a58d9SAndroid Build Coastguard Worker        }
1312*800a58d9SAndroid Build Coastguard Worker        expected_ssh_keys_exists = {
1313*800a58d9SAndroid Build Coastguard Worker            "kind": "compute#metadata",
1314*800a58d9SAndroid Build Coastguard Worker            "fingerprint": "a-23icsyx4E=",
1315*800a58d9SAndroid Build Coastguard Worker            "items": [
1316*800a58d9SAndroid Build Coastguard Worker                {
1317*800a58d9SAndroid Build Coastguard Worker                    "key": "sshKeys",
1318*800a58d9SAndroid Build Coastguard Worker                    "value": "old_user:%s\n%s" % (self.SSHKEY, new_entry)
1319*800a58d9SAndroid Build Coastguard Worker                }
1320*800a58d9SAndroid Build Coastguard Worker            ]
1321*800a58d9SAndroid Build Coastguard Worker        }
1322*800a58d9SAndroid Build Coastguard Worker
1323*800a58d9SAndroid Build Coastguard Worker        self.compute_client.UpdateRsaInMetadata(
1324*800a58d9SAndroid Build Coastguard Worker            "fake_zone",
1325*800a58d9SAndroid Build Coastguard Worker            "fake_instance",
1326*800a58d9SAndroid Build Coastguard Worker            fake_metadata_ssh_keys_exists,
1327*800a58d9SAndroid Build Coastguard Worker            new_entry)
1328*800a58d9SAndroid Build Coastguard Worker        self.compute_client.SetInstanceMetadata.assert_called_with(
1329*800a58d9SAndroid Build Coastguard Worker            "fake_zone",
1330*800a58d9SAndroid Build Coastguard Worker            "fake_instance",
1331*800a58d9SAndroid Build Coastguard Worker            expected_ssh_keys_exists)
1332*800a58d9SAndroid Build Coastguard Worker
1333*800a58d9SAndroid Build Coastguard Worker    def testAddSshRsaToInstance(self):
1334*800a58d9SAndroid Build Coastguard Worker        """Test add ssh rsa key to instance."""
1335*800a58d9SAndroid Build Coastguard Worker        fake_user = "fake_user"
1336*800a58d9SAndroid Build Coastguard Worker        instance_metadata_key_not_exist = {
1337*800a58d9SAndroid Build Coastguard Worker            "metadata": {
1338*800a58d9SAndroid Build Coastguard Worker                "kind": "compute#metadata",
1339*800a58d9SAndroid Build Coastguard Worker                "fingerprint": "a-23icsyx4E=",
1340*800a58d9SAndroid Build Coastguard Worker                "items": [
1341*800a58d9SAndroid Build Coastguard Worker                    {
1342*800a58d9SAndroid Build Coastguard Worker                        "key": "sshKeys",
1343*800a58d9SAndroid Build Coastguard Worker                        "value": ""
1344*800a58d9SAndroid Build Coastguard Worker                    }
1345*800a58d9SAndroid Build Coastguard Worker                ]
1346*800a58d9SAndroid Build Coastguard Worker            }
1347*800a58d9SAndroid Build Coastguard Worker        }
1348*800a58d9SAndroid Build Coastguard Worker        instance_metadata_key_exist = {
1349*800a58d9SAndroid Build Coastguard Worker            "metadata": {
1350*800a58d9SAndroid Build Coastguard Worker                "kind": "compute#metadata",
1351*800a58d9SAndroid Build Coastguard Worker                "fingerprint": "a-23icsyx4E=",
1352*800a58d9SAndroid Build Coastguard Worker                "items": [
1353*800a58d9SAndroid Build Coastguard Worker                    {
1354*800a58d9SAndroid Build Coastguard Worker                        "key": "sshKeys",
1355*800a58d9SAndroid Build Coastguard Worker                        "value": "%s:%s" % (fake_user, self.SSHKEY)
1356*800a58d9SAndroid Build Coastguard Worker                    }
1357*800a58d9SAndroid Build Coastguard Worker                ]
1358*800a58d9SAndroid Build Coastguard Worker            }
1359*800a58d9SAndroid Build Coastguard Worker        }
1360*800a58d9SAndroid Build Coastguard Worker        expected = {
1361*800a58d9SAndroid Build Coastguard Worker            "kind": "compute#metadata",
1362*800a58d9SAndroid Build Coastguard Worker            "fingerprint": "a-23icsyx4E=",
1363*800a58d9SAndroid Build Coastguard Worker            "items": [
1364*800a58d9SAndroid Build Coastguard Worker                {
1365*800a58d9SAndroid Build Coastguard Worker                    "key": "sshKeys",
1366*800a58d9SAndroid Build Coastguard Worker                    "value": "%s:%s" % (fake_user, self.SSHKEY)
1367*800a58d9SAndroid Build Coastguard Worker                }
1368*800a58d9SAndroid Build Coastguard Worker            ]
1369*800a58d9SAndroid Build Coastguard Worker        }
1370*800a58d9SAndroid Build Coastguard Worker
1371*800a58d9SAndroid Build Coastguard Worker        self.Patch(os.path, "exists", return_value=True)
1372*800a58d9SAndroid Build Coastguard Worker        m = mock.mock_open(read_data=self.SSHKEY)
1373*800a58d9SAndroid Build Coastguard Worker        self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
1374*800a58d9SAndroid Build Coastguard Worker        self.Patch(gcompute_client.ComputeClient, "GetZoneByInstance",
1375*800a58d9SAndroid Build Coastguard Worker                   return_value="fake_zone")
1376*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
1377*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.instances = mock.MagicMock(
1378*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
1379*800a58d9SAndroid Build Coastguard Worker        resource_mock.setMetadata = mock.MagicMock()
1380*800a58d9SAndroid Build Coastguard Worker
1381*800a58d9SAndroid Build Coastguard Worker        # Test the key not exists in the metadata.
1382*800a58d9SAndroid Build Coastguard Worker        self.Patch(
1383*800a58d9SAndroid Build Coastguard Worker            gcompute_client.ComputeClient, "GetInstance",
1384*800a58d9SAndroid Build Coastguard Worker            return_value=instance_metadata_key_not_exist)
1385*800a58d9SAndroid Build Coastguard Worker        with mock.patch("builtins.open", m):
1386*800a58d9SAndroid Build Coastguard Worker            self.compute_client.AddSshRsaInstanceMetadata(
1387*800a58d9SAndroid Build Coastguard Worker                fake_user,
1388*800a58d9SAndroid Build Coastguard Worker                "/path/to/test_rsa.pub",
1389*800a58d9SAndroid Build Coastguard Worker                "fake_instance")
1390*800a58d9SAndroid Build Coastguard Worker            resource_mock.setMetadata.assert_called_with(
1391*800a58d9SAndroid Build Coastguard Worker                project=PROJECT,
1392*800a58d9SAndroid Build Coastguard Worker                zone="fake_zone",
1393*800a58d9SAndroid Build Coastguard Worker                instance="fake_instance",
1394*800a58d9SAndroid Build Coastguard Worker                body=expected)
1395*800a58d9SAndroid Build Coastguard Worker
1396*800a58d9SAndroid Build Coastguard Worker        # Test the key already exists in the metadata.
1397*800a58d9SAndroid Build Coastguard Worker        resource_mock.setMetadata.call_count = 0
1398*800a58d9SAndroid Build Coastguard Worker        self.Patch(
1399*800a58d9SAndroid Build Coastguard Worker            gcompute_client.ComputeClient, "GetInstance",
1400*800a58d9SAndroid Build Coastguard Worker            return_value=instance_metadata_key_exist)
1401*800a58d9SAndroid Build Coastguard Worker        with mock.patch("builtins.open", m):
1402*800a58d9SAndroid Build Coastguard Worker            self.compute_client.AddSshRsaInstanceMetadata(
1403*800a58d9SAndroid Build Coastguard Worker                fake_user,
1404*800a58d9SAndroid Build Coastguard Worker                "/path/to/test_rsa.pub",
1405*800a58d9SAndroid Build Coastguard Worker                "fake_instance")
1406*800a58d9SAndroid Build Coastguard Worker            resource_mock.setMetadata.assert_not_called()
1407*800a58d9SAndroid Build Coastguard Worker
1408*800a58d9SAndroid Build Coastguard Worker    @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation")
1409*800a58d9SAndroid Build Coastguard Worker    def testDeleteDisks(self, mock_wait):
1410*800a58d9SAndroid Build Coastguard Worker        """Test DeleteDisks."""
1411*800a58d9SAndroid Build Coastguard Worker        self._SetupBatchHttpRequestMock()
1412*800a58d9SAndroid Build Coastguard Worker        fake_disks = ["fake_disk_1", "fake_disk_2"]
1413*800a58d9SAndroid Build Coastguard Worker        mock_api = mock.MagicMock()
1414*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
1415*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.disks = mock.MagicMock(
1416*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
1417*800a58d9SAndroid Build Coastguard Worker        resource_mock.delete = mock.MagicMock(return_value=mock_api)
1418*800a58d9SAndroid Build Coastguard Worker        # Call the API.
1419*800a58d9SAndroid Build Coastguard Worker        deleted, failed, error_msgs = self.compute_client.DeleteDisks(
1420*800a58d9SAndroid Build Coastguard Worker            fake_disks, zone=self.ZONE)
1421*800a58d9SAndroid Build Coastguard Worker        # Verify
1422*800a58d9SAndroid Build Coastguard Worker        calls = [
1423*800a58d9SAndroid Build Coastguard Worker            mock.call(project=PROJECT, disk="fake_disk_1", zone=self.ZONE),
1424*800a58d9SAndroid Build Coastguard Worker            mock.call(project=PROJECT, disk="fake_disk_2", zone=self.ZONE)
1425*800a58d9SAndroid Build Coastguard Worker        ]
1426*800a58d9SAndroid Build Coastguard Worker        resource_mock.delete.assert_has_calls(calls, any_order=True)
1427*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(mock_wait.call_count, 2)
1428*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(error_msgs, [])
1429*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(failed, [])
1430*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(set(deleted), set(fake_disks))
1431*800a58d9SAndroid Build Coastguard Worker
1432*800a58d9SAndroid Build Coastguard Worker    def testRetryOnFingerPrintError(self):
1433*800a58d9SAndroid Build Coastguard Worker        """Test RetryOnFingerPrintError."""
1434*800a58d9SAndroid Build Coastguard Worker        @utils.RetryOnException(gcompute_client._IsFingerPrintError, 10)
1435*800a58d9SAndroid Build Coastguard Worker        def Raise412(sentinel):
1436*800a58d9SAndroid Build Coastguard Worker            """Raise 412 HTTP exception."""
1437*800a58d9SAndroid Build Coastguard Worker            if not sentinel.hitFingerPrintConflict.called:
1438*800a58d9SAndroid Build Coastguard Worker                sentinel.hitFingerPrintConflict()
1439*800a58d9SAndroid Build Coastguard Worker                raise errors.HttpError(412, "resource labels have changed")
1440*800a58d9SAndroid Build Coastguard Worker            return "Passed"
1441*800a58d9SAndroid Build Coastguard Worker
1442*800a58d9SAndroid Build Coastguard Worker        sentinel = mock.MagicMock()
1443*800a58d9SAndroid Build Coastguard Worker        result = Raise412(sentinel)
1444*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(1, sentinel.hitFingerPrintConflict.call_count)
1445*800a58d9SAndroid Build Coastguard Worker        self.assertEqual("Passed", result)
1446*800a58d9SAndroid Build Coastguard Worker
1447*800a58d9SAndroid Build Coastguard Worker    def testCheckAccess(self):
1448*800a58d9SAndroid Build Coastguard Worker        """Test CheckAccess."""
1449*800a58d9SAndroid Build Coastguard Worker        # Checking non-403 should raise error
1450*800a58d9SAndroid Build Coastguard Worker        error = errors.HttpError(503, "fake retriable error.")
1451*800a58d9SAndroid Build Coastguard Worker        self.Patch(
1452*800a58d9SAndroid Build Coastguard Worker            gcompute_client.ComputeClient, "Execute",
1453*800a58d9SAndroid Build Coastguard Worker            side_effect=error)
1454*800a58d9SAndroid Build Coastguard Worker
1455*800a58d9SAndroid Build Coastguard Worker        with self.assertRaises(errors.HttpError):
1456*800a58d9SAndroid Build Coastguard Worker            self.compute_client.CheckAccess()
1457*800a58d9SAndroid Build Coastguard Worker
1458*800a58d9SAndroid Build Coastguard Worker        # Checking 403 should return False
1459*800a58d9SAndroid Build Coastguard Worker        error = errors.HttpError(403, "fake retriable error.")
1460*800a58d9SAndroid Build Coastguard Worker        self.Patch(
1461*800a58d9SAndroid Build Coastguard Worker            gcompute_client.ComputeClient, "Execute",
1462*800a58d9SAndroid Build Coastguard Worker            side_effect=error)
1463*800a58d9SAndroid Build Coastguard Worker        self.assertFalse(self.compute_client.CheckAccess())
1464*800a58d9SAndroid Build Coastguard Worker
1465*800a58d9SAndroid Build Coastguard Worker    def testEnoughMetricsInZone(self):
1466*800a58d9SAndroid Build Coastguard Worker        """Test EnoughMetricsInZone."""
1467*800a58d9SAndroid Build Coastguard Worker        region_info_enough_quota = {
1468*800a58d9SAndroid Build Coastguard Worker            "items": [{
1469*800a58d9SAndroid Build Coastguard Worker                "name": "asia-east1",
1470*800a58d9SAndroid Build Coastguard Worker                "quotas": [{
1471*800a58d9SAndroid Build Coastguard Worker                    "usage": 50,
1472*800a58d9SAndroid Build Coastguard Worker                    "metric": "CPUS",
1473*800a58d9SAndroid Build Coastguard Worker                    "limit": 100
1474*800a58d9SAndroid Build Coastguard Worker                }, {
1475*800a58d9SAndroid Build Coastguard Worker                    "usage": 640,
1476*800a58d9SAndroid Build Coastguard Worker                    "metric": "DISKS_TOTAL_GB",
1477*800a58d9SAndroid Build Coastguard Worker                    "limit": 10240
1478*800a58d9SAndroid Build Coastguard Worker                }, {
1479*800a58d9SAndroid Build Coastguard Worker                    "usage": 20,
1480*800a58d9SAndroid Build Coastguard Worker                    "metric": "IN_USE_ADDRESSES",
1481*800a58d9SAndroid Build Coastguard Worker                    "limit": 100
1482*800a58d9SAndroid Build Coastguard Worker                }]
1483*800a58d9SAndroid Build Coastguard Worker            }]
1484*800a58d9SAndroid Build Coastguard Worker        }
1485*800a58d9SAndroid Build Coastguard Worker        self.Patch(
1486*800a58d9SAndroid Build Coastguard Worker            gcompute_client.ComputeClient, "GetRegionInfo",
1487*800a58d9SAndroid Build Coastguard Worker            return_value=region_info_enough_quota)
1488*800a58d9SAndroid Build Coastguard Worker        self.assertTrue(self.compute_client.EnoughMetricsInZone("asia-east1-b"))
1489*800a58d9SAndroid Build Coastguard Worker        self.assertFalse(self.compute_client.EnoughMetricsInZone("fake_zone"))
1490*800a58d9SAndroid Build Coastguard Worker
1491*800a58d9SAndroid Build Coastguard Worker        region_info_not_enough_quota = {
1492*800a58d9SAndroid Build Coastguard Worker            "items": [{
1493*800a58d9SAndroid Build Coastguard Worker                "name": "asia-east1",
1494*800a58d9SAndroid Build Coastguard Worker                "quotas": [{
1495*800a58d9SAndroid Build Coastguard Worker                    "usage": 100,
1496*800a58d9SAndroid Build Coastguard Worker                    "metric": "CPUS",
1497*800a58d9SAndroid Build Coastguard Worker                    "limit": 100
1498*800a58d9SAndroid Build Coastguard Worker                }, {
1499*800a58d9SAndroid Build Coastguard Worker                    "usage": 640,
1500*800a58d9SAndroid Build Coastguard Worker                    "metric": "DISKS_TOTAL_GB",
1501*800a58d9SAndroid Build Coastguard Worker                    "limit": 10240
1502*800a58d9SAndroid Build Coastguard Worker                }, {
1503*800a58d9SAndroid Build Coastguard Worker                    "usage": 20,
1504*800a58d9SAndroid Build Coastguard Worker                    "metric": "IN_USE_ADDRESSES",
1505*800a58d9SAndroid Build Coastguard Worker                    "limit": 100
1506*800a58d9SAndroid Build Coastguard Worker                }]
1507*800a58d9SAndroid Build Coastguard Worker            }]
1508*800a58d9SAndroid Build Coastguard Worker        }
1509*800a58d9SAndroid Build Coastguard Worker        self.Patch(
1510*800a58d9SAndroid Build Coastguard Worker            gcompute_client.ComputeClient, "GetRegionInfo",
1511*800a58d9SAndroid Build Coastguard Worker            return_value=region_info_not_enough_quota)
1512*800a58d9SAndroid Build Coastguard Worker        self.assertFalse(self.compute_client.EnoughMetricsInZone("asia-east1-b"))
1513*800a58d9SAndroid Build Coastguard Worker
1514*800a58d9SAndroid Build Coastguard Worker    def testGetDisk(self):
1515*800a58d9SAndroid Build Coastguard Worker        """Test GetDisk."""
1516*800a58d9SAndroid Build Coastguard Worker        resource_mock = mock.MagicMock()
1517*800a58d9SAndroid Build Coastguard Worker        mock_api = mock.MagicMock()
1518*800a58d9SAndroid Build Coastguard Worker        self.compute_client._service.disks = mock.MagicMock(
1519*800a58d9SAndroid Build Coastguard Worker            return_value=resource_mock)
1520*800a58d9SAndroid Build Coastguard Worker        resource_mock.get = mock.MagicMock(return_value=mock_api)
1521*800a58d9SAndroid Build Coastguard Worker        mock_api.execute = mock.MagicMock(return_value={"name": self.DISK})
1522*800a58d9SAndroid Build Coastguard Worker        result = self.compute_client.GetDisk(self.DISK, self.ZONE)
1523*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(result, {"name": self.DISK})
1524*800a58d9SAndroid Build Coastguard Worker        resource_mock.get.assert_called_with(project=PROJECT,
1525*800a58d9SAndroid Build Coastguard Worker                                             zone=self.ZONE,
1526*800a58d9SAndroid Build Coastguard Worker                                             disk=self.DISK)
1527*800a58d9SAndroid Build Coastguard Worker        self.assertTrue(self.compute_client.CheckDiskExists(self.DISK, self.ZONE))
1528*800a58d9SAndroid Build Coastguard Worker
1529*800a58d9SAndroid Build Coastguard Worker    def testGetGCEHostName(self):
1530*800a58d9SAndroid Build Coastguard Worker        """Test GetGCEHostName."""
1531*800a58d9SAndroid Build Coastguard Worker        instance_name = "instance_name"
1532*800a58d9SAndroid Build Coastguard Worker        project = "fake-project"
1533*800a58d9SAndroid Build Coastguard Worker        zone = "fake-zone"
1534*800a58d9SAndroid Build Coastguard Worker        expected = "nic0.instance_name.fake-zone.c.fake-project.internal.gcpnode.com"
1535*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(expected, gcompute_client.GetGCEHostName(
1536*800a58d9SAndroid Build Coastguard Worker                             project, instance_name, zone))
1537*800a58d9SAndroid Build Coastguard Worker
1538*800a58d9SAndroid Build Coastguard Worker        project = "test.com:project"
1539*800a58d9SAndroid Build Coastguard Worker        expected = "nic0.instance_name.fake-zone.c.project.test.com.internal.gcpnode.com"
1540*800a58d9SAndroid Build Coastguard Worker        self.assertEqual(expected, gcompute_client.GetGCEHostName(
1541*800a58d9SAndroid Build Coastguard Worker                             project, instance_name, zone))
1542*800a58d9SAndroid Build Coastguard Worker
1543*800a58d9SAndroid Build Coastguard Worker
1544*800a58d9SAndroid Build Coastguard Workerif __name__ == "__main__":
1545*800a58d9SAndroid Build Coastguard Worker    unittest.main()
1546