diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 551b533..d4b2ee2 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -220,7 +220,9 @@
 
     @Before
     public void setUp() throws Exception {
+        Log.d(TAG, "Before disabling battery defender");
         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")) {
             Log.d(TAG, "Before unlocking the phone");
@@ -228,6 +230,9 @@
         } else {
             Log.d(TAG, "Phone isn't locked");
         }
+        Assert.assertTrue("Keyguard still visible",
+                mDevice.wait(Until.gone(By.res(SYSTEMUI_PACKAGE, "keyguard_status_view")), 10000));
+        Log.d(TAG, "Keyguard is not visible");
 
         final String launcherPackageName = mDevice.getLauncherPackageName();
         try {