Specify event for latch in InputMethodServiceTest

Currently we set and wait on a single countDownLatch for either of
onStartInputView, onFinishInputView or onConfigurationChanged. This can
introduce some noise if a different event than the expected one is
trigerred (e.g. onConfigurationChanged while waiting for
onStartInputView). This also introduces some ambiguity as it's not clear
which event is the expected one for each test.

This introduces an IntDef to specify the expected event, and only
triggeres the latch for that specific event.

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

3 files changed, 127 insertions, 32 deletions

diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/inputmethodservice/InputMethodServiceTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/inputmethodservice/InputMethodServiceTest.java
index a7280c2167ea..0b45e8396bf9 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/inputmethodservice/InputMethodServiceTest.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/inputmethodservice/InputMethodServiceTest.java
@@ -18,6 +18,10 @@ package com.android.inputmethodservice;
 
 import static android.view.WindowInsets.Type.captionBar;
 
+import static com.android.apps.inputmethod.simpleime.ims.InputMethodServiceWrapper.EVENT_CONFIG;
+import static com.android.apps.inputmethod.simpleime.ims.InputMethodServiceWrapper.EVENT_HIDE;
+import static com.android.apps.inputmethod.simpleime.ims.InputMethodServiceWrapper.EVENT_SHOW;
+import static com.android.apps.inputmethod.simpleime.ims.InputMethodServiceWrapper.eventToString;
 import static com.android.compatibility.common.util.SystemUtil.eventually;
 import static com.android.cts.input.injectinputinprocess.InjectInputInProcessKt.clickOnViewCenter;
 import static com.android.internal.inputmethod.InputMethodNavButtonFlags.IME_DRAWS_IME_NAV_BAR;
@@ -56,6 +60,7 @@ import androidx.test.uiautomator.UiObject2;
 import androidx.test.uiautomator.Until;
 
 import com.android.apps.inputmethod.simpleime.ims.InputMethodServiceWrapper;
+import com.android.apps.inputmethod.simpleime.ims.InputMethodServiceWrapper.Event;
 import com.android.apps.inputmethod.simpleime.testing.TestActivity;
 import com.android.compatibility.common.util.GestureNavSwitchHelper;
 import com.android.compatibility.common.util.SystemUtil;
