Merge "acloud: increasing test cases of gce_setup_runner."
am: 062bb4566e

Change-Id: Ide27f01d311c65afe08a42738dd2218b696cd5c2
diff --git a/setup/gcp_setup_runner_test.py b/setup/gcp_setup_runner_test.py
index 343a3c0..9f4a252 100644
--- a/setup/gcp_setup_runner_test.py
+++ b/setup/gcp_setup_runner_test.py
@@ -19,7 +19,9 @@
 import os
 import mock
 
-# pylint: disable=no-name-in-module,import-error
+# pylint: disable=no-name-in-module,import-error,no-member
+from acloud import errors
+from acloud.internal.lib import utils
 from acloud.internal.proto import user_config_pb2
 from acloud.public import config
 from acloud.setup import gcp_setup_runner
@@ -74,18 +76,26 @@
             open(self.cfg_path, "r"), user_config_pb2.UserConfig)
         self.assertEqual(cfg.project, "test_project")
         self.assertEqual(cfg.ssh_private_key_path, "")
-        # Test add ssh key path in config
+        # Test add ssh key path in config.
         gcp_setup_runner.UpdateConfigFile(self.cfg_path,
                                           "ssh_private_key_path", "test_path")
         cfg = config.AcloudConfigManager.LoadConfigFromProtocolBuffer(
             open(self.cfg_path, "r"), user_config_pb2.UserConfig)
         self.assertEqual(cfg.project, "test_project")
         self.assertEqual(cfg.ssh_private_key_path, "test_path")
+        # Test config is not a file
+        with mock.patch("os.path.isfile") as chkfile:
+            chkfile.return_value = False
+            gcp_setup_runner.UpdateConfigFile(self.cfg_path, "project",
+                                              "test_project")
+            cfg = config.AcloudConfigManager.LoadConfigFromProtocolBuffer(
+                open(self.cfg_path, "r"), user_config_pb2.UserConfig)
+            self.assertEqual(cfg.project, "test_project")
 
-    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_CreateBucket")
+    @mock.patch("subprocess.check_output")
     @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_BucketExists")
     @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_BucketInDefaultRegion")
