Include inputViewShown in verifyInputViewStatus

Currently verifyInputViewStatus includes a few checks, but does not
incldue checking isInputViewShown. This refactors the method to also
check this value, to reduce duplication.

Additionally re-orders the methods in TestActivity and
InputMethodServiceWrapper, as well as storing the static instance of the
wrapper in a weak reference.

Flag: EXEMPT testfix
Bug: 394328311
Test: atest InputMethodServiceTest
Change-Id: Idd844caadedec49b0165d3cf4ed6ca3e61638783

3 files changed, 183 insertions, 322 deletions

diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/inputmethodservice/InputMethodServiceTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/inputmethodservice/InputMethodServiceTest.java
index 72e9cc566497..b501dd246363 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/inputmethodservice/InputMethodServiceTest.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/inputmethodservice/InputMethodServiceTest.java
@@ -173,15 +173,14 @@ public class InputMethodServiceTest {
         Log.i(TAG, "Click on EditText");
         verifyInputViewStatus(
                 () -> clickOnViewCenter(mActivity.getEditText()),
-                EVENT_SHOW,
-                true /* expected */,
-                true /* inputViewStarted */);
-        assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
 
         // Press home key to hide IME.
         Log.i(TAG, "Press home");
         if (mFlagsValueProvider.getBoolean(Flags.FLAG_REFACTOR_INSETS_CONTROLLER)) {
             assertWithMessage("Home key press was handled").that(mUiDevice.pressHome()).isTrue();
+            // This doesn't call verifyInputViewStatus, as the refactor delays the events such that
+            // getCurrentInputStarted() would be false, as we would already be in launcher.
             // The IME visibility is only sent at the end of the animation. Therefore, we have to
             // wait until the visibility was sent to the server and the IME window hidden.
             eventually(() -> assertWithMessage("IME is not shown")
@@ -190,11 +189,7 @@ public class InputMethodServiceTest {
             verifyInputViewStatus(
                     () -> assertWithMessage("Home key press was handled")
                             .that(mUiDevice.pressHome()).isTrue(),
-                    EVENT_HIDE,
-                    true /* expected */,
-                    false /* inputViewStarted */);
-            assertWithMessage("IME is not shown")
-                    .that(mInputMethodService.isInputViewShown()).isFalse();
+                    EVENT_HIDE, true /* eventExpected */, false /* shown */, "IME is not shown");
         }
     }
 
@@ -208,26 +203,12 @@ public class InputMethodServiceTest {
         // Triggers to show IME via public API.
         verifyInputViewStatusOnMainSync(
                 () -> assertThat(mActivity.showImeWithInputMethodManager(0 /* flags */)).isTrue(),
-                EVENT_SHOW,
-                true /* expected */,
-                true /* inputViewStarted */);
-        assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
 
         // Triggers to hide IME via public API.
         verifyInputViewStatusOnMainSync(
                 () -> assertThat(mActivity.hideImeWithInputMethodManager(0 /* flags */)).isTrue(),
-                EVENT_HIDE,
-                true /* expected */,
-                false /* inputViewStarted */);
-        if (mFlagsValueProvider.getBoolean(Flags.FLAG_REFACTOR_INSETS_CONTROLLER)) {
-            // The IME visibility is only sent at the end of the animation. Therefore, we have to
-            // wait until the visibility was sent to the server and the IME window hidden.
-            eventually(() -> assertWithMessage("IME is not shown")
-                    .that(mInputMethodService.isInputViewShown()).isFalse());
-        } else {
-            assertWithMessage("IME is not shown")
-                    .that(mInputMethodService.isInputViewShown()).isFalse();
-        }
+                EVENT_HIDE, true /* eventExpected */, false /* shown */, "IME is not shown");
     }
 
     /**
@@ -240,26 +221,12 @@ public class InputMethodServiceTest {
         // Triggers to show IME via public API.
         verifyInputViewStatusOnMainSync(
                 () -> mActivity.showImeWithWindowInsetsController(),
-                EVENT_SHOW,
-                true /* expected */,
-                true /* inputViewStarted */);
-        assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
 
         // Triggers to hide IME via public API.
         verifyInputViewStatusOnMainSync(
                 () -> mActivity.hideImeWithWindowInsetsController(),
-                EVENT_HIDE,
-                true /* expected */,
-                false /* inputViewStarted */);
-        if (mFlagsValueProvider.getBoolean(Flags.FLAG_REFACTOR_INSETS_CONTROLLER)) {
-            // The IME visibility is only sent at the end of the animation. Therefore, we have to
-            // wait until the visibility was sent to the server and the IME window hidden.
-            eventually(() -> assertWithMessage("IME is not shown")
-                    .that(mInputMethodService.isInputViewShown()).isFalse());
-        } else {
-            assertWithMessage("IME is not shown")
-                    .that(mInputMethodService.isInputViewShown()).isFalse();
-        }
+                EVENT_HIDE, true /* eventExpected */, false /* shown */, "IME is not shown");
     }
 
     /**
@@ -277,10 +244,7 @@ public class InputMethodServiceTest {
         Log.i(TAG, "Call IMS#requestShowSelf(0)");
         verifyInputViewStatusOnMainSync(
                 () -> mInputMethodService.requestShowSelf(0 /* flags */),
