Internal change

src/java/com/google/devtools/mobileharness/platform/android/instrumentation/AndroidInstrumentationSetting.java

@@ -37,6 +37,7 @@ public static AndroidInstrumentationSetting create(
       boolean prefixAndroidTest,
       boolean noIsolatedStorage,
       boolean useTestStorageService,
+      @Nullable Map<String, String> testArgs,
       boolean enableCoverage) {
     return new AutoValue_AndroidInstrumentationSetting(
         checkNotNull(packageName),
@@ -48,6 +49,7 @@ public static AndroidInstrumentationSetting create(
         prefixAndroidTest,
         noIsolatedStorage,
         useTestStorageService,
+        testArgs,
         enableCoverage);
   }
 
@@ -83,6 +85,10 @@ public static AndroidInstrumentationSetting create(
   /** Whether to set useTestStorageService in instrumentation command */
   public abstract boolean useTestStorageService();
 
+  /** Test Args for am instrumentation, or null. */
+  @Nullable
+  public abstract Map<String, String> testArgs();
+
   /** Whether to enable coverage for the instrumentation command. */
   public abstract boolean enableCoverage();
 }

src/java/com/google/devtools/mobileharness/platform/android/instrumentation/AndroidInstrumentationUtil.java

@@ -336,6 +336,12 @@ public String instrument(
 
     if (useTestStorageService) {
       command.append(" -e useTestStorageService true");
+    } else {
+      if (instrumentationSetting.testArgs() != null) {
+        for (Entry<String, String> entry : instrumentationSetting.testArgs().entrySet()) {
+          command.append(" -e ").append(entry.getKey()).append(' ').append(entry.getValue());
+        }
+      }
     }
 
     if (enableCoverage) {

src/java/com/google/wireless/qa/mobileharness/shared/api/decorator/AndroidSwitchLanguageDecorator.java

@@ -284,6 +284,7 @@ private void switchLocale(
                 /* prefixAndroidTest= */ false,
                 /* noIsolatedStorage= */ false,
                 /* useTestStorageService= */ false,
+                /* testArgs= */ null,
                 /* enableCoverage= */ false),
             /* timeout= */ Duration.ofMinutes(1));
     try {

src/java/com/google/wireless/qa/mobileharness/shared/api/driver/AndroidInstrumentation.java

@@ -199,7 +199,11 @@ public void run(TestInfo testInfo) throws MobileHarnessException, InterruptedExc
         externalStoragePath.orElse(null));
 
     // Prepares the test arg.
-    prepareTestArgs(testInfo, externalStoragePath.orElse(null));
+    boolean useTestStorageService =
+        job.params().getBool(AndroidInstrumentationDriverSpec.PARAM_USE_TEST_STORAGE_SERVICE, true);
+    if (useTestStorageService) {
+      prepareTestArgs(testInfo, externalStoragePath.orElse(null));
+    }
 
     // Finalizes the test target.
     String testTarget = getTestTarget(testInfo, optionMaps);
@@ -264,7 +268,8 @@ public void run(TestInfo testInfo) throws MobileHarnessException, InterruptedExc
           deviceSdkVersion,
           instrumentTimeoutMs,
           showRawResults,
-          noIsolatedStorage);
+          noIsolatedStorage,
+          useTestStorageService);
     } else {
       try {
         syncRun(
@@ -277,7 +282,8 @@ public void run(TestInfo testInfo) throws MobileHarnessException, InterruptedExc
             instrumentTimeoutMs,
             buildPackageNames,
             showRawResults,
-            noIsolatedStorage);
+            noIsolatedStorage,
+            useTestStorageService);
       } finally {
         androidInstrumentationUtil.pullInstrumentationFilesFromDevice(
             deviceId, testInfo, externalStoragePath.orElse(null));
@@ -308,7 +314,8 @@ private void asyncRun(
       @Nullable Integer deviceSdkVersion,
       @Nullable Long instrumentTimeoutMs,
       boolean showRawResults,
-      boolean noIsolatedStorage)
+      boolean noIsolatedStorage,
+      boolean useTestStorageService)
       throws MobileHarnessException, InterruptedException {
     JobInfo job = testInfo.jobInfo();
     // When async is enable, only one option map is allowed.
@@ -354,7 +361,8 @@ private void asyncRun(
               job.params()
                   .getBool(AndroidInstrumentationDriverSpec.PARAM_PREFIX_ANDROID_TEST, false),
               noIsolatedStorage,
-              /* useTestStorageService= */ true,
+              useTestStorageService,
+              useTestStorageService ? null : androidInstrumentationUtil.getTestArgs(testInfo),
               job.params().getBool(AndroidInstrumentationDriverSpec.PARAM_ENABLE_COVERAGE, false)),
           Duration.ofMillis(testTimeoutMs));
     } catch (MobileHarnessException e) {
@@ -380,7 +388,8 @@ private void syncRun(
       @Nullable Long instrumentTimeoutMs,
       List<String> buildPackageNames,
       boolean showRawResults,
-      boolean noIsolatedStorage)
+      boolean noIsolatedStorage,
+      boolean useTestStorageService)
       throws MobileHarnessException, InterruptedException {
     JobInfo job = testInfo.jobInfo();
     Properties properties = testInfo.properties();
@@ -472,7 +481,8 @@ private void syncRun(
                     job.params()
                         .getBool(AndroidInstrumentationDriverSpec.PARAM_PREFIX_ANDROID_TEST, false),
                     noIsolatedStorage,
-                    /* useTestStorageService= */ true,
+                    useTestStorageService,
+                    useTestStorageService ? null : androidInstrumentationUtil.getTestArgs(testInfo),
                     job.params()
                         .getBool(AndroidInstrumentationDriverSpec.PARAM_ENABLE_COVERAGE, false)),
                 Duration.ofMillis(testTimeoutMs));

