Fix initial iteration in DefaultImeVisibilityTest

Ensure rotation is successfully applied before starting the stress test
iterations, as this would previously lead to the first iteration of the
landscape test still being in portrait mode, and incorrectly passing the
assertions.

Test: atest DefaultImeVisibilityTest
Bug: 291742721
Change-Id: Iae068a92601d6dc632f10373ab6ab4bd41dd742d

2 files changed, 34 insertions, 11 deletions

diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/DefaultImeVisibilityTest.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/DefaultImeVisibilityTest.java
index 03df07acbe37..320daeeb2e54 100644
--- a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/DefaultImeVisibilityTest.java
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/DefaultImeVisibilityTest.java
@@ -19,6 +19,7 @@ package com.android.inputmethod.stresstest;
 import static android.view.WindowManager.LayoutParams.SOFT_INPUT_ADJUST_RESIZE;
 import static android.view.WindowManager.LayoutParams.SOFT_INPUT_STATE_HIDDEN;
 
+import static com.android.compatibility.common.util.SystemUtil.eventually;
 import static com.android.inputmethod.stresstest.ImeStressTestUtil.REQUEST_FOCUS_ON_CREATE;
 import static com.android.inputmethod.stresstest.ImeStressTestUtil.TestActivity.createIntent;
 import static com.android.inputmethod.stresstest.ImeStressTestUtil.callOnMainSync;
@@ -26,11 +27,16 @@ import static com.android.inputmethod.stresstest.ImeStressTestUtil.verifyWindowA
 import static com.android.inputmethod.stresstest.ImeStressTestUtil.waitOnMainUntilImeIsHidden;
 import static com.android.inputmethod.stresstest.ImeStressTestUtil.waitOnMainUntilImeIsShown;
 
+import static com.google.common.truth.Truth.assertWithMessage;
+
 import android.content.Intent;
 import android.platform.test.annotations.RootPermissionTest;
 import android.platform.test.rule.UnlockScreenRule;
+import android.support.test.uiautomator.UiDevice;
 import android.widget.EditText;
 
+import androidx.test.platform.app.InstrumentationRegistry;
+
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -39,6 +45,7 @@ import org.junit.runners.Parameterized;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Test IME visibility by using system default IME to ensure the behavior is consistent
@@ -59,8 +66,12 @@ public final class DefaultImeVisibilityTest {
     public ScreenCaptureRule mScreenCaptureRule =
             new ScreenCaptureRule("/sdcard/InputMethodStressTest");
 
+    private static final long TIMEOUT = TimeUnit.SECONDS.toMillis(3);
+
     private static final int NUM_TEST_ITERATIONS = 10;
 
+    private final boolean mIsPortrait;
+
     @Parameterized.Parameters(name = "isPortrait={0}")
     public static List<Boolean> isPortraitCases() {
         // Test in both portrait and landscape mode.
@@ -68,6 +79,7 @@ public final class DefaultImeVisibilityTest {
     }
 
     public DefaultImeVisibilityTest(boolean isPortrait) {
+        mIsPortrait = isPortrait;
         mImeStressTestRule.setIsPortrait(isPortrait);
     }
 
@@ -80,6 +92,15 @@ public final class DefaultImeVisibilityTest {
                         Collections.singletonList(REQUEST_FOCUS_ON_CREATE));
         ImeStressTestUtil.TestActivity activity = ImeStressTestUtil.TestActivity.start(intent);
         EditText editText = activity.getEditText();
+
+        UiDevice uiDevice = UiDevice.getInstance(InstrumentationRegistry.getInstrumentation());
+        eventually(
+                () ->
+                        assertWithMessage("Display rotation should be updated.")
+                                .that(uiDevice.getDisplayRotation())
+                                .isEqualTo(mIsPortrait ? 0 : 1),
+                TIMEOUT);
+
         for (int i = 0; i < NUM_TEST_ITERATIONS; i++) {
             // TODO(b/291752364): Remove the explicit focus request once the issue with view focus
             //  change between fullscreen IME and actual editText is fixed.
diff --git a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeStressTestRule.java b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeStressTestRule.java
index 12104b298dac..c7463218b646 100644
--- a/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeStressTestRule.java
+++ b/tests/InputMethodStressTest/src/com/android/inputmethod/stresstest/ImeStressTestRule.java
@@ -20,9 +20,9 @@ import android.app.Instrumentation;
 import android.os.RemoteException;
 import android.support.test.uiautomator.UiDevice;
 
+import androidx.annotation.NonNull;
 import androidx.test.platform.app.InstrumentationRegistry;
 
-import org.checkerframework.checker.nullness.qual.NonNull;
 import org.junit.rules.TestWatcher;
 import org.junit.runner.Description;
 
@@ -35,8 +35,6 @@ import java.io.IOException;
 public class ImeStressTestRule extends TestWatcher {
     private static final String LOCK_SCREEN_OFF_COMMAND = "locksettings set-disabled true";
     private static final String LOCK_SCREEN_ON_COMMAND = "locksettings set-disabled false";
-    private static final String SET_PORTRAIT_MODE_COMMAND = "settings put system user_rotation 0";
-    private static final String SET_LANDSCAPE_MODE_COMMAND = "settings put system user_rotation 1";
     private static final String SIMPLE_IME_ID =
             "com.android.apps.inputmethod.simpleime/.SimpleInputMethodService";
     private static final String ENABLE_IME_COMMAND = "ime enable " + SIMPLE_IME_ID;
@@ -44,8 +42,10 @@ public class ImeStressTestRule extends TestWatcher {
     private static final String DISABLE_IME_COMMAND = "ime disable " + SIMPLE_IME_ID;
     private static final String RESET_IME_COMMAND = "ime reset";
 
-    @NonNull private final Instrumentation mInstrumentation;
-    @NonNull private final UiDevice mUiDevice;
+    @NonNull
+    private final Instrumentation mInstrumentation;
+    @NonNull
+    private final UiDevice mUiDevice;
     // Whether the screen orientation is set to portrait.
     private boolean mIsPortrait;
     // Whether to use a simple test Ime or system default Ime for test.
@@ -105,12 +105,13 @@ public class ImeStressTestRule extends TestWatcher {
     private void setOrientation() {
         try {
             mUiDevice.freezeRotation();
-            executeShellCommand(
-                    mIsPortrait ? SET_PORTRAIT_MODE_COMMAND : SET_LANDSCAPE_MODE_COMMAND);
-        } catch (IOException e) {
-            throw new RuntimeException("Could not set screen orientation.", e);
+            if (mIsPortrait) {
+                mUiDevice.setOrientationNatural();
+            } else {
+                mUiDevice.setOrientationLeft();
+            }
         } catch (RemoteException e) {
-            throw new RuntimeException("Could not freeze rotation.", e);
+            throw new RuntimeException("Could not freeze rotation or set screen orientation.", e);
         }
     }
 
@@ -147,7 +148,8 @@ public class ImeStressTestRule extends TestWatcher {
         }
     }
 
-    private @NonNull String executeShellCommand(@NonNull String cmd) throws IOException {
+    @NonNull
+    private String executeShellCommand(@NonNull String cmd) throws IOException {
         return mUiDevice.executeShellCommand(cmd);
     }
 }