diff --git a/src/com/android/settings/notification/NotificationAssistantPreferenceController.java b/src/com/android/settings/notification/NotificationAssistantPreferenceController.java
index 6237326f4e4..a46f1646113 100644
--- a/src/com/android/settings/notification/NotificationAssistantPreferenceController.java
+++ b/src/com/android/settings/notification/NotificationAssistantPreferenceController.java
@@ -77,11 +77,7 @@ public class NotificationAssistantPreferenceController extends TogglePreferenceC
     protected void setNotificationAssistantGranted(ComponentName cn) {
         if (Settings.Secure.getIntForUser(mContext.getContentResolver(),
                 Settings.Secure.NAS_SETTINGS_UPDATED, 0, mUserId) == 0) {
-            for (int profileId : mUserManager.getProfileIds(mUserId, false)) {
-                Settings.Secure.putIntForUser(mContext.getContentResolver(),
-                        Settings.Secure.NAS_SETTINGS_UPDATED, 1, profileId);
-            }
-            mNotificationBackend.resetDefaultNotificationAssistant(cn != null);
+            mNotificationBackend.setNASMigrationDoneAndResetDefault(mUserId, cn != null);
         }
         mNotificationBackend.setNotificationAssistantGranted(cn);
     }
diff --git a/src/com/android/settings/notification/NotificationBackend.java b/src/com/android/settings/notification/NotificationBackend.java
index b08d02cef26..e448dda20aa 100644
--- a/src/com/android/settings/notification/NotificationBackend.java
+++ b/src/com/android/settings/notification/NotificationBackend.java
@@ -570,9 +570,9 @@ public class NotificationBackend {
         }
     }
 
-    public void resetDefaultNotificationAssistant(boolean loadFromConfig) {
+    public void setNASMigrationDoneAndResetDefault(int userId, boolean loadFromConfig) {
         try {
-            sINM.resetDefaultNotificationAssistant(loadFromConfig);
+            sINM.setNASMigrationDoneAndResetDefault(userId, loadFromConfig);
         } catch (Exception e) {
             Log.w(TAG, "Error calling NoMan", e);
         }
diff --git a/tests/robotests/src/com/android/settings/notification/NotificationAssistantPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/notification/NotificationAssistantPreferenceControllerTest.java
index f6b72d963bc..9b45a9808eb 100644
--- a/tests/robotests/src/com/android/settings/notification/NotificationAssistantPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/notification/NotificationAssistantPreferenceControllerTest.java
@@ -122,13 +122,13 @@ public class NotificationAssistantPreferenceControllerTest {
         assertEquals(1, Settings.Secure.getIntForUser(mContext.getContentResolver(),
                 Settings.Secure.NAS_SETTINGS_UPDATED, 0, 10));
         verify(mBackend, times(1))
-                .resetDefaultNotificationAssistant(eq(true));
+                .setNASMigrationDoneAndResetDefault(eq(0), eq(true));
 
         //Test user enable again, migration should not happen
         mPreferenceController.setNotificationAssistantGranted(mNASComponent);
         //Number of invocations should not increase
         verify(mBackend, times(1))
-                .resetDefaultNotificationAssistant(eq(true));
+                .setNASMigrationDoneAndResetDefault(eq(0), eq(true));
     }
 
     @Test
@@ -146,13 +146,13 @@ public class NotificationAssistantPreferenceControllerTest {
         assertEquals(0, Settings.Secure.getIntForUser(mContext.getContentResolver(),
                 Settings.Secure.NAS_SETTINGS_UPDATED, 0, 20));
         verify(mBackend, times(1))
-                .resetDefaultNotificationAssistant(eq(true));
+                .setNASMigrationDoneAndResetDefault(eq(0), eq(true));
 
         //Test user enable again, migration should not happen
         mPreferenceController.setNotificationAssistantGranted(mNASComponent);
         //Number of invocations should not increase
         verify(mBackend, times(1))
-                .resetDefaultNotificationAssistant(eq(true));
+                .setNASMigrationDoneAndResetDefault(eq(0), eq(true));
     }
 
     @Test
@@ -170,13 +170,13 @@ public class NotificationAssistantPreferenceControllerTest {
         assertEquals(1, Settings.Secure.getIntForUser(mContext.getContentResolver(),
                 Settings.Secure.NAS_SETTINGS_UPDATED, 0, 10));
         verify(mBackend, times(1))
-                .resetDefaultNotificationAssistant(eq(false));
+                .setNASMigrationDoneAndResetDefault(eq(0), eq(false));
 
         //Test user disable again, migration should not happen
         mPreferenceController.setChecked(false);
         //Number of invocations should not increase
         verify(mBackend, times(1))
-                .resetDefaultNotificationAssistant(eq(false));
+                .setNASMigrationDoneAndResetDefault(eq(0), eq(false));
     }
 
 }