Merge "Fix NotificationTest on small screen." into udc-dev am: a21f721f6d

Original change: https://googleplex-android-review.googlesource.com/c/platform/frameworks/base/+/21642593

Change-Id: Ifca466fc0aab401e0c4e3f82f0c58be1225f92fe
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/NotificationTest.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/NotificationTest.java
index 573b3b69..d2708ad 100644
--- a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/NotificationTest.java
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/NotificationTest.java
@@ -42,6 +42,7 @@
 import androidx.test.uiautomator.By;
 import androidx.test.uiautomator.BySelector;
 import androidx.test.uiautomator.UiDevice;
+import androidx.test.uiautomator.UiObject2;
 import androidx.test.uiautomator.Until;
 
 import org.junit.After;
@@ -74,7 +75,7 @@
 
     // This is for AOSP System UI for phones. When testing customized System UI, please modify here.
     private static final BySelector REPLY_SEND_BUTTON_SELECTOR =
-            By.res("com.android.systemui", "remote_input_send");
+            By.res("com.android.systemui", "remote_input_send").enabled(true);
 
     @Rule
     public UnlockScreenRule mUnlockScreenRule = new UnlockScreenRule();
@@ -119,7 +120,15 @@
         mUiDevice.pressKeyCode(KeyEvent.KEYCODE_A);
         mUiDevice.pressKeyCode(KeyEvent.KEYCODE_B);
         mUiDevice.pressKeyCode(KeyEvent.KEYCODE_C);
-        mUiDevice.wait(Until.findObject(REPLY_SEND_BUTTON_SELECTOR.enabled(true)), TIMEOUT).click();
+        UiObject2 sendButton = mUiDevice.wait(
+                Until.findObject(REPLY_SEND_BUTTON_SELECTOR), TIMEOUT);
+        if (sendButton == null) {
+            // If the screen is too small, sendButton may be hidden by IME.
+            // Dismiss IME and try again.
+            mUiDevice.pressBack();
+            sendButton = mUiDevice.wait(Until.findObject(REPLY_SEND_BUTTON_SELECTOR), TIMEOUT);
+        }
+        sendButton.click();
         // Verify that IME is gone.
         assertThat(mUiDevice.wait(Until.gone(By.pkg(getImePackage(mContext))), TIMEOUT)).isTrue();
     }