diff --git a/res/xml/configure_notification_settings.xml b/res/xml/configure_notification_settings.xml
index c52140c4711..59c0c218faa 100644
--- a/res/xml/configure_notification_settings.xml
+++ b/res/xml/configure_notification_settings.xml
@@ -174,7 +174,7 @@
             android:title="@string/notification_pulse_title"
             settings:controller="com.android.settings.notification.PulseNotificationPreferenceController"/>
 
-        <SwitchPreference
+        <SwitchPreferenceCompat
             android:key="notification_assistant"
             android:order="25"
             android:title="@string/notification_assistant_title"
diff --git a/src/com/android/settings/notification/NotificationAssistantPreferenceController.java b/src/com/android/settings/notification/NotificationAssistantPreferenceController.java
index 3274aec4585..6f246b730ce 100644
--- a/src/com/android/settings/notification/NotificationAssistantPreferenceController.java
+++ b/src/com/android/settings/notification/NotificationAssistantPreferenceController.java
@@ -18,44 +18,31 @@ package com.android.settings.notification;
 
 import android.content.ComponentName;
 import android.content.Context;
-import android.content.Intent;
-import android.content.pm.PackageManager;
-import android.content.pm.ResolveInfo;
 import android.os.UserHandle;
-import android.os.UserManager;
 import android.provider.Settings;
-import android.service.notification.NotificationAssistantService;
 
 import androidx.fragment.app.Fragment;
 import androidx.preference.Preference;
 
 import com.android.settings.R;
 import com.android.settings.core.TogglePreferenceController;
-import com.android.settingslib.PrimarySwitchPreference;
 
 import com.google.common.annotations.VisibleForTesting;
 
-import java.util.List;
-
 public class NotificationAssistantPreferenceController extends TogglePreferenceController {
     private static final String TAG = "NASPreferenceController";
     static final String KEY_NAS = "notification_assistant";
 
-    private final UserManager mUserManager;
-    private final PackageManager mPackageManager;
     private Fragment mFragment;
     private int mUserId = UserHandle.myUserId();
 
     @VisibleForTesting
     protected NotificationBackend mNotificationBackend;
     private ComponentName mDefaultNASComponent;
-    private Intent mNASSettingIntent;
 
     public NotificationAssistantPreferenceController(Context context) {
         super(context, KEY_NAS);
-        mUserManager = UserManager.get(context);
         mNotificationBackend = new NotificationBackend();
-        mPackageManager = context.getPackageManager();
         getDefaultNASIntent();
     }
 
@@ -118,12 +105,6 @@ public class NotificationAssistantPreferenceController extends TogglePreferenceC
     @VisibleForTesting
     void getDefaultNASIntent() {
         mDefaultNASComponent = mNotificationBackend.getDefaultNotificationAssistant();
-        if (mDefaultNASComponent != null) {
-            mNASSettingIntent = new Intent(
-                    NotificationAssistantService.ACTION_NOTIFICATION_ASSISTANT_DETAIL_SETTINGS);
-            mNASSettingIntent.setPackage(mDefaultNASComponent.getPackageName());
-            mNASSettingIntent.addFlags(Intent.FLAG_ACTIVITY_NEW_TASK);
-        }
     }
 
     @Override
@@ -136,7 +117,6 @@ public class NotificationAssistantPreferenceController extends TogglePreferenceC
         super.updateState(preference);
         if (mDefaultNASComponent == null) {
             preference.setEnabled(false);
-            ((PrimarySwitchPreference) preference).setSwitchEnabled(false);
         }
     }
 }
diff --git a/tests/robotests/src/com/android/settings/notification/NotificationAssistantPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/notification/NotificationAssistantPreferenceControllerTest.java
index 5f506b8ac0b..76060570229 100644
--- a/tests/robotests/src/com/android/settings/notification/NotificationAssistantPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/notification/NotificationAssistantPreferenceControllerTest.java
@@ -16,18 +16,12 @@
 
 package com.android.settings.notification;
 
