diff --git a/tests/inputmethod/mockime/src/com/android/cts/mockime/ImeSettings.java b/tests/inputmethod/mockime/src/com/android/cts/mockime/ImeSettings.java
index be709b4..cfc30ec 100644
--- a/tests/inputmethod/mockime/src/com/android/cts/mockime/ImeSettings.java
+++ b/tests/inputmethod/mockime/src/com/android/cts/mockime/ImeSettings.java
@@ -61,6 +61,11 @@
     private static final String STRICT_MODE_ENABLED = "StrictModeEnabled";
     private static final String VERIFY_CONTEXT_APIS_IN_ON_CREATE = "VerifyContextApisInOnCreate";
 
+    /**
+     * Simulate the manifest flag enableOnBackInvokedCallback being true for the IME.
+     */
+    private static final String ON_BACK_CALLBACK_ENABLED = "onBackCallbackEnabled";
+
     @NonNull
     private final PersistableBundle mBundle;
 
@@ -182,6 +187,10 @@
         return mBundle.getBoolean(VERIFY_CONTEXT_APIS_IN_ON_CREATE, false);
     }
 
+    public boolean isOnBackCallbackEnabled() {
+        return mBundle.getBoolean(ON_BACK_CALLBACK_ENABLED, false);
+    }
+
     static Bundle serializeToBundle(@NonNull String eventCallbackActionName,
             @Nullable Builder builder) {
         final Bundle result = new Bundle();
@@ -363,5 +372,15 @@
             mBundle.putBoolean(VERIFY_CONTEXT_APIS_IN_ON_CREATE, enabled);
             return this;
         }
+
+        /**
+         * Sets whether the IME's
+         * {@link android.content.pm.ApplicationInfo#isOnBackInvokedCallbackEnabled()}
+         * should be set to {@code true}.
+         */
+        public Builder setOnBackCallbackEnabled(boolean enabled) {
+            mBundle.putBoolean(ON_BACK_CALLBACK_ENABLED, enabled);
+            return this;
+        }
     }
 }
diff --git a/tests/inputmethod/mockime/src/com/android/cts/mockime/MockIme.java b/tests/inputmethod/mockime/src/com/android/cts/mockime/MockIme.java
index fc4c05e..528da58 100644
--- a/tests/inputmethod/mockime/src/com/android/cts/mockime/MockIme.java
+++ b/tests/inputmethod/mockime/src/com/android/cts/mockime/MockIme.java
@@ -426,10 +426,6 @@
                             return e;
                         }
                     }
-                    case "setEnableOnBackInvokedCallback":
-                        boolean isEnabled = command.getExtras().getBoolean("isEnabled");
-                        getApplicationInfo().setEnableOnBackInvokedCallback(isEnabled);
-                        return ImeEvent.RETURN_VALUE_UNAVAILABLE;
                     case "getDisplayId":
                         return getDisplay().getDisplayId();
                     case "verifyLayoutInflaterContext":
@@ -679,6 +675,10 @@
                             .build());
         }
 
