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