From 9fd67f43ed1877767d744e45c67bf42baf05701f Mon Sep 17 00:00:00 2001
From: Myeongjun Go <myeongjun.ko@gmail.com>
Date: Wed, 16 Mar 2022 12:39:55 +0000
Subject: [PATCH] Bug 1732376 - Add a --clean option to remove the raptor venv
 before running r=sparky,perftest-reviewers

Differential Revision: https://phabricator.services.mozilla.com/D139830
---
 .hgignore                                        |  3 +++
 .../mozharness/mozilla/testing/raptor.py         | 16 ++++++++++++++++
 testing/raptor/mach_commands.py                  |  2 ++
 testing/raptor/raptor/cmdline.py                 |  7 +++++++
 4 files changed, 28 insertions(+)

diff --git a/.hgignore b/.hgignore
index cd122bc15775..3b5ba52abd64 100644
--- a/.hgignore
+++ b/.hgignore
@@ -214,6 +214,9 @@ tps_result\.json
 ^testing/raptor/raptor/tests/json/
 ^testing/raptor/webext/raptor/auto_gen_test_config.js
 
+# Ignore condprofile build directory
+^testing/condprofile/build
+
 # Ignore browsertime output directory
 ^browsertime-results
 
diff --git a/testing/mozharness/mozharness/mozilla/testing/raptor.py b/testing/mozharness/mozharness/mozilla/testing/raptor.py
index 4945419670e3..e774f143b752 100644
--- a/testing/mozharness/mozharness/mozilla/testing/raptor.py
+++ b/testing/mozharness/mozharness/mozilla/testing/raptor.py
@@ -530,6 +530,18 @@ class Raptor(
                     "help": "Enable marionette tracing",
                 },
             ],
+            [
+                ["--clean"],
+                {
+                    "action": "store_true",
+                    "dest": "clean",
+                    "default": False,
+                    "help": (
+                        "Clean the python virtualenv (remove, and rebuild) for "
+                        "Raptor before running tests."
+                    ),
+                },
+            ],
         ]
         + testing_config_options
         + copy.deepcopy(code_coverage_config_options)
@@ -648,6 +660,7 @@ class Raptor(
         self.browsertime_video = False
         self.enable_marionette_trace = self.config.get("enable_marionette_trace")
         self.browser_cycles = self.config.get("browser_cycles")
+        self.clean = self.config.get("clean")
 
         for (arg,), details in Raptor.browsertime_options:
             # Allow overriding defaults on the `./mach raptor-test ...` command-line.
@@ -987,6 +1000,9 @@ class Raptor(
         # We need it in-path to import jsonschema later when validating output for perfherder.
         _virtualenv_path = self.config.get("virtualenv_path")
 
+        if self.clean:
+            rmtree(_virtualenv_path)
+
         if self.run_local and os.path.exists(_virtualenv_path):
             self.info("Virtualenv already exists, skipping creation")
             # ffmpeg exists outside of this virtual environment so
diff --git a/testing/raptor/mach_commands.py b/testing/raptor/mach_commands.py
index edea8e78ea8a..b8a09854e8ff 100644
--- a/testing/raptor/mach_commands.py
+++ b/testing/raptor/mach_commands.py
@@ -66,6 +66,7 @@ class RaptorRunner(MozbuildObject):
         self.device_name = kwargs["device_name"]
         self.enable_marionette_trace = kwargs["enable_marionette_trace"]
         self.browsertime_visualmetrics = kwargs["browsertime_visualmetrics"]
+        self.clean = kwargs["clean"]
 
         if Conditions.is_android(self) or kwargs["app"] in ANDROID_BROWSERS:
             self.binary_path = None
@@ -178,6 +179,7 @@ class RaptorRunner(MozbuildObject):
             "enable_marionette_trace": self.enable_marionette_trace,
             "browsertime_visualmetrics": self.browsertime_visualmetrics,
             "mozbuild_path": get_state_dir(),
+            "clean": self.clean,
         }
 
         sys.path.insert(0, os.path.join(self.topsrcdir, "tools", "browsertime"))
diff --git a/testing/raptor/raptor/cmdline.py b/testing/raptor/raptor/cmdline.py
index 8a873596b404..50665893fa6a 100644
--- a/testing/raptor/raptor/cmdline.py
+++ b/testing/raptor/raptor/cmdline.py
@@ -429,6 +429,13 @@ def create_parser(mach_interface=False):
         default=False,
         help="Enable marionette tracing",
     )
+    add_arg(
+        "--clean",
+        dest="clean",
+        action="store_true",
+        default=False,
+        help="Clean the python virtualenv (remove, and rebuild) for Raptor before running tests.",
+    )
 
     add_logging_group(parser)
     return parser