+        if (mSettings.isOnBackCallbackEnabled()) {
+            getApplicationInfo().setEnableOnBackInvokedCallback(true);
+        }
+
         getTracer().onCreate(() -> {
             super.onCreate();
             mHandlerThread.start();
diff --git a/tests/inputmethod/src/android/view/inputmethod/cts/KeyboardVisibilityControlTest.java b/tests/inputmethod/src/android/view/inputmethod/cts/KeyboardVisibilityControlTest.java
index ac95dc8..d79dcca 100644
--- a/tests/inputmethod/src/android/view/inputmethod/cts/KeyboardVisibilityControlTest.java
+++ b/tests/inputmethod/src/android/view/inputmethod/cts/KeyboardVisibilityControlTest.java
@@ -51,6 +51,7 @@
 import android.app.AlertDialog;
 import android.app.Instrumentation;
 import android.app.compat.CompatChanges;
+import android.content.Context;
 import android.content.pm.PackageManager;
 import android.graphics.Color;
 import android.os.SystemClock;
@@ -245,27 +246,38 @@
     }
 
     private void verifyHideImeBackPressed(
-            boolean appRequestsLegacy, boolean imeRequestsLegacy) throws Exception {
+            boolean appRequestsBackCallback, boolean imeRequestsBackCallback) throws Exception {
         final Instrumentation instrumentation = InstrumentationRegistry.getInstrumentation();
-        final InputMethodManager imm = InstrumentationRegistry.getInstrumentation()
-                .getTargetContext().getSystemService(InputMethodManager.class);
+        final Context context = instrumentation.getTargetContext();
+        final InputMethodManager imm = context.getSystemService(InputMethodManager.class);
+
+        // Whether 'OnBackInvokedCallback' or 'onBackPressed' (legacy back) is used is defined by
+        // the 'enableOnBackInvokedCallback' flag in the Application manifest.
+        // Registering a callback is only authorized if the flag is set to true. Since the
+        // WindowOnBackDispatcher is created at the same time as the ViewRootImpl, for test purpose,
+        // we need to manually set the flag on ApplicationInfo before the window is created which
+        // happens during the MockIme creation and TestActivity creation.
 
         try (MockImeSession imeSession = MockImeSession.create(
                 instrumentation.getContext(),
                 instrumentation.getUiAutomation(),
-                new ImeSettings.Builder())) {
+                new ImeSettings.Builder()
+                        .setOnBackCallbackEnabled(imeRequestsBackCallback)
+        )) {
             final ImeEventStream stream = imeSession.openEventStream();
 
             final String marker = getTestMarker();
+
+            if (appRequestsBackCallback) {
+                context.getApplicationInfo().setEnableOnBackInvokedCallback(true);
+            }
+
             final EditText editText = launchTestActivity(marker);
             final TestActivity testActivity = (TestActivity) editText.getContext();
-            if (appRequestsLegacy) {
-                testActivity.getApplicationInfo().setEnableOnBackInvokedCallback(true);
+
+            if (!appRequestsBackCallback) {
                 testActivity.setIgnoreBackKey(true);
             }
-            if (imeRequestsLegacy) {
-                imeSession.callSetEnableOnBackInvokedCallback(true);
-            }
 
             expectEvent(stream, editorMatcher("onStartInput", marker), TIMEOUT);
             notExpectEvent(stream, editorMatcher("onStartInputView", marker), TIMEOUT);
@@ -296,22 +308,26 @@
 
     @Test
     public void testHideImeAfterBackPressed_legacyAppLegacyIme() throws Exception {
-        verifyHideImeBackPressed(true /* appRequestsLegacy */, true /* imeRequestsLegacy */);
+        verifyHideImeBackPressed(false/* appRequestsBackCallback */,
+                false/* imeRequestsBackCallback */);
     }
 
     @Test
     public void testHideImeAfterBackPressed_migratedAppLegacyIme() throws Exception {
-        verifyHideImeBackPressed(false /* appRequestsLegacy */, true /* imeRequestsLegacy */);
+        verifyHideImeBackPressed(true/* appRequestsBackCallback */,
+                false/* imeRequestsBackCallback */);
     }
 
     @Test
     public void testHideImeAfterBackPressed_migratedAppMigratedIme() throws Exception {
-        verifyHideImeBackPressed(false /* appRequestsLegacy */, false /* imeRequestsLegacy */);
+        verifyHideImeBackPressed(true/* appRequestsBackCallback */,
+                true/* imeRequestsBackCallback */);
     }
 
     @Test
     public void testHideImeAfterBackPressed_legacyAppMigratedIme() throws Exception {
-        verifyHideImeBackPressed(true /* appRequestsLegacy */, false /* imeRequestsLegacy */);
+        verifyHideImeBackPressed(false/* appRequestsBackCallback */,
+                true/* imeRequestsBackCallback */);
     }
 
     @Test