diff --git a/tests/robotests/src/com/android/settings/display/DeviceStateAutoRotateDetailsFragmentTest.java b/tests/robotests/src/com/android/settings/display/DeviceStateAutoRotateDetailsFragmentTest.java
index c1268b6..e5940b6 100644
--- a/tests/robotests/src/com/android/settings/display/DeviceStateAutoRotateDetailsFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/display/DeviceStateAutoRotateDetailsFragmentTest.java
@@ -16,6 +16,8 @@
 
 package com.android.settings.display;
 
+import static android.provider.Settings.Secure.DEVICE_STATE_ROTATION_LOCK_LOCKED;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.Mockito.spy;
@@ -39,6 +41,10 @@
 
 @RunWith(RobolectricTestRunner.class)
 public class DeviceStateAutoRotateDetailsFragmentTest {
+    private static final int FOLDED_STATE = 0;
+    private static final int HALF_FOLDED_STATE = 1;
+    private static final int UNFOLDED_STATE = 2;
+    private static final int REAR_DISPLAY_STATE = 3;
 
     private final DeviceStateAutoRotateDetailsFragment mFragment =
             spy(new DeviceStateAutoRotateDetailsFragment());
@@ -51,6 +57,7 @@
         when(mContext.getApplicationContext()).thenReturn(mContext);
         when(mFragment.getContext()).thenReturn(mContext);
         when(mFragment.getResources()).thenReturn(mResources);
+        setUpPostureMappings();
     }
 
     @Test
@@ -67,7 +74,9 @@
 
     @Test
     public void createPreferenceControllers_settableDeviceStates_returnsDeviceStateControllers() {
-        enableDeviceStateSettableRotationStates(new String[]{"0:1", "1:1"},
+        enableDeviceStateSettableRotationStates(
+                new String[]{FOLDED_STATE + ":" + DEVICE_STATE_ROTATION_LOCK_LOCKED,
+                        UNFOLDED_STATE + ":" + DEVICE_STATE_ROTATION_LOCK_LOCKED},
                 new String[]{"Folded", "Unfolded"});
 
         List<AbstractPreferenceController> preferenceControllers =
@@ -102,4 +111,19 @@
         DeviceStateRotationLockSettingsManager.getInstance(mContext)
                 .resetStateForTesting(mResources);
     }
+
+    private void setUpPostureMappings() {
+        when(mResources.getIntArray(
+                com.android.internal.R.array.config_foldedDeviceStates)).thenReturn(
+                new int[]{FOLDED_STATE});
+        when(mResources.getIntArray(
+                com.android.internal.R.array.config_halfFoldedDeviceStates)).thenReturn(
+                new int[]{HALF_FOLDED_STATE});
+        when(mResources.getIntArray(
+                com.android.internal.R.array.config_openDeviceStates)).thenReturn(
+                new int[]{UNFOLDED_STATE});
+        when(mResources.getIntArray(
+                com.android.internal.R.array.config_rearDisplayDeviceStates)).thenReturn(
+                new int[]{REAR_DISPLAY_STATE});
+    }
 }
diff --git a/tests/robotests/src/com/android/settings/display/SmartAutoRotateControllerTest.java b/tests/robotests/src/com/android/settings/display/SmartAutoRotateControllerTest.java
index 4fec38b..4596364 100644
--- a/tests/robotests/src/com/android/settings/display/SmartAutoRotateControllerTest.java
+++ b/tests/robotests/src/com/android/settings/display/SmartAutoRotateControllerTest.java
@@ -53,6 +53,7 @@
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
+import org.robolectric.shadow.api.Shadow;
 
 @RunWith(RobolectricTestRunner.class)
 @Config(shadows = ShadowSensorPrivacyManager.class)
@@ -171,16 +172,14 @@
     }
 
     private void lockDeviceStateRotation() {
-        mDeviceStateAutoRotateSettingsManager.updateSetting(
-                /* deviceState= */0, /* rotationLocked= */ true);
-        mDeviceStateAutoRotateSettingsManager.updateSetting(
-                /* deviceState= */1, /* rotationLocked= */ true);
+        ShadowDeviceStateRotationLockSettingsManager shadowManager =
+                Shadow.extract(mDeviceStateAutoRotateSettingsManager);
+        shadowManager.setRotationLockedForAllStates(true);
     }
 
     private void unlockDeviceStateRotation() {
-        mDeviceStateAutoRotateSettingsManager.updateSetting(
-                /* deviceState= */0, /* rotationLocked= */ false);
-        mDeviceStateAutoRotateSettingsManager.updateSetting(
-                /* deviceState= */1, /* rotationLocked= */ true);
+        ShadowDeviceStateRotationLockSettingsManager shadowManager =
+                Shadow.extract(mDeviceStateAutoRotateSettingsManager);
+        shadowManager.setRotationLockedForAllStates(false);
     }
 }
diff --git a/tests/robotests/src/com/android/settings/testutils/shadow/ShadowDeviceStateRotationLockSettingsManager.java b/tests/robotests/src/com/android/settings/testutils/shadow/ShadowDeviceStateRotationLockSettingsManager.java
index 72df3cc..ed266e3 100644
--- a/tests/robotests/src/com/android/settings/testutils/shadow/ShadowDeviceStateRotationLockSettingsManager.java
+++ b/tests/robotests/src/com/android/settings/testutils/shadow/ShadowDeviceStateRotationLockSettingsManager.java
@@ -27,6 +27,7 @@
 public class ShadowDeviceStateRotationLockSettingsManager {
 
     private static boolean sDeviceStateRotationLockEnabled;
+    private boolean mIsRotationLockedForAllStates;
 
     @Implementation
     public static boolean isDeviceStateRotationLockEnabled(Context context) {
@@ -36,4 +37,13 @@
     public static void setDeviceStateRotationLockEnabled(boolean enabled) {
         sDeviceStateRotationLockEnabled = enabled;
     }
+
+    @Implementation
+    public boolean isRotationLockedForAllStates() {
+        return mIsRotationLockedForAllStates;
+    }
+
+    public void setRotationLockedForAllStates(boolean value) {
+        mIsRotationLockedForAllStates = value;
+    }
 }