@@ -168,6 +173,7 @@ 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();
@@ -184,6 +190,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")
@@ -201,6 +208,7 @@ 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();
@@ -208,6 +216,7 @@ public class InputMethodServiceTest {
         // 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)) {
@@ -231,6 +240,7 @@ 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();
@@ -238,6 +248,7 @@ public class InputMethodServiceTest {
         // Triggers to hide IME via public API.
         verifyInputViewStatusOnMainSync(
                 () -> mActivity.hideImeWithWindowInsetsController(),
+                EVENT_HIDE,
                 true /* expected */,
                 false /* inputViewStarted */);
         if (mFlagsValueProvider.getBoolean(Flags.FLAG_REFACTOR_INSETS_CONTROLLER)) {
@@ -266,6 +277,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();
@@ -276,6 +288,7 @@ 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")
@@ -286,6 +299,7 @@ public class InputMethodServiceTest {
         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)) {
@@ -303,6 +317,7 @@ public class InputMethodServiceTest {
             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")
@@ -313,6 +328,7 @@ public class InputMethodServiceTest {
             verifyInputViewStatusOnMainSync(
                     () -> mInputMethodService.requestHideSelf(
                             InputMethodManager.HIDE_IMPLICIT_ONLY),
+                    EVENT_HIDE,
                     true /* expected */,
                     false /* inputViewStarted */);
             assertWithMessage("IME is not shown after HIDE_IMPLICIT_ONLY")
@@ -408,6 +424,7 @@ 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")
@@ -416,6 +433,7 @@ public class InputMethodServiceTest {
             verifyInputViewStatusOnMainSync(
                     () -> assertThat(mActivity.showImeWithInputMethodManager(0 /* flags */))
                             .isTrue(),
+                    EVENT_SHOW,
                     false /* expected */,
                     false /* inputViewStarted */);
             assertWithMessage("IME is not shown after SHOW_EXPLICIT")
@@ -437,6 +455,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();
@@ -453,6 +472,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();
@@ -482,6 +502,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();
@@ -516,6 +537,7 @@ public class InputMethodServiceTest {
 
         verifyInputViewStatusOnMainSync(() -> assertThat(
                 mActivity.showImeWithInputMethodManager(InputMethodManager.SHOW_IMPLICIT)).isTrue(),
+                EVENT_SHOW,
                 false /* expected */,
                 false /* inputViewStarted */);
         assertWithMessage("IME is not shown")
@@ -539,6 +561,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();
@@ -571,6 +594,7 @@ public class InputMethodServiceTest {
             verifyInputViewStatusOnMainSync(() ->assertThat(
                     mActivity.showImeWithInputMethodManager(InputMethodManager.SHOW_IMPLICIT))
                             .isTrue(),
+                    EVENT_SHOW,
                     false /* expected */,
                     false /* inputViewStarted */);
             assertWithMessage("IME is not shown")
@@ -599,6 +623,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();
@@ -612,6 +637,7 @@ public class InputMethodServiceTest {
 
             verifyInputViewStatusOnMainSync(
                     () -> mInputMethodService.onConfigurationChanged(config),
+                    EVENT_CONFIG,
                     true /* expected */,
                     true /* inputViewStarted */);
             assertWithMessage("IME is still shown after a configuration change")
@@ -646,6 +672,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();
@@ -665,6 +692,7 @@ public class InputMethodServiceTest {
             // still alive.
             verifyInputViewStatusOnMainSync(
                     () -> mInputMethodService.onConfigurationChanged(config),
+                    EVENT_CONFIG,
                     true /* expected */,
                     true /* inputViewStarted */);
             assertWithMessage("IME is not shown after a configuration change")
@@ -694,6 +722,7 @@ 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();
@@ -702,6 +731,7 @@ public class InputMethodServiceTest {
             verifyInputViewStatusOnMainSync(() -> assertThat(
                             mActivity.showImeWithInputMethodManager(
                                     InputMethodManager.SHOW_IMPLICIT)).isTrue(),
+                    EVENT_SHOW,
                     false /* expected */,
                     true /* inputViewStarted */);
             assertWithMessage("IME is still shown")
@@ -712,6 +742,7 @@ public class InputMethodServiceTest {
             // 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")
@@ -738,12 +769,14 @@ 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();
 
         verifyInputViewStatusOnMainSync(() -> assertThat(
                         mActivity.showImeWithInputMethodManager(0 /* flags */)).isTrue(),
+                EVENT_SHOW,
                 false /* expected */,
                 true /* inputViewStarted */);
         assertWithMessage("IME is still shown")
@@ -752,6 +785,7 @@ public class InputMethodServiceTest {
         verifyInputViewStatusOnMainSync(() -> assertThat(
                         mActivity.hideImeWithInputMethodManager(InputMethodManager.HIDE_NOT_ALWAYS))
                         .isTrue(),
+                EVENT_HIDE,
                 false /* expected */,
                 true /* inputViewStarted */);
         assertWithMessage("IME is still shown after HIDE_NOT_ALWAYS")
@@ -811,6 +845,7 @@ public class InputMethodServiceTest {
                     setDrawsImeNavBarAndSwitcherButton(true /* enabled */);
                     mActivity.showImeWithWindowInsetsController();
                 },
+                EVENT_SHOW,
                 true /* expected */,
                 true /* inputViewStarted */);
         assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
@@ -848,6 +883,7 @@ public class InputMethodServiceTest {
                     setDrawsImeNavBarAndSwitcherButton(false /* enabled */);
                     mActivity.showImeWithWindowInsetsController();
                 },
+                EVENT_SHOW,
                 true /* expected */,
                 true /* inputViewStarted */);
         assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
@@ -886,6 +922,7 @@ public class InputMethodServiceTest {
                 //  Wait for onConfigurationChanged when changing navigation modes.
                 verifyInputViewStatus(
                         () -> restoreNav[0] = mGestureNavSwitchHelper.withGestureNavigationMode(),
+                        EVENT_CONFIG,
                         true, /* expected */
                         false /* inputViewStarted */
                 );
@@ -893,6 +930,7 @@ public class InputMethodServiceTest {
 
             verifyInputViewStatusOnMainSync(
                     () -> mActivity.showImeWithWindowInsetsController(),
+                    EVENT_SHOW,
                     true /* expected */,
                     true /* inputViewStarted */);
             assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
@@ -934,6 +972,7 @@ public class InputMethodServiceTest {
                 //  Wait for onConfigurationChanged when changing navigation modes.
                 verifyInputViewStatus(
                         () -> restoreNav[0] = mGestureNavSwitchHelper.withGestureNavigationMode(),
+                        EVENT_CONFIG,
                         true, /* expected */
                         false /* inputViewStarted */
                 );
@@ -941,6 +980,7 @@ public class InputMethodServiceTest {
 
             verifyInputViewStatusOnMainSync(
                     () -> mActivity.showImeWithWindowInsetsController(),
+                    EVENT_SHOW,
                     true /* expected */,
                     true /* inputViewStarted */);
             assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
@@ -983,6 +1023,7 @@ public class InputMethodServiceTest {
                 //  Wait for onConfigurationChanged when changing navigation modes.
                 verifyInputViewStatus(
                         () -> restoreNav[0] = mGestureNavSwitchHelper.withGestureNavigationMode(),
+                        EVENT_CONFIG,
                         true, /* expected */
                         false /* inputViewStarted */
                 );
@@ -993,6 +1034,7 @@ public class InputMethodServiceTest {
                         setDrawsImeNavBarAndSwitcherButton(true /* enabled */);
                         mActivity.showImeWithWindowInsetsController();
                     },
+                    EVENT_SHOW,
                     true /* expected */,
                     true /* inputViewStarted */);
             assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
@@ -1038,6 +1080,7 @@ public class InputMethodServiceTest {
                 //  Wait for onConfigurationChanged when changing navigation modes.
                 verifyInputViewStatus(
                         () -> restoreNav[0] = mGestureNavSwitchHelper.withGestureNavigationMode(),
+                        EVENT_CONFIG,
                         true, /* expected */
                         false /* inputViewStarted */
                 );
@@ -1048,6 +1091,7 @@ public class InputMethodServiceTest {
                         setDrawsImeNavBarAndSwitcherButton(true /* enabled */);
                         mActivity.showImeWithWindowInsetsController();
                     },
+                    EVENT_SHOW,
                     true /* expected */,
                     true /* inputViewStarted */);
             assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
@@ -1068,33 +1112,34 @@ public class InputMethodServiceTest {
         }
     }
 
-    private void verifyInputViewStatus(@NonNull Runnable runnable, boolean expected,
-            boolean inputViewStarted) {
-        verifyInputViewStatusInternal(runnable, expected, inputViewStarted,
+    private void verifyInputViewStatus(@NonNull Runnable runnable, @Event int event,
+            boolean expected, boolean inputViewStarted) {
+        verifyInputViewStatusInternal(runnable, event, expected, inputViewStarted,
                 false /* runOnMainSync */);
     }
 
-    private void verifyInputViewStatusOnMainSync(@NonNull Runnable runnable, boolean expected,
-            boolean inputViewStarted) {
-        verifyInputViewStatusInternal(runnable, expected, inputViewStarted,
+    private void verifyInputViewStatusOnMainSync(@NonNull Runnable runnable, @Event int event,
+            boolean expected, boolean inputViewStarted) {
+        verifyInputViewStatusInternal(runnable, event, expected, inputViewStarted,
                 true /* runOnMainSync */);
     }
 
     /**
-     * Verifies the status of the Input View after executing the given runnable.
+     * 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.
      *
-     * @param runnable         the runnable to execute for showing or hiding the IME.
-     * @param expected         whether the runnable is expected to trigger the signal.
+     * @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 runOnMainSync    whether to execute the runnable on the main thread.
      */
-    private void verifyInputViewStatusInternal(@NonNull Runnable runnable, boolean expected,
-            boolean inputViewStarted, boolean runOnMainSync) {
+    private void verifyInputViewStatusInternal(@NonNull Runnable runnable, @Event int event,
+            boolean expected, boolean inputViewStarted, boolean runOnMainSync) {
         final boolean completed;
         try {
             final var latch = new CountDownLatch(1);
-            mInputMethodService.setCountDownLatchForTesting(latch);
-            // Trigger onStartInputView() / onFinishInputView() / onConfigurationChanged()
+            mInputMethodService.setCountDownLatchForTesting(latch, event);
             if (runOnMainSync) {
                 mInstrumentation.runOnMainSync(runnable);
             } else {
@@ -1106,15 +1151,13 @@ public class InputMethodServiceTest {
             fail("Interrupted while waiting for latch: " + e.getMessage());
             return;
         } finally {
-            mInputMethodService.setCountDownLatchForTesting(null);
+            mInputMethodService.setCountDownLatchForTesting(null /* latch */, event);
         }
 
         if (expected && !completed) {
-            fail("Timed out waiting for"
-                    + " onStartInputView() / onFinishInputView() / onConfigurationChanged()");
+            fail("Timed out waiting for " + eventToString(event));
         } else if (!expected && completed) {
-            fail("Unexpected call"
-                    + " onStartInputView() / onFinishInputView() / onConfigurationChanged()");
+            fail("Unexpected call " + eventToString(event));
         }
         // Input is not finished.
         assertWithMessage("Input connection is still started")
@@ -1150,7 +1193,7 @@ public class InputMethodServiceTest {
      */
     private void verifyFullscreenMode(@NonNull Runnable runnable, boolean expected,
             boolean orientationPortrait) {
-        verifyInputViewStatus(runnable, expected, false /* inputViewStarted */);
+        verifyInputViewStatus(runnable, EVENT_CONFIG, expected, false /* inputViewStarted */);
         if (expected) {
             // Wait for the TestActivity to be recreated.
             eventually(() -> assertWithMessage("Activity was re-created after rotation")
@@ -1165,6 +1208,7 @@ public class InputMethodServiceTest {
 
         verifyInputViewStatusOnMainSync(
                 () -> mActivity.showImeWithWindowInsetsController(),
+                EVENT_SHOW,
                 true /* expected */,
                 true /* inputViewStarted */);
         assertWithMessage("IME is shown").that(mInputMethodService.isInputViewShown()).isTrue();
@@ -1190,6 +1234,7 @@ public class InputMethodServiceTest {
         // Hide IME before finishing the run.
         verifyInputViewStatusOnMainSync(
                 () -> mActivity.hideImeWithWindowInsetsController(),
+                EVENT_HIDE,
                 true /* expected */,
                 false /* inputViewStarted */);
 
diff --git a/services/tests/InputMethodSystemServerTests/test-apps/SimpleTestIme/Android.bp b/services/tests/InputMethodSystemServerTests/test-apps/SimpleTestIme/Android.bp
index 8a12dcd0add4..e2362f7fa24f 100644
--- a/services/tests/InputMethodSystemServerTests/test-apps/SimpleTestIme/Android.bp
+++ b/services/tests/InputMethodSystemServerTests/test-apps/SimpleTestIme/Android.bp
@@ -53,6 +53,9 @@ android_library {
     srcs: [
         "src/com/android/apps/inputmethod/simpleime/ims/*.java",
     ],
+    static_libs: [
+        "androidx.annotation_annotation",
+    ],
     sdk_version: "current",
 }
 
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 be59dd2cb7a3..3a7abbb167ec 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
@@ -21,6 +21,12 @@ import android.inputmethodservice.InputMethodService;
 import android.util.Log;
 import android.view.inputmethod.EditorInfo;
 
+import androidx.annotation.IntDef;
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
 import java.util.concurrent.CountDownLatch;
 
 /** Wrapper of {@link InputMethodService} to expose interfaces for testing purpose. */
@@ -29,16 +35,41 @@ public class InputMethodServiceWrapper extends InputMethodService {
     private static final String TAG = "InputMethodServiceWrapper";
 
     /** Last created instance of this wrapper. */
+    @Nullable
     private static InputMethodServiceWrapper sInstance;
 
+    /** IME show event ({@link #onStartInputView}). */
+    public static final int EVENT_SHOW = 0;
+
+    /** IME hide event ({@link #onFinishInputView}). */
+    public static final int EVENT_HIDE = 1;
+
+    /** IME configuration change event ({@link #onConfigurationChanged}). */
+    public static final int EVENT_CONFIG = 2;
+
+    /** The type of event that can be waited with a latch. */
+    @IntDef(value = {
+            EVENT_SHOW,
+            EVENT_HIDE,
+            EVENT_CONFIG,
+    })
+    @Retention(RetentionPolicy.SOURCE)
+    public @interface Event {}
+
+    /** The IME event type that the current latch, if any, waits on. */
+    @Event
+    private int mLatchEvent;
+
     private boolean mInputViewStarted;
 
     /**
      * @see #setCountDownLatchForTesting
      */
-    private CountDownLatch mCountDownLatchForTesting;
+    @Nullable
+    private CountDownLatch mCountDownLatch;
 
     /** Gets the last created instance of this wrapper, if available. */
+    @Nullable
     public static InputMethodServiceWrapper getInstance() {
         return sInstance;
     }
@@ -48,14 +79,14 @@ public class InputMethodServiceWrapper extends InputMethodService {
     }
 
     /**
-     * Sets the latch used to wait for the IME to start showing ({@link #onStartInputView},
-     * start hiding ({@link #onFinishInputView}) or receive a configuration change
-     * ({@link #onConfigurationChanged}).
+     * Sets the latch used to wait for the IME event.
      *
-     * @param countDownLatchForTesting the latch to wait on.
+     * @param latch      the latch to wait on.
+     * @param latchEvent the event to set the latch on.
      */
-    public void setCountDownLatchForTesting(CountDownLatch countDownLatchForTesting) {
-        mCountDownLatchForTesting = countDownLatchForTesting;
+    public void setCountDownLatchForTesting(@Nullable CountDownLatch latch, @Event int latchEvent) {
+        mCountDownLatch = latch;
+        mLatchEvent = latchEvent;
     }
 
     @Override
@@ -77,8 +108,8 @@ public class InputMethodServiceWrapper extends InputMethodService {
                 + ", restarting=" + restarting);
         super.onStartInputView(info, restarting);
         mInputViewStarted = true;
-        if (mCountDownLatchForTesting != null) {
-            mCountDownLatchForTesting.countDown();
+        if (mCountDownLatch != null && mLatchEvent == EVENT_SHOW) {
+            mCountDownLatch.countDown();
         }
     }
 
@@ -94,8 +125,8 @@ public class InputMethodServiceWrapper extends InputMethodService {
         super.onFinishInputView(finishingInput);
         mInputViewStarted = false;
 
-        if (mCountDownLatchForTesting != null) {
-            mCountDownLatchForTesting.countDown();
+        if (mCountDownLatch != null && mLatchEvent == EVENT_HIDE) {
+            mCountDownLatch.countDown();
         }
     }
 
@@ -110,11 +141,27 @@ public class InputMethodServiceWrapper extends InputMethodService {
         Log.i(TAG, "onConfigurationChanged() " + newConfig);
         super.onConfigurationChanged(newConfig);
 
-        if (mCountDownLatchForTesting != null) {
-            mCountDownLatchForTesting.countDown();
+        if (mCountDownLatch != null && mLatchEvent == EVENT_CONFIG) {
+            mCountDownLatch.countDown();
         }
     }
 
+    /**
+     * Gets the string representation of the IME event that is being waited on.
+     *
+     * @param event the IME event.
+     */
+    @NonNull
+    public static String eventToString(@Event int event) {
+        return switch (event) {
+            case EVENT_SHOW -> "onStartInputView";
+            case EVENT_HIDE -> "onFinishInputView";
+            case EVENT_CONFIG -> "onConfigurationChanged";
+            default -> "unknownEvent";
+        };
+    }
+
+    @NonNull
     private String dumpEditorInfo(EditorInfo info) {
         if (info == null) {
             return "null";