Bug 1763188 - Add Snap support using TC builds

mozregression/branches.py

@@ -78,6 +78,12 @@ def create_branches():
     ):
         for alias in aliases:
             branches.set_alias(alias, name)
+
+    branches.set_branch("snap-nightly", "mozilla-central")
+    branches.set_branch("snap-beta", "mozilla-central")
+    branches.set_branch("snap-stable", "mozilla-central")
+    branches.set_branch("snap-esr", "mozilla-central")
+
     return branches
 
 

mozregression/cli.py

@@ -417,6 +417,12 @@ def create_parser(defaults):
         help="Helps to write the configuration file.",
     )
 
+    parser.add_argument(
+        "--allow-sudo",
+        action="store_true",
+        help="Allow the use of sudo for snap install/remove operations (otherwise, you will be prompted on each)",
+    )
+
     parser.add_argument("--debug", "-d", action="store_true", help="Show the debug output.")
 
     return parser
@@ -598,6 +604,11 @@ def validate(self):
                 f"`--arch` required for specified app ({options.app}). "
                 f"Please specify one of {', '.join(arch_options[options.app])}."
             )
+        elif options.allow_sudo == True and options.app != "firefox-snap":
+            raise MozRegressionError(
+                f"--allow-sudo specified for app ({options.app}), but only valid for "
+                f"firefox-snap. Please verify your config."
+            )
 
         fetch_config = create_config(
             options.app, mozinfo.os, options.bits, mozinfo.processor, options.arch

mozregression/fetch_configs.py

@@ -752,3 +752,71 @@ def build_regex(self):
             part = "mac"
         psuffix = "-asan" if "asan" in self.build_type else ""
         return r"jsshell-%s%s\.zip$" % (part, psuffix)
+
+TIMESTAMP_SNAP_UPSTREAM_BUILD = to_utc_timestamp(datetime.datetime(2023, 7, 26, 9, 39, 21))
+TIMESTAMP_SNAP_INDEX_RENAME = to_utc_timestamp(datetime.datetime(2023, 11, 17, 21, 46, 39))
+
+class FirefoxSnapNightlyConfigMixin(NightlyConfigMixin):
+    def _get_nightly_repo(self, date):
+        return "mozilla-central"
+
+# we also have beta, stable and ESR how can we use them?
+class FirefoxSnapIntegrationConfigMixin(IntegrationConfigMixin):
+    def _idx_key(self, date):
+        index_name = ""
+        branch_name = ""
+
+        if self.integration_branch == "snap-nightly":
+            branch_name = "nightly"
+        elif self.integration_branch == "snap-beta":
+            branch_name = "beta"
+        elif self.integration_branch == "snap-stable":
+            branch_name = "stable"
+        elif self.integration_branch == "snap-esr":
+            branch_name = "esr"
+        else:
+            # how can we get the beta/release/esr builds from tc's cron ?
+            raise NotImplementedError
+
+        if date < TIMESTAMP_SNAP_UPSTREAM_BUILD:
+            raise ValueError
+        elif date >= TIMESTAMP_SNAP_UPSTREAM_BUILD and date < TIMESTAMP_SNAP_INDEX_RENAME:
+            index_base = ""
+        elif date >= TIMESTAMP_SNAP_INDEX_RENAME:
+            index_base = "amd64-"
+
+        return "{}{}".format(index_base, branch_name)
+
+    def tk_routes(self, push):
+        for build_type in self.build_types:
+            name = "gecko.v2.mozilla-central.revision.{}.firefox.{}{}".format(
+                push.changeset,
+                self._idx_key(push.timestamp),
+                "-{}".format(build_type) if build_type != "opt" and build_type != "shippable" else "",
+            )
+            yield name
+            self._inc_used_build()
+        return
+
+
+class SnapCommonConfig(CommonConfig):
+    def should_use_archive(self):
+        """
+        We only want to use TaskCluster builds
+        """
+        return False
+
+    def build_regex(self):
+        return r"(firefox_.*)\.snap"
+
+@REGISTRY.register("firefox-snap")
+class FirefoxSnapConfig(SnapCommonConfig, FirefoxSnapNightlyConfigMixin, FirefoxSnapIntegrationConfigMixin):
+    BUILD_TYPES = ("shippable", "opt", "debug")
+    BUILD_TYPE_FALLBACKS = {
+        "shippable": ("opt",),
+        "opt": ("shippable",),
+    }
+
+    def __init__(self, os, bits, processor, arch):
+        super(FirefoxSnapConfig, self).__init__(os, bits, processor, arch)
+        self.set_build_type("shippable")

mozregression/launchers.py

@@ -7,11 +7,14 @@
 import json
 import os
 import stat
+import subprocess
 import sys
+import tempfile
 import time
 import zipfile
 from abc import ABCMeta, abstractmethod
 from enum import Enum
+from shutil import move
 from subprocess import STDOUT, CalledProcessError, call, check_output
 from threading import Thread
 
@@ -22,7 +25,7 @@
 from mozfile import remove
 from mozlog.structured import get_default_logger, get_proxy_logger
 from mozprofile import Profile, ThunderbirdProfile
-from mozrunner import Runner
+from mozrunner import Runner, GeckoRuntimeRunner
 
 from mozregression.class_registry import ClassRegistry
 from mozregression.errors import LauncherError, LauncherNotRunnable
@@ -338,11 +341,11 @@ def get_app_info(self):
 REGISTRY = ClassRegistry("app_name")
 
 
-def create_launcher(buildinfo):
+def create_launcher(buildinfo, allow_sudo):
     """
     Create and returns an instance launcher for the given buildinfo.
     """
-    return REGISTRY.get(buildinfo.app_name)(buildinfo.build_file, task_id=buildinfo.task_id)
+    return REGISTRY.get(buildinfo.app_name)(buildinfo.build_file, task_id=buildinfo.task_id, allow_sudo=allow_sudo)
 
 
 class FirefoxRegressionProfile(Profile):
@@ -616,3 +619,132 @@ def cleanup(self):
             # always remove tempdir
             if self.tempdir is not None:
                 remove(self.tempdir)
+
+
+# Should this be part of mozrunner ?
+class SnapRunner(GeckoRuntimeRunner):
+    _allow_sudo = False
+    _snap_pkg = None
+
+    def __init__(self, binary, cmdargs, allow_sudo=False, snap_pkg=None, **runner_args):
+        self._allow_sudo = allow_sudo
+        self._snap_pkg = snap_pkg
+        GeckoRuntimeRunner.__init__(self, binary, cmdargs, **runner_args)
+
+    @property
+    def command(self):
+        """
+        Rewrite the command for performing the actual execution with "snap run PKG", keeping everything else
+        """
+        self._command = FirefoxSnapLauncher._get_snap_command(self._allow_sudo, "run", [ self._snap_pkg ] + super().command[1:])
+        return self._command
+
+
+@REGISTRY.register("firefox-snap")
+class FirefoxSnapLauncher(MozRunnerLauncher):
+    profile_class = FirefoxRegressionProfile
+    _id = None
+    snap_pkg = None
+    binary = None
+    allow_sudo = False
+    runner = None
+
+    def __init__(self, dest, **kwargs):
+        self.allow_sudo = kwargs["allow_sudo"]
+        super(MozRunnerLauncher, self).__init__(dest)
+
+    def get_snap_command(self, action, extra):
+        return FirefoxSnapLauncher._get_snap_command(self.allow_sudo, action, extra)
+
+    def _get_snap_command(allow_sudo, action, extra):
+        if action not in ("connect", "install", "run", "remove"):
+            raise NotImplementedError
+
+        cmd = []
+        if allow_sudo and action in ("connect", "install", "remove"):
+                cmd += ["sudo"]
+
+        cmd += ["snap", action]
+        cmd += extra
+
+        return cmd
+
+    def _install(self, dest):
+        # From https://snapcraft.io/docs/parallel-installs#heading--naming
+        #  - The instance key needs to be manually appended to the snap name,
+        #    and takes the following format: <snap>_<instance-key>
+        #  - The instance key must match the following regular expression:
+        #    ^[a-z0-9]{1,10}$.
+        self._id = os.path.basename(dest)[0:9]
+
+        self.snap_pkg = "firefox_{}".format(self._id)
+        self.binary = "/snap/{}/current/usr/lib/firefox/firefox".format(self.snap_pkg)
+
+        subprocess.run(self.get_snap_command("install", ["--name", self.snap_pkg, "--dangerous", "{}".format(dest)]), check=True)
+        self._fix_connections()
+
+        self.binarydir = os.path.dirname(self.binary)
+        self.appdir = os.path.normpath(os.path.join(self.binarydir, "..", ".."))
+
+        # On Snap updates are already disabled
+
+    def _fix_connections(self):
+        existing = {}
+        for line in subprocess.getoutput("snap connections {}".format(self.snap_pkg)).splitlines()[1:]:
+            interface, plug, slot, _ = line.split()
+            existing[plug] = slot
+
+        for line in subprocess.getoutput("snap connections firefox").splitlines()[1:]:
+            interface, plug, slot, _ = line.split()
+            plug = plug.replace("firefox:", "{}:".format(self.snap_pkg))
+            slot = slot.replace("firefox:", "{}:".format(self.snap_pkg))
+            if existing[plug] == "-":
+                if plug != "-" and slot != "-":
+                    cmd = self.get_snap_command("connect", ["{}".format(plug), "{}".format(slot)])
+                    print(cmd)
+                    subprocess.run(cmd)
+
+    def _create_profile(self, profile=None, addons=(), preferences=None):
+        """
+        Let's create a profile as usual, but rewrite its path to be in Snap's dir because it looks like
+        MozProfile class will consider a profile=xxx to be a pre-existing one
+        """
+        real_profile = super(MozRunnerLauncher, self)._create_profile(profile, addons, preferences)
+        snap_profile_dir = os.path.abspath(os.path.expanduser("~/snap/{}/common/.mozilla/firefox/".format(self.snap_pkg)))
+        if not os.path.exists(snap_profile_dir):
+            os.makedirs(snap_profile_dir)
+        profile_dir_name = os.path.basename(real_profile.profile)
+        snap_profile = os.path.join(snap_profile_dir, profile_dir_name)
+        move(real_profile.profile, snap_profile_dir)
+        real_profile.profile = snap_profile
+        return real_profile
+
+    def _start(
+        self,
+        profile=None,
+        addons=(),
+        cmdargs=(),
+        preferences=None,
+        adb_profile_dir=None,
+        allow_sudo=False,
+    ):
+        profile = self._create_profile(profile=profile, addons=addons, preferences=preferences)
+
+        LOG.info("Launching %s [%s]" % (self.binary, self.allow_sudo))
+        self.runner = SnapRunner(binary=self.binary, cmdargs=cmdargs, profile=profile, allow_sudo=self.allow_sudo, snap_pkg=self.snap_pkg)
+        self.runner.start()
+
+    def _wait(self):
+        self.runner.wait()
+
+    def _stop(self):
+        self.runner.stop()
+        # release the runner since it holds a profile reference
+        del self.runner
+
+    def cleanup(self):
+        try:
+            Launcher.cleanup(self)
+        finally:
+            subprocess.run(self.get_snap_command("remove", [self.snap_pkg]))
+

mozregression/main.py

@@ -89,6 +89,7 @@ def test_runner(self):
                         cmdargs=self.options.cmdargs,
                         preferences=self.options.preferences,
                         adb_profile_dir=self.options.adb_profile_dir,
+                        allow_sudo=self.options.allow_sudo,
                     )
                 )
             else:

