diff --git a/res/layout/settings_action_buttons.xml b/res/layout/settings_action_buttons.xml
deleted file mode 100644
index 56e1a36..0000000
--- a/res/layout/settings_action_buttons.xml
+++ /dev/null
@@ -1,52 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!--
-  Copyright (C) 2017 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:paddingStart="8dp"
-    android:orientation="horizontal">
-
-    <Button
-        android:id="@+id/button1"
-        style="@style/SettingsActionButton"
-        android:layout_width="0dp"
-        android:layout_height="wrap_content"
-        android:layout_weight="1"/>
-
-    <Button
-        android:id="@+id/button2"
-        style="@style/SettingsActionButton"
-        android:layout_width="0dp"
-        android:layout_height="wrap_content"
-        android:layout_weight="1"/>
-
-    <Button
-        android:id="@+id/button3"
-        style="@style/SettingsActionButton"
-        android:layout_width="0dp"
-        android:layout_height="wrap_content"
-        android:layout_weight="1"/>
-
-    <Button
-        android:id="@+id/button4"
-        style="@style/SettingsActionButton"
-        android:layout_width="0dp"
-        android:layout_height="wrap_content"
-        android:layout_weight="1"/>
-</LinearLayout>
\ No newline at end of file
diff --git a/res/values/styles.xml b/res/values/styles.xml
index 8a92eaa..141426e 100644
--- a/res/values/styles.xml
+++ b/res/values/styles.xml
@@ -401,16 +401,6 @@
 
     <style name="ActionSecondaryButton" parent="android:Widget.DeviceDefault.Button"/>
 
-    <style name="SettingsActionButton"
-           parent="android:Widget.DeviceDefault.Button.Borderless.Colored">
-        <item name="android:drawablePadding">4dp</item>
-        <item name="android:drawableTint">@*android:color/btn_colored_borderless_text_material
-        </item>
-        <item name="android:layout_marginEnd">8dp</item>
-        <item name="android:paddingTop">20dp</item>
-        <item name="android:paddingBottom">20dp</item>
-    </style>
-
     <style name="LockPatternContainerStyle">
         <item name="android:maxHeight">400dp</item>
         <item name="android:maxWidth">420dp</item>
diff --git a/res/xml/app_info_settings.xml b/res/xml/app_info_settings.xml
index 8f99ce1..5731a73 100644
--- a/res/xml/app_info_settings.xml
+++ b/res/xml/app_info_settings.xml
@@ -36,7 +36,7 @@
         settings:allowDividerAbove="true"
         settings:allowDividerBelow="true"/>
 
-    <com.android.settings.widget.ActionButtonPreference
+    <com.android.settingslib.widget.ActionButtonsPreference
         android:key="action_buttons"
         android:order="-9998" />
 
diff --git a/res/xml/app_storage_settings.xml b/res/xml/app_storage_settings.xml
index 7036d27..376f09e 100644
--- a/res/xml/app_storage_settings.xml
+++ b/res/xml/app_storage_settings.xml
@@ -19,7 +19,7 @@
     xmlns:settings="http://schemas.android.com/apk/res-auto"
     android:title="@string/storage_label">
 
-    <com.android.settings.widget.ActionButtonPreference
+    <com.android.settingslib.widget.ActionButtonsPreference
         android:key="header_view" />
 
     <com.android.settings.applications.SpacePreference
diff --git a/res/xml/bluetooth_device_details_fragment.xml b/res/xml/bluetooth_device_details_fragment.xml
index 6449e4b..40ce93d 100644
--- a/res/xml/bluetooth_device_details_fragment.xml
+++ b/res/xml/bluetooth_device_details_fragment.xml
@@ -25,7 +25,7 @@
         android:selectable="false"
         settings:allowDividerBelow="true"/>
 
-    <com.android.settings.widget.ActionButtonPreference
+    <com.android.settingslib.widget.ActionButtonsPreference
         android:key="action_buttons" />
 
     <PreferenceCategory
diff --git a/res/xml/power_usage_detail.xml b/res/xml/power_usage_detail.xml
index 12c7e21..945ceed 100644
--- a/res/xml/power_usage_detail.xml
+++ b/res/xml/power_usage_detail.xml
@@ -26,7 +26,7 @@
         android:order="-10000"
         settings:allowDividerBelow="true"/>
 
-    <com.android.settings.widget.ActionButtonPreference
+    <com.android.settingslib.widget.ActionButtonsPreference
         android:key="action_buttons"
         android:order="-9999"/>
 
diff --git a/res/xml/tts_settings.xml b/res/xml/tts_settings.xml
index e76cdb7..92bf5d7 100644
--- a/res/xml/tts_settings.xml
+++ b/res/xml/tts_settings.xml
@@ -50,7 +50,7 @@
             android:defaultValue="100"
             android:max="400"/>
 
-        <com.android.settings.widget.ActionButtonPreference
+        <com.android.settingslib.widget.ActionButtonsPreference
             android:key="action_buttons" />
 
     </PreferenceCategory>
diff --git a/res/xml/wifi_network_details_fragment.xml b/res/xml/wifi_network_details_fragment.xml
index 8b027e2..e44d0d3 100644
--- a/res/xml/wifi_network_details_fragment.xml
+++ b/res/xml/wifi_network_details_fragment.xml
@@ -27,7 +27,7 @@
         settings:allowDividerBelow="true"/>
 
     <!-- Buttons -->
-    <com.android.settings.widget.ActionButtonPreference
+    <com.android.settingslib.widget.ActionButtonsPreference
         android:key="buttons"
         android:selectable="false" />
 
diff --git a/src/com/android/settings/applications/AppStorageSettings.java b/src/com/android/settings/applications/AppStorageSettings.java
index 61ce33d..61293c5 100644
--- a/src/com/android/settings/applications/AppStorageSettings.java
+++ b/src/com/android/settings/applications/AppStorageSettings.java
@@ -53,11 +53,11 @@
 import com.android.settings.R;
 import com.android.settings.Utils;
 import com.android.settings.deviceinfo.StorageWizardMoveConfirm;
-import com.android.settings.widget.ActionButtonPreference;
 import com.android.settingslib.RestrictedLockUtils;
 import com.android.settingslib.applications.ApplicationsState.Callbacks;
 import com.android.settingslib.applications.StorageStatsSource;
 import com.android.settingslib.applications.StorageStatsSource.AppStorageStats;
+import com.android.settingslib.widget.ActionButtonsPreference;
 import com.android.settingslib.widget.LayoutPreference;
 
 import java.util.Collections;
