Adjust split layout with IME animation in split (2/N)

Add surface utils to help establishing dim surface layer for splits.

Bug: 179262787
Test: atest WMShellUnitTests
Change-Id: I0c268e0b1891fddcf717e8e54e0d29acbf8c5f24

12 files changed, 130 insertions, 34 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/apppairs/AppPair.java b/libs/WindowManager/Shell/src/com/android/wm/shell/apppairs/AppPair.java
index eb82c6d597d4..83af53f962f3 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/apppairs/AppPair.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/apppairs/AppPair.java
@@ -25,6 +25,7 @@ import static com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_TASK_ORG
 import android.app.ActivityManager;
 import android.graphics.Rect;
 import android.view.SurfaceControl;
+import android.view.SurfaceSession;
 import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 
@@ -35,6 +36,7 @@ import com.android.internal.protolog.common.ProtoLog;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.common.DisplayImeController;
+import com.android.wm.shell.common.SurfaceUtils;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.common.split.SplitLayout;
 
@@ -54,6 +56,9 @@ class AppPair implements ShellTaskOrganizer.TaskListener, SplitLayout.LayoutChan
     private SurfaceControl mTaskLeash1;
     private ActivityManager.RunningTaskInfo mTaskInfo2;
     private SurfaceControl mTaskLeash2;
+    private SurfaceControl mDimLayer1;
+    private SurfaceControl mDimLayer2;
+    private final SurfaceSession mSurfaceSession = new SurfaceSession();
 
     private final AppPairsController mController;
     private final SyncTransactionQueue mSyncQueue;
@@ -153,9 +158,13 @@ class AppPair implements ShellTaskOrganizer.TaskListener, SplitLayout.LayoutChan
         } else if (taskInfo.taskId == getTaskId1()) {
             mTaskInfo1 = taskInfo;
             mTaskLeash1 = leash;
+            mSyncQueue.runInSync(t -> mDimLayer1 =
+                    SurfaceUtils.makeDimLayer(t, mTaskLeash1, "Dim layer", mSurfaceSession));
         } else if (taskInfo.taskId == getTaskId2()) {
             mTaskInfo2 = taskInfo;
             mTaskLeash2 = leash;
+            mSyncQueue.runInSync(t -> mDimLayer2 =
+                    SurfaceUtils.makeDimLayer(t, mTaskLeash2, "Dim layer", mSurfaceSession));
         } else {
             throw new IllegalStateException("Unknown task=" + taskInfo.taskId);
         }
