Split DesktopModeWindowDecoration relayout into two stages.

The idea here is to delay ViewHost (SurfaceControlViewHost) creation in
WindowDecoration#relayout() to avoid blocking the animation that will be
posted from TransitionPlayer.
We split WindowDecoration#relayout() in two:
1. View / Surface / Transaction updates.
2. SCVH creation / VRI creation

To delay stage 2 here we simply repost it to the shell-main thread with
a 0ms delay.

We should only split the two stages for app handles; for app headers
it's important to
1. show the header whenever the task/window is visible
2. keep transactions and ViewHost-draw in sync for resize events

Test: WMShellUnitTests
Bug: 346316691
Flag: com.android.window.flags.enable_desktop_windowing_mode

Change-Id: Id1f76cb106e2f9152f4d52c682829fffe53bd026

4 files changed, 297 insertions, 20 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
index 01a479f0bc6c..65f62fb9ccf2 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
@@ -99,10 +99,12 @@ public class DesktopModeWindowDecoration extends WindowDecoration<WindowDecorLin
     private DragPositioningCallback mDragPositioningCallback;
     private DragResizeInputListener mDragResizeListener;
     private DragDetector mDragDetector;
-
+    private Runnable mCurrentViewHostRunnable = null;
     private RelayoutParams mRelayoutParams = new RelayoutParams();
     private final WindowDecoration.RelayoutResult<WindowDecorLinearLayout> mResult =
             new WindowDecoration.RelayoutResult<>();
+    private final Runnable mViewHostRunnable =
+            () -> updateViewHost(mRelayoutParams, null /* onDrawTransaction */, mResult);
 
     private final Point mPositionInParent = new Point();
     private HandleMenu mHandleMenu;
@@ -194,17 +196,88 @@ public class DesktopModeWindowDecoration extends WindowDecoration<WindowDecorLin
         // position and crop are set.
         final boolean shouldSetTaskPositionAndCrop = !DesktopModeStatus.isVeiledResizeEnabled()
                 && mTaskDragResizer.isResizingOrAnimating();