@@ -107,7 +107,7 @@
 
     // Views related to cache info
     @VisibleForTesting
-    ActionButtonPreference mButtonsPref;
+    ActionButtonsPreference mButtonsPref;
 
     private Preference mStorageUsed;
     private Button mChangeStorageButton;
@@ -168,7 +168,7 @@
                 .setComputingString(R.string.computing_size)
                 .setErrorString(R.string.invalid_size_value)
                 .build();
-        mButtonsPref = ((ActionButtonPreference) findPreference(KEY_HEADER_BUTTONS));
+        mButtonsPref = ((ActionButtonsPreference) findPreference(KEY_HEADER_BUTTONS));
         mStorageUsed = findPreference(KEY_STORAGE_USED);
         mChangeStorageButton = (Button) ((LayoutPreference) findPreference(KEY_CHANGE_STORAGE))
                 .findViewById(R.id.button);
diff --git a/src/com/android/settings/applications/appinfo/AppButtonsPreferenceController.java b/src/com/android/settings/applications/appinfo/AppButtonsPreferenceController.java
index 9740ecb..ebcb6ed 100644
--- a/src/com/android/settings/applications/appinfo/AppButtonsPreferenceController.java
+++ b/src/com/android/settings/applications/appinfo/AppButtonsPreferenceController.java
@@ -54,7 +54,6 @@
 import com.android.settings.core.InstrumentedPreferenceFragment;
 import com.android.settings.core.PreferenceControllerMixin;
 import com.android.settings.overlay.FeatureFactory;
-import com.android.settings.widget.ActionButtonPreference;
 import com.android.settingslib.RestrictedLockUtils;
 import com.android.settingslib.RestrictedLockUtilsInternal;
 import com.android.settingslib.applications.AppUtils;
@@ -64,6 +63,7 @@
 import com.android.settingslib.core.lifecycle.LifecycleObserver;
 import com.android.settingslib.core.lifecycle.events.OnDestroy;
 import com.android.settingslib.core.lifecycle.events.OnResume;
+import com.android.settingslib.widget.ActionButtonsPreference;
 
 import java.util.ArrayList;
 import java.util.HashSet;
@@ -99,7 +99,7 @@
     @VisibleForTesting
     boolean mDisableAfterUninstall = false;
     @VisibleForTesting
-    ActionButtonPreference mButtonsPref;
+    ActionButtonsPreference mButtonsPref;
 
     private final int mRequestUninstall;
     private final int mRequestRemoveDeviceAdmin;