-                EVENT_SHOW,
-                true /* expected */,
-                true /* inputViewStarted */);
-        assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
 
         if (!mFlagsValueProvider.getBoolean(Flags.FLAG_REFACTOR_INSETS_CONTROLLER)) {
             // IME request to hide itself with flag HIDE_IMPLICIT_ONLY, expect not hide (shown).
@@ -288,51 +252,31 @@ public class InputMethodServiceTest {
             verifyInputViewStatusOnMainSync(
                     () -> mInputMethodService.requestHideSelf(
                             InputMethodManager.HIDE_IMPLICIT_ONLY),
-                    EVENT_HIDE,
-                    false /* expected */,
-                    true /* inputViewStarted */);
-            assertWithMessage("IME is still shown after HIDE_IMPLICIT_ONLY")
-                    .that(mInputMethodService.isInputViewShown()).isTrue();
+                    EVENT_HIDE, false /* eventExpected */, true /* shown */,
+                    "IME is still shown after HIDE_IMPLICIT_ONLY");
         }
 
         // IME request to hide itself without any flags, expect hidden.
         Log.i(TAG, "Call IMS#requestHideSelf(0)");
         verifyInputViewStatusOnMainSync(
                 () -> mInputMethodService.requestHideSelf(0 /* flags */),
-                EVENT_HIDE,
-                true /* expected */,
-                false /* inputViewStarted */);
-        if (mFlagsValueProvider.getBoolean(Flags.FLAG_REFACTOR_INSETS_CONTROLLER)) {
-            // The IME visibility is only sent at the end of the animation. Therefore, we have to
-            // wait until the visibility was sent to the server and the IME window hidden.
-            eventually(() -> assertWithMessage("IME is not shown")
-                    .that(mInputMethodService.isInputViewShown()).isFalse());
-        } else {
-            assertWithMessage("IME is not shown")
-                    .that(mInputMethodService.isInputViewShown()).isFalse();
-        }
+                EVENT_HIDE, true /* eventExpected */, false /* shown */, "IME is not shown");
 
         if (!mFlagsValueProvider.getBoolean(Flags.FLAG_REFACTOR_INSETS_CONTROLLER)) {
             // IME request to show itself with flag SHOW_IMPLICIT, expect shown.
             Log.i(TAG, "Call IMS#requestShowSelf(InputMethodManager.SHOW_IMPLICIT)");
             verifyInputViewStatusOnMainSync(
                     () -> mInputMethodService.requestShowSelf(InputMethodManager.SHOW_IMPLICIT),
-                    EVENT_SHOW,
-                    true /* expected */,
-                    true /* inputViewStarted */);
-            assertWithMessage("IME is shown with SHOW_IMPLICIT")
-                    .that(mInputMethodService.isInputViewShown()).isTrue();
+                    EVENT_SHOW, true /* eventExpected */, true /* shown */,
+                    "IME is shown with SHOW_IMPLICIT");
 
             // IME request to hide itself with flag HIDE_IMPLICIT_ONLY, expect hidden.
             Log.i(TAG, "Call IMS#requestHideSelf(InputMethodManager.HIDE_IMPLICIT_ONLY)");
             verifyInputViewStatusOnMainSync(
                     () -> mInputMethodService.requestHideSelf(
                             InputMethodManager.HIDE_IMPLICIT_ONLY),
-                    EVENT_HIDE,
-                    true /* expected */,
-                    false /* inputViewStarted */);
-            assertWithMessage("IME is not shown after HIDE_IMPLICIT_ONLY")
-                    .that(mInputMethodService.isInputViewShown()).isFalse();
+                    EVENT_HIDE, true /* eventExpected */, false /* shown */,
+                    "IME is not shown after HIDE_IMPLICIT_ONLY");
         }
     }
 
@@ -424,20 +368,14 @@ public class InputMethodServiceTest {
             verifyInputViewStatusOnMainSync(() -> assertThat(
                             mActivity.showImeWithInputMethodManager(
                                     InputMethodManager.SHOW_IMPLICIT)).isTrue(),
-                    EVENT_SHOW,
-                    false /* expected */,
-                    false /* inputViewStarted */);
-            assertWithMessage("IME is not shown after SHOW_IMPLICIT")
-                    .that(mInputMethodService.isInputViewShown()).isFalse();
+                    EVENT_SHOW, false /* eventExpected */, false /* shown */,
+                    "IME is not shown after SHOW_IMPLICIT");
 
             verifyInputViewStatusOnMainSync(
                     () -> assertThat(mActivity.showImeWithInputMethodManager(0 /* flags */))
                             .isTrue(),
-                    EVENT_SHOW,
-                    false /* expected */,
-                    false /* inputViewStarted */);
-            assertWithMessage("IME is not shown after SHOW_EXPLICIT")
-                    .that(mInputMethodService.isInputViewShown()).isFalse();
+                    EVENT_SHOW, false /* eventExpected */, false /* shown */,
+                    "IME is not shown after SHOW_EXPLICIT");
         } finally {
             mInputMethodService.getResources()
                     .updateConfiguration(initialConfig, null /* metrics */, null /* compat */);
@@ -455,10 +393,7 @@ public class InputMethodServiceTest {
         // IME should be shown.
         verifyInputViewStatusOnMainSync(
                 () -> assertThat(mActivity.showImeWithInputMethodManager(0 /* flags */)).isTrue(),
-                EVENT_SHOW,
-                true /* expected */,
-                true /* inputViewStarted */);
-        assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
     }
 
     /**
@@ -472,10 +407,7 @@ public class InputMethodServiceTest {
         // the IME should be shown.
         verifyInputViewStatusOnMainSync(() -> assertThat(
                 mActivity.showImeWithInputMethodManager(InputMethodManager.SHOW_IMPLICIT)).isTrue(),
-                EVENT_SHOW,
-                true /* expected */,
-                true /* inputViewStarted */);
-        assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
     }
 
     /**
@@ -492,9 +424,9 @@ public class InputMethodServiceTest {
                 .that(mUiDevice.isNaturalOrientation()).isFalse());
         // Wait for the TestActivity to be recreated.
         eventually(() -> assertWithMessage("Activity was re-created after rotation")
-                .that(TestActivity.getLastCreatedInstance()).isNotEqualTo(mActivity));
+                .that(TestActivity.getInstance()).isNotEqualTo(mActivity));
         // Get the new TestActivity.
-        mActivity = TestActivity.getLastCreatedInstance();
+        mActivity = TestActivity.getInstance();
         assertWithMessage("Re-created activity is not null").that(mActivity).isNotNull();
         // Wait for the new EditText to be served by InputMethodManager.
         eventually(() -> assertWithMessage("Has an input connection to the re-created Activity")
@@ -502,10 +434,7 @@ public class InputMethodServiceTest {
 
         verifyInputViewStatusOnMainSync(() -> assertThat(
                         mActivity.showImeWithInputMethodManager(0 /* flags */)).isTrue(),
