diff --git a/src/com/android/settings/notification/ShowOnLockScreenNotificationPreferenceController.java b/src/com/android/settings/notification/ShowOnLockScreenNotificationPreferenceController.java
index df1f4a21eb0..f8f76caecbd 100644
--- a/src/com/android/settings/notification/ShowOnLockScreenNotificationPreferenceController.java
+++ b/src/com/android/settings/notification/ShowOnLockScreenNotificationPreferenceController.java
@@ -35,6 +35,7 @@ import com.google.common.annotations.VisibleForTesting;
 import java.util.ArrayList;
 
 import androidx.preference.Preference;
+import androidx.preference.PreferenceScreen;
 
 public class ShowOnLockScreenNotificationPreferenceController extends AbstractPreferenceController
         implements PreferenceControllerMixin, Preference.OnPreferenceChangeListener {
@@ -66,8 +67,9 @@ public class ShowOnLockScreenNotificationPreferenceController extends AbstractPr
     }
 
     @Override
-    public void updateState(Preference preference) {
-        RestrictedListPreference pref = (RestrictedListPreference) preference;
+    public void displayPreference(PreferenceScreen screen) {
+        super.displayPreference(screen);
+        RestrictedListPreference pref = screen.findPreference(mSettingKey);
         pref.clearRestrictedItems();
         ArrayList<CharSequence> entries = new ArrayList<>();
         ArrayList<CharSequence> values = new ArrayList<>();
@@ -104,7 +106,7 @@ public class ShowOnLockScreenNotificationPreferenceController extends AbstractPr
 
         pref.setOnPreferenceChangeListener(this);
 
-        refreshSummary(preference);
+        refreshSummary(pref);
     }
 
     @Override
diff --git a/tests/robotests/src/com/android/settings/notification/ShowOnLockscreenNotificationPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/notification/ShowOnLockscreenNotificationPreferenceControllerTest.java
index a341def3ed5..496897f87a9 100644
--- a/tests/robotests/src/com/android/settings/notification/ShowOnLockscreenNotificationPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/notification/ShowOnLockscreenNotificationPreferenceControllerTest.java
@@ -89,7 +89,7 @@ public class ShowOnLockscreenNotificationPreferenceControllerTest {
                 LOCK_SCREEN_SHOW_SILENT_NOTIFICATIONS,
                 1);
 
-        mController.updateState(mPreference);
+        mController.displayPreference(mScreen);
 
         assertThat(mPreference.getValue()).isEqualTo(
                 String.valueOf(R.string.lock_screen_notifs_show_none));
@@ -107,7 +107,7 @@ public class ShowOnLockscreenNotificationPreferenceControllerTest {
                 LOCK_SCREEN_SHOW_SILENT_NOTIFICATIONS,
                 0);
 
-        mController.updateState(mPreference);
+        mController.displayPreference(mScreen);
 
         assertThat(mPreference.getValue()).isEqualTo(
                 String.valueOf(R.string.lock_screen_notifs_show_alerting));
@@ -124,7 +124,7 @@ public class ShowOnLockscreenNotificationPreferenceControllerTest {
                 LOCK_SCREEN_SHOW_SILENT_NOTIFICATIONS,
                 1);
 
-        mController.updateState(mPreference);
+        mController.displayPreference(mScreen);
 
         assertThat(mPreference.getValue()).isEqualTo(
                 String.valueOf(R.string.lock_screen_notifs_show_all));
@@ -136,7 +136,7 @@ public class ShowOnLockscreenNotificationPreferenceControllerTest {
     public void updateState_allNotifsOnLockscreen_isDefault() {
         // settings don't exist
 
-        mController.updateState(mPreference);
+        mController.displayPreference(mScreen);
 
         assertThat(mPreference.getValue()).isEqualTo(
                 String.valueOf(R.string.lock_screen_notifs_show_all));
@@ -158,7 +158,7 @@ public class ShowOnLockscreenNotificationPreferenceControllerTest {
         ShadowRestrictedLockUtilsInternal.setKeyguardDisabledFeatures(
                 KEYGUARD_DISABLE_SECURE_NOTIFICATIONS);
 
-        mController.updateState(mPreference);
+        mController.displayPreference(mScreen);
 
         assertThat(mPreference.getValue()).isEqualTo(
                 String.valueOf(R.string.lock_screen_notifs_show_none));