Update device state overlay config usage to new API

Updates CameraManager and LogicalDisplayMapper to use
the new DeviceState#hasProperty API's instead of the
overlay configuration values for determining what
states are folded or should put the device to sleep

Flag: android.hardware.devicestate.feature.flags.device_state_property_migration
Test: LogicalDisplayManagerTest
Test: CameraManagerTest
Bug: 336640888
Change-Id: Ib9ac5fab821958859d4e33bb22c797c405d7e623
diff --git a/core/java/android/hardware/camera2/CameraManager.java b/core/java/android/hardware/camera2/CameraManager.java
index 2dbd4b8..1b0f20d 100644
--- a/core/java/android/hardware/camera2/CameraManager.java
+++ b/core/java/android/hardware/camera2/CameraManager.java
@@ -56,6 +56,8 @@
 import android.hardware.camera2.utils.ExceptionUtils;
 import android.hardware.devicestate.DeviceState;
 import android.hardware.devicestate.DeviceStateManager;
+import android.hardware.devicestate.feature.flags.FeatureFlags;
+import android.hardware.devicestate.feature.flags.FeatureFlagsImpl;
 import android.hardware.display.DisplayManager;
 import android.os.Binder;
 import android.os.Handler;
@@ -247,14 +249,22 @@
         private ArrayList<WeakReference<DeviceStateListener>> mDeviceStateListeners =
                 new ArrayList<>();
         private boolean mFoldedDeviceState;
+        private final FeatureFlags mDeviceStateManagerFlags;
 
         public FoldStateListener(Context context) {
             mFoldedDeviceStates = context.getResources().getIntArray(
                     com.android.internal.R.array.config_foldedDeviceStates);
+            mDeviceStateManagerFlags = new FeatureFlagsImpl();
         }
 