-                EVENT_SHOW,
-                true /* expected */,
-                true /* inputViewStarted */);
-        assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
     }
 
     /**
@@ -527,9 +456,9 @@ public class InputMethodServiceTest {
                 .that(mUiDevice.isNaturalOrientation()).isFalse());
         // Wait for the TestActivity to be recreated.
         eventually(() -> assertWithMessage("Activity was re-created after rotation")
-                .that(TestActivity.getLastCreatedInstance()).isNotEqualTo(mActivity));
+                .that(TestActivity.getInstance()).isNotEqualTo(mActivity));
         // Get the new TestActivity.
-        mActivity = TestActivity.getLastCreatedInstance();
+        mActivity = TestActivity.getInstance();
         assertWithMessage("Re-created activity is not null").that(mActivity).isNotNull();
         // Wait for the new EditText to be served by InputMethodManager.
         eventually(() -> assertWithMessage("Has an input connection to the re-created Activity")
@@ -537,11 +466,7 @@ public class InputMethodServiceTest {
 
         verifyInputViewStatusOnMainSync(() -> assertThat(
                 mActivity.showImeWithInputMethodManager(InputMethodManager.SHOW_IMPLICIT)).isTrue(),
-                EVENT_SHOW,
-                false /* expected */,
-                false /* inputViewStarted */);
-        assertWithMessage("IME is not shown")
-                .that(mInputMethodService.isInputViewShown()).isFalse();
+                EVENT_SHOW, false /* eventExpected */, false /* shown */, "IME is not shown");
     }
 
     /**
@@ -561,10 +486,7 @@ public class InputMethodServiceTest {
 
             verifyInputViewStatusOnMainSync(() -> assertThat(
                             mActivity.showImeWithInputMethodManager(0 /* flags */)).isTrue(),
