crosperf: Add keep_stateful experiment field

Experiment field "keep_stateful=true" flags the image_chromeos script to
drop "--clobber-stateful" flag in `cros flash` which was always on
before.
This is especially helpful when running tests on a DUT connected via
Wi-Fi. There are also other uses cases where we want to keep the
stateful information.

BUG=b:277269271
TEST=tested locally with crosperf

Change-Id: I11bf25a4c84ac6dce0abbf04a340e1bca6007f6f
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/4484929
Reviewed-by: George Burgess <gbiv@chromium.org>
Tested-by: Denis Nikitin <denik@chromium.org>
Commit-Queue: Denis Nikitin <denik@chromium.org>
diff --git a/crosperf/crosperf_unittest.py b/crosperf/crosperf_unittest.py
index 9f5cb0d..bbcb171 100755
--- a/crosperf/crosperf_unittest.py
+++ b/crosperf/crosperf_unittest.py
@@ -67,7 +67,7 @@
         settings = crosperf.ConvertOptionsToSettings(options)
         self.assertIsNotNone(settings)
         self.assertIsInstance(settings, settings_factory.GlobalSettings)
-        self.assertEqual(len(settings.fields), 41)
+        self.assertEqual(len(settings.fields), 42)
         self.assertTrue(settings.GetField("rerun"))
         argv = ["crosperf/crosperf.py", "temp.exp"]
         options, _ = parser.parse_known_args(argv)
diff --git a/crosperf/experiment.py b/crosperf/experiment.py
index 9973f7e..7d35e31 100644
--- a/crosperf/experiment.py
+++ b/crosperf/experiment.py
@@ -44,6 +44,7 @@
         ignore_min_max,
         crosfleet,
         dut_config,
+        keep_stateful: bool,
         no_lock: bool,
     ):
         self.name = name
@@ -101,7 +102,11 @@
         if test_flag.GetTestMode():
             machine_manager_fn = MockMachineManager
         self.machine_manager = machine_manager_fn(
-            chromeos_root, acquire_timeout, log_level, locks_directory
+            chromeos_root,
+            acquire_timeout,
+            log_level,
+            locks_directory,
+            keep_stateful=keep_stateful,
         )
         self.l = logger.GetLogger(log_dir)
 
@@ -137,7 +142,6 @@
         for label in self.labels:
             for benchmark in self.benchmarks:
                 for iteration in range(1, benchmark.iterations + 1):
-
                     benchmark_run_name = "%s: %s (%s)" % (
                         label.name,
                         benchmark.name,
diff --git a/crosperf/experiment_factory.py b/crosperf/experiment_factory.py
index 6d2ca0a..e89adb8 100644
--- a/crosperf/experiment_factory.py
+++ b/crosperf/experiment_factory.py
@@ -177,6 +177,7 @@
             "turbostat": global_settings.GetField("turbostat"),
             "top_interval": global_settings.GetField("top_interval"),
         }
+        keep_stateful = global_settings.GetField("keep_stateful")
 
         # Default cache hit conditions. The image checksum in the cache and the
         # computed checksum of the image must match. Also a cache file must exist.
@@ -521,6 +522,7 @@
             ignore_min_max,
             crosfleet,
             dut_config,
+            keep_stateful,
             no_lock=no_lock,
         )
 
diff --git a/crosperf/experiment_factory_unittest.py b/crosperf/experiment_factory_unittest.py
index 2196fcd..87e8c4f 100755
--- a/crosperf/experiment_factory_unittest.py
+++ b/crosperf/experiment_factory_unittest.py
@@ -420,10 +420,13 @@
         )
         self.assertEqual(exp.labels[0].autotest_path, "/tmp/autotest")
         self.assertEqual(exp.labels[0].board, "lumpy")
+        self.assertEqual(exp.machine_manager.keep_stateful, False)
 
         # Second test: Remotes listed in labels.
         test_flag.SetTestMode(True)
         label_settings.SetField("remote", "chromeos1.cros chromeos2.cros")
+        # Also verify keep_stateful.
+        global_settings.SetField("keep_stateful", "true")
         exp = ef.GetExperiment(mock_experiment_file, "", "")
         self.assertCountEqual(
             exp.remote,
@@ -434,6 +437,9 @@
                 "chromeos2.cros",
             ],
         )
+        # keep_stateful is propagated to machine_manager which flashes the
+        # images.
+        self.assertEqual(exp.machine_manager.keep_stateful, True)
 
         # Third test: Automatic fixing of bad  logging_level param:
         global_settings.SetField("logging_level", "really loud!")
diff --git a/crosperf/machine_manager.py b/crosperf/machine_manager.py
index ffb0b5e..17db64f 100644
--- a/crosperf/machine_manager.py
+++ b/crosperf/machine_manager.py
@@ -221,6 +221,7 @@
         locks_dir,
         cmd_exec=None,
         lgr=None,
+        keep_stateful: bool = False,
     ):
         self._lock = threading.RLock()
         self._all_machines = []
@@ -233,6 +234,7 @@
         self.acquire_timeout = acquire_timeout
         self.log_level = log_level
         self.locks_dir = locks_dir
