diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/ImeAppHelper.java b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/ImeAppHelper.java
index 098fd6d4250bd..80899396ec71a 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/ImeAppHelper.java
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/ImeAppHelper.java
@@ -18,6 +18,8 @@ package com.android.server.wm.flicker.helpers;
 
 import static android.os.SystemClock.sleep;
 
+import static org.junit.Assert.assertNotNull;
+
 import android.app.Instrumentation;
 import android.support.test.uiautomator.By;
 import android.support.test.uiautomator.UiDevice;
@@ -35,6 +37,8 @@ public class ImeAppHelper extends FlickerAppHelper {
 
     public void clickEditTextWidget(UiDevice device) {
         UiObject2 editText = device.findObject(By.res(getPackage(), "plain_text_input"));
+        assertNotNull("Text field not found, this usually happens when the device was left "
+                + "in an unknown state (e.g. in split screen)", editText);
         editText.click();
         sleep(500);
     }
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/PipAppHelper.java b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/PipAppHelper.java
index d5f9a2062a17b..361ca51539e16 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/PipAppHelper.java
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/helpers/PipAppHelper.java
@@ -16,13 +16,14 @@
 
 package com.android.server.wm.flicker.helpers;
 
-import static com.android.server.wm.flicker.helpers.AutomationUtils.getPipWindowSelector;
+import static com.android.server.wm.flicker.helpers.AutomationUtils.hasPipWindow;
+
+import static org.junit.Assert.assertNotNull;
 
 import android.app.Instrumentation;
 import android.support.test.uiautomator.By;
 import android.support.test.uiautomator.UiDevice;
 import android.support.test.uiautomator.UiObject2;
-import android.support.test.uiautomator.Until;
 
 public class PipAppHelper extends FlickerAppHelper {
 
@@ -32,12 +33,10 @@ public class PipAppHelper extends FlickerAppHelper {
 
     public void clickEnterPipButton(UiDevice device) {
         UiObject2 enterPipButton = device.findObject(By.res(getPackage(), "enter_pip"));
+        assertNotNull("Pip button not found, this usually happens when the device was left "
+                + "in an unknown state (e.g. in split screen)", enterPipButton);
         enterPipButton.click();
-        UiObject2 pipWindow = device.wait(Until.findObject(getPipWindowSelector()), sFindTimeout);
-
-        if (pipWindow == null) {
-            throw new RuntimeException("Unable to find PIP window");
-        }
+        hasPipWindow(device);
     }
 
     public void closePipWindow(UiDevice device) {