-                    EVENT_SHOW,
-                    true /* expected */,
-                    true /* inputViewStarted */);
-            assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                    EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
         } finally {
             mInputMethodService.getResources()
                     .updateConfiguration(initialConfig, null /* metrics */, null /* compat */);
@@ -594,11 +516,7 @@ public class InputMethodServiceTest {
             verifyInputViewStatusOnMainSync(() ->assertThat(
                     mActivity.showImeWithInputMethodManager(InputMethodManager.SHOW_IMPLICIT))
                             .isTrue(),
-                    EVENT_SHOW,
-                    false /* expected */,
-                    false /* inputViewStarted */);
-            assertWithMessage("IME is not shown")
-                    .that(mInputMethodService.isInputViewShown()).isFalse();
+                    EVENT_SHOW, false /* eventExpected */, false /* shown */, "IME is not shown");
         } finally {
             mInputMethodService.getResources()
                     .updateConfiguration(initialConfig, null /* metrics */, null /* compat */);
@@ -623,10 +541,7 @@ public class InputMethodServiceTest {
             verifyInputViewStatusOnMainSync(
                     () -> assertThat(mActivity.showImeWithInputMethodManager(0 /* flags */))
                             .isTrue(),
-                    EVENT_SHOW,
-                    true /* expected */,
-                    true /* inputViewStarted */);
-            assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                    EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
 
             // Simulate connecting a hardware keyboard.
             config.keyboard = Configuration.KEYBOARD_QWERTY;
@@ -637,11 +552,8 @@ public class InputMethodServiceTest {
 
             verifyInputViewStatusOnMainSync(
                     () -> mInputMethodService.onConfigurationChanged(config),
-                    EVENT_CONFIG,
-                    true /* expected */,
-                    true /* inputViewStarted */);
-            assertWithMessage("IME is still shown after a configuration change")
-                    .that(mInputMethodService.isInputViewShown()).isTrue();
+                    EVENT_CONFIG, true /* eventExpected */, true /* shown */,
+                    "IME is still shown after a configuration change");
         } finally {
             mInputMethodService.getResources()
                     .updateConfiguration(initialConfig, null /* metrics */, null /* compat */);
@@ -672,10 +584,7 @@ public class InputMethodServiceTest {
             verifyInputViewStatusOnMainSync(() -> assertThat(
                             mActivity.showImeWithInputMethodManager(
                                     InputMethodManager.SHOW_IMPLICIT)).isTrue(),
-                    EVENT_SHOW,
-                    true /* expected */,
-                    true /* inputViewStarted */);
-            assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                    EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
 
             // Simulate connecting a hardware keyboard.
             config.keyboard = Configuration.KEYBOARD_QWERTY;
@@ -692,11 +601,8 @@ public class InputMethodServiceTest {
             // still alive.
             verifyInputViewStatusOnMainSync(
                     () -> mInputMethodService.onConfigurationChanged(config),
-                    EVENT_CONFIG,
-                    true /* expected */,
-                    true /* inputViewStarted */);
-            assertWithMessage("IME is not shown after a configuration change")
-                    .that(mInputMethodService.isInputViewShown()).isFalse();
+                    EVENT_CONFIG, true /* eventExpected */, true /* inputViewStarted */,
+                    false /* shown */, "IME is not shown after a configuration change");
         } finally {
             mInputMethodService.getResources()
                     .updateConfiguration(initialConfig, null /* metrics */, null /* compat */);
@@ -722,31 +628,21 @@ public class InputMethodServiceTest {
             // Explicit show request.
             verifyInputViewStatusOnMainSync(() -> assertThat(
                             mActivity.showImeWithInputMethodManager(0 /* flags */)).isTrue(),
-                    EVENT_SHOW,
-                    true /* expected */,
-                    true /* inputViewStarted */);
-            assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                    EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
 
             // Implicit show request.
             verifyInputViewStatusOnMainSync(() -> assertThat(
                             mActivity.showImeWithInputMethodManager(
                                     InputMethodManager.SHOW_IMPLICIT)).isTrue(),
-                    EVENT_SHOW,
-                    false /* expected */,
-                    true /* inputViewStarted */);
-            assertWithMessage("IME is still shown")
-                    .that(mInputMethodService.isInputViewShown()).isTrue();
+                    EVENT_SHOW, false /* eventExpected */, true /* shown */, "IME is still shown");
 
             // Simulate a fake configuration change to avoid the recreation of TestActivity.
             // This should now consider the implicit show request, but keep the state from the
             // explicit show request, and thus not hide the IME.
             verifyInputViewStatusOnMainSync(
                     () -> mInputMethodService.onConfigurationChanged(config),
-                    EVENT_CONFIG,
-                    true /* expected */,
-                    true /* inputViewStarted */);
-            assertWithMessage("IME is still shown after a configuration change")
-                    .that(mInputMethodService.isInputViewShown()).isTrue();
+                    EVENT_CONFIG, true /* eventExpected */, true /* shown */,
+                    "IME is still shown after a configuration change");
         } finally {
             mInputMethodService.getResources()
                     .updateConfiguration(initialConfig, null /* metrics */, null /* compat */);
@@ -769,27 +665,17 @@ public class InputMethodServiceTest {
 
         verifyInputViewStatusOnMainSync(() -> assertThat(
                 mActivity.showImeWithInputMethodManager(InputMethodManager.SHOW_FORCED)).isTrue(),
-                EVENT_SHOW,
-                true /* expected */,
-                true /* inputViewStarted */);
-        assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
 
         verifyInputViewStatusOnMainSync(() -> assertThat(
                         mActivity.showImeWithInputMethodManager(0 /* flags */)).isTrue(),
-                EVENT_SHOW,
-                false /* expected */,
-                true /* inputViewStarted */);
-        assertWithMessage("IME is still shown")
-                .that(mInputMethodService.isInputViewShown()).isTrue();
+                EVENT_SHOW, false /* eventExpected */, true /* shown */, "IME is still shown");
 
         verifyInputViewStatusOnMainSync(() -> assertThat(
                         mActivity.hideImeWithInputMethodManager(InputMethodManager.HIDE_NOT_ALWAYS))
                         .isTrue(),
-                EVENT_HIDE,
-                false /* expected */,
-                true /* inputViewStarted */);
-        assertWithMessage("IME is still shown after HIDE_NOT_ALWAYS")
-                .that(mInputMethodService.isInputViewShown()).isTrue();
+                EVENT_HIDE, false /* eventExpected */, true /* shown */,
+                "IME is still shown after HIDE_NOT_ALWAYS");
     }
 
     /**
@@ -800,18 +686,15 @@ public class InputMethodServiceTest {
         setShowImeWithHardKeyboard(true /* enabled */);
 
         Log.i(TAG, "Set orientation natural");
-        verifyFullscreenMode(() -> setOrientation(0),
-                false /* expected */,
+        verifyFullscreenMode(() -> setOrientation(0), false /* eventExpected */,
                 true /* orientationPortrait */);
 
         Log.i(TAG, "Set orientation left");
-        verifyFullscreenMode(() -> setOrientation(1),
-                true /* expected */,
+        verifyFullscreenMode(() -> setOrientation(1), true /* eventExpected */,
                 false /* orientationPortrait */);
 
         Log.i(TAG, "Set orientation right");
-        verifyFullscreenMode(() -> setOrientation(2),
-                false /* expected */,
+        verifyFullscreenMode(() -> setOrientation(2), false /* eventExpected */,
                 false /* orientationPortrait */);
     }
 
@@ -845,10 +728,7 @@ public class InputMethodServiceTest {
                     setDrawsImeNavBarAndSwitcherButton(true /* enabled */);
                     mActivity.showImeWithWindowInsetsController();
                 },
-                EVENT_SHOW,
-                true /* expected */,
-                true /* inputViewStarted */);
-        assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
         assertWithMessage("IME navigation bar is initially shown")
                 .that(mInputMethodService.isImeNavigationBarShownForTesting()).isTrue();
 
@@ -883,10 +763,7 @@ public class InputMethodServiceTest {
                     setDrawsImeNavBarAndSwitcherButton(false /* enabled */);
                     mActivity.showImeWithWindowInsetsController();
                 },
-                EVENT_SHOW,
-                true /* expected */,
-                true /* inputViewStarted */);
-        assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
         assertWithMessage("IME navigation bar is initially not shown")
                 .that(mInputMethodService.isImeNavigationBarShownForTesting()).isFalse();
 