-        // Use |applyStartTransactionOnDraw| so that the transaction (that applies task crop) is
-        // synced with the buffer transaction (that draws the View). Both will be shown on screen
-        // at the same, whereas applying them independently causes flickering. See b/270202228.
-        relayout(taskInfo, t, t, true /* applyStartTransactionOnDraw */,
-                shouldSetTaskPositionAndCrop);
+        // For headers only (i.e. in freeform): use |applyStartTransactionOnDraw| so that the
+        // transaction (that applies task crop) is synced with the buffer transaction (that draws
+        // the View). Both will be shown on screen at the same, whereas applying them independently
+        // causes flickering. See b/270202228.
+        final boolean applyTransactionOnDraw =
+                taskInfo.getWindowingMode() == WINDOWING_MODE_FREEFORM;
+        relayout(taskInfo, t, t, applyTransactionOnDraw, shouldSetTaskPositionAndCrop);
+        if (!applyTransactionOnDraw) {
+            t.apply();
+        }
     }
 
     void relayout(ActivityManager.RunningTaskInfo taskInfo,
             SurfaceControl.Transaction startT, SurfaceControl.Transaction finishT,
             boolean applyStartTransactionOnDraw, boolean shouldSetTaskPositionAndCrop) {
         Trace.beginSection("DesktopModeWindowDecoration#relayout");
+        if (taskInfo.getWindowingMode() == WINDOWING_MODE_FREEFORM) {
+            // The Task is in Freeform mode -> show its header in sync since it's an integral part
+            // of the window itself - a delayed header might cause bad UX.
+            relayoutInSync(taskInfo, startT, finishT, applyStartTransactionOnDraw,
+                    shouldSetTaskPositionAndCrop);
+        } else {
+            // The Task is outside Freeform mode -> allow the handle view to be delayed since the
+            // handle is just a small addition to the window.
+            relayoutWithDelayedViewHost(taskInfo, startT, finishT, applyStartTransactionOnDraw,
+                    shouldSetTaskPositionAndCrop);
+        }
+        Trace.endSection();
+    }
+
+    /** Run the whole relayout phase immediately without delay. */
+    private void relayoutInSync(ActivityManager.RunningTaskInfo taskInfo,
+            SurfaceControl.Transaction startT, SurfaceControl.Transaction finishT,
+            boolean applyStartTransactionOnDraw, boolean shouldSetTaskPositionAndCrop) {
+        // Clear the current ViewHost runnable as we will update the ViewHost here
+        clearCurrentViewHostRunnable();
+        updateRelayoutParamsAndSurfaces(taskInfo, startT, finishT, applyStartTransactionOnDraw,
+                shouldSetTaskPositionAndCrop);
+        if (mResult.mRootView != null) {
+            updateViewHost(mRelayoutParams, startT, mResult);
+        }
+    }
+
+    /**
+     * Clear the current ViewHost runnable - to ensure it doesn't run once relayout params have been
+     * updated.
+     */
+    private void clearCurrentViewHostRunnable() {
+        if (mCurrentViewHostRunnable != null) {
+            mHandler.removeCallbacks(mCurrentViewHostRunnable);
+            mCurrentViewHostRunnable = null;
+        }
+    }
+
+    /**
+     * Relayout the window decoration but repost some of the work, to unblock the current callstack.
+     */
+    private void relayoutWithDelayedViewHost(ActivityManager.RunningTaskInfo taskInfo,
+            SurfaceControl.Transaction startT, SurfaceControl.Transaction finishT,
+            boolean applyStartTransactionOnDraw, boolean shouldSetTaskPositionAndCrop) {
+        if (applyStartTransactionOnDraw) {
+            throw new IllegalArgumentException(
+                    "We cannot both sync viewhost ondraw and delay viewhost creation.");
+        }
+        // Clear the current ViewHost runnable as we will update the ViewHost here
+        clearCurrentViewHostRunnable();
+        updateRelayoutParamsAndSurfaces(taskInfo, startT, finishT,
+                false /* applyStartTransactionOnDraw */, shouldSetTaskPositionAndCrop);
+        if (mResult.mRootView == null) {
+            // This means something blocks the window decor from showing, e.g. the task is hidden.
+            // Nothing is set up in this case including the decoration surface.
+            return;
+        }
+        // Store the current runnable so it can be removed if we start a new relayout.
+        mCurrentViewHostRunnable = mViewHostRunnable;
+        mHandler.post(mCurrentViewHostRunnable);
+    }
+
+    private void updateRelayoutParamsAndSurfaces(ActivityManager.RunningTaskInfo taskInfo,
+            SurfaceControl.Transaction startT, SurfaceControl.Transaction finishT,
+            boolean applyStartTransactionOnDraw, boolean shouldSetTaskPositionAndCrop) {
+        Trace.beginSection("DesktopModeWindowDecoration#updateRelayoutParamsAndSurfaces");
         if (isHandleMenuActive()) {
             mHandleMenu.relayout(startT);
         }
@@ -216,8 +289,8 @@ public class DesktopModeWindowDecoration extends WindowDecoration<WindowDecorLin
         final SurfaceControl oldDecorationSurface = mDecorationContainerSurface;
         final WindowContainerTransaction wct = new WindowContainerTransaction();
 
-        Trace.beginSection("DesktopModeWindowDecoration#relayout-inner");
-        relayout(mRelayoutParams, startT, finishT, wct, oldRootView, mResult);
+        Trace.beginSection("DesktopModeWindowDecoration#relayout-updateViewsAndSurfaces");
+        updateViewsAndSurfaces(mRelayoutParams, startT, finishT, wct, oldRootView, mResult);
         Trace.endSection();
         // After this line, mTaskInfo is up-to-date and should be used instead of taskInfo
 
@@ -228,7 +301,7 @@ public class DesktopModeWindowDecoration extends WindowDecoration<WindowDecorLin
         if (mResult.mRootView == null) {
             // This means something blocks the window decor from showing, e.g. the task is hidden.
             // Nothing is set up in this case including the decoration surface.
-            Trace.endSection(); // DesktopModeWindowDecoration#relayout
+            Trace.endSection(); // DesktopModeWindowDecoration#updateRelayoutParamsAndSurfaces
             return;
         }
 
@@ -246,7 +319,7 @@ public class DesktopModeWindowDecoration extends WindowDecoration<WindowDecorLin
 
         updateDragResizeListener(oldDecorationSurface);
         updateMaximizeMenu(startT);
-        Trace.endSection(); // DesktopModeWindowDecoration#relayout
+        Trace.endSection(); // DesktopModeWindowDecoration#updateRelayoutParamsAndSurfaces
     }
 
     private void updateDragResizeListener(SurfaceControl oldDecorationSurface) {
@@ -851,6 +924,7 @@ public class DesktopModeWindowDecoration extends WindowDecoration<WindowDecorLin
         closeHandleMenu();
         mExclusionRegionListener.onExclusionRegionDismissed(mTaskInfo.taskId);
         disposeResizeVeil();
+        clearCurrentViewHostRunnable();
         super.close();
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/WindowDecoration.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/WindowDecoration.java
index b9532dd9e681..216990c35247 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/WindowDecoration.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/WindowDecoration.java
@@ -199,8 +199,16 @@ public abstract class WindowDecoration<T extends View & TaskFocusStateConsumer>
     void relayout(RelayoutParams params, SurfaceControl.Transaction startT,
             SurfaceControl.Transaction finishT, WindowContainerTransaction wct, T rootView,
             RelayoutResult<T> outResult) {
-        outResult.reset();
+        updateViewsAndSurfaces(params, startT, finishT, wct, rootView, outResult);
+        if (outResult.mRootView != null) {
+            updateViewHost(params, startT, outResult);
+        }
+    }
 
+    protected void updateViewsAndSurfaces(RelayoutParams params,
+            SurfaceControl.Transaction startT, SurfaceControl.Transaction finishT,
+            WindowContainerTransaction wct, T rootView, RelayoutResult<T> outResult) {
+        outResult.reset();
         if (params.mRunningTaskInfo != null) {
             mTaskInfo = params.mRunningTaskInfo;
         }
@@ -236,7 +244,6 @@ public abstract class WindowDecoration<T extends View & TaskFocusStateConsumer>
         updateCaptionContainerSurface(startT, outResult);
         updateCaptionInsets(params, wct, outResult, taskBounds);
         updateTaskSurface(params, startT, finishT, outResult);
-        updateViewHost(params, startT, outResult);
     }
 
     private void inflateIfNeeded(RelayoutParams params, WindowContainerTransaction wct,
@@ -410,8 +417,17 @@ public abstract class WindowDecoration<T extends View & TaskFocusStateConsumer>
         }
     }
 
