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 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().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 self.assertRaisesRegex( 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 self.assertRaisesRegex( 345 errors.DriverError, 346 "Expected fake error", 347 self.compute_client.CreateImage, 348 image_name=self.IMAGE, 349 source_uri=GS_IMAGE_SOURCE_URI) 350 resource_mock.insert.assert_called_with( 351 project=PROJECT, body=expected_body) 352 mock_wait.assert_called_with( 353 operation=mock.ANY, 354 operation_scope=gcompute_client.OperationScope.GLOBAL) 355 mock_check.assert_called_with(self.IMAGE) 356 mock_delete.assert_called_with(self.IMAGE) 357 358 def testCheckImageExistsTrue(self): 359 """Test CheckImageExists return True.""" 360 resource_mock = mock.MagicMock() 361 mock_api = mock.MagicMock() 362 self.compute_client._service.images = mock.MagicMock( 363 return_value=resource_mock) 364 resource_mock.get = mock.MagicMock(return_value=mock_api) 365 mock_api.execute = mock.MagicMock(return_value={"name": self.IMAGE}) 366 self.assertTrue(self.compute_client.CheckImageExists(self.IMAGE)) 367 368 def testCheckImageExistsFalse(self): 369 """Test CheckImageExists return False.""" 370 resource_mock = mock.MagicMock() 371 mock_api = mock.MagicMock() 372 self.compute_client._service.images = mock.MagicMock( 373 return_value=resource_mock) 374 resource_mock.get = mock.MagicMock(return_value=mock_api) 375 mock_api.execute = mock.MagicMock( 376 side_effect=errors.ResourceNotFoundError(404, "no image")) 377 self.assertFalse(self.compute_client.CheckImageExists(self.IMAGE)) 378 379 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 380 def testDeleteImage(self, mock_wait): 381 """Test DeleteImage.""" 382 resource_mock = mock.MagicMock() 383 self.compute_client._service.images = mock.MagicMock( 384 return_value=resource_mock) 385 resource_mock.delete = mock.MagicMock() 386 self.compute_client.DeleteImage(self.IMAGE) 387 resource_mock.delete.assert_called_with( 388 project=PROJECT, image=self.IMAGE) 389 self.assertTrue(mock_wait.called) 390 391 def _SetupBatchHttpRequestMock(self): 392 """Setup BatchHttpRequest mock.""" 393 requests = {} 394 395 def _Add(request, callback, request_id): 396 requests[request_id] = (request, callback) 397 398 def _Execute(): 399 for rid in requests: 400 _, callback = requests[rid] 401 callback( 402 request_id=rid, response=mock.MagicMock(), exception=None) 403 mock_batch = mock.MagicMock() 404 mock_batch.add = _Add 405 mock_batch.execute = _Execute 406 self.Patch(self.compute_client._service, 407 "new_batch_http_request", 408 return_value=mock_batch) 409 410 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 411 def testDeleteImages(self, mock_wait): 412 """Test DeleteImages.""" 413 self._SetupBatchHttpRequestMock() 414 fake_images = ["fake_image_1", "fake_image_2"] 415 mock_api = mock.MagicMock() 416 resource_mock = mock.MagicMock() 417 self.compute_client._service.images = mock.MagicMock( 418 return_value=resource_mock) 419 resource_mock.delete = mock.MagicMock(return_value=mock_api) 420 # Call the API. 421 deleted, failed, error_msgs = self.compute_client.DeleteImages( 422 fake_images) 423 # Verify 424 calls = [ 425 mock.call(project=PROJECT, image="fake_image_1"), 426 mock.call(project=PROJECT, image="fake_image_2") 427 ] 428 resource_mock.delete.assert_has_calls(calls, any_order=True) 429 self.assertEqual(mock_wait.call_count, 2) 430 self.assertEqual(error_msgs, []) 431 self.assertEqual(failed, []) 432 self.assertEqual(set(deleted), set(fake_images)) 433 434 def testListImages(self): 435 """Test ListImages.""" 436 fake_token = "fake_next_page_token" 437 image_1 = "image_1" 438 image_2 = "image_2" 439 response_1 = {"items": [image_1], "nextPageToken": fake_token} 440 response_2 = {"items": [image_2]} 441 self.Patch( 442 gcompute_client.ComputeClient, 443 "Execute", 444 side_effect=[response_1, response_2]) 445 resource_mock = mock.MagicMock() 446 self.compute_client._service.images = mock.MagicMock( 447 return_value=resource_mock) 448 resource_mock.list = mock.MagicMock() 449 images = self.compute_client.ListImages() 450 calls = [ 451 mock.call(project=PROJECT, filter=None, pageToken=None), 452 mock.call(project=PROJECT, filter=None, pageToken=fake_token) 453 ] 454 resource_mock.list.assert_has_calls(calls) 455 self.assertEqual(images, [image_1, image_2]) 456 457 def testListImagesFromExternalProject(self): 458 """Test ListImages which accepts different project.""" 459 image = "image_1" 460 response = {"items": [image]} 461 self.Patch(gcompute_client.ComputeClient, "Execute", side_effect=[response]) 462 resource_mock = mock.MagicMock() 463 self.compute_client._service.images = mock.MagicMock( 464 return_value=resource_mock) 465 resource_mock.list = mock.MagicMock() 466 images = self.compute_client.ListImages( 467 image_project="fake-project-2") 468 calls = [ 469 mock.call(project="fake-project-2", filter=None, pageToken=None)] 470 resource_mock.list.assert_has_calls(calls) 471 self.assertEqual(images, [image]) 472 473 def testGetInstance(self): 474 """Test GetInstance.""" 475 resource_mock = mock.MagicMock() 476 mock_api = mock.MagicMock() 477 self.compute_client._service.instances = mock.MagicMock( 478 return_value=resource_mock) 479 resource_mock.get = mock.MagicMock(return_value=mock_api) 480 mock_api.execute = mock.MagicMock(return_value={"name": self.INSTANCE}) 481 result = self.compute_client.GetInstance(self.INSTANCE, self.ZONE) 482 self.assertEqual(result, {"name": self.INSTANCE}) 483 resource_mock.get.assert_called_with( 484 project=PROJECT, zone=self.ZONE, instance=self.INSTANCE) 485 486 def testListInstances(self): 487 """Test ListInstances.""" 488 instance_1 = "instance_1" 489 instance_2 = "instance_2" 490 response = {"items": {'zones/fake_zone': {"instances": [instance_1, instance_2]}}} 491 self.Patch( 492 gcompute_client.ComputeClient, 493 "Execute", 494 side_effect=[response]) 495 resource_mock = mock.MagicMock() 496 self.compute_client._service.instances = mock.MagicMock( 497 return_value=resource_mock) 498 resource_mock.aggregatedList = mock.MagicMock() 499 instances = self.compute_client.ListInstances() 500 calls = [ 501 mock.call( 502 project=PROJECT, 503 filter=None, 504 pageToken=None), 505 ] 506 resource_mock.aggregatedList.assert_has_calls(calls) 507 self.assertEqual(instances, [instance_1, instance_2]) 508 509 def testGetZoneByInstance(self): 510 """Test GetZoneByInstance.""" 511 instance_1 = "instance_1" 512 response = {"items": {'zones/fake_zone': {"instances": [instance_1]}}} 513 self.Patch( 514 gcompute_client.ComputeClient, 515 "Execute", 516 side_effect=[response]) 517 expected_zone = "fake_zone" 518 self.assertEqual(self.compute_client.GetZoneByInstance(instance_1), 519 expected_zone) 520 521 # Test unable to find 'zone' from instance name. 522 response = {"items": {'zones/fake_zone': {"warning": "No instances."}}} 523 self.Patch( 524 gcompute_client.ComputeClient, 525 "Execute", 526 side_effect=[response]) 527 with self.assertRaises(errors.GetGceZoneError): 528 self.compute_client.GetZoneByInstance(instance_1) 529 530 def testGetZonesByInstances(self): 531 """Test GetZonesByInstances.""" 532 instances = ["instance_1", "instance_2"] 533 # Test instances in the same zone. 534 self.Patch( 535 gcompute_client.ComputeClient, 536 "GetZoneByInstance", 537 side_effect=["zone_1", "zone_1"]) 538 expected_result = {"zone_1": ["instance_1", "instance_2"]} 539 self.assertEqual(self.compute_client.GetZonesByInstances(instances), 540 expected_result) 541 542 # Test instances in different zones. 543 self.Patch( 544 gcompute_client.ComputeClient, 545 "GetZoneByInstance", 546 side_effect=["zone_1", "zone_2"]) 547 expected_result = {"zone_1": ["instance_1"], 548 "zone_2": ["instance_2"]} 549 self.assertEqual(self.compute_client.GetZonesByInstances(instances), 550 expected_result) 551 552 @mock.patch.object(gcompute_client.ComputeClient, "GetImage") 553 @mock.patch.object(gcompute_client.ComputeClient, "GetNetworkUrl") 554 @mock.patch.object(gcompute_client.ComputeClient, "GetSubnetworkUrl") 555 @mock.patch.object(gcompute_client.ComputeClient, "GetMachineType") 556 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 557 @mock.patch("getpass.getuser", return_value="fake_user") 558 def testCreateInstance(self, _get_user, mock_wait, mock_get_mach_type, 559 mock_get_subnetwork_url, mock_get_network_url, 560 mock_get_image): 561 """Test CreateInstance.""" 562 mock_get_mach_type.return_value = {"selfLink": self.MACHINE_TYPE_URL} 563 mock_get_network_url.return_value = self.NETWORK_URL 564 mock_get_subnetwork_url.return_value = self.SUBNETWORK_URL 565 mock_get_image.return_value = {"selfLink": self.IMAGE_URL} 566 resource_mock = mock.MagicMock() 567 self.compute_client._service.instances = mock.MagicMock( 568 return_value=resource_mock) 569 resource_mock.insert = mock.MagicMock() 570 self.Patch( 571 self.compute_client, 572 "_GetExtraDiskArgs", 573 return_value=[{"fake_extra_arg": "fake_extra_value"}]) 574 extra_disk_name = "gce-x86-userdebug-2345-abcd-data" 575 expected_disk_args = [self._disk_args] 576 expected_disk_args.extend([{"fake_extra_arg": "fake_extra_value"}]) 577 expected_scope = [] 578 expected_scope.extend(self.compute_client.DEFAULT_INSTANCE_SCOPE) 579 expected_scope.extend(self.EXTRA_SCOPES) 580 581 expected_body = { 582 "machineType": self.MACHINE_TYPE_URL, 583 "name": self.INSTANCE, 584 "networkInterfaces": [ 585 { 586 "network": self.NETWORK_URL, 587 "subnetwork": self.SUBNETWORK_URL, 588 "accessConfigs": [ 589 {"name": "External NAT", 590 "type": "ONE_TO_ONE_NAT"} 591 ], 592 } 593 ], 594 "disks": expected_disk_args, 595 "serviceAccounts": [ 596 {"email": "default", 597 "scopes": expected_scope} 598 ], 599 "metadata": { 600 "items": [{"key": self.METADATA[0], 601 "value": self.METADATA[1]}], 602 }, 603 "labels":{constants.LABEL_CREATE_BY: "fake_user"}, 604 "enableVtpm": True, 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 "enableVtpm": True, 683 } 684 685 self.compute_client.CreateInstance( 686 instance=self.INSTANCE, 687 image_name=self.IMAGE, 688 machine_type=self.MACHINE_TYPE, 689 metadata={self.METADATA[0]: self.METADATA[1]}, 690 network=self.NETWORK, 691 zone=self.ZONE, 692 extra_disk_name=extra_disk_name, 693 tags=["https-server"], 694 extra_scopes=self.EXTRA_SCOPES) 695 696 resource_mock.insert.assert_called_with( 697 project=PROJECT, zone=self.ZONE, body=expected_body) 698 mock_wait.assert_called_with( 699 mock.ANY, 700 operation_scope=gcompute_client.OperationScope.ZONE, 701 scope_name=self.ZONE) 702 703 @mock.patch.object(gcompute_client.ComputeClient, "GetAcceleratorUrl") 704 @mock.patch.object(gcompute_client.ComputeClient, "GetImage") 705 @mock.patch.object(gcompute_client.ComputeClient, "GetNetworkUrl") 706 @mock.patch.object(gcompute_client.ComputeClient, "GetSubnetworkUrl") 707 @mock.patch.object(gcompute_client.ComputeClient, "GetMachineType") 708 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 709 @mock.patch("getpass.getuser", return_value="fake_user") 710 def testCreateInstanceWithGpu(self, _get_user, mock_wait, mock_get_mach, 711 mock_get_subnetwork, mock_get_network, 712 mock_get_image, mock_get_accel): 713 """Test CreateInstance with a GPU parameter not set to None.""" 714 mock_get_mach.return_value = {"selfLink": self.MACHINE_TYPE_URL} 715 mock_get_network.return_value = self.NETWORK_URL 716 mock_get_subnetwork.return_value = self.SUBNETWORK_URL 717 mock_get_accel.return_value = self.ACCELERATOR_URL 718 mock_get_image.return_value = {"selfLink": self.IMAGE_URL} 719 720 resource_mock = mock.MagicMock() 721 self.compute_client._service.instances = mock.MagicMock( 722 return_value=resource_mock) 723 resource_mock.insert = mock.MagicMock() 724 725 expected_body = { 726 "machineType": 727 self.MACHINE_TYPE_URL, 728 "name": 729 self.INSTANCE, 730 "networkInterfaces": [{ 731 "network": self.NETWORK_URL, 732 "subnetwork": self.SUBNETWORK_URL, 733 "accessConfigs": [{ 734 "name": "External NAT", 735 "type": "ONE_TO_ONE_NAT" 736 }], 737 }], 738 "disks": [self._disk_args], 739 "serviceAccounts": [{ 740 "email": "default", 741 "scopes": self.compute_client.DEFAULT_INSTANCE_SCOPE 742 }], 743 "scheduling": { 744 "onHostMaintenance": "terminate" 745 }, 746 "guestAccelerators": [{ 747 "acceleratorCount": 1, 748 "acceleratorType": "http://speedy-gpu" 749 }], 750 "metadata": { 751 "items": [{ 752 "key": self.METADATA[0], 753 "value": self.METADATA[1] 754 }], 755 }, 756 "labels":{'created_by': "fake_user"}, 757 "enableVtpm": True, 758 } 759 760 self.compute_client.CreateInstance( 761 instance=self.INSTANCE, 762 image_name=self.IMAGE, 763 machine_type=self.MACHINE_TYPE, 764 metadata={self.METADATA[0]: self.METADATA[1]}, 765 network=self.NETWORK, 766 zone=self.ZONE, 767 gpu=self.GPU, 768 extra_scopes=None) 769 770 resource_mock.insert.assert_called_with( 771 project=PROJECT, zone=self.ZONE, body=expected_body) 772 mock_wait.assert_called_with( 773 mock.ANY, operation_scope=gcompute_client.OperationScope.ZONE, 774 scope_name=self.ZONE) 775 776 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 777 def testDeleteInstance(self, mock_wait): 778 """Test DeleteInstance.""" 779 resource_mock = mock.MagicMock() 780 self.compute_client._service.instances = mock.MagicMock( 781 return_value=resource_mock) 782 resource_mock.delete = mock.MagicMock() 783 self.compute_client.DeleteInstance( 784 instance=self.INSTANCE, zone=self.ZONE) 785 resource_mock.delete.assert_called_with( 786 project=PROJECT, zone=self.ZONE, instance=self.INSTANCE) 787 mock_wait.assert_called_with( 788 mock.ANY, 789 operation_scope=gcompute_client.OperationScope.ZONE, 790 scope_name=self.ZONE) 791 792 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 793 def testDeleteInstances(self, mock_wait): 794 """Test DeleteInstances.""" 795 self._SetupBatchHttpRequestMock() 796 fake_instances = ["fake_instance_1", "fake_instance_2"] 797 mock_api = mock.MagicMock() 798 resource_mock = mock.MagicMock() 799 self.compute_client._service.instances = mock.MagicMock( 800 return_value=resource_mock) 801 resource_mock.delete = mock.MagicMock(return_value=mock_api) 802 deleted, failed, error_msgs = self.compute_client.DeleteInstances( 803 fake_instances, self.ZONE) 804 calls = [ 805 mock.call( 806 project=PROJECT, 807 instance="fake_instance_1", 808 zone=self.ZONE), 809 mock.call( 810 project=PROJECT, 811 instance="fake_instance_2", 812 zone=self.ZONE) 813 ] 814 resource_mock.delete.assert_has_calls(calls, any_order=True) 815 self.assertEqual(mock_wait.call_count, 2) 816 self.assertEqual(error_msgs, []) 817 self.assertEqual(failed, []) 818 self.assertEqual(set(deleted), set(fake_instances)) 819 820 def testCreateDiskWithProject(self): 821 """Test CreateDisk with images using a set project.""" 822 source_project = "fake-image-project" 823 expected_project_to_use = "fake-image-project" 824 mock_wait = self.Patch(gcompute_client.ComputeClient, "WaitOnOperation") 825 resource_mock = mock.MagicMock() 826 self.compute_client._service.disks = mock.MagicMock( 827 return_value=resource_mock) 828 resource_mock.insert = mock.MagicMock() 829 self.compute_client.CreateDisk( 830 "fake_disk", "fake_image", 10, self.ZONE, source_project=source_project) 831 resource_mock.insert.assert_called_with( 832 project=PROJECT, 833 zone=self.ZONE, 834 sourceImage="projects/%s/global/images/fake_image" % 835 expected_project_to_use, 836 body={ 837 "name": 838 "fake_disk", 839 "sizeGb": 840 10, 841 "type": 842 "projects/%s/zones/%s/diskTypes/pd-standard" % (PROJECT, 843 self.ZONE) 844 }) 845 self.assertTrue(mock_wait.called) 846 847 def testCreateDiskWithNoSourceProject(self): 848 """Test CreateDisk with images with no set project.""" 849 source_project = None 850 expected_project_to_use = PROJECT 851 mock_wait = self.Patch(gcompute_client.ComputeClient, "WaitOnOperation") 852 resource_mock = mock.MagicMock() 853 self.compute_client._service.disks = mock.MagicMock( 854 return_value=resource_mock) 855 resource_mock.insert = mock.MagicMock() 856 self.compute_client.CreateDisk( 857 "fake_disk", "fake_image", 10, self.ZONE, source_project=source_project) 858 resource_mock.insert.assert_called_with( 859 project=PROJECT, 860 zone=self.ZONE, 861 sourceImage="projects/%s/global/images/fake_image" % 862 expected_project_to_use, 863 body={ 864 "name": 865 "fake_disk", 866 "sizeGb": 867 10, 868 "type": 869 "projects/%s/zones/%s/diskTypes/pd-standard" % (PROJECT, 870 self.ZONE) 871 }) 872 self.assertTrue(mock_wait.called) 873 874 def testCreateDiskWithTypeStandard(self): 875 """Test CreateDisk with images using standard.""" 876 disk_type = gcompute_client.PersistentDiskType.STANDARD 877 expected_disk_type_string = "pd-standard" 878 mock_wait = self.Patch(gcompute_client.ComputeClient, "WaitOnOperation") 879 resource_mock = mock.MagicMock() 880 self.compute_client._service.disks = mock.MagicMock( 881 return_value=resource_mock) 882 resource_mock.insert = mock.MagicMock() 883 self.compute_client.CreateDisk( 884 "fake_disk", 885 "fake_image", 886 10, 887 self.ZONE, 888 source_project="fake-project", 889 disk_type=disk_type) 890 resource_mock.insert.assert_called_with( 891 project=PROJECT, 892 zone=self.ZONE, 893 sourceImage="projects/%s/global/images/fake_image" % "fake-project", 894 body={ 895 "name": 896 "fake_disk", 897 "sizeGb": 898 10, 899 "type": 900 "projects/%s/zones/%s/diskTypes/%s" % 901 (PROJECT, self.ZONE, expected_disk_type_string) 902 }) 903 self.assertTrue(mock_wait.called) 904 905 def testCreateDiskWithTypeSSD(self): 906 """Test CreateDisk with images using standard.""" 907 disk_type = gcompute_client.PersistentDiskType.SSD 908 expected_disk_type_string = "pd-ssd" 909 mock_wait = self.Patch(gcompute_client.ComputeClient, "WaitOnOperation") 910 resource_mock = mock.MagicMock() 911 self.compute_client._service.disks = mock.MagicMock( 912 return_value=resource_mock) 913 resource_mock.insert = mock.MagicMock() 914 self.compute_client.CreateDisk( 915 "fake_disk", 916 "fake_image", 917 10, 918 self.ZONE, 919 source_project="fake-project", 920 disk_type=disk_type) 921 resource_mock.insert.assert_called_with( 922 project=PROJECT, 923 zone=self.ZONE, 924 sourceImage="projects/%s/global/images/fake_image" % "fake-project", 925 body={ 926 "name": 927 "fake_disk", 928 "sizeGb": 929 10, 930 "type": 931 "projects/%s/zones/%s/diskTypes/%s" % 932 (PROJECT, self.ZONE, expected_disk_type_string) 933 }) 934 self.assertTrue(mock_wait.called) 935 936 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 937 def testAttachDisk(self, mock_wait): 938 """Test AttachDisk.""" 939 resource_mock = mock.MagicMock() 940 self.compute_client._service.instances = mock.MagicMock( 941 return_value=resource_mock) 942 resource_mock.attachDisk = mock.MagicMock() 943 self.compute_client.AttachDisk( 944 "fake_instance_1", self.ZONE, deviceName="fake_disk", 945 source="fake-selfLink") 946 resource_mock.attachDisk.assert_called_with( 947 project=PROJECT, 948 zone=self.ZONE, 949 instance="fake_instance_1", 950 body={ 951 "deviceName": "fake_disk", 952 "source": "fake-selfLink" 953 }) 954 self.assertTrue(mock_wait.called) 955 956 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 957 def testDetachDisk(self, mock_wait): 958 """Test DetachDisk.""" 959 resource_mock = mock.MagicMock() 960 self.compute_client._service.instances = mock.MagicMock( 961 return_value=resource_mock) 962 resource_mock.detachDisk = mock.MagicMock() 963 self.compute_client.DetachDisk("fake_instance_1", self.ZONE, "fake_disk") 964 resource_mock.detachDisk.assert_called_with( 965 project=PROJECT, 966 zone=self.ZONE, 967 instance="fake_instance_1", 968 deviceName="fake_disk") 969 self.assertTrue(mock_wait.called) 970 971 @mock.patch.object(gcompute_client.ComputeClient, "GetAcceleratorUrl") 972 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 973 def testAttachAccelerator(self, mock_wait, mock_get_accel): 974 """Test AttachAccelerator.""" 975 mock_get_accel.return_value = self.ACCELERATOR_URL 976 resource_mock = mock.MagicMock() 977 self.compute_client._service.instances = mock.MagicMock( 978 return_value=resource_mock) 979 resource_mock.attachAccelerator = mock.MagicMock() 980 self.compute_client.AttachAccelerator("fake_instance_1", self.ZONE, 1, 981 "nvidia-tesla-k80") 982 resource_mock.setMachineResources.assert_called_with( 983 project=PROJECT, 984 zone=self.ZONE, 985 instance="fake_instance_1", 986 body={ 987 "guestAccelerators": [{ 988 "acceleratorType": self.ACCELERATOR_URL, 989 "acceleratorCount": 1 990 }] 991 }) 992 self.assertTrue(mock_wait.called) 993 994 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 995 def testBatchExecuteOnInstances(self, mock_wait): 996 """Test BatchExecuteOnInstances.""" 997 self._SetupBatchHttpRequestMock() 998 action = mock.MagicMock(return_value=mock.MagicMock()) 999 fake_instances = ["fake_instance_1", "fake_instance_2"] 1000 done, failed, error_msgs = self.compute_client._BatchExecuteOnInstances( 1001 fake_instances, self.ZONE, action) 1002 calls = [mock.call(instance="fake_instance_1"), 1003 mock.call(instance="fake_instance_2")] 1004 action.assert_has_calls(calls, any_order=True) 1005 self.assertEqual(mock_wait.call_count, 2) 1006 self.assertEqual(set(done), set(fake_instances)) 1007 self.assertEqual(error_msgs, []) 1008 self.assertEqual(failed, []) 1009 1010 @mock.patch.object(gcompute_client.ComputeClient, "WaitOnOperation") 1011 def testResetInstance(self, mock_wait): 1012 """Test ResetInstance.""" 1013 resource_mock = mock.MagicMock() 1014 self.compute_client._service.instances = mock.MagicMock( 1015 return_value=resource_mock) 1016 resource_mock.reset = mock.MagicMock() 1017 self.compute_client.ResetInstance( 1018 instance=self.INSTANCE, zone=self.ZONE) 1019 resource_mock.reset.assert_called_with( 1020 project=PROJECT, zone=self.ZONE, instance=self.INSTANCE) 1021 mock_wait.assert_called_with( 1022 mock.ANY, 1023 operation_scope=gcompute_client.OperationScope.ZONE, 1024 scope_name=self.ZONE) 1025 1026 def _CompareMachineSizeTestHelper(self, 1027 machine_info_1, 1028 machine_info_2, 1029 expected_result=None, 1030 expected_error_type=None): 1031 """Helper class for testing CompareMachineSize. 1032 1033 Args: 1034 machine_info_1: A dictionary representing the first machine size. 1035 machine_info_2: A dictionary representing the second machine size. 1036 expected_result: An integer, 0, 1 or -1, or None if not set. 1037 expected_error_type: An exception type, if set will check for exception. 1038 """ 1039 mock_get_mach_type = self.Patch( 1040 gcompute_client.ComputeClient, 1041 "GetMachineType", 1042 side_effect=[machine_info_1, machine_info_2]) 1043 if expected_error_type: 1044 self.assertRaises(expected_error_type, 1045 self.compute_client.CompareMachineSize, "name1", 1046 "name2", self.ZONE) 1047 else: 1048 result = self.compute_client.CompareMachineSize("name1", "name2", 1049 self.ZONE) 1050 self.assertEqual(result, expected_result) 1051 1052 mock_get_mach_type.assert_has_calls( 1053 [mock.call("name1", self.ZONE), mock.call("name2", self.ZONE)]) 1054 1055 def testCompareMachineSizeSmall(self): 1056 """Test CompareMachineSize where the first one is smaller.""" 1057 machine_info_1 = {"guestCpus": 10, "memoryMb": 100} 1058 machine_info_2 = {"guestCpus": 10, "memoryMb": 200} 1059 self._CompareMachineSizeTestHelper(machine_info_1, machine_info_2, -1) 1060 1061 def testCompareMachineSizeSmallSmallerOnSecond(self): 1062 """Test CompareMachineSize where the first one is smaller.""" 1063 machine_info_1 = {"guestCpus": 11, "memoryMb": 100} 1064 machine_info_2 = {"guestCpus": 10, "memoryMb": 200} 1065 self._CompareMachineSizeTestHelper(machine_info_1, machine_info_2, -1) 1066 1067 def testCompareMachineSizeLarge(self): 1068 """Test CompareMachineSize where the first one is larger.""" 1069 machine_info_1 = {"guestCpus": 11, "memoryMb": 200} 1070 machine_info_2 = {"guestCpus": 10, "memoryMb": 100} 1071 self._CompareMachineSizeTestHelper(machine_info_1, machine_info_2, 1) 1072 1073 def testCompareMachineSizeLargeWithEqualElement(self): 1074 """Test CompareMachineSize where the first one is larger.""" 1075 machine_info_1 = {"guestCpus": 10, "memoryMb": 200} 1076 machine_info_2 = {"guestCpus": 10, "memoryMb": 100} 1077 self._CompareMachineSizeTestHelper(machine_info_1, machine_info_2, 1) 1078 1079 def testCompareMachineSizeEqual(self): 1080 """Test CompareMachineSize where two machine sizes are equal.""" 1081 machine_info = {"guestCpus": 10, "memoryMb": 100} 1082 self._CompareMachineSizeTestHelper(machine_info, machine_info, 0) 1083 1084 def testCompareMachineSizeBadMetric(self): 1085 """Test CompareMachineSize with bad metric.""" 1086 machine_info = {"unknown_metric": 10, "memoryMb": 100} 1087 self._CompareMachineSizeTestHelper( 1088 machine_info, machine_info, expected_error_type=errors.DriverError) 1089 1090 def testGetMachineType(self): 1091 """Test GetMachineType.""" 1092 resource_mock = mock.MagicMock() 1093 mock_api = mock.MagicMock() 1094 self.compute_client._service.machineTypes = mock.MagicMock( 1095 return_value=resource_mock) 1096 resource_mock.get = mock.MagicMock(return_value=mock_api) 1097 mock_api.execute = mock.MagicMock( 1098 return_value={"name": self.MACHINE_TYPE}) 1099 result = self.compute_client.GetMachineType(self.MACHINE_TYPE, 1100 self.ZONE) 1101 self.assertEqual(result, {"name": self.MACHINE_TYPE}) 1102 resource_mock.get.assert_called_with( 1103 project=PROJECT, 1104 zone=self.ZONE, 1105 machineType=self.MACHINE_TYPE) 1106 1107 def _GetSerialPortOutputTestHelper(self, response): 1108 """Helper function for testing GetSerialPortOutput. 1109 1110 Args: 1111 response: A dictionary representing a fake response. 1112 """ 1113 resource_mock = mock.MagicMock() 1114 mock_api = mock.MagicMock() 1115 self.compute_client._service.instances = mock.MagicMock( 1116 return_value=resource_mock) 1117 resource_mock.getSerialPortOutput = mock.MagicMock( 1118 return_value=mock_api) 1119 mock_api.execute = mock.MagicMock(return_value=response) 1120 1121 if "contents" in response: 1122 result = self.compute_client.GetSerialPortOutput( 1123 instance=self.INSTANCE, zone=self.ZONE) 1124 self.assertEqual(result, "fake contents") 1125 else: 1126 self.assertRaisesRegex( 1127 errors.DriverError, 1128 "Malformed response.*", 1129 self.compute_client.GetSerialPortOutput, 1130 instance=self.INSTANCE, 1131 zone=self.ZONE) 1132 resource_mock.getSerialPortOutput.assert_called_with( 1133 project=PROJECT, 1134 zone=self.ZONE, 1135 instance=self.INSTANCE, 1136 port=1) 1137 1138 def testGetSerialPortOutput(self): 1139 """Test GetSerialPortOutput.""" 1140 response = {"contents": "fake contents"} 1141 self._GetSerialPortOutputTestHelper(response) 1142 1143 def testGetSerialPortOutputFail(self): 1144 """Test GetSerialPortOutputFail.""" 1145 response = {"malformed": "fake contents"} 1146 self._GetSerialPortOutputTestHelper(response) 1147 1148 def testGetInstanceNamesByIPs(self): 1149 """Test GetInstanceNamesByIPs.""" 1150 good_instance = { 1151 "name": "instance_1", 1152 "networkInterfaces": [ 1153 { 1154 "accessConfigs": [ 1155 {"natIP": "172.22.22.22"}, 1156 ], 1157 }, 1158 ], 1159 } 1160 bad_instance = {"name": "instance_2"} 1161 self.Patch( 1162 gcompute_client.ComputeClient, 1163 "ListInstances", 1164 return_value=[good_instance, bad_instance]) 1165 ip_name_map = self.compute_client.GetInstanceNamesByIPs( 1166 ips=["172.22.22.22", "172.22.22.23"]) 1167 self.assertEqual(ip_name_map, {"172.22.22.22": "instance_1", 1168 "172.22.22.23": None}) 1169 1170 def testRsaNotInMetadata(self): 1171 """Test rsa not in metadata.""" 1172 fake_user = "fake_user" 1173 fake_ssh_key = "fake_ssh" 1174 metadata = { 1175 "kind": "compute#metadata", 1176 "fingerprint": "a-23icsyx4E=", 1177 "items": [ 1178 { 1179 "key": "sshKeys", 1180 "value": "%s:%s" % (fake_user, self.SSHKEY) 1181 } 1182 ] 1183 } 1184 # Test rsa doesn't exist in metadata. 1185 new_entry = "%s:%s" % (fake_user, fake_ssh_key) 1186 self.assertEqual(True, gcompute_client.RsaNotInMetadata(metadata, new_entry)) 1187 1188 # Test rsa exists in metadata. 1189 exist_entry = "%s:%s" %(fake_user, self.SSHKEY) 1190 self.assertEqual(False, gcompute_client.RsaNotInMetadata(metadata, exist_entry)) 1191 1192 def testGetSshKeyFromMetadata(self): 1193 """Test get ssh key from metadata.""" 1194 fake_user = "fake_user" 1195 metadata_key_exist_value_is_empty = { 1196 "kind": "compute#metadata", 1197 "fingerprint": "a-23icsyx4E=", 1198 "items": [ 1199 { 1200 "key": "sshKeys", 1201 "value": "" 1202 } 1203 ] 1204 } 1205 metadata_key_exist = { 1206 "kind": "compute#metadata", 1207 "fingerprint": "a-23icsyx4E=", 1208 "items": [ 1209 { 1210 "key": "sshKeys", 1211 "value": "%s:%s" % (fake_user, self.SSHKEY) 1212 } 1213 ] 1214 } 1215 metadata_key_not_exist = { 1216 "kind": "compute#metadata", 1217 "fingerprint": "a-23icsyx4E=", 1218 "items": [ 1219 { 1220 } 1221 ] 1222 } 1223 expected_key_exist_value_is_empty = { 1224 "key": "sshKeys", 1225 "value": "" 1226 } 1227 expected_key_exist = { 1228 "key": "sshKeys", 1229 "value": "%s:%s" % (fake_user, self.SSHKEY) 1230 } 1231 self.assertEqual(expected_key_exist_value_is_empty, 1232 gcompute_client.GetSshKeyFromMetadata(metadata_key_exist_value_is_empty)) 1233 self.assertEqual(expected_key_exist, 1234 gcompute_client.GetSshKeyFromMetadata(metadata_key_exist)) 1235 self.assertEqual(None, 1236 gcompute_client.GetSshKeyFromMetadata(metadata_key_not_exist)) 1237 1238 1239 def testGetRsaKeyPathExistsFalse(self): 1240 """Test the rsa key path not exists.""" 1241 fake_ssh_rsa_path = "/path/to/test_rsa.pub" 1242 self.Patch(os.path, "exists", return_value=False) 1243 self.assertRaisesRegex(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("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("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("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 "usage": 20, 1480 "metric": "IN_USE_ADDRESSES", 1481 "limit": 100 1482 }] 1483 }] 1484 } 1485 self.Patch( 1486 gcompute_client.ComputeClient, "GetRegionInfo", 1487 return_value=region_info_enough_quota) 1488 self.assertTrue(self.compute_client.EnoughMetricsInZone("asia-east1-b")) 1489 self.assertFalse(self.compute_client.EnoughMetricsInZone("fake_zone")) 1490 1491 region_info_not_enough_quota = { 1492 "items": [{ 1493 "name": "asia-east1", 1494 "quotas": [{ 1495 "usage": 100, 1496 "metric": "CPUS", 1497 "limit": 100 1498 }, { 1499 "usage": 640, 1500 "metric": "DISKS_TOTAL_GB", 1501 "limit": 10240 1502 }, { 1503 "usage": 20, 1504 "metric": "IN_USE_ADDRESSES", 1505 "limit": 100 1506 }] 1507 }] 1508 } 1509 self.Patch( 1510 gcompute_client.ComputeClient, "GetRegionInfo", 1511 return_value=region_info_not_enough_quota) 1512 self.assertFalse(self.compute_client.EnoughMetricsInZone("asia-east1-b")) 1513 1514 def testGetDisk(self): 1515 """Test GetDisk.""" 1516 resource_mock = mock.MagicMock() 1517 mock_api = mock.MagicMock() 1518 self.compute_client._service.disks = mock.MagicMock( 1519 return_value=resource_mock) 1520 resource_mock.get = mock.MagicMock(return_value=mock_api) 1521 mock_api.execute = mock.MagicMock(return_value={"name": self.DISK}) 1522 result = self.compute_client.GetDisk(self.DISK, self.ZONE) 1523 self.assertEqual(result, {"name": self.DISK}) 1524 resource_mock.get.assert_called_with(project=PROJECT, 1525 zone=self.ZONE, 1526 disk=self.DISK) 1527 self.assertTrue(self.compute_client.CheckDiskExists(self.DISK, self.ZONE)) 1528 1529 1530if __name__ == "__main__": 1531 unittest.main() 1532