• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python
2#
3# Copyright 2018 - 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.setup.gcp_setup_runner."""
17
18import unittest
19import os
20
21from unittest import mock
22import six
23
24# pylint: disable=no-name-in-module,import-error,no-member
25from acloud import errors
26from acloud.internal.lib import utils
27from acloud.internal.proto import user_config_pb2
28from acloud.public import config
29from acloud.setup import gcp_setup_runner
30
31_GCP_USER_CONFIG = """
32[compute]
33region = new_region
34zone = new_zone
35[core]
36account = new@google.com
37disable_usage_reporting = False
38project = new_project
39"""
40
41def _CreateCfgFile():
42    """A helper method that creates a mock configuration object."""
43    default_cfg = """
44project: "fake_project"
45zone: "fake_zone"
46storage_bucket_name: "fake_bucket"
47client_id: "fake_client_id"
48client_secret: "fake_client_secret"
49"""
50    return default_cfg
51
52
53# pylint: disable=protected-access
54class AcloudGCPSetupTest(unittest.TestCase):
55    """Test GCP Setup steps."""
56
57    def setUp(self):
58        """Create config and gcp_env_runner."""
59        self.cfg_path = "acloud_unittest.config"
60        file_write = open(self.cfg_path, 'w')
61        file_write.write(_CreateCfgFile().strip())
62        file_write.close()
63        self.gcp_env_runner = gcp_setup_runner.GcpTaskRunner(self.cfg_path)
64        self.gcloud_runner = gcp_setup_runner.GoogleSDKBins("")
65
66    def tearDown(self):
67        """Remove temp file."""
68        if os.path.isfile(self.cfg_path):
69            os.remove(self.cfg_path)
70
71    def testUpdateConfigFile(self):
72        """Test update config file."""
73        # Test update project field.
74        gcp_setup_runner.UpdateConfigFile(self.cfg_path, "project",
75                                          "test_project")
76        cfg = config.AcloudConfigManager.LoadConfigFromProtocolBuffer(
77            open(self.cfg_path, "r"), user_config_pb2.UserConfig)
78        self.assertEqual(cfg.project, "test_project")
79        self.assertEqual(cfg.ssh_private_key_path, "")
80        # Test add ssh key path in config.
81        gcp_setup_runner.UpdateConfigFile(self.cfg_path,
82                                          "ssh_private_key_path", "test_path")
83        cfg = config.AcloudConfigManager.LoadConfigFromProtocolBuffer(
84            open(self.cfg_path, "r"), user_config_pb2.UserConfig)
85        self.assertEqual(cfg.project, "test_project")
86        self.assertEqual(cfg.ssh_private_key_path, "test_path")
87        # Test config is not a file
88        with mock.patch("os.path.isfile") as chkfile:
89            chkfile.return_value = False
90            gcp_setup_runner.UpdateConfigFile(self.cfg_path, "project",
91                                              "test_project")
92            cfg = config.AcloudConfigManager.LoadConfigFromProtocolBuffer(
93                open(self.cfg_path, "r"), user_config_pb2.UserConfig)
94            self.assertEqual(cfg.project, "test_project")
95
96    @mock.patch("os.path.dirname", return_value="")
97    @mock.patch.object(utils, "CheckOutput")
98    def testSeupProjectZone(self, mock_runner, mock_path):
99        """Test setup project and zone."""
100        gcloud_runner = gcp_setup_runner.GoogleSDKBins(mock_path)
101        self.gcp_env_runner.project = "fake_project"
102        self.gcp_env_runner.zone = "fake_zone"
103        mock_runner.side_effect = [0, _GCP_USER_CONFIG]
104        self.gcp_env_runner._UpdateProject(gcloud_runner)
105        self.assertEqual(self.gcp_env_runner.project, "new_project")
106        self.assertEqual(self.gcp_env_runner.zone, "new_zone")
107
108    @mock.patch.object(six.moves, "input")
109    def testSetupClientIDSecret(self, mock_id):
110        """Test setup client ID and client secret."""
111        self.gcp_env_runner.client_id = "fake_client_id"
112        self.gcp_env_runner.client_secret = "fake_client_secret"
113        mock_id.side_effect = ["new_id", "new_secret"]
114        self.gcp_env_runner._SetupClientIDSecret()
115        self.assertEqual(self.gcp_env_runner.client_id, "new_id")
116        self.assertEqual(self.gcp_env_runner.client_secret, "new_secret")
117
118    @mock.patch.object(gcp_setup_runner, "UpdateConfigFile")
119    @mock.patch.object(utils, "CreateSshKeyPairIfNotExist")
120    def testSetupSSHKeys(self, mock_check, mock_update):
121        """Test setup the pair of the ssh key for acloud.config."""
122        # Test ssh key has already setup
123        gcp_setup_runner.SetupSSHKeys(self.cfg_path,
124                                      "fake_private_key_path",
125                                      "fake_public_key_path")
126        self.assertEqual(mock_update.call_count, 0)
127        # Test if private_key_path is empty string
128        with mock.patch('os.path.expanduser') as ssh_path:
129            ssh_path.return_value = ""
130            gcp_setup_runner.SetupSSHKeys(self.cfg_path,
131                                          "fake_private_key_path",
132                                          "fake_public_key_path")
133            mock_check.assert_called_once_with(
134                gcp_setup_runner._DEFAULT_SSH_PRIVATE_KEY,
135                gcp_setup_runner._DEFAULT_SSH_PUBLIC_KEY)
136
137            mock_update.assert_has_calls([
138                mock.call(self.cfg_path, "ssh_private_key_path",
139                          gcp_setup_runner._DEFAULT_SSH_PRIVATE_KEY),
140                mock.call(self.cfg_path, "ssh_public_key_path",
141                          gcp_setup_runner._DEFAULT_SSH_PUBLIC_KEY)])
142
143    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_CreateStableHostImage")
144    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_EnableGcloudServices")
145    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_SetupProject")
146    @mock.patch.object(gcp_setup_runner, "GoogleSDKBins")
147    def testSetupGcloudInfo(self, mock_sdk, mock_set, mock_run, mock_create):
148        """test setup gcloud info"""
149        with mock.patch("acloud.setup.google_sdk.GoogleSDK"):
150            self.gcp_env_runner._SetupGcloudInfo()
151            mock_sdk.assert_called_once()
152            mock_set.assert_called_once()
153            mock_run.assert_called_once()
154            mock_create.assert_called_once()
155
156    @mock.patch.object(gcp_setup_runner, "UpdateConfigFile")
157    def testCreateStableHostImage(self, mock_update):
158        """test create stable hostimage."""
159        # Test no need to create stable hose image name.
160        self.gcp_env_runner.stable_host_image_name = "fake_host_image_name"
161        self.gcp_env_runner._CreateStableHostImage()
162        self.assertEqual(mock_update.call_count, 0)
163        # Test need to reset stable hose image name.
164        self.gcp_env_runner.stable_host_image_name = ""
165        self.gcp_env_runner._CreateStableHostImage()
166        self.assertEqual(mock_update.call_count, 1)
167
168    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_CheckBillingEnable")
169    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_NeedProjectSetup")
170    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_SetupClientIDSecret")
171    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_UpdateProject")
172    def testSetupProjectNoChange(self, mock_setproj, mock_setid,
173                                 mock_chkproj, mock_check_billing):
174        """test setup project and project not be changed."""
175        # Test project didn't change, and no need to setup client id/secret
176        mock_chkproj.return_value = False
177        self.gcp_env_runner.client_id = "test_client_id"
178        self.gcp_env_runner._SetupProject(self.gcloud_runner)
179        self.assertEqual(mock_setproj.call_count, 0)
180        self.assertEqual(mock_setid.call_count, 0)
181        mock_check_billing.assert_called_once()
182        # Test project didn't change, but client_id is empty
183        self.gcp_env_runner.client_id = ""
184        self.gcp_env_runner._SetupProject(self.gcloud_runner)
185        self.assertEqual(mock_setproj.call_count, 0)
186        mock_setid.assert_called_once()
187        self.assertEqual(mock_check_billing.call_count, 2)
188
189    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_CheckBillingEnable")
190    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_NeedProjectSetup")
191    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_SetupClientIDSecret")
192    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_UpdateProject")
193    def testSetupProjectChanged(self, mock_setproj, mock_setid,
194                                mock_chkproj, mock_check_billing):
195        """test setup project when project changed."""
196        mock_chkproj.return_value = True
197        mock_setproj.return_value = True
198        self.gcp_env_runner._SetupProject(self.gcloud_runner)
199        mock_setproj.assert_called_once()
200        mock_setid.assert_called_once()
201        mock_check_billing.assert_called_once()
202
203    @mock.patch.object(utils, "GetUserAnswerYes")
204    def testNeedProjectSetup(self, mock_ans):
205        """test need project setup."""
206        # Test need project setup.
207        self.gcp_env_runner.project = ""
208        self.gcp_env_runner.zone = ""
209        self.assertTrue(self.gcp_env_runner._NeedProjectSetup())
210        # Test no need project setup and get user's answer.
211        self.gcp_env_runner.project = "test_project"
212        self.gcp_env_runner.zone = "test_zone"
213        self.gcp_env_runner._NeedProjectSetup()
214        mock_ans.assert_called_once()
215
216    def testNeedClientIDSetup(self):
217        """test need client_id setup."""
218        # Test project changed.
219        self.assertTrue(self.gcp_env_runner._NeedClientIDSetup(True))
220        # Test project is not changed but client_id or client_secret is empty.
221        self.gcp_env_runner.client_id = ""
222        self.gcp_env_runner.client_secret = ""
223        self.assertTrue(self.gcp_env_runner._NeedClientIDSetup(False))
224        # Test no need client_id setup.
225        self.gcp_env_runner.client_id = "test_client_id"
226        self.gcp_env_runner.client_secret = "test_client_secret"
227        self.assertFalse(self.gcp_env_runner._NeedClientIDSetup(False))
228
229    @mock.patch.object(utils, "CheckOutput")
230    def testEnableGcloudServices(self, mock_run):
231        """test enable Gcloud services."""
232        mock_run.return_value = ""
233        self.gcp_env_runner._EnableGcloudServices(self.gcloud_runner)
234        mock_run.assert_has_calls([
235            mock.call(["gcloud", "services", "enable",
236                       gcp_setup_runner._ANDROID_BUILD_SERVICE],
237                      env=self.gcloud_runner._env, stderr=-2),
238            mock.call(["gcloud", "services", "enable",
239                       gcp_setup_runner._COMPUTE_ENGINE_SERVICE],
240                      env=self.gcloud_runner._env, stderr=-2)])
241
242    @mock.patch.object(utils, "CheckOutput")
243    def testGoogleAPIService(self, mock_run):
244        """Test GoogleAPIService"""
245        api_service = gcp_setup_runner.GoogleAPIService("service_name",
246                                                        "error_message")
247        api_service.EnableService(self.gcloud_runner)
248        mock_run.assert_has_calls([
249            mock.call(["gcloud", "services", "enable", "service_name"],
250                      env=self.gcloud_runner._env, stderr=-2)])
251
252    @mock.patch.object(utils, "CheckOutput")
253    def testCheckBillingEnable(self, mock_run):
254        """Test CheckBillingEnable"""
255        # Test billing account in gcp project already enabled.
256        mock_run.return_value = "billingEnabled: true"
257        self.gcp_env_runner._CheckBillingEnable(self.gcloud_runner)
258        mock_run.assert_has_calls([
259            mock.call(
260                [
261                    "gcloud", "alpha", "billing", "projects", "describe",
262                    self.gcp_env_runner.project
263                ],
264                env=self.gcloud_runner._env)
265        ])
266
267        # Test billing account in gcp project was not enabled.
268        mock_run.return_value = "billingEnabled: false"
269        with self.assertRaises(errors.NoBillingError):
270            self.gcp_env_runner._CheckBillingEnable(self.gcloud_runner)
271
272
273if __name__ == "__main__":
274    unittest.main()
275