Associate ResetGestureInputConsumer and NO_OP input consumers with display IDs

All input consumers should be associated with display IDs, especially since these are used by TouchInteractionService.onConsumerInactive and they will be using per-display objects.

Flag: com.android.launcher3.enable_gesture_nav_on_connected_displays
Bug: 382130680
Test: InputConsumerUtilsTest
Change-Id: Ic14121db2361da1f0a819221b85256b1b3926774
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index 0fff0aa..0579314 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -1168,7 +1168,7 @@
      * @return The {@link TaskbarActivityContext} for the specified display, or
      * {@code null} if no taskbar is associated with that display.
      */
-    private TaskbarActivityContext getTaskbarForDisplay(int displayId) {
+    public TaskbarActivityContext getTaskbarForDisplay(int displayId) {
         return mTaskbars.get(displayId);
     }
 
diff --git a/quickstep/src/com/android/quickstep/InputConsumer.java b/quickstep/src/com/android/quickstep/InputConsumer.java
index 081ed9d..e5cbc66 100644
--- a/quickstep/src/com/android/quickstep/InputConsumer.java
+++ b/quickstep/src/com/android/quickstep/InputConsumer.java
@@ -63,13 +63,25 @@
             "TYPE_BUBBLE_BAR",              // 16
     };
 
-    InputConsumer NO_OP = () -> TYPE_NO_OP;
+    InputConsumer DEFAULT_NO_OP = createNoOpInputConsumer(Display.DEFAULT_DISPLAY);
+
+    static InputConsumer createNoOpInputConsumer(int displayId) {
+        return new InputConsumer() {
+            @Override
+            public int getType() {
+                return TYPE_NO_OP;
+            }
+
+            @Override
+            public int getDisplayId() {
+                return displayId;
+            }
+        };
+    }
 
     int getType();
 
-    default int getDisplayId() {
-        return Display.DEFAULT_DISPLAY;
-    }
+    int getDisplayId();
 
     /**
      * Returns true if the user has crossed the threshold for it to be an explicit action.
diff --git a/quickstep/src/com/android/quickstep/InputConsumerUtils.kt b/quickstep/src/com/android/quickstep/InputConsumerUtils.kt
index cd3ac12..c8ca534 100644
--- a/quickstep/src/com/android/quickstep/InputConsumerUtils.kt
+++ b/quickstep/src/com/android/quickstep/InputConsumerUtils.kt
@@ -57,7 +57,7 @@
     @JvmStatic
     fun <S : BaseState<S>, T> newConsumer(
         context: Context,
-        resetGestureInputConsumer: ResetGestureInputConsumer?,
+        userUnlocked: Boolean,
         overviewComponentObserver: OverviewComponentObserver,
         deviceState: RecentsAnimationDeviceState,
         previousGestureState: GestureState,
@@ -122,7 +122,8 @@
                     // camera).
                     createDeviceLockedInputConsumer(
                         context,
-                        resetGestureInputConsumer,
+                        userUnlocked,
+                        taskbarManager,
                         deviceState,
                         gestureState,
                         taskAnimationManager,
@@ -131,7 +132,10 @@
                     )
                 } else {
                     getDefaultInputConsumer(
-                        resetGestureInputConsumer,
+                        gestureState.displayId,
+                        userUnlocked,
+                        taskAnimationManager,
+                        taskbarManager,
                         reasonString.append("%scannot start system gesture", SUBSTRING_PREFIX),
                     )
                 }
@@ -153,7 +157,8 @@
             base =
                 newBaseConsumer<S, T>(
                     context,
-                    resetGestureInputConsumer,
+                    userUnlocked,
+                    taskbarManager,
                     overviewComponentObserver,
                     deviceState,
                     previousGestureState,
@@ -172,7 +177,14 @@
                     "cannot start system gesture and recents " +
                         "animation was not running, trying to use default input consumer"
                 )
-            base = getDefaultInputConsumer(resetGestureInputConsumer, reasonString)
+            base =
+                getDefaultInputConsumer(
+                    gestureState.displayId,
+                    userUnlocked,
+                    taskAnimationManager,
+                    taskbarManager,
+                    reasonString,
+                )
         }
         if (deviceState.isGesturalNavMode || gestureState.isTrackpadGesture) {
             handleOrientationSetup(base)
@@ -237,7 +249,14 @@
                                 SUBSTRING_PREFIX,
                             )
                     // Bubbles can handle home gesture itself.
-                    base = getDefaultInputConsumer(resetGestureInputConsumer, reasonString)
+                    base =
+                        getDefaultInputConsumer(
+                            gestureState.displayId,
+                            userUnlocked,
+                            taskAnimationManager,
+                            taskbarManager,
+                            reasonString,
+                        )
                 }
             }
 
@@ -279,7 +298,14 @@
                                 SUBSTRING_PREFIX,
                             )
                     // Bubbles can handle home gesture itself.
-                    base = getDefaultInputConsumer(resetGestureInputConsumer, reasonString)
+                    base =
+                        getDefaultInputConsumer(
+                            gestureState.displayId,
+                            userUnlocked,
+                            taskAnimationManager,
+                            taskbarManager,
+                            reasonString,
+                        )
                 }
             }
 
@@ -374,7 +400,14 @@
                             "%sscreen pinning is active, trying to use default input consumer",
                             SUBSTRING_PREFIX,
                         )
-                base = getDefaultInputConsumer(resetGestureInputConsumer, reasonString)
+                base =
+                    getDefaultInputConsumer(
+                        gestureState.displayId,
+                        userUnlocked,
+                        taskAnimationManager,
+                        taskbarManager,
+                        reasonString,
+                    )
             }
 
             if (deviceState.canTriggerOneHandedAction(event)) {
@@ -411,7 +444,7 @@
             context,
             deviceState,
             inputMonitorCompat,
-            InputConsumer.NO_OP,
+            InputConsumer.createNoOpInputConsumer(gestureState.displayId),
             gestureState,
             motionEvent,
             CompoundString.NO_OP,
@@ -450,7 +483,8 @@
     @JvmStatic
     fun <S : BaseState<S>, T> newBaseConsumer(
         context: Context,
-        resetGestureInputConsumer: ResetGestureInputConsumer?,
+        userUnlocked: Boolean,
+        taskbarManager: TaskbarManager,
         overviewComponentObserver: OverviewComponentObserver,
         deviceState: RecentsAnimationDeviceState,
         previousGestureState: GestureState,
@@ -467,7 +501,8 @@
             // This handles apps showing over the lockscreen (e.g. camera)
             return createDeviceLockedInputConsumer(
                 context,
-                resetGestureInputConsumer,
+                userUnlocked,
+                taskbarManager,
                 deviceState,
                 gestureState,
                 taskAnimationManager,
@@ -521,7 +556,9 @@
 
         return if (gestureState.getContainerInterface<S, T>().isInLiveTileMode()) {
             createOverviewInputConsumer<S, T>(
-                resetGestureInputConsumer,
+                userUnlocked,
+                taskAnimationManager,
+                taskbarManager,
                 deviceState,
                 inputMonitorCompat,
                 previousGestureState,
@@ -534,7 +571,10 @@
             )
         } else if (runningTask == null) {
             getDefaultInputConsumer(
-                resetGestureInputConsumer,
+                gestureState.displayId,
+                userUnlocked,
+                taskAnimationManager,
+                taskbarManager,
                 reasonString.append("%srunning task == null", SUBSTRING_PREFIX),
             )
         } else if (
@@ -543,7 +583,9 @@
                 forceOverviewInputConsumer
         ) {
             createOverviewInputConsumer<S, T>(
-                resetGestureInputConsumer,
+                userUnlocked,
+                taskAnimationManager,
+                taskbarManager,
                 deviceState,
                 inputMonitorCompat,
                 previousGestureState,
@@ -565,7 +607,10 @@
             )
         } else if (deviceState.isGestureBlockedTask(runningTask) || launcherChildActivityResumed) {
             getDefaultInputConsumer(
-                resetGestureInputConsumer,
+                gestureState.displayId,
+                userUnlocked,
+                taskAnimationManager,
+                taskbarManager,
                 reasonString.append(
                     if (launcherChildActivityResumed)
                         "%sis launcher child-task, trying to use default input consumer"
@@ -592,7 +637,8 @@
 
     private fun createDeviceLockedInputConsumer(
         context: Context,
-        resetGestureInputConsumer: ResetGestureInputConsumer?,
+        userUnlocked: Boolean,
+        taskbarManager: TaskbarManager,
         deviceState: RecentsAnimationDeviceState,
         gestureState: GestureState,
         taskAnimationManager: TaskAnimationManager,
@@ -617,7 +663,10 @@
             )
         } else {
             getDefaultInputConsumer(
-                resetGestureInputConsumer,
+                gestureState.displayId,
+                userUnlocked,
+                taskAnimationManager,
+                taskbarManager,
                 reasonString.append(
                     if (deviceState.isFullyGesturalNavMode || gestureState.isTrackpadGesture)
                         "%srunning task == null, trying to use default input consumer"
@@ -631,7 +680,9 @@
     }
 
     private fun <S : BaseState<S>, T> createOverviewInputConsumer(
-        resetGestureInputConsumer: ResetGestureInputConsumer?,
+        userUnlocked: Boolean,
+        taskAnimationManager: TaskAnimationManager,
+        taskbarManager: TaskbarManager,
         deviceState: RecentsAnimationDeviceState,
         inputMonitorCompat: InputMonitorCompat,
         previousGestureState: GestureState,
@@ -642,7 +693,10 @@
         val container: T =
             gestureState.getContainerInterface<S, T>().getCreatedContainer()
                 ?: return getDefaultInputConsumer(
-                    resetGestureInputConsumer,
+                    gestureState.displayId,
+                    userUnlocked,
+                    taskAnimationManager,
+                    taskbarManager,
                     reasonString.append(
                         "%sactivity == null, trying to use default input consumer",
                         SUBSTRING_PREFIX,
@@ -694,24 +748,34 @@
     }
 
     /** Returns the [ResetGestureInputConsumer] if user is unlocked, else NO_OP. */
