diff --git a/res/layout/reset_importance_button.xml b/res/layout/reset_importance_button.xml
deleted file mode 100644
index 6d023b9..0000000
--- a/res/layout/reset_importance_button.xml
+++ /dev/null
@@ -1,33 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!--
-  ~ Copyright (C) 2020 The Android Open Source Project
-  ~
-  ~ Licensed under the Apache License, Version 2.0 (the "License");
-  ~ you may not use this file except in compliance with the License.
-  ~ You may obtain a copy of the License at
-  ~
-  ~      http://www.apache.org/licenses/LICENSE-2.0
-  ~
-  ~ Unless required by applicable law or agreed to in writing, software
-  ~ distributed under the License is distributed on an "AS IS" BASIS,
-  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-  ~ See the License for the specific language governing permissions and
-  ~ limitations under the License.
-  -->
-
-<LinearLayout
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    android:layout_width="match_parent"
-    android:layout_height="wrap_content"
-    android:orientation="vertical">
-
-    <Button
-        android:id="@+id/reset_importance_button"
-        style="@style/ActionPrimaryButton"
-        android:layout_marginStart="@dimen/screen_margin_sides"
-        android:text="@string/asst_importance_reset_title"
-        android:layout_gravity="start"
-        android:layout_width="wrap_content"
-        android:layout_height="wrap_content"/>
-
-</LinearLayout>
\ No newline at end of file
diff --git a/res/xml/development_settings.xml b/res/xml/development_settings.xml
index 19c1209..b1b9017 100644
--- a/res/xml/development_settings.xml
+++ b/res/xml/development_settings.xml
@@ -609,11 +609,9 @@
             android:title="@string/show_notification_channel_warnings"
             android:summary="@string/show_notification_channel_warnings_summary" />
 
-        <com.android.settingslib.widget.LayoutPreference
+        <Preference
             android:key="asst_importance_reset"
             android:title="@string/asst_importance_reset_title"
-            android:selectable="false"
-            android:layout="@layout/reset_importance_button"
             settings:controller="com.android.settings.notification.ImportanceResetPreferenceController" />
 
         <Preference
diff --git a/src/com/android/settings/notification/ImportanceResetPreferenceController.java b/src/com/android/settings/notification/ImportanceResetPreferenceController.java
index 3c66f1a..2d03f23 100644
--- a/src/com/android/settings/notification/ImportanceResetPreferenceController.java
+++ b/src/com/android/settings/notification/ImportanceResetPreferenceController.java
@@ -17,24 +17,20 @@
 package com.android.settings.notification;
 
 import android.content.Context;
-import android.view.View;
-import android.widget.Button;
+import android.text.TextUtils;
 import android.widget.Toast;
 
 import androidx.preference.Preference;
 
 import com.android.settings.R;
 import com.android.settings.core.BasePreferenceController;
-import com.android.settingslib.widget.LayoutPreference;
 
-public class ImportanceResetPreferenceController extends BasePreferenceController implements
-        View.OnClickListener {
+public class ImportanceResetPreferenceController extends BasePreferenceController {
 
     public static final String KEY = "asst_importance_reset";
     private static final String TAG = "ResetImportanceButton";
 
     private NotificationBackend mBackend;
-    private Button mButton;
 
     public ImportanceResetPreferenceController(Context context, String key) {
         super(context, key);
@@ -47,19 +43,14 @@
     }
 
     @Override
-    public void updateState(Preference preference) {
-        super.updateState(preference);
-
-        mButton = ((LayoutPreference) preference)
-                .findViewById(R.id.reset_importance_button);
-        mButton.setOnClickListener(this);
-    }
-
-    @Override
-    public void onClick(View v) {
+    public boolean handlePreferenceTreeClick(Preference preference) {
+        if (!TextUtils.equals(preference.getKey(), getPreferenceKey())) {
+            return false;
+        }
         mBackend.resetNotificationImportance();
         Toast.makeText(mContext, R.string.reset_importance_completed, Toast.LENGTH_SHORT)
-                            .show();
+                .show();
+        return true;
     }
 
     @Override
diff --git a/tests/robotests/src/com/android/settings/notification/ImportanceResetPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/notification/ImportanceResetPreferenceControllerTest.java
index 1dfa65d..687ef74 100644
--- a/tests/robotests/src/com/android/settings/notification/ImportanceResetPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/notification/ImportanceResetPreferenceControllerTest.java
@@ -26,10 +26,10 @@
 import android.content.Context;
 import android.view.View;
 
+import androidx.preference.Preference;
 import androidx.preference.PreferenceScreen;
 
 import com.android.settings.R;
-import com.android.settingslib.widget.LayoutPreference;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -43,12 +43,11 @@
 @RunWith(RobolectricTestRunner.class)
 public class ImportanceResetPreferenceControllerTest {
 
-    private static final String KEY = "asst_importance_reset";
     private ImportanceResetPreferenceController mController;
 
     @Mock
     private PreferenceScreen mScreen;
-    private LayoutPreference mPreference;
+    private Preference mPreference;
     private Context mContext;
     @Mock
     private NotificationBackend mBackend;
@@ -58,9 +57,10 @@
         MockitoAnnotations.initMocks(this);
 
         mContext = RuntimeEnvironment.application;
-        mPreference = new LayoutPreference(mContext, R.layout.reset_importance_button);
+        mPreference = new Preference(mContext);
+        mPreference.setKey(ImportanceResetPreferenceController.KEY);
         when(mScreen.findPreference(anyString())).thenReturn(mPreference);
-        mController = new ImportanceResetPreferenceController(mContext, KEY);
+        mController = new ImportanceResetPreferenceController(mContext, "some_key");
         mController.displayPreference(mScreen);
 
         ReflectionHelpers.setField(mController, "mBackend", mBackend);
@@ -68,10 +68,7 @@
 
     @Test
     public void onClick_callReset() {
-        final View view = mPreference.findViewById(R.id.reset_importance_button);
-        mController.updateState(mPreference);
-        assertThat(view.getVisibility()).isEqualTo(View.VISIBLE);
-        view.performClick();
+        mController.handlePreferenceTreeClick(mPreference);
 
         verify(mBackend, times(1)).resetNotificationImportance();
     }