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