@@ -158,14 +158,15 @@
     public int getAvailabilityStatus() {
         // TODO(b/37313605): Re-enable once this controller supports instant apps
         return mAppEntry != null && !AppUtils.isInstant(mAppEntry.info)
-            ? AVAILABLE : DISABLED_FOR_USER ;
+                ? AVAILABLE : DISABLED_FOR_USER;
     }
 
     @Override
     public void displayPreference(PreferenceScreen screen) {
         super.displayPreference(screen);
         if (isAvailable()) {
-            mButtonsPref = ((ActionButtonPreference) screen.findPreference(KEY_ACTION_BUTTONS))
+            mButtonsPref = ((ActionButtonsPreference) screen.findPreference(
+                    KEY_ACTION_BUTTONS))
                     .setButton1Text(R.string.uninstall_text)
                     .setButton1Icon(R.drawable.ic_settings_delete)
                     .setButton2Text(R.string.force_stop)
@@ -484,7 +485,7 @@
         } else {
             Intent intent = new Intent(Intent.ACTION_QUERY_PACKAGE_RESTART,
                     Uri.fromParts("package", mAppEntry.info.packageName, null));
-            intent.putExtra(Intent.EXTRA_PACKAGES, new String[] {mAppEntry.info.packageName});
+            intent.putExtra(Intent.EXTRA_PACKAGES, new String[]{mAppEntry.info.packageName});
             intent.putExtra(Intent.EXTRA_UID, mAppEntry.info.uid);
             intent.putExtra(Intent.EXTRA_USER_HANDLE, UserHandle.getUserId(mAppEntry.info.uid));
             Log.d(TAG, "Sending broadcast to query restart status for "
diff --git a/src/com/android/settings/bluetooth/BluetoothDetailsButtonsController.java b/src/com/android/settings/bluetooth/BluetoothDetailsButtonsController.java
index 4bb2103..a438f09 100644
--- a/src/com/android/settings/bluetooth/BluetoothDetailsButtonsController.java
+++ b/src/com/android/settings/bluetooth/BluetoothDetailsButtonsController.java
@@ -22,9 +22,9 @@
 import androidx.preference.PreferenceScreen;
 
 import com.android.settings.R;
-import com.android.settings.widget.ActionButtonPreference;
 import com.android.settingslib.bluetooth.CachedBluetoothDevice;
 import com.android.settingslib.core.lifecycle.Lifecycle;
+import com.android.settingslib.widget.ActionButtonsPreference;
 
 /**
  * This class adds two buttons: one to connect/disconnect from a device (depending on the current
@@ -35,7 +35,7 @@
     private boolean mIsConnected;
 
     private boolean mConnectButtonInitialized;
-    private ActionButtonPreference mActionButtons;
+    private ActionButtonsPreference mActionButtons;
 
     public BluetoothDetailsButtonsController(Context context, PreferenceFragmentCompat fragment,
             CachedBluetoothDevice device, Lifecycle lifecycle) {
@@ -51,7 +51,8 @@
 
     @Override
     protected void init(PreferenceScreen screen) {
-        mActionButtons = ((ActionButtonPreference) screen.findPreference(getPreferenceKey()))
+        mActionButtons = ((ActionButtonsPreference) screen.findPreference(
+                getPreferenceKey()))
                 .setButton1Text(R.string.forget)
                 .setButton1Icon(R.drawable.ic_settings_delete)
                 .setButton1OnClickListener((view) -> onForgetButtonPressed())
diff --git a/src/com/android/settings/tts/TextToSpeechSettings.java b/src/com/android/settings/tts/TextToSpeechSettings.java
index b1a1f8f..bd85ef3 100644
--- a/src/com/android/settings/tts/TextToSpeechSettings.java
+++ b/src/com/android/settings/tts/TextToSpeechSettings.java
@@ -44,15 +44,14 @@
 import com.android.settings.SettingsPreferenceFragment;
 import com.android.settings.search.BaseSearchIndexProvider;
 import com.android.settings.search.Indexable;
-import com.android.settings.widget.ActionButtonPreference;
 import com.android.settings.widget.GearPreference;
 import com.android.settings.widget.SeekBarPreference;
 import com.android.settingslib.search.SearchIndexable;
+import com.android.settingslib.widget.ActionButtonsPreference;
 
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Locale;
@@ -113,7 +112,7 @@
 
     private SeekBarPreference mDefaultPitchPref;
     private SeekBarPreference mDefaultRatePref;
-    private ActionButtonPreference mActionButtons;
+    private ActionButtonsPreference mActionButtons;
 
     private int mDefaultPitch = TextToSpeech.Engine.DEFAULT_PITCH;
     private int mDefaultRate = TextToSpeech.Engine.DEFAULT_RATE;
@@ -169,7 +168,7 @@
         mDefaultPitchPref = (SeekBarPreference) findPreference(KEY_DEFAULT_PITCH);
         mDefaultRatePref = (SeekBarPreference) findPreference(KEY_DEFAULT_RATE);
 
-        mActionButtons = ((ActionButtonPreference) findPreference(KEY_ACTION_BUTTONS))
+        mActionButtons = ((ActionButtonsPreference) findPreference(KEY_ACTION_BUTTONS))
                 .setButton1Text(R.string.tts_play)
                 .setButton1OnClickListener(v -> speakSampleText())
                 .setButton1Enabled(false)
diff --git a/src/com/android/settings/widget/ActionButtonPreference.java b/src/com/android/settings/widget/ActionButtonPreference.java
deleted file mode 100644
index a34d735..0000000
--- a/src/com/android/settings/widget/ActionButtonPreference.java
+++ /dev/null
@@ -1,333 +0,0 @@
-/*
- * Copyright (C) 2017 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.
- */
-
-package com.android.settings.widget;
-
-import android.content.Context;
-import android.content.res.Resources;
-import android.graphics.drawable.Drawable;
-import android.text.TextUtils;
-import android.util.AttributeSet;
-import android.util.Log;
-import android.view.View;
-import android.widget.Button;
-
-import androidx.annotation.DrawableRes;
-import androidx.annotation.StringRes;
-import androidx.preference.Preference;
-import androidx.preference.PreferenceViewHolder;
-
-import com.android.settings.R;
-
-/**
- * This preference provides a four buttons layout with Settings style.
- * It looks like below
- *
- * --------------------------------------------------
- *  button1     | button2   | button3   | button4   |
- * --------------------------------------------------
- *
- * User can set title / icon / click listener for each button.
- *
- * By default, four buttons are visible.
- * However, there are two cases which button should be invisible(View.GONE).
- *
- * 1. User sets invisible for button. ex: ActionButtonPreference.setButton1Visible(false)
- * 2. User doesn't set any title or icon for button.
- *
- */
-public class ActionButtonPreference extends Preference {
-
-    private final String TAG = "ActionButtonPreference";
-    private final ButtonInfo mButton1Info = new ButtonInfo();
-    private final ButtonInfo mButton2Info = new ButtonInfo();
-    private final ButtonInfo mButton3Info = new ButtonInfo();
-    private final ButtonInfo mButton4Info = new ButtonInfo();
-
-    public ActionButtonPreference(Context context, AttributeSet attrs,
-            int defStyleAttr, int defStyleRes) {
-        super(context, attrs, defStyleAttr, defStyleRes);
-        init();
-    }
-
-    public ActionButtonPreference(Context context, AttributeSet attrs, int defStyleAttr) {
-        super(context, attrs, defStyleAttr);
-        init();
-    }
-
-    public ActionButtonPreference(Context context, AttributeSet attrs) {
-        super(context, attrs);
-        init();
-    }
-
-    public ActionButtonPreference(Context context) {
-        super(context);
-        init();
-    }
-
-    private void init() {
-        setLayoutResource(R.layout.settings_action_buttons);
-        setSelectable(false);
-    }
-
-    @Override
-    public void onBindViewHolder(PreferenceViewHolder holder) {
-        super.onBindViewHolder(holder);
-        holder.setDividerAllowedAbove(true);
-        holder.setDividerAllowedBelow(true);
-
-        mButton1Info.mButton = (Button) holder.findViewById(R.id.button1);
-        mButton2Info.mButton = (Button) holder.findViewById(R.id.button2);
-        mButton3Info.mButton = (Button) holder.findViewById(R.id.button3);
-        mButton4Info.mButton = (Button) holder.findViewById(R.id.button4);
-
-        mButton1Info.setUpButton();
-        mButton2Info.setUpButton();
-        mButton3Info.setUpButton();
-        mButton4Info.setUpButton();
-    }
-
-    public ActionButtonPreference setButton1Visible(boolean isVisible) {
-        if (isVisible != mButton1Info.mIsVisible) {
-            mButton1Info.mIsVisible = isVisible;
-            notifyChanged();
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton1Text(@StringRes int textResId) {
-        final String newText = getContext().getString(textResId);
-        if (!TextUtils.equals(newText, mButton1Info.mText)) {
-            mButton1Info.mText = newText;
-            notifyChanged();
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton1Icon(@DrawableRes int iconResId) {
-        if (iconResId == 0) {
-            return this;
-        }
-
-        final Drawable icon;
-        try {
-            icon = getContext().getDrawable(iconResId);
-            mButton1Info.mIcon = icon;
-            notifyChanged();
-        } catch (Resources.NotFoundException exception) {
-            Log.e(TAG, "Resource does not exist: " + iconResId);
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton1Enabled(boolean isEnabled) {
-        if (isEnabled != mButton1Info.mIsEnabled) {
-            mButton1Info.mIsEnabled = isEnabled;
-            notifyChanged();
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton1OnClickListener(View.OnClickListener listener) {
-        if (listener != mButton1Info.mListener) {
-            mButton1Info.mListener = listener;
-            notifyChanged();
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton2Visible(boolean isVisible) {
-        if (isVisible != mButton2Info.mIsVisible) {
-            mButton2Info.mIsVisible = isVisible;
-            notifyChanged();
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton2Text(@StringRes int textResId) {
-        final String newText = getContext().getString(textResId);
-        if (!TextUtils.equals(newText, mButton2Info.mText)) {
-            mButton2Info.mText = newText;
-            notifyChanged();
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton2Icon(@DrawableRes int iconResId) {
-        if (iconResId == 0) {
-            return this;
-        }
-
-        final Drawable icon;
-        try {
-            icon = getContext().getDrawable(iconResId);
-            mButton2Info.mIcon = icon;
-            notifyChanged();
-        } catch (Resources.NotFoundException exception) {
-            Log.e(TAG, "Resource does not exist: " + iconResId);
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton2Enabled(boolean isEnabled) {
-        if (isEnabled != mButton2Info.mIsEnabled) {
-            mButton2Info.mIsEnabled = isEnabled;
-            notifyChanged();
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton2OnClickListener(View.OnClickListener listener) {
-        if (listener != mButton2Info.mListener) {
-            mButton2Info.mListener = listener;
-            notifyChanged();
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton3Visible(boolean isVisible) {
-        if (isVisible != mButton3Info.mIsVisible) {
-            mButton3Info.mIsVisible = isVisible;
-            notifyChanged();
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton3Text(@StringRes int textResId) {
-        final String newText = getContext().getString(textResId);
-        if (!TextUtils.equals(newText, mButton3Info.mText)) {
-            mButton3Info.mText = newText;
-            notifyChanged();
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton3Icon(@DrawableRes int iconResId) {
-        if (iconResId == 0) {
-            return this;
-        }
-
-        final Drawable icon;
-        try {
-            icon = getContext().getDrawable(iconResId);
-            mButton3Info.mIcon = icon;
-            notifyChanged();
-        } catch (Resources.NotFoundException exception) {
-            Log.e(TAG, "Resource does not exist: " + iconResId);
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton3Enabled(boolean isEnabled) {
-        if (isEnabled != mButton3Info.mIsEnabled) {
-            mButton3Info.mIsEnabled = isEnabled;
-            notifyChanged();
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton3OnClickListener(View.OnClickListener listener) {
-        if (listener != mButton3Info.mListener) {
-            mButton3Info.mListener = listener;
-            notifyChanged();
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton4Visible(boolean isVisible) {
-        if (isVisible != mButton4Info.mIsVisible) {
-            mButton4Info.mIsVisible = isVisible;
-            notifyChanged();
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton4Text(@StringRes int textResId) {
-        final String newText = getContext().getString(textResId);
-        if (!TextUtils.equals(newText, mButton4Info.mText)) {
-            mButton4Info.mText = newText;
-            notifyChanged();
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton4Icon(@DrawableRes int iconResId) {
-        if (iconResId == 0) {
-            return this;
-        }
-
-        final Drawable icon;
-        try {
-            icon = getContext().getDrawable(iconResId);
-            mButton4Info.mIcon = icon;
-            notifyChanged();
-        } catch (Resources.NotFoundException exception) {
-            Log.e(TAG, "Resource does not exist: " + iconResId);
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton4Enabled(boolean isEnabled) {
-        if (isEnabled != mButton4Info.mIsEnabled) {
-            mButton4Info.mIsEnabled = isEnabled;
-            notifyChanged();
-        }
-        return this;
-    }
-
-    public ActionButtonPreference setButton4OnClickListener(View.OnClickListener listener) {
-        if (listener != mButton4Info.mListener) {
-            mButton4Info.mListener = listener;
-            notifyChanged();
-        }
-        return this;
-    }
-
-    static class ButtonInfo {
-        private Button mButton;
-        private CharSequence mText;
-        private Drawable mIcon;
-        private View.OnClickListener mListener;
-        private boolean mIsEnabled = true;
-        private boolean mIsVisible = true;
-
-        void setUpButton() {
-            mButton.setText(mText);
-            mButton.setOnClickListener(mListener);
-            mButton.setEnabled(mIsEnabled);
-            mButton.setCompoundDrawablesWithIntrinsicBounds(
-                    null /* left */, mIcon /* top */, null /* right */, null /* bottom */);
-
-            if (shouldBeVisible()) {
-                mButton.setVisibility(View.VISIBLE);
-            } else {
-                mButton.setVisibility(View.GONE);
-            }
-        }
-
-        /**
-         * By default, four buttons are visible.
-         * However, there are two cases which button should be invisible.
-         *
-         * 1. User set invisible for this button. ex: mIsVisible = false.
-         * 2. User didn't set any title or icon.
-         *
-         */
-        private boolean shouldBeVisible() {
-            return mIsVisible && (!TextUtils.isEmpty(mText) || mIcon != null);
-        }
-    }
-}
\ No newline at end of file
diff --git a/src/com/android/settings/wifi/details/WifiDetailPreferenceController.java b/src/com/android/settings/wifi/details/WifiDetailPreferenceController.java
index ba8241a..a676bfa 100644
--- a/src/com/android/settings/wifi/details/WifiDetailPreferenceController.java
+++ b/src/com/android/settings/wifi/details/WifiDetailPreferenceController.java
@@ -55,7 +55,6 @@
 import com.android.settings.R;
 import com.android.settings.Utils;
 import com.android.settings.core.PreferenceControllerMixin;
-import com.android.settings.widget.ActionButtonPreference;
 import com.android.settings.widget.EntityHeaderController;
 import com.android.settings.wifi.WifiDialog;
 import com.android.settings.wifi.WifiDialog.WifiDialogListener;
@@ -66,6 +65,7 @@
 import com.android.settingslib.core.lifecycle.LifecycleObserver;
 import com.android.settingslib.core.lifecycle.events.OnPause;
 import com.android.settingslib.core.lifecycle.events.OnResume;
+import com.android.settingslib.widget.ActionButtonsPreference;
 import com.android.settingslib.widget.LayoutPreference;
 import com.android.settingslib.wifi.AccessPoint;
 
@@ -130,7 +130,7 @@
     private final MetricsFeatureProvider mMetricsFeatureProvider;
 
     // UI elements - in order of appearance
-    private ActionButtonPreference mButtonsPref;
+    private ActionButtonsPreference mButtonsPref;
     private EntityHeaderController mEntityHeaderController;
     private Preference mSignalStrengthPref;
     private Preference mLinkSpeedPref;
@@ -275,7 +275,7 @@
 
         setupEntityHeader(screen);
 
-        mButtonsPref = ((ActionButtonPreference) screen.findPreference(KEY_BUTTONS_PREF))
+        mButtonsPref = ((ActionButtonsPreference) screen.findPreference(KEY_BUTTONS_PREF))
                 .setButton1Text(R.string.forget)
                 .setButton1Icon(R.drawable.ic_settings_delete)
                 .setButton1OnClickListener(view -> forgetNetwork())
diff --git a/tests/robotests/src/com/android/settings/applications/AppStorageSettingsTest.java b/tests/robotests/src/com/android/settings/applications/AppStorageSettingsTest.java
index 87eb6a6..114000c 100644
--- a/tests/robotests/src/com/android/settings/applications/AppStorageSettingsTest.java
+++ b/tests/robotests/src/com/android/settings/applications/AppStorageSettingsTest.java
@@ -16,6 +16,8 @@
 
 package com.android.settings.applications;
 
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.nullable;
 import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.doNothing;
@@ -30,9 +32,8 @@
 import android.widget.Button;
 
 import com.android.settings.testutils.SettingsRobolectricTestRunner;
-import com.android.settings.widget.ActionButtonPreference;
-import com.android.settings.widget.ActionButtonPreferenceTest;
 import com.android.settingslib.applications.StorageStatsSource.AppStorageStats;
+import com.android.settingslib.widget.ActionButtonsPreference;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -46,7 +47,7 @@
 
     @Mock
     private AppStorageSizesController mSizesController;
-    private ActionButtonPreference mButtonsPref;
+    private ActionButtonsPreference mButtonsPref;
     private AppStorageSettings mSettings;
     private Button mLeftButton;
     private Button mRightButton;
@@ -58,7 +59,7 @@
         mRightButton = new Button(RuntimeEnvironment.application);
         mSettings = spy(new AppStorageSettings());
         mSettings.mSizeController = mSizesController;
-        mButtonsPref = ActionButtonPreferenceTest.createMock();
+        mButtonsPref = createMock();
         mSettings.mButtonsPref = mButtonsPref;
 
         when(mButtonsPref.setButton1OnClickListener(any(View.OnClickListener.class)))
@@ -103,5 +104,22 @@
         verify(mSettings).handleClearDataClick();
         verify(mSettings).handleClearCacheClick();
     }
+
+    private ActionButtonsPreference createMock() {
+        final ActionButtonsPreference pref = mock(ActionButtonsPreference.class);
+        when(pref.setButton1Text(anyInt())).thenReturn(pref);
+        when(pref.setButton1Icon(anyInt())).thenReturn(pref);
+        when(pref.setButton1Enabled(anyBoolean())).thenReturn(pref);
+        when(pref.setButton1Visible(anyBoolean())).thenReturn(pref);
+        when(pref.setButton1OnClickListener(any(View.OnClickListener.class))).thenReturn(pref);
+
+        when(pref.setButton2Text(anyInt())).thenReturn(pref);
+        when(pref.setButton2Icon(anyInt())).thenReturn(pref);
+        when(pref.setButton2Enabled(anyBoolean())).thenReturn(pref);
+        when(pref.setButton2Visible(anyBoolean())).thenReturn(pref);
+        when(pref.setButton2OnClickListener(any(View.OnClickListener.class))).thenReturn(pref);
+
+        return pref;
+    }
 }
 
diff --git a/tests/robotests/src/com/android/settings/applications/appinfo/AppButtonsPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/applications/appinfo/AppButtonsPreferenceControllerTest.java
index 815c76b..e40a885 100644
--- a/tests/robotests/src/com/android/settings/applications/appinfo/AppButtonsPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/applications/appinfo/AppButtonsPreferenceControllerTest.java
@@ -18,6 +18,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
 import static org.mockito.Matchers.anyString;
@@ -42,18 +43,18 @@
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.os.UserManager;
+import android.view.View;
 
 import com.android.settings.R;
 import com.android.settings.SettingsActivity;
 import com.android.settings.core.InstrumentedPreferenceFragment;
 import com.android.settings.testutils.FakeFeatureFactory;
 import com.android.settings.testutils.SettingsRobolectricTestRunner;
-import com.android.settings.widget.ActionButtonPreference;
-import com.android.settings.widget.ActionButtonPreferenceTest;
 import com.android.settingslib.applications.AppUtils;
 import com.android.settingslib.applications.ApplicationsState;
 import com.android.settingslib.applications.instantapps.InstantAppDataProvider;
 import com.android.settingslib.core.lifecycle.Lifecycle;
+import com.android.settingslib.widget.ActionButtonsPreference;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -99,7 +100,7 @@
     @Mock
     private PackageInfo mPackageInfo;
 
-    private ActionButtonPreference mButtonPrefs;
+    private ActionButtonsPreference mButtonPrefs;
 
     private Intent mUninstallIntent;
     private AppButtonsPreferenceController mController;
@@ -127,7 +128,7 @@
         mPackageInfo.packageName = PACKAGE_NAME;
         mPackageInfo.applicationInfo = mAppInfo;
 
-        mButtonPrefs = ActionButtonPreferenceTest.createMock();
+        mButtonPrefs = createMock();
         mController.mButtonsPref = mButtonPrefs;
         mController.mPackageInfo = mPackageInfo;
 
@@ -390,4 +391,15 @@
             return SettingsEnums.PAGE_UNKNOWN;
         }
     }
+
+    private ActionButtonsPreference createMock() {
+        final ActionButtonsPreference pref = mock(ActionButtonsPreference.class);
+        when(pref.setButton1Text(anyInt())).thenReturn(pref);
+        when(pref.setButton1Icon(anyInt())).thenReturn(pref);
+        when(pref.setButton1Enabled(anyBoolean())).thenReturn(pref);
+        when(pref.setButton1Visible(anyBoolean())).thenReturn(pref);
+        when(pref.setButton1OnClickListener(any(View.OnClickListener.class))).thenReturn(pref);
+
+        return pref;
+    }
 }
diff --git a/tests/robotests/src/com/android/settings/bluetooth/BluetoothDetailsButtonsControllerTest.java b/tests/robotests/src/com/android/settings/bluetooth/BluetoothDetailsButtonsControllerTest.java
index 094ec80..d8be01f 100644
--- a/tests/robotests/src/com/android/settings/bluetooth/BluetoothDetailsButtonsControllerTest.java
+++ b/tests/robotests/src/com/android/settings/bluetooth/BluetoothDetailsButtonsControllerTest.java
@@ -18,6 +18,8 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyString;
 import static org.mockito.Matchers.eq;
@@ -34,8 +36,7 @@
 import com.android.settings.R;
 import com.android.settings.testutils.SettingsRobolectricTestRunner;
 import com.android.settings.testutils.shadow.SettingsShadowBluetoothDevice;
-import com.android.settings.widget.ActionButtonPreference;
-import com.android.settings.widget.ActionButtonPreferenceTest;
+import com.android.settingslib.widget.ActionButtonsPreference;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -47,7 +48,7 @@
 @Config(shadows = SettingsShadowBluetoothDevice.class)
 public class BluetoothDetailsButtonsControllerTest extends BluetoothDetailsControllerTestBase {
     private BluetoothDetailsButtonsController mController;
-    private ActionButtonPreference mButtonsPref;
+    private ActionButtonsPreference mButtonsPref;
     private Button mConnectButton;
     private Button mForgetButton;
 
@@ -60,8 +61,9 @@
         mConnectButton = buttons.findViewById(R.id.button2);
         mForgetButton = buttons.findViewById(R.id.button1);
         mController =
-            new BluetoothDetailsButtonsController(mContext, mFragment, mCachedDevice, mLifecycle);
-        mButtonsPref = ActionButtonPreferenceTest.createMock();
+                new BluetoothDetailsButtonsController(mContext, mFragment, mCachedDevice,
+                        mLifecycle);
+        mButtonsPref = createMock();
         when(mButtonsPref.getKey()).thenReturn(mController.getPreferenceKey());
         when(mButtonsPref.setButton2OnClickListener(any(View.OnClickListener.class)))
                 .thenAnswer(invocation -> {
@@ -186,4 +188,21 @@
 
         verify(mButtonsPref).setButton2Enabled(false);
     }
+
+    private ActionButtonsPreference createMock() {
+        final ActionButtonsPreference pref = mock(ActionButtonsPreference.class);
+        when(pref.setButton1Text(anyInt())).thenReturn(pref);
+        when(pref.setButton1Icon(anyInt())).thenReturn(pref);
+        when(pref.setButton1Enabled(anyBoolean())).thenReturn(pref);
+        when(pref.setButton1Visible(anyBoolean())).thenReturn(pref);
+        when(pref.setButton1OnClickListener(any(View.OnClickListener.class))).thenReturn(pref);
+
+        when(pref.setButton2Text(anyInt())).thenReturn(pref);
+        when(pref.setButton2Icon(anyInt())).thenReturn(pref);
+        when(pref.setButton2Enabled(anyBoolean())).thenReturn(pref);
+        when(pref.setButton2Visible(anyBoolean())).thenReturn(pref);
+        when(pref.setButton2OnClickListener(any(View.OnClickListener.class))).thenReturn(pref);
+
+        return pref;
+    }
 }
diff --git a/tests/robotests/src/com/android/settings/testutils/SettingsRobolectricTestRunner.java b/tests/robotests/src/com/android/settings/testutils/SettingsRobolectricTestRunner.java
index 9417e9c..74914f4 100644
--- a/tests/robotests/src/com/android/settings/testutils/SettingsRobolectricTestRunner.java
+++ b/tests/robotests/src/com/android/settings/testutils/SettingsRobolectricTestRunner.java
@@ -84,6 +84,8 @@
             paths.add(new ResourcePath(null,
                     Fs.fromURL(new URL("file:frameworks/base/packages/SettingsLib/SearchWidget/res/")), null));
             paths.add(new ResourcePath(null,
+                    Fs.fromURL(new URL("file:frameworks/base/packages/SettingsLib/ActionButtonsPreference/res")), null));
+            paths.add(new ResourcePath(null,
                     Fs.fromURL(new URL("file:frameworks/base/packages/SettingsLib/SettingsLayoutPreference/res")), null));
             paths.add(new ResourcePath(null,
                 Fs.fromURL(new URL("file:frameworks/base/core/res/res")), null));
diff --git a/tests/robotests/src/com/android/settings/widget/ActionButtonPreferenceTest.java b/tests/robotests/src/com/android/settings/widget/ActionButtonPreferenceTest.java
deleted file mode 100644
index 34a9872..0000000
--- a/tests/robotests/src/com/android/settings/widget/ActionButtonPreferenceTest.java
+++ /dev/null
@@ -1,283 +0,0 @@
-/*
- * Copyright (C) 2017 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.
- */
-
-package com.android.settings.widget;
-
-import static com.google.common.truth.Truth.assertThat;
-
-import static org.mockito.Matchers.any;
-import static org.mockito.Matchers.anyBoolean;
-import static org.mockito.Matchers.anyInt;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
-
-import android.content.Context;
-import android.graphics.drawable.Drawable;
-import android.view.View;
-import android.widget.Button;
-
-import androidx.preference.PreferenceViewHolder;
-
-import com.android.settings.R;
-import com.android.settings.testutils.SettingsRobolectricTestRunner;
-
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.robolectric.RuntimeEnvironment;
-
-@RunWith(SettingsRobolectricTestRunner.class)
-public class ActionButtonPreferenceTest {
-
-    private Context mContext;
-    private View mRootView;
-    private ActionButtonPreference mPref;
-    private PreferenceViewHolder mHolder;
-
-    @Before
-    public void setUp() {
-        mContext = RuntimeEnvironment.application;
-        mRootView = View.inflate(mContext, R.layout.settings_action_buttons, null /* parent */);
-        mHolder = PreferenceViewHolder.createInstanceForTests(mRootView);
-        mPref = new ActionButtonPreference(mContext);
-    }
-
-    @Test
-    public void onBindViewHolder_setTitle_shouldShowButtonByDefault() {
-        mPref.setButton1Text(R.string.settings_label);
-        mPref.setButton2Text(R.string.settings_label);
-        mPref.setButton3Text(R.string.settings_label);
-        mPref.setButton4Text(R.string.settings_label);
-
-        mPref.onBindViewHolder(mHolder);
-
-        assertThat(mRootView.findViewById(R.id.button1).getVisibility())
-                .isEqualTo(View.VISIBLE);
-        assertThat(mRootView.findViewById(R.id.button2).getVisibility())
-                .isEqualTo(View.VISIBLE);
-        assertThat(mRootView.findViewById(R.id.button3).getVisibility())
-                .isEqualTo(View.VISIBLE);
-        assertThat(mRootView.findViewById(R.id.button4).getVisibility())
-                .isEqualTo(View.VISIBLE);
-    }
-
-    @Test
-    public void onBindViewHolder_setIcon_shouldShowButtonByDefault() {
-        mPref.setButton1Icon(R.drawable.ic_settings);
-        mPref.setButton2Icon(R.drawable.ic_settings);
-        mPref.setButton3Icon(R.drawable.ic_settings);
-        mPref.setButton4Icon(R.drawable.ic_settings);
-
-        mPref.onBindViewHolder(mHolder);
-
-        assertThat(mRootView.findViewById(R.id.button1).getVisibility())
-                .isEqualTo(View.VISIBLE);
-        assertThat(mRootView.findViewById(R.id.button2).getVisibility())
-                .isEqualTo(View.VISIBLE);
-        assertThat(mRootView.findViewById(R.id.button3).getVisibility())
-                .isEqualTo(View.VISIBLE);
-        assertThat(mRootView.findViewById(R.id.button4).getVisibility())
-                .isEqualTo(View.VISIBLE);
-    }
-
-    @Test
-    public void onBindViewHolder_notSetTitleOrIcon_shouldNotShowButtonByDefault() {
-        mPref.onBindViewHolder(mHolder);
-
-        assertThat(mRootView.findViewById(R.id.button1).getVisibility())
-                .isEqualTo(View.GONE);
-        assertThat(mRootView.findViewById(R.id.button2).getVisibility())
-                .isEqualTo(View.GONE);
-        assertThat(mRootView.findViewById(R.id.button3).getVisibility())
-                .isEqualTo(View.GONE);
-        assertThat(mRootView.findViewById(R.id.button4).getVisibility())
-                .isEqualTo(View.GONE);
-    }
-
-    @Test
-    public void onBindViewHolder_setVisibleIsGoneAndSetTitle_shouldNotShowButton() {
-        mPref.setButton1Text(R.string.settings_label).setButton1Visible(false);
-        mPref.setButton2Text(R.string.settings_label).setButton2Visible(false);
-        mPref.setButton3Text(R.string.settings_label).setButton3Visible(false);
-        mPref.setButton4Text(R.string.settings_label).setButton4Visible(false);
-
-        mPref.onBindViewHolder(mHolder);
-
-        assertThat(mRootView.findViewById(R.id.button1).getVisibility())
-                .isEqualTo(View.GONE);
-        assertThat(mRootView.findViewById(R.id.button2).getVisibility())
-                .isEqualTo(View.GONE);
-        assertThat(mRootView.findViewById(R.id.button3).getVisibility())
-                .isEqualTo(View.GONE);
-        assertThat(mRootView.findViewById(R.id.button4).getVisibility())
-                .isEqualTo(View.GONE);
-    }
-
-    @Test
-    public void onBindViewHolder_setVisibleIsGoneAndSetIcon_shouldNotShowButton() {
-        mPref.setButton1Icon(R.drawable.ic_settings).setButton1Visible(false);
-        mPref.setButton2Icon(R.drawable.ic_settings).setButton2Visible(false);
-        mPref.setButton3Icon(R.drawable.ic_settings).setButton3Visible(false);
-        mPref.setButton4Icon(R.drawable.ic_settings).setButton4Visible(false);
-
-        mPref.onBindViewHolder(mHolder);
-
-        assertThat(mRootView.findViewById(R.id.button1).getVisibility())
-                .isEqualTo(View.GONE);
-        assertThat(mRootView.findViewById(R.id.button2).getVisibility())
-                .isEqualTo(View.GONE);
-        assertThat(mRootView.findViewById(R.id.button3).getVisibility())
-                .isEqualTo(View.GONE);
-        assertThat(mRootView.findViewById(R.id.button4).getVisibility())
-                .isEqualTo(View.GONE);
-    }
-
-    @Test
-    public void onBindViewHolder_setVisibility_shouldUpdateButtonVisibility() {
-        mPref.setButton1Text(R.string.settings_label).setButton1Visible(false);
-        mPref.setButton2Text(R.string.settings_label).setButton2Visible(false);
-        mPref.setButton3Text(R.string.settings_label).setButton3Visible(false);
-        mPref.setButton4Text(R.string.settings_label).setButton4Visible(false);
-
-        mPref.onBindViewHolder(mHolder);
-
-        assertThat(mRootView.findViewById(R.id.button1).getVisibility())
-                .isEqualTo(View.GONE);
-        assertThat(mRootView.findViewById(R.id.button2).getVisibility())
-                .isEqualTo(View.GONE);
-        assertThat(mRootView.findViewById(R.id.button3).getVisibility())
-                .isEqualTo(View.GONE);
-        assertThat(mRootView.findViewById(R.id.button4).getVisibility())
-                .isEqualTo(View.GONE);
-
-        mPref.setButton1Visible(true);
-        mPref.setButton2Visible(true);
-        mPref.setButton3Visible(true);
-        mPref.setButton4Visible(true);
-
-        mPref.onBindViewHolder(mHolder);
-
-        assertThat(mRootView.findViewById(R.id.button1).getVisibility())
-                .isEqualTo(View.VISIBLE);
-        assertThat(mRootView.findViewById(R.id.button2).getVisibility())
-                .isEqualTo(View.VISIBLE);
-        assertThat(mRootView.findViewById(R.id.button3).getVisibility())
-                .isEqualTo(View.VISIBLE);
-        assertThat(mRootView.findViewById(R.id.button4).getVisibility())
-                .isEqualTo(View.VISIBLE);
-    }
-
-    @Test
-    public void onBindViewHolder_setEnabled_shouldEnableButton() {
-        mPref.setButton1Enabled(true);
-        mPref.setButton2Enabled(false);
-        mPref.setButton3Enabled(true);
-        mPref.setButton4Enabled(false);
-
-        mPref.onBindViewHolder(mHolder);
-
-        assertThat(mRootView.findViewById(R.id.button1).isEnabled()).isTrue();
-        assertThat(mRootView.findViewById(R.id.button2).isEnabled()).isFalse();
-        assertThat(mRootView.findViewById(R.id.button3).isEnabled()).isTrue();
-        assertThat(mRootView.findViewById(R.id.button4).isEnabled()).isFalse();
-    }
-
-    @Test
-    public void onBindViewHolder_setText_shouldShowSameText() {
-        mPref.setButton1Text(R.string.settings_label);
-        mPref.setButton2Text(R.string.settings_label);
-        mPref.setButton3Text(R.string.settings_label);
-        mPref.setButton4Text(R.string.settings_label);
-
-        mPref.onBindViewHolder(mHolder);
-
-        assertThat(((Button) mRootView.findViewById(R.id.button1)).getText())
-                .isEqualTo(mContext.getText(R.string.settings_label));
-        assertThat(((Button) mRootView.findViewById(R.id.button2)).getText())
-                .isEqualTo(mContext.getText(R.string.settings_label));
-        assertThat(((Button) mRootView.findViewById(R.id.button3)).getText())
-                .isEqualTo(mContext.getText(R.string.settings_label));
-        assertThat(((Button) mRootView.findViewById(R.id.button4)).getText())
-                .isEqualTo(mContext.getText(R.string.settings_label));
-    }
-
-    @Test
-    public void onBindViewHolder_setButtonIcon_iconMustDisplayAboveText() {
-        mPref.setButton1Text(R.string.settings_label);
-        mPref.setButton1Icon(R.drawable.ic_settings);
-
-        mPref.onBindViewHolder(mHolder);
-        final Drawable[] drawablesAroundText =
-                ((Button) mRootView.findViewById(R.id.button1))
-                        .getCompoundDrawables();
-
-        assertThat(drawablesAroundText[1 /* top */]).isNotNull();
-    }
-
-    @Test
-    public void setButtonIcon_iconResourceIdIsZero_shouldNotDisplayIcon() {
-        mPref.setButton1Text(R.string.settings_label);
-        mPref.setButton1Icon(0);
-
-        mPref.onBindViewHolder(mHolder);
-        final Drawable[] drawablesAroundText =
-                ((Button) mRootView.findViewById(R.id.button1))
-                        .getCompoundDrawables();
-
-        assertThat(drawablesAroundText[1 /* top */]).isNull();
-    }
-
-    @Test
-    public void setButtonIcon_iconResourceIdNotExisting_shouldNotDisplayIconAndCrash() {
-        mPref.setButton1Text(R.string.settings_label);
-        mPref.setButton1Icon(999999999 /* not existing id */);
-        // Should not crash here
-        mPref.onBindViewHolder(mHolder);
-        final Drawable[] drawablesAroundText =
-                ((Button) mRootView.findViewById(R.id.button1))
-                        .getCompoundDrawables();
-
-        assertThat(drawablesAroundText[1 /* top */]).isNull();
-    }
-
-    public static ActionButtonPreference createMock() {
-        final ActionButtonPreference pref = mock(ActionButtonPreference.class);
-        when(pref.setButton1Text(anyInt())).thenReturn(pref);
-        when(pref.setButton1Icon(anyInt())).thenReturn(pref);
-        when(pref.setButton1Enabled(anyBoolean())).thenReturn(pref);
-        when(pref.setButton1Visible(anyBoolean())).thenReturn(pref);
-        when(pref.setButton1OnClickListener(any(View.OnClickListener.class))).thenReturn(pref);
-
-        when(pref.setButton2Text(anyInt())).thenReturn(pref);
-        when(pref.setButton2Icon(anyInt())).thenReturn(pref);
-        when(pref.setButton2Enabled(anyBoolean())).thenReturn(pref);
-        when(pref.setButton2Visible(anyBoolean())).thenReturn(pref);
-        when(pref.setButton2OnClickListener(any(View.OnClickListener.class))).thenReturn(pref);
-
-        when(pref.setButton3Text(anyInt())).thenReturn(pref);
-        when(pref.setButton3Icon(anyInt())).thenReturn(pref);
-        when(pref.setButton3Enabled(anyBoolean())).thenReturn(pref);
-        when(pref.setButton3Visible(anyBoolean())).thenReturn(pref);
-        when(pref.setButton3OnClickListener(any(View.OnClickListener.class))).thenReturn(pref);
-
-        when(pref.setButton4Text(anyInt())).thenReturn(pref);
-        when(pref.setButton4Icon(anyInt())).thenReturn(pref);
-        when(pref.setButton4Enabled(anyBoolean())).thenReturn(pref);
-        when(pref.setButton4Visible(anyBoolean())).thenReturn(pref);
-        when(pref.setButton4OnClickListener(any(View.OnClickListener.class))).thenReturn(pref);
-        return pref;
-    }
-}
\ No newline at end of file
diff --git a/tests/robotests/src/com/android/settings/wifi/details/WifiDetailPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/wifi/details/WifiDetailPreferenceControllerTest.java
index d39b55f..39215c1 100644
--- a/tests/robotests/src/com/android/settings/wifi/details/WifiDetailPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/wifi/details/WifiDetailPreferenceControllerTest.java
@@ -17,12 +17,14 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.nullable;
 import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.spy;
@@ -67,11 +69,10 @@
 import com.android.settings.testutils.shadow.ShadowBidiFormatter;
 import com.android.settings.testutils.shadow.ShadowDevicePolicyManager;
 import com.android.settings.testutils.shadow.ShadowEntityHeaderController;
-import com.android.settings.widget.ActionButtonPreference;
-import com.android.settings.widget.ActionButtonPreferenceTest;
 import com.android.settings.widget.EntityHeaderController;
 import com.android.settingslib.core.instrumentation.MetricsFeatureProvider;
 import com.android.settingslib.core.lifecycle.Lifecycle;
+import com.android.settingslib.widget.ActionButtonsPreference;
 import com.android.settingslib.widget.LayoutPreference;
 import com.android.settingslib.wifi.AccessPoint;
 
@@ -142,7 +143,7 @@
     private ImageView mockHeaderIcon;
 
     @Mock
-    private ActionButtonPreference mockButtonsPref;
+    private ActionButtonsPreference mockButtonsPref;
     @Mock
     private Preference mockSignalStrengthPref;
     @Mock
@@ -247,7 +248,7 @@
                 .thenReturn(mockNetworkInfo);
         doNothing().when(mockConnectivityManager).registerNetworkCallback(
                 nullable(NetworkRequest.class), mCallbackCaptor.capture(), nullable(Handler.class));
-        mockButtonsPref = ActionButtonPreferenceTest.createMock();
+        mockButtonsPref = createMock();
         when(mockButtonsPref.setButton1OnClickListener(mForgetClickListener.capture()))
                 .thenReturn(mockButtonsPref);
 
@@ -849,4 +850,21 @@
         verify(mockAccessPoint, times(2)).getLevel();
         verify(mockIconInjector, times(2)).getIcon(anyInt());
     }
+
+    private ActionButtonsPreference createMock() {
+        final ActionButtonsPreference pref = mock(ActionButtonsPreference.class);
+        when(pref.setButton1Text(anyInt())).thenReturn(pref);
+        when(pref.setButton1Icon(anyInt())).thenReturn(pref);
+        when(pref.setButton1Enabled(anyBoolean())).thenReturn(pref);
+        when(pref.setButton1Visible(anyBoolean())).thenReturn(pref);
+        when(pref.setButton1OnClickListener(any(View.OnClickListener.class))).thenReturn(pref);
+
+        when(pref.setButton2Text(anyInt())).thenReturn(pref);
+        when(pref.setButton2Icon(anyInt())).thenReturn(pref);
+        when(pref.setButton2Enabled(anyBoolean())).thenReturn(pref);
+        when(pref.setButton2Visible(anyBoolean())).thenReturn(pref);
+        when(pref.setButton2OnClickListener(any(View.OnClickListener.class))).thenReturn(pref);
+
+        return pref;
+    }
 }