Stop assuming the existence of displays when tasks appear

WM shell may not receive the display added notifications when it
receives tasks appearing, so we need to wait for the display to show up
when it doesn't exist.

It's fairly difficult to test the case in real life, but it's a good
case for a unit test. Therefore this CL also makes WindowDecoration unit
tests friendly.

Bug: 236215336
Test: Caption in shell looks OK.
Test: atest WMShellUnitTests:WindowDecorationTests
Test: atest WMShellUnitTests:PipTaskOrganizerTest
Test: atest WMShellUnitTests:PipAnimationControllerTest
Change-Id: Ife0f85fdbd59b84303fc50f4ec896c53a259cec6

libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecoration.java

@@ -95,6 +95,7 @@ public class CaptionWindowDecoration extends WindowDecoration<WindowDecorLinearL
         mDragResizeCallback = dragResizeCallback;
     }
 
+    @Override
     void relayout(ActivityManager.RunningTaskInfo taskInfo) {
         final int shadowRadiusDp = taskInfo.isFocused
                 ? DECOR_SHADOW_FOCUSED_THICKNESS_IN_DIP : DECOR_SHADOW_UNFOCUSED_THICKNESS_IN_DIP;
@@ -119,8 +120,8 @@ public class CaptionWindowDecoration extends WindowDecoration<WindowDecorLinearL
         });
 
         if (mResult.mRootView == null) {
-            // This means the task is hidden. Nothing is set up in this case including the
-            // decoration surface.
+            // 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;
         }
         if (oldRootView != mResult.mRootView) {

libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/WindowDecoration.java

@@ -38,6 +38,8 @@ import android.window.WindowContainerTransaction;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.DisplayController;
 
+import java.util.function.Supplier;
+
 /**
  * Manages a container surface and a windowless window to show window decoration. Responsible to
  * update window decoration window state and layout parameters on task info changes and so that
@@ -53,7 +55,8 @@ import com.android.wm.shell.common.DisplayController;
  *
  * @param <T> The type of the root view
  */
-public class WindowDecoration<T extends View & TaskFocusStateConsumer> implements AutoCloseable {
+public abstract class WindowDecoration<T extends View & TaskFocusStateConsumer>
+        implements AutoCloseable {
     private static final int[] CAPTION_INSETS_TYPES = { InsetsState.ITYPE_CAPTION_BAR };
 
     /**
@@ -62,6 +65,20 @@ public class WindowDecoration<T extends View & TaskFocusStateConsumer> implement
     final Context mContext;
     final DisplayController mDisplayController;
     final ShellTaskOrganizer mTaskOrganizer;
+    final Supplier<SurfaceControl.Builder> mSurfaceControlBuilderSupplier;
+    final SurfaceControlViewHostFactory mSurfaceControlViewHostFactory;
+    private final DisplayController.OnDisplaysChangedListener mOnDisplaysChangedListener =
+            new DisplayController.OnDisplaysChangedListener() {
+                @Override
+                public void onDisplayAdded(int displayId) {
+                    if (mTaskInfo.displayId != displayId) {
+                        return;
+                    }
+
+                    mDisplayController.removeDisplayWindowListener(this);
+                    relayout(mTaskInfo);
+                }
+            };
 
     RunningTaskInfo mTaskInfo;
     final SurfaceControl mTaskSurface;
@@ -71,7 +88,7 @@ public class WindowDecoration<T extends View & TaskFocusStateConsumer> implement
     SurfaceControl mDecorationContainerSurface;
     SurfaceControl mTaskBackgroundSurface;
 
-    private CaptionWindowManager mCaptionWindowManager;
+    private final CaptionWindowManager mCaptionWindowManager;
     private SurfaceControlViewHost mViewHost;
 
     private final Rect mCaptionInsetsRect = new Rect();
@@ -84,11 +101,25 @@ public class WindowDecoration<T extends View & TaskFocusStateConsumer> implement
             ShellTaskOrganizer taskOrganizer,
             RunningTaskInfo taskInfo,
             SurfaceControl taskSurface) {
+        this(context, displayController, taskOrganizer, taskInfo, taskSurface,
+                SurfaceControl.Builder::new, new SurfaceControlViewHostFactory() {});
+    }
+
+    WindowDecoration(
+            Context context,
+            DisplayController displayController,
+            ShellTaskOrganizer taskOrganizer,
+            RunningTaskInfo taskInfo,
+            SurfaceControl taskSurface,
+            Supplier<SurfaceControl.Builder> surfaceControlBuilderSupplier,
+            SurfaceControlViewHostFactory surfaceControlViewHostFactory) {
         mContext = context;
         mDisplayController = displayController;
         mTaskOrganizer = taskOrganizer;
         mTaskInfo = taskInfo;
         mTaskSurface = taskSurface;
+        mSurfaceControlBuilderSupplier = surfaceControlBuilderSupplier;
+        mSurfaceControlViewHostFactory = surfaceControlViewHostFactory;
 
         mDisplay = mDisplayController.getDisplay(mTaskInfo.displayId);
         mDecorWindowContext = mContext.createConfigurationContext(mTaskInfo.getConfiguration());
@@ -99,6 +130,15 @@ public class WindowDecoration<T extends View & TaskFocusStateConsumer> implement
                 new CaptionWindowManager(mTaskInfo.getConfiguration(), mTaskSurface);
     }
 
+    /**
+     * Used by {@link WindowDecoration} to trigger a new relayout because the requirements for a
+     * relayout weren't satisfied are satisfied now.
+     *
+     * @param taskInfo The previous {@link RunningTaskInfo} passed into {@link #relayout} or the
+     *                 constructor.
+     */
+    abstract void relayout(RunningTaskInfo taskInfo);
+
     void relayout(RunningTaskInfo taskInfo, int layoutResId, T rootView, float captionHeightDp,
             Rect outsetsDp, float shadowRadiusDp, SurfaceControl.Transaction t,
             WindowContainerTransaction wct, RelayoutResult<T> outResult) {
@@ -110,7 +150,7 @@ public class WindowDecoration<T extends View & TaskFocusStateConsumer> implement
         }
 
         if (!mTaskInfo.isVisible) {
-            close();
+            releaseViews();
             t.hide(mTaskSurface);
             return;
         }
@@ -123,10 +163,14 @@ public class WindowDecoration<T extends View & TaskFocusStateConsumer> implement
         rootView = null; // Clear it just in case we use it accidentally
         final Configuration taskConfig = mTaskInfo.getConfiguration();
         if (oldTaskConfig.densityDpi != taskConfig.densityDpi
+                || mDisplay == null
                 || mDisplay.getDisplayId() != mTaskInfo.displayId) {
-            close();
+            releaseViews();
 
-            mDisplay = mDisplayController.getDisplay(mTaskInfo.displayId);
+            if (!obtainDisplayOrRegisterListener()) {
+                outResult.mRootView = null;
+                return;
+            }
             mDecorWindowContext = mContext.createConfigurationContext(taskConfig);
             if (layoutResId != 0) {
                 outResult.mRootView =
@@ -141,7 +185,7 @@ public class WindowDecoration<T extends View & TaskFocusStateConsumer> implement
 
         // DecorationContainerSurface
         if (mDecorationContainerSurface == null) {
-            final SurfaceControl.Builder builder = new SurfaceControl.Builder();
+            final SurfaceControl.Builder builder = mSurfaceControlBuilderSupplier.get();
             mDecorationContainerSurface = builder
                     .setName("Decor container of Task=" + mTaskInfo.taskId)
                     .setContainerLayer()
@@ -168,7 +212,7 @@ public class WindowDecoration<T extends View & TaskFocusStateConsumer> implement
 
         // TaskBackgroundSurface
         if (mTaskBackgroundSurface == null) {
-            final SurfaceControl.Builder builder = new SurfaceControl.Builder();
+            final SurfaceControl.Builder builder = mSurfaceControlBuilderSupplier.get();
             mTaskBackgroundSurface = builder
                     .setName("Background of Task=" + mTaskInfo.taskId)
                     .setEffectLayer()
@@ -195,7 +239,7 @@ public class WindowDecoration<T extends View & TaskFocusStateConsumer> implement
         lp.setTitle("Caption of Task=" + mTaskInfo.taskId);
         lp.setTrustedOverlay();
         if (mViewHost == null) {
-            mViewHost = new SurfaceControlViewHost(mDecorWindowContext, mDisplay,
+            mViewHost = mSurfaceControlViewHostFactory.create(mDecorWindowContext, mDisplay,
                     mCaptionWindowManager, true);
             mViewHost.setView(outResult.mRootView, lp);
         } else {
@@ -225,8 +269,22 @@ public class WindowDecoration<T extends View & TaskFocusStateConsumer> implement
                 .show(mTaskSurface);
     }
 
-    @Override
-    public void close() {
+    /**
+     * Obtains the {@link Display} instance for the display ID in {@link #mTaskInfo} if it exists or
+     * registers {@link #mOnDisplaysChangedListener} if it doesn't.
+     *
+     * @return {@code true} if the {@link Display} instance exists; or {@code false} otherwise
+     */
+    private boolean obtainDisplayOrRegisterListener() {
+        mDisplay = mDisplayController.getDisplay(mTaskInfo.displayId);
+        if (mDisplay == null) {
+            mDisplayController.addDisplayWindowListener(mOnDisplaysChangedListener);
+            return false;
+        }
+        return true;
+    }
+
+    private void releaseViews() {
         if (mViewHost != null) {
             mViewHost.release();
             mViewHost = null;
@@ -243,6 +301,12 @@ public class WindowDecoration<T extends View & TaskFocusStateConsumer> implement
         }
     }
 
+    @Override
+    public void close() {
+        mDisplayController.removeDisplayWindowListener(mOnDisplaysChangedListener);
+        releaseViews();
+    }
+
     static class RelayoutResult<T extends View & TaskFocusStateConsumer> {
         int mWidth;
         int mHeight;
@@ -267,4 +331,11 @@ public class WindowDecoration<T extends View & TaskFocusStateConsumer> implement
             super.setConfiguration(configuration);
         }
     }
+
+    interface SurfaceControlViewHostFactory {
+        default SurfaceControlViewHost create(
+                Context c, Display d, WindowlessWindowManager wmm, boolean useSfChoreographer) {
+            return new SurfaceControlViewHost(c, d, wmm, useSfChoreographer);
+        }
+    }
 }

libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/MockSurfaceControlHelper.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2022 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;
+
+import static org.mockito.Mockito.RETURNS_SELF;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+
+import android.view.SurfaceControl;
+
+/**
+ * Helper class to provide mocks for {@link SurfaceControl.Builder} and
+ * {@link SurfaceControl.Transaction} with method chaining support.
+ */
+public class MockSurfaceControlHelper {
+    private MockSurfaceControlHelper() {}
+
+    /**
+     * Creates a mock {@link SurfaceControl.Builder} that supports method chaining and return the
+     * given {@link SurfaceControl} when calling {@link SurfaceControl.Builder#build()}.
+     *
+     * @param mockSurfaceControl the first {@link SurfaceControl} to return
+     * @param mockSurfaceControls following {@link SurfaceControl} to return
+     * @return the mock of {@link SurfaceControl.Builder}
+     */
+    public static SurfaceControl.Builder createMockSurfaceControlBuilder(
+            SurfaceControl mockSurfaceControl, SurfaceControl... mockSurfaceControls) {
+        final SurfaceControl.Builder mockBuilder = mock(SurfaceControl.Builder.class, RETURNS_SELF);
+        doReturn(mockSurfaceControl, (Object[]) mockSurfaceControls)
+                .when(mockBuilder)
+                .build();
+        return mockBuilder;
+    }
+
+    /**
+     * Creates a mock {@link SurfaceControl.Transaction} that supports method chaining.
+     * @return the mock of {@link SurfaceControl.Transaction}
+     */
+    public static SurfaceControl.Transaction createMockSurfaceControlTransaction() {
+        return mock(SurfaceControl.Transaction.class, RETURNS_SELF);
+    }
+}

libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/TestRunningTaskInfoBuilder.java

@@ -25,8 +25,10 @@ import static org.mockito.Mockito.mock;
 
 import android.app.ActivityManager;
 import android.app.WindowConfiguration;
+import android.graphics.Point;
 import android.graphics.Rect;
 import android.os.IBinder;
+import android.view.Display;
 import android.window.IWindowContainerToken;
 import android.window.WindowContainerToken;
 
@@ -38,6 +40,10 @@ public final class TestRunningTaskInfoBuilder {
     private int mParentTaskId = INVALID_TASK_ID;
     private @WindowConfiguration.ActivityType int mActivityType = ACTIVITY_TYPE_STANDARD;
     private @WindowConfiguration.WindowingMode int mWindowingMode = WINDOWING_MODE_UNDEFINED;
+    private int mDisplayId = Display.DEFAULT_DISPLAY;
+    private ActivityManager.TaskDescription.Builder mTaskDescriptionBuilder = null;
+    private final Point mPositionInParent = new Point();
+    private boolean mIsVisible = false;
 
     public static WindowContainerToken createMockWCToken() {
         final IWindowContainerToken itoken = mock(IWindowContainerToken.class);
@@ -68,17 +74,42 @@ public final class TestRunningTaskInfoBuilder {
         return this;
     }
 
+    public TestRunningTaskInfoBuilder setDisplayId(int displayId) {
+        mDisplayId = displayId;
+        return this;
+    }
+
+    public TestRunningTaskInfoBuilder setTaskDescriptionBuilder(
+            ActivityManager.TaskDescription.Builder builder) {
+        mTaskDescriptionBuilder = builder;
+        return this;
+    }
+
+    public TestRunningTaskInfoBuilder setPositionInParent(int x, int y) {
+        mPositionInParent.set(x, y);
+        return this;
+    }
+
+    public TestRunningTaskInfoBuilder setVisible(boolean isVisible) {
+        mIsVisible = isVisible;
+        return this;
+    }
+
     public ActivityManager.RunningTaskInfo build() {
         final ActivityManager.RunningTaskInfo info = new ActivityManager.RunningTaskInfo();
-        info.parentTaskId = INVALID_TASK_ID;
         info.taskId = sNextTaskId++;
         info.parentTaskId = mParentTaskId;
+        info.displayId = mDisplayId;
         info.configuration.windowConfiguration.setBounds(mBounds);
         info.configuration.windowConfiguration.setActivityType(mActivityType);
         info.configuration.windowConfiguration.setWindowingMode(mWindowingMode);
         info.token = mToken;
         info.isResizeable = true;
         info.supportsMultiWindow = true;
+        info.taskDescription =
+                mTaskDescriptionBuilder != null ? mTaskDescriptionBuilder.build() : null;
+        info.positionInParent = mPositionInParent;
+        info.isVisible = mIsVisible;
         return info;
     }
 }

libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/PipAnimationControllerTest.java

@@ -21,6 +21,7 @@ import static android.view.Surface.ROTATION_0;
 import static android.view.Surface.ROTATION_270;
 import static android.view.Surface.ROTATION_90;
 
+import static com.android.wm.shell.MockSurfaceControlHelper.createMockSurfaceControlTransaction;
 import static com.android.wm.shell.pip.PipAnimationController.TRANSITION_DIRECTION_LEAVE_PIP;
 import static com.android.wm.shell.pip.PipAnimationController.TRANSITION_DIRECTION_TO_PIP;
 
@@ -37,6 +38,7 @@ import android.view.SurfaceControl;
 
 import androidx.test.filters.SmallTest;
 
+import com.android.wm.shell.MockSurfaceControlHelper;
 import com.android.wm.shell.ShellTestCase;
 
 import org.junit.Before;
@@ -103,7 +105,8 @@ public class PipAnimationControllerTest extends ShellTestCase {
         final PipAnimationController.PipTransitionAnimator oldAnimator = mPipAnimationController
                 .getAnimator(mTaskInfo, mLeash, baseValue, startValue, endValue1, null,
                         TRANSITION_DIRECTION_TO_PIP, 0, ROTATION_0);
-        oldAnimator.setSurfaceControlTransactionFactory(PipDummySurfaceControlTx::new);
+        oldAnimator.setSurfaceControlTransactionFactory(
+                MockSurfaceControlHelper::createMockSurfaceControlTransaction);
         oldAnimator.start();
 
         final PipAnimationController.PipTransitionAnimator newAnimator = mPipAnimationController
@@ -133,7 +136,7 @@ public class PipAnimationControllerTest extends ShellTestCase {
 
     @Test
     public void pipTransitionAnimator_rotatedEndValue() {
-        final PipDummySurfaceControlTx tx = new PipDummySurfaceControlTx();
+        final SurfaceControl.Transaction tx = createMockSurfaceControlTransaction();
         final Rect startBounds = new Rect(200, 700, 400, 800);
         final Rect endBounds = new Rect(0, 0, 500, 1000);
         // Fullscreen to PiP.
@@ -183,7 +186,8 @@ public class PipAnimationControllerTest extends ShellTestCase {
         final PipAnimationController.PipTransitionAnimator animator = mPipAnimationController
                 .getAnimator(mTaskInfo, mLeash, baseValue, startValue, endValue, null,
                         TRANSITION_DIRECTION_TO_PIP, 0, ROTATION_0);
-        animator.setSurfaceControlTransactionFactory(PipDummySurfaceControlTx::new);
+        animator.setSurfaceControlTransactionFactory(
+                MockSurfaceControlHelper::createMockSurfaceControlTransaction);
 
         animator.setPipAnimationCallback(mPipAnimationCallback);
 

libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/PipDummySurfaceControlTx.java

@@ -1,66 +0,0 @@
-/*
- * Copyright (C) 2022 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.pip;
-
-import android.graphics.Matrix;
-import android.view.SurfaceControl;
-
-/**
- * A dummy {@link SurfaceControl.Transaction} class for testing purpose and supports
- * method chaining.
- */
-public class PipDummySurfaceControlTx extends SurfaceControl.Transaction {
-    @Override
-    public SurfaceControl.Transaction setAlpha(SurfaceControl leash, float alpha) {
-        return this;
-    }
-
-    @Override
-    public SurfaceControl.Transaction setPosition(SurfaceControl leash, float x, float y) {
-        return this;
-    }
-
-    @Override
-    public SurfaceControl.Transaction setWindowCrop(SurfaceControl leash, int w, int h) {
-        return this;
-    }
-
-    @Override
-    public SurfaceControl.Transaction setCornerRadius(SurfaceControl leash, float radius) {
-        return this;
-    }
-
-    @Override
-    public SurfaceControl.Transaction setShadowRadius(SurfaceControl leash, float radius) {
-        return this;
-    }
-
-    @Override
-    public SurfaceControl.Transaction setMatrix(SurfaceControl leash, Matrix matrix,
-            float[] float9) {
-        return this;
-    }
-
-    @Override
-    public SurfaceControl.Transaction setFrameTimelineVsync(long frameTimelineVsyncId) {
-        return this;
-    }
-
-    @Override
-    public void apply() {}
-}
-

libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/PipTaskOrganizerTest.java

@@ -44,6 +44,7 @@ import android.util.Size;
 import android.view.DisplayInfo;
 import android.window.WindowContainerToken;
 
+import com.android.wm.shell.MockSurfaceControlHelper;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.TestShellExecutor;
@@ -246,7 +247,8 @@ public class PipTaskOrganizerTest extends ShellTestCase {
         mPipBoundsState.setDisplayLayout(new DisplayLayout(info,
                 mContext.getResources(), true, true));
         mSpiedPipTaskOrganizer.setOneShotAnimationType(PipAnimationController.ANIM_TYPE_ALPHA);
-        mSpiedPipTaskOrganizer.setSurfaceControlTransactionFactory(PipDummySurfaceControlTx::new);
+        mSpiedPipTaskOrganizer.setSurfaceControlTransactionFactory(
+                MockSurfaceControlHelper::createMockSurfaceControlTransaction);
         doNothing().when(mSpiedPipTaskOrganizer).enterPipWithAlphaAnimation(any(), anyLong());
         doNothing().when(mSpiedPipTaskOrganizer).scheduleAnimateResizePip(any(), anyInt(), any());
     }

libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/WindowDecorationTests.java

@@ -0,0 +1,178 @@
+/*
+ * Copyright (C) 2022 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.windowdecor;
+
+import static com.android.wm.shell.MockSurfaceControlHelper.createMockSurfaceControlBuilder;
+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.anyBoolean;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.same;
+import static org.mockito.Mockito.verify;
+
+import android.app.ActivityManager;
+import android.content.Context;
+import android.graphics.Color;
+import android.graphics.Rect;
+import android.testing.AndroidTestingRunner;
+import android.view.Display;
+import android.view.SurfaceControl;
+import android.view.SurfaceControlViewHost;
+import android.view.View;
+import android.window.WindowContainerTransaction;
+
+import androidx.test.filters.SmallTest;
+
+import com.android.wm.shell.ShellTaskOrganizer;
+import com.android.wm.shell.ShellTestCase;
+import com.android.wm.shell.TestRunningTaskInfoBuilder;
+import com.android.wm.shell.common.DisplayController;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+
+import java.util.function.Supplier;
+
+/**
+ * Tests for {@link WindowDecoration}.
+ *
+ * Build/Install/Run:
+ * atest WMShellUnitTests:WindowDecorationTests
+ */
+@SmallTest
+@RunWith(AndroidTestingRunner.class)
+public class WindowDecorationTests extends ShellTestCase {
+    private static final int CAPTION_HEIGHT_DP = 32;
+    private static final int SHADOW_RADIUS_DP = 5;
+
+    private final Rect mOutsetsDp = new Rect();
+    private final WindowDecoration.RelayoutResult<TestView> mRelayoutResult =
+            new WindowDecoration.RelayoutResult<>();
+
+    @Mock
+    private DisplayController mMockDisplayController;
+    @Mock
+    private ShellTaskOrganizer mMockShellTaskOrganizer;
+    @Mock
+    private WindowDecoration.SurfaceControlViewHostFactory mMockSurfaceControlViewHostFactory;
+    @Mock
+    private SurfaceControlViewHost mMockSurfaceControlViewHost;
+    @Mock
+    private TestView mMockView;
+    @Mock
+    private WindowContainerTransaction mMockWindowContainerTransaction;
+
+    private SurfaceControl.Builder mMockSurfaceControlBuilder;
+    private SurfaceControl.Transaction mMockSurfaceControlTransaction;
+
+    @Before
+    public void setUp() {
+        mMockSurfaceControlBuilder = createMockSurfaceControlBuilder(mock(SurfaceControl.class));
+        mMockSurfaceControlTransaction = createMockSurfaceControlTransaction();
+
+        doReturn(mMockSurfaceControlViewHost).when(mMockSurfaceControlViewHostFactory)
+                .create(any(), any(), any(), anyBoolean());
+    }
+
+    @Test
+    public void testNotCrashWhenDisplayAppearsAfterTask() {
+        doReturn(mock(Display.class)).when(mMockDisplayController)
+                .getDisplay(Display.DEFAULT_DISPLAY);
+
+        final int displayId = Display.DEFAULT_DISPLAY + 1;
+        final ActivityManager.TaskDescription.Builder taskDescriptionBuilder =
+                new ActivityManager.TaskDescription.Builder()
+                        .setBackgroundColor(Color.BLACK);
+        final ActivityManager.RunningTaskInfo taskInfo = new TestRunningTaskInfoBuilder()
+                .setDisplayId(displayId)
+                .setTaskDescriptionBuilder(taskDescriptionBuilder)
+                .setVisible(true)
+                .build();
+
+        final TestWindowDecoration windowDecor =
+                createWindowDecoration(taskInfo, new SurfaceControl());
+        windowDecor.relayout(taskInfo);
+
+        // It shouldn't show the window decoration when it can't obtain the display instance.
+        assertThat(mRelayoutResult.mRootView).isNull();
+
+        final ArgumentCaptor<DisplayController.OnDisplaysChangedListener> listenerArgumentCaptor =
+                ArgumentCaptor.forClass(DisplayController.OnDisplaysChangedListener.class);
+        verify(mMockDisplayController).addDisplayWindowListener(listenerArgumentCaptor.capture());
+        final DisplayController.OnDisplaysChangedListener listener =
+                listenerArgumentCaptor.getValue();
+
+        // Adding an irrelevant display shouldn't change the result.
+        listener.onDisplayAdded(Display.DEFAULT_DISPLAY);
+        assertThat(mRelayoutResult.mRootView).isNull();
+
+        final Display mockDisplay = mock(Display.class);
+        doReturn(mockDisplay).when(mMockDisplayController).getDisplay(displayId);
+
+        listener.onDisplayAdded(displayId);
+
+        // The listener should be removed when the display shows up.
+        verify(mMockDisplayController).removeDisplayWindowListener(same(listener));
+
+        assertThat(mRelayoutResult.mRootView).isSameInstanceAs(mMockView);
+        verify(mMockSurfaceControlViewHostFactory)
+                .create(any(), eq(mockDisplay), any(), anyBoolean());
+        verify(mMockSurfaceControlViewHost).setView(same(mMockView), any());
+    }
+
+    private TestWindowDecoration createWindowDecoration(
+            ActivityManager.RunningTaskInfo taskInfo, SurfaceControl testSurface) {
+        return new TestWindowDecoration(mContext, mMockDisplayController, mMockShellTaskOrganizer,
+                taskInfo, testSurface, () -> mMockSurfaceControlBuilder,
+                mMockSurfaceControlViewHostFactory);
+    }
+
+    private static class TestView extends View implements TaskFocusStateConsumer {
+        private TestView(Context context) {
+            super(context);
+        }
+
+        @Override
+        public void setTaskFocusState(boolean focused) {}
+    }
+
+    private class TestWindowDecoration extends WindowDecoration<TestView> {
+        TestWindowDecoration(Context context, DisplayController displayController,
+                ShellTaskOrganizer taskOrganizer, ActivityManager.RunningTaskInfo taskInfo,
+                SurfaceControl taskSurface,
+                Supplier<SurfaceControl.Builder> surfaceControlBuilderSupplier,
+                SurfaceControlViewHostFactory surfaceControlViewHostFactory) {
+            super(context, displayController, taskOrganizer, taskInfo, taskSurface,
+                    surfaceControlBuilderSupplier, surfaceControlViewHostFactory);
+        }
+
+        @Override
+        void relayout(ActivityManager.RunningTaskInfo taskInfo) {
+            relayout(null /* taskInfo */, 0 /* layoutResId */, mMockView, CAPTION_HEIGHT_DP,
+                    mOutsetsDp, SHADOW_RADIUS_DP, mMockSurfaceControlTransaction,
+                    mMockWindowContainerTransaction, mRelayoutResult);
+        }
+    }
+}