@@ -917,24 +794,15 @@ public class InputMethodServiceTest {
         try (var ignored = mGestureNavSwitchHelper.withGestureNavigationMode()) {
             verifyInputViewStatusOnMainSync(
                     () -> mActivity.showImeWithWindowInsetsController(),
-                    EVENT_SHOW,
-                    true /* expected */,
-                    true /* inputViewStarted */);
-            assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                    EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
 
             final var backButton = getUiObject(By.res(INPUT_METHOD_NAV_BACK_ID));
-            backButton.click();
-            mInstrumentation.waitForIdleSync();
-
-            if (mFlagsValueProvider.getBoolean(Flags.FLAG_REFACTOR_INSETS_CONTROLLER)) {
-                // The IME visibility is only sent at the end of the animation. Therefore, we have
-                // to wait until the visibility was sent to the server and the IME window hidden.
-                eventually(() -> assertWithMessage("IME is not shown")
-                        .that(mInputMethodService.isInputViewShown()).isFalse());
-            } else {
-                assertWithMessage("IME is not shown")
-                        .that(mInputMethodService.isInputViewShown()).isFalse();
-            }
+            verifyInputViewStatus(
+                    () -> {
+                        backButton.click();
+                        mInstrumentation.waitForIdleSync();
+                    },
+                    EVENT_HIDE, true /* eventExpected */, false /* shown */, "IME is not shown");
         }
     }
 
@@ -952,24 +820,15 @@ public class InputMethodServiceTest {
         try (var ignored = mGestureNavSwitchHelper.withGestureNavigationMode()) {
             verifyInputViewStatusOnMainSync(
                     () -> mActivity.showImeWithWindowInsetsController(),
-                    EVENT_SHOW,
-                    true /* expected */,
-                    true /* inputViewStarted */);
-            assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                    EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
 
             final var backButton = getUiObject(By.res(INPUT_METHOD_NAV_BACK_ID));
-            backButton.longClick();
-            mInstrumentation.waitForIdleSync();
-
-            if (mFlagsValueProvider.getBoolean(Flags.FLAG_REFACTOR_INSETS_CONTROLLER)) {
-                // The IME visibility is only sent at the end of the animation. Therefore, we have
-                // to wait until the visibility was sent to the server and the IME window hidden.
-                eventually(() -> assertWithMessage("IME is not shown")
-                        .that(mInputMethodService.isInputViewShown()).isFalse());
-            } else {
-                assertWithMessage("IME is not shown")
-                        .that(mInputMethodService.isInputViewShown()).isFalse();
-            }
+            verifyInputViewStatus(
+                    () -> {
+                        backButton.longClick();
+                        mInstrumentation.waitForIdleSync();
+                    },
+                    EVENT_HIDE, true /* eventExpected */, false /* shown */, "IME is not shown");
         }
     }
 
@@ -991,10 +850,7 @@ public class InputMethodServiceTest {
                         setDrawsImeNavBarAndSwitcherButton(true /* enabled */);
                         mActivity.showImeWithWindowInsetsController();
                     },
-                    EVENT_SHOW,
-                    true /* expected */,
-                    true /* inputViewStarted */);
-            assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                    EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
 
             final var initialInfo = mImm.getCurrentInputMethodInfo();
 
@@ -1033,10 +889,7 @@ public class InputMethodServiceTest {
                         setDrawsImeNavBarAndSwitcherButton(true /* enabled */);
                         mActivity.showImeWithWindowInsetsController();
                     },
-                    EVENT_SHOW,
-                    true /* expected */,
-                    true /* inputViewStarted */);
-            assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+                    EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
 
             final var imeSwitcherButton = getUiObject(By.res(INPUT_METHOD_NAV_IME_SWITCHER_ID));
             imeSwitcherButton.longClick();
@@ -1052,58 +905,78 @@ public class InputMethodServiceTest {
         }
     }
 
