diff --git a/packages/SettingsLib/MainSwitchPreference/src/com/android/settingslib/widget/MainSwitchBar.java b/packages/SettingsLib/MainSwitchPreference/src/com/android/settingslib/widget/MainSwitchBar.java
index 6001155..2b5fcd8 100644
--- a/packages/SettingsLib/MainSwitchPreference/src/com/android/settingslib/widget/MainSwitchBar.java
+++ b/packages/SettingsLib/MainSwitchPreference/src/com/android/settingslib/widget/MainSwitchBar.java
@@ -24,8 +24,8 @@
 import android.view.LayoutInflater;
 import android.view.View;
 import android.widget.CompoundButton;
+import android.widget.CompoundButton.OnCheckedChangeListener;
 import android.widget.LinearLayout;
-import android.widget.Switch;
 import android.widget.TextView;
 
 import androidx.annotation.ColorInt;
@@ -41,9 +41,9 @@
  * This component is used as the main switch of the page
  * to enable or disable the prefereces on the page.
  */
-public class MainSwitchBar extends LinearLayout implements CompoundButton.OnCheckedChangeListener {
+public class MainSwitchBar extends LinearLayout implements OnCheckedChangeListener {
 
-    private final List<OnMainSwitchChangeListener> mSwitchChangeListeners = new ArrayList<>();
+    private final List<OnCheckedChangeListener> mSwitchChangeListeners = new ArrayList<>();
 
     @ColorInt
     private int mBackgroundColor;
@@ -51,8 +51,8 @@
     private int mBackgroundActivatedColor;
 
     protected TextView mTextView;
-    protected Switch mSwitch;
-    private View mFrameView;
+    protected CompoundButton mSwitch;
+    private final View mFrameView;
 
     public MainSwitchBar(Context context) {
         this(context, null);
@@ -84,8 +84,8 @@
         setClickable(true);
 
         mFrameView = findViewById(R.id.frame);
-        mTextView = (TextView) findViewById(R.id.switch_text);
-        mSwitch = (Switch) findViewById(android.R.id.switch_widget);
+        mTextView = findViewById(R.id.switch_text);
+        mSwitch = findViewById(android.R.id.switch_widget);
         addOnSwitchChangeListener((switchView, isChecked) -> setChecked(isChecked));
 
         if (mSwitch.getVisibility() == VISIBLE) {
@@ -136,13 +136,6 @@
     }
 
     /**
-     * Return the Switch
-     */
-    public final Switch getSwitch() {
-        return mSwitch;
-    }
-
-    /**
      * Set the title text
      */
     public void setTitle(CharSequence text) {
@@ -192,7 +185,7 @@
     /**
      * Adds a listener for switch changes
      */
-    public void addOnSwitchChangeListener(OnMainSwitchChangeListener listener) {
+    public void addOnSwitchChangeListener(OnCheckedChangeListener listener) {
         if (!mSwitchChangeListeners.contains(listener)) {
             mSwitchChangeListeners.add(listener);
         }
@@ -201,7 +194,7 @@
     /**
      * Remove a listener for switch changes
      */
-    public void removeOnSwitchChangeListener(OnMainSwitchChangeListener listener) {
+    public void removeOnSwitchChangeListener(OnCheckedChangeListener listener) {
         mSwitchChangeListeners.remove(listener);
     }
 
@@ -223,9 +216,8 @@
     private void propagateChecked(boolean isChecked) {
         setBackground(isChecked);
 
-        final int count = mSwitchChangeListeners.size();
-        for (int n = 0; n < count; n++) {
-            mSwitchChangeListeners.get(n).onSwitchChanged(mSwitch, isChecked);
+        for (OnCheckedChangeListener changeListener : mSwitchChangeListeners) {
+            changeListener.onCheckedChanged(mSwitch, isChecked);
         }
     }
 
diff --git a/packages/SettingsLib/MainSwitchPreference/src/com/android/settingslib/widget/MainSwitchPreference.java b/packages/SettingsLib/MainSwitchPreference/src/com/android/settingslib/widget/MainSwitchPreference.java
index 11a6804..b294d4e 100644
--- a/packages/SettingsLib/MainSwitchPreference/src/com/android/settingslib/widget/MainSwitchPreference.java
+++ b/packages/SettingsLib/MainSwitchPreference/src/com/android/settingslib/widget/MainSwitchPreference.java
@@ -19,24 +19,25 @@
 import android.content.Context;
 import android.content.res.TypedArray;
 import android.util.AttributeSet;
-import android.widget.Switch;
+import android.widget.CompoundButton;
+import android.widget.CompoundButton.OnCheckedChangeListener;
 
 import androidx.preference.PreferenceViewHolder;
 import androidx.preference.TwoStatePreference;
 
+import com.android.settingslib.widget.mainswitch.R;
+
 import java.util.ArrayList;
 import java.util.List;
 
-import com.android.settingslib.widget.mainswitch.R;
-
 /**
  * MainSwitchPreference is a Preference with a customized Switch.
  * This component is used as the main switch of the page
  * to enable or disable the prefereces on the page.
  */
-public class MainSwitchPreference extends TwoStatePreference implements OnMainSwitchChangeListener {
+public class MainSwitchPreference extends TwoStatePreference implements OnCheckedChangeListener {
 
-    private final List<OnMainSwitchChangeListener> mSwitchChangeListeners = new ArrayList<>();
+    private final List<OnCheckedChangeListener> mSwitchChangeListeners = new ArrayList<>();
 
     private MainSwitchBar mMainSwitchBar;
 
@@ -120,7 +121,7 @@
     }
 
     @Override
-    public void onSwitchChanged(Switch switchView, boolean isChecked) {
+    public void onCheckedChanged(CompoundButton buttonView, boolean isChecked) {
         super.setChecked(isChecked);
     }
 
@@ -138,7 +139,7 @@
     /**
      * Adds a listener for switch changes
      */
-    public void addOnSwitchChangeListener(OnMainSwitchChangeListener listener) {
+    public void addOnSwitchChangeListener(OnCheckedChangeListener listener) {
         if (!mSwitchChangeListeners.contains(listener)) {
             mSwitchChangeListeners.add(listener);
         }
@@ -151,7 +152,7 @@
     /**
      * Remove a listener for switch changes
      */
-    public void removeOnSwitchChangeListener(OnMainSwitchChangeListener listener) {
+    public void removeOnSwitchChangeListener(OnCheckedChangeListener listener) {
         mSwitchChangeListeners.remove(listener);
         if (mMainSwitchBar != null) {
             mMainSwitchBar.removeOnSwitchChangeListener(listener);
@@ -159,7 +160,7 @@
     }
 
     private void registerListenerToSwitchBar() {
-        for (OnMainSwitchChangeListener listener : mSwitchChangeListeners) {
+        for (OnCheckedChangeListener listener : mSwitchChangeListeners) {
             mMainSwitchBar.addOnSwitchChangeListener(listener);
         }
     }
diff --git a/packages/SettingsLib/MainSwitchPreference/src/com/android/settingslib/widget/OnMainSwitchChangeListener.java b/packages/SettingsLib/MainSwitchPreference/src/com/android/settingslib/widget/OnMainSwitchChangeListener.java
deleted file mode 100644
index 03868f9..0000000
--- a/packages/SettingsLib/MainSwitchPreference/src/com/android/settingslib/widget/OnMainSwitchChangeListener.java
+++ /dev/null
@@ -1,32 +0,0 @@
-/*
- * 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.
- */
-
-package com.android.settingslib.widget;
-
-import android.widget.Switch;
-
-import com.android.settingslib.widget.mainswitch.R;
-
-/**
- * Called when the checked state of the Switch has changed.
- */
-public interface OnMainSwitchChangeListener {
-    /**
-     * @param switchView The Switch view whose state has changed.
-     * @param isChecked  The new checked state of switchView.
-     */
-    void onSwitchChanged(Switch switchView, boolean isChecked);
-}
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/widget/MainSwitchBarTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/widget/MainSwitchBarTest.java
index 942e915..74a282f 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/widget/MainSwitchBarTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/widget/MainSwitchBarTest.java
@@ -21,30 +21,25 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import android.content.Context;
-import android.text.TextUtils;
 import android.view.View;
-import android.widget.Switch;
+import android.widget.CompoundButton;
 import android.widget.TextView;
 
+import androidx.test.core.app.ApplicationProvider;
+
 import com.android.settingslib.widget.mainswitch.R;
 
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
 
 @RunWith(RobolectricTestRunner.class)
 public class MainSwitchBarTest {
 
-    private Context mContext;
-    private MainSwitchBar mBar;
+    private final Context mContext = ApplicationProvider.getApplicationContext();
+    private final MainSwitchBar mBar = new MainSwitchBar(mContext);
 
-    @Before
-    public void setUp() {
-        mContext = RuntimeEnvironment.application;
-        mBar = new MainSwitchBar(mContext);
-    }
+    private final CompoundButton mSwitch = mBar.findViewById(android.R.id.switch_widget);
 
     @Test
     public void setChecked_true_shouldChecked() {
@@ -60,7 +55,7 @@
         mBar.setTitle(title);
         final TextView textView = ((TextView) mBar.findViewById(R.id.switch_text));
 
-        assertThat(textView.getText()).isEqualTo(title);
+        assertThat(textView.getText().toString()).isEqualTo(title);
     }
 
     @Test
@@ -69,23 +64,18 @@
 
         mBar.setTitle(title);
 
-        final Switch switchObj = mBar.getSwitch();
-        assertThat(TextUtils.isEmpty(switchObj.getContentDescription())).isTrue();
+        assertThat(mSwitch.getContentDescription()).isNull();
     }
 
     @Test
     public void getSwitch_shouldNotNull() {
-        final Switch switchObj = mBar.getSwitch();
-
-        assertThat(switchObj).isNotNull();
+        assertThat(mSwitch).isNotNull();
     }
 
     @Test
     public void getSwitch_shouldNotFocusableAndClickable() {
-        final Switch switchObj = mBar.getSwitch();
-
-        assertThat(switchObj.isFocusable()).isFalse();
-        assertThat(switchObj.isClickable()).isFalse();
+        assertThat(mSwitch.isFocusable()).isFalse();
+        assertThat(mSwitch.isClickable()).isFalse();
     }
 
     @Test