diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index a870ed3faa284..974340ee2510e 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -4407,8 +4407,8 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
             boolean imeLayeringTargetMayUseIme =
                     LayoutParams.mayUseInputMethod(mImeLayeringTarget.mAttrs.flags)
                     || mImeLayeringTarget.mAttrs.type == TYPE_APPLICATION_STARTING;
-            if (imeLayeringTargetMayUseIme && mImeInputTarget != null
-                    && mImeLayeringTarget.mActivityRecord != mImeInputTarget.getActivityRecord()) {
+            if (imeLayeringTargetMayUseIme && (mImeInputTarget == null
+                    || mImeLayeringTarget.mActivityRecord != mImeInputTarget.getActivityRecord())) {
                 // Do not change parent if the window hasn't requested IME.
                 return null;
             }
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
index b27060872baa5..8f5da44212800 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
@@ -643,6 +643,7 @@ public class DisplayContentTests extends WindowTestsBase {
         final DisplayContent dc = mDisplayContent;
         final WindowState ws = createWindow(null, TYPE_APPLICATION, dc, "app window");
         dc.setImeLayeringTarget(ws);
+        dc.setImeInputTarget(ws);
 
         // Adjust bounds so that matchesRootDisplayAreaBounds() returns false.
         final Rect bounds = new Rect(dc.getBounds());
@@ -1154,6 +1155,7 @@ public class DisplayContentTests extends WindowTestsBase {
         dc.setImeLayeringTarget(createWindow(null, TYPE_STATUS_BAR, "app"));
         dc.getImeTarget(IME_TARGET_LAYERING).getWindow().setWindowingMode(
                 WindowConfiguration.WINDOWING_MODE_MULTI_WINDOW);
+        dc.setImeInputTarget(dc.getImeTarget(IME_TARGET_LAYERING).getWindow());
         assertEquals(dc.getImeContainer().getParentSurfaceControl(), dc.computeImeParent());
     }
 
@@ -1172,6 +1174,7 @@ public class DisplayContentTests extends WindowTestsBase {
     public void testComputeImeParent_noApp() throws Exception {
         final DisplayContent dc = createNewDisplay();
         dc.setImeLayeringTarget(createWindow(null, TYPE_STATUS_BAR, "statusBar"));
+        dc.setImeInputTarget(dc.getImeTarget(IME_TARGET_LAYERING).getWindow());
         assertEquals(dc.getImeContainer().getParentSurfaceControl(), dc.computeImeParent());
     }
 
diff --git a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeEditorPopupDialogTest.kt b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeEditorPopupDialogTest.kt
index 6257484be9bde..2e29b3e314ca7 100644
--- a/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeEditorPopupDialogTest.kt
+++ b/tests/FlickerTests/src/com/android/server/wm/flicker/ime/CloseImeEditorPopupDialogTest.kt
@@ -59,7 +59,7 @@ class CloseImeEditorPopupDialogTest(private val testSpec: FlickerTestParameter)
             }
             transitions {
                 imeTestApp.dismissDialog(wmHelper)
-                instrumentation.uiAutomation.syncInputTransactions()
+                wmHelper.waitImeGone()
             }
             teardown {
                 eachRun {
@@ -91,7 +91,7 @@ class CloseImeEditorPopupDialogTest(private val testSpec: FlickerTestParameter)
                     .then()
                     .isVisible(FlickerComponentName.IME_SNAPSHOT)
                     .then()
-                    .isInvisible(FlickerComponentName.IME_SNAPSHOT)
+                    .isInvisible(FlickerComponentName.IME_SNAPSHOT, isOptional = true)
                     .isInvisible(FlickerComponentName.IME)
         }
     }