@@ -212,8 +221,12 @@ class AppPair implements ShellTaskOrganizer.TaskListener, SplitLayout.LayoutChan
         if (taskInfo.taskId == getRootTaskId()) {
             // We don't want to release this object back to the pool since the root task went away.
             mController.unpair(mRootTaskInfo.taskId, false /* releaseToPool */);
-        } else if (taskInfo.taskId == getTaskId1() || taskInfo.taskId == getTaskId2()) {
+        } else if (taskInfo.taskId == getTaskId1()) {
+            mController.unpair(mRootTaskInfo.taskId);
+            mSyncQueue.runInSync(t -> t.remove(mDimLayer1));
+        } else if (taskInfo.taskId == getTaskId2()) {
             mController.unpair(mRootTaskInfo.taskId);
+            mSyncQueue.runInSync(t -> t.remove(mDimLayer2));
         }
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/SurfaceUtils.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/SurfaceUtils.java
new file mode 100644
index 000000000000..55c5125f0a00
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/SurfaceUtils.java
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2021 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.common;
+
+import android.view.SurfaceControl;
+import android.view.SurfaceSession;
+
+/**
+ * Helpers for handling surface.
+ */
+public class SurfaceUtils {
+    /** Creates a dim layer above indicated host surface. */
+    public static SurfaceControl makeDimLayer(SurfaceControl.Transaction t, SurfaceControl host,
+            String name, SurfaceSession surfaceSession) {
+        SurfaceControl dimLayer = new SurfaceControl.Builder(surfaceSession)
+                .setParent(host)
+                .setColorLayer()
+                .setName(name)
+                .setCallsite("SurfaceUtils.makeDimLayer")
+                .build();
+        t.setLayer(dimLayer, Integer.MAX_VALUE).setColor(dimLayer, new float[]{0f, 0f, 0f});
+        return dimLayer;
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/legacysplitscreen/LegacySplitScreenController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/legacysplitscreen/LegacySplitScreenController.java
index 9eacaecc4f50..ee2202a48bf2 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/legacysplitscreen/LegacySplitScreenController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/legacysplitscreen/LegacySplitScreenController.java
@@ -577,7 +577,7 @@ public class LegacySplitScreenController implements DisplayController.OnDisplays
             mSplits.getSplitTransitions().dismissSplit(
                     mSplits, mSplitLayout, !toPrimaryTask, snapped);
         } else {
-        mWindowManagerProxy.applyDismissSplit(mSplits, mSplitLayout, !toPrimaryTask);
+            mWindowManagerProxy.applyDismissSplit(mSplits, mSplitLayout, !toPrimaryTask);
             onDismissSplit();
         }
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/legacysplitscreen/LegacySplitScreenTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/legacysplitscreen/LegacySplitScreenTaskListener.java
index cf35656a395b..912c9605197f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/legacysplitscreen/LegacySplitScreenTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/legacysplitscreen/LegacySplitScreenTaskListener.java
@@ -38,6 +38,7 @@ import androidx.annotation.NonNull;
 
 import com.android.internal.protolog.common.ProtoLog;
 import com.android.wm.shell.ShellTaskOrganizer;
+import com.android.wm.shell.common.SurfaceUtils;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.transition.Transitions;
 
@@ -70,9 +71,9 @@ class LegacySplitScreenTaskListener implements ShellTaskOrganizer.TaskListener {
     private final LegacySplitScreenTransitions mSplitTransitions;
 
     LegacySplitScreenTaskListener(LegacySplitScreenController splitScreenController,
-                    ShellTaskOrganizer shellTaskOrganizer,
-                    Transitions transitions,
-                    SyncTransactionQueue syncQueue) {
+            ShellTaskOrganizer shellTaskOrganizer,
+            Transitions transitions,
+            SyncTransactionQueue syncQueue) {
         mSplitScreenController = splitScreenController;
         mTaskOrganizer = shellTaskOrganizer;
         mSplitTransitions = new LegacySplitScreenTransitions(splitScreenController.mTransactionPool,
@@ -146,21 +147,11 @@ class LegacySplitScreenTaskListener implements ShellTaskOrganizer.TaskListener {
                 ProtoLog.v(WM_SHELL_TASK_ORG, "%s onTaskAppeared Supported", TAG);
 
                 // Initialize dim surfaces:
-                mPrimaryDim = new SurfaceControl.Builder(mSurfaceSession)
-                        .setParent(mPrimarySurface).setColorLayer()
-                        .setName("Primary Divider Dim")
-                        .setCallsite("SplitScreenTaskOrganizer.onTaskAppeared")
-                        .build();
-                mSecondaryDim = new SurfaceControl.Builder(mSurfaceSession)
-                        .setParent(mSecondarySurface).setColorLayer()
-                        .setName("Secondary Divider Dim")
-                        .setCallsite("SplitScreenTaskOrganizer.onTaskAppeared")
-                        .build();
                 SurfaceControl.Transaction t = getTransaction();
-                t.setLayer(mPrimaryDim, Integer.MAX_VALUE);
-                t.setColor(mPrimaryDim, new float[]{0f, 0f, 0f});
-                t.setLayer(mSecondaryDim, Integer.MAX_VALUE);
-                t.setColor(mSecondaryDim, new float[]{0f, 0f, 0f});
+                mPrimaryDim = SurfaceUtils.makeDimLayer(
+                        t, mPrimarySurface, "Primary Divider Dim", mSurfaceSession);
+                mSecondaryDim = SurfaceUtils.makeDimLayer(
+                        t, mSecondarySurface, "Secondary Divider Dim", mSurfaceSession);
                 t.apply();
                 releaseTransaction(t);
             }
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 66a4a60d4be6..d0998eb57633 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
@@ -19,6 +19,7 @@ package com.android.wm.shell.splitscreen;
 import static android.app.WindowConfiguration.WINDOWING_MODE_MULTI_WINDOW;
 
 import android.graphics.Rect;
+import android.view.SurfaceSession;
 import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 
@@ -36,8 +37,9 @@ class MainStage extends StageTaskListener {
     private boolean mIsActive = false;
 
     MainStage(ShellTaskOrganizer taskOrganizer, int displayId,
-            StageListenerCallbacks callbacks, SyncTransactionQueue syncQueue) {
-        super(taskOrganizer, displayId, callbacks, syncQueue);
+            StageListenerCallbacks callbacks, SyncTransactionQueue syncQueue,
+            SurfaceSession surfaceSession) {
+        super(taskOrganizer, displayId, callbacks, syncQueue, surfaceSession);
     }
 
     boolean isActive() {
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
index 01a81d2ff5e0..82f95a4f32ea 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SideStage.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SideStage.java
@@ -18,6 +18,7 @@ package com.android.wm.shell.splitscreen;
 
 import android.app.ActivityManager;
 import android.graphics.Rect;
+import android.view.SurfaceSession;
 import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 
@@ -33,8 +34,9 @@ class SideStage extends StageTaskListener {
     private static final String TAG = SideStage.class.getSimpleName();
 
     SideStage(ShellTaskOrganizer taskOrganizer, int displayId,
-            StageListenerCallbacks callbacks, SyncTransactionQueue syncQueue) {
-        super(taskOrganizer, displayId, callbacks, syncQueue);
+            StageListenerCallbacks callbacks, SyncTransactionQueue syncQueue,
+            SurfaceSession surfaceSession) {
+        super(taskOrganizer, displayId, callbacks, syncQueue, surfaceSession);
     }
 
     void addTask(ActivityManager.RunningTaskInfo task, Rect rootBounds,
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 ea302691af8e..e20feae2d522 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
@@ -47,6 +47,7 @@ import android.os.Bundle;
 import android.os.IBinder;
 import android.util.Log;
 import android.view.SurfaceControl;
+import android.view.SurfaceSession;
 import android.view.WindowManager;
 import android.window.DisplayAreaInfo;
 import android.window.IRemoteTransition;
@@ -92,6 +93,8 @@ class StageCoordinator implements SplitLayout.LayoutChangeListener,
     /** internal value for mDismissTop that represents no dismiss */
     private static final int NO_DISMISS = -2;
 
+    private final SurfaceSession mSurfaceSession = new SurfaceSession();
+
     private final MainStage mMainStage;
     private final StageListenerImpl mMainStageListener = new StageListenerImpl();
     private final SideStage mSideStage;
@@ -116,8 +119,8 @@ class StageCoordinator implements SplitLayout.LayoutChangeListener,
     /** Whether the device is supporting legacy split or not. */
     private boolean mUseLegacySplit;
 
-
     @SplitScreen.StageType int mDismissTop = NO_DISMISS;
+
     private final Runnable mOnTransitionAnimationComplete = () -> {
         // If still playing, let it finish.
         if (!isSplitScreenVisible()) {
@@ -138,8 +141,18 @@ class StageCoordinator implements SplitLayout.LayoutChangeListener,
         mSyncQueue = syncQueue;
         mRootTDAOrganizer = rootTDAOrganizer;
         mTaskOrganizer = taskOrganizer;
-        mMainStage = new MainStage(mTaskOrganizer, mDisplayId, mMainStageListener, mSyncQueue);
-        mSideStage = new SideStage(mTaskOrganizer, mDisplayId, mSideStageListener, mSyncQueue);
+        mMainStage = new MainStage(
+                mTaskOrganizer,
+                mDisplayId,
+                mMainStageListener,
+                mSyncQueue,
+                mSurfaceSession);
+        mSideStage = new SideStage(
+                mTaskOrganizer,
+                mDisplayId,
+                mSideStageListener,
+                mSyncQueue,
+                mSurfaceSession);
         mDisplayImeController = displayImeController;
         mRootTDAOrganizer.registerListener(displayId, this);
         mSplitTransitions = new SplitScreenTransitions(transactionPool, transitions,
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 1da0a2d82766..a6d8a0b098df 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
@@ -29,11 +29,13 @@ import android.graphics.Point;
 import android.graphics.Rect;
 import android.util.SparseArray;
 import android.view.SurfaceControl;
+import android.view.SurfaceSession;
 import android.window.WindowContainerTransaction;
 
 import androidx.annotation.NonNull;
 
 import com.android.wm.shell.ShellTaskOrganizer;
+import com.android.wm.shell.common.SurfaceUtils;
 import com.android.wm.shell.common.SyncTransactionQueue;
 
 import java.io.PrintWriter;
@@ -44,6 +46,7 @@ import java.io.PrintWriter;
  * They only serve to hold a collection of tasks and provide APIs like
  * {@link #setBounds(Rect, WindowContainerTransaction)} for the centralized {@link StageCoordinator}
  * to perform operations in-sync with other containers.
+ *
  * @see StageCoordinator
  */
 class StageTaskListener implements ShellTaskOrganizer.TaskListener {
@@ -58,22 +61,30 @@ class StageTaskListener implements ShellTaskOrganizer.TaskListener {
     /** Callback interface for listening to changes in a split-screen stage. */
     public interface StageListenerCallbacks {
         void onRootTaskAppeared();
+
         void onStatusChanged(boolean visible, boolean hasChildren);
+
         void onChildTaskStatusChanged(int taskId, boolean present, boolean visible);
+
         void onRootTaskVanished();
     }
+
     private final StageListenerCallbacks mCallbacks;
     private final SyncTransactionQueue mSyncQueue;
+    private final SurfaceSession mSurfaceSession;
 
     protected ActivityManager.RunningTaskInfo mRootTaskInfo;
     protected SurfaceControl mRootLeash;
+    protected SurfaceControl mDimLayer;
     protected SparseArray<ActivityManager.RunningTaskInfo> mChildrenTaskInfo = new SparseArray<>();
     private final SparseArray<SurfaceControl> mChildrenLeashes = new SparseArray<>();
 
     StageTaskListener(ShellTaskOrganizer taskOrganizer, int displayId,
-            StageListenerCallbacks callbacks, SyncTransactionQueue syncQueue) {
+            StageListenerCallbacks callbacks, SyncTransactionQueue syncQueue,
+            SurfaceSession surfaceSession) {
         mCallbacks = callbacks;
         mSyncQueue = syncQueue;
+        mSurfaceSession = surfaceSession;
         taskOrganizer.createRootTask(displayId, WINDOWING_MODE_MULTI_WINDOW, this);
     }
 
@@ -93,6 +104,8 @@ class StageTaskListener implements ShellTaskOrganizer.TaskListener {
             mRootTaskInfo = taskInfo;
             mCallbacks.onRootTaskAppeared();
             sendStatusChanged();
+            mSyncQueue.runInSync(t -> mDimLayer =
+                    SurfaceUtils.makeDimLayer(t, mRootLeash, "Dim layer", mSurfaceSession));
         } else if (taskInfo.parentTaskId == mRootTaskInfo.taskId) {
             final int taskId = taskInfo.taskId;
             mChildrenLeashes.put(taskId, leash);
@@ -140,6 +153,7 @@ class StageTaskListener implements ShellTaskOrganizer.TaskListener {
         final int taskId = taskInfo.taskId;
         if (mRootTaskInfo.taskId == taskId) {
             mCallbacks.onRootTaskVanished();
+            mSyncQueue.runInSync(t -> t.remove(mDimLayer));
             mRootTaskInfo = null;
         } else if (mChildrenTaskInfo.contains(taskId)) {
             mChildrenTaskInfo.remove(taskId);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/MainStageTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/MainStageTests.java
index 702e8945de01..1bb5fd1e49e7 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/MainStageTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/MainStageTests.java
@@ -22,6 +22,7 @@ import static com.google.common.truth.Truth.assertThat;
 
 import android.app.ActivityManager;
 import android.view.SurfaceControl;
+import android.view.SurfaceSession;
 import android.window.WindowContainerTransaction;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
@@ -36,7 +37,6 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
-import org.mockito.Spy;
 
 /** Tests for {@link MainStage} */
 @SmallTest
@@ -47,14 +47,16 @@ public class MainStageTests {
     @Mock private SyncTransactionQueue mSyncQueue;
     @Mock private ActivityManager.RunningTaskInfo mRootTaskInfo;
     @Mock private SurfaceControl mRootLeash;
-    @Spy private WindowContainerTransaction mWct;
+    private WindowContainerTransaction mWct = new WindowContainerTransaction();
+    private SurfaceSession mSurfaceSession = new SurfaceSession();
     private MainStage mMainStage;
 
     @Before
     public void setup() {
         MockitoAnnotations.initMocks(this);
         mRootTaskInfo = new TestRunningTaskInfoBuilder().build();
-        mMainStage = new MainStage(mTaskOrganizer, DEFAULT_DISPLAY, mCallbacks, mSyncQueue);
+        mMainStage = new MainStage(mTaskOrganizer, DEFAULT_DISPLAY, mCallbacks, mSyncQueue,
+                mSurfaceSession);
         mMainStage.onTaskAppeared(mRootTaskInfo, mRootLeash);
     }
 
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
index 01888b758bf6..56a005642ce2 100644
--- 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
@@ -26,6 +26,7 @@ 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.ext.junit.runners.AndroidJUnit4;
@@ -52,13 +53,15 @@ public class SideStageTests {
     @Mock private ActivityManager.RunningTaskInfo mRootTask;
     @Mock private SurfaceControl mRootLeash;
     @Spy private WindowContainerTransaction mWct;
+    private SurfaceSession mSurfaceSession = new SurfaceSession();
     private SideStage mSideStage;
 
     @Before
     public void setup() {
         MockitoAnnotations.initMocks(this);
         mRootTask = new TestRunningTaskInfoBuilder().build();
-        mSideStage = new SideStage(mTaskOrganizer, DEFAULT_DISPLAY, mCallbacks, mSyncQueue);
+        mSideStage = new SideStage(mTaskOrganizer, DEFAULT_DISPLAY, mCallbacks, mSyncQueue,
+                mSurfaceSession);
         mSideStage.onTaskAppeared(mRootTask, mRootLeash);
     }
 
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 08ac2a6cfa77..aca80f3556b9 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
@@ -43,6 +43,7 @@ import android.graphics.Rect;
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.view.SurfaceControl;
+import android.view.SurfaceSession;
 import android.window.IRemoteTransition;
 import android.window.IRemoteTransitionFinishedCallback;
 import android.window.TransitionInfo;
@@ -80,6 +81,7 @@ public class SplitTransitionTests extends ShellTestCase {
     @Mock private DisplayImeController mDisplayImeController;
     @Mock private TransactionPool mTransactionPool;
     @Mock private Transitions mTransitions;
+    @Mock private SurfaceSession mSurfaceSession;
     private SplitLayout mSplitLayout;
     private MainStage mMainStage;
     private SideStage mSideStage;
@@ -98,10 +100,10 @@ public class SplitTransitionTests extends ShellTestCase {
         doReturn(mock(SurfaceControl.Transaction.class)).when(mTransactionPool).acquire();
         mSplitLayout = SplitTestUtils.createMockSplitLayout();
         mMainStage = new MainStage(mTaskOrganizer, DEFAULT_DISPLAY, mock(
-                StageTaskListener.StageListenerCallbacks.class), mSyncQueue);
+                StageTaskListener.StageListenerCallbacks.class), mSyncQueue, mSurfaceSession);
         mMainStage.onTaskAppeared(new TestRunningTaskInfoBuilder().build(), createMockSurface());
         mSideStage = new SideStage(mTaskOrganizer, DEFAULT_DISPLAY, mock(
-                StageTaskListener.StageListenerCallbacks.class), mSyncQueue);
+                StageTaskListener.StageListenerCallbacks.class), mSyncQueue, mSurfaceSession);
         mSideStage.onTaskAppeared(new TestRunningTaskInfoBuilder().build(), createMockSurface());
         mStageCoordinator = new SplitTestUtils.TestStageCoordinator(mContext, DEFAULT_DISPLAY,
                     mSyncQueue, mRootTDAOrganizer, mTaskOrganizer, mMainStage, mSideStage,
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 c66e0730422c..afb26948f194 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
@@ -27,6 +27,7 @@ import static org.mockito.Mockito.verify;
 
 import android.app.ActivityManager;
 import android.view.SurfaceControl;
+import android.view.SurfaceSession;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
@@ -38,6 +39,8 @@ import com.android.wm.shell.common.SyncTransactionQueue;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -48,6 +51,8 @@ public final class StageTaskListenerTests {
     @Mock private ShellTaskOrganizer mTaskOrganizer;
     @Mock private StageTaskListener.StageListenerCallbacks mCallbacks;
     @Mock private SyncTransactionQueue mSyncQueue;
+    @Captor private ArgumentCaptor<SyncTransactionQueue.TransactionRunnable> mRunnableCaptor;
+    private SurfaceSession mSurfaceSession = new SurfaceSession();
     private ActivityManager.RunningTaskInfo mRootTask;
     private StageTaskListener mStageTaskListener;
 
@@ -58,13 +63,24 @@ public final class StageTaskListenerTests {
                 mTaskOrganizer,
                 DEFAULT_DISPLAY,
                 mCallbacks,
-                mSyncQueue);
+                mSyncQueue,
+                mSurfaceSession);
         mRootTask = new TestRunningTaskInfoBuilder().build();
         mRootTask.parentTaskId = INVALID_TASK_ID;
         mStageTaskListener.onTaskAppeared(mRootTask, new SurfaceControl());
     }
 
     @Test
+    public void testInitsDimLayer() {
+        verify(mSyncQueue).runInSync(mRunnableCaptor.capture());
+        final SurfaceControl.Transaction t = new SurfaceControl.Transaction();
+        mRunnableCaptor.getValue().runWithTransaction(t);
+        t.apply();
+
+        assertThat(mStageTaskListener.mDimLayer).isNotNull();
+    }
+
+    @Test
     public void testRootTaskAppeared() {
         assertThat(mStageTaskListener.mRootTaskInfo.taskId).isEqualTo(mRootTask.taskId);
         verify(mCallbacks).onRootTaskAppeared();