Internal change

src/devtools/mobileharness/infra/ats/console/controller/proto/session_plugin.proto

@@ -188,6 +188,9 @@ message RunCommand {
 
   // Optional. Whether to enable default logs for TF tests.
   optional bool enable_default_logs = 37;
+
+  // Optional. Whether to enable Mobly resultstore upload.
+  optional bool enable_mobly_resultstore_upload = 38;
 }
 
 message RunCommandState {

src/java/com/google/devtools/mobileharness/infra/ats/common/BUILD

@@ -155,6 +155,7 @@ java_library(
         "//src/java/com/google/devtools/mobileharness/platform/android/sdktool/adb:enums",
         "//src/java/com/google/devtools/mobileharness/platform/android/xts/common/util:abi_util",
         "//src/java/com/google/devtools/mobileharness/platform/android/xts/common/util:mobly_test_loader",
+        "//src/java/com/google/devtools/mobileharness/platform/android/xts/common/util:xts_constants",
         "//src/java/com/google/devtools/mobileharness/platform/android/xts/common/util:xts_dir_util",
         "//src/java/com/google/devtools/mobileharness/platform/android/xts/config:configuration_util",
         "//src/java/com/google/devtools/mobileharness/platform/android/xts/config:module_configuration_helper",

src/java/com/google/devtools/mobileharness/infra/ats/common/SessionRequestHandlerUtil.java

@@ -65,6 +65,7 @@
 import com.google.devtools.mobileharness.platform.android.sdktool.adb.AndroidProperty;
 import com.google.devtools.mobileharness.platform.android.xts.common.util.AbiUtil;
 import com.google.devtools.mobileharness.platform.android.xts.common.util.MoblyTestLoader;
+import com.google.devtools.mobileharness.platform.android.xts.common.util.XtsConstants;
 import com.google.devtools.mobileharness.platform.android.xts.common.util.XtsDirUtil;
 import com.google.devtools.mobileharness.platform.android.xts.config.ConfigurationUtil;
 import com.google.devtools.mobileharness.platform.android.xts.config.ModuleConfigurationHelper;
@@ -1165,6 +1166,9 @@ public ImmutableList<JobInfo> createXtsNonTradefedJobs(
         }
         addSessionClientIdToJobInfo(jobInfo, sessionRequestInfo);
         extraJobProperties.forEach((key, value) -> jobInfo.properties().add(key, value));
+        if (sessionRequestInfo.isMoblyResultstoreUploadEnabled().orElse(false)) {
+          jobInfo.properties().add(XtsConstants.IS_MOBLY_RESULTSTORE_UPLOAD_ENABLED, "true");
+        }
         printCreatedJobInfo(jobInfo, /* isTf= */ false);
         jobInfos.add(jobInfo);
       }

src/java/com/google/devtools/mobileharness/infra/ats/common/SessionRequestInfo.java

@@ -159,6 +159,8 @@ public abstract class SessionRequestInfo {
 
   public abstract Optional<Boolean> isXtsDynamicDownloadEnabled();
 
+  public abstract Optional<Boolean> isMoblyResultstoreUploadEnabled();
+
   public abstract ImmutableMap<String, String> xtsSuiteInfo();
 
   public abstract Optional<TestSuiteInfo> testSuiteInfo();
@@ -316,6 +318,9 @@ public abstract Builder setGivenMatchedNonTfModules(
 
     public abstract Builder setIsXtsDynamicDownloadEnabled(boolean isXtsDynamicDownloadEnabled);
 
+    public abstract Builder setIsMoblyResultstoreUploadEnabled(
+        boolean isMoblyResultstoreUploadEnabled);
+
     public abstract Builder setXtsSuiteInfo(ImmutableMap<String, String> xtsSuiteInfo);
 
     public abstract Builder setTestSuiteInfo(TestSuiteInfo testSuiteInfo);

src/java/com/google/devtools/mobileharness/infra/ats/console/command/RunCommand.java

@@ -817,6 +817,9 @@ int runWithCommand(ImmutableList<String> command)
     if (Flags.instance().enableXtsDynamicDownloader.getNonNull()) {
       runCommand.setEnableXtsDynamicDownload(true);
     }
+    if (Flags.instance().enableMoblyResultstoreUpload.getNonNull()) {
+      runCommand.setEnableMoblyResultstoreUpload(true);
+    }
     if (deviceTypeOptionsGroup != null) {
       if (deviceTypeOptionsGroup.runTestOnEmulator) {
         runCommand.setDeviceType(DeviceType.EMULATOR);

src/java/com/google/devtools/mobileharness/infra/ats/console/controller/sessionplugin/RunCommandHandler.java

@@ -414,6 +414,9 @@ SessionRequestInfo generateSessionRequestInfo(RunCommand runCommand)
     if (runCommand.getEnableXtsDynamicDownload()) {
       builder.setIsXtsDynamicDownloadEnabled(true);
     }
+    if (runCommand.getEnableMoblyResultstoreUpload()) {
+      builder.setIsMoblyResultstoreUploadEnabled(true);
+    }
     builder.setEnableDefaultLogs(runCommand.getEnableDefaultLogs());
 
     sessionInfo

src/java/com/google/devtools/mobileharness/infra/controller/test/local/BUILD

@@ -60,6 +60,7 @@ java_library(
     name = "local_test_builtin_plugins",
     srcs = ["LocalTestBuiltinPlugins.java"],
     runtime_deps = [
+        "//src/java/com/google/devtools/mobileharness/platform/android/xts/plugin:mobly_resultstore_upload_plugin",
     ],
     deps = [
         "//src/java/com/google/devtools/mobileharness/api/model/error",

src/java/com/google/devtools/mobileharness/infra/controller/test/local/LocalTestBuiltinPlugins.java

@@ -88,6 +88,13 @@ ImmutableList<PluginItem<?>> loadBuiltInPlugin(TestInfo testInfo, DirectTestRunn
                   "com.google.devtools.mobileharness.platform.android.xts.plugin.XtsDeviceCompatibilityChecker"),
               EventScope.INTERNAL_PLUGIN));
     }
+    if (isMoblyResultstoreUploadEnabled(testInfo)) {
+      builtinPluginsBuilder.add(
+          PluginItem.create(
+              createBuiltinPlugin(
+                  "com.google.devtools.mobileharness.platform.android.xts.plugin.MoblyResultstoreUploadPlugin"),
+              EventScope.INTERNAL_PLUGIN));
+    }
 
     ImmutableList<PluginItem<?>> builtinPlugins = builtinPluginsBuilder.build();
     for (PluginItem<?> pluginItem : builtinPlugins) {
@@ -137,6 +144,14 @@ private static boolean isXtsDynamicDownloaderEnabled(TestInfo testInfo) {
         .orElse(false);
   }
 
+  private static boolean isMoblyResultstoreUploadEnabled(TestInfo testInfo) {
+    return testInfo
+        .jobInfo()
+        .properties()
+        .getBoolean(XtsConstants.IS_MOBLY_RESULTSTORE_UPLOAD_ENABLED)
+        .orElse(false);
+  }
+
   private static boolean isAtsFileServerUploaderEnabled(TestInfo testInfo) {
     return isAtsModeEnabled()
         && Objects.equals(testInfo.jobInfo().properties().get(Job.CLIENT_TYPE), "olc");

src/java/com/google/devtools/mobileharness/platform/android/xts/common/util/XtsConstants.java

@@ -48,6 +48,10 @@ public class XtsConstants {
 
   public static final String INVOCATION_SUMMARY_FILE_NAME = "invocation_summary.txt";
 
+  /** A MH job property key to indicate whether Mobly resultstore upload is enabled. */
+  public static final String IS_MOBLY_RESULTSTORE_UPLOAD_ENABLED =
+      "is_mobly_resultstore_upload_enabled";
+
   /** A MH job property key to indicate whether xTS dynamic download is enabled. */
   public static final String IS_XTS_DYNAMIC_DOWNLOAD_ENABLED = "is_xts_dynamic_download_enabled";
 

src/java/com/google/devtools/mobileharness/platform/android/xts/plugin/BUILD

@@ -96,3 +96,21 @@ java_library(
         "@maven//:com_google_guava_guava",
     ],
 )
+
+java_library(
+    name = "mobly_resultstore_upload_plugin",
+    srcs = ["MoblyResultstoreUploadPlugin.java"],
+    visibility = [
+        "//src/java/com/google/devtools/mobileharness/infra/controller/test/local:__pkg__",
+    ],
+    deps = [
+        "//src/java/com/google/devtools/mobileharness/api/model/error",
+        "//src/java/com/google/devtools/mobileharness/platform/testbed/mobly/util:mobly_python_venv_util",
+        "//src/java/com/google/devtools/mobileharness/shared/util/command",
+        "//src/java/com/google/devtools/mobileharness/shared/util/file/local",
+        "//src/java/com/google/devtools/mobileharness/shared/util/logging:google_logger",
+        "//src/java/com/google/wireless/qa/mobileharness/shared/controller/event:client_test_events",
+        "//src/java/com/google/wireless/qa/mobileharness/shared/controller/plugin",
+        "@maven//:com_google_guava_guava",
+    ],
+)

src/java/com/google/devtools/mobileharness/platform/android/xts/plugin/MoblyResultstoreUploadPlugin.java

@@ -0,0 +1,109 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.devtools.mobileharness.platform.android.xts.plugin;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.eventbus.Subscribe;
+import com.google.common.flogger.FluentLogger;
+import com.google.devtools.mobileharness.api.model.error.MobileHarnessException;
+import com.google.devtools.mobileharness.platform.testbed.mobly.util.MoblyPythonVenvUtil;
+import com.google.devtools.mobileharness.shared.util.command.Command;
+import com.google.devtools.mobileharness.shared.util.command.CommandException;
+import com.google.devtools.mobileharness.shared.util.command.CommandExecutor;
+import com.google.devtools.mobileharness.shared.util.file.local.LocalFileUtil;
+import com.google.wireless.qa.mobileharness.shared.controller.event.TestEndingEvent;
+import com.google.wireless.qa.mobileharness.shared.controller.plugin.Plugin;
+import com.google.wireless.qa.mobileharness.shared.controller.plugin.Plugin.PluginType;
+import java.nio.file.Path;
+import java.time.Duration;
+
+/** Lab plugin to upload Mobly test results to Resultstore. */
+@Plugin(type = PluginType.LAB)
+public final class MoblyResultstoreUploadPlugin {
+
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
+  private static final Duration UPLOAD_TIMEOUT = Duration.ofMinutes(20);
+
+  private static final String MOBLY_ANDROID_PARTNER_TOOLS_PACKAGE = "mobly-android-partner-tools";
+  private static final String RESULTS_UPLOADER_BIN = "results_uploader";
+
+  private final MoblyPythonVenvUtil moblyPythonVenvUtil;
+  private final CommandExecutor commandExecutor;
+  private final LocalFileUtil localFileUtil;
+
+  public MoblyResultstoreUploadPlugin() {
+    this(new MoblyPythonVenvUtil(), new CommandExecutor(), new LocalFileUtil());
+  }
+
+  @VisibleForTesting
+  MoblyResultstoreUploadPlugin(
+      MoblyPythonVenvUtil moblyPythonVenvUtil,
+      CommandExecutor commandExecutor,
+      LocalFileUtil localFileUtil) {
+    this.moblyPythonVenvUtil = moblyPythonVenvUtil;
+    this.commandExecutor = commandExecutor;
+    this.localFileUtil = localFileUtil;
+  }
+
+  @Subscribe
+  public void onTestEnding(TestEndingEvent event) throws InterruptedException {
+    logger.atInfo().log("Starting Mobly result upload to Resultstore.");
+    Path genFileDir;
+    Path tmpFileDir;
+    try {
+      genFileDir = Path.of(event.getTest().getGenFileDir());
+      tmpFileDir = Path.of(event.getTest().getTmpFileDir());
+    } catch (MobileHarnessException e) {
+      logger.atWarning().withCause(e).log("Failed to get test gen or tmp file dir.");
+      return;
+    }
+
+    Path venvPath = tmpFileDir.resolve("venv");
+    try {
+      localFileUtil.prepareDir(venvPath);
+      Path sysPythonBin = moblyPythonVenvUtil.getPythonPath(null);
+      Path venvPythonBin = moblyPythonVenvUtil.createVenv(sysPythonBin, venvPath);
+      moblyPythonVenvUtil.installPackageFromPypi(
+          venvPythonBin, MOBLY_ANDROID_PARTNER_TOOLS_PACKAGE);
+
+      Path resultsUploaderBin = moblyPythonVenvUtil.resolveVenvBin(venvPath, RESULTS_UPLOADER_BIN);
+      Command uploadCmd =
+          Command.of(
+                  resultsUploaderBin.toString(),
+                  genFileDir.toString(),
+                  "--abort_if_no_creds",
+                  "--label",
+                  event.getTest().locator().getId())
+              .timeout(UPLOAD_TIMEOUT);
+      if (event.getTest().jobInfo().properties().has("xts-module-name")) {
+        uploadCmd =
+            uploadCmd.argsAppended(
+                "--label", event.getTest().jobInfo().properties().get("xts-module-name"));
+      }
+      try {
+        commandExecutor.run(uploadCmd);
+        logger.atInfo().log("Successfully uploaded results to Resultstore.");
+      } catch (CommandException e) {
+        logger.atWarning().withCause(e).log("Failed to run results_uploader command.");
+      }
+    } catch (MobileHarnessException e) {
+      logger.atWarning().withCause(e).log(
+          "Failed to set up Python venv and install packages for result uploading.");
+    }
+  }
+}

src/java/com/google/devtools/mobileharness/platform/testbed/mobly/util/BUILD

@@ -84,7 +84,9 @@ java_library(
     name = "mobly_python_venv_util",
     srcs = ["MoblyPythonVenvUtil.java"],
     visibility = [
+        "//src/java/com/google/devtools/mobileharness/platform/android/xts/plugin:__pkg__",
         "//src/java/com/google/wireless/qa/mobileharness/shared/api/driver:__pkg__",
+        "//src/javatests/com/google/devtools/mobileharness/platform/testbed/mobly/util:__pkg__",
         "//src/javatests/com/google/wireless/qa/mobileharness/shared/api/driver:__pkg__",
     ],
     deps = [

src/java/com/google/devtools/mobileharness/platform/testbed/mobly/util/MoblyPythonVenvUtil.java

@@ -57,8 +57,7 @@ public MoblyPythonVenvUtil() {
   }
 
   /** Locate the path to the system Python binary. */
-  @VisibleForTesting
-  Path getPythonPath(@Nullable String pythonVersion)
+  public Path getPythonPath(@Nullable String pythonVersion)
       throws MobileHarnessException, InterruptedException {
     if (pythonVersion == null) {
       pythonVersion = "python3";
@@ -82,8 +81,7 @@ Path getPythonPath(@Nullable String pythonVersion)
   }
 
   /** Creates the Python virtual environment for installing package dependencies. */
-  @VisibleForTesting
-  Path createVenv(Path sysPythonBin, Path venvPath)
+  public Path createVenv(Path sysPythonBin, Path venvPath)
       throws MobileHarnessException, InterruptedException {
     logger.atInfo().log("Creating Python venv at %s", venvPath);
     List<String> venvCmd = new ArrayList<>();
@@ -155,4 +153,32 @@ void installPythonPkgDeps(
           e);
     }
   }
+
+  /** Installs a Python package from PyPI. */
+  public void installPackageFromPypi(Path venvPythonBin, String pypiPkgName)
+      throws MobileHarnessException, InterruptedException {
+    Duration cmdTimeout = Duration.ofMinutes(10);
+    List<String> pipCmd = new ArrayList<>();
+    pipCmd.add(venvPythonBin.toString());
+    pipCmd.add("-m");
+    pipCmd.add("pip");
+    pipCmd.add("install");
+    pipCmd.add(pypiPkgName);
+    logger.atInfo().log(
+        "Installing Python package %s from PyPI with command: %s.",
+        pypiPkgName, Joiner.on(" ").join(pipCmd));
+    try {
+      executor.run(Command.of(pipCmd).timeout(cmdTimeout.plusMinutes(1)));
+    } catch (CommandException e) {
+      throw new MobileHarnessException(
+          ExtErrorId.MOBLY_AOSP_PIP_INSTALL_ERROR,
+          String.format("Failed to install Python package %s from PyPI.", pypiPkgName),
+          e);
+    }
+  }
+
+  /** Resolves the path to an executable in the virtual environment's bin folder. */
+  public Path resolveVenvBin(Path venvPath, String binName) {
+    return venvPath.resolve("bin").resolve(binName);
+  }
 }

src/java/com/google/devtools/mobileharness/shared/util/flags/Flags.java

@@ -1173,6 +1173,14 @@ public enum ConfigServiceStorageType {
       converter = Flag.BooleanConverter.class)
   public Flag<Boolean> enableMessagingService = enableMessagingServiceDefault;
 
+  private static final Flag<Boolean> enableMoblyResultstoreUploadDefault = Flag.value(false);
+
+  @com.beust.jcommander.Parameter(
+      names = "--enable_mobly_resultstore_upload",
+      description = "Whether to enable Mobly result store upload. Default is false.",
+      converter = Flag.BooleanConverter.class)
+  public Flag<Boolean> enableMoblyResultstoreUpload = enableMoblyResultstoreUploadDefault;
+
   private static final Flag<Boolean> enablePersistentCacheDefault = Flag.value(false);
 
   @com.beust.jcommander.Parameter(