diff --git a/AndroidManifest.xml b/AndroidManifest.xml
index e678983c972..da438473a57 100644
--- a/AndroidManifest.xml
+++ b/AndroidManifest.xml
@@ -2631,7 +2631,7 @@
                 <category android:name="android.intent.category.DEFAULT" />
             </intent-filter>
             <meta-data android:name="com.android.settings.FRAGMENT_CLASS"
-                       android:value="com.android.settings.notification.NotificationAssistantPicker" />
+                       android:value="com.android.settings.notification.ConfigureNotificationSettings" />
         </activity>
 
         <activity
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 655edb1f61a..2e525a3edbe 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -9032,8 +9032,10 @@
         <item quantity="other">%d apps can read notifications</item>
     </plurals>
 
-    <!-- Title for Notification Assistant Picker screen [CHAR LIMIT=30]-->
-    <string name="notification_assistant_title">Adaptive Notifications</string>
+    <!-- Title for Notification Assistant setting  [CHAR LIMIT=30]-->
+    <string name="notification_assistant_title">Enhanced notifications</string>
+    <!-- Summary of Notification Assistant provided features [CHAR LIMIT=NONE]-->
+    <string name="notification_assistant_summary">Get suggested actions, replies, and more</string>
 
     <!-- Label for no NotificationAssistantService [CHAR_LIMIT=NONE] -->
     <string name="no_notification_assistant">None</string>
@@ -9051,10 +9053,11 @@
     <!-- Summary for a warning message about security implications of enabling a notification
           listener, displayed as a dialog message. [CHAR LIMIT=NONE] -->
     <string name="notification_assistant_security_warning_summary">
-        <xliff:g id="notification_assistant_name" example="Notification Assistant">%1$s</xliff:g> will be able to read all notifications,
-        including personal information such as contact names and the text of messages you receive.
-        This app will also be able to dismiss notifications or take action on buttons in notifications, including answering phone calls.
-        \n\nThis will also give the app the ability to turn Do Not Disturb on or off and change related settings.
+        Enhanced notifications can read all notification content,
+        including personal information like contact names and messages.
+        This feature can also dismiss notifications or take actions on buttons in notifications,
+        such as answering phone calls.
+        \n\nThis feature can also turn Priority mode on or off and change related settings.
     </string>
 
     <!-- Title for a warning message about security implications of enabling a notification
diff --git a/res/xml/configure_notification_settings.xml b/res/xml/configure_notification_settings.xml
index c7925fb6945..529fbea4312 100644
--- a/res/xml/configure_notification_settings.xml
+++ b/res/xml/configure_notification_settings.xml
@@ -115,6 +115,11 @@
             android:title="@string/snooze_options_title"
             settings:controller="com.android.settings.notification.SnoozeNotificationPreferenceController" />
 
+        <SwitchPreference
+            android:key="notification_assistant"
+            android:title="@string/notification_assistant_title"
+            android:summary="@string/notification_assistant_summary"/>
+
         <!-- Notification badging -->
         <SwitchPreference
             android:key="notification_badging"
diff --git a/res/xml/configure_notification_settings_v2.xml b/res/xml/configure_notification_settings_v2.xml
index b7cc2c86a0d..98768f87be5 100644
--- a/res/xml/configure_notification_settings_v2.xml
+++ b/res/xml/configure_notification_settings_v2.xml
@@ -155,5 +155,11 @@
             android:order="22"
             android:title="@string/notification_pulse_title"
             settings:controller="com.android.settings.notification.PulseNotificationPreferenceController"/>
+
+        <SwitchPreference
+            android:key="notification_assistant"
+            android:order="23"
+            android:title="@string/notification_assistant_title"
+            android:summary="@string/notification_assistant_summary"/>
     </PreferenceCategory>
 </PreferenceScreen>
diff --git a/res/xml/special_access.xml b/res/xml/special_access.xml
index 83c23b5305a..892c3eb7455 100644
--- a/res/xml/special_access.xml
+++ b/res/xml/special_access.xml
@@ -80,13 +80,6 @@
             android:value="com.android.settings.Settings$WriteSettingsActivity" />
     </Preference>
 
-    <com.android.settingslib.widget.AppPreference
-        android:key="notification_assistant"
-        android:title="@string/notification_assistant_title"
-        android:summary="@string/summary_placeholder"
-        settings:fragment="com.android.settings.notification.NotificationAssistantPicker"
-        settings:controller="com.android.settings.notification.NotificationAssistantPreferenceController"/>
-
     <Preference
         android:key="notification_access"
         android:title="@string/manage_notification_access_title"
