From 0fe753af2a993fbb61201e9ee569d66fa8848573 Mon Sep 17 00:00:00 2001
From: Vadim Tryshev <vadimt@google.com>
Date: Fri, 3 Nov 2023 12:27:00 -0700
Subject: [PATCH] Waiting for the setup wizard to go away

This is a workaround for b/309008042 where it appears that after being
dismissed by the test infra, setup wizard doesn't immediately go away.

This CL will wait 2 minutes for the dismissal, and provide a better
diags when the wizard doesn't go away after 2 min.

Bug: 309008042
Test: presubmit
Flag: N/A
Change-Id: I26eb2f202842702af08c424929d74fbd21eb872d
---
 .../quickstep/FallbackRecentsTest.java        |  1 +
 .../launcher3/ui/AbstractLauncherUiTest.java  | 30 ++++++++++++++++++-
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
index 16235e0117..12568ea85d 100644
--- a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
+++ b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
@@ -162,6 +162,7 @@ public class FallbackRecentsTest {
     @Before
     public void setUp() {
         mLauncher.onTestStart();
+        AbstractLauncherUiTest.waitForSetupWizardDismissal();
         AbstractLauncherUiTest.verifyKeyguardInvisible();
     }
 
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index ed846edc28..ed3a76bba4 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -187,7 +187,7 @@ public abstract class AbstractLauncherUiTest {
             Utilities.enableRunningInTestHarnessForTests();
             mLauncher.setSystemHealthSupplier(startTime -> TestCommandReceiver.callCommand(
                             TestCommandReceiver.GET_SYSTEM_HEALTH_MESSAGE, startTime.toString())
-                            .getString("result"));
+                    .getString("result"));
             mLauncher.setOnSettledStateAction(
                     containerType -> executeOnLauncher(
                             launcher ->
@@ -248,6 +248,7 @@ public abstract class AbstractLauncherUiTest {
     public void setUp() throws Exception {
         mLauncher.onTestStart();
 
+        waitForSetupWizardDismissal();
         if (TestStabilityRule.isPresubmit()) {
             aggressivelyUnlockSysUi();
         } else {
@@ -308,6 +309,32 @@ public abstract class AbstractLauncherUiTest {
         Log.d(TAG, "Keyguard is not visible");
     }
 
+    // b/309008042
+    private static boolean sFirstTimeWaitingForWizard = true;
+
+    static {
+        waitForSetupWizardDismissal();
+    }
+
+    /** Waits for setup wizard to go away. */
+    public static void waitForSetupWizardDismissal() {
+        if (sFirstTimeWaitingForWizard && TestStabilityRule.isPresubmit()) {
+            try {
+                UiDevice.getInstance(getInstrumentation()).executeShellCommand(
+                        "am force-stop com.google.android.setupwizard");
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
+        final boolean wizardDismissed = TestHelpers.wait(
+                Until.gone(By.pkg("com.google.android.setupwizard").depth(0)),
+                sFirstTimeWaitingForWizard ? 120000 : 0);
+        sFirstTimeWaitingForWizard = false;
+        Assert.assertTrue("Setup wizard is still visible",
+                wizardDismissed);
+    }
+
     public static void verifyKeyguardInvisible() {
         final boolean keyguardAlreadyVisible = sSeenKeyguard;
 
@@ -422,6 +449,7 @@ public abstract class AbstractLauncherUiTest {
     // flakiness.
     protected void waitForLauncherCondition(
             String message, Function<Launcher, Boolean> condition, long timeout) {
+        waitForSetupWizardDismissal();
         verifyKeyguardInvisible();
         if (!TestHelpers.isInLauncherProcess()) return;
         Wait.atMost(message, () -> getFromLauncher(condition), timeout, mLauncher);