-        private synchronized void handleStateChange(int state) {
-            boolean folded = ArrayUtils.contains(mFoldedDeviceStates, state);
+        private synchronized void handleStateChange(DeviceState state) {
+            final boolean folded;
+            if (mDeviceStateManagerFlags.deviceStatePropertyMigration()) {
+                folded = state.hasProperty(
+                        DeviceState.PROPERTY_FOLDABLE_DISPLAY_CONFIGURATION_OUTER_PRIMARY);
+            } else {
+                folded = ArrayUtils.contains(mFoldedDeviceStates, state.getIdentifier());
+            }
 
             mFoldedDeviceState = folded;
             Iterator<WeakReference<DeviceStateListener>> it = mDeviceStateListeners.iterator();
@@ -276,10 +286,8 @@
 
         @SuppressWarnings("FlaggedApi")
         @Override
-        public void onDeviceStateChanged(DeviceState state) {
-            // Suppressing the FlaggedAPI warning as this specific API isn't new, just moved to
-            // system API which requires it to be flagged.
-            handleStateChange(state.getIdentifier());
+        public void onDeviceStateChanged(@NonNull DeviceState state) {
+            handleStateChange(state);
         }
     }
 
diff --git a/core/java/android/hardware/devicestate/DeviceStateManager.java b/core/java/android/hardware/devicestate/DeviceStateManager.java
index febc24c1..5b511cc 100644
--- a/core/java/android/hardware/devicestate/DeviceStateManager.java
+++ b/core/java/android/hardware/devicestate/DeviceStateManager.java
@@ -67,6 +67,14 @@
     public static final int MAXIMUM_DEVICE_STATE_IDENTIFIER = 10000;
 
     /**
+     * {@link DeviceState} to represent an invalid device state.
+     * @hide
+     */
+    public static final DeviceState INVALID_DEVICE_STATE = new DeviceState(
+            new DeviceState.Configuration.Builder(INVALID_DEVICE_STATE_IDENTIFIER,
+                    "INVALID").build());
+
+    /**
      * Intent needed to launch the rear display overlay activity from SysUI
      *
      * @hide
diff --git a/services/core/java/com/android/server/display/DisplayManagerService.java b/services/core/java/com/android/server/display/DisplayManagerService.java
index 6928b33..3cab08a 100644
--- a/services/core/java/com/android/server/display/DisplayManagerService.java
+++ b/services/core/java/com/android/server/display/DisplayManagerService.java
@@ -5236,7 +5236,7 @@
                 mHandler.sendMessage(msg);
 
                 mLogicalDisplayMapper
-                        .setDeviceStateLocked(deviceState.getIdentifier());
+                        .setDeviceStateLocked(deviceState);
             }
         }
     };
diff --git a/services/core/java/com/android/server/display/LogicalDisplayMapper.java b/services/core/java/com/android/server/display/LogicalDisplayMapper.java
index e645e98..71dbce00 100644
--- a/services/core/java/com/android/server/display/LogicalDisplayMapper.java
+++ b/services/core/java/com/android/server/display/LogicalDisplayMapper.java
@@ -16,12 +16,17 @@
 
 package com.android.server.display;
 
+import static android.hardware.devicestate.DeviceState.PROPERTY_POWER_CONFIGURATION_TRIGGER_SLEEP;
+import static android.hardware.devicestate.DeviceState.PROPERTY_POWER_CONFIGURATION_TRIGGER_WAKE;
+import static android.hardware.devicestate.DeviceStateManager.INVALID_DEVICE_STATE;
 import static android.view.Display.DEFAULT_DISPLAY;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.content.Context;
-import android.hardware.devicestate.DeviceStateManager;
+import android.hardware.devicestate.DeviceState;
+import android.hardware.devicestate.feature.flags.FeatureFlags;
+import android.hardware.devicestate.feature.flags.FeatureFlagsImpl;
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
@@ -198,14 +203,14 @@
     private final DisplayIdProducer mIdProducer = (isDefault) ->
             isDefault ? DEFAULT_DISPLAY : sNextNonDefaultDisplayId++;
     private Layout mCurrentLayout = null;
-    private int mDeviceState = DeviceStateManager.INVALID_DEVICE_STATE_IDENTIFIER;
-    private int mPendingDeviceState = DeviceStateManager.INVALID_DEVICE_STATE_IDENTIFIER;
-    private int mDeviceStateToBeAppliedAfterBoot =
-            DeviceStateManager.INVALID_DEVICE_STATE_IDENTIFIER;
+    private DeviceState mDeviceState = INVALID_DEVICE_STATE;
+    private DeviceState mPendingDeviceState = INVALID_DEVICE_STATE;
+    private DeviceState mDeviceStateToBeAppliedAfterBoot = INVALID_DEVICE_STATE;
     private boolean mBootCompleted = false;
     private boolean mInteractive;
     private final DisplayManagerFlags mFlags;
     private final SyntheticModeManager mSyntheticModeManager;
+    private final FeatureFlags mDeviceStateManagerFlags;
 
     LogicalDisplayMapper(@NonNull Context context, FoldSettingProvider foldSettingProvider,
             FoldGracePeriodProvider foldGracePeriodProvider,
@@ -245,6 +250,7 @@
         mDeviceStateToLayoutMap = deviceStateToLayoutMap;
         mFlags = flags;
         mSyntheticModeManager = syntheticModeManager;
+        mDeviceStateManagerFlags = new FeatureFlagsImpl();
     }
 
     @Override
@@ -399,8 +405,8 @@
         // Retrieve the display info for the display that matches the display id.
         final DisplayDevice device = mDisplayDeviceRepo.getByAddressLocked(display.getAddress());
         if (device == null) {
-            Slog.w(TAG, "The display device (" + display.getAddress() + "), is not available"
-                    + " for the display state " + mDeviceState);
+            Slog.w(TAG, "The display device (" + display.getAddress()
+                    + "), is not available for the display state " + mDeviceState.getIdentifier());
             return null;
         }
         LogicalDisplay logicalDisplay = getDisplayLocked(device, /* includeDisabled= */ true);
@@ -427,9 +433,11 @@
         ipw.println("mBootCompleted=" + mBootCompleted);
 
         ipw.println();
-        ipw.println("mDeviceState=" + mDeviceState);
-        ipw.println("mPendingDeviceState=" + mPendingDeviceState);
-        ipw.println("mDeviceStateToBeAppliedAfterBoot=" + mDeviceStateToBeAppliedAfterBoot);
+
+        ipw.println("mDeviceState=" + mDeviceState.getIdentifier());
+        ipw.println("mPendingDeviceState=" + mPendingDeviceState.getIdentifier());
+        ipw.println("mDeviceStateToBeAppliedAfterBoot="
+                + mDeviceStateToBeAppliedAfterBoot.getIdentifier());
 
         final int logicalDisplayCount = mLogicalDisplays.size();
         ipw.println();
@@ -459,7 +467,7 @@
         mVirtualDeviceDisplayMapping.put(displayDevice.getUniqueId(), virtualDeviceUniqueId);
     }
 
