diff --git a/tests/robotests/src/com/android/settings/display/SmartAutoRotatePreferenceFragmentTest.java b/tests/robotests/src/com/android/settings/display/SmartAutoRotatePreferenceFragmentTest.java
index bdff55e6200..3fa4582b7d6 100644
--- a/tests/robotests/src/com/android/settings/display/SmartAutoRotatePreferenceFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/display/SmartAutoRotatePreferenceFragmentTest.java
@@ -16,6 +16,8 @@
 
 package com.android.settings.display;
 
+import static android.provider.Settings.Secure.DEVICE_STATE_ROTATION_LOCK_LOCKED;
+
 import static com.android.settings.display.SmartAutoRotatePreferenceFragment.AUTO_ROTATE_MAIN_SWITCH_PREFERENCE_KEY;
 import static com.android.settings.display.SmartAutoRotatePreferenceFragment.AUTO_ROTATE_SWITCH_PREFERENCE_KEY;
 
@@ -68,6 +70,11 @@ import java.util.List;
 })
 public class SmartAutoRotatePreferenceFragmentTest {
 
+    private static final int STATE_FOLDED = 0;
+    private static final int STATE_HALF_FOLDED = 1;
+    private static final int STATE_UNFOLDED = 2;
+    private static final int STATE_REAR_DISPLAY = 3;
+
     private static final String PACKAGE_NAME = "package_name";
 
     private SmartAutoRotatePreferenceFragment mFragment;
@@ -97,6 +104,7 @@ public class SmartAutoRotatePreferenceFragmentTest {
         ContentResolver mContentResolver = RuntimeEnvironment.application.getContentResolver();
         when(mContext.getPackageManager()).thenReturn(mPackageManager);
         when(mContext.getContentResolver()).thenReturn(mContentResolver);
+        when(mContext.getApplicationContext()).thenReturn(mContext);
         doReturn(PACKAGE_NAME).when(mPackageManager).getRotationResolverPackageName();
         doReturn(PackageManager.PERMISSION_GRANTED).when(mPackageManager).checkPermission(
                 Manifest.permission.CAMERA, PACKAGE_NAME);
@@ -119,12 +127,19 @@ public class SmartAutoRotatePreferenceFragmentTest {
         when(mFragment.findPreference(AUTO_ROTATE_SWITCH_PREFERENCE_KEY)).thenReturn(
                 mRotateSwitchPreference);
 
+        when(mFragment.findPreference(AUTO_ROTATE_MAIN_SWITCH_PREFERENCE_KEY))
+                .thenReturn(mRotateMainSwitchPreference);
 
-        when(mFragment.findPreference(AUTO_ROTATE_MAIN_SWITCH_PREFERENCE_KEY)).thenReturn(
-                mRotateMainSwitchPreference);
+        when(mResources.getIntArray(com.android.internal.R.array.config_foldedDeviceStates))
+                .thenReturn(new int[] {STATE_FOLDED});
+        when(mResources.getIntArray(com.android.internal.R.array.config_halfFoldedDeviceStates))
+                .thenReturn(new int[] {STATE_HALF_FOLDED});
+        when(mResources.getIntArray(com.android.internal.R.array.config_openDeviceStates))
+                .thenReturn(new int[] {STATE_UNFOLDED});
+        when(mResources.getIntArray(com.android.internal.R.array.config_rearDisplayDeviceStates))
+                .thenReturn(new int[] {STATE_REAR_DISPLAY});
     }
 
-
     @Test
     public void createHeader_faceDetectionSupported_switchBarIsEnabled() {
         ShadowDeviceStateRotationLockSettingsManager.setDeviceStateRotationLockEnabled(false);
@@ -168,7 +183,7 @@ public class SmartAutoRotatePreferenceFragmentTest {
 
     @Test
     public void createPreferenceControllers_noSettableDeviceStates_returnsEmptyList() {
-        enableDeviceStateSettableRotationStates(new String[]{}, new String[]{});
+        enableDeviceStateSettableRotationStates(new String[] {}, new String[] {});
 
         List<AbstractPreferenceController> preferenceControllers =
                 mFragment.createPreferenceControllers(mContext);
@@ -178,17 +193,21 @@ public class SmartAutoRotatePreferenceFragmentTest {
 
     @Test
     public void createPreferenceControllers_settableDeviceStates_returnsDeviceStateControllers() {
-        enableDeviceStateSettableRotationStates(new String[]{"0:1", "1:1"},
-                new String[]{"Folded", "Unfolded"});
+        enableDeviceStateSettableRotationStates(
+                new String[] {
+                    STATE_FOLDED + ":" + DEVICE_STATE_ROTATION_LOCK_LOCKED,
+                    STATE_UNFOLDED + ":" + DEVICE_STATE_ROTATION_LOCK_LOCKED
+                },
+                new String[] {"Folded", "Unfolded"});
 
         List<AbstractPreferenceController> preferenceControllers =
                 mFragment.createPreferenceControllers(mContext);
 
         assertThat(preferenceControllers).hasSize(2);
-        assertThat(preferenceControllers.get(0)).isInstanceOf(
-                DeviceStateAutoRotateSettingController.class);
-        assertThat(preferenceControllers.get(1)).isInstanceOf(
-                DeviceStateAutoRotateSettingController.class);
+        assertThat(preferenceControllers.get(0))
+                .isInstanceOf(DeviceStateAutoRotateSettingController.class);
+        assertThat(preferenceControllers.get(1))
+                .isInstanceOf(DeviceStateAutoRotateSettingController.class);
     }
 
     @Test
@@ -204,16 +223,15 @@ public class SmartAutoRotatePreferenceFragmentTest {
         verify(mFragment, times(1)).addHelpLink();
     }
 
-    private void enableDeviceStateSettableRotationStates(String[] settableStates,
-            String[] settableStatesDescriptions) {
+    private void enableDeviceStateSettableRotationStates(
+            String[] settableStates, String[] settableStatesDescriptions) {
         when(mResources.getStringArray(
-                com.android.internal.R.array.config_perDeviceStateRotationLockDefaults)).thenReturn(
-                settableStates);
-        when(mResources.getStringArray(
-                R.array.config_settableAutoRotationDeviceStatesDescriptions)).thenReturn(
-                settableStatesDescriptions);
-        when(mResources.getBoolean(R.bool.config_auto_rotate_face_detection_available)).thenReturn(
-                true);
+                        com.android.internal.R.array.config_perDeviceStateRotationLockDefaults))
+                .thenReturn(settableStates);
+        when(mResources.getStringArray(R.array.config_settableAutoRotationDeviceStatesDescriptions))
+                .thenReturn(settableStatesDescriptions);
+        when(mResources.getBoolean(R.bool.config_auto_rotate_face_detection_available))
+                .thenReturn(true);
         DeviceStateRotationLockSettingsManager.resetInstance();
         DeviceStateRotationLockSettingsManager.getInstance(mContext)
                 .resetStateForTesting(mResources);