1 files changed, 4 insertions, 2 deletions
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/inputmethodservice/InputMethodServiceTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/inputmethodservice/InputMethodServiceTest.java
index 2d3f7231cc5c..03d8da525f1e 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/inputmethodservice/InputMethodServiceTest.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/inputmethodservice/InputMethodServiceTest.java
@@ -139,10 +139,9 @@ public class InputMethodServiceTest {
         if (!mOriginalVerboseImeTrackerLoggingEnabled) {
             setVerboseImeTrackerLogging(true);
         }
+        mUiDevice.setOrientationNatural();
         prepareIme();
         prepareActivity();
-        mUiDevice.freezeRotation();
-        mUiDevice.setOrientationNatural();
         // Waits for input binding ready.
         eventually(() -> {
             mInputMethodService = InputMethodServiceWrapper.getInstance();
@@ -169,6 +168,9 @@ public class InputMethodServiceTest {
 
     @After
     public void tearDown() throws Exception {
+        if (!mUiDevice.isNaturalOrientation()) {
+            mUiDevice.setOrientationNatural();
+        }
         mUiDevice.unfreezeRotation();
         if (!mOriginalVerboseImeTrackerLoggingEnabled) {
             setVerboseImeTrackerLogging(false);