src/java/com/google/wireless/qa/mobileharness/shared/api/spec/AndroidInstrumentationDriverSpec.java

@@ -218,6 +218,19 @@ public interface AndroidInstrumentationDriverSpec {
               + " is false.")
   String PARAM_ENABLE_COVERAGE = "enable_coverage";
 
+  @ParamAnnotation(
+      required = false,
+      help =
+          "Whether to use TestStorage service when running the test. By default it is true. When"
+              + " set to false, the test args will be passed to the test via the command line in"
+              + " key-value paris. We still recommend to use TestStorage service, but if you"
+              + " encounter some issues on multi-user scenario, you can try to disable TestStorage"
+              + " service and give it a try. And based on"
+              + " https://developer.android.com/studio/test/command-line#am-instrument-flags, only"
+              + " limited number of test runners support it. The args can then be fetched by"
+              + " InstrumentationRegistry.getArguments().keySet() in the test apk.")
+  String PARAM_USE_TEST_STORAGE_SERVICE = "use_test_storage_service";
+
   /** Resource path of the Android basic_services.apk which is needed for reading test_args. */
   String BASIC_SERVICE_APK_PATH =
       "/com/google/android/apps/common/testing/services/basic_services.apk";

src/javatests/com/google/wireless/qa/mobileharness/shared/api/driver/AndroidInstrumentationTest.java

@@ -108,6 +108,8 @@ public class AndroidInstrumentationTest {
   // NOTICE: the map entries here should match the `OPTIONS` value above.
   private static final ImmutableMap<String, String> OPTION_MAP =
       ImmutableMap.of("key1", "value1", "key2", "value2");
+  private static final ImmutableMap<String, String> TEST_ARGS_MAP =
+      ImmutableMap.of("key1", "value1", "key2", "value2");
 
   private static final String[] BUILD_PACKAGES =
       new String[] {"com.google.spinner", "com.google.dependency"};
@@ -373,6 +375,43 @@ public void run_instrumentStatusError_fail() throws Exception {
             instrumentationLog);
   }
 
+  @Test
+  public void run_instrument_noUseTestStorageService_pass() throws Exception {
+    mockRunTestBasicSteps(TEST_NAME, OPTIONS, false, false);
+    jobInfo.params().add(AndroidInstrumentationDriverSpec.PARAM_PREFIX_ANDROID_TEST, "true");
+    jobInfo.params().add(AndroidInstrumentationDriverSpec.PARAM_USE_TEST_STORAGE_SERVICE, "false");
+    String instrumentationLog =
+        "com.google.codelab.mobileharness.android.hellomobileharness"
+            + ".HelloMobileHarnessTest:.......\n\nTime: 31.281\n\nOK (7 tests)";
+
+    AndroidInstrumentationSetting setting =
+        mockInstrumentSetting(
+            TEST_NAME,
+            OPTION_MAP,
+            /* async= */ false,
+            /* showRawResults= */ false,
+            /* prefixAndroidTest= */ true,
+            /* noIsolatedStorage= */ false,
+            /* useTestStorageService= */ false,
+            /* enableCoverage= */ false,
+            TEST_ARGS_MAP);
+    when(androidInstrumentationUtil.instrument(
+            DEVICE_ID, DEFAULT_SDK_VERSION, setting, TEST_TIMEOUT))
+        .thenReturn(instrumentationLog);
+    when(androidInstrumentationUtil.getTestArgs(testInfo)).thenReturn(TEST_ARGS_MAP);
+
+    driver.run(testInfo);
+
+    assertThat(testInfo.resultWithCause().get().type()).isEqualTo(TestResult.PASS);
+    verify(fileUtil)
+        .writeToFile(
+            PathUtil.join(GEN_FILE_DIR_PATH, DEFAULT_INSTRUMENTATION_LOG_FILE_NAME),
+            instrumentationLog);
+    verify(androidInstrumentationUtil)
+        .prepareServicesApks(eq(testInfo), eq(device), eq(DEFAULT_SDK_VERSION), any(), any());
+    verify(androidInstrumentationUtil).enableLegacyStorageForApk(DEVICE_ID, TEST_PACKAGE);
+  }
+
   @Test
   public void runBroadcastInstallMessage() throws Exception {
     mockRunTestBasicSteps(TEST_NAME, OPTIONS, false, false);
@@ -676,6 +715,28 @@ private AndroidInstrumentationSetting mockInstrumentSetting(
       boolean noIsolatedStorage,
       boolean useTestStorageService,
       boolean enableCoverage) {
+    return mockInstrumentSetting(
+        className,
+        otherOptions,
+        async,
+        showRawResults,
+        prefixAndroidTest,
+        noIsolatedStorage,
+        useTestStorageService,
+        enableCoverage,
+        /* testArgs= */ null);
+  }
+
+  private AndroidInstrumentationSetting mockInstrumentSetting(
+      @Nullable String className,
+      @Nullable Map<String, String> otherOptions,
+      boolean async,
+      boolean showRawResults,
+      boolean prefixAndroidTest,
+      boolean noIsolatedStorage,
+      boolean useTestStorageService,
+      boolean enableCoverage,
+      @Nullable Map<String, String> testArgs) {
     return AndroidInstrumentationSetting.create(
         TEST_PACKAGE,
         RUNNER,
@@ -686,6 +747,7 @@ private AndroidInstrumentationSetting mockInstrumentSetting(
         prefixAndroidTest,
         noIsolatedStorage,
         useTestStorageService,
+        testArgs,
         enableCoverage);
   }
 }