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