diff --git a/src/com/android/settings/notification/ConfigureNotificationSettings.java b/src/com/android/settings/notification/ConfigureNotificationSettings.java
index dcba273c6a2..22b4311e1fa 100644
--- a/src/com/android/settings/notification/ConfigureNotificationSettings.java
+++ b/src/com/android/settings/notification/ConfigureNotificationSettings.java
@@ -22,6 +22,7 @@ import android.app.Activity;
 import android.app.Application;
 import android.app.settings.SettingsEnums;
 import android.app.usage.IUsageStatsManager;
+import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
 import android.os.Bundle;
@@ -62,6 +63,7 @@ public class ConfigureNotificationSettings extends DashboardFragment implements
     private static final int REQUEST_CODE = 200;
     private static final String SELECTED_PREFERENCE_KEY = "selected_preference";
     private static final String KEY_ADVANCED_CATEGORY = "configure_notifications_advanced";
+    private static final String KEY_NAS = "notification_assistant";
 
     private RingtonePreference mRequestPreference;
 
@@ -116,6 +118,8 @@ public class ConfigureNotificationSettings extends DashboardFragment implements
             }
 
         });
+        controllers.add(new NotificationAssistantPreferenceController(context,
+                new NotificationBackend(), host, KEY_NAS));
 
         if (FeatureFlagUtils.isEnabled(context, FeatureFlags.SILKY_HOME)) {
             controllers.add(new EmergencyBroadcastPreferenceController(context,
@@ -199,4 +203,14 @@ public class ConfigureNotificationSettings extends DashboardFragment implements
                     return keys;
                 }
             };
+
+    // Dialogs only have access to the parent fragment, not the controller, so pass the information
+    // along to keep business logic out of this file
+    protected void enableNAS(ComponentName cn) {
+        final PreferenceScreen screen = getPreferenceScreen();
+        NotificationAssistantPreferenceController napc =
+                use(NotificationAssistantPreferenceController.class);
+        napc.setNotificationAssistantGranted(cn);
+        napc.updateState(screen.findPreference(napc.getPreferenceKey()));
+    }
 }
