diff --git a/tests/robotests/Android.bp b/tests/robotests/Android.bp
index 5648290..1c6794d 100644
--- a/tests/robotests/Android.bp
+++ b/tests/robotests/Android.bp
@@ -71,10 +71,12 @@
         "Settings-testutils2",
         "notification_flags_lib",
         "com_android_server_accessibility_flags_lib",
+        "testables",
     ],
 
     libs: [
         "ims-common",
+        "android.test.mock",
     ],
 
     java_resource_dirs: [
diff --git a/tests/robotests/src/com/android/settings/development/ForcePeakRefreshRatePreferenceControllerTest.java b/tests/robotests/src/com/android/settings/development/ForcePeakRefreshRatePreferenceControllerTest.java
index 314120a..c7c76b1 100644
--- a/tests/robotests/src/com/android/settings/development/ForcePeakRefreshRatePreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/development/ForcePeakRefreshRatePreferenceControllerTest.java
@@ -16,6 +16,8 @@
 
 package com.android.settings.development;
 
+import static android.hardware.display.DisplayManager.DISPLAY_CATEGORY_ALL_INCLUDING_DISABLED;
+
 import static com.android.internal.display.RefreshRateSettingsUtils.DEFAULT_REFRESH_RATE;
 import static com.android.settings.development.ForcePeakRefreshRatePreferenceController.NO_CONFIG;
 
@@ -24,15 +26,18 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import android.content.Context;
+import android.hardware.display.DisplayManager;
 import android.platform.test.annotations.RequiresFlagsDisabled;
 import android.platform.test.annotations.RequiresFlagsEnabled;
 import android.platform.test.flag.junit.CheckFlagsRule;
 import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.provider.Settings;
+import android.testing.TestableContext;
+import android.view.Display;
 
 import androidx.preference.PreferenceScreen;
 import androidx.preference.SwitchPreference;
+import androidx.test.platform.app.InstrumentationRegistry;
 
 import com.android.server.display.feature.flags.Flags;
 
@@ -43,7 +48,6 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
 
 @RunWith(RobolectricTestRunner.class)
@@ -53,17 +57,51 @@
     private SwitchPreference mPreference;
     @Mock
     private PreferenceScreen mScreen;
+    @Mock
+    private DisplayManager mDisplayManagerMock;
+    @Mock
+    private Display mDisplayMock;
+    @Mock
+    private Display mDisplayMock2;
 
-    private Context mContext;
     private ForcePeakRefreshRatePreferenceController mController;
 
     @Rule
     public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
+    @Rule
+    public final TestableContext mContext = new TestableContext(
+            InstrumentationRegistry.getInstrumentation().getContext());
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mContext = RuntimeEnvironment.application;
+        mContext.addMockSystemService(DisplayManager.class, mDisplayManagerMock);
+
+        Display.Mode[] modes = new Display.Mode[]{
+                new Display.Mode(/* modeId= */ 0, /* width= */ 800, /* height= */ 600,
+                        /* refreshRate= */ 60),
+                new Display.Mode(/* modeId= */ 0, /* width= */ 800, /* height= */ 600,
+                        /* refreshRate= */ 120),
+                new Display.Mode(/* modeId= */ 0, /* width= */ 800, /* height= */ 600,
+                        /* refreshRate= */ 90)
+        };
+        when(mDisplayManagerMock.getDisplay(Display.DEFAULT_DISPLAY)).thenReturn(mDisplayMock);
+        when(mDisplayMock.getSupportedModes()).thenReturn(modes);
+
+        Display.Mode[] modes2 = new Display.Mode[]{
+                new Display.Mode(/* modeId= */ 0, /* width= */ 800, /* height= */ 600,
+                        /* refreshRate= */ 70),
+                new Display.Mode(/* modeId= */ 0, /* width= */ 800, /* height= */ 600,
+                        /* refreshRate= */ 130),
+                new Display.Mode(/* modeId= */ 0, /* width= */ 800, /* height= */ 600,
+                        /* refreshRate= */ 80)
+        };
+        when(mDisplayManagerMock.getDisplay(Display.DEFAULT_DISPLAY + 1)).thenReturn(mDisplayMock2);
+        when(mDisplayMock2.getSupportedModes()).thenReturn(modes2);
+
+        when(mDisplayManagerMock.getDisplays(DISPLAY_CATEGORY_ALL_INCLUDING_DISABLED))
+                .thenReturn(new Display[]{ mDisplayMock, mDisplayMock2 });
+
         mController = new ForcePeakRefreshRatePreferenceController(mContext);
         when(mScreen.findPreference(mController.getPreferenceKey())).thenReturn(mPreference);
         when(mPreference.getKey()).thenReturn(mController.getPreferenceKey());
@@ -153,4 +191,16 @@
         assertThat(mPreference.isChecked()).isFalse();
         assertThat(mPreference.isEnabled()).isFalse();
     }
+
+    @Test
+    @RequiresFlagsDisabled(Flags.FLAG_BACK_UP_SMOOTH_DISPLAY_AND_FORCE_PEAK_REFRESH_RATE)
+    public void peakRefreshRate_highestOfDefaultDisplay_featureFlagOff() {
+        assertThat(mController.mPeakRefreshRate).isEqualTo(120);
+    }
+
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_BACK_UP_SMOOTH_DISPLAY_AND_FORCE_PEAK_REFRESH_RATE)
+    public void peakRefreshRate_highestOfAllDisplays_featureFlagOn() {
+        assertThat(mController.mPeakRefreshRate).isEqualTo(130);
+    }
 }
