diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 4f01794..8847b8b 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -332,7 +332,9 @@
     // TODO(309471958) Productize killing/dismissal of setup wizard.
     /** Waits for setup wizard to go away. */
     public static void waitForSetupWizardDismissal() {
-        if (sFirstTimeWaitingForWizard && TestStabilityRule.isPresubmit()) {
+        if (!TestStabilityRule.isPresubmit()) return;
+
+        if (sFirstTimeWaitingForWizard) {
             try {
                 getUiDevice().executeShellCommand(
                         "am force-stop com.google.android.setupwizard");
@@ -345,9 +347,7 @@
                 Until.gone(By.pkg("com.google.android.setupwizard").depth(0)),
                 sFirstTimeWaitingForWizard ? 120000 : 0);
         sFirstTimeWaitingForWizard = false;
-        // b/309496273
-//        Assert.assertTrue("Setup wizard is still visible",
-//                wizardDismissed);
+        Assert.assertTrue("Setup wizard is still visible", wizardDismissed);
     }
 
     private static void verifyKeyguardInvisible() {