Waiting until the launcher change is propagater before test

Bug: 143488140
Test: Presubmit
Change-Id: Ie3f35a1e58b784d078a305814bd682beaae1f4cb
diff --git a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
index b9e0f62..475953b 100644
--- a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
+++ b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
@@ -23,6 +23,7 @@
 import static com.android.launcher3.tapl.TestHelpers.getHomeIntentInPackage;
 import static com.android.launcher3.tapl.TestHelpers.getLauncherInMyProcess;
 import static com.android.launcher3.ui.AbstractLauncherUiTest.DEFAULT_ACTIVITY_TIMEOUT;
+import static com.android.launcher3.ui.AbstractLauncherUiTest.DEFAULT_BROADCAST_TIMEOUT_SECS;
 import static com.android.launcher3.ui.AbstractLauncherUiTest.DEFAULT_UI_TIMEOUT;
 import static com.android.launcher3.ui.AbstractLauncherUiTest.resolveSystemApp;
 import static com.android.launcher3.ui.AbstractLauncherUiTest.startAppFast;
@@ -66,6 +67,8 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.model.Statement;
 
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 import java.util.function.Function;
 
@@ -112,11 +115,16 @@
             @Override
             public void evaluate() throws Throwable {
                 TestCommandReceiver.callCommand(TestCommandReceiver.ENABLE_TEST_LAUNCHER);
+                OverviewUpdateHandler updateHandler =
+                        MAIN_EXECUTOR.submit(OverviewUpdateHandler::new).get();
                 UiDevice.getInstance(getInstrumentation()).executeShellCommand(
                         getLauncherCommand(mOtherLauncherActivity));
+                updateHandler.mChangeCounter
+                        .await(DEFAULT_BROADCAST_TIMEOUT_SECS, TimeUnit.SECONDS);
                 try {
                     base.evaluate();
                 } finally {
+                    MAIN_EXECUTOR.submit(updateHandler::destroy).get();
                     TestCommandReceiver.callCommand(TestCommandReceiver.DISABLE_TEST_LAUNCHER);
                     UiDevice.getInstance(getInstrumentation()).executeShellCommand(
                             getLauncherCommand(getLauncherInMyProcess()));
@@ -241,4 +249,30 @@
     private int getTaskCount(RecentsActivity recents) {
         return recents.<RecentsView>getOverviewPanel().getTaskViewCount();
     }
+
+    private class OverviewUpdateHandler {
+
+        final RecentsAnimationDeviceState mRads;
+        final OverviewComponentObserver mObserver;
+        final CountDownLatch mChangeCounter;
+
+        OverviewUpdateHandler() {
+            Context ctx = getInstrumentation().getTargetContext();
+            mRads = new RecentsAnimationDeviceState(ctx);
+            mObserver = new OverviewComponentObserver(ctx, mRads);
+            mChangeCounter = new CountDownLatch(1);
+            if (mObserver.getHomeIntent().getComponent()
+                    .getPackageName().equals(mOtherLauncherActivity.packageName)) {
+                // Home already same
+                mChangeCounter.countDown();
+            } else {
+                mObserver.setOverviewChangeListener(b -> mChangeCounter.countDown());
+            }
+        }
+
+        void destroy() {
+            mObserver.onDestroy();
+            mRads.destroy();
+        }
+    }
 }