diff --git a/tests/robotests/src/com/android/settings/display/PeakRefreshRatePreferenceControllerTest.java b/tests/robotests/src/com/android/settings/display/PeakRefreshRatePreferenceControllerTest.java
index cb0963b..f8e91bd 100644
--- a/tests/robotests/src/com/android/settings/display/PeakRefreshRatePreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/display/PeakRefreshRatePreferenceControllerTest.java
@@ -16,6 +16,8 @@
 
 package com.android.settings.display;
 
+import static android.hardware.display.DisplayManager.DISPLAY_CATEGORY_ALL_INCLUDING_DISABLED;
+
 import static com.android.internal.display.RefreshRateSettingsUtils.DEFAULT_REFRESH_RATE;
 import static com.android.settings.core.BasePreferenceController.AVAILABLE;
 import static com.android.settings.core.BasePreferenceController.UNSUPPORTED_ON_DEVICE;
@@ -24,14 +26,17 @@
 
 import static org.mockito.Mockito.when;
 
-import android.content.Context;
+import android.hardware.display.DisplayManager;
 import android.platform.test.annotations.RequiresFlagsDisabled;
 import android.platform.test.annotations.RequiresFlagsEnabled;
 import android.platform.test.flag.junit.CheckFlagsRule;
 import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.provider.Settings;
+import android.testing.TestableContext;
+import android.view.Display;
 
 import androidx.preference.SwitchPreference;
+import androidx.test.platform.app.InstrumentationRegistry;
 
 import com.android.server.display.feature.flags.Flags;
 
@@ -48,21 +53,55 @@
 @RunWith(RobolectricTestRunner.class)
 public class PeakRefreshRatePreferenceControllerTest {
 
-    private Context mContext;
     private PeakRefreshRatePreferenceController mController;
     private SwitchPreference mPreference;
 
     @Mock
     private PeakRefreshRatePreferenceController.DeviceConfigDisplaySettings
             mDeviceConfigDisplaySettings;
+    @Mock
+    private DisplayManager mDisplayManagerMock;
+    @Mock
+    private Display mDisplayMock;
+    @Mock
+    private Display mDisplayMock2;
 
     @Rule
     public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
+    @Rule
+    public final TestableContext mContext = new TestableContext(
+            InstrumentationRegistry.getInstrumentation().getContext());
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mContext = RuntimeEnvironment.application;
+        mContext.addMockSystemService(DisplayManager.class, mDisplayManagerMock);
+
+        Display.Mode[] modes = new Display.Mode[]{
+                new Display.Mode(/* modeId= */ 0, /* width= */ 800, /* height= */ 600,
+                        /* refreshRate= */ 60),
+                new Display.Mode(/* modeId= */ 0, /* width= */ 800, /* height= */ 600,
+                        /* refreshRate= */ 120),
+                new Display.Mode(/* modeId= */ 0, /* width= */ 800, /* height= */ 600,
+                        /* refreshRate= */ 90)
+        };
+        when(mDisplayManagerMock.getDisplay(Display.DEFAULT_DISPLAY)).thenReturn(mDisplayMock);
+        when(mDisplayMock.getSupportedModes()).thenReturn(modes);
+
+        Display.Mode[] modes2 = new Display.Mode[]{
+                new Display.Mode(/* modeId= */ 0, /* width= */ 800, /* height= */ 600,
+                        /* refreshRate= */ 70),
+                new Display.Mode(/* modeId= */ 0, /* width= */ 800, /* height= */ 600,
+                        /* refreshRate= */ 130),
+                new Display.Mode(/* modeId= */ 0, /* width= */ 800, /* height= */ 600,
+                        /* refreshRate= */ 80)
+        };
+        when(mDisplayManagerMock.getDisplay(Display.DEFAULT_DISPLAY + 1)).thenReturn(mDisplayMock2);
+        when(mDisplayMock2.getSupportedModes()).thenReturn(modes2);
+
+        when(mDisplayManagerMock.getDisplays(DISPLAY_CATEGORY_ALL_INCLUDING_DISABLED))
+                .thenReturn(new Display[]{ mDisplayMock, mDisplayMock2 });
+
         mController = new PeakRefreshRatePreferenceController(mContext, "key");
         mController.injectDeviceConfigDisplaySettings(mDeviceConfigDisplaySettings);
         mPreference = new SwitchPreference(RuntimeEnvironment.application);
@@ -152,4 +191,16 @@
 
         assertThat(mController.isChecked()).isFalse();
     }
+
+    @Test
+    @RequiresFlagsDisabled(Flags.FLAG_BACK_UP_SMOOTH_DISPLAY_AND_FORCE_PEAK_REFRESH_RATE)
+    public void peakRefreshRate_highestOfDefaultDisplay_featureFlagOff() {
+        assertThat(mController.mPeakRefreshRate).isEqualTo(120);
+    }
+
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_BACK_UP_SMOOTH_DISPLAY_AND_FORCE_PEAK_REFRESH_RATE)
+    public void peakRefreshRate_highestOfAllDisplays_featureFlagOn() {
+        assertThat(mController.mPeakRefreshRate).isEqualTo(130);
+    }
 }