diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 5c9fe83..7cd656e 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -18,7 +18,6 @@
 import static androidx.test.InstrumentationRegistry.getInstrumentation;
 
 import static com.android.launcher3.WorkspaceLayoutManager.FIRST_SCREEN_ID;
-import static com.android.launcher3.tapl.LauncherInstrumentation.ContainerType;
 import static com.android.launcher3.ui.TaplTestsLauncher3.getAppPackageName;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
@@ -57,6 +56,7 @@
 import com.android.launcher3.Utilities;
 import com.android.launcher3.model.AppLaunchTracker;
 import com.android.launcher3.tapl.LauncherInstrumentation;
+import com.android.launcher3.tapl.LauncherInstrumentation.ContainerType;
 import com.android.launcher3.tapl.TestHelpers;
 import com.android.launcher3.testcomponent.TestCommandReceiver;
 import com.android.launcher3.testing.TestProtocol;
@@ -224,14 +224,14 @@
         mDevice.executeShellCommand("setprop vendor.battery.defender.disable 1");
         Log.d(TAG, "Before enabling stay awake");
         mDevice.executeShellCommand("settings put global stay_on_while_plugged_in 3");
-        if (hasSystemUiObject("keyguard_status_view")) {
+        for (int i = 0; i < 10 && hasSystemUiObject("keyguard_status_view"); ++i) {
             Log.d(TAG, "Before unlocking the phone");
             mDevice.executeShellCommand("input keyevent 82");
-        } else {
-            Log.d(TAG, "Phone isn't locked");
+            mDevice.waitForIdle();
         }
         Assert.assertTrue("Keyguard still visible",
-                mDevice.wait(Until.gone(By.res(SYSTEMUI_PACKAGE, "keyguard_status_view")), 60000));
+                mDevice.wait(
+                        Until.gone(By.res(SYSTEMUI_PACKAGE, "keyguard_status_view")), 60000));
         Log.d(TAG, "Keyguard is not visible");
 
         final String launcherPackageName = mDevice.getLauncherPackageName();
@@ -294,7 +294,8 @@
     protected void resetLoaderState() {
         try {
             mMainThreadExecutor.execute(
-                    () -> LauncherAppState.getInstance(mTargetContext).getModel().forceReload());
+                    () -> LauncherAppState.getInstance(
+                            mTargetContext).getModel().forceReload());
         } catch (Throwable t) {
             throw new IllegalArgumentException(t);
         }
@@ -308,7 +309,8 @@
         ContentResolver resolver = mTargetContext.getContentResolver();
         int screenId = FIRST_SCREEN_ID;
         // Update the screen id counter for the provider.
-        LauncherSettings.Settings.call(resolver, LauncherSettings.Settings.METHOD_NEW_SCREEN_ID);
+        LauncherSettings.Settings.call(resolver,
+                LauncherSettings.Settings.METHOD_NEW_SCREEN_ID);
 
         if (screenId > FIRST_SCREEN_ID) {
             screenId = FIRST_SCREEN_ID;
@@ -322,7 +324,8 @@
         item.screenId = screenId;
         item.onAddToDatabase(writer);
         writer.put(LauncherSettings.Favorites._ID, item.id);
-        resolver.insert(LauncherSettings.Favorites.CONTENT_URI, writer.getValues(mTargetContext));
+        resolver.insert(LauncherSettings.Favorites.CONTENT_URI,
+                writer.getValues(mTargetContext));
         resetLoaderState();
 
         // Launch the home activity
@@ -353,7 +356,8 @@
         });
     }
 
-    // Cannot be used in TaplTests between a Tapl call injecting a gesture and a tapl call expecting
+    // Cannot be used in TaplTests between a Tapl call injecting a gesture and a tapl call
+    // expecting
     // the results of that gesture because the wait can hide flakeness.
     protected void waitForState(String message, Supplier<LauncherState> state) {
         waitForLauncherCondition(message,
@@ -366,7 +370,8 @@
 
     // Cannot be used in TaplTests after injecting any gesture using Tapl because this can hide
     // flakiness.
-    protected void waitForLauncherCondition(String message, Function<Launcher, Boolean> condition) {
+    protected void waitForLauncherCondition(String
+            message, Function<Launcher, Boolean> condition) {
         waitForLauncherCondition(message, condition, DEFAULT_ACTIVITY_TIMEOUT);
     }
 
@@ -442,7 +447,8 @@
 
         public Intent blockingGetExtraIntent() throws InterruptedException {
             Intent intent = blockingGetIntent();
-            return intent == null ? null : (Intent) intent.getParcelableExtra(Intent.EXTRA_INTENT);
+            return intent == null ? null : (Intent) intent.getParcelableExtra(
+                    Intent.EXTRA_INTENT);
         }
     }
 
@@ -469,7 +475,8 @@
         if (newTask) {
             intent.addFlags(Intent.FLAG_ACTIVITY_NEW_TASK | Intent.FLAG_ACTIVITY_CLEAR_TASK);
         } else {
-            intent.addFlags(Intent.FLAG_ACTIVITY_MULTIPLE_TASK | Intent.FLAG_ACTIVITY_NEW_DOCUMENT);
+            intent.addFlags(
+                    Intent.FLAG_ACTIVITY_MULTIPLE_TASK | Intent.FLAG_ACTIVITY_NEW_DOCUMENT);
         }
         getInstrumentation().getTargetContext().startActivity(intent);
         assertTrue("App didn't start: " + selector,
@@ -506,7 +513,8 @@
 
     protected boolean isInState(Supplier<LauncherState> state) {
         if (!TestHelpers.isInLauncherProcess()) return true;
-        return getFromLauncher(launcher -> launcher.getStateManager().getState() == state.get());
+        return getFromLauncher(
+                launcher -> launcher.getStateManager().getState() == state.get());
     }
 
     protected int getAllAppsScroll(Launcher launcher) {