Merge changes I6333b08f,I4186bb54 into main

* changes:
  Remove SideStage class
  Migrate individual calls off mMainStage

10 files changed, 91 insertions, 202 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/MainStage.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/MainStage.java
index 1cbb8bbe5f75..1dad4137a9df 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/MainStage.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/MainStage.java
@@ -33,7 +33,7 @@ import java.util.Optional;
 
 /**
  * Main stage for split-screen mode. When split-screen is active all standard activity types launch
- * on the main stage, except for task that are explicitly pinned to the {@link SideStage}.
+ * on the main stage, except for task that are explicitly pinned to the {@link StageTaskListener}.
  * @see StageCoordinator
  */
 class MainStage extends StageTaskListener {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SideStage.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SideStage.java
deleted file mode 100644
index 27fd309c09b0..000000000000
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SideStage.java
+++ /dev/null
@@ -1,73 +0,0 @@
-/*
- * Copyright (C) 2020 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.wm.shell.splitscreen;
-
-import static com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN;
-
-import android.app.ActivityManager;
-import android.content.Context;
-import android.view.SurfaceSession;
-import android.window.WindowContainerToken;
-import android.window.WindowContainerTransaction;
-
-import com.android.internal.protolog.ProtoLog;
-import com.android.launcher3.icons.IconProvider;
-import com.android.wm.shell.ShellTaskOrganizer;
-import com.android.wm.shell.common.SyncTransactionQueue;
-import com.android.wm.shell.windowdecor.WindowDecorViewModel;
-
-import java.util.Optional;
-
-/**
- * Side stage for split-screen mode. Only tasks that are explicitly pinned to this stage show up
- * here. All other task are launch in the {@link MainStage}.
- *
- * @see StageCoordinator
- */
-class SideStage extends StageTaskListener {
-    private static final String TAG = SideStage.class.getSimpleName();
-
-    SideStage(Context context, ShellTaskOrganizer taskOrganizer, int displayId,
-            StageListenerCallbacks callbacks, SyncTransactionQueue syncQueue,
-            SurfaceSession surfaceSession, IconProvider iconProvider,
-            Optional<WindowDecorViewModel> windowDecorViewModel) {
-        super(context, taskOrganizer, displayId, callbacks, syncQueue, surfaceSession,
-                iconProvider, windowDecorViewModel);
-    }
-
-    boolean removeAllTasks(WindowContainerTransaction wct, boolean toTop) {
-        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "remove all side stage tasks: childCount=%d toTop=%b",
-                mChildrenTaskInfo.size(), toTop);
-        if (mChildrenTaskInfo.size() == 0) return false;
-        wct.reparentTasks(
-                mRootTaskInfo.token,
-                null /* newParent */,
-                null /* windowingModes */,
-                null /* activityTypes */,
-                toTop);
-        return true;
-    }
-
-    boolean removeTask(int taskId, WindowContainerToken newParent, WindowContainerTransaction wct) {
-        final ActivityManager.RunningTaskInfo task = mChildrenTaskInfo.get(taskId);
-        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "remove side stage task: task=%d exists=%b", taskId,
-                task != null);
-        if (task == null) return false;
-        wct.reparent(task.token, newParent, false /* onTop */);
-        return true;
-    }
-}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
index a6233dc927a5..526c1d4a179d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
@@ -50,7 +50,7 @@ public interface SplitScreen {
 
     /**
      * The side stage type.
-     * @see SideStage
+     * @see StageTaskListener
      */
     int STAGE_TYPE_SIDE = 1;
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
index 8921ceb6175d..a10c1db20317 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
@@ -155,13 +155,13 @@ import java.util.concurrent.Executor;
 
 /**
  * Coordinates the staging (visibility, sizing, ...) of the split-screen {@link MainStage} and
- * {@link SideStage} stages.
+ * other stages.
  * Some high-level rules:
- * - The {@link StageCoordinator} is only considered active if the {@link SideStage} contains at
+ * - The {@link StageCoordinator} is only considered active if the other stages contain at
  * least one child task.
  * - The {@link MainStage} should only have children if the coordinator is active.
  * - The {@link SplitLayout} divider is only visible if both the {@link MainStage}
- * and {@link SideStage} are visible.
+ * and other stages are visible.
  * - Both stages are put under a single-top root task.
  * This rules are mostly implemented in {@link #onStageVisibilityChanged(StageListenerImpl)} and
  * {@link #onStageHasChildrenChanged(StageListenerImpl).}
@@ -176,7 +176,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
 
     private final MainStage mMainStage;
     private final StageListenerImpl mMainStageListener = new StageListenerImpl();
-    private final SideStage mSideStage;
+    private final StageTaskListener mSideStage;
     private final StageListenerImpl mSideStageListener = new StageListenerImpl();
     @SplitPosition
     private int mSideStagePosition = SPLIT_POSITION_BOTTOM_OR_RIGHT;
@@ -338,7 +338,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
                 mSurfaceSession,
                 iconProvider,
                 mWindowDecorViewModel);
-        mSideStage = new SideStage(
+        mSideStage = new StageTaskListener(
                 mContext,
                 mTaskOrganizer,
                 mDisplayId,
@@ -367,7 +367,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
 
     @VisibleForTesting
     StageCoordinator(Context context, int displayId, SyncTransactionQueue syncQueue,
-            ShellTaskOrganizer taskOrganizer, MainStage mainStage, SideStage sideStage,
+            ShellTaskOrganizer taskOrganizer, MainStage mainStage, StageTaskListener sideStage,
             DisplayController displayController, DisplayImeController displayImeController,
             DisplayInsetsController displayInsetsController, SplitLayout splitLayout,
             Transitions transitions, TransactionPool transactionPool, ShellExecutor mainExecutor,
@@ -424,6 +424,15 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         return mMainStage.isActive();
     }
 
+    /**
+     * Deactivates main stage by removing the stage from the top level split root (usually when a
+     * task underneath gets removed from the stage root).
+     * @param reparentToTop whether we want to put the stage root back on top
+     */
+    private void deactivateSplit(WindowContainerTransaction wct, boolean reparentToTop) {
+        mMainStage.deactivate(wct, reparentToTop);
+    }
+
     /** @return whether this transition-request has the launch-adjacent flag. */
     public boolean requestHasLaunchAdjacentFlag(TransitionRequestInfo request) {
         final ActivityManager.RunningTaskInfo triggerTask = request.getTriggerTask();
@@ -498,10 +507,10 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
 
         /**
          * {@link MainStage} will be deactivated in {@link #onStageHasChildrenChanged} if the
-         * {@link SideStage} no longer has children.
+         * other stages no longer have children.
          */
         final boolean result = mSideStage.removeTask(taskId,
-                mMainStage.isActive() ? mMainStage.mRootTaskInfo.token : null,
+                isSplitActive() ? mMainStage.mRootTaskInfo.token : null,
                 wct);
         mTaskOrganizer.applyTransaction(wct);
         return result;
@@ -618,7 +627,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         }
 
         // If split screen is not activated, we're expecting to open a pair of apps to split.
-        final int extraTransitType = mMainStage.isActive()
+        final int extraTransitType = isSplitActive()
                 ? TRANSIT_SPLIT_SCREEN_OPEN_TO_SIDE : TRANSIT_SPLIT_SCREEN_PAIR_OPEN;
         prepareEnterSplitScreen(wct, null /* taskInfo */, position, !mIsDropEntering);
 
@@ -661,7 +670,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         }
 
         // If split screen is not activated, we're expecting to open a pair of apps to split.
