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 23 24from unittest import mock 25import six 26 27# pylint: disable=import-error 28from acloud import errors 29from acloud.internal import constants 30from acloud.internal.lib import driver_test_lib 31from acloud.internal.lib import gcompute_client 32from acloud.internal.lib import utils 33 34 35GS_IMAGE_SOURCE_URI = "https://storage.googleapis.com/fake-bucket/fake.tar.gz" 36GS_IMAGE_SOURCE_DISK = ( 37 "https://www.googleapis.com/compute/v1/projects/fake-project/zones/" 38 "us-east1-d/disks/fake-disk") 39PROJECT = "fake-project" 40 41 42# pylint: disable=protected-access, too-many-public-methods 43class ComputeClientTest(driver_test_lib.BaseDriverTest): 44 """Test ComputeClient.""" 45 46 PROJECT_OTHER = "fake-project-other" 47 INSTANCE = "fake-instance" 48 IMAGE = "fake-image" 49 IMAGE_URL = "http://fake-image-url" 50 IMAGE_OTHER = "fake-image-other" 51 DISK = "fake-disk" 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(self.compute_client._service, 409 "new_batch_http_request", 410 return_value=mock_batch) 411 412 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 413 def testDeleteImages(self, mock_wait): 414 """Test DeleteImages.""" 415 self._SetupBatchHttpRequestMock() 416 fake_images = ["fake_image_1", "fake_image_2"] 417 mock_api = mock.MagicMock() 418 resource_mock = mock.MagicMock() 419 self.compute_client._service.images = mock.MagicMock( 420 return_value=resource_mock) 421 resource_mock.delete = mock.MagicMock(return_value=mock_api) 422 # Call the API. 423 deleted, failed, error_msgs = self.compute_client.DeleteImages( 424 fake_images) 425 # Verify 426 calls = [ 427 mock.call(project=PROJECT, image="fake_image_1"), 428 mock.call(project=PROJECT, image="fake_image_2") 429 ] 430 resource_mock.delete.assert_has_calls(calls, any_order=True) 431 self.assertEqual(mock_wait.call_count, 2) 432 self.assertEqual(error_msgs, []) 433 self.assertEqual(failed, []) 434 self.assertEqual(set(deleted), set(fake_images)) 435 436 def testListImages(self): 437 """Test ListImages.""" 438 fake_token = "fake_next_page_token" 439 image_1 = "image_1" 440 image_2 = "image_2" 441 response_1 = {"items": [image_1], "nextPageToken": fake_token} 442 response_2 = {"items": [image_2]} 443 self.Patch( 444 gcompute_client.ComputeClient, 445 "Execute", 446 side_effect=[response_1, response_2]) 447 resource_mock = mock.MagicMock() 448 self.compute_client._service.images = mock.MagicMock( 449 return_value=resource_mock) 450 resource_mock.list = mock.MagicMock() 451 images = self.compute_client.ListImages() 452 calls = [ 453 mock.call(project=PROJECT, filter=None, pageToken=None), 454 mock.call(project=PROJECT, filter=None, pageToken=fake_token) 455 ] 456 resource_mock.list.assert_has_calls(calls) 457 self.assertEqual(images, [image_1, image_2]) 458 459 def testListImagesFromExternalProject(self): 460 """Test ListImages which accepts different project.""" 461 image = "image_1" 462 response = {"items": [image]} 463 self.Patch(gcompute_client.ComputeClient, "Execute", side_effect=[response]) 464 resource_mock = mock.MagicMock() 465 self.compute_client._service.images = mock.MagicMock( 466 return_value=resource_mock) 467 resource_mock.list = mock.MagicMock() 468 images = self.compute_client.ListImages( 469 image_project="fake-project-2") 470 calls = [ 471 mock.call(project="fake-project-2", filter=None, pageToken=None)] 472 resource_mock.list.assert_has_calls(calls) 473 self.assertEqual(images, [image]) 474 475 def testGetInstance(self): 476 """Test GetInstance.""" 477 resource_mock = mock.MagicMock() 478 mock_api = mock.MagicMock() 479 self.compute_client._service.instances = mock.MagicMock( 480 return_value=resource_mock) 481 resource_mock.get = mock.MagicMock(return_value=mock_api) 482 mock_api.execute = mock.MagicMock(return_value={"name": self.INSTANCE}) 483 result = self.compute_client.GetInstance(self.INSTANCE, self.ZONE) 484 self.assertEqual(result, {"name": self.INSTANCE}) 485 resource_mock.get.assert_called_with( 486 project=PROJECT, zone=self.ZONE, instance=self.INSTANCE) 487 488 def testListInstances(self): 489 """Test ListInstances.""" 490 instance_1 = "instance_1" 491 instance_2 = "instance_2" 492 response = {"items": {'zones/fake_zone': {"instances": [instance_1, instance_2]}}} 493 self.Patch( 494 gcompute_client.ComputeClient, 495 "Execute", 496 side_effect=[response]) 497 resource_mock = mock.MagicMock() 498 self.compute_client._service.instances = mock.MagicMock( 499 return_value=resource_mock) 500 resource_mock.aggregatedList = mock.MagicMock() 501 instances = self.compute_client.ListInstances() 502 calls = [ 503 mock.call( 504 project=PROJECT, 505 filter=None, 506 pageToken=None), 507 ] 508 resource_mock.aggregatedList.assert_has_calls(calls) 509 self.assertEqual(instances, [instance_1, instance_2]) 510 511 def testGetZoneByInstance(self): 512 """Test GetZoneByInstance.""" 513 instance_1 = "instance_1" 514 response = {"items": {'zones/fake_zone': {"instances": [instance_1]}}} 515 self.Patch( 516 gcompute_client.ComputeClient, 517 "Execute", 518 side_effect=[response]) 519 expected_zone = "fake_zone" 520 self.assertEqual(self.compute_client.GetZoneByInstance(instance_1), 521 expected_zone) 522 523 # Test unable to find 'zone' from instance name. 524 response = {"items": {'zones/fake_zone': {"warning": "No instances."}}} 525 self.Patch( 526 gcompute_client.ComputeClient, 527 "Execute", 528 side_effect=[response]) 529 with self.assertRaises(errors.GetGceZoneError): 530 self.compute_client.GetZoneByInstance(instance_1) 531 532 def testGetZonesByInstances(self): 533 """Test GetZonesByInstances.""" 534 instances = ["instance_1", "instance_2"] 535 # Test instances in the same zone. 536 self.Patch( 537 gcompute_client.ComputeClient, 538 "GetZoneByInstance", 539 side_effect=["zone_1", "zone_1"]) 540 expected_result = {"zone_1": ["instance_1", "instance_2"]} 541 self.assertEqual(self.compute_client.GetZonesByInstances(instances), 542 expected_result) 543 544 # Test instances in different zones. 545 self.Patch( 546 gcompute_client.ComputeClient, 547 "GetZoneByInstance", 548 side_effect=["zone_1", "zone_2"]) 549 expected_result = {"zone_1": ["instance_1"], 550 "zone_2": ["instance_2"]} 551 self.assertEqual(self.compute_client.GetZonesByInstances(instances), 552 expected_result) 553 554 @mock.patch.object(gcompute_client.ComputeClient, "GetImage") 555 @mock.patch.object(gcompute_client.ComputeClient, "GetNetworkUrl") 556 @mock.patch.object(gcompute_client.ComputeClient, "GetSubnetworkUrl") 557 @mock.patch.object(gcompute_client.ComputeClient, "GetMachineType") 558 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 559 @mock.patch("getpass.getuser", return_value="fake_user") 560 def testCreateInstance(self, _get_user, mock_wait, mock_get_mach_type, 561 mock_get_subnetwork_url, mock_get_network_url, 562 mock_get_image): 563 """Test CreateInstance.""" 564 mock_get_mach_type.return_value = {"selfLink": self.MACHINE_TYPE_URL} 565 mock_get_network_url.return_value = self.NETWORK_URL 566 mock_get_subnetwork_url.return_value = self.SUBNETWORK_URL 567 mock_get_image.return_value = {"selfLink": self.IMAGE_URL} 568 resource_mock = mock.MagicMock() 569 self.compute_client._service.instances = mock.MagicMock( 570 return_value=resource_mock) 571 resource_mock.insert = mock.MagicMock() 572 self.Patch( 573 self.compute_client, 574 "_GetExtraDiskArgs", 575 return_value=[{"fake_extra_arg": "fake_extra_value"}]) 576 extra_disk_name = "gce-x86-userdebug-2345-abcd-data" 577 expected_disk_args = [self._disk_args] 578 expected_disk_args.extend([{"fake_extra_arg": "fake_extra_value"}]) 579 expected_scope = [] 580 expected_scope.extend(self.compute_client.DEFAULT_INSTANCE_SCOPE) 581 expected_scope.extend(self.EXTRA_SCOPES) 582 583 expected_body = { 584 "machineType": self.MACHINE_TYPE_URL, 585 "name": self.INSTANCE, 586 "networkInterfaces": [ 587 { 588 "network": self.NETWORK_URL, 589 "subnetwork": self.SUBNETWORK_URL, 590 "accessConfigs": [ 591 {"name": "External NAT", 592 "type": "ONE_TO_ONE_NAT"} 593 ], 594 } 595 ], 596 "disks": expected_disk_args, 597 "serviceAccounts": [ 598 {"email": "default", 599 "scopes": expected_scope} 600 ], 601 "metadata": { 602 "items": [{"key": self.METADATA[0], 603 "value": self.METADATA[1]}], 604 }, 605 "labels":{constants.LABEL_CREATE_BY: "fake_user"}, 606 "enableVtpm": True, 607 } 608 609 self.compute_client.CreateInstance( 610 instance=self.INSTANCE, 611 image_name=self.IMAGE, 612 machine_type=self.MACHINE_TYPE, 613 metadata={self.METADATA[0]: self.METADATA[1]}, 614 network=self.NETWORK, 615 zone=self.ZONE, 616 extra_disk_name=extra_disk_name, 617 extra_scopes=self.EXTRA_SCOPES) 618 619 resource_mock.insert.assert_called_with( 620 project=PROJECT, zone=self.ZONE, body=expected_body) 621 mock_wait.assert_called_with( 622 mock.ANY, 623 operation_scope=gcompute_client.OperationScope.ZONE, 624 scope_name=self.ZONE) 625 626 627 @mock.patch.object(gcompute_client.ComputeClient, "GetImage") 628 @mock.patch.object(gcompute_client.ComputeClient, "GetNetworkUrl") 629 @mock.patch.object(gcompute_client.ComputeClient, "GetSubnetworkUrl") 630 @mock.patch.object(gcompute_client.ComputeClient, "GetMachineType") 631 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 632 @mock.patch("getpass.getuser", return_value="fake_user") 633 def testCreateInstanceWithTags(self, 634 _get_user, 635 mock_wait, 636 mock_get_mach_type, 637 mock_get_subnetwork_url, 638 mock_get_network_url, 639 mock_get_image): 640 """Test CreateInstance.""" 641 mock_get_mach_type.return_value = {"selfLink": self.MACHINE_TYPE_URL} 642 mock_get_network_url.return_value = self.NETWORK_URL 643 mock_get_subnetwork_url.return_value = self.SUBNETWORK_URL 644 mock_get_image.return_value = {"selfLink": self.IMAGE_URL} 645 resource_mock = mock.MagicMock() 646 self.compute_client._service.instances = mock.MagicMock( 647 return_value=resource_mock) 648 resource_mock.insert = mock.MagicMock() 649 self.Patch( 650 self.compute_client, 651 "_GetExtraDiskArgs", 652 return_value=[{"fake_extra_arg": "fake_extra_value"}]) 653 extra_disk_name = "gce-x86-userdebug-2345-abcd-data" 654 expected_disk_args = [self._disk_args] 655 expected_disk_args.extend([{"fake_extra_arg": "fake_extra_value"}]) 656 expected_scope = [] 657 expected_scope.extend(self.compute_client.DEFAULT_INSTANCE_SCOPE) 658 expected_scope.extend(self.EXTRA_SCOPES) 659 660 expected_body = { 661 "machineType": self.MACHINE_TYPE_URL, 662 "name": self.INSTANCE, 663 "networkInterfaces": [ 664 { 665 "network": self.NETWORK_URL, 666 "subnetwork": self.SUBNETWORK_URL, 667 "accessConfigs": [ 668 {"name": "External NAT", 669 "type": "ONE_TO_ONE_NAT"} 670 ], 671 } 672 ], 673 'tags': {'items': ['https-server']}, 674 "disks": expected_disk_args, 675 "serviceAccounts": [ 676 {"email": "default", 677 "scopes": expected_scope} 678 ], 679 "metadata": { 680 "items": [{"key": self.METADATA[0], 681 "value": self.METADATA[1]}], 682 }, 683 "labels":{'created_by': "fake_user"}, 684 "enableVtpm": True, 685 } 686 687 self.compute_client.CreateInstance( 688 instance=self.INSTANCE, 689 image_name=self.IMAGE, 690 machine_type=self.MACHINE_TYPE, 691 metadata={self.METADATA[0]: self.METADATA[1]}, 692 network=self.NETWORK, 693 zone=self.ZONE, 694 extra_disk_name=extra_disk_name, 695 tags=["https-server"], 696 extra_scopes=self.EXTRA_SCOPES) 697 698 resource_mock.insert.assert_called_with( 699 project=PROJECT, zone=self.ZONE, body=expected_body) 700 mock_wait.assert_called_with( 701 mock.ANY, 702 operation_scope=gcompute_client.OperationScope.ZONE, 703 scope_name=self.ZONE) 704 705 @mock.patch.object(gcompute_client.ComputeClient, "GetAcceleratorUrl") 706 @mock.patch.object(gcompute_client.ComputeClient, "GetImage") 707 @mock.patch.object(gcompute_client.ComputeClient, "GetNetworkUrl") 708 @mock.patch.object(gcompute_client.ComputeClient, "GetSubnetworkUrl") 709 @mock.patch.object(gcompute_client.ComputeClient, "GetMachineType") 710 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 711 @mock.patch("getpass.getuser", return_value="fake_user") 712 def testCreateInstanceWithGpu(self, _get_user, mock_wait, mock_get_mach, 713 mock_get_subnetwork, mock_get_network, 714 mock_get_image, mock_get_accel): 715 """Test CreateInstance with a GPU parameter not set to None.""" 716 mock_get_mach.return_value = {"selfLink": self.MACHINE_TYPE_URL} 717 mock_get_network.return_value = self.NETWORK_URL 718 mock_get_subnetwork.return_value = self.SUBNETWORK_URL 719 mock_get_accel.return_value = self.ACCELERATOR_URL 720 mock_get_image.return_value = {"selfLink": self.IMAGE_URL} 721 722 resource_mock = mock.MagicMock() 723 self.compute_client._service.instances = mock.MagicMock( 724 return_value=resource_mock) 725 resource_mock.insert = mock.MagicMock() 726 727 expected_body = { 728 "machineType": 729 self.MACHINE_TYPE_URL, 730 "name": 731 self.INSTANCE, 732 "networkInterfaces": [{ 733 "network": self.NETWORK_URL, 734 "subnetwork": self.SUBNETWORK_URL, 735 "accessConfigs": [{ 736 "name": "External NAT", 737 "type": "ONE_TO_ONE_NAT" 738 }], 739 }], 740 "disks": [self._disk_args], 741 "serviceAccounts": [{ 742 "email": "default", 743 "scopes": self.compute_client.DEFAULT_INSTANCE_SCOPE 744 }], 745 "scheduling": { 746 "onHostMaintenance": "terminate" 747 }, 748 "guestAccelerators": [{ 749 "acceleratorCount": 1, 750 "acceleratorType": "http://speedy-gpu" 751 }], 752 "metadata": { 753 "items": [{ 754 "key": self.METADATA[0], 755 "value": self.METADATA[1] 756 }], 757 }, 758 "labels":{'created_by': "fake_user"}, 759 "enableVtpm": True, 760 } 761 762 self.compute_client.CreateInstance( 763 instance=self.INSTANCE, 764 image_name=self.IMAGE, 765 machine_type=self.MACHINE_TYPE, 766 metadata={self.METADATA[0]: self.METADATA[1]}, 767 network=self.NETWORK, 768 zone=self.ZONE, 769 gpu=self.GPU, 770 extra_scopes=None) 771 772 resource_mock.insert.assert_called_with( 773 project=PROJECT, zone=self.ZONE, body=expected_body) 774 mock_wait.assert_called_with( 775 mock.ANY, operation_scope=gcompute_client.OperationScope.ZONE, 776 scope_name=self.ZONE) 777 778 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 779 def testDeleteInstance(self, mock_wait): 780 """Test DeleteInstance.""" 781 resource_mock = mock.MagicMock() 782 self.compute_client._service.instances = mock.MagicMock( 783 return_value=resource_mock) 784 resource_mock.delete = mock.MagicMock() 785 self.compute_client.DeleteInstance( 786 instance=self.INSTANCE, zone=self.ZONE) 787 resource_mock.delete.assert_called_with( 788 project=PROJECT, zone=self.ZONE, instance=self.INSTANCE) 789 mock_wait.assert_called_with( 790 mock.ANY, 791 operation_scope=gcompute_client.OperationScope.ZONE, 792 scope_name=self.ZONE) 793 794 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 795 def testDeleteInstances(self, mock_wait): 796 """Test DeleteInstances.""" 797 self._SetupBatchHttpRequestMock() 798 fake_instances = ["fake_instance_1", "fake_instance_2"] 799 mock_api = mock.MagicMock() 800 resource_mock = mock.MagicMock() 801 self.compute_client._service.instances = mock.MagicMock( 802 return_value=resource_mock) 803 resource_mock.delete = mock.MagicMock(return_value=mock_api) 804 deleted, failed, error_msgs = self.compute_client.DeleteInstances( 805 fake_instances, self.ZONE) 806 calls = [ 807 mock.call( 808 project=PROJECT, 809 instance="fake_instance_1", 810 zone=self.ZONE), 811 mock.call( 812 project=PROJECT, 813 instance="fake_instance_2", 814 zone=self.ZONE) 815 ] 816 resource_mock.delete.assert_has_calls(calls, any_order=True) 817 self.assertEqual(mock_wait.call_count, 2) 818 self.assertEqual(error_msgs, []) 819 self.assertEqual(failed, []) 820 self.assertEqual(set(deleted), set(fake_instances)) 821 822 def testCreateDiskWithProject(self): 823 """Test CreateDisk with images using a set project.""" 824 source_project = "fake-image-project" 825 expected_project_to_use = "fake-image-project" 826 mock_wait = self.Patch(gcompute_client.ComputeClient, "WaitOnOperation") 827 resource_mock = mock.MagicMock() 828 self.compute_client._service.disks = mock.MagicMock( 829 return_value=resource_mock) 830 resource_mock.insert = mock.MagicMock() 831 self.compute_client.CreateDisk( 832 "fake_disk", "fake_image", 10, self.ZONE, source_project=source_project) 833 resource_mock.insert.assert_called_with( 834 project=PROJECT, 835 zone=self.ZONE, 836 sourceImage="projects/%s/global/images/fake_image" % 837 expected_project_to_use, 838 body={ 839 "name": 840 "fake_disk", 841 "sizeGb": 842 10, 843 "type": 844 "projects/%s/zones/%s/diskTypes/pd-standard" % (PROJECT, 845 self.ZONE) 846 }) 847 self.assertTrue(mock_wait.called) 848 849 def testCreateDiskWithNoSourceProject(self): 850 """Test CreateDisk with images with no set project.""" 851 source_project = None 852 expected_project_to_use = PROJECT 853 mock_wait = self.Patch(gcompute_client.ComputeClient, "WaitOnOperation") 854 resource_mock = mock.MagicMock() 855 self.compute_client._service.disks = mock.MagicMock( 856 return_value=resource_mock) 857 resource_mock.insert = mock.MagicMock() 858 self.compute_client.CreateDisk( 859 "fake_disk", "fake_image", 10, self.ZONE, source_project=source_project) 860 resource_mock.insert.assert_called_with( 861 project=PROJECT, 862 zone=self.ZONE, 863 sourceImage="projects/%s/global/images/fake_image" % 864 expected_project_to_use, 865 body={ 866 "name": 867 "fake_disk", 868 "sizeGb": 869 10, 870 "type": 871 "projects/%s/zones/%s/diskTypes/pd-standard" % (PROJECT, 872 self.ZONE) 873 }) 874 self.assertTrue(mock_wait.called) 875 876 def testCreateDiskWithTypeStandard(self): 877 """Test CreateDisk with images using standard.""" 878 disk_type = gcompute_client.PersistentDiskType.STANDARD 879 expected_disk_type_string = "pd-standard" 880 mock_wait = self.Patch(gcompute_client.ComputeClient, "WaitOnOperation") 881 resource_mock = mock.MagicMock() 882 self.compute_client._service.disks = mock.MagicMock( 883 return_value=resource_mock) 884 resource_mock.insert = mock.MagicMock() 885 self.compute_client.CreateDisk( 886 "fake_disk", 887 "fake_image", 888 10, 889 self.ZONE, 890 source_project="fake-project", 891 disk_type=disk_type) 892 resource_mock.insert.assert_called_with( 893 project=PROJECT, 894 zone=self.ZONE, 895 sourceImage="projects/%s/global/images/fake_image" % "fake-project", 896 body={ 897 "name": 898 "fake_disk", 899 "sizeGb": 900 10, 901 "type": 902 "projects/%s/zones/%s/diskTypes/%s" % 903 (PROJECT, self.ZONE, expected_disk_type_string) 904 }) 905 self.assertTrue(mock_wait.called) 906 907 def testCreateDiskWithTypeSSD(self): 908 """Test CreateDisk with images using standard.""" 909 disk_type = gcompute_client.PersistentDiskType.SSD 910 expected_disk_type_string = "pd-ssd" 911 mock_wait = self.Patch(gcompute_client.ComputeClient, "WaitOnOperation") 912 resource_mock = mock.MagicMock() 913 self.compute_client._service.disks = mock.MagicMock( 914 return_value=resource_mock) 915 resource_mock.insert = mock.MagicMock() 916 self.compute_client.CreateDisk( 917 "fake_disk", 918 "fake_image", 919 10, 920 self.ZONE, 921 source_project="fake-project", 922 disk_type=disk_type) 923 resource_mock.insert.assert_called_with( 924 project=PROJECT, 925 zone=self.ZONE, 926 sourceImage="projects/%s/global/images/fake_image" % "fake-project", 927 body={ 928 "name": 929 "fake_disk", 930 "sizeGb": 931 10, 932 "type": 933 "projects/%s/zones/%s/diskTypes/%s" % 934 (PROJECT, self.ZONE, expected_disk_type_string) 935 }) 936 self.assertTrue(mock_wait.called) 937 938 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 939 def testAttachDisk(self, mock_wait): 940 """Test AttachDisk.""" 941 resource_mock = mock.MagicMock() 942 self.compute_client._service.instances = mock.MagicMock( 943 return_value=resource_mock) 944 resource_mock.attachDisk = mock.MagicMock() 945 self.compute_client.AttachDisk( 946 "fake_instance_1", self.ZONE, deviceName="fake_disk", 947 source="fake-selfLink") 948 resource_mock.attachDisk.assert_called_with( 949 project=PROJECT, 950 zone=self.ZONE, 951 instance="fake_instance_1", 952 body={ 953 "deviceName": "fake_disk", 954 "source": "fake-selfLink" 955 }) 956 self.assertTrue(mock_wait.called) 957 958 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 959 def testDetachDisk(self, mock_wait): 960 """Test DetachDisk.""" 961 resource_mock = mock.MagicMock() 962 self.compute_client._service.instances = mock.MagicMock( 963 return_value=resource_mock) 964 resource_mock.detachDisk = mock.MagicMock() 965 self.compute_client.DetachDisk("fake_instance_1", self.ZONE, "fake_disk") 966 resource_mock.detachDisk.assert_called_with( 967 project=PROJECT, 968 zone=self.ZONE, 969 instance="fake_instance_1", 970 deviceName="fake_disk") 971 self.assertTrue(mock_wait.called) 972 973 @mock.patch.object(gcompute_client.ComputeClient, "GetAcceleratorUrl") 974 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 975 def testAttachAccelerator(self, mock_wait, mock_get_accel): 976 """Test AttachAccelerator.""" 977 mock_get_accel.return_value = self.ACCELERATOR_URL 978 resource_mock = mock.MagicMock() 979 self.compute_client._service.instances = mock.MagicMock( 980 return_value=resource_mock) 981 resource_mock.attachAccelerator = mock.MagicMock() 982 self.compute_client.AttachAccelerator("fake_instance_1", self.ZONE, 1, 983 "nvidia-tesla-k80") 984 resource_mock.setMachineResources.assert_called_with( 985 project=PROJECT, 986 zone=self.ZONE, 987 instance="fake_instance_1", 988 body={ 989 "guestAccelerators": [{ 990 "acceleratorType": self.ACCELERATOR_URL, 991 "acceleratorCount": 1 992 }] 993 }) 994 self.assertTrue(mock_wait.called) 995 996 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 997 def testBatchExecuteOnInstances(self, mock_wait): 998 """Test BatchExecuteOnInstances.""" 999 self._SetupBatchHttpRequestMock() 1000 action = mock.MagicMock(return_value=mock.MagicMock()) 1001 fake_instances = ["fake_instance_1", "fake_instance_2"] 1002 done, failed, error_msgs = self.compute_client._BatchExecuteOnInstances( 1003 fake_instances, self.ZONE, action) 1004 calls = [mock.call(instance="fake_instance_1"), 1005 mock.call(instance="fake_instance_2")] 1006 action.assert_has_calls(calls, any_order=True) 1007 self.assertEqual(mock_wait.call_count, 2) 1008 self.assertEqual(set(done), set(fake_instances)) 1009 self.assertEqual(error_msgs, []) 1010 self.assertEqual(failed, []) 1011 1012 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 1013 def testResetInstance(self, mock_wait): 1014 """Test ResetInstance.""" 1015 resource_mock = mock.MagicMock() 1016 self.compute_client._service.instances = mock.MagicMock( 1017 return_value=resource_mock) 1018 resource_mock.reset = mock.MagicMock() 1019 self.compute_client.ResetInstance( 1020 instance=self.INSTANCE, zone=self.ZONE) 1021 resource_mock.reset.assert_called_with( 1022 project=PROJECT, zone=self.ZONE, instance=self.INSTANCE) 1023 mock_wait.assert_called_with( 1024 mock.ANY, 1025 operation_scope=gcompute_client.OperationScope.ZONE, 1026 scope_name=self.ZONE) 1027 1028 def _CompareMachineSizeTestHelper(self, 1029 machine_info_1, 1030 machine_info_2, 1031 expected_result=None, 1032 expected_error_type=None): 1033 """Helper class for testing CompareMachineSize. 1034 1035 Args: 1036 machine_info_1: A dictionary representing the first machine size. 1037 machine_info_2: A dictionary representing the second machine size. 1038 expected_result: An integer, 0, 1 or -1, or None if not set. 1039 expected_error_type: An exception type, if set will check for exception. 1040 """ 1041 mock_get_mach_type = self.Patch( 1042 gcompute_client.ComputeClient, 1043 "GetMachineType", 1044 side_effect=[machine_info_1, machine_info_2]) 1045 if expected_error_type: 1046 self.assertRaises(expected_error_type, 1047 self.compute_client.CompareMachineSize, "name1", 1048 "name2", self.ZONE) 1049 else: 1050 result = self.compute_client.CompareMachineSize("name1", "name2", 1051 self.ZONE) 1052 self.assertEqual(result, expected_result) 1053 1054 mock_get_mach_type.assert_has_calls( 1055 [mock.call("name1", self.ZONE), mock.call("name2", self.ZONE)]) 1056 1057 def testCompareMachineSizeSmall(self): 1058 """Test CompareMachineSize where the first one is smaller.""" 1059 machine_info_1 = {"guestCpus": 10, "memoryMb": 100} 1060 machine_info_2 = {"guestCpus": 10, "memoryMb": 200} 1061 self._CompareMachineSizeTestHelper(machine_info_1, machine_info_2, -1) 1062 1063 def testCompareMachineSizeSmallSmallerOnSecond(self): 1064 """Test CompareMachineSize where the first one is smaller.""" 1065 machine_info_1 = {"guestCpus": 11, "memoryMb": 100} 1066 machine_info_2 = {"guestCpus": 10, "memoryMb": 200} 1067 self._CompareMachineSizeTestHelper(machine_info_1, machine_info_2, -1) 1068 1069 def testCompareMachineSizeLarge(self): 1070 """Test CompareMachineSize where the first one is larger.""" 1071 machine_info_1 = {"guestCpus": 11, "memoryMb": 200} 1072 machine_info_2 = {"guestCpus": 10, "memoryMb": 100} 1073 self._CompareMachineSizeTestHelper(machine_info_1, machine_info_2, 1) 1074 1075 def testCompareMachineSizeLargeWithEqualElement(self): 1076 """Test CompareMachineSize where the first one is larger.""" 1077 machine_info_1 = {"guestCpus": 10, "memoryMb": 200} 1078 machine_info_2 = {"guestCpus": 10, "memoryMb": 100} 1079 self._CompareMachineSizeTestHelper(machine_info_1, machine_info_2, 1) 1080 1081 def testCompareMachineSizeEqual(self): 1082 """Test CompareMachineSize where two machine sizes are equal.""" 1083 machine_info = {"guestCpus": 10, "memoryMb": 100} 1084 self._CompareMachineSizeTestHelper(machine_info, machine_info, 0) 1085 1086 def testCompareMachineSizeBadMetric(self): 1087 """Test CompareMachineSize with bad metric.""" 1088 machine_info = {"unknown_metric": 10, "memoryMb": 100} 1089 self._CompareMachineSizeTestHelper( 1090 machine_info, machine_info, expected_error_type=errors.DriverError) 1091 1092 def testGetMachineType(self): 1093 """Test GetMachineType.""" 1094 resource_mock = mock.MagicMock() 1095 mock_api = mock.MagicMock() 1096 self.compute_client._service.machineTypes = mock.MagicMock( 1097 return_value=resource_mock) 1098 resource_mock.get = mock.MagicMock(return_value=mock_api) 1099 mock_api.execute = mock.MagicMock( 1100 return_value={"name": self.MACHINE_TYPE}) 1101 result = self.compute_client.GetMachineType(self.MACHINE_TYPE, 1102 self.ZONE) 1103 self.assertEqual(result, {"name": self.MACHINE_TYPE}) 1104 resource_mock.get.assert_called_with( 1105 project=PROJECT, 1106 zone=self.ZONE, 1107 machineType=self.MACHINE_TYPE) 1108 1109 def _GetSerialPortOutputTestHelper(self, response): 1110 """Helper function for testing GetSerialPortOutput. 1111 1112 Args: 1113 response: A dictionary representing a fake response. 1114 """ 1115 resource_mock = mock.MagicMock() 1116 mock_api = mock.MagicMock() 1117 self.compute_client._service.instances = mock.MagicMock( 1118 return_value=resource_mock) 1119 resource_mock.getSerialPortOutput = mock.MagicMock( 1120 return_value=mock_api) 1121 mock_api.execute = mock.MagicMock(return_value=response) 1122 1123 if "contents" in response: 1124 result = self.compute_client.GetSerialPortOutput( 1125 instance=self.INSTANCE, zone=self.ZONE) 1126 self.assertEqual(result, "fake contents") 1127 else: 1128 six.assertRaisesRegex( 1129 self, 1130 errors.DriverError, 1131 "Malformed response.*", 1132 self.compute_client.GetSerialPortOutput, 1133 instance=self.INSTANCE, 1134 zone=self.ZONE) 1135 resource_mock.getSerialPortOutput.assert_called_with( 1136 project=PROJECT, 1137 zone=self.ZONE, 1138 instance=self.INSTANCE, 1139 port=1) 1140 1141 def testGetSerialPortOutput(self): 1142 """Test GetSerialPortOutput.""" 1143 response = {"contents": "fake contents"} 1144 self._GetSerialPortOutputTestHelper(response) 1145 1146 def testGetSerialPortOutputFail(self): 1147 """Test GetSerialPortOutputFail.""" 1148 response = {"malformed": "fake contents"} 1149 self._GetSerialPortOutputTestHelper(response) 1150 1151 def testGetInstanceNamesByIPs(self): 1152 """Test GetInstanceNamesByIPs.""" 1153 good_instance = { 1154 "name": "instance_1", 1155 "networkInterfaces": [ 1156 { 1157 "accessConfigs": [ 1158 {"natIP": "172.22.22.22"}, 1159 ], 1160 }, 1161 ], 1162 } 1163 bad_instance = {"name": "instance_2"} 1164 self.Patch( 1165 gcompute_client.ComputeClient, 1166 "ListInstances", 1167 return_value=[good_instance, bad_instance]) 1168 ip_name_map = self.compute_client.GetInstanceNamesByIPs( 1169 ips=["172.22.22.22", "172.22.22.23"]) 1170 self.assertEqual(ip_name_map, {"172.22.22.22": "instance_1", 1171 "172.22.22.23": None}) 1172 1173 def testRsaNotInMetadata(self): 1174 """Test rsa not in metadata.""" 1175 fake_user = "fake_user" 1176 fake_ssh_key = "fake_ssh" 1177 metadata = { 1178 "kind": "compute#metadata", 1179 "fingerprint": "a-23icsyx4E=", 1180 "items": [ 1181 { 1182 "key": "sshKeys", 1183 "value": "%s:%s" % (fake_user, self.SSHKEY) 1184 } 1185 ] 1186 } 1187 # Test rsa doesn't exist in metadata. 1188 new_entry = "%s:%s" % (fake_user, fake_ssh_key) 1189 self.assertEqual(True, gcompute_client.RsaNotInMetadata(metadata, new_entry)) 1190 1191 # Test rsa exists in metadata. 1192 exist_entry = "%s:%s" %(fake_user, self.SSHKEY) 1193 self.assertEqual(False, gcompute_client.RsaNotInMetadata(metadata, exist_entry)) 1194 1195 def testGetSshKeyFromMetadata(self): 1196 """Test get ssh key from metadata.""" 1197 fake_user = "fake_user" 1198 metadata_key_exist_value_is_empty = { 1199 "kind": "compute#metadata", 1200 "fingerprint": "a-23icsyx4E=", 1201 "items": [ 1202 { 1203 "key": "sshKeys", 1204 "value": "" 1205 } 1206 ] 1207 } 1208 metadata_key_exist = { 1209 "kind": "compute#metadata", 1210 "fingerprint": "a-23icsyx4E=", 1211 "items": [ 1212 { 1213 "key": "sshKeys", 1214 "value": "%s:%s" % (fake_user, self.SSHKEY) 1215 } 1216 ] 1217 } 1218 metadata_key_not_exist = { 1219 "kind": "compute#metadata", 1220 "fingerprint": "a-23icsyx4E=", 1221 "items": [ 1222 { 1223 } 1224 ] 1225 } 1226 expected_key_exist_value_is_empty = { 1227 "key": "sshKeys", 1228 "value": "" 1229 } 1230 expected_key_exist = { 1231 "key": "sshKeys", 1232 "value": "%s:%s" % (fake_user, self.SSHKEY) 1233 } 1234 self.assertEqual(expected_key_exist_value_is_empty, 1235 gcompute_client.GetSshKeyFromMetadata(metadata_key_exist_value_is_empty)) 1236 self.assertEqual(expected_key_exist, 1237 gcompute_client.GetSshKeyFromMetadata(metadata_key_exist)) 1238 self.assertEqual(None, 1239 gcompute_client.GetSshKeyFromMetadata(metadata_key_not_exist)) 1240 1241 1242 def testGetRsaKeyPathExistsFalse(self): 1243 """Test the rsa key path not exists.""" 1244 fake_ssh_rsa_path = "/path/to/test_rsa.pub" 1245 self.Patch(os.path, "exists", return_value=False) 1246 six.assertRaisesRegex(self, 1247 errors.DriverError, 1248 "RSA file %s does not exist." % fake_ssh_rsa_path, 1249 gcompute_client.GetRsaKey, 1250 ssh_rsa_path=fake_ssh_rsa_path) 1251 1252 def testGetRsaKey(self): 1253 """Test get the rsa key.""" 1254 fake_ssh_rsa_path = "/path/to/test_rsa.pub" 1255 self.Patch(os.path, "exists", return_value=True) 1256 m = mock.mock_open(read_data=self.SSHKEY) 1257 with mock.patch.object(six.moves.builtins, "open", m): 1258 result = gcompute_client.GetRsaKey(fake_ssh_rsa_path) 1259 self.assertEqual(self.SSHKEY, result) 1260 1261 def testUpdateRsaInMetadata(self): 1262 """Test update rsa in metadata.""" 1263 fake_ssh_key = "fake_ssh" 1264 fake_metadata_sshkeys_not_exist = { 1265 "kind": "compute#metadata", 1266 "fingerprint": "a-23icsyx4E=", 1267 "items": [ 1268 { 1269 "key": "not_sshKeys", 1270 "value": "" 1271 } 1272 ] 1273 } 1274 new_entry = "new_user:%s" % fake_ssh_key 1275 expected = { 1276 "kind": "compute#metadata", 1277 "fingerprint": "a-23icsyx4E=", 1278 "items": [ 1279 { 1280 "key": "not_sshKeys", 1281 "value": "" 1282 }, 1283 { 1284 "key": "sshKeys", 1285 "value": new_entry 1286 } 1287 ] 1288 } 1289 self.Patch(os.path, "exists", return_value=True) 1290 self.Patch(gcompute_client.ComputeClient, "WaitOnOperation") 1291 resource_mock = mock.MagicMock() 1292 self.compute_client.SetInstanceMetadata = mock.MagicMock( 1293 return_value=resource_mock) 1294 # Test the key item not exists in the metadata. 1295 self.compute_client.UpdateRsaInMetadata( 1296 "fake_zone", 1297 "fake_instance", 1298 fake_metadata_sshkeys_not_exist, 1299 new_entry) 1300 self.compute_client.SetInstanceMetadata.assert_called_with( 1301 "fake_zone", 1302 "fake_instance", 1303 expected) 1304 1305 # Test the key item exists in the metadata. 1306 fake_metadata_ssh_keys_exists = { 1307 "kind": "compute#metadata", 1308 "fingerprint": "a-23icsyx4E=", 1309 "items": [ 1310 { 1311 "key": "sshKeys", 1312 "value": "old_user:%s" % self.SSHKEY 1313 } 1314 ] 1315 } 1316 expected_ssh_keys_exists = { 1317 "kind": "compute#metadata", 1318 "fingerprint": "a-23icsyx4E=", 1319 "items": [ 1320 { 1321 "key": "sshKeys", 1322 "value": "old_user:%s\n%s" % (self.SSHKEY, new_entry) 1323 } 1324 ] 1325 } 1326 1327 self.compute_client.UpdateRsaInMetadata( 1328 "fake_zone", 1329 "fake_instance", 1330 fake_metadata_ssh_keys_exists, 1331 new_entry) 1332 self.compute_client.SetInstanceMetadata.assert_called_with( 1333 "fake_zone", 1334 "fake_instance", 1335 expected_ssh_keys_exists) 1336 1337 def testAddSshRsaToInstance(self): 1338 """Test add ssh rsa key to instance.""" 1339 fake_user = "fake_user" 1340 instance_metadata_key_not_exist = { 1341 "metadata": { 1342 "kind": "compute#metadata", 1343 "fingerprint": "a-23icsyx4E=", 1344 "items": [ 1345 { 1346 "key": "sshKeys", 1347 "value": "" 1348 } 1349 ] 1350 } 1351 } 1352 instance_metadata_key_exist = { 1353 "metadata": { 1354 "kind": "compute#metadata", 1355 "fingerprint": "a-23icsyx4E=", 1356 "items": [ 1357 { 1358 "key": "sshKeys", 1359 "value": "%s:%s" % (fake_user, self.SSHKEY) 1360 } 1361 ] 1362 } 1363 } 1364 expected = { 1365 "kind": "compute#metadata", 1366 "fingerprint": "a-23icsyx4E=", 1367 "items": [ 1368 { 1369 "key": "sshKeys", 1370 "value": "%s:%s" % (fake_user, self.SSHKEY) 1371 } 1372 ] 1373 } 1374 1375 self.Patch(os.path, "exists", return_value=True) 1376 m = mock.mock_open(read_data=self.SSHKEY) 1377 self.Patch(gcompute_client.ComputeClient, "WaitOnOperation") 1378 self.Patch(gcompute_client.ComputeClient, "GetZoneByInstance", 1379 return_value="fake_zone") 1380 resource_mock = mock.MagicMock() 1381 self.compute_client._service.instances = mock.MagicMock( 1382 return_value=resource_mock) 1383 resource_mock.setMetadata = mock.MagicMock() 1384 1385 # Test the key not exists in the metadata. 1386 self.Patch( 1387 gcompute_client.ComputeClient, "GetInstance", 1388 return_value=instance_metadata_key_not_exist) 1389 with mock.patch.object(six.moves.builtins, "open", m): 1390 self.compute_client.AddSshRsaInstanceMetadata( 1391 fake_user, 1392 "/path/to/test_rsa.pub", 1393 "fake_instance") 1394 resource_mock.setMetadata.assert_called_with( 1395 project=PROJECT, 1396 zone="fake_zone", 1397 instance="fake_instance", 1398 body=expected) 1399 1400 # Test the key already exists in the metadata. 1401 resource_mock.setMetadata.call_count = 0 1402 self.Patch( 1403 gcompute_client.ComputeClient, "GetInstance", 1404 return_value=instance_metadata_key_exist) 1405 with mock.patch.object(six.moves.builtins, "open", m): 1406 self.compute_client.AddSshRsaInstanceMetadata( 1407 fake_user, 1408 "/path/to/test_rsa.pub", 1409 "fake_instance") 1410 resource_mock.setMetadata.assert_not_called() 1411 1412 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 1413 def testDeleteDisks(self, mock_wait): 1414 """Test DeleteDisks.""" 1415 self._SetupBatchHttpRequestMock() 1416 fake_disks = ["fake_disk_1", "fake_disk_2"] 1417 mock_api = mock.MagicMock() 1418 resource_mock = mock.MagicMock() 1419 self.compute_client._service.disks = mock.MagicMock( 1420 return_value=resource_mock) 1421 resource_mock.delete = mock.MagicMock(return_value=mock_api) 1422 # Call the API. 1423 deleted, failed, error_msgs = self.compute_client.DeleteDisks( 1424 fake_disks, zone=self.ZONE) 1425 # Verify 1426 calls = [ 1427 mock.call(project=PROJECT, disk="fake_disk_1", zone=self.ZONE), 1428 mock.call(project=PROJECT, disk="fake_disk_2", zone=self.ZONE) 1429 ] 1430 resource_mock.delete.assert_has_calls(calls, any_order=True) 1431 self.assertEqual(mock_wait.call_count, 2) 1432 self.assertEqual(error_msgs, []) 1433 self.assertEqual(failed, []) 1434 self.assertEqual(set(deleted), set(fake_disks)) 1435 1436 def testRetryOnFingerPrintError(self): 1437 """Test RetryOnFingerPrintError.""" 1438 @utils.RetryOnException(gcompute_client._IsFingerPrintError, 10) 1439 def Raise412(sentinel): 1440 """Raise 412 HTTP exception.""" 1441 if not sentinel.hitFingerPrintConflict.called: 1442 sentinel.hitFingerPrintConflict() 1443 raise errors.HttpError(412, "resource labels have changed") 1444 return "Passed" 1445 1446 sentinel = mock.MagicMock() 1447 result = Raise412(sentinel) 1448 self.assertEqual(1, sentinel.hitFingerPrintConflict.call_count) 1449 self.assertEqual("Passed", result) 1450 1451 def testCheckAccess(self): 1452 """Test CheckAccess.""" 1453 # Checking non-403 should raise error 1454 error = errors.HttpError(503, "fake retriable error.") 1455 self.Patch( 1456 gcompute_client.ComputeClient, "Execute", 1457 side_effect=error) 1458 1459 with self.assertRaises(errors.HttpError): 1460 self.compute_client.CheckAccess() 1461 1462 # Checking 403 should return False 1463 error = errors.HttpError(403, "fake retriable error.") 1464 self.Patch( 1465 gcompute_client.ComputeClient, "Execute", 1466 side_effect=error) 1467 self.assertFalse(self.compute_client.CheckAccess()) 1468 1469 def testEnoughMetricsInZone(self): 1470 """Test EnoughMetricsInZone.""" 1471 region_info_enough_quota = { 1472 "items": [{ 1473 "name": "asia-east1", 1474 "quotas": [{ 1475 "usage": 50, 1476 "metric": "CPUS", 1477 "limit": 100 1478 }, { 1479 "usage": 640, 1480 "metric": "DISKS_TOTAL_GB", 1481 "limit": 10240 1482 }, { 1483 "usage": 20, 1484 "metric": "IN_USE_ADDRESSES", 1485 "limit": 100 1486 }] 1487 }] 1488 } 1489 self.Patch( 1490 gcompute_client.ComputeClient, "GetRegionInfo", 1491 return_value=region_info_enough_quota) 1492 self.assertTrue(self.compute_client.EnoughMetricsInZone("asia-east1-b")) 1493 self.assertFalse(self.compute_client.EnoughMetricsInZone("fake_zone")) 1494 1495 region_info_not_enough_quota = { 1496 "items": [{ 1497 "name": "asia-east1", 1498 "quotas": [{ 1499 "usage": 100, 1500 "metric": "CPUS", 1501 "limit": 100 1502 }, { 1503 "usage": 640, 1504 "metric": "DISKS_TOTAL_GB", 1505 "limit": 10240 1506 }, { 1507 "usage": 20, 1508 "metric": "IN_USE_ADDRESSES", 1509 "limit": 100 1510 }] 1511 }] 1512 } 1513 self.Patch( 1514 gcompute_client.ComputeClient, "GetRegionInfo", 1515 return_value=region_info_not_enough_quota) 1516 self.assertFalse(self.compute_client.EnoughMetricsInZone("asia-east1-b")) 1517 1518 def testGetDisk(self): 1519 """Test GetDisk.""" 1520 resource_mock = mock.MagicMock() 1521 mock_api = mock.MagicMock() 1522 self.compute_client._service.disks = mock.MagicMock( 1523 return_value=resource_mock) 1524 resource_mock.get = mock.MagicMock(return_value=mock_api) 1525 mock_api.execute = mock.MagicMock(return_value={"name": self.DISK}) 1526 result = self.compute_client.GetDisk(self.DISK, self.ZONE) 1527 self.assertEqual(result, {"name": self.DISK}) 1528 resource_mock.get.assert_called_with(project=PROJECT, 1529 zone=self.ZONE, 1530 disk=self.DISK) 1531 self.assertTrue(self.compute_client.CheckDiskExists(self.DISK, self.ZONE)) 1532 1533 1534if __name__ == "__main__": 1535 unittest.main() 1536