Merge "Add powerwash option into restart function."

5 files changed, 23 insertions, 7 deletions

diff --git a/Android.bp b/Android.bp
index b0ab7d48..ec2c2c16 100644
--- a/Android.bp
+++ b/Android.bp
@@ -106,9 +106,10 @@ python_test_host {
         "acloud_reconnect",
         "acloud_internal",
         "acloud_list",
+        "acloud_powerwash",
+        "acloud_public",
         "acloud_pull",
         "acloud_proto",
-        "acloud_public",
         "acloud_restart",
         "acloud_setup",
         "asuite_cc_client",
diff --git a/powerwash/powerwash.py b/powerwash/powerwash.py
index fcd8853a..05ab4e58 100644
--- a/powerwash/powerwash.py
+++ b/powerwash/powerwash.py
@@ -53,6 +53,7 @@ def PowerwashFromInstance(cfg, instance, instance_id):
     return report.Report(command="powerwash")
 
 
+@utils.TimeExecute(function_description="Waiting for AVD to powerwash")
 def PowerwashDevice(ssh, instance_id):
     """Powerwash AVD with the instance id.
 
diff --git a/restart/restart.py b/restart/restart.py
index ec2e6e70..5e148941 100644
--- a/restart/restart.py
+++ b/restart/restart.py
@@ -25,6 +25,7 @@ from acloud.internal.lib import utils
 from acloud.internal.lib.ssh import Ssh
 from acloud.internal.lib.ssh import IP
 from acloud.list import list as list_instances
+from acloud.powerwash import powerwash
 from acloud.public import config
 from acloud.public import report
 from acloud.reconnect import reconnect
@@ -33,13 +34,14 @@ from acloud.reconnect import reconnect
 logger = logging.getLogger(__name__)
 
 
-def RestartFromInstance(cfg, instance, instance_id):
+def RestartFromInstance(cfg, instance, instance_id, powerwash_data):
     """Restart AVD from remote CF instance.
 
     Args:
         cfg: AcloudConfig object.
         instance: list.Instance() object.
         instance_id: Integer of the instance id.
+        powerwash_data: Boolean, True to powerwash AVD data.
 
     Returns:
         A Report instance.
@@ -50,7 +52,10 @@ def RestartFromInstance(cfg, instance, instance_id):
               extra_args_ssh_tunnel=cfg.extra_args_ssh_tunnel)
     logger.info("Start to restart AVD id (%s) from the instance: %s.",
                 instance_id, instance.name)
-    RestartDevice(ssh, instance_id)
+    if powerwash_data:
+        powerwash.PowerwashDevice(ssh, instance_id)
+    else:
+        RestartDevice(ssh, instance_id)
     reconnect.ReconnectInstance(cfg.ssh_private_key_path,
                                 instance,
                                 report.Report(command="reconnect"),
@@ -89,7 +94,9 @@ def Run(args):
     if args.instance_name:
         instance = list_instances.GetInstancesFromInstanceNames(
             cfg, [args.instance_name])
-        return RestartFromInstance(cfg, instance[0], args.instance_id)
+        return RestartFromInstance(
+            cfg, instance[0], args.instance_id, args.powerwash)
     return RestartFromInstance(cfg,
                                list_instances.ChooseOneRemoteInstance(cfg),
-                               args.instance_id)
+                               args.instance_id,
+                               args.powerwash)
diff --git a/restart/restart_args.py b/restart/restart_args.py
index bc114d36..b63904a3 100644
--- a/restart/restart_args.py
+++ b/restart/restart_args.py
@@ -55,5 +55,11 @@ def GetRestartArgParser(subparser):
         required=False,
         default=1,
         help="The instance id of the remote instance that need to be restart.")
+    restart_parser.add_argument(
+        "--powerwash",
+        dest="powerwash",
+        action="store_true",
+        required=False,
+        help="Erase all userdata in the AVD.")
 
     return restart_parser
diff --git a/restart/restart_test.py b/restart/restart_test.py
index 659bdbe6..1448a841 100644
--- a/restart/restart_test.py
+++ b/restart/restart_test.py
@@ -34,12 +34,13 @@ class RestartTest(driver_test_lib.BaseDriverTest):
         # Test case with provided instance name.
         args.instance_name = "instance_1"
         args.instance_id = 1
+        args.powerwash = False
         self.Patch(config, "GetAcloudConfig", return_value=cfg)
         self.Patch(list_instances, "GetInstancesFromInstanceNames",
                    return_value=[instance_obj])
         restart.Run(args)
         mock_restart.assert_has_calls([
-            mock.call(cfg, instance_obj, args.instance_id)])
+            mock.call(cfg, instance_obj, args.instance_id, args.powerwash)])
 
         # Test case for user select one instance to restart AVD.
         selected_instance = mock.MagicMock()
@@ -48,7 +49,7 @@ class RestartTest(driver_test_lib.BaseDriverTest):
         args.instance_name = None
         restart.Run(args)
         mock_restart.assert_has_calls([
-            mock.call(cfg, selected_instance, args.instance_id)])
+            mock.call(cfg, selected_instance, args.instance_id, args.powerwash)])
 
 
 if __name__ == '__main__':