Update ExtService to use suggestConversationActions

1. Use suggestConversationActions for both replies and actions
2. Make existing flags configurable via DeviceConfig

Test: atest SmartActionHelperTest.java
Test: atest AssistantSettingsTest.java

BUG: 123745079

Change-Id: I9b84edf9818d5839202c337ac3c3d48378adbf55

7 files changed, 382 insertions, 260 deletions

diff --git a/api/system-current.txt b/api/system-current.txt
index 4cb5eb0c343c..391728b39939 100644
--- a/api/system-current.txt
+++ b/api/system-current.txt
@@ -5753,6 +5753,8 @@ package android.provider {
   public static interface DeviceConfig.NotificationAssistant {
     field public static final String GENERATE_ACTIONS = "generate_actions";
     field public static final String GENERATE_REPLIES = "generate_replies";
+    field public static final String MAX_MESSAGES_TO_EXTRACT = "max_messages_to_extract";
+    field public static final String MAX_SUGGESTIONS = "max_suggestions";
     field public static final String NAMESPACE = "notification_assistant";
   }
 
diff --git a/core/java/android/provider/DeviceConfig.java b/core/java/android/provider/DeviceConfig.java
index ed3026c03810..dbf6b6c9be72 100644
--- a/core/java/android/provider/DeviceConfig.java
+++ b/core/java/android/provider/DeviceConfig.java
@@ -147,6 +147,10 @@ public final class DeviceConfig {
          * Whether the Notification Assistant should generate contextual actions for notifications.
          */
         String GENERATE_ACTIONS = "generate_actions";
+
+        String MAX_MESSAGES_TO_EXTRACT = "max_messages_to_extract";
+
+        String MAX_SUGGESTIONS = "max_suggestions";
     }
 
     /**
diff --git a/packages/ExtServices/src/android/ext/services/notification/AssistantSettings.java b/packages/ExtServices/src/android/ext/services/notification/AssistantSettings.java
index d99c356b8194..6cf72a45a06b 100644
--- a/packages/ExtServices/src/android/ext/services/notification/AssistantSettings.java
+++ b/packages/ExtServices/src/android/ext/services/notification/AssistantSettings.java
@@ -23,7 +23,6 @@ import android.os.Handler;
 import android.provider.DeviceConfig;
 import android.provider.Settings;
 import android.text.TextUtils;
-import android.util.KeyValueListParser;
 import android.util.Log;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -37,6 +36,9 @@ final class AssistantSettings extends ContentObserver {
     private static final boolean DEFAULT_GENERATE_REPLIES = true;
     private static final boolean DEFAULT_GENERATE_ACTIONS = true;
     private static final int DEFAULT_NEW_INTERRUPTION_MODEL_INT = 1;
+    private static final int DEFAULT_MAX_MESSAGES_TO_EXTRACT = 5;
+    @VisibleForTesting
+    static final int DEFAULT_MAX_SUGGESTIONS = 3;
 
     private static final Uri STREAK_LIMIT_URI =
             Settings.Global.getUriFor(Settings.Global.BLOCKING_HELPER_STREAK_LIMIT);
@@ -46,7 +48,6 @@ final class AssistantSettings extends ContentObserver {
     private static final Uri NOTIFICATION_NEW_INTERRUPTION_MODEL_URI =
             Settings.Secure.getUriFor(Settings.Secure.NOTIFICATION_NEW_INTERRUPTION_MODEL);
 
-    private final KeyValueListParser mParser = new KeyValueListParser(',');
     private final ContentResolver mResolver;
     private final int mUserId;
 
@@ -55,12 +56,14 @@ final class AssistantSettings extends ContentObserver {
     @VisibleForTesting
     protected final Runnable mOnUpdateRunnable;
 
-    // Actuall configuration settings.
+    // Actual configuration settings.
     float mDismissToViewRatioLimit;
     int mStreakLimit;
     boolean mGenerateReplies = DEFAULT_GENERATE_REPLIES;
     boolean mGenerateActions = DEFAULT_GENERATE_ACTIONS;
     boolean mNewInterruptionModel;
+    int mMaxMessagesToExtract = DEFAULT_MAX_MESSAGES_TO_EXTRACT;
+    int mMaxSuggestions = DEFAULT_MAX_SUGGESTIONS;
 
     private AssistantSettings(Handler handler, ContentResolver resolver, int userId,
             Runnable onUpdateRunnable) {
@@ -124,27 +127,18 @@ final class AssistantSettings extends ContentObserver {
     }
 
     private void updateFromDeviceConfigFlags() {
-        String generateRepliesFlag = DeviceConfig.getProperty(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_REPLIES);
-        if (TextUtils.isEmpty(generateRepliesFlag)) {
-            mGenerateReplies = DEFAULT_GENERATE_REPLIES;
-        } else {
-            // parseBoolean returns false for everything that isn't 'true' so there's no need to
-            // sanitise the flag string here.
-            mGenerateReplies = Boolean.parseBoolean(generateRepliesFlag);
-        }
+        mGenerateReplies = DeviceConfigHelper.getBoolean(
+                DeviceConfig.NotificationAssistant.GENERATE_REPLIES, DEFAULT_GENERATE_REPLIES);
 
-        String generateActionsFlag = DeviceConfig.getProperty(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_ACTIONS);
-        if (TextUtils.isEmpty(generateActionsFlag)) {
-            mGenerateActions = DEFAULT_GENERATE_ACTIONS;
-        } else {
-            // parseBoolean returns false for everything that isn't 'true' so there's no need to
-            // sanitise the flag string here.
-            mGenerateActions = Boolean.parseBoolean(generateActionsFlag);
-        }
+        mGenerateActions = DeviceConfigHelper.getBoolean(
+                DeviceConfig.NotificationAssistant.GENERATE_ACTIONS, DEFAULT_GENERATE_ACTIONS);
+
+        mMaxMessagesToExtract = DeviceConfigHelper.getInteger(
+                DeviceConfig.NotificationAssistant.MAX_MESSAGES_TO_EXTRACT,
+                DEFAULT_MAX_MESSAGES_TO_EXTRACT);
+
+        mMaxSuggestions = DeviceConfigHelper.getInteger(
+                DeviceConfig.NotificationAssistant.MAX_SUGGESTIONS, DEFAULT_MAX_SUGGESTIONS);
 
         mOnUpdateRunnable.run();
     }
@@ -175,8 +169,37 @@ final class AssistantSettings extends ContentObserver {
         mOnUpdateRunnable.run();
     }
 
+    static class DeviceConfigHelper {
+
+        static int getInteger(String key, int defaultValue) {
+            String value = getValue(key);
+            if (TextUtils.isEmpty(value)) {
+                return defaultValue;
+            }
+            try {
+                return Integer.parseInt(value);
+            } catch (NumberFormatException ex) {
+                return defaultValue;
+            }
+        }
+
+        static boolean getBoolean(String key, boolean defaultValue) {
+            String value = getValue(key);
+            if (TextUtils.isEmpty(value)) {
+                return defaultValue;
+            }
+            return Boolean.parseBoolean(value);
+        }
+
+        private static String getValue(String key) {
+            return DeviceConfig.getProperty(
+                    DeviceConfig.NotificationAssistant.NAMESPACE,
+                    key);
+        }
+    }
+
     public interface Factory {
         AssistantSettings createAndRegister(Handler handler, ContentResolver resolver, int userId,
                 Runnable onUpdateRunnable);
     }
-}
+}
\ No newline at end of file
diff --git a/packages/ExtServices/src/android/ext/services/notification/SmartActionsHelper.java b/packages/ExtServices/src/android/ext/services/notification/SmartActionsHelper.java
index 48a3974e913b..08cc39fc4935 100644
--- a/packages/ExtServices/src/android/ext/services/notification/SmartActionsHelper.java
+++ b/packages/ExtServices/src/android/ext/services/notification/SmartActionsHelper.java
@@ -29,28 +29,22 @@ import android.text.TextUtils;
 import android.util.LruCache;
 import android.view.textclassifier.ConversationAction;
 import android.view.textclassifier.ConversationActions;
-import android.view.textclassifier.TextClassification;
 import android.view.textclassifier.TextClassificationContext;
 import android.view.textclassifier.TextClassificationManager;
 import android.view.textclassifier.TextClassifier;
 import android.view.textclassifier.TextClassifierEvent;
-import android.view.textclassifier.TextLinks;
 
 import java.time.Instant;
 import java.time.ZoneOffset;
 import java.time.ZonedDateTime;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collections;
 import java.util.Deque;
 import java.util.List;
 import java.util.stream.Collectors;
 
 public class SmartActionsHelper {
-    private static final ArrayList<Notification.Action> EMPTY_ACTION_LIST = new ArrayList<>();
-    private static final ArrayList<CharSequence> EMPTY_REPLY_LIST = new ArrayList<>();
-
     private static final String KEY_ACTION_TYPE = "action_type";
     // If a notification has any of these flags set, it's inelgibile for actions being added.
     private static final int FLAG_MASK_INELGIBILE_FOR_ACTIONS =
@@ -58,19 +52,8 @@ public class SmartActionsHelper {
                     | Notification.FLAG_FOREGROUND_SERVICE
                     | Notification.FLAG_GROUP_SUMMARY
                     | Notification.FLAG_NO_CLEAR;
-    private static final int MAX_ACTION_EXTRACTION_TEXT_LENGTH = 400;
-    private static final int MAX_ACTIONS_PER_LINK = 1;
-    private static final int MAX_SMART_ACTIONS = 3;
-    private static final int MAX_SUGGESTED_REPLIES = 3;
-    // TODO: Make this configurable.
-    private static final int MAX_MESSAGES_TO_EXTRACT = 5;
     private static final int MAX_RESULT_ID_TO_CACHE = 20;
 
-    private static final TextClassifier.EntityConfig TYPE_CONFIG =
-            new TextClassifier.EntityConfig.Builder().setIncludedTypes(
-                    Collections.singletonList(ConversationAction.TYPE_TEXT_REPLY))
-                    .includeTypesFromTextClassifier(false)
-                    .build();
     private static final List<String> HINTS =
             Collections.singletonList(ConversationActions.Request.HINT_FOR_NOTIFICATION);
 
@@ -92,20 +75,30 @@ public class SmartActionsHelper {
         mSettings = settings;
     }
 
-    @NonNull
     SmartSuggestions suggest(@NonNull NotificationEntry entry) {
         // Whenever suggest() is called on a notification, its previous session is ended.
         mNotificationKeyToResultIdCache.remove(entry.getSbn().getKey());
 
-        ArrayList<Notification.Action> actions = suggestActions(entry);
-        ArrayList<CharSequence> replies = suggestReplies(entry);
+        boolean eligibleForReplyAdjustment =
+                mSettings.mGenerateReplies && isEligibleForReplyAdjustment(entry);
+        boolean eligibleForActionAdjustment =
+                mSettings.mGenerateActions && isEligibleForActionAdjustment(entry);
 
-        // Not logging subsequent events of this notification if we didn't generate any suggestion
-        // for it.
-        if (replies.isEmpty() && actions.isEmpty()) {
-            mNotificationKeyToResultIdCache.remove(entry.getSbn().getKey());
-        }
+        List<ConversationAction> conversationActions =
+                suggestConversationActions(
+                        entry,
+                        eligibleForReplyAdjustment,
+                        eligibleForActionAdjustment);
 
+        ArrayList<CharSequence> replies = conversationActions.stream()
+                .map(ConversationAction::getTextReply)
+                .filter(textReply -> !TextUtils.isEmpty(textReply))
+                .collect(Collectors.toCollection(ArrayList::new));
+
+        ArrayList<Notification.Action> actions = conversationActions.stream()
+                .filter(conversationAction -> conversationAction.getAction() != null)
+                .map(action -> createNotificationAction(action.getAction(), action.getType()))
+                .collect(Collectors.toCollection(ArrayList::new));
         return new SmartSuggestions(replies, actions);
     }
 
@@ -113,61 +106,48 @@ public class SmartActionsHelper {
      * Adds action adjustments based on the notification contents.
      */
     @NonNull
-    ArrayList<Notification.Action> suggestActions(@NonNull NotificationEntry entry) {
-        if (!mSettings.mGenerateActions) {
-            return EMPTY_ACTION_LIST;
-        }
-        if (!isEligibleForActionAdjustment(entry)) {
-            return EMPTY_ACTION_LIST;
+    private List<ConversationAction> suggestConversationActions(
+            @NonNull NotificationEntry entry,
+            boolean includeReplies,
+            boolean includeActions) {
+        if (!includeReplies && !includeActions) {
+            return Collections.emptyList();
         }
         if (mTextClassifier == null) {
-            return EMPTY_ACTION_LIST;
+            return Collections.emptyList();
         }
         List<ConversationActions.Message> messages = extractMessages(entry.getNotification());
         if (messages.isEmpty()) {
-            return EMPTY_ACTION_LIST;
+            return Collections.emptyList();
         }
-        // TODO: Move to TextClassifier.suggestConversationActions once it is ready.
-        return suggestActionsFromText(
-                messages.get(messages.size() - 1).getText(), MAX_SMART_ACTIONS);
-    }
 
-    @NonNull
-    ArrayList<CharSequence> suggestReplies(@NonNull NotificationEntry entry) {
-        if (!mSettings.mGenerateReplies) {
-            return EMPTY_REPLY_LIST;
-        }
-        if (!isEligibleForReplyAdjustment(entry)) {
-            return EMPTY_REPLY_LIST;
-        }
-        if (mTextClassifier == null) {
-            return EMPTY_REPLY_LIST;
-        }
-        List<ConversationActions.Message> messages = extractMessages(entry.getNotification());
-        if (messages.isEmpty()) {
-            return EMPTY_REPLY_LIST;
+        TextClassifier.EntityConfig.Builder typeConfigBuilder =
+                new TextClassifier.EntityConfig.Builder();
+        if (!includeReplies) {
+            typeConfigBuilder.setExcludedTypes(
+                    Collections.singletonList(ConversationAction.TYPE_TEXT_REPLY));
+        } else if (!includeActions) {
+            typeConfigBuilder
+                    .setIncludedTypes(
+                            Collections.singletonList(ConversationAction.TYPE_TEXT_REPLY))
+                    .includeTypesFromTextClassifier(false);
         }
         ConversationActions.Request request =
                 new ConversationActions.Request.Builder(messages)
-                        .setMaxSuggestions(MAX_SUGGESTED_REPLIES)
+                        .setMaxSuggestions(mSettings.mMaxSuggestions)
                         .setHints(HINTS)
-                        .setTypeConfig(TYPE_CONFIG)
+                        .setTypeConfig(typeConfigBuilder.build())
                         .build();
 
         ConversationActions conversationActionsResult =
                 mTextClassifier.suggestConversationActions(request);
-        List<ConversationAction> conversationActions =
-                conversationActionsResult.getConversationActions();
-        ArrayList<CharSequence> replies = conversationActions.stream()
-                .map(conversationAction -> conversationAction.getTextReply())
-                .filter(textReply -> !TextUtils.isEmpty(textReply))
-                .collect(Collectors.toCollection(ArrayList::new));
 
         String resultId = conversationActionsResult.getId();
-        if (resultId != null) {
+        if (!TextUtils.isEmpty(resultId)
+                && !conversationActionsResult.getConversationActions().isEmpty()) {
             mNotificationKeyToResultIdCache.put(entry.getSbn().getKey(), resultId);
         }
-        return replies;
+        return conversationActionsResult.getConversationActions();
     }
 
     void onNotificationExpansionChanged(@NonNull NotificationEntry entry, boolean isUserAction,
@@ -248,6 +228,17 @@ public class SmartActionsHelper {
         mTextClassifier.onTextClassifierEvent(textClassifierEvent);
     }
 
+    private Notification.Action createNotificationAction(
+            RemoteAction remoteAction, String actionType) {
+        return new Notification.Action.Builder(
+                remoteAction.getIcon(),
+                remoteAction.getTitle(),
+                remoteAction.getActionIntent())
+                .setContextual(true)
+                .addExtras(Bundle.forPair(KEY_ACTION_TYPE, actionType))
+                .build();
+    }
+
     private TextClassifierEvent.Builder createTextClassifierEventBuilder(
             int eventType, @NonNull String resultId) {
         return new TextClassifierEvent.Builder(
@@ -308,7 +299,7 @@ public class SmartActionsHelper {
     private List<ConversationActions.Message> extractMessages(@NonNull Notification notification) {
         Parcelable[] messages = notification.extras.getParcelableArray(Notification.EXTRA_MESSAGES);
         if (messages == null || messages.length == 0) {
-            return Arrays.asList(new ConversationActions.Message.Builder(
+            return Collections.singletonList(new ConversationActions.Message.Builder(
                     ConversationActions.Message.PERSON_USER_OTHERS)
                     .setText(notification.extras.getCharSequence(Notification.EXTRA_TEXT))
                     .build());
@@ -335,75 +326,13 @@ public class SmartActionsHelper {
                             ZonedDateTime.ofInstant(Instant.ofEpochMilli(message.getTimestamp()),
                                     ZoneOffset.systemDefault()))
                     .build());
-            if (extractMessages.size() >= MAX_MESSAGES_TO_EXTRACT) {
+            if (extractMessages.size() >= mSettings.mMaxMessagesToExtract) {
                 break;
             }
         }
         return new ArrayList<>(extractMessages);
     }
 
-    /** Returns a list of actions to act on entities in a given piece of text. */
-    @NonNull
-    private ArrayList<Notification.Action> suggestActionsFromText(
-            @Nullable CharSequence text, int maxSmartActions) {
-        if (TextUtils.isEmpty(text)) {
-            return EMPTY_ACTION_LIST;
-        }
-        // We want to process only text visible to the user to avoid confusing suggestions, so we
-        // truncate the text to a reasonable length. This is particularly important for e.g.
-        // email apps that sometimes include the text for the entire thread.
-        text = text.subSequence(0, Math.min(text.length(), MAX_ACTION_EXTRACTION_TEXT_LENGTH));
-
-        // Extract all entities.
-        TextLinks.Request textLinksRequest = new TextLinks.Request.Builder(text)
-                .setEntityConfig(
-                        TextClassifier.EntityConfig.createWithHints(
-                                Collections.singletonList(
-                                        TextClassifier.HINT_TEXT_IS_NOT_EDITABLE)))
-                .build();
-        TextLinks links = mTextClassifier.generateLinks(textLinksRequest);
-        EntityTypeCounter entityTypeCounter = EntityTypeCounter.fromTextLinks(links);
-
-        ArrayList<Notification.Action> actions = new ArrayList<>();
-        for (TextLinks.TextLink link : links.getLinks()) {
-            // Ignore any entity type for which we have too many entities. This is to handle the
-            // case where a notification contains e.g. a list of phone numbers. In such cases, the
-            // user likely wants to act on the whole list rather than an individual entity.
-            if (link.getEntityCount() == 0
-                    || entityTypeCounter.getCount(link.getEntity(0)) != 1) {
-                continue;
-            }
-
-            // Generate the actions, and add the most prominent ones to the action bar.
-            TextClassification classification =
-                    mTextClassifier.classifyText(
-                            new TextClassification.Request.Builder(
-                                    text, link.getStart(), link.getEnd()).build());
-            if (classification.getEntityCount() == 0) {
-                continue;
-            }
-            int numOfActions = Math.min(
-                    MAX_ACTIONS_PER_LINK, classification.getActions().size());
-            for (int i = 0; i < numOfActions; ++i) {
-                RemoteAction remoteAction = classification.getActions().get(i);
-                Notification.Action action = new Notification.Action.Builder(
-                        remoteAction.getIcon(),
-                        remoteAction.getTitle(),
-                        remoteAction.getActionIntent())
-                        .setContextual(true)
-                        .addExtras(Bundle.forPair(KEY_ACTION_TYPE, classification.getEntity(0)))
-                        .build();
-                actions.add(action);
-
-                // We have enough smart actions.
-                if (actions.size() >= maxSmartActions) {
-                    return actions;
-                }
-            }
-        }
-        return actions;
-    }
-
     static class SmartSuggestions {
         public final ArrayList<CharSequence> replies;
         public final ArrayList<Notification.Action> actions;
diff --git a/packages/ExtServices/tests/Android.bp b/packages/ExtServices/tests/Android.bp
index 5de454836c2a..930b783aaa98 100644
--- a/packages/ExtServices/tests/Android.bp
+++ b/packages/ExtServices/tests/Android.bp
@@ -15,6 +15,7 @@ android_test {
     static_libs: [
         "ExtServices-core",
         "android-support-test",
+        "compatibility-device-util",
         "mockito-target-minus-junit4",
         "espresso-core",
         "truth-prebuilt",
diff --git a/packages/ExtServices/tests/src/android/ext/services/notification/AssistantSettingsTest.java b/packages/ExtServices/tests/src/android/ext/services/notification/AssistantSettingsTest.java
index 597051a8b744..d890c1ae81fb 100644
--- a/packages/ExtServices/tests/src/android/ext/services/notification/AssistantSettingsTest.java
+++ b/packages/ExtServices/tests/src/android/ext/services/notification/AssistantSettingsTest.java
@@ -16,6 +16,10 @@
 
 package android.ext.services.notification;
 
+import static android.ext.services.notification.AssistantSettings.DEFAULT_MAX_SUGGESTIONS;
+import static android.provider.DeviceConfig.NotificationAssistant;
+import static android.provider.DeviceConfig.setProperty;
+
 import static junit.framework.Assert.assertFalse;
 import static junit.framework.Assert.assertTrue;
 
@@ -26,12 +30,13 @@ import static org.mockito.Mockito.verify;
 import android.content.ContentResolver;
 import android.os.Handler;
 import android.os.Looper;
-import android.provider.DeviceConfig;
 import android.provider.Settings;
 import android.support.test.InstrumentationRegistry;
 import android.support.test.runner.AndroidJUnit4;
+import android.support.test.uiautomator.UiDevice;
 import android.testing.TestableContext;
 
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -39,8 +44,13 @@ import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.io.IOException;
+
 @RunWith(AndroidJUnit4.class)
 public class AssistantSettingsTest {
+    private static final String CLEAR_DEVICE_CONFIG_KEY_CMD =
+            "device_config delete " + NotificationAssistant.NAMESPACE;
+
     private static final int USER_ID = 5;
 
     @Rule
@@ -69,16 +79,21 @@ public class AssistantSettingsTest {
                 handler, mResolver, USER_ID, mOnUpdateRunnable);
     }
 
+    @After
+    public void tearDown() throws IOException {
+        clearDeviceConfig();
+    }
+
     @Test
     public void testGenerateRepliesDisabled() {
-        DeviceConfig.setProperty(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_REPLIES,
+        setProperty(
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.GENERATE_REPLIES,
                 "false",
                 false /* makeDefault */);
         mAssistantSettings.onDeviceConfigPropertyChanged(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_REPLIES,
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.GENERATE_REPLIES,
                 "false");
 
         assertFalse(mAssistantSettings.mGenerateReplies);
@@ -86,14 +101,14 @@ public class AssistantSettingsTest {
 
     @Test
     public void testGenerateRepliesEnabled() {
-        DeviceConfig.setProperty(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_REPLIES,
+        setProperty(
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.GENERATE_REPLIES,
                 "true",
                 false /* makeDefault */);
         mAssistantSettings.onDeviceConfigPropertyChanged(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_REPLIES,
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.GENERATE_REPLIES,
                 "true");
 
         assertTrue(mAssistantSettings.mGenerateReplies);
@@ -101,26 +116,26 @@ public class AssistantSettingsTest {
 
     @Test
     public void testGenerateRepliesEmptyFlag() {
-        DeviceConfig.setProperty(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_REPLIES,
+        setProperty(
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.GENERATE_REPLIES,
                 "false",
                 false /* makeDefault */);
         mAssistantSettings.onDeviceConfigPropertyChanged(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_REPLIES,
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.GENERATE_REPLIES,
                 "false");
 
         assertFalse(mAssistantSettings.mGenerateReplies);
 
-        DeviceConfig.setProperty(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_REPLIES,
+        setProperty(
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.GENERATE_REPLIES,
                 "",
                 false /* makeDefault */);
         mAssistantSettings.onDeviceConfigPropertyChanged(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_REPLIES,
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.GENERATE_REPLIES,
                 "");
 
         // Go back to the default value.
@@ -129,14 +144,14 @@ public class AssistantSettingsTest {
 
     @Test
     public void testGenerateActionsDisabled() {
-        DeviceConfig.setProperty(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_ACTIONS,
+        setProperty(
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.GENERATE_ACTIONS,
                 "false",
                 false /* makeDefault */);
         mAssistantSettings.onDeviceConfigPropertyChanged(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_ACTIONS,
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.GENERATE_ACTIONS,
                 "false");
 
         assertFalse(mAssistantSettings.mGenerateActions);
@@ -144,14 +159,14 @@ public class AssistantSettingsTest {
 
     @Test
     public void testGenerateActionsEnabled() {
-        DeviceConfig.setProperty(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_ACTIONS,
+        setProperty(
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.GENERATE_ACTIONS,
                 "true",
                 false /* makeDefault */);
         mAssistantSettings.onDeviceConfigPropertyChanged(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_ACTIONS,
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.GENERATE_ACTIONS,
                 "true");
 
         assertTrue(mAssistantSettings.mGenerateActions);
@@ -159,26 +174,26 @@ public class AssistantSettingsTest {
 
     @Test
     public void testGenerateActionsEmptyFlag() {
-        DeviceConfig.setProperty(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_ACTIONS,
+        setProperty(
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.GENERATE_ACTIONS,
                 "false",
                 false /* makeDefault */);
         mAssistantSettings.onDeviceConfigPropertyChanged(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_ACTIONS,
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.GENERATE_ACTIONS,
                 "false");
 
         assertFalse(mAssistantSettings.mGenerateActions);
 
-        DeviceConfig.setProperty(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_ACTIONS,
+        setProperty(
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.GENERATE_ACTIONS,
                 "",
                 false /* makeDefault */);
         mAssistantSettings.onDeviceConfigPropertyChanged(
-                DeviceConfig.NotificationAssistant.NAMESPACE,
-                DeviceConfig.NotificationAssistant.GENERATE_ACTIONS,
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.GENERATE_ACTIONS,
                 "");
 
         // Go back to the default value.
@@ -186,6 +201,46 @@ public class AssistantSettingsTest {
     }
 
     @Test
+    public void testMaxMessagesToExtract() {
+        setProperty(
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.MAX_MESSAGES_TO_EXTRACT,
+                "10",
+                false /* makeDefault */);
+        mAssistantSettings.onDeviceConfigPropertyChanged(
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.MAX_MESSAGES_TO_EXTRACT,
+                "10");
+
+        assertEquals(10, mAssistantSettings.mMaxMessagesToExtract);
+    }
+
+    @Test
+    public void testMaxSuggestions() {
+        setProperty(
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.MAX_SUGGESTIONS,
+                "5",
+                false /* makeDefault */);
+        mAssistantSettings.onDeviceConfigPropertyChanged(
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.MAX_SUGGESTIONS,
+                "5");
+
+        assertEquals(5, mAssistantSettings.mMaxSuggestions);
+    }
+
+    @Test
+    public void testMaxSuggestionsEmpty() {
+        mAssistantSettings.onDeviceConfigPropertyChanged(
+                NotificationAssistant.NAMESPACE,
+                NotificationAssistant.MAX_SUGGESTIONS,
+                "");
+
+        assertEquals(DEFAULT_MAX_SUGGESTIONS, mAssistantSettings.mMaxSuggestions);
+    }
+
+    @Test
     public void testStreakLimit() {
         verify(mOnUpdateRunnable, never()).run();
 
@@ -219,4 +274,17 @@ public class AssistantSettingsTest {
         assertEquals(newDismissToViewRatioLimit, mAssistantSettings.mDismissToViewRatioLimit, 1e-6);
         verify(mOnUpdateRunnable).run();
     }
+
+    private static void clearDeviceConfig() throws IOException {
+        UiDevice uiDevice = UiDevice.getInstance(InstrumentationRegistry.getInstrumentation());
+        uiDevice.executeShellCommand(
+                CLEAR_DEVICE_CONFIG_KEY_CMD + " " + NotificationAssistant.GENERATE_ACTIONS);
+        uiDevice.executeShellCommand(
+                CLEAR_DEVICE_CONFIG_KEY_CMD + " " + NotificationAssistant.GENERATE_REPLIES);
+        uiDevice.executeShellCommand(
+                CLEAR_DEVICE_CONFIG_KEY_CMD + " " + NotificationAssistant.MAX_MESSAGES_TO_EXTRACT);
+        uiDevice.executeShellCommand(
+                CLEAR_DEVICE_CONFIG_KEY_CMD + " " + NotificationAssistant.MAX_SUGGESTIONS);
+    }
+
 }
diff --git a/packages/ExtServices/tests/src/android/ext/services/notification/SmartActionHelperTest.java b/packages/ExtServices/tests/src/android/ext/services/notification/SmartActionsHelperTest.java
index 7f8127aa43a8..ebbd961b6f23 100644
--- a/packages/ExtServices/tests/src/android/ext/services/notification/SmartActionHelperTest.java
+++ b/packages/ExtServices/tests/src/android/ext/services/notification/SmartActionsHelperTest.java
@@ -25,8 +25,15 @@ import static org.mockito.Mockito.when;
 
 import android.annotation.NonNull;
 import android.app.Notification;
+import android.app.NotificationChannel;
+import android.app.NotificationManager;
+import android.app.PendingIntent;
 import android.app.Person;
+import android.app.RemoteInput;
 import android.content.Context;
+import android.content.Intent;
+import android.content.pm.IPackageManager;
+import android.graphics.drawable.Icon;
 import android.os.Process;
 import android.service.notification.NotificationAssistantService;
 import android.service.notification.StatusBarNotification;
@@ -53,7 +60,6 @@ import org.mockito.MockitoAnnotations;
 import java.time.Instant;
 import java.time.ZoneOffset;
 import java.time.ZonedDateTime;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
@@ -62,20 +68,26 @@ import java.util.Objects;
 import javax.annotation.Nullable;
 
 @RunWith(AndroidJUnit4.class)
-public class SmartActionHelperTest {
+public class SmartActionsHelperTest {
     private static final String NOTIFICATION_KEY = "key";
     private static final String RESULT_ID = "id";
-
     private static final ConversationAction REPLY_ACTION =
             new ConversationAction.Builder(ConversationAction.TYPE_TEXT_REPLY)
-            .setTextReply("Home")
-            .build();
+                    .setTextReply("Home")
+                    .build();
+    private static final String MESSAGE = "Where are you?";
+
+    @Mock
+    IPackageManager mIPackageManager;
+    @Mock
+    private TextClassifier mTextClassifier;
+    @Mock
+    private StatusBarNotification mStatusBarNotification;
+    @Mock
+    private SmsHelper mSmsHelper;
 
     private SmartActionsHelper mSmartActionsHelper;
     private Context mContext;
-    @Mock private TextClassifier mTextClassifier;
-    @Mock private NotificationEntry mNotificationEntry;
-    @Mock private StatusBarNotification mStatusBarNotification;
     private Notification.Builder mNotificationBuilder;
     private AssistantSettings mSettings;
 
@@ -89,10 +101,6 @@ public class SmartActionHelperTest {
         when(mTextClassifier.suggestConversationActions(any(ConversationActions.Request.class)))
                 .thenReturn(new ConversationActions(Arrays.asList(REPLY_ACTION), RESULT_ID));
 
-        when(mNotificationEntry.getSbn()).thenReturn(mStatusBarNotification);
-        // The notification is eligible to have smart suggestions.
-        when(mNotificationEntry.hasInlineReply()).thenReturn(true);
-        when(mNotificationEntry.isMessaging()).thenReturn(true);
         when(mStatusBarNotification.getPackageName()).thenReturn("random.app");
         when(mStatusBarNotification.getUser()).thenReturn(Process.myUserHandle());
         when(mStatusBarNotification.getKey()).thenReturn(NOTIFICATION_KEY);
@@ -105,33 +113,96 @@ public class SmartActionHelperTest {
     }
 
     @Test
-    public void testSuggestReplies_notMessagingApp() {
-        when(mNotificationEntry.isMessaging()).thenReturn(false);
-        ArrayList<CharSequence> textReplies =
-                mSmartActionsHelper.suggestReplies(mNotificationEntry);
-        assertThat(textReplies).isEmpty();
+    public void testSuggest_notMessageNotification() {
+        Notification notification = mNotificationBuilder.setContentText(MESSAGE).build();
+        when(mStatusBarNotification.getNotification()).thenReturn(notification);
+
+        mSmartActionsHelper.suggest(createNotificationEntry());
+
+        verify(mTextClassifier, never())
+                .suggestConversationActions(any(ConversationActions.Request.class));
+    }
+
+    @Test
+    public void testSuggest_noInlineReply() {
+        Notification notification =
+                mNotificationBuilder
+                        .setContentText(MESSAGE)
+                        .setCategory(Notification.CATEGORY_MESSAGE)
+                        .build();
+        when(mStatusBarNotification.getNotification()).thenReturn(notification);
+
+        ConversationActions.Request request = runSuggestAndCaptureRequest();
+
+        // actions are enabled, but replies are not.
+        assertThat(
+                request.getTypeConfig().resolveEntityListModifications(
+                        Arrays.asList(ConversationAction.TYPE_TEXT_REPLY,
+                                ConversationAction.TYPE_OPEN_URL)))
+                .containsExactly(ConversationAction.TYPE_OPEN_URL);
+    }
+
+    @Test
+    public void testSuggest_settingsOff() {
+        mSettings.mGenerateActions = false;
+        mSettings.mGenerateReplies = false;
+        Notification notification = createMessageNotification();
+        when(mStatusBarNotification.getNotification()).thenReturn(notification);
+
+        mSmartActionsHelper.suggest(createNotificationEntry());
+
+        verify(mTextClassifier, never())
+                .suggestConversationActions(any(ConversationActions.Request.class));
+    }
+
+    @Test
+    public void testSuggest_settings_repliesOnActionsOff() {
+        mSettings.mGenerateReplies = true;
+        mSettings.mGenerateActions = false;
+        Notification notification = createMessageNotification();
+        when(mStatusBarNotification.getNotification()).thenReturn(notification);
+
+        ConversationActions.Request request = runSuggestAndCaptureRequest();
+
+        // replies are enabled, but actions are not.
+        assertThat(
+                request.getTypeConfig().resolveEntityListModifications(
+                        Arrays.asList(ConversationAction.TYPE_TEXT_REPLY,
+                                ConversationAction.TYPE_OPEN_URL)))
+                .containsExactly(ConversationAction.TYPE_TEXT_REPLY);
     }
 
     @Test
-    public void testSuggestReplies_noInlineReply() {
-        when(mNotificationEntry.hasInlineReply()).thenReturn(false);
-        ArrayList<CharSequence> textReplies =
-                mSmartActionsHelper.suggestReplies(mNotificationEntry);
-        assertThat(textReplies).isEmpty();
+    public void testSuggest_settings_repliesOffActionsOn() {
+        mSettings.mGenerateReplies = false;
+        mSettings.mGenerateActions = true;
+        Notification notification = createMessageNotification();
+        when(mStatusBarNotification.getNotification()).thenReturn(notification);
+
+        ConversationActions.Request request = runSuggestAndCaptureRequest();
+
+        // actions are enabled, but replies are not.
+        assertThat(
+                request.getTypeConfig().resolveEntityListModifications(
+                        Arrays.asList(ConversationAction.TYPE_TEXT_REPLY,
+                                ConversationAction.TYPE_OPEN_URL)))
+                .containsExactly(ConversationAction.TYPE_OPEN_URL);
     }
 
+
     @Test
-    public void testSuggestReplies_nonMessageStyle() {
-        Notification notification = mNotificationBuilder.setContentText("Where are you?").build();
-        when(mNotificationEntry.getNotification()).thenReturn(notification);
+    public void testSuggest_nonMessageStyleMessageNotification() {
+        Notification notification = createMessageNotification();
+        when(mStatusBarNotification.getNotification()).thenReturn(notification);
 
-        List<ConversationActions.Message> messages = getMessagesInRequest();
+        List<ConversationActions.Message> messages =
+                runSuggestAndCaptureRequest().getConversation();
         assertThat(messages).hasSize(1);
-        MessageSubject.assertThat(messages.get(0)).hasText("Where are you?");
+        MessageSubject.assertThat(messages.get(0)).hasText(MESSAGE);
     }
 
     @Test
-    public void testSuggestReplies_messageStyle() {
+    public void testSuggest_messageStyle() {
         Person me = new Person.Builder().setName("Me").build();
         Person userA = new Person.Builder().setName("A").build();
         Person userB = new Person.Builder().setName("B").build();
@@ -145,10 +216,12 @@ public class SmartActionHelperTest {
                 mNotificationBuilder
                         .setContentText("You have three new messages")
                         .setStyle(style)
+                        .setActions(createReplyAction())
                         .build();
-        when(mNotificationEntry.getNotification()).thenReturn(notification);
+        when(mStatusBarNotification.getNotification()).thenReturn(notification);
 
-        List<ConversationActions.Message> messages = getMessagesInRequest();
+        List<ConversationActions.Message> messages =
+                runSuggestAndCaptureRequest().getConversation();
         assertThat(messages).hasSize(3);
 
         ConversationActions.Message secondMessage = messages.get(0);
@@ -172,7 +245,7 @@ public class SmartActionHelperTest {
     }
 
     @Test
-    public void testSuggestReplies_messageStyle_noPerson() {
+    public void testSuggest_messageStyle_noPerson() {
         Person me = new Person.Builder().setName("Me").build();
         Notification.MessagingStyle style =
                 new Notification.MessagingStyle(me).addMessage("message", 1000, (Person) null);
@@ -180,10 +253,11 @@ public class SmartActionHelperTest {
                 mNotificationBuilder
                         .setContentText("You have one new message")
                         .setStyle(style)
+                        .setActions(createReplyAction())
                         .build();
-        when(mNotificationEntry.getNotification()).thenReturn(notification);
+        when(mStatusBarNotification.getNotification()).thenReturn(notification);
 
-        mSmartActionsHelper.suggestReplies(mNotificationEntry);
+        mSmartActionsHelper.suggest(createNotificationEntry());
 
         verify(mTextClassifier, never())
                 .suggestConversationActions(any(ConversationActions.Request.class));
@@ -191,13 +265,12 @@ public class SmartActionHelperTest {
 
     @Test
     public void testOnSuggestedReplySent() {
-        final String message = "Where are you?";
-        Notification notification = mNotificationBuilder.setContentText(message).build();
-        when(mNotificationEntry.getNotification()).thenReturn(notification);
+        Notification notification = createMessageNotification();
+        when(mStatusBarNotification.getNotification()).thenReturn(notification);
 
-        mSmartActionsHelper.suggestReplies(mNotificationEntry);
+        mSmartActionsHelper.suggest(createNotificationEntry());
         mSmartActionsHelper.onSuggestedReplySent(
-                NOTIFICATION_KEY, message, NotificationAssistantService.SOURCE_FROM_ASSISTANT);
+                NOTIFICATION_KEY, MESSAGE, NotificationAssistantService.SOURCE_FROM_ASSISTANT);
 
         ArgumentCaptor<TextClassifierEvent> argumentCaptor =
                 ArgumentCaptor.forClass(TextClassifierEvent.class);
@@ -208,13 +281,12 @@ public class SmartActionHelperTest {
 
     @Test
     public void testOnSuggestedReplySent_anotherNotification() {
-        final String message = "Where are you?";
-        Notification notification = mNotificationBuilder.setContentText(message).build();
-        when(mNotificationEntry.getNotification()).thenReturn(notification);
+        Notification notification = createMessageNotification();
+        when(mStatusBarNotification.getNotification()).thenReturn(notification);
 
-        mSmartActionsHelper.suggestReplies(mNotificationEntry);
+        mSmartActionsHelper.suggest(createNotificationEntry());
         mSmartActionsHelper.onSuggestedReplySent(
-                "something_else", message, NotificationAssistantService.SOURCE_FROM_ASSISTANT);
+                "something_else", MESSAGE, NotificationAssistantService.SOURCE_FROM_ASSISTANT);
 
         verify(mTextClassifier, never())
                 .onTextClassifierEvent(Mockito.any(TextClassifierEvent.class));
@@ -225,13 +297,12 @@ public class SmartActionHelperTest {
         when(mTextClassifier.suggestConversationActions(any(ConversationActions.Request.class)))
                 .thenReturn(new ConversationActions(Collections.emptyList(), null));
 
-        final String message = "Where are you?";
-        Notification notification = mNotificationBuilder.setContentText(message).build();
-        when(mNotificationEntry.getNotification()).thenReturn(notification);
+        Notification notification = createMessageNotification();
+        when(mStatusBarNotification.getNotification()).thenReturn(notification);
 
-        mSmartActionsHelper.suggestReplies(mNotificationEntry);
+        mSmartActionsHelper.suggest(createNotificationEntry());
         mSmartActionsHelper.onSuggestedReplySent(
-                "something_else", message, NotificationAssistantService.SOURCE_FROM_ASSISTANT);
+                "something_else", MESSAGE, NotificationAssistantService.SOURCE_FROM_ASSISTANT);
 
         verify(mTextClassifier, never())
                 .onTextClassifierEvent(Mockito.any(TextClassifierEvent.class));
@@ -239,10 +310,10 @@ public class SmartActionHelperTest {
 
     @Test
     public void testOnNotificationDirectReply() {
-        Notification notification = mNotificationBuilder.setContentText("Where are you?").build();
-        when(mNotificationEntry.getNotification()).thenReturn(notification);
+        Notification notification = createMessageNotification();
+        when(mStatusBarNotification.getNotification()).thenReturn(notification);
 
-        mSmartActionsHelper.suggestReplies(mNotificationEntry);
+        mSmartActionsHelper.suggest(createNotificationEntry());
         mSmartActionsHelper.onNotificationDirectReplied(NOTIFICATION_KEY);
 
         ArgumentCaptor<TextClassifierEvent> argumentCaptor =
@@ -254,12 +325,11 @@ public class SmartActionHelperTest {
 
     @Test
     public void testOnNotificationExpansionChanged() {
-        final String message = "Where are you?";
-        Notification notification = mNotificationBuilder.setContentText(message).build();
-        when(mNotificationEntry.getNotification()).thenReturn(notification);
+        Notification notification = createMessageNotification();
+        when(mStatusBarNotification.getNotification()).thenReturn(notification);
 
-        mSmartActionsHelper.suggestReplies(mNotificationEntry);
-        mSmartActionsHelper.onNotificationExpansionChanged(mNotificationEntry, true, true);
+        mSmartActionsHelper.suggest(createNotificationEntry());
+        mSmartActionsHelper.onNotificationExpansionChanged(createNotificationEntry(), true, true);
 
         ArgumentCaptor<TextClassifierEvent> argumentCaptor =
                 ArgumentCaptor.forClass(TextClassifierEvent.class);
@@ -270,12 +340,11 @@ public class SmartActionHelperTest {
 
     @Test
     public void testOnNotificationsSeen_notExpanded() {
-        final String message = "Where are you?";
-        Notification notification = mNotificationBuilder.setContentText(message).build();
-        when(mNotificationEntry.getNotification()).thenReturn(notification);
+        Notification notification = createMessageNotification();
+        when(mStatusBarNotification.getNotification()).thenReturn(notification);
 
-        mSmartActionsHelper.suggestReplies(mNotificationEntry);
-        mSmartActionsHelper.onNotificationExpansionChanged(mNotificationEntry, false, false);
+        mSmartActionsHelper.suggest(createNotificationEntry());
+        mSmartActionsHelper.onNotificationExpansionChanged(createNotificationEntry(), false, false);
 
         verify(mTextClassifier, never()).onTextClassifierEvent(
                 Mockito.any(TextClassifierEvent.class));
@@ -283,12 +352,11 @@ public class SmartActionHelperTest {
 
     @Test
     public void testOnNotifications_expanded() {
-        final String message = "Where are you?";
-        Notification notification = mNotificationBuilder.setContentText(message).build();
-        when(mNotificationEntry.getNotification()).thenReturn(notification);
+        Notification notification = createMessageNotification();
+        when(mStatusBarNotification.getNotification()).thenReturn(notification);
 
-        mSmartActionsHelper.suggestReplies(mNotificationEntry);
-        mSmartActionsHelper.onNotificationExpansionChanged(mNotificationEntry, false, true);
+        mSmartActionsHelper.suggest(createNotificationEntry());
+        mSmartActionsHelper.onNotificationExpansionChanged(createNotificationEntry(), false, true);
 
         ArgumentCaptor<TextClassifierEvent> argumentCaptor =
                 ArgumentCaptor.forClass(TextClassifierEvent.class);
@@ -301,14 +369,41 @@ public class SmartActionHelperTest {
         return ZonedDateTime.ofInstant(Instant.ofEpochMilli(msUtc), ZoneOffset.systemDefault());
     }
 
-    private List<ConversationActions.Message> getMessagesInRequest() {
-        mSmartActionsHelper.suggestReplies(mNotificationEntry);
+    private ConversationActions.Request runSuggestAndCaptureRequest() {
+        mSmartActionsHelper.suggest(createNotificationEntry());
 
         ArgumentCaptor<ConversationActions.Request> argumentCaptor =
                 ArgumentCaptor.forClass(ConversationActions.Request.class);
         verify(mTextClassifier).suggestConversationActions(argumentCaptor.capture());
-        ConversationActions.Request request = argumentCaptor.getValue();
-        return request.getConversation();
+        return argumentCaptor.getValue();
+    }
+
+    private Notification.Action createReplyAction() {
+        PendingIntent pendingIntent =
+                PendingIntent.getActivity(mContext, 0, new Intent(mContext, this.getClass()), 0);
+        RemoteInput remoteInput = new RemoteInput.Builder("result")
+                .setAllowFreeFormInput(true)
+                .build();
+        return new Notification.Action.Builder(
+                Icon.createWithResource(mContext.getResources(),
+                        android.R.drawable.stat_sys_warning),
+                "Reply", pendingIntent)
+                .addRemoteInput(remoteInput)
+                .build();
+    }
+
+    private NotificationEntry createNotificationEntry() {
+        NotificationChannel channel =
+                new NotificationChannel("id", "name", NotificationManager.IMPORTANCE_DEFAULT);
+        return new NotificationEntry(mIPackageManager, mStatusBarNotification, channel, mSmsHelper);
+    }
+
+    private Notification createMessageNotification() {
+        return mNotificationBuilder
+                .setContentText(MESSAGE)
+                .setCategory(Notification.CATEGORY_MESSAGE)
+                .setActions(createReplyAction())
+                .build();
     }
 
     private void assertTextClassifierEvent(