Waiting for app install to finish before procedding with the test

Bug: 256659409
Test: Presubmit
Change-Id: Ia0f4cdd072c4c439d09070b0395fcfd6909c2a8f
diff --git a/tests/AndroidManifest-common.xml b/tests/AndroidManifest-common.xml
index 4af8468..bedf277 100644
--- a/tests/AndroidManifest-common.xml
+++ b/tests/AndroidManifest-common.xml
@@ -24,6 +24,7 @@
     <uses-permission android:name="android.permission.PACKAGE_USAGE_STATS"/>
     <uses-permission android:name="android.permission.READ_LOGS"/>
     <uses-permission android:name="android.permission.WRITE_SECURE_SETTINGS"/>
+    <uses-permission android:name="android.permission.QUERY_ALL_PACKAGES" />
 
     <application android:debuggable="true" android:extractNativeLibs="true">
         <uses-library android:name="android.test.runner"/>
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
index 01e6ed7..978e84c 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
@@ -29,7 +29,6 @@
 
 import android.content.Intent;
 import android.graphics.Point;
-import android.os.SystemClock;
 import android.platform.test.annotations.IwTest;
 
 import androidx.test.filters.LargeTest;
@@ -479,7 +478,7 @@
     @Test
     @PortraitLandscape
     public void testUninstallFromWorkspace() throws Exception {
-        TestUtil.installDummyApp();
+        installDummyAppAndWaitForUIUpdate();
         try {
             verifyAppUninstalledFromAllApps(
                     createShortcutInCenterIfNotExist(DUMMY_APP_NAME).uninstall(), DUMMY_APP_NAME);
@@ -492,10 +491,8 @@
     @PortraitLandscape
     @ScreenRecord // (b/256659409)
     public void testUninstallFromAllApps() throws Exception {
-        TestUtil.installDummyApp();
+        installDummyAppAndWaitForUIUpdate();
         try {
-            // b/256659409
-            SystemClock.sleep(5000);
             Workspace workspace = mLauncher.getWorkspace();
             final HomeAllApps allApps = workspace.switchToAllApps();
             allApps.freeze();
@@ -539,7 +536,7 @@
         Point[] gridPositions = getCornersAndCenterPositions();
         createShortcutIfNotExist(STORE_APP_NAME, gridPositions[0]);
         createShortcutIfNotExist(MAPS_APP_NAME, gridPositions[1]);
-        TestUtil.installDummyApp();
+        installDummyAppAndWaitForUIUpdate();
         try {
             createShortcutIfNotExist(DUMMY_APP_NAME, gridPositions[2]);
             Map<String, Point> initialPositions =
@@ -590,6 +587,17 @@
                 mLauncher.getWorkspace().getHotseatAppIcon(APP_NAME));
     }
 
+    private void installDummyAppAndWaitForUIUpdate() throws IOException {
+        TestUtil.installDummyApp();
+        // Wait for model thread completion as it may be processing
+        // the install event from the SystemService
+        mLauncher.waitForModelQueueCleared();
+        // Wait for Launcher UI thread completion, as it may be processing updating the UI in
+        // response to the model update. Not that `waitForLauncherInitialized` is just a proxy
+        // method, we can use any method which touches Launcher UI thread,
+        mLauncher.waitForLauncherInitialized();
+    }
+
     /**
      * @return List of workspace grid coordinates. Those are not pixels. See {@link
      *     Workspace#getIconGridDimensions()}
diff --git a/tests/src/com/android/launcher3/util/TestUtil.java b/tests/src/com/android/launcher3/util/TestUtil.java
index 67f3902..d7c6c4f 100644
--- a/tests/src/com/android/launcher3/util/TestUtil.java
+++ b/tests/src/com/android/launcher3/util/TestUtil.java
@@ -17,8 +17,13 @@
 
 import static androidx.test.InstrumentationRegistry.getContext;
 import static androidx.test.InstrumentationRegistry.getInstrumentation;
+import static androidx.test.InstrumentationRegistry.getTargetContext;
 
+import android.content.pm.LauncherApps;
 import android.content.res.Resources;
+import android.os.Handler;
+import android.os.Looper;
+import android.os.UserHandle;
 
 import androidx.test.uiautomator.UiDevice;
 
@@ -27,6 +32,7 @@
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.util.concurrent.CountDownLatch;
 
 public class TestUtil {
     public static final String DUMMY_PACKAGE = "com.example.android.aardwolf";
@@ -40,24 +46,77 @@
         final String apkFilename = getInstrumentation().getTargetContext().
                 getFilesDir().getPath() + "/dummy_app.apk";
 
-        final FileOutputStream out = new FileOutputStream(apkFilename);
-        byte[] buff = new byte[1024];
-        int read;
+        try (PackageInstallCheck pic = new PackageInstallCheck()) {
+            final FileOutputStream out = new FileOutputStream(apkFilename);
+            byte[] buff = new byte[1024];
+            int read;
 
-        while ((read = in.read(buff)) > 0) {
-            out.write(buff, 0, read);
+            while ((read = in.read(buff)) > 0) {
+                out.write(buff, 0, read);
+            }
+            in.close();
+            out.close();
+
+            final String result = UiDevice.getInstance(getInstrumentation())
+                    .executeShellCommand("pm install " + apkFilename);
+            Assert.assertTrue(
+                    "Failed to install wellbeing test apk; make sure the device is rooted",
+                    "Success".equals(result.replaceAll("\\s+", "")));
+            pic.mAddWait.await();
+        } catch (InterruptedException e) {
+            throw new IOException(e);
         }
-        in.close();
-        out.close();
-
-        final String result = UiDevice.getInstance(getInstrumentation())
-                .executeShellCommand("pm install " + apkFilename);
-        Assert.assertTrue("Failed to install wellbeing test apk; make sure the device is rooted",
-                "Success".equals(result.replaceAll("\\s+", "")));
     }
 
     public static void uninstallDummyApp() throws IOException {
         UiDevice.getInstance(getInstrumentation()).executeShellCommand(
                 "pm uninstall " + DUMMY_PACKAGE);
     }
+
+    private static class PackageInstallCheck extends LauncherApps.Callback
+            implements AutoCloseable {
+
+        final CountDownLatch mAddWait = new CountDownLatch(1);
+        final LauncherApps mLauncherApps;
+
+        PackageInstallCheck() {
+            mLauncherApps = getTargetContext().getSystemService(LauncherApps.class);
+            mLauncherApps.registerCallback(this, new Handler(Looper.getMainLooper()));
+        }
+
+        private void verifyPackage(String packageName) {
+            if (DUMMY_PACKAGE.equals(packageName)) {
+                mAddWait.countDown();
+            }
+        }
+
+        @Override
+        public void onPackageAdded(String packageName, UserHandle user) {
+            verifyPackage(packageName);
+        }
+
+        @Override
+        public void onPackageChanged(String packageName, UserHandle user) {
+            verifyPackage(packageName);
+        }
+
+        @Override
+        public void onPackageRemoved(String packageName, UserHandle user) { }
+
+        @Override
+        public void onPackagesAvailable(String[] packageNames, UserHandle user, boolean replacing) {
+            for (String packageName : packageNames) {
+                verifyPackage(packageName);
+            }
+        }
+
+        @Override
+        public void onPackagesUnavailable(String[] packageNames, UserHandle user,
+                boolean replacing) { }
+
+        @Override
+        public void close() {
+            mLauncherApps.unregisterCallback(this);
+        }
+    }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 449b7b7..1c5c5fa 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -848,6 +848,10 @@
         }
     }
 
+    public void waitForModelQueueCleared() {
+        getTestInfo(TestProtocol.REQUEST_MODEL_QUEUE_CLEARED);
+    }
+
     public void waitForLauncherInitialized() {
         for (int i = 0; i < 100; ++i) {
             if (getTestInfo(