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
17"""Tests for acloud.internal.lib.gcompute_client."""
18
19import os
20
21import apiclient.http
22import mock
23
24import unittest
25from acloud.internal.lib import driver_test_lib
26from acloud.internal.lib import gcompute_client
27from acloud.internal.lib import utils
28from acloud.public import errors
29
30
31class ComputeClientTest(driver_test_lib.BaseDriverTest):
32    """Test ComputeClient."""
33
34    PROJECT = "fake-project"
35    INSTANCE = "fake-instance"
36    IMAGE = "fake-image"
37    IMAGE_URL = "http://fake-image-url"
38    GS_IMAGE_SOURCE_URI = "https://storage.googleapis.com/fake-bucket/fake.tar.gz"
39    MACHINE_TYPE = "fake-machine-type"
40    MACHINE_TYPE_URL = "http://fake-machine-type-url"
41    METADATA = ("metadata_key", "metadata_value")
42    NETWORK = "fake-network"
43    NETWORK_URL = "http://fake-network-url"
44    ZONE = "fake-zone"
45    REGION = "fake-region"
46    OPERATION_NAME = "fake-op"
47
48    def setUp(self):
49        """Set up test."""
50        super(ComputeClientTest, self).setUp()
51        self.Patch(gcompute_client.ComputeClient, "InitResourceHandle")
52        fake_cfg = mock.MagicMock()
53        fake_cfg.project = self.PROJECT
54        self.compute_client = gcompute_client.ComputeClient(fake_cfg,
55                                                            mock.MagicMock())
56        self.compute_client._service = mock.MagicMock()
57
58    def _SetupMocksForGetOperationStatus(self, mock_result, operation_scope):
59        """A helper class for setting up mocks for testGetOperationStatus*.
60
61    Args:
62      mock_result: The result to return by _GetOperationStatus.
63      operation_scope: A value of OperationScope.
64
65    Returns:
66      A mock for Resource object.
67    """
68        resource_mock = mock.MagicMock()
69        mock_api = mock.MagicMock()
70        if operation_scope == gcompute_client.OperationScope.GLOBAL:
71            self.compute_client._service.globalOperations = mock.MagicMock(
72                return_value=resource_mock)
73        elif operation_scope == gcompute_client.OperationScope.ZONE:
74            self.compute_client._service.zoneOperations = mock.MagicMock(
75                return_value=resource_mock)
76        elif operation_scope == gcompute_client.OperationScope.REGION:
77            self.compute_client._service.regionOperations = mock.MagicMock(
78                return_value=resource_mock)
79        resource_mock.get = mock.MagicMock(return_value=mock_api)
80        mock_api.execute = mock.MagicMock(return_value=mock_result)
81        return resource_mock
82
83    def testGetOperationStatusGlobal(self):
84        """Test _GetOperationStatus for global."""
85        resource_mock = self._SetupMocksForGetOperationStatus(
86            {"status": "GOOD"}, gcompute_client.OperationScope.GLOBAL)
87        status = self.compute_client._GetOperationStatus(
88            {"name": self.OPERATION_NAME},
89            gcompute_client.OperationScope.GLOBAL)
90        self.assertEqual(status, "GOOD")
91        resource_mock.get.assert_called_with(
92            project=self.PROJECT, operation=self.OPERATION_NAME)
93
94    def testGetOperationStatusZone(self):
95        """Test _GetOperationStatus for zone."""
96        resource_mock = self._SetupMocksForGetOperationStatus(
97            {"status": "GOOD"}, gcompute_client.OperationScope.ZONE)
98        status = self.compute_client._GetOperationStatus(
99            {"name": self.OPERATION_NAME}, gcompute_client.OperationScope.ZONE,
100            self.ZONE)
101        self.assertEqual(status, "GOOD")
102        resource_mock.get.assert_called_with(
103            project=self.PROJECT,
104            operation=self.OPERATION_NAME,
105            zone=self.ZONE)
106
107    def testGetOperationStatusRegion(self):
108        """Test _GetOperationStatus for region."""
109        resource_mock = self._SetupMocksForGetOperationStatus(
110            {"status": "GOOD"}, gcompute_client.OperationScope.REGION)
111        self.compute_client._GetOperationStatus(
112            {"name": self.OPERATION_NAME},
113            gcompute_client.OperationScope.REGION, self.REGION)
114        resource_mock.get.assert_called_with(
115            project=self.PROJECT,
116            operation=self.OPERATION_NAME,
117            region=self.REGION)
118
119    def testGetOperationStatusError(self):
120        """Test _GetOperationStatus failed."""
121        self._SetupMocksForGetOperationStatus(
122            {"error": {"errors": ["error1", "error2"]}},
123            gcompute_client.OperationScope.GLOBAL)
124        self.assertRaisesRegexp(errors.DriverError,
125                                "Get operation state failed.*error1.*error2",
126                                self.compute_client._GetOperationStatus,
127                                {"name": self.OPERATION_NAME},
128                                gcompute_client.OperationScope.GLOBAL)
129
130    def testWaitOnOperation(self):
131        """Test WaitOnOperation."""
132        mock_error = mock.MagicMock()
133        self.Patch(utils, "PollAndWait")
134        self.Patch(errors, "GceOperationTimeoutError", return_value=mock_error)
135        self.compute_client.WaitOnOperation(
136            operation={"name": self.OPERATION_NAME},
137            operation_scope=gcompute_client.OperationScope.REGION,
138            scope_name=self.REGION)
139        utils.PollAndWait.assert_called_with(
140            func=self.compute_client._GetOperationStatus,
141            expected_return="DONE",
142            timeout_exception=mock_error,
143            timeout_secs=self.compute_client.OPERATION_TIMEOUT_SECS,
144            sleep_interval_secs=self.compute_client.
145            OPERATION_POLL_INTERVAL_SECS,
146            operation={"name": self.OPERATION_NAME},
147            operation_scope=gcompute_client.OperationScope.REGION,
148            scope_name=self.REGION)
149
150    def testCreateImage(self):
151        """Test CreateImage."""
152        self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
153        resource_mock = mock.MagicMock()
154        self.compute_client._service.images = mock.MagicMock(
155            return_value=resource_mock)
156        resource_mock.insert = mock.MagicMock()
157
158        expected_body = {
159            "name": self.IMAGE,
160            "rawDisk": {
161                "source": self.GS_IMAGE_SOURCE_URI,
162            },
163        }
164        self.compute_client.CreateImage(
165            image_name=self.IMAGE, source_uri=self.GS_IMAGE_SOURCE_URI)
166        resource_mock.insert.assert_called_with(
167            project=self.PROJECT, body=expected_body)
168        self.compute_client.WaitOnOperation.assert_called_with(
169            operation=mock.ANY,
170            operation_scope=gcompute_client.OperationScope.GLOBAL)
171
172    def testCreateImageFail(self):
173        """Test CreateImage fails."""
174        self.Patch(
175            gcompute_client.ComputeClient,
176            "WaitOnOperation",
177            side_effect=errors.DriverError("Expected fake error"))
178        self.Patch(
179            gcompute_client.ComputeClient,
180            "CheckImageExists",
181            return_value=True)
182        self.Patch(gcompute_client.ComputeClient, "DeleteImage")
183
184        resource_mock = mock.MagicMock()
185        self.compute_client._service.images = mock.MagicMock(
186            return_value=resource_mock)
187        resource_mock.insert = mock.MagicMock()
188
189        expected_body = {
190            "name": self.IMAGE,
191            "rawDisk": {
192                "source": self.GS_IMAGE_SOURCE_URI,
193            },
194        }
195        self.assertRaisesRegexp(
196            errors.DriverError,
197            "Expected fake error",
198            self.compute_client.CreateImage,
199            image_name=self.IMAGE,
200            source_uri=self.GS_IMAGE_SOURCE_URI)
201        resource_mock.insert.assert_called_with(
202            project=self.PROJECT, body=expected_body)
203        self.compute_client.WaitOnOperation.assert_called_with(
204            operation=mock.ANY,
205            operation_scope=gcompute_client.OperationScope.GLOBAL)
206        self.compute_client.CheckImageExists.assert_called_with(self.IMAGE)
207        self.compute_client.DeleteImage.assert_called_with(self.IMAGE)
208
209    def testCheckImageExistsTrue(self):
210        """Test CheckImageExists return True."""
211        resource_mock = mock.MagicMock()
212        mock_api = mock.MagicMock()
213        self.compute_client._service.images = mock.MagicMock(
214            return_value=resource_mock)
215        resource_mock.get = mock.MagicMock(return_value=mock_api)
216        mock_api.execute = mock.MagicMock(return_value={"name": self.IMAGE})
217        self.assertTrue(self.compute_client.CheckImageExists(self.IMAGE))
218
219    def testCheckImageExistsFalse(self):
220        """Test CheckImageExists return False."""
221        resource_mock = mock.MagicMock()
222        mock_api = mock.MagicMock()
223        self.compute_client._service.images = mock.MagicMock(
224            return_value=resource_mock)
225        resource_mock.get = mock.MagicMock(return_value=mock_api)
226        mock_api.execute = mock.MagicMock(
227            side_effect=errors.ResourceNotFoundError(404, "no image"))
228        self.assertFalse(self.compute_client.CheckImageExists(self.IMAGE))
229
230    def testDeleteImage(self):
231        """Test DeleteImage."""
232        self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
233        resource_mock = mock.MagicMock()
234        self.compute_client._service.images = mock.MagicMock(
235            return_value=resource_mock)
236        resource_mock.delete = mock.MagicMock()
237        self.compute_client.DeleteImage(self.IMAGE)
238        resource_mock.delete.assert_called_with(
239            project=self.PROJECT, image=self.IMAGE)
240        self.assertTrue(self.compute_client.WaitOnOperation.called)
241
242    def _SetupBatchHttpRequestMock(self):
243        """Setup BatchHttpRequest mock."""
244        requests = {}
245
246        def _Add(request, callback, request_id):
247            requests[request_id] = (request, callback)
248
249        def _Execute():
250            for rid in requests:
251                _, callback = requests[rid]
252                callback(
253                    request_id=rid, response=mock.MagicMock(), exception=None)
254
255        mock_batch = mock.MagicMock()
256        mock_batch.add = _Add
257        mock_batch.execute = _Execute
258        self.Patch(apiclient.http, "BatchHttpRequest", return_value=mock_batch)
259
260    def testDeleteImages(self):
261        """Test DeleteImages."""
262        self._SetupBatchHttpRequestMock()
263        self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
264        fake_images = ["fake_image_1", "fake_image_2"]
265        mock_api = mock.MagicMock()
266        resource_mock = mock.MagicMock()
267        self.compute_client._service.images = mock.MagicMock(
268            return_value=resource_mock)
269        resource_mock.delete = mock.MagicMock(return_value=mock_api)
270        # Call the API.
271        deleted, failed, error_msgs = self.compute_client.DeleteImages(
272            fake_images)
273        # Verify
274        calls = [mock.call(
275            project=self.PROJECT, image="fake_image_1"), mock.call(
276                project=self.PROJECT, image="fake_image_2")]
277        resource_mock.delete.assert_has_calls(calls, any_order=True)
278        self.assertEqual(
279            gcompute_client.ComputeClient.WaitOnOperation.call_count, 2)
280        self.assertEqual(error_msgs, [])
281        self.assertEqual(failed, [])
282        self.assertEqual(set(deleted), set(fake_images))
283
284    def testListImages(self):
285        """Test ListImages."""
286        fake_token = "fake_next_page_token"
287        image_1 = "image_1"
288        image_2 = "image_2"
289        response_1 = {"items": [image_1], "nextPageToken": fake_token}
290        response_2 = {"items": [image_2]}
291        self.Patch(
292            gcompute_client.ComputeClient,
293            "Execute",
294            side_effect=[response_1, response_2])
295        resource_mock = mock.MagicMock()
296        self.compute_client._service.images = mock.MagicMock(
297            return_value=resource_mock)
298        resource_mock.list = mock.MagicMock()
299        images = self.compute_client.ListImages()
300        calls = [
301            mock.call(
302                project=self.PROJECT, filter=None, pageToken=None), mock.call(
303                    project=self.PROJECT, filter=None, pageToken=fake_token)
304        ]
305        resource_mock.list.assert_has_calls(calls)
306        self.assertEqual(images, [image_1, image_2])
307
308    def testGetInstance(self):
309        """Test GetInstance."""
310        resource_mock = mock.MagicMock()
311        mock_api = mock.MagicMock()
312        self.compute_client._service.instances = mock.MagicMock(
313            return_value=resource_mock)
314        resource_mock.get = mock.MagicMock(return_value=mock_api)
315        mock_api.execute = mock.MagicMock(return_value={"name": self.INSTANCE})
316        result = self.compute_client.GetInstance(self.INSTANCE, self.ZONE)
317        self.assertEqual(result, {"name": self.INSTANCE})
318        resource_mock.get.assert_called_with(
319            project=self.PROJECT, zone=self.ZONE, instance=self.INSTANCE)
320
321    def testListInstances(self):
322        """Test ListInstances."""
323        fake_token = "fake_next_page_token"
324        instance_1 = "instance_1"
325        instance_2 = "instance_2"
326        response_1 = {"items": [instance_1], "nextPageToken": fake_token}
327        response_2 = {"items": [instance_2]}
328        self.Patch(
329            gcompute_client.ComputeClient,
330            "Execute",
331            side_effect=[response_1, response_2])
332        resource_mock = mock.MagicMock()
333        self.compute_client._service.instances = mock.MagicMock(
334            return_value=resource_mock)
335        resource_mock.list = mock.MagicMock()
336        instances = self.compute_client.ListInstances(self.ZONE)
337        calls = [
338            mock.call(
339                project=self.PROJECT,
340                zone=self.ZONE,
341                filter=None,
342                pageToken=None),
343            mock.call(
344                project=self.PROJECT,
345                zone=self.ZONE,
346                filter=None,
347                pageToken=fake_token),
348        ]
349        resource_mock.list.assert_has_calls(calls)
350        self.assertEqual(instances, [instance_1, instance_2])
351
352    def testCreateInstance(self):
353        """Test CreateInstance."""
354        self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
355        self.Patch(
356            gcompute_client.ComputeClient,
357            "GetMachineType",
358            return_value={"selfLink": self.MACHINE_TYPE_URL})
359        self.Patch(
360            gcompute_client.ComputeClient,
361            "GetNetworkUrl",
362            return_value=self.NETWORK_URL)
363        self.Patch(
364            gcompute_client.ComputeClient,
365            "GetImage",
366            return_value={"selfLink": self.IMAGE_URL})
367
368        resource_mock = mock.MagicMock()
369        self.compute_client._service.instances = mock.MagicMock(
370            return_value=resource_mock)
371        resource_mock.insert = mock.MagicMock()
372
373        expected_body = {
374            "machineType": self.MACHINE_TYPE_URL,
375            "name": self.INSTANCE,
376            "networkInterfaces": [
377                {
378                    "network": self.NETWORK_URL,
379                    "accessConfigs": [
380                        {"name": "External NAT",
381                         "type": "ONE_TO_ONE_NAT"}
382                    ],
383                }
384            ],
385            "disks": [
386                {
387                    "type": "PERSISTENT",
388                    "boot": True,
389                    "mode": "READ_WRITE",
390                    "autoDelete": True,
391                    "initializeParams": {
392                        "diskName": self.INSTANCE,
393                        "sourceImage": self.IMAGE_URL,
394                    },
395                }
396            ],
397            "serviceAccounts": [
398                {"email": "default",
399                 "scopes": self.compute_client.DEFAULT_INSTANCE_SCOPE}
400            ],
401            "metadata": {
402                "items": [{"key": self.METADATA[0],
403                           "value": self.METADATA[1]}],
404            },
405        }
406
407        self.compute_client.CreateInstance(
408            instance=self.INSTANCE,
409            image_name=self.IMAGE,
410            machine_type=self.MACHINE_TYPE,
411            metadata={self.METADATA[0]: self.METADATA[1]},
412            network=self.NETWORK,
413            zone=self.ZONE)
414
415        resource_mock.insert.assert_called_with(
416            project=self.PROJECT, zone=self.ZONE, body=expected_body)
417        self.compute_client.WaitOnOperation.assert_called_with(
418            mock.ANY,
419            operation_scope=gcompute_client.OperationScope.ZONE,
420            scope_name=self.ZONE)
421
422    def testDeleteInstance(self):
423        """Test DeleteInstance."""
424        self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
425        resource_mock = mock.MagicMock()
426        self.compute_client._service.instances = mock.MagicMock(
427            return_value=resource_mock)
428        resource_mock.delete = mock.MagicMock()
429        self.compute_client.DeleteInstance(
430            instance=self.INSTANCE, zone=self.ZONE)
431        resource_mock.delete.assert_called_with(
432            project=self.PROJECT, zone=self.ZONE, instance=self.INSTANCE)
433        self.compute_client.WaitOnOperation.assert_called_with(
434            mock.ANY,
435            operation_scope=gcompute_client.OperationScope.ZONE,
436            scope_name=self.ZONE)
437
438    def testDeleteInstances(self):
439        """Test DeleteInstances."""
440        self._SetupBatchHttpRequestMock()
441        self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
442        fake_instances = ["fake_instance_1", "fake_instance_2"]
443        mock_api = mock.MagicMock()
444        resource_mock = mock.MagicMock()
445        self.compute_client._service.instances = mock.MagicMock(
446            return_value=resource_mock)
447        resource_mock.delete = mock.MagicMock(return_value=mock_api)
448        deleted, failed, error_msgs = self.compute_client.DeleteInstances(
449            fake_instances, self.ZONE)
450        calls = [
451            mock.call(
452                project=self.PROJECT,
453                instance="fake_instance_1",
454                zone=self.ZONE), mock.call(
455                    project=self.PROJECT,
456                    instance="fake_instance_2",
457                    zone=self.ZONE)
458        ]
459        resource_mock.delete.assert_has_calls(calls, any_order=True)
460        self.assertEqual(
461            gcompute_client.ComputeClient.WaitOnOperation.call_count, 2)
462        self.assertEqual(error_msgs, [])
463        self.assertEqual(failed, [])
464        self.assertEqual(set(deleted), set(fake_instances))
465
466    def testBatchExecuteOnInstances(self):
467        self._SetupBatchHttpRequestMock()
468        self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
469        action = mock.MagicMock(return_value=mock.MagicMock())
470        fake_instances = ["fake_instance_1", "fake_instance_2"]
471        done, failed, error_msgs = self.compute_client._BatchExecuteOnInstances(
472            fake_instances, self.ZONE, action)
473        calls = [mock.call(instance="fake_instance_1"),
474                 mock.call(instance="fake_instance_2")]
475        action.assert_has_calls(calls, any_order=True)
476        self.assertEqual(
477            gcompute_client.ComputeClient.WaitOnOperation.call_count, 2)
478        self.assertEqual(set(done), set(fake_instances))
479        self.assertEqual(error_msgs, [])
480        self.assertEqual(failed, [])
481
482    def testResetInstance(self):
483        """Test ResetInstance."""
484        self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
485        resource_mock = mock.MagicMock()
486        self.compute_client._service.instances = mock.MagicMock(
487            return_value=resource_mock)
488        resource_mock.reset = mock.MagicMock()
489        self.compute_client.ResetInstance(
490            instance=self.INSTANCE, zone=self.ZONE)
491        resource_mock.reset.assert_called_with(
492            project=self.PROJECT, zone=self.ZONE, instance=self.INSTANCE)
493        self.compute_client.WaitOnOperation.assert_called_with(
494            mock.ANY,
495            operation_scope=gcompute_client.OperationScope.ZONE,
496            scope_name=self.ZONE)
497
498    def _CompareMachineSizeTestHelper(self,
499                                      machine_info_1,
500                                      machine_info_2,
501                                      expected_result=None,
502                                      expected_error_type=None):
503        """Helper class for testing CompareMachineSize.
504
505    Args:
506      machine_info_1: A dictionary representing the first machine size.
507      machine_info_2: A dictionary representing the second machine size.
508      expected_result: An integer, 0, 1 or -1, or None if not set.
509      expected_error_type: An exception type, if set will check for exception.
510    """
511        self.Patch(
512            gcompute_client.ComputeClient,
513            "GetMachineType",
514            side_effect=[machine_info_1, machine_info_2])
515        if expected_error_type:
516            self.assertRaises(expected_error_type,
517                              self.compute_client.CompareMachineSize, "name1",
518                              "name2", self.ZONE)
519        else:
520            result = self.compute_client.CompareMachineSize("name1", "name2",
521                                                            self.ZONE)
522            self.assertEqual(result, expected_result)
523
524        self.compute_client.GetMachineType.assert_has_calls(
525            [mock.call("name1", self.ZONE), mock.call("name2", self.ZONE)])
526
527    def testCompareMachineSizeSmall(self):
528        """Test CompareMachineSize where the first one is smaller."""
529        machine_info_1 = {"guestCpus": 10, "memoryMb": 100}
530        machine_info_2 = {"guestCpus": 10, "memoryMb": 200}
531        self._CompareMachineSizeTestHelper(machine_info_1, machine_info_2, -1)
532
533    def testCompareMachineSizeLarge(self):
534        """Test CompareMachineSize where the first one is larger."""
535        machine_info_1 = {"guestCpus": 10, "memoryMb": 200}
536        machine_info_2 = {"guestCpus": 10, "memoryMb": 100}
537        self._CompareMachineSizeTestHelper(machine_info_1, machine_info_2, 1)
538
539    def testCompareMachineSizeEqual(self):
540        """Test CompareMachineSize where two machine sizes are equal."""
541        machine_info = {"guestCpus": 10, "memoryMb": 100}
542        self._CompareMachineSizeTestHelper(machine_info, machine_info, 0)
543
544    def testCompareMachineSizeBadMetric(self):
545        """Test CompareMachineSize with bad metric."""
546        machine_info = {"unkown_metric": 10, "memoryMb": 100}
547        self._CompareMachineSizeTestHelper(
548            machine_info, machine_info, expected_error_type=errors.DriverError)
549
550    def testGetMachineType(self):
551        """Test GetMachineType."""
552        resource_mock = mock.MagicMock()
553        mock_api = mock.MagicMock()
554        self.compute_client._service.machineTypes = mock.MagicMock(
555            return_value=resource_mock)
556        resource_mock.get = mock.MagicMock(return_value=mock_api)
557        mock_api.execute = mock.MagicMock(
558            return_value={"name": self.MACHINE_TYPE})
559        result = self.compute_client.GetMachineType(self.MACHINE_TYPE,
560                                                    self.ZONE)
561        self.assertEqual(result, {"name": self.MACHINE_TYPE})
562        resource_mock.get.assert_called_with(
563            project=self.PROJECT,
564            zone=self.ZONE,
565            machineType=self.MACHINE_TYPE)
566
567    def _GetSerialPortOutputTestHelper(self, response):
568        """Helper function for testing GetSerialPortOutput.
569
570    Args:
571      response: A dictionary representing a fake response.
572    """
573        resource_mock = mock.MagicMock()
574        mock_api = mock.MagicMock()
575        self.compute_client._service.instances = mock.MagicMock(
576            return_value=resource_mock)
577        resource_mock.getSerialPortOutput = mock.MagicMock(
578            return_value=mock_api)
579        mock_api.execute = mock.MagicMock(return_value=response)
580
581        if "contents" in response:
582            result = self.compute_client.GetSerialPortOutput(
583                instance=self.INSTANCE, zone=self.ZONE)
584            self.assertEqual(result, "fake contents")
585        else:
586            self.assertRaisesRegexp(
587                errors.DriverError,
588                "Malformed response.*",
589                self.compute_client.GetSerialPortOutput,
590                instance=self.INSTANCE,
591                zone=self.ZONE)
592        resource_mock.getSerialPortOutput.assert_called_with(
593            project=self.PROJECT,
594            zone=self.ZONE,
595            instance=self.INSTANCE,
596            port=1)
597
598    def testGetSerialPortOutput(self):
599        response = {"contents": "fake contents"}
600        self._GetSerialPortOutputTestHelper(response)
601
602    def testGetSerialPortOutputFail(self):
603        response = {"malformed": "fake contents"}
604        self._GetSerialPortOutputTestHelper(response)
605
606    def testGetInstanceNamesByIPs(self):
607        """Test GetInstanceNamesByIPs."""
608        good_instance = {
609            "name": "instance_1",
610            "networkInterfaces": [
611                {
612                    "accessConfigs": [
613                        {"natIP": "172.22.22.22"},
614                    ],
615                },
616            ],
617        }
618        bad_instance = {"name": "instance_2"}
619        self.Patch(
620            gcompute_client.ComputeClient,
621            "ListInstances",
622            return_value=[good_instance, bad_instance])
623        ip_name_map = self.compute_client.GetInstanceNamesByIPs(
624            ips=["172.22.22.22", "172.22.22.23"], zone=self.ZONE)
625        self.assertEqual(ip_name_map, {"172.22.22.22": "instance_1",
626                                       "172.22.22.23": None})
627
628    def testAddSshRsa(self):
629        """Test AddSshRsa.."""
630        fake_user = "fake_user"
631        sshkey = (
632            "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDBkTOTRze9v2VOqkkf7RG"
633            "jSkg6Z2kb9Q9UHsDGatvend3fmjIw1Tugg0O7nnjlPkskmlgyd4a/j99WOeLL"
634            "CPk6xPyoVjrPUVBU/pAk09ORTC4Zqk6YjlW7LOfzvqmXhmIZfYu6Q4Yt50pZzhl"
635            "lllfu26nYjY7Tg12D019nJi/kqPX5+NKgt0LGXTu8T1r2Gav/q4V7QRWQrB8Eiu"
636            "pxXR7I2YhynqovkEt/OXG4qWgvLEXGsWtSQs0CtCzqEVxz0Y9ECr7er4VdjSQxV"
637            "AaeLAsQsK9ROae8hMBFZ3//8zLVapBwpuffCu+fUoql9qeV9xagZcc9zj8XOUOW"
638            "ApiihqNL1111 test@test1.org")
639        project = {
640            "commonInstanceMetadata": {
641                "kind": "compute#metadata",
642                "fingerprint": "a-23icsyx4E=",
643                "items": [
644                    {
645                        "key": "sshKeys",
646                        "value": "user:key"
647                    }
648                ]
649            }
650        }
651        expected = {
652            "kind": "compute#metadata",
653            "fingerprint": "a-23icsyx4E=",
654            "items": [
655                {
656                    "key": "sshKeys",
657                    "value": "user:key\n%s:%s" % (fake_user, sshkey)
658                }
659            ]
660        }
661
662        self.Patch(os.path, "exists", return_value=True)
663        m = mock.mock_open(read_data=sshkey)
664        self.Patch(__builtins__, "open", m, create=True)
665        self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
666        self.Patch(
667            gcompute_client.ComputeClient, "GetProject", return_value=project)
668        resource_mock = mock.MagicMock()
669        self.compute_client._service.projects = mock.MagicMock(
670            return_value=resource_mock)
671        resource_mock.setCommonInstanceMetadata = mock.MagicMock()
672
673        self.compute_client.AddSshRsa(fake_user, "/path/to/test_rsa.pub")
674        resource_mock.setCommonInstanceMetadata.assert_called_with(
675            project=self.PROJECT, body=expected)
676
677    def testAddSshRsaInvalidKey(self):
678        """Test AddSshRsa.."""
679        fake_user = "fake_user"
680        sshkey = "ssh-rsa v2VOqkkf7RGL1111 test@test1.org"
681        project = {
682            "commonInstanceMetadata": {
683                "kind": "compute#metadata",
684                "fingerprint": "a-23icsyx4E=",
685                "items": [
686                    {
687                        "key": "sshKeys",
688                        "value": "user:key"
689                    }
690                ]
691            }
692        }
693        self.Patch(os.path, "exists", return_value=True)
694        m = mock.mock_open(read_data=sshkey)
695        self.Patch(__builtins__, "open", m, create=True)
696        self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
697        self.Patch(
698            gcompute_client.ComputeClient, "GetProject", return_value=project)
699        self.assertRaisesRegexp(errors.DriverError, "rsa key is invalid:*",
700                                self.compute_client.AddSshRsa, fake_user,
701                                "/path/to/test_rsa.pub")
702
703    def testDeleteDisks(self):
704        """Test DeleteDisks."""
705        self._SetupBatchHttpRequestMock()
706        self.Patch(gcompute_client.ComputeClient, "WaitOnOperation")
707        fake_disks = ["fake_disk_1", "fake_disk_2"]
708        mock_api = mock.MagicMock()
709        resource_mock = mock.MagicMock()
710        self.compute_client._service.disks = mock.MagicMock(
711            return_value=resource_mock)
712        resource_mock.delete = mock.MagicMock(return_value=mock_api)
713        # Call the API.
714        deleted, failed, error_msgs = self.compute_client.DeleteDisks(
715            fake_disks, zone=self.ZONE)
716        # Verify
717        calls = [mock.call(
718            project=self.PROJECT, disk="fake_disk_1", zone=self.ZONE),
719                 mock.call(
720                     project=self.PROJECT, disk="fake_disk_2", zone=self.ZONE)]
721        resource_mock.delete.assert_has_calls(calls, any_order=True)
722        self.assertEqual(
723            gcompute_client.ComputeClient.WaitOnOperation.call_count, 2)
724        self.assertEqual(error_msgs, [])
725        self.assertEqual(failed, [])
726        self.assertEqual(set(deleted), set(fake_disks))
727
728
729if __name__ == "__main__":
730    unittest.main()
731