Snap for 9940335 from 7931131bd45dcbb180824661cfae667a06f5ec95 to udc-d1-release

Change-Id: I7ce7a9468468eab26bba8297a20fd6e1916a2ebb

2 files changed, 22 insertions, 2 deletions

diff --git a/restart/restart.py b/restart/restart.py
index 5e148941..b10901e7 100644
--- a/restart/restart.py
+++ b/restart/restart.py
@@ -18,6 +18,7 @@ This command will restart the CF AVD from a remote instance.
 
 import logging
 import subprocess
+import sys
 
 from acloud import errors
 from acloud.internal import constants
@@ -32,6 +33,8 @@ from acloud.reconnect import reconnect
 
 
 logger = logging.getLogger(__name__)
+_NOT_SUPPORT_MSG = ("Currently the restart function doesn't support local "
+                    "instances. Please try to create one new instance.")
 
 
 def RestartFromInstance(cfg, instance, instance_id, powerwash_data):
@@ -96,6 +99,10 @@ def Run(args):
             cfg, [args.instance_name])
         return RestartFromInstance(
             cfg, instance[0], args.instance_id, args.powerwash)
+    if (not list_instances.GetCFRemoteInstances(cfg)
+            and list_instances.GetLocalInstances()):
+        utils.PrintColorString(_NOT_SUPPORT_MSG, utils.TextColors.FAIL)
+        sys.exit()
     return RestartFromInstance(cfg,
                                list_instances.ChooseOneRemoteInstance(cfg),
                                args.instance_id,
diff --git a/restart/restart_test.py b/restart/restart_test.py
index d74dee75..00256f4c 100644
--- a/restart/restart_test.py
+++ b/restart/restart_test.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Tests for restart."""
+import sys
+
 import unittest
 
 from unittest import mock
@@ -29,9 +31,9 @@ from acloud.restart import restart
 
 class RestartTest(driver_test_lib.BaseDriverTest):
     """Test restart."""
-
+    @mock.patch.object(sys, "exit")
     @mock.patch.object(restart, "RestartFromInstance")
-    def testRun(self, mock_restart):
+    def testRun(self, mock_restart, mock_exit):
         """test Run."""
         cfg = mock.MagicMock()
         args = mock.MagicMock()
@@ -49,6 +51,8 @@ class RestartTest(driver_test_lib.BaseDriverTest):
 
         # Test case for user select one instance to restart AVD.
         selected_instance = mock.MagicMock()
+        self.Patch(list_instances, "GetCFRemoteInstances",
+                   return_value=selected_instance)
         self.Patch(list_instances, "ChooseOneRemoteInstance",
                    return_value=selected_instance)
         args.instance_name = None
@@ -56,6 +60,15 @@ class RestartTest(driver_test_lib.BaseDriverTest):
         mock_restart.assert_has_calls([
             mock.call(cfg, selected_instance, args.instance_id, args.powerwash)])
 
+        # Test case for not support local instances.
+        local_instances = mock.MagicMock()
+        self.Patch(list_instances, "GetCFRemoteInstances",
+                   return_value=None)
+        self.Patch(list_instances, "GetLocalInstances",
+                   return_value=local_instances)
+        restart.Run(args)
+        mock_exit.assert_called_once()
+
     # pylint: disable=no-member
     def testRestartFromInstance(self):
         """test RestartFromInstance."""