Internal change

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

@@ -103,6 +103,7 @@ java_library(
         "//src/java/com/google/devtools/mobileharness/infra/ats/common:session_request_info",
         "//src/java/com/google/devtools/mobileharness/infra/ats/common:xts_property_name",
         "//src/java/com/google/devtools/mobileharness/infra/ats/common/plan:test_plan_parser",
+        "//src/java/com/google/devtools/mobileharness/infra/ats/server/util:ats_server_session_util",
         "//src/java/com/google/devtools/mobileharness/platform/android/xts/suite:suite_test_filter",
         "//src/java/com/google/devtools/mobileharness/platform/android/xts/suite/retry:previous_result_loader",
         "//src/java/com/google/devtools/mobileharness/platform/android/xts/suite/retry:retry_args",

src/java/com/google/devtools/mobileharness/infra/ats/common/jobcreator/ServerJobCreator.java

@@ -28,6 +28,7 @@
 import com.google.devtools.mobileharness.infra.ats.common.SessionRequestInfo;
 import com.google.devtools.mobileharness.infra.ats.common.XtsPropertyName;
 import com.google.devtools.mobileharness.infra.ats.common.plan.TestPlanParser;
+import com.google.devtools.mobileharness.infra.ats.server.util.AtsServerSessionUtil;
 import com.google.devtools.mobileharness.platform.android.xts.suite.SuiteTestFilter;
 import com.google.devtools.mobileharness.platform.android.xts.suite.retry.PreviousResultLoader;
 import com.google.devtools.mobileharness.platform.android.xts.suite.retry.RetryArgs;
@@ -47,6 +48,7 @@
 public class ServerJobCreator extends XtsJobCreator {
 
   private final PreviousResultLoader previousResultLoader;
+  private final AtsServerSessionUtil atsServerSessionUtil;
 
   @Inject
   ServerJobCreator(
@@ -55,7 +57,8 @@ public class ServerJobCreator extends XtsJobCreator {
       TestPlanParser testPlanParser,
       PreviousResultLoader previousResultLoader,
       RetryGenerator retryGenerator,
-      ModuleShardingArgsGenerator moduleShardingArgsGenerator) {
+      ModuleShardingArgsGenerator moduleShardingArgsGenerator,
+      AtsServerSessionUtil atsServerSessionUtil) {
     super(
         sessionRequestHandlerUtil,
         localFileUtil,
@@ -64,12 +67,18 @@ public class ServerJobCreator extends XtsJobCreator {
         moduleShardingArgsGenerator);
 
     this.previousResultLoader = previousResultLoader;
+    this.atsServerSessionUtil = atsServerSessionUtil;
   }
 
   @Override
   protected void injectEnvSpecificProperties(
-      SessionRequestInfo sessionRequestInfo, Map<String, String> driverParams) {
-    driverParams.put("android_xts_zip", sessionRequestInfo.androidXtsZip().get());
+      SessionRequestInfo sessionRequestInfo, Map<String, String> driverParams)
+      throws InterruptedException {
+    if (atsServerSessionUtil.isLocalMode()) {
+      driverParams.put("xts_root_dir", sessionRequestInfo.xtsRootDir());
+    } else {
+      driverParams.put("android_xts_zip", sessionRequestInfo.androidXtsZip().get());
+    }
     if (sessionRequestInfo.remoteRunnerFilePathPrefix().isPresent()
         && driverParams.containsKey("subplan_xml")) {
       driverParams.put(

src/java/com/google/devtools/mobileharness/infra/ats/common/jobcreator/XtsJobCreator.java

@@ -630,7 +630,8 @@ protected abstract Optional<Path> getPrevSessionTestReportProperties(
       SessionRequestInfo sessionRequestInfo) throws MobileHarnessException;
 
   protected abstract void injectEnvSpecificProperties(
-      SessionRequestInfo sessionRequestInfo, Map<String, String> driverParams);
+      SessionRequestInfo sessionRequestInfo, Map<String, String> driverParams)
+      throws InterruptedException;
 
   protected abstract Path prepareRunRetryTfSubPlanXmlFile(
       SessionRequestInfo sessionRequestInfo, SubPlan subPlan) throws MobileHarnessException;

src/java/com/google/devtools/mobileharness/infra/ats/server/sessionplugin/BUILD

@@ -80,6 +80,7 @@ java_library(
         "//src/java/com/google/devtools/mobileharness/infra/ats/common:xts_type_loader",
         "//src/java/com/google/devtools/mobileharness/infra/ats/common/jobcreator:xts_job_creator",
         "//src/java/com/google/devtools/mobileharness/infra/ats/console/command/parser:command_line_parser",
+        "//src/java/com/google/devtools/mobileharness/infra/ats/server/util:ats_server_session_util",
         "//src/java/com/google/devtools/mobileharness/infra/client/longrunningservice/constant:session_properties",
         "//src/java/com/google/devtools/mobileharness/infra/client/longrunningservice/model:session_info",
         "//src/java/com/google/devtools/mobileharness/infra/lab/common/dir",

src/java/com/google/devtools/mobileharness/infra/ats/server/sessionplugin/NewMultiCommandRequestHandler.java

@@ -60,6 +60,7 @@
 import com.google.devtools.mobileharness.infra.ats.server.proto.ServiceProto.RequestDetail.RequestState;
 import com.google.devtools.mobileharness.infra.ats.server.proto.ServiceProto.TestContext;
 import com.google.devtools.mobileharness.infra.ats.server.proto.ServiceProto.TestResource;
+import com.google.devtools.mobileharness.infra.ats.server.util.AtsServerSessionUtil;
 import com.google.devtools.mobileharness.infra.client.longrunningservice.constant.SessionProperties;
 import com.google.devtools.mobileharness.infra.client.longrunningservice.model.SessionInfo;
 import com.google.devtools.mobileharness.infra.lab.common.dir.DirUtil;
@@ -114,6 +115,8 @@ final class NewMultiCommandRequestHandler {
   /** Timeout setting for slow commands. */
   private static final Duration SLOW_CMD_TIMEOUT = Duration.ofMinutes(10);
 
+  private static final Duration UNZIP_TIMEOUT = Duration.ofHours(1);
+
   private static final DateTimeFormatter TIMESTAMP_DIR_NAME_FORMATTER =
       DateTimeFormatter.ofPattern("uuuu.MM.dd_HH.mm.ss.SSS").withZone(ZoneId.systemDefault());
   private static final String OUTPUT_MANIFEST_FILE_NAME = "FILES";
@@ -141,6 +144,7 @@ final class NewMultiCommandRequestHandler {
   private final XtsTypeLoader xtsTypeLoader;
   private final XtsJobCreator xtsJobCreator;
   private final Sleeper sleeper;
+  private final AtsServerSessionUtil atsServerSessionUtil;
 
   // Cache for storing validated sessionRequestInfo for each command generated by addTradefedJobs(),
   // and mainly used by addNonTradefedJobs() to reduce duplicate sessionRequestInfo generation that
@@ -160,7 +164,8 @@ final class NewMultiCommandRequestHandler {
       Clock clock,
       XtsTypeLoader xtsTypeLoader,
       XtsJobCreator xtsJobCreator,
-      Sleeper sleeper) {
+      Sleeper sleeper,
+      AtsServerSessionUtil atsServerSessionUtil) {
     this.sessionRequestHandlerUtil = sessionRequestHandlerUtil;
     this.sessionResultHandlerUtil = sessionResultHandlerUtil;
     this.localFileUtil = localFileUtil;
@@ -170,6 +175,7 @@ final class NewMultiCommandRequestHandler {
     this.xtsTypeLoader = xtsTypeLoader;
     this.xtsJobCreator = xtsJobCreator;
     this.sleeper = sleeper;
+    this.atsServerSessionUtil = atsServerSessionUtil;
   }
 
   CreateJobsResult createTradefedJobs(NewMultiCommandRequest request, SessionInfo sessionInfo)
@@ -980,6 +986,9 @@ private boolean hasCommandFailed(CommandDetail commandDetail) {
   @CanIgnoreReturnValue
   private String mountOrUnzipXtsZip(String zipFilePath, String targetDirPath)
       throws MobileHarnessException, InterruptedException {
+    if (atsServerSessionUtil.isLocalMode()) {
+      return localFileUtil.unzipFile(zipFilePath, targetDirPath, UNZIP_TIMEOUT);
+    }
     Command command =
         Command.of("fuse-zip", "-r", zipFilePath, targetDirPath).timeout(SLOW_CMD_TIMEOUT);
     try {
@@ -996,13 +1005,17 @@ private String mountOrUnzipXtsZip(String zipFilePath, String targetDirPath)
         logger.atWarning().withCause(e).log(
             "Failed to mount XTS zip file %s to %s. Trying to unzip it instead.",
             zipFilePath, targetDirPath);
-        return localFileUtil.unzipFile(zipFilePath, targetDirPath, SLOW_CMD_TIMEOUT);
+        return localFileUtil.unzipFile(zipFilePath, targetDirPath, UNZIP_TIMEOUT);
       }
     }
   }
 
   private void unmountOrRemoveZipDir(String targetDirPath)
       throws MobileHarnessException, InterruptedException {
+    if (atsServerSessionUtil.isLocalMode()) {
+      localFileUtil.removeFileOrDir(targetDirPath);
+      return;
+    }
     Command command = Command.of("fusermount", "-u", targetDirPath).timeout(SLOW_CMD_TIMEOUT);
     try {
       // Add a 5 seconds delay before unmounting the zip file to avoid race condition and unmount

src/java/com/google/devtools/mobileharness/infra/ats/server/util/AtsServerSessionUtil.java

@@ -0,0 +1,82 @@
+/*
+ * 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.infra.ats.server.util;
+
+import com.google.common.base.Ascii;
+import com.google.common.flogger.FluentLogger;
+import com.google.devtools.mobileharness.api.model.error.MobileHarnessException;
+import com.google.devtools.mobileharness.infra.client.api.controller.device.DeviceQuerier;
+import com.google.devtools.mobileharness.shared.util.network.NetworkUtil;
+import com.google.wireless.qa.mobileharness.shared.constant.Dimension;
+import com.google.wireless.qa.mobileharness.shared.proto.query.DeviceQuery;
+import com.google.wireless.qa.mobileharness.shared.proto.query.DeviceQuery.DeviceInfo;
+import com.google.wireless.qa.mobileharness.shared.proto.query.DeviceQuery.DeviceQueryFilter;
+import com.google.wireless.qa.mobileharness.shared.proto.query.DeviceQuery.DeviceQueryResult;
+import javax.inject.Inject;
+
+/** Utility class for Ats Server Session. */
+public class AtsServerSessionUtil {
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+  private final DeviceQuerier deviceQuerier;
+  private final NetworkUtil networkUtil;
+
+  @Inject
+  AtsServerSessionUtil(DeviceQuerier deviceQuerier, NetworkUtil networkUtil) {
+    this.deviceQuerier = deviceQuerier;
+    this.networkUtil = networkUtil;
+  }
+
+  /**
+   * Checks if the current OLCS instance is running on the same host as the lab.
+   *
+   * @return true if the current OLCS instance is running on the same host as the lab.
+   */
+  public boolean isLocalMode() throws InterruptedException {
+    DeviceQueryResult queryResult;
+    try {
+      queryResult = deviceQuerier.queryDevice(DeviceQueryFilter.getDefaultInstance());
+    } catch (MobileHarnessException e) {
+      logger.atWarning().withCause(e).log("Failed to query device");
+      return false;
+    }
+    String olcsHostName = "";
+    try {
+      olcsHostName = networkUtil.getLocalHostName();
+    } catch (MobileHarnessException ignored) {
+      return false;
+    }
+    if (olcsHostName.isEmpty()) {
+      return false;
+    }
+    for (DeviceInfo deviceInfo : queryResult.getDeviceInfoList()) {
+      String labHostName =
+          deviceInfo.getDimensionList().stream()
+              .filter(
+                  dimension ->
+                      dimension
+                          .getName()
+                          .equals(Ascii.toLowerCase(Dimension.Name.HOST_NAME.name())))
+              .findFirst()
+              .map(DeviceQuery.Dimension::getValue)
+              .orElse("");
+      if (!labHostName.equals(olcsHostName)) {
+        return false;
+      }
+    }
+    return true;
+  }
+}

src/java/com/google/devtools/mobileharness/infra/ats/server/util/BUILD

@@ -0,0 +1,40 @@
+# 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.
+#
+
+load("@rules_java//java:defs.bzl", "java_library")
+
+package(
+    default_applicable_licenses = ["//:license"],
+    default_visibility = [
+        "//src/java/com/google/devtools/mobileharness/infra/ats:__subpackages__",
+        "//src/java/com/google/devtools/mobileharness/platform/android/xts/suite/retry:__pkg__",
+        "//src/javatests/com/google/devtools/mobileharness/infra/ats:__subpackages__",
+    ],
+)
+
+java_library(
+    name = "ats_server_session_util",
+    srcs = ["AtsServerSessionUtil.java"],
+    deps = [
+        "//src/devtools/mobileharness/api/query/proto:device_query_java_proto",
+        "//src/java/com/google/devtools/mobileharness/api/model/error",
+        "//src/java/com/google/devtools/mobileharness/infra/client/api/controller/device:querier",
+        "//src/java/com/google/devtools/mobileharness/shared/util/logging:google_logger",
+        "//src/java/com/google/devtools/mobileharness/shared/util/network",
+        "//src/java/com/google/wireless/qa/mobileharness/shared/constant:dimension",
+        "@maven//:com_google_guava_guava",
+        "@maven//:javax_inject_jsr330_api",
+    ],
+)

src/javatests/com/google/devtools/mobileharness/infra/ats/common/jobcreator/BUILD

@@ -37,6 +37,7 @@ java_library(
         "//src/java/com/google/devtools/mobileharness/infra/ats/common/jobcreator:module_sharding_args_generator",
         "//src/java/com/google/devtools/mobileharness/infra/ats/common/jobcreator:server_job_creator",
         "//src/java/com/google/devtools/mobileharness/infra/ats/common/plan:test_plan_parser",
+        "//src/java/com/google/devtools/mobileharness/infra/ats/server/util:ats_server_session_util",
         "//src/java/com/google/devtools/mobileharness/platform/android/xts/suite:suite_test_filter",
         "//src/java/com/google/devtools/mobileharness/platform/android/xts/suite/retry:previous_result_loader",
         "//src/java/com/google/devtools/mobileharness/platform/android/xts/suite/retry:retry_args",

src/javatests/com/google/devtools/mobileharness/infra/ats/common/jobcreator/ServerJobCreatorTest.java

@@ -37,6 +37,7 @@
 import com.google.devtools.mobileharness.infra.ats.common.XtsPropertyName.Job;
 import com.google.devtools.mobileharness.infra.ats.common.plan.TestPlanParser;
 import com.google.devtools.mobileharness.infra.ats.common.proto.XtsCommonProto.ShardingMode;
+import com.google.devtools.mobileharness.infra.ats.server.util.AtsServerSessionUtil;
 import com.google.devtools.mobileharness.platform.android.xts.suite.retry.PreviousResultLoader;
 import com.google.devtools.mobileharness.platform.android.xts.suite.retry.RetryArgs;
 import com.google.devtools.mobileharness.platform.android.xts.suite.retry.RetryGenerator;
@@ -88,6 +89,7 @@ public final class ServerJobCreatorTest {
   @Rule public final SetFlagsOss flags = new SetFlagsOss();
 
   @Bind @Mock private SessionRequestHandlerUtil sessionRequestHandlerUtil;
+  @Bind @Mock private AtsServerSessionUtil atsServerSessionUtil;
   @Bind @Mock private LocalFileUtil localFileUtil;
   @Bind @Mock private TestPlanParser testPlanParser;
   @Bind @Mock private PreviousResultLoader previousResultLoader;
@@ -148,6 +150,45 @@ public void createXtsTradefedTestJob() throws Exception {
         .containsExactly(Job.XTS_TEST_PLAN, "cts");
   }
 
+  @SuppressWarnings("unchecked")
+  @Test
+  public void createXtsTradefedTestJob_localMode() throws Exception {
+    when(atsServerSessionUtil.isLocalMode()).thenReturn(true);
+    SessionRequestInfo sessionRequestInfo =
+        SessionRequestInfo.builder()
+            .setTestPlan("cts")
+            .setCommandLineArgs("cts")
+            .setXtsType("cts")
+            .setXtsRootDir(XTS_ROOT_DIR_PATH)
+            .setAndroidXtsZip(ANDROID_XTS_ZIP_PATH)
+            .setModuleNames(ImmutableList.of("mock_module"))
+            .build();
+    ArgumentCaptor<Map<String, String>> driverParamsCaptor = ArgumentCaptor.forClass(Map.class);
+
+    when(sessionRequestHandlerUtil.initializeJobConfig(eq(sessionRequestInfo), any()))
+        .thenReturn(JobConfig.getDefaultInstance());
+
+    ImmutableList<TradefedJobInfo> tradefedJobInfoList =
+        jobCreator.createXtsTradefedTestJobInfo(
+            sessionRequestInfo, ImmutableList.of("mock_module"));
+
+    assertThat(tradefedJobInfoList).hasSize(1);
+    verify(sessionRequestHandlerUtil)
+        .initializeJobConfig(eq(sessionRequestInfo), driverParamsCaptor.capture());
+    assertThat(driverParamsCaptor.getValue())
+        .containsExactly(
+            "run_command_args",
+            "-m mock_module",
+            "xts_type",
+            "cts",
+            "xts_root_dir",
+            XTS_ROOT_DIR_PATH,
+            "xts_test_plan",
+            "cts");
+    assertThat(tradefedJobInfoList.get(0).extraJobProperties())
+        .containsExactly(Job.XTS_TEST_PLAN, "cts");
+  }
+
   @Test
   public void createXtsTradefedTestJob_moduleSharding() throws Exception {
     SessionRequestInfo sessionRequestInfo =

src/javatests/com/google/devtools/mobileharness/infra/ats/server/sessionplugin/BUILD

@@ -53,6 +53,7 @@ java_library(
         "//src/java/com/google/devtools/mobileharness/infra/ats/server/sessionplugin:ats_server_session_plugin_lib",
         "//src/java/com/google/devtools/mobileharness/infra/ats/server/sessionplugin:new_multi_command_request_handler",
         "//src/java/com/google/devtools/mobileharness/infra/ats/server/sessionplugin:tradefed_config_generator",
+        "//src/java/com/google/devtools/mobileharness/infra/ats/server/util:ats_server_session_util",
         "//src/java/com/google/devtools/mobileharness/infra/client/api/controller/device:querier",
         "//src/java/com/google/devtools/mobileharness/infra/client/longrunningservice/constant:session_properties",
         "//src/java/com/google/devtools/mobileharness/infra/client/longrunningservice/model:session_event",