+        self.keep_stateful = keep_stateful
         self.ce = cmd_exec or command_executer.GetCommandExecuter(
             log_level=self.log_level
         )
@@ -282,14 +284,16 @@
         image_chromeos_args = [
             image_chromeos.__file__,
             "--no_lock",
-            "--chromeos_root=%s" % chromeos_root,
-            "--image=%s" % label.chromeos_image,
-            "--image_args=%s" % label.image_args,
-            "--remote=%s" % machine.name,
-            "--logging_level=%s" % self.log_level,
+            f"--chromeos_root={chromeos_root}",
+            f"--image={label.chromeos_image}",
+            f"--image_args={label.image_args}",
+            f"--remote={machine.name}",
+            f"--logging_level={self.log_level}",
         ]
         if label.board:
-            image_chromeos_args.append("--board=%s" % label.board)
+            image_chromeos_args.append(f"--board={label.board}")
+        if self.keep_stateful:
+            image_chromeos_args.append("--keep_stateful")
 
         # Currently can't image two machines at once.
         # So have to serialized on this lock.
@@ -729,9 +733,20 @@
 class MockMachineManager(MachineManager):
     """Mock machine manager class."""
 
-    def __init__(self, chromeos_root, acquire_timeout, log_level, locks_dir):
+    def __init__(
+        self,
+        chromeos_root,
+        acquire_timeout,
+        log_level,
+        locks_dir,
+        keep_stateful: bool = False,
+    ):
         super(MockMachineManager, self).__init__(
-            chromeos_root, acquire_timeout, log_level, locks_dir
+            chromeos_root,
+            acquire_timeout,
+            log_level,
+            locks_dir,
+            keep_stateful=keep_stateful,
         )
 
     def _TryToLockMachine(self, cros_machine):
diff --git a/crosperf/settings_factory.py b/crosperf/settings_factory.py
index 2bf4119..b34f0b1 100644
--- a/crosperf/settings_factory.py
+++ b/crosperf/settings_factory.py
@@ -559,6 +559,15 @@
                 " Useful when lock is held externally, say with crosfleet.",
             )
         )
+        self.AddField(
+            BooleanField(
+                "keep_stateful",
+                default=False,
+                description="When flashing a ChromeOS image keep the stateful"
+                " partition, i.e. don't use --clobber-stateful. This option"
+                " is useful to keep ssh keys, wi-fi settings and so on.",
+            )
+        )
 
 
 class SettingsFactory(object):
diff --git a/crosperf/settings_factory_unittest.py b/crosperf/settings_factory_unittest.py
index 4b7c725..a6771c0 100755
--- a/crosperf/settings_factory_unittest.py
+++ b/crosperf/settings_factory_unittest.py
@@ -49,7 +49,7 @@
     def test_init(self):
         res = settings_factory.GlobalSettings("g_settings")
         self.assertIsNotNone(res)
-        self.assertEqual(len(res.fields), 41)
+        self.assertEqual(len(res.fields), 42)
         self.assertEqual(res.GetField("name"), "")
         self.assertEqual(res.GetField("board"), "")
         self.assertEqual(res.GetField("crosfleet"), False)
@@ -115,7 +115,7 @@
             "global", "global"
         )
         self.assertIsInstance(g_settings, settings_factory.GlobalSettings)
-        self.assertEqual(len(g_settings.fields), 41)
+        self.assertEqual(len(g_settings.fields), 42)
 
 
 if __name__ == "__main__":
diff --git a/image_chromeos.py b/image_chromeos.py
index 5d22269..3830059 100755
--- a/image_chromeos.py
+++ b/image_chromeos.py
@@ -49,12 +49,12 @@
         command, chromeos_root=chromeos_root, machine=remote
     )
     logger.GetLogger().LogFatalIf(
-        ret == 255, "Failed ssh to %s (for checking cherrypy)" % remote
+        ret == 255, f"Failed ssh to {remote} (for checking cherrypy)"
     )
     logger.GetLogger().LogFatalIf(
         ret != 0,
-        "Failed to find cherrypy or ctypes on remote '{}', "
-        "cros flash cannot work.".format(remote),
+        f"Failed to find cherrypy or ctypes on '{remote}', "
+        "cros flash cannot work.",
     )
 
 
@@ -142,6 +142,13 @@
         "'quiet', 'average', and 'verbose'.",
     )
     parser.add_argument("-a", "--image_args", dest="image_args")
+    parser.add_argument(
+        "--keep_stateful",
+        dest="keep_stateful",
+        default=False,
+        action="store_true",
+        help="Do not clobber the stateful partition.",
+    )
 
     options = parser.parse_args(argv[1:])
 
@@ -261,9 +268,13 @@
                 "cros",
                 "flash",
                 "--board=%s" % board,
-                "--clobber-stateful",
-                options.remote,
             ]
+            if not options.keep_stateful:
+                cros_flash_args.append("--clobber-stateful")
+            # New arguments should be added here.
+
+            # The last two arguments are positional and have to be at the end.
+            cros_flash_args.append(options.remote)
             if local_image:
                 cros_flash_args.append(chroot_image)
             else: