From a3e139c68a5734c29c7c129dd7a387a185f8781c Mon Sep 17 00:00:00 2001
From: Vinit Nayak <peanutbutter@google.com>
Date: Thu, 18 May 2023 14:33:05 -0700
Subject: [PATCH] Install test app for work profile user

Bug: 243688989
Change-Id: I3ecd0c9eefdbd84d2b9da45680b9efb078c3fb67
---
 tests/src/com/android/launcher3/ui/WorkProfileTest.java | 5 ++++-
 tests/src/com/android/launcher3/util/TestUtil.java      | 7 ++++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/tests/src/com/android/launcher3/ui/WorkProfileTest.java b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
index c83820b3ac..026766c7f4 100644
--- a/tests/src/com/android/launcher3/ui/WorkProfileTest.java
+++ b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
@@ -19,6 +19,7 @@ import static com.android.launcher3.LauncherState.ALL_APPS;
 import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.allapps.AllAppsStore.DEFER_UPDATES_TEST;
 import static com.android.launcher3.testing.shared.TestProtocol.WORK_TAB_MISSING;
+import static com.android.launcher3.util.TestUtil.installDummyAppForUser;
 import static com.android.launcher3.util.rule.TestStabilityRule.LOCAL;
 import static com.android.launcher3.util.rule.TestStabilityRule.PLATFORM_POSTSUBMIT;
 
@@ -39,6 +40,7 @@ import com.android.launcher3.allapps.WorkPausedCard;
 import com.android.launcher3.allapps.WorkProfileManager;
 import com.android.launcher3.tapl.LauncherInstrumentation;
 import com.android.launcher3.testing.shared.TestProtocol;
+import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.rule.TestStabilityRule.Stability;
 
 import org.junit.After;
@@ -69,11 +71,11 @@ public class WorkProfileTest extends AbstractLauncherUiTest {
 
         String[] tokens = output.split("\\s+");
         mProfileUserId = Integer.parseInt(tokens[tokens.length - 1]);
-        output = mDevice.executeShellCommand("am start-user " + mProfileUserId);
         StringBuilder logStr = new StringBuilder().append("profileId: ").append(mProfileUserId);
         for (String str : tokens) {
             logStr.append(str).append("\n");
         }
+        installDummyAppForUser(mProfileUserId);
         updateWorkProfileSetupSuccessful("am start-user", output);
 
         Log.d(WORK_TAB_MISSING, "workProfileSuccessful? " + mWorkProfileSetupSuccessful +
@@ -101,6 +103,7 @@ public class WorkProfileTest extends AbstractLauncherUiTest {
             }
             launcher.getAppsView().getAppsStore().disableDeferUpdates(DEFER_UPDATES_TEST);
         });
+        TestUtil.uninstallDummyApp();
         mDevice.executeShellCommand("pm remove-user " + mProfileUserId);
     }
 
diff --git a/tests/src/com/android/launcher3/util/TestUtil.java b/tests/src/com/android/launcher3/util/TestUtil.java
index 4981795fe5..f8cd995aab 100644
--- a/tests/src/com/android/launcher3/util/TestUtil.java
+++ b/tests/src/com/android/launcher3/util/TestUtil.java
@@ -58,8 +58,13 @@ import java.util.function.ToIntFunction;
 
 public class TestUtil {
     public static final String DUMMY_PACKAGE = "com.example.android.aardwolf";
+    public static final int DEFAULT_USER_ID = 0;
 
     public static void installDummyApp() throws IOException {
+        installDummyAppForUser(DEFAULT_USER_ID);
+    }
+
+    public static void installDummyAppForUser(int userId) throws IOException {
         // Copy apk from resources to a local file and install from there.
         final Resources resources = getContext().getResources();
         final InputStream in = resources.openRawResource(
@@ -80,7 +85,7 @@ public class TestUtil {
             out.close();
 
             final String result = UiDevice.getInstance(getInstrumentation())
-                    .executeShellCommand("pm install " + apkFilename);
+                    .executeShellCommand("pm install --user " + userId + " " + apkFilename);
             Assert.assertTrue(
                     "Failed to install wellbeing test apk; make sure the device is rooted",
                     "Success".equals(result.replaceAll("\\s+", "")));