diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/ImeAppHelper.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/ImeAppHelper.kt
index d580104ade198..c546a4d180276 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/ImeAppHelper.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/helpers/ImeAppHelper.kt
@@ -20,7 +20,6 @@ import android.app.Instrumentation
 import androidx.test.uiautomator.By
 import androidx.test.uiautomator.Until
 import com.android.server.wm.flicker.helpers.FIND_TIMEOUT
-import com.android.server.wm.flicker.helpers.waitForIME
 import com.android.wm.shell.flicker.TEST_APP_IME_ACTIVITY_LABEL
 import com.android.wm.shell.flicker.testapp.Components
 import org.junit.Assert
@@ -39,14 +38,9 @@ open class ImeAppHelper(
         Assert.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()
-        if (!uiDevice.waitForIME()) {
-            Assert.fail("IME did not appear")
-        }
     }
 
     fun closeIME() {
         uiDevice.pressBack()
-        // Using only the AccessibilityInfo it is not possible to identify if the IME is active
-        uiDevice.waitForIdle(1000)
     }
 }
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/PipKeyboardTest.kt b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/PipKeyboardTest.kt
index 6b44ce6ace0ca..866d654362ca0 100644
--- a/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/PipKeyboardTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/src/com/android/wm/shell/flicker/pip/PipKeyboardTest.kt
@@ -25,8 +25,10 @@ import com.android.server.wm.flicker.helpers.WindowUtils
 import com.android.server.wm.flicker.helpers.closePipWindow
 import com.android.server.wm.flicker.helpers.hasPipWindow
 import com.android.server.wm.flicker.helpers.wakeUpAndGoToHomeScreen
+import com.android.server.wm.traces.parser.windowmanager.WindowManagerStateHelper
 import com.android.wm.shell.flicker.IME_WINDOW_NAME
 import com.android.wm.shell.flicker.helpers.ImeAppHelper
+import com.android.wm.shell.flicker.testapp.Components
 import org.junit.FixMethodOrder
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -46,6 +48,8 @@ class PipKeyboardTest(
     rotation: Int
 ) : PipTestBase(rotationName, rotation) {
     private val keyboardApp = ImeAppHelper(instrumentation)
+    private val keyboardComponent = Components.ImeActivity().componentName
+    private val helper = WindowManagerStateHelper()
 
     private val keyboardScenario: FlickerBuilder
         get() = FlickerBuilder(instrumentation).apply {
@@ -64,6 +68,8 @@ class PipKeyboardTest(
                     // UiAutomator doesn't support to launch the multiple Activities in a task.
                     // So use launchActivity() for the Keyboard Activity.
                     keyboardApp.launchViaIntent()
+                    helper.waitForAppTransitionIdle()
+                    helper.waitForFullScreenApp(keyboardComponent)
                 }
             }
             teardown {
@@ -88,9 +94,11 @@ class PipKeyboardTest(
             transitions {
                 // open the soft keyboard
                 keyboardApp.openIME()
+                helper.waitImeWindowShown()
 
                 // then close it again
                 keyboardApp.closeIME()
+                helper.waitImeWindowGone()
             }
             assertions {
                 windowManagerTrace {
@@ -112,11 +120,13 @@ class PipKeyboardTest(
             transitions {
                 // open the soft keyboard
                 keyboardApp.openIME()
+                helper.waitImeWindowShown()
             }
             teardown {
                 eachRun {
                     // close the keyboard
                     keyboardApp.closeIME()
+                    helper.waitImeWindowGone()
                 }
             }
             assertions {