-        final int extraTransitType = mMainStage.isActive()
+        final int extraTransitType = isSplitActive()
                 ? TRANSIT_SPLIT_SCREEN_OPEN_TO_SIDE : TRANSIT_SPLIT_SCREEN_PAIR_OPEN;
         prepareEnterSplitScreen(wct, null /* taskInfo */, position, !mIsDropEntering);
 
@@ -793,7 +802,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
     private void startWithTask(WindowContainerTransaction wct, int mainTaskId,
             @Nullable Bundle mainOptions, @PersistentSnapPosition int snapPosition,
             @Nullable RemoteTransition remoteTransition, InstanceId instanceId) {
-        if (!mMainStage.isActive()) {
+        if (!isSplitActive()) {
             // Build a request WCT that will launch both apps such that task 0 is on the main stage
             // while task 1 is on the side stage.
             mMainStage.activate(wct, false /* reparent */);
@@ -860,7 +869,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             return;
         }
 
-        if (!mMainStage.isActive()) {
+        if (!isSplitActive()) {
             // Build a request WCT that will launch both apps such that task 0 is on the main stage
             // while task 1 is on the side stage.
             mMainStage.activate(wct, false /* reparent */);
@@ -1110,7 +1119,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
      */
     void onKeyguardStateChanged(boolean active, boolean occludingTaskRunning) {
         mKeyguardActive = active;
-        if (!mMainStage.isActive()) {
+        if (!isSplitActive()) {
             return;
         }
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN,
@@ -1154,7 +1163,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
      * will do a no-op.
      */
     void dismissSplitKeepingLastActiveStage(@ExitReason int reason) {
-        if (!mMainStage.isActive() || mLastActiveStage == STAGE_TYPE_UNDEFINED) {
+        if (!isSplitActive() || mLastActiveStage == STAGE_TYPE_UNDEFINED) {
             // no-op
             return;
         }
@@ -1177,8 +1186,8 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
     private void exitSplitScreen(@Nullable StageTaskListener childrenToTop,
             @ExitReason int exitReason) {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "exitSplitScreen: mainStageToTop=%b reason=%s active=%b",
-                childrenToTop == mMainStage, exitReasonToString(exitReason), mMainStage.isActive());
-        if (!mMainStage.isActive()) return;
+                childrenToTop == mMainStage, exitReasonToString(exitReason), isSplitActive());
+        if (!isSplitActive()) return;
 
         final WindowContainerTransaction wct = new WindowContainerTransaction();
         applyExitSplitScreen(childrenToTop, wct, exitReason);
@@ -1188,7 +1197,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             WindowContainerTransaction wct, @ExitReason int exitReason) {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "applyExitSplitScreen: reason=%s",
                 exitReasonToString(exitReason));
-        if (!mMainStage.isActive() || mIsExiting) return;
+        if (!isSplitActive() || mIsExiting) return;
 
         onSplitScreenExit();
         clearSplitPairedInRecents(exitReason);
@@ -1200,7 +1209,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         mSplitLayout.getInvisibleBounds(mTempRect1);
         if (childrenToTop == null || childrenToTop.getTopVisibleChildTaskId() == INVALID_TASK_ID) {
             mSideStage.removeAllTasks(wct, false /* toTop */);
-            mMainStage.deactivate(wct, false /* toTop */);
+            deactivateSplit(wct, false /* reparentToTop */);
             wct.reorder(mRootTaskInfo.token, false /* onTop */);
             setRootForceTranslucent(true, wct);
             wct.setBounds(mSideStage.mRootTaskInfo.token, mTempRect1);
@@ -1229,7 +1238,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
                 childrenToTop.fadeOutDecor(() -> {
                     WindowContainerTransaction finishedWCT = new WindowContainerTransaction();
                     mIsExiting = false;
-                    mMainStage.deactivate(finishedWCT, childrenToTop == mMainStage /* toTop */);
+                    deactivateSplit(finishedWCT, childrenToTop == mMainStage /* reparentToTop */);
                     mSideStage.removeAllTasks(finishedWCT, childrenToTop == mSideStage /* toTop */);
                     finishedWCT.reorder(mRootTaskInfo.token, false /* toTop */);
                     setRootForceTranslucent(true, finishedWCT);
@@ -1252,7 +1261,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
     }
 
     void dismissSplitScreen(int toTopTaskId, @ExitReason int exitReason) {
-        if (!mMainStage.isActive()) return;
+        if (!isSplitActive()) return;
         final int stage = getStageOfTask(toTopTaskId);
         final WindowContainerTransaction wct = new WindowContainerTransaction();
         prepareExitSplitScreen(stage, wct);
@@ -1362,10 +1371,10 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
      */
     void prepareExitSplitScreen(@StageType int stageToTop,
             @NonNull WindowContainerTransaction wct) {
-        if (!mMainStage.isActive()) return;
+        if (!isSplitActive()) return;
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "prepareExitSplitScreen: stageToTop=%d", stageToTop);
         mSideStage.removeAllTasks(wct, stageToTop == STAGE_TYPE_SIDE);
-        mMainStage.deactivate(wct, stageToTop == STAGE_TYPE_MAIN);
+        deactivateSplit(wct, stageToTop == STAGE_TYPE_MAIN);
     }
 
     private void prepareEnterSplitScreen(WindowContainerTransaction wct) {
@@ -1662,7 +1671,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         mRootTaskInfo = taskInfo;
         if (mSplitLayout != null
                 && mSplitLayout.updateConfiguration(mRootTaskInfo.configuration)
-                && mMainStage.isActive()) {
+                && isSplitActive()) {
             ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onTaskInfoChanged: task=%d updating",
                     taskInfo.taskId);
             // Clear the divider remote animating flag as the divider will be re-rendered to apply
@@ -1916,7 +1925,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
                 stageListener == mMainStageListener);
         final boolean hasChildren = stageListener.mHasChildren;
         final boolean isSideStage = stageListener == mSideStageListener;
-        if (!hasChildren && !mIsExiting && mMainStage.isActive()) {
+        if (!hasChildren && !mIsExiting && isSplitActive()) {
             if (isSideStage && mMainStageListener.mVisible) {
                 // Exit to main stage if side stage no longer has children.
                 mSplitLayout.flingDividerToDismiss(
@@ -1931,7 +1940,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
                 // Dismiss split screen in the background once any sides of the split become empty.
                 exitSplitScreen(null /* childrenToTop */, EXIT_REASON_APP_FINISHED);
             }
-        } else if (isSideStage && hasChildren && !mMainStage.isActive()) {
+        } else if (isSideStage && hasChildren && !isSplitActive()) {
             final WindowContainerTransaction wct = new WindowContainerTransaction();
             prepareEnterSplitScreen(wct);
 
@@ -2146,7 +2155,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
 
     private void onDisplayChange(int displayId, int fromRotation, int toRotation,
             @Nullable DisplayAreaInfo newDisplayAreaInfo, WindowContainerTransaction wct) {
-        if (displayId != DEFAULT_DISPLAY || !mMainStage.isActive()) {
+        if (displayId != DEFAULT_DISPLAY || !isSplitActive()) {
             return;
         }
 
@@ -2441,7 +2450,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             // Not entering or exiting, so just do some house-keeping and validation.
 
             // If we're not in split-mode, just abort so something else can handle it.
-            if (!mMainStage.isActive()) return false;
+            if (!isSplitActive()) return false;
 
             ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "startAnimation: transition=%d", info.getDebugId());
             mSplitLayout.setFreezeDividerWindow(false);
@@ -2683,7 +2692,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
     public void onTransitionAnimationComplete() {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onTransitionAnimationComplete");
         // If still playing, let it finish.
-        if (!mMainStage.isActive() && !mIsExiting) {
+        if (!isSplitActive() && !mIsExiting) {
             // Update divider state after animation so that it is still around and positioned
             // properly for the animation itself.
             mSplitLayout.release();
@@ -3147,7 +3156,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
                 + (mSplitLayout != null ? mSplitLayout.isLeftRightSplit() : "null"));
         pw.println(innerPrefix + "MainStage");
         pw.println(childPrefix + "stagePosition=" + splitPositionToString(getMainStagePosition()));
-        pw.println(childPrefix + "isActive=" + mMainStage.isActive());
+        pw.println(childPrefix + "isActive=" + isSplitActive());
         mMainStage.dump(pw, childPrefix);
         pw.println(innerPrefix + "MainStageListener");
         mMainStageListener.dump(pw, childPrefix);
@@ -3267,7 +3276,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         @Override
         public void onNoLongerSupportMultiWindow(ActivityManager.RunningTaskInfo taskInfo) {
             ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onNoLongerSupportMultiWindow: task=%s", taskInfo);
-            if (mMainStage.isActive()) {
+            if (isSplitActive()) {
                 final boolean isMainStage = mMainStageListener == this;
 
                 // If visible, we preserve the app and keep it running. If an app becomes
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
index 99f3832192e0..459355305280 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
@@ -475,6 +475,30 @@ public class StageTaskListener implements ShellTaskOrganizer.TaskListener {
         });
     }
 
+    // --------
+    // Previously only used in SideStage
+    boolean removeAllTasks(WindowContainerTransaction wct, boolean toTop) {
+        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "remove all side stage tasks: childCount=%d toTop=%b",
+                mChildrenTaskInfo.size(), toTop);
+        if (mChildrenTaskInfo.size() == 0) return false;
+        wct.reparentTasks(
+                mRootTaskInfo.token,
+                null /* newParent */,
+                null /* windowingModes */,
+                null /* activityTypes */,
+                toTop);
+        return true;
+    }
+
+    boolean removeTask(int taskId, WindowContainerToken newParent, WindowContainerTransaction wct) {
+        final ActivityManager.RunningTaskInfo task = mChildrenTaskInfo.get(taskId);
+        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "remove side stage task: task=%d exists=%b", taskId,
+                task != null);
+        if (task == null) return false;
+        wct.reparent(task.token, newParent, false /* onTop */);
+        return true;
+    }
+
     private void sendStatusChanged() {
         mCallbacks.onStatusChanged(mRootTaskInfo.isVisible, mChildrenTaskInfo.size() > 0);
     }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SideStageTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SideStageTests.java
deleted file mode 100644
index 549bd3fcabfb..000000000000
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SideStageTests.java
+++ /dev/null
@@ -1,93 +0,0 @@
-/*
- * Copyright (C) 2020 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.wm.shell.splitscreen;
-
-import static android.view.Display.DEFAULT_DISPLAY;
-
-import static com.google.common.truth.Truth.assertThat;
-
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.ArgumentMatchers.isNull;
-import static org.mockito.Mockito.verify;
-
-import android.app.ActivityManager;
-import android.view.SurfaceControl;
-import android.view.SurfaceSession;
-import android.window.WindowContainerTransaction;
-
-import androidx.test.annotation.UiThreadTest;
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
-import com.android.launcher3.icons.IconProvider;
-import com.android.wm.shell.ShellTaskOrganizer;
-import com.android.wm.shell.ShellTestCase;
-import com.android.wm.shell.TestRunningTaskInfoBuilder;
-import com.android.wm.shell.common.SyncTransactionQueue;
-
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
-import org.mockito.Spy;
-
-import java.util.Optional;
-
-/** Tests for {@link SideStage} */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
-public class SideStageTests extends ShellTestCase {
-    @Mock private ShellTaskOrganizer mTaskOrganizer;
-    @Mock private StageTaskListener.StageListenerCallbacks mCallbacks;
-    @Mock private SyncTransactionQueue mSyncQueue;
-    @Mock private ActivityManager.RunningTaskInfo mRootTask;
-    @Mock private SurfaceControl mRootLeash;
-    @Mock private IconProvider mIconProvider;
-    @Spy private WindowContainerTransaction mWct;
-    private SurfaceSession mSurfaceSession = new SurfaceSession();
-    private SideStage mSideStage;
-
-    @Before
-    @UiThreadTest
-    public void setup() {
-        MockitoAnnotations.initMocks(this);
-        mRootTask = new TestRunningTaskInfoBuilder().build();
-        mSideStage = new SideStage(mContext, mTaskOrganizer, DEFAULT_DISPLAY, mCallbacks,
-                mSyncQueue, mSurfaceSession, mIconProvider, Optional.empty());
-        mSideStage.onTaskAppeared(mRootTask, mRootLeash);
-    }
-
-    @Test
-    public void testAddTask() {
-        final ActivityManager.RunningTaskInfo task = new TestRunningTaskInfoBuilder().build();
-
-        mSideStage.addTask(task, mWct);
-
-        verify(mWct).reparent(eq(task.token), eq(mRootTask.token), eq(true));
-    }
-
-    @Test
-    public void testRemoveTask() {
-        final ActivityManager.RunningTaskInfo task = new TestRunningTaskInfoBuilder().build();
-        assertThat(mSideStage.removeTask(task.taskId, null, mWct)).isFalse();
-
-        mSideStage.mChildrenTaskInfo.put(task.taskId, task);
-        assertThat(mSideStage.removeTask(task.taskId, null, mWct)).isTrue();
-        verify(mWct).reparent(eq(task.token), isNull(), eq(false));
-    }
-}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTestUtils.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTestUtils.java
index aa96c45489dd..4de227836104 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTestUtils.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTestUtils.java
@@ -74,7 +74,7 @@ public class SplitTestUtils {
         final SurfaceControl mRootLeash;
 
         TestStageCoordinator(Context context, int displayId, SyncTransactionQueue syncQueue,
-                ShellTaskOrganizer taskOrganizer, MainStage mainStage, SideStage sideStage,
+                ShellTaskOrganizer taskOrganizer, MainStage mainStage, StageTaskListener sideStage,
                 DisplayController displayController, DisplayImeController imeController,
                 DisplayInsetsController insetsController, SplitLayout splitLayout,
                 Transitions transitions, TransactionPool transactionPool,
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
index abe3dcc1eb80..e16743386489 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
@@ -117,7 +117,7 @@ public class SplitTransitionTests extends ShellTestCase {
     private final TestShellExecutor mTestShellExecutor = new TestShellExecutor();
     private SplitLayout mSplitLayout;
     private MainStage mMainStage;
-    private SideStage mSideStage;
+    private StageTaskListener mSideStage;
     private StageCoordinator mStageCoordinator;
     private SplitScreenTransitions mSplitScreenTransitions;
 
@@ -137,7 +137,7 @@ public class SplitTransitionTests extends ShellTestCase {
                 StageTaskListener.StageListenerCallbacks.class), mSyncQueue, mSurfaceSession,
                 mIconProvider, Optional.of(mWindowDecorViewModel)));
         mMainStage.onTaskAppeared(new TestRunningTaskInfoBuilder().build(), createMockSurface());
-        mSideStage = spy(new SideStage(mContext, mTaskOrganizer, DEFAULT_DISPLAY, mock(
+        mSideStage = spy(new StageTaskListener(mContext, mTaskOrganizer, DEFAULT_DISPLAY, mock(
                 StageTaskListener.StageListenerCallbacks.class), mSyncQueue, mSurfaceSession,
                 mIconProvider, Optional.of(mWindowDecorViewModel)));
         mSideStage.onTaskAppeared(new TestRunningTaskInfoBuilder().build(), createMockSurface());
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
index 0054cb6ccc8c..c9e1414c39b9 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
@@ -99,7 +99,7 @@ public class StageCoordinatorTests extends ShellTestCase {
     @Mock
     private MainStage mMainStage;
     @Mock
-    private SideStage mSideStage;
+    private StageTaskListener mSideStage;
     @Mock
     private SplitLayout mSplitLayout;
     @Mock
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageTaskListenerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageTaskListenerTests.java
index 946a7ef7d8c3..acd612eb34d9 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageTaskListenerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageTaskListenerTests.java
@@ -25,6 +25,7 @@ import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assume.assumeFalse;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.isNull;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
@@ -52,6 +53,7 @@ import org.mockito.ArgumentCaptor;
 import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.mockito.Spy;
 
 import java.util.Optional;
 
@@ -76,6 +78,8 @@ public final class StageTaskListenerTests extends ShellTestCase {
     private IconProvider mIconProvider;
     @Mock
     private WindowDecorViewModel mWindowDecorViewModel;
+    @Spy
+    private WindowContainerTransaction mWct;
     @Captor
     private ArgumentCaptor<SyncTransactionQueue.TransactionRunnable> mRunnableCaptor;
     private SurfaceSession mSurfaceSession = new SurfaceSession();
@@ -177,4 +181,22 @@ public final class StageTaskListenerTests extends ShellTestCase {
         mStageTaskListener.evictAllChildren(wct);
         assertFalse(wct.isEmpty());
     }
+
+    @Test
+    public void testAddTask() {
+        final ActivityManager.RunningTaskInfo task = new TestRunningTaskInfoBuilder().build();
+        mStageTaskListener.addTask(task, mWct);
+
+        verify(mWct).reparent(eq(task.token), eq(mRootTask.token), eq(true));
+    }
+
+    @Test
+    public void testRemoveTask() {
+        final ActivityManager.RunningTaskInfo task = new TestRunningTaskInfoBuilder().build();
+        assertThat(mStageTaskListener.removeTask(task.taskId, null, mWct)).isFalse();
+
+        mStageTaskListener.mChildrenTaskInfo.put(task.taskId, task);
+        assertThat(mStageTaskListener.removeTask(task.taskId, null, mWct)).isTrue();
+        verify(mWct).reparent(eq(task.token), isNull(), eq(false));
+    }
 }