Internal change

PiperOrigin-RevId: 838713400

Author: j-gao

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

@@ -191,6 +191,7 @@ java_library(
         "//src/devtools/mobileharness/infra/ats/common/proto:xts_common_java_proto",
         "//src/devtools/mobileharness/infra/ats/server/proto:service_java_proto",
         "//src/devtools/mobileharness/platform/android/xts/proto:configuration_java_proto",
+        "//src/java/com/google/devtools/mobileharness/platform/android/xts/suite:test_suite_helper",
         "//src/java/com/google/devtools/mobileharness/platform/android/xts/suite:test_suite_info",
         "//src/java/com/google/devtools/mobileharness/platform/android/xts/suite/retry:retry_type",
         "//src/java/com/google/devtools/mobileharness/shared/util/auto:auto_value",

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

@@ -705,7 +705,7 @@ public SessionRequestInfo addNonTradefedModuleInfo(SessionRequestInfo sessionReq
         getTestSuiteHelper(xtsRootDir.toString(), xtsType, sessionRequestInfo);
     updatedSessionRequestInfo.setExpandedModules(
         ImmutableMap.copyOf(
-            testSuiteHelper.loadTests(getDeviceInfo(sessionRequestInfo).orElse(null))));
+            testSuiteHelper.loadTests(sessionRequestInfo.deviceInfo().orElse(null))));
 
     ImmutableList<String> modules = sessionRequestInfo.moduleNames();
     ImmutableSet<String> allNonTfModules = getNonTfModules(configsMap);
@@ -719,7 +719,7 @@ public static boolean shouldEnableModuleSharding(SessionRequestInfo sessionReque
         && !SessionRequestHandlerUtil.isRunRetry(sessionRequestInfo.testPlan());
   }
 
-  private Optional<DeviceInfo> getDeviceInfo(SessionRequestInfo sessionRequestInfo)
+  public Optional<DeviceInfo> getDeviceInfo(SessionRequestInfo sessionRequestInfo)
       throws MobileHarnessException, InterruptedException {
     Optional<DeviceInfo> deviceInfo =
         Flags.instance().enableAtsMode.getNonNull()

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

@@ -26,6 +26,7 @@
 import com.google.devtools.mobileharness.infra.ats.server.proto.ServiceProto.TestEnvironment;
 import com.google.devtools.mobileharness.infra.ats.server.proto.ServiceProto.TestResource;
 import com.google.devtools.mobileharness.platform.android.xts.config.proto.ConfigurationProto.Configuration;
+import com.google.devtools.mobileharness.platform.android.xts.suite.TestSuiteHelper.DeviceInfo;
 import com.google.devtools.mobileharness.platform.android.xts.suite.TestSuiteInfo;
 import com.google.devtools.mobileharness.platform.android.xts.suite.retry.RetryType;
 import java.time.Duration;
@@ -51,6 +52,8 @@ public abstract class SessionRequestInfo {
 
   public abstract ImmutableList<String> deviceSerials();
 
+  public abstract Optional<DeviceInfo> deviceInfo();
+
   public abstract ImmutableList<String> excludeDeviceSerials();
 
   public abstract ImmutableList<String> productTypes();
@@ -107,6 +110,7 @@ public abstract class SessionRequestInfo {
   /** Key is config file absolute path. */
   public abstract ImmutableMap<String, Configuration> v2ConfigsMap();
 
+  /** Key is expanded module name (e.g. `arm64-v8a CtsBatteryHealthTestCases`). */
   public abstract ImmutableMap<String, Configuration> expandedModules();
 
   public abstract boolean enableModuleParameter();
@@ -173,6 +177,7 @@ public static Builder builder() {
     return new AutoValue_SessionRequestInfo.Builder()
         .setModuleNames(ImmutableList.of())
         .setDeviceSerials(ImmutableList.of())
+        .setDeviceInfo(Optional.empty())
         .setExcludeDeviceSerials(ImmutableList.of())
         .setProductTypes(ImmutableList.of())
         .setDeviceProperties(ImmutableMap.of())
@@ -216,6 +221,8 @@ public abstract static class Builder {
 
     public abstract Builder setExcludeDeviceSerials(List<String> excludeDeviceSerials);
 
+    public abstract Builder setDeviceInfo(Optional<DeviceInfo> deviceInfo);
+
     public abstract Builder setProductTypes(List<String> productTypes);
 
     public abstract Builder setDeviceProperties(ImmutableMap<String, String> deviceProperties);

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

@@ -196,6 +196,7 @@ java_library(
         "//src/java/com/google/devtools/mobileharness/infra/client/longrunningservice/model:session_info",
         "//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/suite:test_suite_helper",
         "//src/java/com/google/devtools/mobileharness/platform/android/xts/suite:test_suite_info_provider",
         "//src/java/com/google/devtools/mobileharness/platform/android/xts/suite/retry:retry_type",
         "//src/java/com/google/devtools/mobileharness/shared/constant:log_record_importance",

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

@@ -51,6 +51,7 @@
 import com.google.devtools.mobileharness.infra.client.longrunningservice.model.SessionInfo;
 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.suite.TestSuiteHelper.DeviceInfo;
 import com.google.devtools.mobileharness.platform.android.xts.suite.TestSuiteInfoProvider;
 import com.google.devtools.mobileharness.platform.android.xts.suite.retry.RetryType;
 import com.google.devtools.mobileharness.shared.util.file.local.LocalFileUtil;
@@ -417,6 +418,9 @@ SessionRequestInfo generateSessionRequestInfo(RunCommand runCommand)
     sessionInfo
         .getSessionProperty(SessionProperties.PROPERTY_KEY_SESSION_CLIENT_ID)
         .ifPresent(builder::setSessionClientId);
+
+    Optional<DeviceInfo> deviceInfo = sessionRequestHandlerUtil.getDeviceInfo(builder.build());
+    builder.setDeviceInfo(deviceInfo);
     return sessionRequestHandlerUtil.addNonTradefedModuleInfo(builder.build());
   }
 }

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

@@ -598,6 +598,7 @@ private SessionRequestInfo generateSessionRequestInfo(
     }
 
     return sessionRequestInfoBuilder
+        .setDeviceInfo(sessionRequestHandlerUtil.getDeviceInfo(sessionRequestInfoBuilder.build()))
         .setAtsServerTestResources(fileTestResources.build())
         .setAtsServerTestEnvironment(request.getTestEnvironment())
         .build();