diff --git a/src/com/android/settings/notification/NotificationAssistantDialogFragment.java b/src/com/android/settings/notification/NotificationAssistantDialogFragment.java
new file mode 100644
index 00000000000..48b1cd988ea
--- /dev/null
+++ b/src/com/android/settings/notification/NotificationAssistantDialogFragment.java
@@ -0,0 +1,71 @@
+/*
+ * Copyright (C) 2021 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.notification;
+
+import android.app.Dialog;
+import android.app.settings.SettingsEnums;
+import android.content.ComponentName;
+import android.content.DialogInterface;
+import android.os.Bundle;
+
+import androidx.appcompat.app.AlertDialog;
+import androidx.fragment.app.Fragment;
+
+import com.android.settings.R;
+import com.android.settings.core.instrumentation.InstrumentedDialogFragment;
+
+public class NotificationAssistantDialogFragment extends InstrumentedDialogFragment
+        implements DialogInterface.OnClickListener {
+    static final String KEY_COMPONENT = "c";
+
+    public static NotificationAssistantDialogFragment newInstance(Fragment target,
+            ComponentName cn) {
+        final NotificationAssistantDialogFragment dialogFragment =
+                new NotificationAssistantDialogFragment();
+        final Bundle args = new Bundle();
+        args.putString(KEY_COMPONENT, cn.flattenToString());
+        dialogFragment.setArguments(args);
+        dialogFragment.setTargetFragment(target, 0);
+
+        return dialogFragment;
+    }
+
+    @Override
+    public Dialog onCreateDialog(Bundle savedInstanceState) {
+        final String summary = getResources()
+                .getString(R.string.notification_assistant_security_warning_summary);
+        return new AlertDialog.Builder(getContext())
+                .setMessage(summary)
+                .setCancelable(true)
+                .setPositiveButton(R.string.okay, this)
+                .create();
+    }
+
+    @Override
+    public int getMetricsCategory() {
+        return SettingsEnums.DEFAULT_NOTIFICATION_ASSISTANT;
+    }
+
+    @Override
+    public void onClick(DialogInterface dialog, int which) {
+        final Bundle args = getArguments();
+        final ComponentName cn = ComponentName.unflattenFromString(args
+                .getString(KEY_COMPONENT));
+        ConfigureNotificationSettings parent = (ConfigureNotificationSettings) getTargetFragment();
+        parent.enableNAS(cn);
+    }
+}
diff --git a/src/com/android/settings/notification/NotificationAssistantPreferenceController.java b/src/com/android/settings/notification/NotificationAssistantPreferenceController.java
index 66f27fef953..637e4b0d435 100644
--- a/src/com/android/settings/notification/NotificationAssistantPreferenceController.java
+++ b/src/com/android/settings/notification/NotificationAssistantPreferenceController.java
@@ -18,44 +18,72 @@ package com.android.settings.notification;
 
 import android.content.ComponentName;
 import android.content.Context;
-import android.content.pm.PackageManager;
 import android.os.UserHandle;
+import android.provider.Settings;
 
-import com.android.settings.core.BasePreferenceController;
-import com.android.settingslib.applications.DefaultAppInfo;
-import com.android.settingslib.widget.CandidateInfo;
+import androidx.fragment.app.Fragment;
+
+import com.android.settings.core.TogglePreferenceController;
 
 import com.google.common.annotations.VisibleForTesting;
 
-public class NotificationAssistantPreferenceController extends BasePreferenceController {
+public class NotificationAssistantPreferenceController extends TogglePreferenceController {
+    private static final String TAG = "NASPreferenceController";
+    private static final int AVAILABLE = 1;
+    private Fragment mFragment;
+    private int mUserId = UserHandle.myUserId();
 
     @VisibleForTesting
     protected NotificationBackend mNotificationBackend;
-    private PackageManager mPackageManager;
 
-    public NotificationAssistantPreferenceController(Context context, String preferenceKey) {
+    public NotificationAssistantPreferenceController(Context context, NotificationBackend backend,
+            Fragment fragment, String preferenceKey) {
         super(context, preferenceKey);
-        mNotificationBackend = new NotificationBackend();
-        mPackageManager = mContext.getPackageManager();
+        mNotificationBackend = backend;
+        mFragment = fragment;
     }
 
     @Override
     public int getAvailabilityStatus() {
-        return BasePreferenceController.AVAILABLE;
+        return AVAILABLE;
     }
 
     @Override
-    public CharSequence getSummary() {
-        CandidateInfo appSelected = new NotificationAssistantPicker.CandidateNone(mContext);
-        ComponentName assistant = mNotificationBackend.getAllowedNotificationAssistant();
-        if (assistant != null) {
-            appSelected = createCandidateInfo(assistant);
-        }
-        return appSelected.loadLabel();
+    public boolean isChecked() {
+        ComponentName acn = mNotificationBackend.getAllowedNotificationAssistant();
+        ComponentName dcn = mNotificationBackend.getDefaultNotificationAssistant();
+        return (acn != null && acn.equals(dcn));
     }
 
-    @VisibleForTesting
-    protected CandidateInfo createCandidateInfo(ComponentName cn) {
-        return new DefaultAppInfo(mContext, mPackageManager, UserHandle.myUserId(), cn);
+    @Override
+    public boolean setChecked(boolean isChecked) {
+        ComponentName cn = isChecked
+                ? mNotificationBackend.getDefaultNotificationAssistant() : null;
+        if (isChecked) {
+            if (mFragment == null) {
+                throw new IllegalStateException("No fragment to start activity");
+            }
+            showDialog(cn);
+            return false;
+        } else {
+            setNotificationAssistantGranted(null);
+            return true;
+        }
     }
-}
+
+    protected void setNotificationAssistantGranted(ComponentName cn) {
+        if (Settings.Secure.getIntForUser(mContext.getContentResolver(),
+                Settings.Secure.NAS_SETTINGS_UPDATED, 0, mUserId) == 0) {
+            Settings.Secure.putIntForUser(mContext.getContentResolver(),
+                    Settings.Secure.NAS_SETTINGS_UPDATED, 1, mUserId);
+            mNotificationBackend.resetDefaultNotificationAssistant(cn != null);
+        }
+        mNotificationBackend.setNotificationAssistantGranted(cn);
+    }
+
+    protected void showDialog(ComponentName cn) {
+        NotificationAssistantDialogFragment dialogFragment =
+                NotificationAssistantDialogFragment.newInstance(mFragment, cn);
+        dialogFragment.show(mFragment.getFragmentManager(), TAG);
+    }
+}
\ No newline at end of file
diff --git a/src/com/android/settings/notification/NotificationBackend.java b/src/com/android/settings/notification/NotificationBackend.java
index 4347ca55104..b08d02cef26 100644
--- a/src/com/android/settings/notification/NotificationBackend.java
+++ b/src/com/android/settings/notification/NotificationBackend.java
@@ -19,7 +19,6 @@ import static android.app.NotificationManager.IMPORTANCE_NONE;
 import static android.app.NotificationManager.IMPORTANCE_UNSPECIFIED;
 import static android.content.pm.LauncherApps.ShortcutQuery.FLAG_MATCH_CACHED;
 import static android.content.pm.LauncherApps.ShortcutQuery.FLAG_MATCH_DYNAMIC;
-import static android.content.pm.LauncherApps.ShortcutQuery.FLAG_MATCH_PINNED;
 import static android.content.pm.LauncherApps.ShortcutQuery.FLAG_MATCH_PINNED_BY_ANY_LAUNCHER;
 
 import android.app.INotificationManager;
@@ -50,7 +49,6 @@ import android.service.notification.NotificationListenerFilter;
 import android.text.format.DateUtils;
 import android.util.IconDrawableFactory;
 import android.util.Log;
-import android.util.Slog;
 
 import androidx.annotation.VisibleForTesting;
 
@@ -563,6 +561,23 @@ public class NotificationBackend {
         }
     }
 
+    public ComponentName getDefaultNotificationAssistant() {
+        try {
+            return sINM.getDefaultNotificationAssistant();
+        } catch (Exception e) {
+            Log.w(TAG, "Error calling NoMan", e);
+            return null;
+        }
+    }
+
+    public void resetDefaultNotificationAssistant(boolean loadFromConfig) {
+        try {
+            sINM.resetDefaultNotificationAssistant(loadFromConfig);
+        } catch (Exception e) {
+            Log.w(TAG, "Error calling NoMan", e);
+        }
+    }
+
     public boolean setNotificationAssistantGranted(ComponentName cn) {
         try {
             sINM.setNotificationAssistantAccessGranted(cn, true);
diff --git a/tests/robotests/src/com/android/settings/notification/NotificationAssistantDialogFragmentTest.java b/tests/robotests/src/com/android/settings/notification/NotificationAssistantDialogFragmentTest.java
new file mode 100644
index 00000000000..eef3f043bf8
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/notification/NotificationAssistantDialogFragmentTest.java
@@ -0,0 +1,69 @@
+/*
+ * Copyright (C) 2021 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.notification;
+
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import android.content.ComponentName;
+import android.content.Context;
+import android.content.DialogInterface;
+
+import androidx.fragment.app.FragmentActivity;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+
+@RunWith(RobolectricTestRunner.class)
+public class NotificationAssistantDialogFragmentTest {
+
+    private Context mContext;
+    @Mock
+    private ConfigureNotificationSettings mFragment;
+    private NotificationAssistantDialogFragment mDialogFragment;
+    @Mock
+    private FragmentActivity mActivity;
+
+    ComponentName mComponentName = new ComponentName("a", "b");
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mContext = spy(RuntimeEnvironment.application);
+        mDialogFragment =
+                spy(NotificationAssistantDialogFragment.newInstance(mFragment, mComponentName));
+        doReturn(mActivity).when(mDialogFragment).getActivity();
+        doReturn(mContext).when(mDialogFragment).getContext();
+
+    }
+
+
+    @Test
+    public void testClickOK_callEnableNAS() {
+        mDialogFragment.onClick(null, DialogInterface.BUTTON_POSITIVE);
+
+        verify(mFragment, times(1)).enableNAS(eq(mComponentName));
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/notification/NotificationAssistantPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/notification/NotificationAssistantPreferenceControllerTest.java
index b2f65e06d7b..4f2145cb07b 100644
--- a/tests/robotests/src/com/android/settings/notification/NotificationAssistantPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/notification/NotificationAssistantPreferenceControllerTest.java
@@ -16,17 +16,25 @@
 
 package com.android.settings.notification;
 
-import static junit.framework.TestCase.assertEquals;
-
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.content.ComponentName;
 import android.content.Context;
-import android.content.pm.PackageManager;
-import android.graphics.drawable.Drawable;
-import android.os.Debug;
+import android.provider.Settings;
 
-import com.android.settingslib.widget.CandidateInfo;
+import androidx.fragment.app.FragmentManager;
+import androidx.fragment.app.FragmentTransaction;
+import androidx.test.core.app.ApplicationProvider;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -35,7 +43,6 @@ import org.mockito.Answers;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
 
 @RunWith(RobolectricTestRunner.class)
 public class NotificationAssistantPreferenceControllerTest {
@@ -44,57 +51,86 @@ public class NotificationAssistantPreferenceControllerTest {
     @Mock(answer = Answers.RETURNS_DEEP_STUBS)
     private Context mContext;
     @Mock
+    private ConfigureNotificationSettings mFragment;
+    @Mock
+    private FragmentManager mFragmentManager;
+    @Mock
+    private FragmentTransaction mFragmentTransaction;
+    @Mock
     private NotificationBackend mBackend;
     private NotificationAssistantPreferenceController mPreferenceController;
+    ComponentName mNASComponent = new ComponentName("a", "b");
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mContext = RuntimeEnvironment.application;
-        mPreferenceController = new TestPreferenceController(mContext, mBackend);
+        mContext = spy(ApplicationProvider.getApplicationContext());
+        doReturn(mContext).when(mFragment).getContext();
+        when(mFragment.getFragmentManager()).thenReturn(mFragmentManager);
+        when(mFragmentManager.beginTransaction()).thenReturn(mFragmentTransaction);
+        when(mBackend.getDefaultNotificationAssistant()).thenReturn(mNASComponent);
+        mPreferenceController = new NotificationAssistantPreferenceController(mContext,
+                mBackend, mFragment, KEY);
     }
 
     @Test
-    public void testGetSummary_noAssistant() {
+    public void testIsChecked() throws Exception {
+        when(mBackend.getAllowedNotificationAssistant()).thenReturn(mNASComponent);
+        assertTrue(mPreferenceController.isChecked());
+
         when(mBackend.getAllowedNotificationAssistant()).thenReturn(null);
-        CharSequence noneLabel = new NotificationAssistantPicker.CandidateNone(mContext)
-                .loadLabel();
-        assertEquals(noneLabel, mPreferenceController.getSummary());
+        assertFalse(mPreferenceController.isChecked());
     }
 
     @Test
-    public void testGetSummary_TestAssistant() {
-        String testName = "test_pkg/test_cls";
-        when(mBackend.getAllowedNotificationAssistant()).thenReturn(
-                ComponentName.unflattenFromString(testName));
-        assertEquals(testName, mPreferenceController.getSummary());
+    public void testSetChecked() throws Exception {
+        // Verify a dialog is shown when the switch is to be enabled.
+        assertFalse(mPreferenceController.setChecked(true));
+        verify(mFragmentTransaction).add(
+                any(NotificationAssistantDialogFragment.class), anyString());
+        verify(mBackend, times(0)).setNotificationAssistantGranted(any());
+
+        // Verify no dialog is shown and NAS set to null when disabled
+        assertTrue(mPreferenceController.setChecked(false));
+        verify(mBackend, times(1)).setNotificationAssistantGranted(null);
     }
 
-    private final class TestPreferenceController extends NotificationAssistantPreferenceController {
+    @Test
+    public void testMigrationFromSetting_userEnable() throws Exception {
+        Settings.Secure.putIntForUser(mContext.getContentResolver(),
+                Settings.Secure.NAS_SETTINGS_UPDATED, 0, 0);
 
-        private TestPreferenceController(Context context, NotificationBackend backend) {
-            super(context, KEY);
-            mNotificationBackend = backend;
-        }
+        //Test user enable for the first time
+        mPreferenceController.setNotificationAssistantGranted(mNASComponent);
+        assertEquals(1, Settings.Secure.getIntForUser(mContext.getContentResolver(),
+                Settings.Secure.NAS_SETTINGS_UPDATED, 0, 0));
+        verify(mBackend, times(1))
+                .resetDefaultNotificationAssistant(eq(true));
 
-        @Override
-        public String getPreferenceKey() {
-            return KEY;
-        }
+        //Test user enable again, migration should not happen
+        mPreferenceController.setNotificationAssistantGranted(mNASComponent);
+        //Number of invocations should not increase
+        verify(mBackend, times(1))
+                .resetDefaultNotificationAssistant(eq(true));
+    }
 
-        @Override
-        protected CandidateInfo createCandidateInfo(ComponentName cn) {
-            return new CandidateInfo(true) {
-                @Override
-                public CharSequence loadLabel() { return cn.flattenToString(); }
+    @Test
+    public void testMigrationFromSetting_userDisable() throws Exception {
+        Settings.Secure.putIntForUser(mContext.getContentResolver(),
+                Settings.Secure.NAS_SETTINGS_UPDATED, 0, 0);
 
-                @Override
-                public Drawable loadIcon() { return null; }
+        //Test user disable for the first time
+        mPreferenceController.setChecked(false);
+        assertEquals(1, Settings.Secure.getIntForUser(mContext.getContentResolver(),
+                Settings.Secure.NAS_SETTINGS_UPDATED, 0, 0));
+        verify(mBackend, times(1))
+                .resetDefaultNotificationAssistant(eq(false));
 
-                @Override
-                public String getKey() { return null; }
-            };
-        }
+        //Test user disable again, migration should not happen
+        mPreferenceController.setChecked(false);
+        //Number of invocations should not increase
+        verify(mBackend, times(1))
+                .resetDefaultNotificationAssistant(eq(false));
     }
 
 }