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