-    private void updateViewHost(RelayoutParams params, SurfaceControl.Transaction onDrawTransaction,
-            RelayoutResult<T> outResult) {
+    /**
+     * Updates a {@link SurfaceControlViewHost} to connect the window decoration surfaces with our
+     * View hierarchy.
+     *
+     * @param params parameters to use from the last relayout
+     * @param onDrawTransaction a transaction to apply in sync with #onDraw
+     * @param outResult results to use from the last relayout
+     *
+     */
+    protected void updateViewHost(RelayoutParams params,
+            SurfaceControl.Transaction onDrawTransaction, RelayoutResult<T> outResult) {
         Trace.beginSection("CaptionViewHostLayout");
         if (mCaptionWindowManager == null) {
             // Put caption under a container surface because ViewRootImpl sets the destination frame
@@ -433,6 +449,9 @@ public abstract class WindowDecoration<T extends View & TaskFocusStateConsumer>
             mViewHost = mSurfaceControlViewHostFactory.create(mDecorWindowContext, mDisplay,
                     mCaptionWindowManager);
             if (params.mApplyStartTransactionOnDraw) {
+                if (onDrawTransaction == null) {
+                    throw new IllegalArgumentException("Trying to sync a null Transaction");
+                }
                 mViewHost.getRootSurfaceControl().applyTransactionOnDraw(onDrawTransaction);
             }
             mViewHost.setView(outResult.mRootView, lp);
@@ -440,6 +459,9 @@ public abstract class WindowDecoration<T extends View & TaskFocusStateConsumer>
         } else {
             Trace.beginSection("CaptionViewHostLayout-relayout");
             if (params.mApplyStartTransactionOnDraw) {
+                if (onDrawTransaction == null) {
+                    throw new IllegalArgumentException("Trying to sync a null Transaction");
+                }
                 mViewHost.getRootSurfaceControl().applyTransactionOnDraw(onDrawTransaction);
             }
             mViewHost.relayout(lp);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java
index 1b223cf2bd59..46c158908226 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java
@@ -23,12 +23,15 @@ import static android.platform.test.flag.junit.SetFlagsRule.DefaultInitValueType
 import static android.view.WindowInsetsController.APPEARANCE_TRANSPARENT_CAPTION_BAR_BACKGROUND;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
+import static com.android.wm.shell.MockSurfaceControlHelper.createMockSurfaceControlTransaction;
 
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.anyInt;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -37,7 +40,7 @@ import android.app.ActivityManager;
 import android.content.ComponentName;
 import android.content.pm.ActivityInfo;
 import android.content.pm.ApplicationInfo;
-import android.content.res.Configuration;
+import android.content.pm.PackageManager;
 import android.content.res.Resources;
 import android.content.res.TypedArray;
 import android.os.Handler;
@@ -47,13 +50,19 @@ import android.platform.test.annotations.EnableFlags;
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableContext;
+import android.view.AttachedSurfaceControl;
 import android.view.Choreographer;
 import android.view.Display;
+import android.view.GestureDetector;
+import android.view.InsetsState;
+import android.view.MotionEvent;
 import android.view.SurfaceControl;
 import android.view.SurfaceControlViewHost;
+import android.view.View;
 import android.view.WindowManager;
 import android.window.WindowContainerTransaction;
 
+import androidx.annotation.Nullable;
 import androidx.test.filters.SmallTest;
 
 import com.android.dx.mockito.inline.extended.StaticMockitoSession;
@@ -74,6 +83,7 @@ import org.junit.BeforeClass;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.quality.Strictness;
 
@@ -112,18 +122,25 @@ public class DesktopModeWindowDecorationTests extends ShellTestCase {
     @Mock
     private Supplier<SurfaceControl.Transaction> mMockTransactionSupplier;
     @Mock
-    private SurfaceControl.Transaction mMockTransaction;
-    @Mock
     private SurfaceControl mMockSurfaceControl;
     @Mock
     private SurfaceControlViewHost mMockSurfaceControlViewHost;
     @Mock
+    private AttachedSurfaceControl mMockRootSurfaceControl;
+    @Mock
     private WindowDecoration.SurfaceControlViewHostFactory mMockSurfaceControlViewHostFactory;
     @Mock
     private TypedArray mMockRoundedCornersRadiusArray;
 
-    private final Configuration mConfiguration = new Configuration();
+    @Mock
+    private TestTouchEventListener mMockTouchEventListener;
+    @Mock
+    private DesktopModeWindowDecoration.ExclusionRegionListener mMockExclusionRegionListener;
+    @Mock
+    private PackageManager mMockPackageManager;
 
+    private final InsetsState mInsetsState = new InsetsState();
+    private SurfaceControl.Transaction mMockTransaction;
     private StaticMockitoSession mMockitoSession;
     private TestableContext mTestableContext;
 
@@ -145,9 +162,17 @@ public class DesktopModeWindowDecorationTests extends ShellTestCase {
         when(DesktopModeStatus.useDesktopOverrideDensity()).thenReturn(false);
         doReturn(mMockSurfaceControlViewHost).when(mMockSurfaceControlViewHostFactory).create(
                 any(), any(), any());
+        when(mMockSurfaceControlViewHost.getRootSurfaceControl())
+                .thenReturn(mMockRootSurfaceControl);
+        mMockTransaction = createMockSurfaceControlTransaction();
         doReturn(mMockTransaction).when(mMockTransactionSupplier).get();
         mTestableContext = new TestableContext(mContext);
         mTestableContext.ensureTestableResources();
+        mContext.setMockPackageManager(mMockPackageManager);
+        when(mMockPackageManager.getApplicationLabel(any())).thenReturn("applicationLabel");
+        final Display defaultDisplay = mock(Display.class);
+        doReturn(defaultDisplay).when(mMockDisplayController).getDisplay(Display.DEFAULT_DISPLAY);
+        doReturn(mInsetsState).when(mMockDisplayController).getInsetsState(anyInt());
     }
 
     @After
@@ -341,6 +366,99 @@ public class DesktopModeWindowDecorationTests extends ShellTestCase {
         assertThat(hasNoInputChannelFeature(relayoutParams)).isTrue();
     }
 
+    @Test
+    public void relayout_fullscreenTask_appliesTransactionImmediately() {
+        final ActivityManager.RunningTaskInfo taskInfo = createTaskInfo(/* visible= */ true);
+        final DesktopModeWindowDecoration spyWindowDecor = spy(createWindowDecoration(taskInfo));
+        taskInfo.configuration.windowConfiguration.setWindowingMode(WINDOWING_MODE_FULLSCREEN);
+
+        spyWindowDecor.relayout(taskInfo);
+
+        verify(mMockTransaction).apply();
+        verify(mMockRootSurfaceControl, never()).applyTransactionOnDraw(any());
+    }
+
+    @Test
+    public void relayout_freeformTask_appliesTransactionOnDraw() {
+        final ActivityManager.RunningTaskInfo taskInfo = createTaskInfo(/* visible= */ true);
+        final DesktopModeWindowDecoration spyWindowDecor = spy(createWindowDecoration(taskInfo));
+        taskInfo.configuration.windowConfiguration.setWindowingMode(WINDOWING_MODE_FREEFORM);
+        // Make non-resizable to avoid dealing with input-permissions (MONITOR_INPUT)
+        taskInfo.isResizeable = false;
+
+        spyWindowDecor.relayout(taskInfo);
+
+        verify(mMockTransaction, never()).apply();
+        verify(mMockRootSurfaceControl).applyTransactionOnDraw(mMockTransaction);
+    }
+
+    @Test
+    public void relayout_fullscreenTask_doesNotCreateViewHostImmediately() {
+        final ActivityManager.RunningTaskInfo taskInfo = createTaskInfo(/* visible= */ true);
+        final DesktopModeWindowDecoration spyWindowDecor = spy(createWindowDecoration(taskInfo));
+        taskInfo.configuration.windowConfiguration.setWindowingMode(WINDOWING_MODE_FULLSCREEN);
+
+        spyWindowDecor.relayout(taskInfo);
+
+        verify(mMockSurfaceControlViewHostFactory, never()).create(any(), any(), any());
+    }
+
+    @Test
+    public void relayout_fullscreenTask_postsViewHostCreation() {
+        final ActivityManager.RunningTaskInfo taskInfo = createTaskInfo(/* visible= */ true);
+        final DesktopModeWindowDecoration spyWindowDecor = spy(createWindowDecoration(taskInfo));
+        taskInfo.configuration.windowConfiguration.setWindowingMode(WINDOWING_MODE_FULLSCREEN);
+
+        ArgumentCaptor<Runnable> runnableArgument = ArgumentCaptor.forClass(Runnable.class);
+        spyWindowDecor.relayout(taskInfo);
+
+        verify(mMockHandler).post(runnableArgument.capture());
+        runnableArgument.getValue().run();
+        verify(mMockSurfaceControlViewHostFactory).create(any(), any(), any());
+    }
+
+    @Test
+    public void relayout_freeformTask_createsViewHostImmediately() {
+        final ActivityManager.RunningTaskInfo taskInfo = createTaskInfo(/* visible= */ true);
+        final DesktopModeWindowDecoration spyWindowDecor = spy(createWindowDecoration(taskInfo));
+        taskInfo.configuration.windowConfiguration.setWindowingMode(WINDOWING_MODE_FREEFORM);
+        // Make non-resizable to avoid dealing with input-permissions (MONITOR_INPUT)
+        taskInfo.isResizeable = false;
+
+        spyWindowDecor.relayout(taskInfo);
+
+        verify(mMockSurfaceControlViewHostFactory).create(any(), any(), any());
+        verify(mMockHandler, never()).post(any());
+    }
+
+    @Test
+    public void relayout_removesExistingHandlerCallback() {
+        final ActivityManager.RunningTaskInfo taskInfo = createTaskInfo(/* visible= */ true);
+        final DesktopModeWindowDecoration spyWindowDecor = spy(createWindowDecoration(taskInfo));
+        taskInfo.configuration.windowConfiguration.setWindowingMode(WINDOWING_MODE_FULLSCREEN);
+        ArgumentCaptor<Runnable> runnableArgument = ArgumentCaptor.forClass(Runnable.class);
+        spyWindowDecor.relayout(taskInfo);
+        verify(mMockHandler).post(runnableArgument.capture());
+
+        spyWindowDecor.relayout(taskInfo);
+
+        verify(mMockHandler).removeCallbacks(runnableArgument.getValue());
+    }
+
+    @Test
+    public void close_removesExistingHandlerCallback() {
+        final ActivityManager.RunningTaskInfo taskInfo = createTaskInfo(/* visible= */ true);
+        final DesktopModeWindowDecoration spyWindowDecor = spy(createWindowDecoration(taskInfo));
+        taskInfo.configuration.windowConfiguration.setWindowingMode(WINDOWING_MODE_FULLSCREEN);
+        ArgumentCaptor<Runnable> runnableArgument = ArgumentCaptor.forClass(Runnable.class);
+        spyWindowDecor.relayout(taskInfo);
+        verify(mMockHandler).post(runnableArgument.capture());
+
+        spyWindowDecor.close();
+
+        verify(mMockHandler).removeCallbacks(runnableArgument.getValue());
+    }
+
     private void fillRoundedCornersResources(int fillValue) {
         when(mMockRoundedCornersRadiusArray.getDimensionPixelSize(anyInt(), anyInt()))
                 .thenReturn(fillValue);
@@ -361,12 +479,16 @@ public class DesktopModeWindowDecorationTests extends ShellTestCase {
 
     private DesktopModeWindowDecoration createWindowDecoration(
             ActivityManager.RunningTaskInfo taskInfo) {
-        return new DesktopModeWindowDecoration(mContext, mMockDisplayController,
-                mMockShellTaskOrganizer, taskInfo, mMockSurfaceControl,
+        DesktopModeWindowDecoration windowDecor = new DesktopModeWindowDecoration(mContext,
+                mMockDisplayController, mMockShellTaskOrganizer, taskInfo, mMockSurfaceControl,
                 mMockHandler, mMockChoreographer, mMockSyncQueue, mMockRootTaskDisplayAreaOrganizer,
                 SurfaceControl.Builder::new, mMockTransactionSupplier,
                 WindowContainerTransaction::new, SurfaceControl::new,
                 mMockSurfaceControlViewHostFactory);
+        windowDecor.setCaptionListeners(mMockTouchEventListener, mMockTouchEventListener,
+                mMockTouchEventListener, mMockTouchEventListener);
+        windowDecor.setExclusionRegionListener(mMockExclusionRegionListener);
+        return windowDecor;
     }
 
     private ActivityManager.RunningTaskInfo createTaskInfo(boolean visible) {
@@ -391,4 +513,32 @@ public class DesktopModeWindowDecorationTests extends ShellTestCase {
         return (params.mInputFeatures & WindowManager.LayoutParams.INPUT_FEATURE_NO_INPUT_CHANNEL)
                 != 0;
     }
+
+    private static class TestTouchEventListener extends GestureDetector.SimpleOnGestureListener
+            implements View.OnClickListener, View.OnTouchListener, View.OnLongClickListener,
+            View.OnGenericMotionListener, DragDetector.MotionEventHandler {
+
+        @Override
+        public void onClick(View v) {}
+
+        @Override
+        public boolean onGenericMotion(View v, MotionEvent event) {
+            return false;
+        }
+
+        @Override
+        public boolean onLongClick(View v) {
+            return false;
+        }
+
+        @Override
+        public boolean onTouch(View v, MotionEvent event) {
+            return false;
+        }
+
+        @Override
+        public boolean handleMotionEvent(@Nullable View v, MotionEvent ev) {
+            return false;
+        }
+    }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/WindowDecorationTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/WindowDecorationTests.java
index e73069ab52a7..f3603e1d9b46 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/WindowDecorationTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/WindowDecorationTests.java
@@ -32,6 +32,7 @@ import static junit.framework.Assert.assertTrue;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
 import static org.mockito.ArgumentMatchers.anyFloat;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.any;
@@ -828,6 +829,36 @@ public class WindowDecorationTests extends ShellTestCase {
                 eq(mMockTaskSurface), anyInt(), anyInt());
     }
 
+    @Test
+    public void updateViewHost_applyTransactionOnDrawIsTrue_surfaceControlIsUpdated() {
+        final TestWindowDecoration windowDecor = createWindowDecoration(
+                new TestRunningTaskInfoBuilder().build());
+        mRelayoutParams.mApplyStartTransactionOnDraw = true;
+
+        windowDecor.updateViewHost(mRelayoutParams, mMockSurfaceControlStartT, mRelayoutResult);
+
+        verify(mMockRootSurfaceControl).applyTransactionOnDraw(mMockSurfaceControlStartT);
+    }
+
+    @Test
+    public void updateViewHost_nullDrawTransaction_applyTransactionOnDrawIsTrue_throwsException() {
+        final TestWindowDecoration windowDecor = createWindowDecoration(
+                new TestRunningTaskInfoBuilder().build());
+        mRelayoutParams.mApplyStartTransactionOnDraw = true;
+
+        assertThrows(IllegalArgumentException.class,
+                () -> windowDecor.updateViewHost(
+                        mRelayoutParams, null /* onDrawTransaction */, mRelayoutResult));
+    }
+
+    @Test
+    public void updateViewHost_nullDrawTransaction_applyTransactionOnDrawIsFalse_doesNotThrow() {
+        final TestWindowDecoration windowDecor = createWindowDecoration(
+                new TestRunningTaskInfoBuilder().build());
+        mRelayoutParams.mApplyStartTransactionOnDraw = false;
+
+        windowDecor.updateViewHost(mRelayoutParams, null /* onDrawTransaction */, mRelayoutResult);
+    }
 
     private TestWindowDecoration createWindowDecoration(ActivityManager.RunningTaskInfo taskInfo) {
         return new TestWindowDecoration(mContext, mMockDisplayController, mMockShellTaskOrganizer,