-    void setDeviceStateLocked(int state) {
+    void setDeviceStateLocked(DeviceState state) {
         if (!mBootCompleted) {
             // The boot animation might still be in progress, we do not want to switch states now
             // as the boot animation would end up with an incorrect size.
@@ -471,15 +479,17 @@
             return;
         }
 
-        Slog.i(TAG, "Requesting Transition to state: " + state + ", from state=" + mDeviceState
-                + ", interactive=" + mInteractive + ", mBootCompleted=" + mBootCompleted);
+        Slog.i(TAG, "Requesting Transition to state: " + state + ", from state="
+                + mDeviceState.getIdentifier() + ", interactive=" + mInteractive
+                + ", mBootCompleted=" + mBootCompleted);
         // As part of a state transition, we may need to turn off some displays temporarily so that
         // the transition is smooth. Plus, on some devices, only one internal displays can be
         // on at a time. We use LogicalDisplay.setIsInTransition to mark a display that needs to be
         // temporarily turned off.
-        resetLayoutLocked(mDeviceState, state, /* transitionValue= */ true);
+        resetLayoutLocked(mDeviceState.getIdentifier(),
+                state.getIdentifier(), /* transitionValue= */ true);
         mPendingDeviceState = state;
-        mDeviceStateToBeAppliedAfterBoot = DeviceStateManager.INVALID_DEVICE_STATE_IDENTIFIER;
+        mDeviceStateToBeAppliedAfterBoot = INVALID_DEVICE_STATE;
         final boolean wakeDevice = shouldDeviceBeWoken(mPendingDeviceState, mDeviceState,
                 mInteractive, mBootCompleted);
         final boolean sleepDevice = shouldDeviceBePutToSleep(mPendingDeviceState, mDeviceState,
@@ -494,7 +504,7 @@
         }
 
         if (DEBUG) {
-            Slog.d(TAG, "Postponing transition to state: " + mPendingDeviceState);
+            Slog.d(TAG, "Postponing transition to state: " + mPendingDeviceState.getIdentifier());
         }
         // Send the transitioning phase updates to DisplayManager so that the displays can
         // start turning OFF in preparation for the new layout.
@@ -529,8 +539,7 @@
     void onBootCompleted() {
         synchronized (mSyncRoot) {
             mBootCompleted = true;
-            if (mDeviceStateToBeAppliedAfterBoot
-                    != DeviceStateManager.INVALID_DEVICE_STATE_IDENTIFIER) {
+            if (!mDeviceStateToBeAppliedAfterBoot.equals(INVALID_DEVICE_STATE)) {
                 setDeviceStateLocked(mDeviceStateToBeAppliedAfterBoot);
             }
         }
@@ -559,11 +568,18 @@
      * @see #setDeviceStateLocked
      */
     @VisibleForTesting
-    boolean shouldDeviceBeWoken(int pendingState, int currentState, boolean isInteractive,
-            boolean isBootCompleted) {
-        return mDeviceStatesOnWhichToWakeUp.get(pendingState)
-                && !mDeviceStatesOnWhichToWakeUp.get(currentState)
-                && !isInteractive && isBootCompleted;
+    boolean shouldDeviceBeWoken(DeviceState pendingState, DeviceState currentState,
+            boolean isInteractive, boolean isBootCompleted) {
+        if (mDeviceStateManagerFlags.deviceStatePropertyMigration()) {
+            return pendingState.hasProperty(PROPERTY_POWER_CONFIGURATION_TRIGGER_WAKE)
+                    && !currentState.equals(INVALID_DEVICE_STATE)
+                    && !currentState.hasProperty(PROPERTY_POWER_CONFIGURATION_TRIGGER_WAKE)
+                    && !isInteractive && isBootCompleted;
+        } else {
+            return mDeviceStatesOnWhichToWakeUp.get(pendingState.getIdentifier())
+                    && !mDeviceStatesOnWhichToWakeUp.get(currentState.getIdentifier())
+                    && !isInteractive && isBootCompleted;
+        }
     }
 
     /**
@@ -584,14 +600,22 @@
      * @see #setDeviceStateLocked
      */
     @VisibleForTesting
-    boolean shouldDeviceBePutToSleep(int pendingState, int currentState, boolean isInteractive,
-            boolean isBootCompleted) {
-        return currentState != DeviceStateManager.INVALID_DEVICE_STATE_IDENTIFIER
-                && mDeviceStatesOnWhichToSelectiveSleep.get(pendingState)
-                && !mDeviceStatesOnWhichToSelectiveSleep.get(currentState)
-                && isInteractive
-                && isBootCompleted
-                && !mFoldSettingProvider.shouldStayAwakeOnFold();
+    boolean shouldDeviceBePutToSleep(DeviceState pendingState, DeviceState currentState,
+            boolean isInteractive, boolean isBootCompleted) {
+        if (mDeviceStateManagerFlags.deviceStatePropertyMigration()) {
+            return pendingState.hasProperty(PROPERTY_POWER_CONFIGURATION_TRIGGER_SLEEP)
+                    && !currentState.equals(INVALID_DEVICE_STATE)
+                    && !currentState.hasProperty(PROPERTY_POWER_CONFIGURATION_TRIGGER_SLEEP)
+                    && isInteractive
+                    && isBootCompleted
+                    && !mFoldSettingProvider.shouldStayAwakeOnFold();
+        } else {
+            return mDeviceStatesOnWhichToSelectiveSleep.get(pendingState.getIdentifier())
+                    && !mDeviceStatesOnWhichToSelectiveSleep.get(currentState.getIdentifier())
+                    && isInteractive
+                    && isBootCompleted
+                    && !mFoldSettingProvider.shouldStayAwakeOnFold();
+        }
     }
 
     private boolean areAllTransitioningDisplaysOffLocked() {
@@ -614,27 +638,25 @@
     }
 
     private void transitionToPendingStateLocked() {
-        resetLayoutLocked(mDeviceState, mPendingDeviceState, /* transitionValue= */ false);
+        resetLayoutLocked(mDeviceState.getIdentifier(),
+                mPendingDeviceState.getIdentifier(), /* transitionValue= */ false);
         mDeviceState = mPendingDeviceState;
-        mPendingDeviceState = DeviceStateManager.INVALID_DEVICE_STATE_IDENTIFIER;
+        mPendingDeviceState = INVALID_DEVICE_STATE;
         applyLayoutLocked();
         updateLogicalDisplaysLocked();
     }
 
     private void finishStateTransitionLocked(boolean force) {
-        if (mPendingDeviceState == DeviceStateManager.INVALID_DEVICE_STATE_IDENTIFIER) {
+        if (mPendingDeviceState.equals(INVALID_DEVICE_STATE)) {
             return;
         }
 
-        final boolean waitingToWakeDevice = mDeviceStatesOnWhichToWakeUp.get(mPendingDeviceState)
-                && !mDeviceStatesOnWhichToWakeUp.get(mDeviceState)
-                && !mInteractive && mBootCompleted;
+        final boolean waitingToWakeDevice = shouldDeviceBeWoken(mPendingDeviceState, 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 waitingToSleepDevice = shouldDeviceBePutToSleep(mPendingDeviceState,
+                mDeviceState, mInteractive, mBootCompleted) && !shouldStayAwakeOnFold();
 
         final boolean displaysOff = areAllTransitioningDisplaysOffLocked();
         final boolean isReadyToTransition = displaysOff && !waitingToWakeDevice
@@ -1100,7 +1122,7 @@
      */
     private void applyLayoutLocked() {
         final Layout oldLayout = mCurrentLayout;
-        mCurrentLayout = mDeviceStateToLayoutMap.get(mDeviceState);
+        mCurrentLayout = mDeviceStateToLayoutMap.get(mDeviceState.getIdentifier());
         Slog.i(TAG, "Applying layout: " + mCurrentLayout + ", Previous layout: " + oldLayout);
 
         // Go through each of the displays in the current layout set.
@@ -1116,7 +1138,7 @@
             final DisplayDevice device = mDisplayDeviceRepo.getByAddressLocked(address);
             if (device == null) {
                 Slog.w(TAG, "applyLayoutLocked: The display device (" + address + "), is not "
-                        + "available for the display state " + mDeviceState);
+                        + "available for the display state " + mDeviceState.getIdentifier());
                 continue;
             }
 
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 6d138c5..1729ad5 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/LogicalDisplayMapperTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/LogicalDisplayMapperTest.java
@@ -16,7 +16,7 @@
 
 package com.android.server.display;
 
-import static android.hardware.devicestate.DeviceStateManager.INVALID_DEVICE_STATE_IDENTIFIER;
+import static android.hardware.devicestate.DeviceStateManager.INVALID_DEVICE_STATE;
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.Display.DEFAULT_DISPLAY_GROUP;
 import static android.view.Display.FLAG_REAR;
@@ -62,9 +62,11 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.annotation.NonNull;
 import android.app.PropertyInvalidatedCache;
 import android.content.Context;
 import android.content.res.Resources;
+import android.hardware.devicestate.DeviceState;
 import android.os.Handler;
 import android.os.IPowerManager;
 import android.os.IThermalService;
@@ -103,7 +105,9 @@
 import java.io.OutputStream;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
+import java.util.Set;
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
@@ -111,9 +115,12 @@
     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 DeviceState DEVICE_STATE_CLOSED = createDeviceState(0, "Zero",
+            Set.of(DeviceState.PROPERTY_POWER_CONFIGURATION_TRIGGER_SLEEP), Collections.emptySet());
+    private static final DeviceState DEVICE_STATE_HALF_OPEN = createDeviceState(1, "One",
+            Set.of(DeviceState.PROPERTY_POWER_CONFIGURATION_TRIGGER_WAKE), Collections.emptySet());
+    private static final DeviceState DEVICE_STATE_OPEN = createDeviceState(2, "Two",
+            Set.of(DeviceState.PROPERTY_POWER_CONFIGURATION_TRIGGER_WAKE), Collections.emptySet());
     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;
@@ -703,8 +710,7 @@
                 /* isInteractive= */true,
                 /* isBootCompleted= */true));
         assertFalse(mLogicalDisplayMapper.shouldDeviceBePutToSleep(DEVICE_STATE_CLOSED,
-                INVALID_DEVICE_STATE_IDENTIFIER,
-                /* isInteractive= */true,
+                INVALID_DEVICE_STATE /* currentState */, /* isInteractive= */true,
                 /* isBootCompleted= */true));
     }
 
@@ -932,7 +938,7 @@
         // We can only have one default display
         assertEquals(DEFAULT_DISPLAY, id(display1));
 
-        mLogicalDisplayMapper.setDeviceStateLocked(0);
+        mLogicalDisplayMapper.setDeviceStateLocked(DEVICE_STATE_CLOSED);
         advanceTime(1000);
         // The new state is not applied until the boot is completed
         assertTrue(mLogicalDisplayMapper.getDisplayLocked(device1).isEnabledLocked());
@@ -953,7 +959,7 @@
         assertEquals("concurrent", mLogicalDisplayMapper.getDisplayLocked(device2)
                 .getDisplayInfoLocked().thermalBrightnessThrottlingDataId);
 
-        mLogicalDisplayMapper.setDeviceStateLocked(1);
+        mLogicalDisplayMapper.setDeviceStateLocked(DEVICE_STATE_HALF_OPEN);
         advanceTime(1000);
         assertFalse(mLogicalDisplayMapper.getDisplayLocked(device1).isEnabledLocked());
         assertTrue(mLogicalDisplayMapper.getDisplayLocked(device2).isEnabledLocked());
@@ -966,7 +972,7 @@
                 mLogicalDisplayMapper.getDisplayLocked(device2)
                         .getDisplayInfoLocked().thermalBrightnessThrottlingDataId);
 
-        mLogicalDisplayMapper.setDeviceStateLocked(2);
+        mLogicalDisplayMapper.setDeviceStateLocked(DEVICE_STATE_OPEN);
         advanceTime(1000);
         assertFalse(mLogicalDisplayMapper.getDisplayLocked(device1).isEnabledLocked());
         assertTrue(mLogicalDisplayMapper.getDisplayLocked(device2).isEnabledLocked());
@@ -1043,7 +1049,7 @@
         // 3) Send DISPLAY_DEVICE_EVENT_CHANGE to inform the mapper of the new display state
         // 4) Dispatch handler events.
         mLogicalDisplayMapper.onBootCompleted();
-        mLogicalDisplayMapper.setDeviceStateLocked(0);
+        mLogicalDisplayMapper.setDeviceStateLocked(DEVICE_STATE_CLOSED);
         mDisplayDeviceRepo.onDisplayDeviceEvent(device3, DISPLAY_DEVICE_EVENT_CHANGED);
         advanceTime(1000);
         final int[] allDisplayIds = mLogicalDisplayMapper.getDisplayIdsLocked(
@@ -1073,7 +1079,7 @@
                 /* includeDisabled= */ false));
 
         // Now do it again to go back to state 1
-        mLogicalDisplayMapper.setDeviceStateLocked(1);
+        mLogicalDisplayMapper.setDeviceStateLocked(DEVICE_STATE_HALF_OPEN);
         mDisplayDeviceRepo.onDisplayDeviceEvent(device3, DISPLAY_DEVICE_EVENT_CHANGED);
         advanceTime(1000);
         final int[] threeDisplaysEnabled = mLogicalDisplayMapper.getDisplayIdsLocked(
@@ -1127,7 +1133,7 @@
         // We can only have one default display
         assertEquals(DEFAULT_DISPLAY, id(display1));
 
-        mLogicalDisplayMapper.setDeviceStateLocked(0);
+        mLogicalDisplayMapper.setDeviceStateLocked(DEVICE_STATE_CLOSED);
         advanceTime(1000);
         mLogicalDisplayMapper.onBootCompleted();
         advanceTime(1000);
@@ -1180,13 +1186,15 @@
         Layout layout = new Layout();
         createDefaultDisplay(layout, outer);
         createNonDefaultDisplay(layout, inner, /* enabled= */ false, /* group= */ null);
-        when(mDeviceStateToLayoutMapSpy.get(DEVICE_STATE_CLOSED)).thenReturn(layout);
+        when(mDeviceStateToLayoutMapSpy.get(DEVICE_STATE_CLOSED.getIdentifier())).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.get(DEVICE_STATE_HALF_OPEN.getIdentifier())).thenReturn(
+                layout);
+        when(mDeviceStateToLayoutMapSpy.get(DEVICE_STATE_OPEN.getIdentifier())).thenReturn(layout);
         when(mDeviceStateToLayoutMapSpy.size()).thenReturn(4);
 
         add(outer);
@@ -1317,6 +1325,15 @@
         assertNotEquals(DEFAULT_DISPLAY, id(displayRemoved));
     }
 
+    private static DeviceState createDeviceState(int identifier, @NonNull String name,
+            @NonNull Set<@DeviceState.SystemDeviceStateProperties Integer> systemProperties,
+            @NonNull Set<@DeviceState.PhysicalDeviceStateProperties Integer> physicalProperties) {
+        DeviceState.Configuration deviceStateConfiguration = new DeviceState.Configuration.Builder(
+                identifier, name).setSystemProperties(systemProperties).setPhysicalProperties(
+                physicalProperties).build();
+        return new DeviceState(deviceStateConfiguration);
+    }
+
     private final static class FoldableDisplayDevices {
         final TestDisplayDevice mOuter;
         final TestDisplayDevice mInner;