diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
index 4edc642..d9b4f47 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
@@ -295,8 +295,8 @@
                         mImeSourceControl.release(SurfaceControl::release);
                     }
                 }
+                mImeSourceControl = imeSourceControl;
             }
-            mImeSourceControl = imeSourceControl;
         }
 
         private void applyVisibilityToLeash(InsetsSourceControl imeSourceControl) {
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayImeControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayImeControllerTest.java
index 8641541..40f2e88 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayImeControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayImeControllerTest.java
@@ -21,9 +21,7 @@
 import static android.view.Surface.ROTATION_0;
 import static android.view.WindowInsets.Type.ime;
 
-import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
@@ -92,22 +90,6 @@
     }
 
     @Test
-    public void insetsControlChanged_updateExpectedImeSourceControl() {
-        final InsetsSourceControl[] insetsSourceControls = new InsetsSourceControl[]{
-                new InsetsSourceControl(ITYPE_IME, mock(SurfaceControl.class), false,
-                        new Point(0, 0), Insets.NONE)};
-        final InsetsSourceControl imeSourceControl = insetsSourceControls[0];
-
-        mPerDisplay.insetsControlChanged(insetsStateWithIme(false), insetsSourceControls);
-
-        assertEquals(imeSourceControl, mPerDisplay.mImeSourceControl);
-
-        mPerDisplay.insetsControlChanged(insetsStateWithIme(false), null);
-
-        assertNull(mPerDisplay.mImeSourceControl);
-    }
-
-    @Test
     public void insetsChanged_schedulesNoWorkOnExecutor() {
         mPerDisplay.insetsChanged(insetsStateWithIme(false));
         verifyZeroInteractions(mExecutor);