-    def testCreateDefaultBucket(self, mock_valid, mock_exist, mock_create):
+    def testCreateDefaultBucket(self, mock_valid, mock_exist, mock_out):
         """Test default bucket name.
 
         Default bucket name is "acloud-{project}".
@@ -95,10 +105,13 @@
         self.gcp_env_runner.project = "fake_project"
         mock_exist.return_value = False
         mock_valid.return_value = False
-        mock_create.return_value = True
         self.assertEqual(
             "acloud-fake_project",
             self.gcp_env_runner._CreateDefaultBucket(self.gcloud_runner))
+        self.gcloud_runner.gsutil_command_path = "gsutil"
+        mock_out.assert_called_once_with(
+            ["gsutil", "mb", "gs://acloud-fake_project"])
+
         mock_exist.return_value = True
         mock_valid.return_value = False
         self.assertEqual(
@@ -173,6 +186,147 @@
         self.assertTrue(
             self.gcp_env_runner._BucketInDefaultRegion("test-bucket",
                                                        self.gcloud_runner))
+        # Test parsing error.
+        mock_region.return_value = "Wrong Lable:ASIA"
+        self.assertRaises(errors.ParseBucketRegionError,
+                          self.gcp_env_runner._BucketInDefaultRegion,
+                          "test-bucket", self.gcloud_runner)
+
+    @mock.patch.object(gcp_setup_runner, "UpdateConfigFile")
+    @mock.patch.object(utils, "CreateSshKeyPairIfNotExist")
+    def testSetupSSHKeys(self, mock_check, mock_update):
+        """Test setup the pair of the ssh key for acloud.config."""
+        # Test ssh key has already setup
+        gcp_setup_runner.SetupSSHKeys(self.cfg_path,
+                                      "fake_private_key_path",
+                                      "fake_public_key_path")
+        self.assertEqual(mock_update.call_count, 0)
+        # Test if private_key_path is empty string
+        with mock.patch('os.path.expanduser') as ssh_path:
+            ssh_path.return_value = ""
+            gcp_setup_runner.SetupSSHKeys(self.cfg_path,
+                                          "fake_private_key_path",
+                                          "fake_public_key_path")
+            mock_check.assert_called_once_with(
+                gcp_setup_runner._DEFAULT_SSH_PRIVATE_KEY,
+                gcp_setup_runner._DEFAULT_SSH_PUBLIC_KEY)
+
+            mock_update.assert_has_calls([
+                mock.call(self.cfg_path, "ssh_private_key_path",
+                          gcp_setup_runner._DEFAULT_SSH_PRIVATE_KEY),
+                mock.call(self.cfg_path, "ssh_public_key_path",
+                          gcp_setup_runner._DEFAULT_SSH_PUBLIC_KEY)])
+
+    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_CreateStableHostImage")
+    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_EnableGcloudServices")
+    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_SetupProject")
+    @mock.patch.object(gcp_setup_runner, "GoogleSDKBins")
+    def testSetupGcloudInfo(self, mock_sdk, mock_set, mock_run, mock_create):
+        """test setup gcloud info"""
+        with mock.patch("google_sdk.GoogleSDK"):
+            self.gcp_env_runner._SetupGcloudInfo()
+            mock_sdk.assert_called_once()
+            mock_set.assert_called_once()
+            mock_run.assert_called_once()
+            mock_create.assert_called_once()
+
+    @mock.patch.object(gcp_setup_runner, "UpdateConfigFile")
+    def testCreateStableHostImage(self, mock_update):
+        """test create stable hostimage."""
+        # Test no need to create stable hose image name.
+        self.gcp_env_runner.stable_host_image_name = "fake_host_image_name"
+        self.gcp_env_runner._CreateStableHostImage()
+        self.assertEqual(mock_update.call_count, 0)
+        # Test need to reset stable hose image name.
+        self.gcp_env_runner.stable_host_image_name = ""
+        self.gcp_env_runner._CreateStableHostImage()
+        self.assertEqual(mock_update.call_count, 1)
+
+    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_NeedProjectSetup")
+    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_SetupStorageBucket")
+    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_SetupClientIDSecret")
+    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_UpdateProject")
+    def testSetupProjectNoChange(self, mock_setproj, mock_setid,
+                                 mock_setstorage, mock_chkproj):
+        """test setup project and project not be changed."""
+        # Test project didn't change, and no need to setup client id/secret
+        mock_chkproj.return_value = False
+        self.gcp_env_runner.client_id = "test_client_id"
+        self.gcp_env_runner._SetupProject(self.gcloud_runner)
+        self.assertEqual(mock_setproj.call_count, 0)
+        self.assertEqual(mock_setid.call_count, 0)
+        mock_setstorage.assert_called_once()
+        # Test project didn't change, but client_id is empty
+        self.gcp_env_runner.client_id = ""
+        self.gcp_env_runner._SetupProject(self.gcloud_runner)
+        self.assertEqual(mock_setproj.call_count, 0)
+        mock_setid.assert_called_once()
+
+    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_NeedProjectSetup")
+    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_SetupStorageBucket")
+    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_SetupClientIDSecret")
+    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_UpdateProject")
+    def testSetupProjectChanged(self, mock_setproj, mock_setid,
+                                mock_setstorage, mock_chkproj):
+        """test setup project when project changed."""
+        mock_chkproj.return_value = True
+        mock_setproj.return_value = True
+        self.gcp_env_runner._SetupProject(self.gcloud_runner)
+        mock_setproj.assert_called_once()
+        mock_setid.assert_called_once()
+        mock_setstorage.assert_called_once()
+
+    @mock.patch.object(utils, "GetUserAnswerYes")
+    def testNeedProjectSetup(self, mock_ans):
+        """test need project setup."""
+        # Test need project setup.
+        self.gcp_env_runner.project = ""
+        self.gcp_env_runner.zone = ""
+        self.assertTrue(self.gcp_env_runner._NeedProjectSetup())
+        # Test no need project setup and get user's answer.
+        self.gcp_env_runner.project = "test_project"
+        self.gcp_env_runner.zone = "test_zone"
+        self.gcp_env_runner._NeedProjectSetup()
+        mock_ans.assert_called_once()
+
+    def testNeedClientIDSetup(self):
+        """test need client_id setup."""
+        # Test project changed.
+        self.assertTrue(self.gcp_env_runner._NeedClientIDSetup(True))
+        # Test project is not changed but client_id or client_secret is empty.
+        self.gcp_env_runner.client_id = ""
+        self.gcp_env_runner.client_secret = ""
+        self.assertTrue(self.gcp_env_runner._NeedClientIDSetup(False))
+        # Test no need client_id setup.
+        self.gcp_env_runner.client_id = "test_client_id"
+        self.gcp_env_runner.client_secret = "test_client_secret"
+        self.assertFalse(self.gcp_env_runner._NeedClientIDSetup(False))
+
+    @mock.patch("subprocess.check_output")
+    @mock.patch.object(gcp_setup_runner.GcpTaskRunner, "_CreateDefaultBucket")
+    @mock.patch.object(gcp_setup_runner, "UpdateConfigFile")
+    def testSetupStorageBucket(self, mock_update, mock_create, mock_check):
+        """test setup storage bucket."""
+        self.gcp_env_runner.storage_bucket_name = ""
+        mock_create.return_value = "fake_create_default_bucket"
+        self.gcp_env_runner._SetupStorageBucket(self.gcloud_runner)
+        mock_check.assert_has_calls([
+            mock.call(["gsutil", "acl", "ch", "-u",
+                       "android-build-prod@system.gserviceaccount.com:W",
+                       "gs://fake_create_default_bucket"], stderr=-2)])
+        mock_update.assert_called_once()
+
+    @mock.patch("subprocess.check_output")
+    def testEnableGcloudServices(self, mock_run):
+        """test enable Gcloud services."""
+        self.gcp_env_runner._EnableGcloudServices(self.gcloud_runner)
+        mock_run.assert_has_calls([
+            mock.call(["gcloud", "services", "enable",
+                       "storage-component.googleapis.com"], stderr=-2),
+            mock.call(["gcloud", "services", "enable",
+                       "androidbuildinternal.googleapis.com"], stderr=-2),
+            mock.call(["gcloud", "services", "enable",
+                       "compute.googleapis.com"], stderr=-2)])
 
 
 if __name__ == "__main__":