mozregression/test_runner.py

@@ -20,7 +20,7 @@
 LOG = get_proxy_logger("Test Runner")
 
 
-def create_launcher(build_info):
+def create_launcher(build_info, allow_sudo=False):
     """
     Create and returns a :class:`mozregression.launchers.Launcher`.
     """
@@ -36,7 +36,7 @@ def create_launcher(build_info):
         )
     LOG.info("Running %s build %s" % (build_info.repo_name, desc))
 
-    return mozlauncher(build_info)
+    return mozlauncher(build_info, allow_sudo)
 
 
 class TestRunner(metaclass=ABCMeta):
@@ -117,7 +117,7 @@ def get_verdict(self, build_info, allow_back):
         return verdict[0]
 
     def evaluate(self, build_info, allow_back=False):
-        with create_launcher(build_info) as launcher:
+        with create_launcher(build_info, self.launcher_kwargs["allow_sudo"]) as launcher:
             launcher.start(**self.launcher_kwargs)
             build_info.update_from_app_info(launcher.get_app_info())
             verdict = self.get_verdict(build_info, allow_back)
@@ -131,7 +131,7 @@ def evaluate(self, build_info, allow_back=False):
         return verdict
 
     def run_once(self, build_info):
-        with create_launcher(build_info) as launcher:
+        with create_launcher(build_info, self.launcher_kwargs["allow_sudo"]) as launcher:
             launcher.start(**self.launcher_kwargs)
             build_info.update_from_app_info(launcher.get_app_info())
             return launcher.wait()
@@ -190,7 +190,7 @@ def __init__(self, command):
         self.command = command
 
     def evaluate(self, build_info, allow_back=False):
-        with create_launcher(build_info) as launcher:
+        with create_launcher(build_info, self.launcher_kwargs["allow_sudo"]) as launcher:
             build_info.update_from_app_info(launcher.get_app_info())
             variables = {k: v for k, v in build_info.to_dict().items()}
             if hasattr(launcher, "binary"):