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
20import apiclient
21import mock
22
23import unittest
24from acloud.internal.lib import base_cloud_client
25from acloud.internal.lib import driver_test_lib
26from acloud.public import errors
27
28
29class FakeError(Exception):
30    """Fake Error for testing retries."""
31
32
33class BaseCloudApiClientTest(driver_test_lib.BaseDriverTest):
34    """Test BaseCloudApiClient."""
35
36    def setUp(self):
37        """Set up test."""
38        super(BaseCloudApiClientTest, self).setUp()
39
40    def testInitResourceHandle(self):
41        """Test InitResourceHandle."""
42        # Setup mocks
43        mock_credentials = mock.MagicMock()
44        self.Patch(base_cloud_client, "build")
45        # Call the method
46        base_cloud_client.BaseCloudApiClient(mock.MagicMock())
47        base_cloud_client.build.assert_called_once_with(
48            serviceName=base_cloud_client.BaseCloudApiClient.API_NAME,
49            version=base_cloud_client.BaseCloudApiClient.API_VERSION,
50            http=mock.ANY)
51
52    def _SetupInitMocks(self):
53        """Setup mocks required to initialize a base cloud client.
54
55    Returns:
56      A base_cloud_client.BaseCloudApiClient mock.
57    """
58        self.Patch(
59            base_cloud_client.BaseCloudApiClient,
60            "InitResourceHandle",
61            return_value=mock.MagicMock())
62        return base_cloud_client.BaseCloudApiClient(mock.MagicMock())
63
64    def _SetupBatchHttpRequestMock(self, rid_to_responses, rid_to_exceptions):
65        """Setup BatchHttpRequest mock."""
66
67        rid_to_exceptions = rid_to_exceptions or {}
68        rid_to_responses = rid_to_responses or {}
69
70        def _CreatMockBatchHttpRequest():
71            """Create a mock BatchHttpRequest object."""
72            requests = {}
73
74            def _Add(request, callback, request_id):
75                requests[request_id] = (request, callback)
76
77            def _Execute():
78                for rid in requests:
79                    requests[rid][0].execute()
80                    _, callback = requests[rid]
81                    callback(
82                        request_id=rid,
83                        response=rid_to_responses.get(rid),
84                        exception=rid_to_exceptions.get(rid))
85
86            mock_batch = mock.MagicMock()
87            mock_batch.add = _Add
88            mock_batch.execute = _Execute
89            return mock_batch
90
91        self.Patch(
92            apiclient.http,
93            "BatchHttpRequest",
94            side_effect=_CreatMockBatchHttpRequest)
95
96    def testBatchExecute(self):
97        """Test BatchExecute."""
98        self.Patch(time, "sleep")
99        client = self._SetupInitMocks()
100        requests = {"r1": mock.MagicMock(),
101                    "r2": mock.MagicMock(),
102                    "r3": mock.MagicMock()}
103        response = {"name": "fake_response"}
104        error_1 = errors.HttpError(503, "fake retriable error.")
105        error_2 = FakeError("fake retriable error.")
106        responses = {"r1": response, "r2": None, "r3": None}
107        exceptions = {"r1": None, "r2": error_1, "r3": error_2}
108        self._SetupBatchHttpRequestMock(responses, exceptions)
109        results = client.BatchExecute(
110            requests, other_retriable_errors=(FakeError, ))
111        expected_results = {
112            "r1": (response, None),
113            "r2": (None, error_1),
114            "r3": (None, error_2)
115        }
116        self.assertEqual(results, expected_results)
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        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