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.base_cloud_client.""" 18 19import time 20 21import unittest 22 23from unittest import mock 24 25from acloud import errors 26from acloud.internal.lib import base_cloud_client 27from acloud.internal.lib import driver_test_lib 28 29 30class FakeError(Exception): 31 """Fake Error for testing retries.""" 32 33 34class BaseCloudApiClientTest(driver_test_lib.BaseDriverTest): 35 """Test BaseCloudApiClient.""" 36 37 def testInitResourceHandle(self): 38 """Test InitResourceHandle.""" 39 # Setup mocks 40 self.Patch(base_cloud_client, "build") 41 # Call the method 42 base_cloud_client.BaseCloudApiClient(mock.MagicMock()) 43 base_cloud_client.build.assert_called_once_with( 44 serviceName=base_cloud_client.BaseCloudApiClient.API_NAME, 45 version=base_cloud_client.BaseCloudApiClient.API_VERSION, 46 cache_discovery=False, 47 static_discovery=False, 48 http=mock.ANY) 49 50 def _SetupInitMocks(self): 51 """Setup mocks required to initialize a base cloud client. 52 53 Returns: 54 A base_cloud_client.BaseCloudApiClient mock. 55 """ 56 self.Patch( 57 base_cloud_client.BaseCloudApiClient, 58 "InitResourceHandle", 59 return_value=mock.MagicMock()) 60 return base_cloud_client.BaseCloudApiClient(mock.MagicMock()) 61 62 def _SetupBatchHttpRequestMock(self, rid_to_responses, rid_to_exceptions, 63 client): 64 """Setup BatchHttpRequest mock.""" 65 66 rid_to_exceptions = rid_to_exceptions or {} 67 rid_to_responses = rid_to_responses or {} 68 69 def _CreatMockBatchHttpRequest(): 70 """Create a mock BatchHttpRequest object.""" 71 requests = {} 72 73 def _Add(request, callback, request_id): 74 requests[request_id] = (request, callback) 75 76 def _Execute(): 77 for rid in requests: 78 requests[rid][0].execute() 79 _, callback = requests[rid] 80 callback( 81 request_id=rid, 82 response=rid_to_responses.get(rid), 83 exception=rid_to_exceptions.get(rid)) 84 85 mock_batch = mock.MagicMock() 86 mock_batch.add = _Add 87 mock_batch.execute = _Execute 88 return mock_batch 89 90 self.Patch(client.service, "new_batch_http_request", 91 side_effect=_CreatMockBatchHttpRequest) 92 93 def testBatchExecute(self): 94 """Test BatchExecute.""" 95 self.Patch(time, "sleep") 96 client = self._SetupInitMocks() 97 requests = {"r1": mock.MagicMock(), 98 "r2": mock.MagicMock(), 99 "r3": mock.MagicMock()} 100 response = {"name": "fake_response"} 101 error_1 = errors.HttpError(503, "fake retriable error.") 102 error_2 = FakeError("fake retriable error.") 103 responses = {"r1": response, "r2": None, "r3": None} 104 exceptions = {"r1": None, "r2": error_1, "r3": error_2} 105 self._SetupBatchHttpRequestMock(responses, exceptions, client) 106 results = client.BatchExecute( 107 requests, other_retriable_errors=(FakeError, )) 108 expected_results = { 109 "r1": (response, None), 110 "r2": (None, error_1), 111 "r3": (None, error_2) 112 } 113 114 self.assertEqual(results, expected_results) 115 # pylint: disable=no-member 116 self.assertEqual(requests["r1"].execute.call_count, 1) 117 self.assertEqual(requests["r2"].execute.call_count, 118 client.RETRY_COUNT + 1) 119 self.assertEqual(requests["r3"].execute.call_count, 120 client.RETRY_COUNT + 1) 121 122 def testListWithMultiPages(self): 123 """Test ListWithMultiPages.""" 124 fake_token = "fake_next_page_token" 125 item_1 = "item_1" 126 item_2 = "item_2" 127 response_1 = {"items": [item_1], "nextPageToken": fake_token} 128 response_2 = {"items": [item_2]} 129 130 api_mock = mock.MagicMock() 131 api_mock.execute.side_effect = [response_1, response_2] 132 resource_mock = mock.MagicMock(return_value=api_mock) 133 client = self._SetupInitMocks() 134 items = client.ListWithMultiPages( 135 api_resource=resource_mock, fake_arg="fake_arg") 136 self.assertEqual(items, [item_1, item_2]) 137 138 def testExecuteWithRetry(self): 139 """Test Execute is called and retries are triggered.""" 140 self.Patch(time, "sleep") 141 client = self._SetupInitMocks() 142 api_mock = mock.MagicMock() 143 error = errors.HttpError(503, "fake retriable error.") 144 api_mock.execute.side_effect = error 145 self.assertRaises(errors.HttpError, client.Execute, api_mock) 146 147 api_mock = mock.MagicMock() 148 api_mock.execute.side_effect = FakeError("fake retriable error.") 149 # pylint: disable=no-member 150 self.assertRaises( 151 FakeError, 152 client.Execute, 153 api_mock, 154 other_retriable_errors=(FakeError, )) 155 self.assertEqual(api_mock.execute.call_count, client.RETRY_COUNT + 1) 156 157 158if __name__ == "__main__": 159 unittest.main() 160