-import static android.service.notification.NotificationAssistantService.ACTION_NOTIFICATION_ASSISTANT_DETAIL_SETTINGS;
-
-import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
@@ -35,24 +29,19 @@ import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import android.app.Application;
 import android.content.ComponentName;
 import android.content.Context;
-import android.content.Intent;
-import android.content.pm.ActivityInfo;
-import android.content.pm.PackageManager;
-import android.content.pm.ResolveInfo;
-import android.os.UserManager;
 import android.provider.Settings;
 
 import androidx.fragment.app.FragmentManager;
 import androidx.fragment.app.FragmentTransaction;
 import androidx.preference.PreferenceManager;
 import androidx.preference.PreferenceScreen;
+import androidx.preference.SwitchPreferenceCompat;
+import androidx.preference.TwoStatePreference;
 import androidx.test.core.app.ApplicationProvider;
 
 import com.android.settings.testutils.shadow.ShadowSecureSettings;
-import com.android.settingslib.PrimarySwitchPreference;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -61,12 +50,7 @@ import org.mockito.Answers;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
-import org.robolectric.Shadows;
 import org.robolectric.annotation.Config;
-import org.robolectric.shadows.ShadowApplication;
-
-import java.util.ArrayList;
-import java.util.List;
 
 
 @RunWith(RobolectricTestRunner.class)
@@ -86,27 +70,19 @@ public class NotificationAssistantPreferenceControllerTest {
     private FragmentTransaction mFragmentTransaction;
     @Mock
     private NotificationBackend mBackend;
-    @Mock
-    private UserManager mUserManager;
-    @Mock
-    private PackageManager mPackageManager;
     private NotificationAssistantPreferenceController mPreferenceController;
     ComponentName mNASComponent = new ComponentName("pkgname", "clsname");
-    private PrimarySwitchPreference mPreference;
-    private ShadowApplication mShadowApplication;
+    private TwoStatePreference mPreference;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         mContext = spy(ApplicationProvider.getApplicationContext());
-        mPreference = spy(new PrimarySwitchPreference(mContext));
-        mShadowApplication = ShadowApplication.getInstance();
-        mShadowApplication.setSystemService(Context.USER_SERVICE, mUserManager);
+        mPreference = spy(new SwitchPreferenceCompat(mContext));
         doReturn(mContext).when(mFragment).getContext();
         when(mFragment.getFragmentManager()).thenReturn(mFragmentManager);
         when(mFragmentManager.beginTransaction()).thenReturn(mFragmentTransaction);
         when(mBackend.getDefaultNotificationAssistant()).thenReturn(mNASComponent);
-        when(mContext.getPackageManager()).thenReturn(mPackageManager);
         mPreferenceController = new NotificationAssistantPreferenceController(mContext);
         mPreferenceController.setBackend(mBackend);
         mPreferenceController.setFragment(mFragment);
@@ -117,19 +93,6 @@ public class NotificationAssistantPreferenceControllerTest {
         mPreference.setKey(NotificationAssistantPreferenceController.KEY_NAS);
         screen.addPreference(mPreference);
         mPreferenceController.displayPreference(screen);
-
-        when(mUserManager.getProfileIds(eq(0), anyBoolean())).thenReturn(new int[] {0, 10});
-        when(mUserManager.getProfileIds(eq(20), anyBoolean())).thenReturn(new int[] {20});
-
-        ActivityInfo activityInfo1 = new ActivityInfo();
-        activityInfo1.packageName = "pkgname";
-        activityInfo1.name = "name";
-        ResolveInfo resolveInfo1 = new ResolveInfo();
-        resolveInfo1.activityInfo = activityInfo1;
-        List<ResolveInfo> resolvers1 = new ArrayList<>();
-        resolvers1.add(resolveInfo1);
-        when(mPackageManager.queryIntentActivities(any(Intent.class), any()))
-                .thenReturn(resolvers1);
     }
 
     @Test