Merge "Selectively expect device to stay awake on fold" into 24D1-dev

3 files changed, 256 insertions, 12 deletions

diff --git a/services/core/java/com/android/server/display/DisplayManagerService.java b/services/core/java/com/android/server/display/DisplayManagerService.java
index ad89444edcfd..28664dd66c19 100644
--- a/services/core/java/com/android/server/display/DisplayManagerService.java
+++ b/services/core/java/com/android/server/display/DisplayManagerService.java
@@ -145,6 +145,7 @@ import android.window.ScreenCapture;
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.display.BrightnessSynchronizer;
+import com.android.internal.foldables.FoldGracePeriodProvider;
 import com.android.internal.foldables.FoldLockSettingAvailabilityProvider;
 import com.android.internal.os.BackgroundThread;
 import com.android.internal.util.ArrayUtils;
@@ -573,7 +574,7 @@ public final class DisplayManagerService extends SystemService {
         mUiHandler = UiThread.getHandler();
         mDisplayDeviceRepo = new DisplayDeviceRepository(mSyncRoot, mPersistentDataStore);
         mLogicalDisplayMapper = new LogicalDisplayMapper(mContext,
-                foldSettingProvider,
+                foldSettingProvider, new FoldGracePeriodProvider(),
                 mDisplayDeviceRepo, new LogicalDisplayListener(), mSyncRoot, mHandler, mFlags);
         mDisplayModeDirector = new DisplayModeDirector(context, mHandler, mFlags);
         mBrightnessSynchronizer = new BrightnessSynchronizer(mContext,
diff --git a/services/core/java/com/android/server/display/LogicalDisplayMapper.java b/services/core/java/com/android/server/display/LogicalDisplayMapper.java
index 2e8de31f2af1..86afa9b0630e 100644
--- a/services/core/java/com/android/server/display/LogicalDisplayMapper.java
+++ b/services/core/java/com/android/server/display/LogicalDisplayMapper.java
@@ -40,6 +40,7 @@ import android.view.DisplayAddress;
 import android.view.DisplayInfo;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.foldables.FoldGracePeriodProvider;
 import com.android.server.display.feature.DisplayManagerFlags;
 import com.android.server.display.layout.DisplayIdProducer;
 import com.android.server.display.layout.Layout;
@@ -120,7 +121,7 @@ class LogicalDisplayMapper implements DisplayDeviceRepository.Listener {
     /**
      * Sleep the device when transitioning into these device state.
      */
-    private final SparseBooleanArray mDeviceStatesOnWhichToSleep;
+    private final SparseBooleanArray mDeviceStatesOnWhichToSelectiveSleep;
 
     /**
      * Map of all logical displays indexed by logical display id.
@@ -153,6 +154,7 @@ class LogicalDisplayMapper implements DisplayDeviceRepository.Listener {
     private final DisplayManagerService.SyncRoot mSyncRoot;
     private final LogicalDisplayMapperHandler mHandler;
     private final FoldSettingProvider mFoldSettingProvider;
+    private final FoldGracePeriodProvider mFoldGracePeriodProvider;
     private final PowerManager mPowerManager;
 
     /**
@@ -200,15 +202,18 @@ class LogicalDisplayMapper implements DisplayDeviceRepository.Listener {
     private final DisplayManagerFlags mFlags;
 
     LogicalDisplayMapper(@NonNull Context context, FoldSettingProvider foldSettingProvider,
+            FoldGracePeriodProvider foldGracePeriodProvider,
             @NonNull DisplayDeviceRepository repo,
             @NonNull Listener listener, @NonNull DisplayManagerService.SyncRoot syncRoot,
             @NonNull Handler handler, DisplayManagerFlags flags) {
-        this(context, foldSettingProvider, repo, listener, syncRoot, handler,
+        this(context, foldSettingProvider, foldGracePeriodProvider, repo, listener, syncRoot,
+                handler,
                 new DeviceStateToLayoutMap((isDefault) -> isDefault ? DEFAULT_DISPLAY
                         : sNextNonDefaultDisplayId++, flags), flags);
     }
 
     LogicalDisplayMapper(@NonNull Context context, FoldSettingProvider foldSettingProvider,
+            FoldGracePeriodProvider foldGracePeriodProvider,
             @NonNull DisplayDeviceRepository repo,
             @NonNull Listener listener, @NonNull DisplayManagerService.SyncRoot syncRoot,
             @NonNull Handler handler, @NonNull DeviceStateToLayoutMap deviceStateToLayoutMap,
@@ -220,13 +225,15 @@ class LogicalDisplayMapper implements DisplayDeviceRepository.Listener {
         mDisplayDeviceRepo = repo;
         mListener = listener;
         mFoldSettingProvider = foldSettingProvider;
+        mFoldGracePeriodProvider = foldGracePeriodProvider;
         mSingleDisplayDemoMode = SystemProperties.getBoolean("persist.demo.singledisplay", false);
         mSupportsConcurrentInternalDisplays = context.getResources().getBoolean(
                 com.android.internal.R.bool.config_supportsConcurrentInternalDisplays);
         mDeviceStatesOnWhichToWakeUp = toSparseBooleanArray(context.getResources().getIntArray(
                 com.android.internal.R.array.config_deviceStatesOnWhichToWakeUp));
-        mDeviceStatesOnWhichToSleep = toSparseBooleanArray(context.getResources().getIntArray(
-                com.android.internal.R.array.config_deviceStatesOnWhichToSleep));
+        mDeviceStatesOnWhichToSelectiveSleep = toSparseBooleanArray(
+                context.getResources().getIntArray(
+                        com.android.internal.R.array.config_deviceStatesOnWhichToSleep));
         mDisplayDeviceRepo.addListener(this);
         mDeviceStateToLayoutMap = deviceStateToLayoutMap;
         mFlags = flags;
@@ -403,7 +410,7 @@ class LogicalDisplayMapper implements DisplayDeviceRepository.Listener {
         ipw.println("mSingleDisplayDemoMode=" + mSingleDisplayDemoMode);
         ipw.println("mCurrentLayout=" + mCurrentLayout);
         ipw.println("mDeviceStatesOnWhichToWakeUp=" + mDeviceStatesOnWhichToWakeUp);
-        ipw.println("mDeviceStatesOnWhichToSleep=" + mDeviceStatesOnWhichToSleep);
+        ipw.println("mDeviceStatesOnWhichSelectiveSleep=" + mDeviceStatesOnWhichToSelectiveSleep);
         ipw.println("mInteractive=" + mInteractive);
         ipw.println("mBootCompleted=" + mBootCompleted);
 
@@ -569,8 +576,8 @@ class LogicalDisplayMapper implements DisplayDeviceRepository.Listener {
     boolean shouldDeviceBePutToSleep(int pendingState, int currentState, boolean isOverrideActive,
             boolean isInteractive, boolean isBootCompleted) {
         return currentState != DeviceStateManager.INVALID_DEVICE_STATE
-                && mDeviceStatesOnWhichToSleep.get(pendingState)
-                && !mDeviceStatesOnWhichToSleep.get(currentState)
+                && mDeviceStatesOnWhichToSelectiveSleep.get(pendingState)
+                && !mDeviceStatesOnWhichToSelectiveSleep.get(currentState)
                 && !isOverrideActive
                 && isInteractive && isBootCompleted
                 && !mFoldSettingProvider.shouldStayAwakeOnFold();
@@ -611,9 +618,12 @@ class LogicalDisplayMapper implements DisplayDeviceRepository.Listener {
         final boolean waitingToWakeDevice = mDeviceStatesOnWhichToWakeUp.get(mPendingDeviceState)
                 && !mDeviceStatesOnWhichToWakeUp.get(mDeviceState)
                 && !mInteractive && mBootCompleted;
-        final boolean waitingToSleepDevice = mDeviceStatesOnWhichToSleep.get(mPendingDeviceState)
-                && !mDeviceStatesOnWhichToSleep.get(mDeviceState)
-                && mInteractive && mBootCompleted;
+        // The device should only wait for sleep if #shouldStayAwakeOnFold method returns false.
+        // If not, device should be marked ready for transition immediately.
+        final boolean waitingToSleepDevice = mDeviceStatesOnWhichToSelectiveSleep.get(
+                mPendingDeviceState)
+                && !mDeviceStatesOnWhichToSelectiveSleep.get(mDeviceState)
+                && mInteractive && mBootCompleted && !shouldStayAwakeOnFold();
 
         final boolean displaysOff = areAllTransitioningDisplaysOffLocked();
         final boolean isReadyToTransition = displaysOff && !waitingToWakeDevice
@@ -1231,6 +1241,16 @@ class LogicalDisplayMapper implements DisplayDeviceRepository.Listener {
         return retval;
     }
 
+    /**
+     * Returns true if the device would definitely have outer display ON/Stay Awake on fold based on
+     * the value of `Continue using app on fold` setting
+     */
+    private boolean shouldStayAwakeOnFold() {
+        return mFoldSettingProvider.shouldStayAwakeOnFold() || (
+                mFoldSettingProvider.shouldSelectiveStayAwakeOnFold()
+                        && mFoldGracePeriodProvider.isEnabled());
+    }
+
     private String displayEventToString(int msg) {
         switch(msg) {
             case LOGICAL_DISPLAY_EVENT_ADDED:
diff --git a/services/tests/displayservicetests/src/com/android/server/display/LogicalDisplayMapperTest.java b/services/tests/displayservicetests/src/com/android/server/display/LogicalDisplayMapperTest.java
index bed6f928a55c..357999ab65c9 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/LogicalDisplayMapperTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/LogicalDisplayMapperTest.java
@@ -20,6 +20,8 @@ import static android.hardware.devicestate.DeviceStateManager.INVALID_DEVICE_STA
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.Display.DEFAULT_DISPLAY_GROUP;
 import static android.view.Display.FLAG_REAR;
+import static android.view.Display.STATE_OFF;
+import static android.view.Display.STATE_ON;
 import static android.view.Display.TYPE_EXTERNAL;
 import static android.view.Display.TYPE_INTERNAL;
 import static android.view.Display.TYPE_VIRTUAL;
@@ -28,6 +30,7 @@ import static com.android.server.display.DeviceStateToLayoutMap.STATE_DEFAULT;
 import static com.android.server.display.DisplayAdapter.DISPLAY_DEVICE_EVENT_ADDED;
 import static com.android.server.display.DisplayAdapter.DISPLAY_DEVICE_EVENT_CHANGED;
 import static com.android.server.display.DisplayAdapter.DISPLAY_DEVICE_EVENT_REMOVED;
+import static com.android.server.display.DisplayDeviceInfo.DIFF_EVERYTHING;
 import static com.android.server.display.DisplayDeviceInfo.FLAG_ALLOWED_TO_BE_DEFAULT_DISPLAY;
 import static com.android.server.display.LogicalDisplayMapper.LOGICAL_DISPLAY_EVENT_ADDED;
 import static com.android.server.display.LogicalDisplayMapper.LOGICAL_DISPLAY_EVENT_CONNECTED;
@@ -35,6 +38,9 @@ import static com.android.server.display.LogicalDisplayMapper.LOGICAL_DISPLAY_EV
 import static com.android.server.display.LogicalDisplayMapper.LOGICAL_DISPLAY_EVENT_REMOVED;
 import static com.android.server.display.layout.Layout.Display.POSITION_REAR;
 import static com.android.server.display.layout.Layout.Display.POSITION_UNKNOWN;
+import static com.android.server.utils.FoldSettingProvider.SETTING_VALUE_SELECTIVE_STAY_AWAKE;
+import static com.android.server.utils.FoldSettingProvider.SETTING_VALUE_SLEEP_ON_FOLD;
+import static com.android.server.utils.FoldSettingProvider.SETTING_VALUE_STAY_AWAKE_ON_FOLD;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -72,6 +78,7 @@ import android.view.DisplayInfo;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 
+import com.android.internal.foldables.FoldGracePeriodProvider;
 import com.android.server.display.feature.DisplayManagerFlags;
 import com.android.server.display.layout.DisplayIdProducer;
 import com.android.server.display.layout.Layout;
@@ -96,9 +103,13 @@ import java.util.List;
 @RunWith(AndroidJUnit4.class)
 public class LogicalDisplayMapperTest {
     private static int sUniqueTestDisplayId = 0;
+    private static final int TIMEOUT_STATE_TRANSITION_MILLIS = 500;
+    private static final int FOLD_SETTLE_DELAY = 1000;
     private static final int DEVICE_STATE_CLOSED = 0;
+    private static final int DEVICE_STATE_HALF_OPEN = 1;
     private static final int DEVICE_STATE_OPEN = 2;
     private static final int FLAG_GO_TO_SLEEP_ON_FOLD = 0;
+    private static final int FLAG_GO_TO_SLEEP_FLAG_SOFT_SLEEP = 2;
     private static int sNextNonDefaultDisplayId = DEFAULT_DISPLAY + 1;
     private static final File NON_EXISTING_FILE = new File("/non_existing_folder/should_not_exist");
 
@@ -116,6 +127,7 @@ public class LogicalDisplayMapperTest {
     @Mock LogicalDisplayMapper.Listener mListenerMock;
     @Mock Context mContextMock;
     @Mock FoldSettingProvider mFoldSettingProviderMock;
+    @Mock FoldGracePeriodProvider mFoldGracePeriodProvider;
     @Mock Resources mResourcesMock;
     @Mock IPowerManager mIPowerManagerMock;
     @Mock IThermalService mIThermalServiceMock;
@@ -160,6 +172,7 @@ public class LogicalDisplayMapperTest {
                 .thenReturn(Context.POWER_SERVICE);
         when(mFoldSettingProviderMock.shouldStayAwakeOnFold()).thenReturn(false);
         when(mFoldSettingProviderMock.shouldSleepOnFold()).thenReturn(false);
+        when(mFoldSettingProviderMock.shouldSelectiveStayAwakeOnFold()).thenReturn(true);
         when(mIPowerManagerMock.isInteractive()).thenReturn(true);
         when(mContextMock.getSystemService(PowerManager.class)).thenReturn(mPowerManager);
         when(mContextMock.getResources()).thenReturn(mResourcesMock);
@@ -177,6 +190,7 @@ public class LogicalDisplayMapperTest {
         mLooper = new TestLooper();
         mHandler = new Handler(mLooper.getLooper());
         mLogicalDisplayMapper = new LogicalDisplayMapper(mContextMock, mFoldSettingProviderMock,
+                mFoldGracePeriodProvider,
                 mDisplayDeviceRepo,
                 mListenerMock, new DisplayManagerService.SyncRoot(), mHandler,
                 mDeviceStateToLayoutMapSpy, mFlagsMock);
@@ -694,22 +708,147 @@ public class LogicalDisplayMapperTest {
         when(mFoldSettingProviderMock.shouldSleepOnFold()).thenReturn(true);
 
         finishBootAndFoldDevice();
+        advanceTime(FOLD_SETTLE_DELAY);
 
         verify(mIPowerManagerMock, atLeastOnce()).goToSleep(anyLong(), anyInt(),
                 eq(FLAG_GO_TO_SLEEP_ON_FOLD));
     }
 
     @Test
+    public void testDeviceShouldPutToSleepWhenFoldSettingSelective() throws RemoteException {
+        when(mFoldSettingProviderMock.shouldSelectiveStayAwakeOnFold()).thenReturn(true);
+
+        finishBootAndFoldDevice();
+        advanceTime(FOLD_SETTLE_DELAY);
+
+        verify(mIPowerManagerMock, atLeastOnce()).goToSleep(anyLong(), anyInt(),
+                eq(FLAG_GO_TO_SLEEP_FLAG_SOFT_SLEEP));
+    }
+
+    @Test
     public void testDeviceShouldNotBePutToSleepWhenSleepSettingFalse() throws RemoteException {
         when(mFoldSettingProviderMock.shouldSleepOnFold()).thenReturn(false);
 
         finishBootAndFoldDevice();
+        advanceTime(FOLD_SETTLE_DELAY);
 
         verify(mIPowerManagerMock, never()).goToSleep(anyLong(), anyInt(),
                 eq(FLAG_GO_TO_SLEEP_ON_FOLD));
     }
 
     @Test
+    public void testWaitForSleepWhenFoldSettingSleep() {
+        // Test device should not be marked ready for transition immediately, when 'Continue
+        // using app on fold' set to 'Never'
+        setFoldLockBehaviorSettingValue(SETTING_VALUE_SLEEP_ON_FOLD);
+        setGracePeriodAvailability(false);
+        FoldableDisplayDevices foldableDisplayDevices = createFoldableDeviceStateToLayoutMap();
+
+        finishBootAndFoldDevice();
+        foldableDisplayDevices.mInner.setState(STATE_OFF);
+        notifyDisplayChanges(foldableDisplayDevices.mOuter);
+
+        assertDisplayDisabled(foldableDisplayDevices.mOuter);
+        assertDisplayEnabled(foldableDisplayDevices.mInner);
+    }
+
+    @Test
+    public void testSwapDeviceStateWithDelayWhenFoldSettingSleep() {
+        // Test device should be marked ready for transition after a delay when 'Continue using
+        // app on fold' set to 'Never'
+        setFoldLockBehaviorSettingValue(SETTING_VALUE_SLEEP_ON_FOLD);
+        setGracePeriodAvailability(false);
+        FoldableDisplayDevices foldableDisplayDevices = createFoldableDeviceStateToLayoutMap();
+
+        finishBootAndFoldDevice();
+        foldableDisplayDevices.mInner.setState(STATE_OFF);
+        notifyDisplayChanges(foldableDisplayDevices.mOuter);
+        advanceTime(TIMEOUT_STATE_TRANSITION_MILLIS);
+
+        assertDisplayEnabled(foldableDisplayDevices.mOuter);
+        assertDisplayDisabled(foldableDisplayDevices.mInner);
+    }
+
+    @Test
+    public void testDoNotWaitForSleepWhenFoldSettingStayAwake() {
+        // Test device should be marked ready for transition immediately when 'Continue using app
+        // on fold' set to 'Always'
+        setFoldLockBehaviorSettingValue(SETTING_VALUE_STAY_AWAKE_ON_FOLD);
+        setGracePeriodAvailability(false);
+        FoldableDisplayDevices foldableDisplayDevices = createFoldableDeviceStateToLayoutMap();
+
+        finishBootAndFoldDevice();
+        foldableDisplayDevices.mInner.setState(STATE_OFF);
+        notifyDisplayChanges(foldableDisplayDevices.mOuter);
+
+        assertDisplayEnabled(foldableDisplayDevices.mOuter);
+        assertDisplayDisabled(foldableDisplayDevices.mInner);
+    }
+
+    @Test
+    public void testDoNotWaitForSleepWhenFoldSettingSelectiveStayAwake() {
+        // Test device should be marked ready for transition immediately when 'Continue using app
+        // on fold' set to 'Swipe up to continue'
+        setFoldLockBehaviorSettingValue(SETTING_VALUE_SELECTIVE_STAY_AWAKE);
+        setGracePeriodAvailability(true);
+        FoldableDisplayDevices foldableDisplayDevices = createFoldableDeviceStateToLayoutMap();
+
+        finishBootAndFoldDevice();
+        foldableDisplayDevices.mInner.setState(STATE_OFF);
+        notifyDisplayChanges(foldableDisplayDevices.mOuter);
+
+        assertDisplayEnabled(foldableDisplayDevices.mOuter);
+        assertDisplayDisabled(foldableDisplayDevices.mInner);
+    }
+
+    @Test
+    public void testWaitForSleepWhenGracePeriodSettingDisabled() {
+        // Test device should not be marked ready for transition immediately when 'Continue using
+        // app on fold' set to 'Swipe up to continue' but Grace Period flag is disabled
+        setFoldLockBehaviorSettingValue(SETTING_VALUE_SELECTIVE_STAY_AWAKE);
+        setGracePeriodAvailability(false);
+        FoldableDisplayDevices foldableDisplayDevices = createFoldableDeviceStateToLayoutMap();
+
+        finishBootAndFoldDevice();
+        foldableDisplayDevices.mInner.setState(STATE_OFF);
+        notifyDisplayChanges(foldableDisplayDevices.mOuter);
+
+        assertDisplayDisabled(foldableDisplayDevices.mOuter);
+        assertDisplayEnabled(foldableDisplayDevices.mInner);
+    }
+
+    @Test
+    public void testWaitForSleepWhenTransitionDisplayStaysOn() {
+        // Test device should not be marked ready for transition immediately, when 'Continue
+        // using app on fold' set to 'Always' but not all transitioning displays are OFF.
+        setFoldLockBehaviorSettingValue(SETTING_VALUE_STAY_AWAKE_ON_FOLD);
+        setGracePeriodAvailability(false);
+        FoldableDisplayDevices foldableDisplayDevices = createFoldableDeviceStateToLayoutMap();
+
+        finishBootAndFoldDevice();
+        notifyDisplayChanges(foldableDisplayDevices.mOuter);
+
+        assertDisplayDisabled(foldableDisplayDevices.mOuter);
+        assertDisplayEnabled(foldableDisplayDevices.mInner);
+    }
+
+    @Test
+    public void testSwapDeviceStateWithDelayWhenTransitionDisplayStaysOn() {
+        // Test device should be marked ready for transition after a delay, when 'Continue using
+        // app on fold' set to 'Never' but not all transitioning displays are OFF.
+        setFoldLockBehaviorSettingValue(SETTING_VALUE_SLEEP_ON_FOLD);
+        setGracePeriodAvailability(false);
+        FoldableDisplayDevices foldableDisplayDevices = createFoldableDeviceStateToLayoutMap();
+
+        finishBootAndFoldDevice();
+        notifyDisplayChanges(foldableDisplayDevices.mOuter);
+        advanceTime(TIMEOUT_STATE_TRANSITION_MILLIS);
+
+        assertDisplayEnabled(foldableDisplayDevices.mOuter);
+        assertDisplayDisabled(foldableDisplayDevices.mInner);
+    }
+
+    @Test
     public void testDeviceStateLocked() {
         DisplayDevice device1 = createDisplayDevice(TYPE_INTERNAL, 600, 800,
                 FLAG_ALLOWED_TO_BE_DEFAULT_DISPLAY);
@@ -963,13 +1102,77 @@ public class LogicalDisplayMapperTest {
     // Helper Methods
     /////////////////
 
+    private void setGracePeriodAvailability(boolean isGracePeriodEnabled) {
+        when(mFoldGracePeriodProvider.isEnabled()).thenReturn(isGracePeriodEnabled);
+    }
+
+    private void setFoldLockBehaviorSettingValue(String foldLockBehaviorSettingValue) {
+        when(mFoldSettingProviderMock.shouldSleepOnFold()).thenReturn(false);
+        when(mFoldSettingProviderMock.shouldStayAwakeOnFold()).thenReturn(false);
+        when(mFoldSettingProviderMock.shouldSelectiveStayAwakeOnFold()).thenReturn(false);
+
+        switch (foldLockBehaviorSettingValue) {
+            case SETTING_VALUE_STAY_AWAKE_ON_FOLD:
+                when(mFoldSettingProviderMock.shouldStayAwakeOnFold()).thenReturn(true);
+                break;
+
+            case SETTING_VALUE_SLEEP_ON_FOLD:
+                when(mFoldSettingProviderMock.shouldSleepOnFold()).thenReturn(true);
+                break;
+
+            default:
+                when(mFoldSettingProviderMock.shouldSelectiveStayAwakeOnFold()).thenReturn(true);
+                break;
+        }
+    }
+
+    private FoldableDisplayDevices createFoldableDeviceStateToLayoutMap() {
+        TestDisplayDevice outer = createDisplayDevice(TYPE_INTERNAL, 600, 800,
+                FLAG_ALLOWED_TO_BE_DEFAULT_DISPLAY);
+        TestDisplayDevice inner = createDisplayDevice(TYPE_INTERNAL, 600, 800,
+                FLAG_ALLOWED_TO_BE_DEFAULT_DISPLAY);
+        outer.setState(STATE_OFF);
+        inner.setState(STATE_ON);
+
+        Layout layout = new Layout();
+        createDefaultDisplay(layout, outer);
+        createNonDefaultDisplay(layout, inner, /* enabled= */ false, /* group= */ null);
+        when(mDeviceStateToLayoutMapSpy.get(DEVICE_STATE_CLOSED)).thenReturn(layout);
+
+        layout = new Layout();
+        createNonDefaultDisplay(layout, outer, /* enabled= */ false, /* group= */ null);
+        createDefaultDisplay(layout, inner);
+        when(mDeviceStateToLayoutMapSpy.get(DEVICE_STATE_HALF_OPEN)).thenReturn(layout);
+        when(mDeviceStateToLayoutMapSpy.get(DEVICE_STATE_OPEN)).thenReturn(layout);
+        when(mDeviceStateToLayoutMapSpy.size()).thenReturn(4);
+
+        add(outer);
+        add(inner);
+
+        return new FoldableDisplayDevices(outer, inner);
+    }
+
     private void finishBootAndFoldDevice() {
         mLogicalDisplayMapper.setDeviceStateLocked(DEVICE_STATE_OPEN, false);
+        mLogicalDisplayMapper.onEarlyInteractivityChange(true);
         advanceTime(1000);
         mLogicalDisplayMapper.onBootCompleted();
         advanceTime(1000);
         mLogicalDisplayMapper.setDeviceStateLocked(DEVICE_STATE_CLOSED, false);
-        advanceTime(1000);
+    }
+
+    private void notifyDisplayChanges(TestDisplayDevice displayDevice) {
+        mLogicalDisplayMapper.onDisplayDeviceChangedLocked(displayDevice, DIFF_EVERYTHING);
+    }
+
+    private void assertDisplayEnabled(DisplayDevice displayDevice) {
+        assertThat(
+                mLogicalDisplayMapper.getDisplayLocked(displayDevice).isEnabledLocked()).isTrue();
+    }
+
+    private void assertDisplayDisabled(DisplayDevice displayDevice) {
+        assertThat(
+                mLogicalDisplayMapper.getDisplayLocked(displayDevice).isEnabledLocked()).isFalse();
     }
 
     private void createDefaultDisplay(Layout layout, DisplayDevice device) {
@@ -1071,6 +1274,16 @@ public class LogicalDisplayMapperTest {
         assertNotEquals(DEFAULT_DISPLAY, id(displayRemoved));
     }
 
+    private final static class FoldableDisplayDevices {
+        final TestDisplayDevice mOuter;
+        final TestDisplayDevice mInner;
+
+        FoldableDisplayDevices(TestDisplayDevice outer, TestDisplayDevice inner) {
+            this.mOuter = outer;
+            this.mInner = inner;
+        }
+    }
+
     class TestDisplayDevice extends DisplayDevice {
         private DisplayDeviceInfo mInfo;
         private DisplayDeviceInfo mSentInfo;
@@ -1096,6 +1309,16 @@ public class LogicalDisplayMapperTest {
             mSentInfo = null;
         }
 
+        public void setState(int state) {
+            mState = state;
+            if (mSentInfo == null) {
+                mInfo.state = state;
+            } else {
+                mInfo.state = state;
+                mSentInfo.state = state;
+            }
+        }
+
         @Override
         public boolean hasStableUniqueId() {
             return true;