-    private fun getDefaultInputConsumer(
-        resetGestureInputConsumer: ResetGestureInputConsumer?,
+    @JvmStatic
+    fun getDefaultInputConsumer(
+        displayId: Int,
+        userUnlocked: Boolean,
+        taskAnimationManager: TaskAnimationManager?,
+        taskbarManager: TaskbarManager?,
         reasonString: CompoundString,
     ): InputConsumer {
-        return if (resetGestureInputConsumer != null) {
+        return if (userUnlocked && taskAnimationManager != null && taskbarManager != null) {
             reasonString.append(
-                "%smResetGestureInputConsumer initialized, using ResetGestureInputConsumer",
+                "%sResetGestureInputConsumer available, using ResetGestureInputConsumer",
                 SUBSTRING_PREFIX,
             )
-            resetGestureInputConsumer
+            ResetGestureInputConsumer(displayId, taskAnimationManager) {
+                taskbarManager.getTaskbarForDisplay(displayId)
+            }
         } else {
             reasonString.append(
-                "%smResetGestureInputConsumer not initialized, using no-op input consumer",
+                "%s${
+                    if (userUnlocked) "user is locked"
+                    else if (taskAnimationManager == null) "taskAnimationManager is null"
+                    else "taskbarManager is null"
+                }, using no-op input consumer",
                 SUBSTRING_PREFIX,
             )
-            // mResetGestureInputConsumer isn't initialized until onUserUnlocked(), so reset to
+            // ResetGestureInputConsumer isn't available until onUserUnlocked(), so reset to
             // NO_OP until then (we never want these to be null).
-            InputConsumer.NO_OP
+            InputConsumer.createNoOpInputConsumer(displayId)
         }
     }
 
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 741ae7d..2bc5626 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -35,6 +35,7 @@
 import static com.android.quickstep.GestureState.DEFAULT_STATE;
 import static com.android.quickstep.GestureState.TrackpadGestureType.getTrackpadGestureType;
 import static com.android.quickstep.InputConsumer.TYPE_CURSOR_HOVER;
+import static com.android.quickstep.InputConsumer.createNoOpInputConsumer;
 import static com.android.quickstep.InputConsumerUtils.newConsumer;
 import static com.android.quickstep.InputConsumerUtils.tryCreateAssistantInputConsumer;
 import static com.android.systemui.shared.system.ActivityManagerWrapper.CLOSE_SYSTEM_WINDOWS_REASON_RECENTS;
@@ -98,7 +99,6 @@
 import com.android.quickstep.fallback.window.RecentsWindowSwipeHandler;
 import com.android.quickstep.inputconsumers.BubbleBarInputConsumer;
 import com.android.quickstep.inputconsumers.OneHandedModeInputConsumer;
-import com.android.quickstep.inputconsumers.ResetGestureInputConsumer;
 import com.android.quickstep.util.ActiveGestureLog;
 import com.android.quickstep.util.ActiveGestureLog.CompoundString;
 import com.android.quickstep.util.ActiveGestureProtoLogProxy;
@@ -559,10 +559,10 @@
     private RecentsAnimationDeviceState mDeviceState;
     private TaskAnimationManager mTaskAnimationManager;
 
-    private @NonNull InputConsumer mUncheckedConsumer = InputConsumer.NO_OP;
-    private @NonNull InputConsumer mConsumer = InputConsumer.NO_OP;
+    private @NonNull InputConsumer mUncheckedConsumer = InputConsumer.DEFAULT_NO_OP;
+    private @NonNull InputConsumer mConsumer = InputConsumer.DEFAULT_NO_OP;
     private Choreographer mMainChoreographer;
-    private @Nullable ResetGestureInputConsumer mResetGestureInputConsumer;
+    private boolean mUserUnlocked = false;
     private GestureState mGestureState = DEFAULT_STATE;
 
     private InputMonitorDisplayModel mInputMonitorDisplayModel;
@@ -692,8 +692,7 @@
         mOverviewCommandHelper = new OverviewCommandHelper(this,
                 mOverviewComponentObserver, mTaskAnimationManager, mRecentsDisplayModel,
                 SystemUiProxy.INSTANCE.get(this).getFocusState(), mTaskbarManager);
-        mResetGestureInputConsumer = new ResetGestureInputConsumer(
-                mTaskAnimationManager, mTaskbarManager::getCurrentActivityContext);
+        mUserUnlocked = true;
         mInputConsumer.registerInputConsumer();
         for (int displayId : mDeviceState.getDisplaysWithSysUIState()) {
             onSystemUiFlagsChanged(mDeviceState.getSystemUiStateFlags(displayId), displayId);
@@ -917,7 +916,7 @@
                 } else {
                     reasonString.append(" but event cannot trigger Assistant, "
                             + "consuming gesture as no-op");
-                    mUncheckedConsumer = InputConsumer.NO_OP;
+                    mUncheckedConsumer = createNoOpInputConsumer(displayId);
                 }
             } else if ((!isOneHandedModeActive && isInSwipeUpTouchRegion)
                     || isHoverActionWithoutConsumer || isOnBubbles) {
@@ -934,7 +933,7 @@
                 mGestureState = newGestureState;
                 mConsumer = newConsumer(
                         this,
-                        mResetGestureInputConsumer,
+                        mUserUnlocked,
                         mOverviewComponentObserver,
                         mDeviceState,
                         prevGestureState,
@@ -968,19 +967,22 @@
                         + "consuming gesture for one-handed action");
                 // Consume gesture event for triggering one handed feature.
                 mUncheckedConsumer = new OneHandedModeInputConsumer(
-                        this, displayId, mDeviceState, InputConsumer.NO_OP, inputMonitorCompat);
+                        this,
+                        displayId,
+                        mDeviceState,
+                        InputConsumer.createNoOpInputConsumer(displayId), inputMonitorCompat);
             } else {
-                mUncheckedConsumer = InputConsumer.NO_OP;
+                mUncheckedConsumer = InputConsumer.createNoOpInputConsumer(displayId);
             }
         } else {
             // Other events
-            if (mUncheckedConsumer != InputConsumer.NO_OP) {
+            if (mUncheckedConsumer.getType() != InputConsumer.TYPE_NO_OP) {
                 // Only transform the event if we are handling it in a proper consumer
                 mRotationTouchHelper.setOrientationTransformIfNeeded(event);
             }
         }
 
-        if (mUncheckedConsumer != InputConsumer.NO_OP) {
+        if (mUncheckedConsumer.getType() != InputConsumer.TYPE_NO_OP) {
             switch (action) {
                 case ACTION_DOWN:
                     ActiveGestureProtoLogProxy.logOnInputEventActionDown(displayId, reasonString);
@@ -1100,7 +1102,12 @@
     }
 
     private void reset(int displayId) {
-        mConsumer = mUncheckedConsumer = getDefaultInputConsumer();
+        mConsumer = mUncheckedConsumer = InputConsumerUtils.getDefaultInputConsumer(
+                displayId,
+                mUserUnlocked,
+                mTaskAnimationManager,
+                mTaskbarManager,
+                CompoundString.NO_OP);
         mGestureState = DEFAULT_STATE;
         // By default, use batching of the input events, but check receiver before using in the rare
         // case that the monitor was disposed before the swipe settled
@@ -1110,29 +1117,6 @@
         }
     }
 
-    private @NonNull InputConsumer getDefaultInputConsumer() {
-        return getDefaultInputConsumer(CompoundString.NO_OP);
-    }
-
-    /**
-     * Returns the {@link ResetGestureInputConsumer} if user is unlocked, else NO_OP.
-     */
-    private @NonNull InputConsumer getDefaultInputConsumer(@NonNull CompoundString reasonString) {
-        if (mResetGestureInputConsumer != null) {
-            reasonString.append(
-                    "%smResetGestureInputConsumer initialized, using ResetGestureInputConsumer",
-                    SUBSTRING_PREFIX);
-            return mResetGestureInputConsumer;
-        } else {
-            reasonString.append(
-                    "%smResetGestureInputConsumer not initialized, using no-op input consumer",
-                    SUBSTRING_PREFIX);
-            // mResetGestureInputConsumer isn't initialized until onUserUnlocked(), so reset to
-            // NO_OP until then (we never want these to be null).
-            return InputConsumer.NO_OP;
-        }
-    }
-
     @Override
     public void onConfigurationChanged(Configuration newConfig) {
         if (!LockedUserState.get(this).isUserUnlocked()) {
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/ResetGestureInputConsumer.java b/quickstep/src/com/android/quickstep/inputconsumers/ResetGestureInputConsumer.java
deleted file mode 100644
index 349f4d2..0000000
--- a/quickstep/src/com/android/quickstep/inputconsumers/ResetGestureInputConsumer.java
+++ /dev/null
@@ -1,55 +0,0 @@
-/*
- * Copyright (C) 2019 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.quickstep.inputconsumers;
-
-import android.view.MotionEvent;
-
-import com.android.launcher3.taskbar.TaskbarActivityContext;
-import com.android.quickstep.InputConsumer;
-import com.android.quickstep.TaskAnimationManager;
-
-import java.util.function.Supplier;
-
-/**
- * A NO_OP input consumer which also resets any pending gesture
- */
-public class ResetGestureInputConsumer implements InputConsumer {
-
-    private final TaskAnimationManager mTaskAnimationManager;
-    private final Supplier<TaskbarActivityContext> mActivityContextSupplier;
-
-    public ResetGestureInputConsumer(
-            TaskAnimationManager taskAnimationManager,
-            Supplier<TaskbarActivityContext> activityContextSupplier) {
-        mTaskAnimationManager = taskAnimationManager;
-        mActivityContextSupplier = activityContextSupplier;
-    }
-
-    @Override
-    public int getType() {
-        return TYPE_RESET_GESTURE;
-    }
-
-    @Override
-    public void onMotionEvent(MotionEvent ev) {
-        if (ev.getAction() == MotionEvent.ACTION_DOWN
-                && mTaskAnimationManager.isRecentsAnimationRunning()) {
-            TaskbarActivityContext tac = mActivityContextSupplier.get();
-            mTaskAnimationManager.finishRunningRecentsAnimation(
-                    /* toHome= */ tac != null && !tac.isInApp());
-        }
-    }
-}
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/ResetGestureInputConsumer.kt b/quickstep/src/com/android/quickstep/inputconsumers/ResetGestureInputConsumer.kt
new file mode 100644
index 0000000..96e7943
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/inputconsumers/ResetGestureInputConsumer.kt
@@ -0,0 +1,44 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.quickstep.inputconsumers
+
+import android.view.MotionEvent
+import com.android.launcher3.taskbar.TaskbarActivityContext
+import com.android.quickstep.InputConsumer
+import com.android.quickstep.TaskAnimationManager
+import java.util.function.Supplier
+
+/** A NO_OP input consumer which also resets any pending gesture */
+class ResetGestureInputConsumer(
+    private val displayId: Int,
+    private val taskAnimationManager: TaskAnimationManager,
+    private val activityContextSupplier: Supplier<TaskbarActivityContext?>,
+) : InputConsumer {
+    override fun getType() = InputConsumer.TYPE_RESET_GESTURE
+
+    override fun getDisplayId() = displayId
+
+    override fun onMotionEvent(ev: MotionEvent) {
+        if (
+            ev.action == MotionEvent.ACTION_DOWN && taskAnimationManager.isRecentsAnimationRunning
+        ) {
+            val tac = activityContextSupplier.get()
+            taskAnimationManager.finishRunningRecentsAnimation(
+                /* toHome= */ tac != null && !tac.isInApp
+            )
+        }
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/util/InputProxyHandlerFactory.java b/quickstep/src/com/android/quickstep/util/InputProxyHandlerFactory.java
index 843619d..9aded89 100644
--- a/quickstep/src/com/android/quickstep/util/InputProxyHandlerFactory.java
+++ b/quickstep/src/com/android/quickstep/util/InputProxyHandlerFactory.java
@@ -35,8 +35,8 @@
     private final GestureState mGestureState;
 
     @UiThread
-    public InputProxyHandlerFactory(BaseContainerInterface activityInterface,
-            GestureState gestureState) {
+    public InputProxyHandlerFactory(
+            BaseContainerInterface activityInterface, GestureState gestureState) {
         mContainerInterface = activityInterface;
         mGestureState = gestureState;
     }
@@ -47,7 +47,8 @@
     @Override
     public InputConsumer get() {
         RecentsViewContainer container = mContainerInterface.getCreatedContainer();
-        return container == null ? InputConsumer.NO_OP
+        return container == null
+                ? InputConsumer.createNoOpInputConsumer(mGestureState.getDisplayId())
                 : new OverviewInputConsumer(mGestureState, container, null, true);
     }
 }
diff --git a/quickstep/tests/src/com/android/quickstep/InputConsumerUtilsTest.java b/quickstep/tests/src/com/android/quickstep/InputConsumerUtilsTest.java
index ef6f55e..3db4398 100644
--- a/quickstep/tests/src/com/android/quickstep/InputConsumerUtilsTest.java
+++ b/quickstep/tests/src/com/android/quickstep/InputConsumerUtilsTest.java
@@ -27,7 +27,6 @@
 import static org.mockito.Mockito.when;
 
 import android.annotation.NonNull;
-import android.annotation.Nullable;
 import android.os.Looper;
 import android.view.Choreographer;
 import android.view.Display;
@@ -100,12 +99,13 @@
 
     @Rule public final SandboxApplication mContext = new SandboxApplication();
 
+    private final int mDisplayId = Display.DEFAULT_DISPLAY;
     @NonNull private final InputMonitorCompat mInputMonitorCompat =
-            new InputMonitorCompat("", Display.DEFAULT_DISPLAY);
+            new InputMonitorCompat("", mDisplayId);
 
     private TaskAnimationManager mTaskAnimationManager;
     private InputChannelCompat.InputEventReceiver mInputEventReceiver;
-    @Nullable private ResetGestureInputConsumer mResetGestureInputConsumer;
+    private boolean mUserUnlocked = true;
     @NonNull private Function<GestureState, AnimatedFloat> mSwipeUpProxyProvider = (state) -> null;
 
     @NonNull @Mock private TaskbarActivityContext mTaskbarActivityContext;
@@ -164,12 +164,6 @@
     }
 
     @Before
-    public void setUpResetGestureInputConsumer() {
-        mResetGestureInputConsumer = new ResetGestureInputConsumer(
-                mTaskAnimationManager, mTaskbarManager::getCurrentActivityContext);
-    }
-
-    @Before
     public void setupLockedUserState() {
         when(mLockedUserState.isUserUnlocked()).thenReturn(true);
     }
@@ -314,8 +308,6 @@
 
     @Test
     public void testNewBaseConsumer_containsOtherActivityInputConsumer() {
-        // OtherActivityInputConsumer needs to be initialized on the main thread because of
-        // MotionPauseDetector.mForcePauseTimeout
         assertCorrectInputConsumer(
                 this::createBaseInputConsumer,
                 OtherActivityInputConsumer.class,
@@ -486,7 +478,7 @@
         MotionEvent event = MotionEvent.obtain(0, 0, MotionEvent.ACTION_DOWN, 0, 0, 0);
         InputConsumer inputConsumer = newConsumer(
                 mContext,
-                mResetGestureInputConsumer,
+                mUserUnlocked,
                 mOverviewComponentObserver,
                 mDeviceState,
                 mPreviousGestureState,
@@ -510,7 +502,8 @@
         MotionEvent event = MotionEvent.obtain(0, 0, MotionEvent.ACTION_DOWN, 0, 0, 0);
         InputConsumer inputConsumer = newBaseConsumer(
                 mContext,
-                mResetGestureInputConsumer,
+                mUserUnlocked,
+                mTaskbarManager,
                 mOverviewComponentObserver,
                 mDeviceState,
                 mPreviousGestureState,
@@ -535,12 +528,15 @@
                 ResetGestureInputConsumer.class,
                 InputConsumer.TYPE_RESET_GESTURE);
 
-        mResetGestureInputConsumer = null;
+        mUserUnlocked = false;
 
-        runOnMainSync(() -> assertThat(inputConsumerProvider.get()).isEqualTo(InputConsumer.NO_OP));
+        assertCorrectInputConsumer(
+                inputConsumerProvider,
+                InputConsumer.class,
+                InputConsumer.TYPE_NO_OP);
     }
 
-    private static void assertCorrectInputConsumer(
+    private void assertCorrectInputConsumer(
             @NonNull Provider<InputConsumer> inputConsumerProvider,
             @NonNull Class<? extends InputConsumer> expectedOutputConsumer,
             int expectedType) {
@@ -551,11 +547,13 @@
                 expectedType);
     }
 
-    private static void assertCorrectInputConsumer(
+    private void assertCorrectInputConsumer(
             @NonNull Provider<InputConsumer> inputConsumerProvider,
             @NonNull Class<? extends InputConsumer> expectedOutputConsumer,
             @NonNull Class<? extends InputConsumer> expectedActiveConsumer,
             int expectedType) {
+        when(mCurrentGestureState.getDisplayId()).thenReturn(mDisplayId);
+
         runOnMainSync(() -> {
             InputConsumer inputConsumer = inputConsumerProvider.get();
 
@@ -563,7 +561,14 @@
             assertThat(inputConsumer.getActiveConsumerInHierarchy())
                     .isInstanceOf(expectedActiveConsumer);
             assertThat(inputConsumer.getType()).isEqualTo(expectedType);
+            assertThat(inputConsumer.getDisplayId()).isEqualTo(mDisplayId);
         });
+        int expectedDisplayId = mDisplayId + 1;
+
+        when(mCurrentGestureState.getDisplayId()).thenReturn(expectedDisplayId);
+
+        runOnMainSync(() -> assertThat(inputConsumerProvider.get().getDisplayId())
+                .isEqualTo(expectedDisplayId));
     }
 
     private static void runOnMainSync(@NonNull Runnable runnable) {