-    private void verifyInputViewStatus(@NonNull Runnable runnable, @Event int event,
-            boolean expected, boolean inputViewStarted) {
-        verifyInputViewStatusInternal(runnable, event, expected, inputViewStarted,
-                false /* runOnMainSync */);
+    private void verifyInputViewStatus(@NonNull Runnable runnable, @Event int eventType,
+            boolean eventExpected, boolean shown, @NonNull String message) {
+        verifyInputViewStatusInternal(runnable, eventType, eventExpected,
+                shown /* inputViewStarted */, shown, message, false /* runOnMainSync */);
     }
 
-    private void verifyInputViewStatusOnMainSync(@NonNull Runnable runnable, @Event int event,
-            boolean expected, boolean inputViewStarted) {
-        verifyInputViewStatusInternal(runnable, event, expected, inputViewStarted,
-                true /* runOnMainSync */);
+    private void verifyInputViewStatusOnMainSync(@NonNull Runnable runnable, @Event int eventType,
+            boolean eventExpected, boolean shown, @NonNull String message) {
+        verifyInputViewStatusInternal(runnable, eventType, eventExpected,
+                shown /* inputViewStarted */, shown, message, true /* runOnMainSync */);
+    }
+
+    private void verifyInputViewStatusOnMainSync(@NonNull Runnable runnable, @Event int eventType,
+            boolean eventExpected, boolean inputViewStarted, boolean shown,
+            @NonNull String message) {
+        verifyInputViewStatusInternal(runnable, eventType, eventExpected, inputViewStarted, shown,
+                message, true /* runOnMainSync */);
     }
 
     /**
      * Verifies the status of the Input View after executing the given runnable, and waiting that
-     * the event was either triggered or not, based on the given expectation.
+     * the event was either called or not.
      *
-     * @param runnable         the runnable to trigger the event
-     * @param event            the event to await.
-     * @param expected         whether the event is expected to be triggered.
-     * @param inputViewStarted the expected state of the Input View after executing the runnable.
+     * @param runnable         the runnable to call the event.
+     * @param eventType        the event type to wait for.
+     * @param eventExpected    whether the event is expected to be called.
+     * @param inputViewStarted whether the input view is expected to be started.
+     * @param shown            whether the input view is expected to be shown.
+     * @param message          the message for the input view shown assertion.
      * @param runOnMainSync    whether to execute the runnable on the main thread.
      */
-    private void verifyInputViewStatusInternal(@NonNull Runnable runnable, @Event int event,
-            boolean expected, boolean inputViewStarted, boolean runOnMainSync) {
-        final boolean completed;
+    private void verifyInputViewStatusInternal(@NonNull Runnable runnable, @Event int eventType,
+            boolean eventExpected, boolean inputViewStarted, boolean shown, @NonNull String message,
+            boolean runOnMainSync) {
+        final boolean eventCalled;
         try {
             final var latch = new CountDownLatch(1);
-            mInputMethodService.setCountDownLatchForTesting(latch, event);
+            mInputMethodService.setCountDownLatchForTesting(latch, eventType);
             if (runOnMainSync) {
                 mInstrumentation.runOnMainSync(runnable);
             } else {
                 runnable.run();
             }
             mInstrumentation.waitForIdleSync();
-            completed = latch.await(TIMEOUT_IN_SECONDS, TimeUnit.SECONDS);
+            eventCalled = latch.await(TIMEOUT_IN_SECONDS, TimeUnit.SECONDS);
         } catch (InterruptedException e) {
             fail("Interrupted while waiting for latch: " + e.getMessage());
             return;
         } finally {
-            mInputMethodService.setCountDownLatchForTesting(null /* latch */, event);
+            mInputMethodService.setCountDownLatchForTesting(null /* latch */, eventType);
         }
 
-        if (expected && !completed) {
-            fail("Timed out waiting for " + eventToString(event));
-        } else if (!expected && completed) {
-            fail("Unexpected call " + eventToString(event));
+        if (eventExpected && !eventCalled) {
+            fail("Timed out waiting for " + eventToString(eventType));
+        } else if (!eventExpected && eventCalled) {
+            fail("Unexpected call " + eventToString(eventType));
         }
-        // Input is not finished.
+        // Input connection is not finished.
         assertWithMessage("Input connection is still started")
                 .that(mInputMethodService.getCurrentInputStarted()).isTrue();
-        assertWithMessage("IME visibility matches expected")
+        assertWithMessage("Input view started matches expected")
                 .that(mInputMethodService.getCurrentInputViewStarted()).isEqualTo(inputViewStarted);
+
+        if (mFlagsValueProvider.getBoolean(Flags.FLAG_REFACTOR_INSETS_CONTROLLER)) {
+            // The IME visibility is only sent at the end of the hide animation. Therefore, we have
+            // to wait until the visibility was sent to the server and the IME window hidden.
+            eventually(() -> assertWithMessage(message).that(mInputMethodService.isInputViewShown())
+                    .isEqualTo(shown));
+        } else {
+            assertWithMessage(message).that(mInputMethodService.isInputViewShown())
+                    .isEqualTo(shown);
+        }
     }
 
     private void setOrientation(int orientation) {
@@ -1127,31 +1000,28 @@ public class InputMethodServiceTest {
     /**
      * Verifies the IME fullscreen mode state after executing the given runnable.
      *
-     * @param runnable            the runnable to execute for setting the orientation.
-     * @param expected            whether the runnable is expected to trigger the signal.
+     * @param runnable            the runnable to set the orientation.
+     * @param eventExpected       whether the event is expected to be called.
      * @param orientationPortrait whether the orientation is expected to be portrait.
      */
-    private void verifyFullscreenMode(@NonNull Runnable runnable, boolean expected,
+    private void verifyFullscreenMode(@NonNull Runnable runnable, boolean eventExpected,
             boolean orientationPortrait) {
-        verifyInputViewStatus(runnable, EVENT_CONFIG, expected, false /* inputViewStarted */);
-        if (expected) {
+        verifyInputViewStatus(runnable, EVENT_CONFIG, eventExpected, false /* shown */,
+                "IME is not shown");
+        if (eventExpected) {
             // Wait for the TestActivity to be recreated.
             eventually(() -> assertWithMessage("Activity was re-created after rotation")
-                    .that(TestActivity.getLastCreatedInstance()).isNotEqualTo(mActivity));
+                    .that(TestActivity.getInstance()).isNotEqualTo(mActivity));
             // Get the new TestActivity.
-            mActivity = TestActivity.getLastCreatedInstance();
+            mActivity = TestActivity.getInstance();
             assertWithMessage("Re-created activity is not null").that(mActivity).isNotNull();
             // Wait for the new EditText to be served by InputMethodManager.
             eventually(() -> assertWithMessage("Has an input connection to the re-created Activity")
                     .that(mImm.hasActiveInputConnection(mActivity.getEditText())).isTrue());
         }
 
-        verifyInputViewStatusOnMainSync(
-                () -> mActivity.showImeWithWindowInsetsController(),
-                EVENT_SHOW,
-                true /* expected */,
-                true /* inputViewStarted */);
-        assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
+        verifyInputViewStatusOnMainSync(() -> mActivity.showImeWithWindowInsetsController(),
+                EVENT_SHOW, true /* eventExpected */, true /* shown */, "IME is shown");
 
         assertWithMessage("IME orientation matches expected")
                 .that(mInputMethodService.getResources().getConfiguration().orientation)
@@ -1172,21 +1042,8 @@ public class InputMethodServiceTest {
                 .that(mInputMethodService.isFullscreenMode()).isEqualTo(!orientationPortrait);
 
         // Hide IME before finishing the run.
-        verifyInputViewStatusOnMainSync(
-                () -> mActivity.hideImeWithWindowInsetsController(),
-                EVENT_HIDE,
-                true /* expected */,
-                false /* inputViewStarted */);
-
-        if (mFlagsValueProvider.getBoolean(Flags.FLAG_REFACTOR_INSETS_CONTROLLER)) {
-            // The IME visibility is only sent at the end of the animation. Therefore, we have to
-            // wait until the visibility was sent to the server and the IME window hidden.
-            eventually(() -> assertWithMessage("IME is not shown")
-                    .that(mInputMethodService.isInputViewShown()).isFalse());
-        } else {
-            assertWithMessage("IME is not shown")
-                    .that(mInputMethodService.isInputViewShown()).isFalse();
-        }
+        verifyInputViewStatusOnMainSync(() -> mActivity.hideImeWithWindowInsetsController(),
+                EVENT_HIDE, true /* eventExpected */, false /* shown */, "IME is not shown");
     }
 
     private void prepareIme() {
diff --git a/services/tests/InputMethodSystemServerTests/test-apps/SimpleTestIme/src/com/android/apps/inputmethod/simpleime/ims/InputMethodServiceWrapper.java b/services/tests/InputMethodSystemServerTests/test-apps/SimpleTestIme/src/com/android/apps/inputmethod/simpleime/ims/InputMethodServiceWrapper.java
index 3a7abbb167ec..558d1a7c4e8b 100644
--- a/services/tests/InputMethodSystemServerTests/test-apps/SimpleTestIme/src/com/android/apps/inputmethod/simpleime/ims/InputMethodServiceWrapper.java
+++ b/services/tests/InputMethodSystemServerTests/test-apps/SimpleTestIme/src/com/android/apps/inputmethod/simpleime/ims/InputMethodServiceWrapper.java
@@ -27,6 +27,7 @@ import androidx.annotation.Nullable;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
+import java.lang.ref.WeakReference;
 import java.util.concurrent.CountDownLatch;
 
 /** Wrapper of {@link InputMethodService} to expose interfaces for testing purpose. */
@@ -35,8 +36,8 @@ public class InputMethodServiceWrapper extends InputMethodService {
     private static final String TAG = "InputMethodServiceWrapper";
 
     /** Last created instance of this wrapper. */
-    @Nullable
-    private static InputMethodServiceWrapper sInstance;
+    @NonNull
+    private static WeakReference<InputMethodServiceWrapper> sInstance = new WeakReference<>(null);
 
     /** IME show event ({@link #onStartInputView}). */
     public static final int EVENT_SHOW = 0;
@@ -68,32 +69,11 @@ public class InputMethodServiceWrapper extends InputMethodService {
     @Nullable
     private CountDownLatch mCountDownLatch;
 
-    /** Gets the last created instance of this wrapper, if available. */
-    @Nullable
-    public static InputMethodServiceWrapper getInstance() {
-        return sInstance;
-    }
-
-    public boolean getCurrentInputViewStarted() {
-        return mInputViewStarted;
-    }
-
-    /**
-     * Sets the latch used to wait for the IME event.
-     *
-     * @param latch      the latch to wait on.
-     * @param latchEvent the event to set the latch on.
-     */
-    public void setCountDownLatchForTesting(@Nullable CountDownLatch latch, @Event int latchEvent) {
-        mCountDownLatch = latch;
-        mLatchEvent = latchEvent;
-    }
-
     @Override
     public void onCreate() {
         Log.i(TAG, "onCreate()");
         super.onCreate();
-        sInstance = this;
+        sInstance = new WeakReference<>(this);
     }
 
     @Override
@@ -103,6 +83,12 @@ public class InputMethodServiceWrapper extends InputMethodService {
     }
 
     @Override
+    public void onFinishInput() {
+        Log.i(TAG, "onFinishInput()");
+        super.onFinishInput();
+    }
+
+    @Override
     public void onStartInputView(EditorInfo info, boolean restarting) {
         Log.i(TAG, "onStartInputView() editor=" + dumpEditorInfo(info)
                 + ", restarting=" + restarting);
@@ -114,12 +100,6 @@ public class InputMethodServiceWrapper extends InputMethodService {
     }
 
     @Override
-    public void onFinishInput() {
-        Log.i(TAG, "onFinishInput()");
-        super.onFinishInput();
-    }
-
-    @Override
     public void onFinishInputView(boolean finishingInput) {
         Log.i(TAG, "onFinishInputView()");
         super.onFinishInputView(finishingInput);
@@ -146,14 +126,35 @@ public class InputMethodServiceWrapper extends InputMethodService {
         }
     }
 
+    public boolean getCurrentInputViewStarted() {
+        return mInputViewStarted;
+    }
+
+    /**
+     * Sets the latch used to wait for the IME event.
+     *
+     * @param latch      the latch to wait on.
+     * @param latchEvent the event to set the latch on.
+     */
+    public void setCountDownLatchForTesting(@Nullable CountDownLatch latch, @Event int latchEvent) {
+        mCountDownLatch = latch;
+        mLatchEvent = latchEvent;
+    }
+
+    /** Gets the last created instance of this wrapper, if available. */
+    @Nullable
+    public static InputMethodServiceWrapper getInstance() {
+        return sInstance.get();
+    }
+
     /**
      * Gets the string representation of the IME event that is being waited on.
      *
-     * @param event the IME event.
+     * @param eventType the IME event type.
      */
     @NonNull
-    public static String eventToString(@Event int event) {
-        return switch (event) {
+    public static String eventToString(@Event int eventType) {
+        return switch (eventType) {
             case EVENT_SHOW -> "onStartInputView";
             case EVENT_HIDE -> "onFinishInputView";
             case EVENT_CONFIG -> "onConfigurationChanged";
diff --git a/services/tests/InputMethodSystemServerTests/test-apps/SimpleTestIme/src/com/android/apps/inputmethod/simpleime/testing/TestActivity.java b/services/tests/InputMethodSystemServerTests/test-apps/SimpleTestIme/src/com/android/apps/inputmethod/simpleime/testing/TestActivity.java
index eadbac3ca74b..ad90b32848b0 100644
--- a/services/tests/InputMethodSystemServerTests/test-apps/SimpleTestIme/src/com/android/apps/inputmethod/simpleime/testing/TestActivity.java
+++ b/services/tests/InputMethodSystemServerTests/test-apps/SimpleTestIme/src/com/android/apps/inputmethod/simpleime/testing/TestActivity.java
@@ -45,30 +45,13 @@ import java.lang.ref.WeakReference;
 public final class TestActivity extends Activity {
 
     private static final String TAG = "TestActivity";
-    private static WeakReference<TestActivity> sLastCreatedInstance = new WeakReference<>(null);
 
-    /**
-     * Start a new test activity with an editor and wait for it to begin running before returning.
-     *
-     * @param instrumentation application instrumentation
-     * @return the newly started activity
-     */
+    /** Last created instance of this activity. */
     @NonNull
-    public static TestActivity startSync(@NonNull Instrumentation instrumentation) {
-        final var intent = new Intent()
-                .setAction(Intent.ACTION_MAIN)
-                .setClass(instrumentation.getTargetContext(), TestActivity.class)
-                .addFlags(Intent.FLAG_ACTIVITY_NEW_TASK)
-                .addFlags(Intent.FLAG_ACTIVITY_CLEAR_TOP);
-        return (TestActivity) instrumentation.startActivitySync(intent);
-    }
+    private static WeakReference<TestActivity> sInstance = new WeakReference<>(null);
 
     private EditText mEditText;
 
-    public EditText getEditText() {
-        return mEditText;
-    }
-
     @Override
     protected void onCreate(@Nullable Bundle savedInstanceState) {
         super.onCreate(savedInstanceState);
@@ -80,13 +63,11 @@ public final class TestActivity extends Activity {
         rootView.setFitsSystemWindows(true);
         setContentView(rootView);
         mEditText.requestFocus();
-        sLastCreatedInstance = new WeakReference<>(this);
+        sInstance = new WeakReference<>(this);
     }
 
-    /** Get the last created TestActivity instance, if available. */
-    @Nullable
-    public static TestActivity getLastCreatedInstance() {
-        return sLastCreatedInstance.get();
+    public EditText getEditText() {
+        return mEditText;
     }
 
     /** Shows soft keyboard via InputMethodManager. */
@@ -118,4 +99,26 @@ public final class TestActivity extends Activity {
         controller.hide(WindowInsets.Type.ime());
         Log.i(TAG, "hideIme() via WindowInsetsController");
     }
+
+    /** Gets the last created instance of this activity, if available. */
+    @Nullable
+    public static TestActivity getInstance() {
+        return sInstance.get();
+    }
+
+    /**
+     * Start a new test activity with an editor and wait for it to begin running before returning.
+     *
+     * @param instrumentation application instrumentation.
+     * @return the newly started activity.
+     */
+    @NonNull
+    public static TestActivity startSync(@NonNull Instrumentation instrumentation) {
+        final var intent = new Intent()
+                .setAction(Intent.ACTION_MAIN)
+                .setClass(instrumentation.getTargetContext(), TestActivity.class)
+                .addFlags(Intent.FLAG_ACTIVITY_NEW_TASK)
+                .addFlags(Intent.FLAG_ACTIVITY_CLEAR_TOP);
+        return (TestActivity) instrumentation.startActivitySync(intent);
+    }
 }