Merge "Add spinner for smart replies." into pi-dev
diff --git a/core/java/android/app/Notification.java b/core/java/android/app/Notification.java
index 4f88a03..4b15378 100644
--- a/core/java/android/app/Notification.java
+++ b/core/java/android/app/Notification.java
@@ -972,6 +972,18 @@
     public static final String EXTRA_REMOTE_INPUT_HISTORY = "android.remoteInputHistory";
 
     /**
+     * {@link #extras} key: boolean as supplied to
+     * {@link Builder#setShowRemoteInputSpinner(boolean)}.
+     *
+     * If set to true, then the view displaying the remote input history from
+     * {@link Builder#setRemoteInputHistory(CharSequence[])} will have a progress spinner.
+     *
+     * @see Builder#setShowRemoteInputSpinner(boolean)
+     * @hide
+     */
+    public static final String EXTRA_SHOW_REMOTE_INPUT_SPINNER = "android.remoteInputSpinner";
+
+    /**
      * {@link #extras} key: this is a small piece of additional text as supplied to
      * {@link Builder#setContentInfo(CharSequence)}.
      */
@@ -3537,6 +3549,15 @@
         }
 
         /**
+         * Sets whether remote history entries view should have a spinner.
+         * @hide
+         */
+        public Builder setShowRemoteInputSpinner(boolean showSpinner) {
+            mN.extras.putBoolean(EXTRA_SHOW_REMOTE_INPUT_SPINNER, showSpinner);
+            return this;
+        }
+
+        /**
          * Sets the number of items this notification represents. May be displayed as a badge count
          * for Launchers that support badging.
          */
@@ -4760,6 +4781,8 @@
 
             big.setViewVisibility(R.id.notification_material_reply_container, View.GONE);
             big.setTextViewText(R.id.notification_material_reply_text_1, null);
+            big.setViewVisibility(R.id.notification_material_reply_text_1_container, View.GONE);
+            big.setViewVisibility(R.id.notification_material_reply_progress, View.GONE);
 
             big.setViewVisibility(R.id.notification_material_reply_text_2, View.GONE);
             big.setTextViewText(R.id.notification_material_reply_text_2, null);
@@ -4810,10 +4833,19 @@
             CharSequence[] replyText = mN.extras.getCharSequenceArray(EXTRA_REMOTE_INPUT_HISTORY);
             if (!p.ambient && validRemoteInput && replyText != null
                     && replyText.length > 0 && !TextUtils.isEmpty(replyText[0])) {
+                boolean showSpinner = mN.extras.getBoolean(EXTRA_SHOW_REMOTE_INPUT_SPINNER);
                 big.setViewVisibility(R.id.notification_material_reply_container, View.VISIBLE);
+                big.setViewVisibility(R.id.notification_material_reply_text_1_container,
+                        View.VISIBLE);
                 big.setTextViewText(R.id.notification_material_reply_text_1,
                         processTextSpans(replyText[0]));
                 setTextViewColorSecondary(big, R.id.notification_material_reply_text_1);
+                big.setViewVisibility(R.id.notification_material_reply_progress,
+                        showSpinner ? View.VISIBLE : View.GONE);
+                big.setProgressIndeterminateTintList(
+                        R.id.notification_material_reply_progress,
+                        ColorStateList.valueOf(
+                                isColorized() ? getPrimaryTextColor() : resolveContrastColor()));
 
                 if (replyText.length > 1 && !TextUtils.isEmpty(replyText[1])) {
                     big.setViewVisibility(R.id.notification_material_reply_text_2, View.VISIBLE);
@@ -6953,11 +6985,16 @@
             static final String KEY_DATA_MIME_TYPE = "type";
             static final String KEY_DATA_URI= "uri";
             static final String KEY_EXTRAS_BUNDLE = "extras";
+            static final String KEY_REMOTE_INPUT_HISTORY = "remote_input_history";
 
             private final CharSequence mText;
             private final long mTimestamp;
             @Nullable
             private final Person mSender;
+            /** True if this message was generated from the extra
+             *  {@link Notification#EXTRA_REMOTE_INPUT_HISTORY}
+             */
+            private final boolean mRemoteInputHistory;
 
             private Bundle mExtras = new Bundle();
             private String mDataMimeType;
@@ -6996,9 +7033,33 @@
              * </p>
              */
             public Message(@NonNull CharSequence text, long timestamp, @Nullable Person sender) {
+                this(text, timestamp, sender, false /* remoteHistory */);
+            }
+
+            /**
+             * Constructor
+             * @param text A {@link CharSequence} to be displayed as the message content
+             * @param timestamp Time at which the message arrived
+             * @param sender The {@link Person} who sent the message.
+             * Should be <code>null</code> for messages by the current user, in which case
+             * the platform will insert the user set in {@code MessagingStyle(Person)}.
+             * @param remoteInputHistory True if the messages was generated from the extra
+             * {@link Notification#EXTRA_REMOTE_INPUT_HISTORY}.
+             * <p>
+             * The person provided should contain an Icon, set with
+             * {@link Person.Builder#setIcon(Icon)} and also have a name provided
+             * with {@link Person.Builder#setName(CharSequence)}. If multiple users have the same
+             * name, consider providing a key with {@link Person.Builder#setKey(String)} in order
+             * to differentiate between the different users.
+             * </p>
+             * @hide
+             */
+            public Message(@NonNull CharSequence text, long timestamp, @Nullable Person sender,
+                    boolean remoteInputHistory) {
                 mText = text;
                 mTimestamp = timestamp;
                 mSender = sender;
+                mRemoteInputHistory = remoteInputHistory;
             }
 
             /**
@@ -7088,6 +7149,15 @@
                 return mDataUri;
             }
 
+            /**
+             * @return True if the message was generated from
+             * {@link Notification#EXTRA_REMOTE_INPUT_HISTORY}.
+             * @hide
+             */
+            public boolean isRemoteInputHistory() {
+                return mRemoteInputHistory;
+            }
+
             private Bundle toBundle() {
                 Bundle bundle = new Bundle();
                 if (mText != null) {
@@ -7108,6 +7178,9 @@
                 if (mExtras != null) {
                     bundle.putBundle(KEY_EXTRAS_BUNDLE, mExtras);
                 }
+                if (mRemoteInputHistory) {
+                    bundle.putBoolean(KEY_REMOTE_INPUT_HISTORY, mRemoteInputHistory);
+                }
                 return bundle;
             }
 
@@ -7159,7 +7232,8 @@
                         }
                         Message message = new Message(bundle.getCharSequence(KEY_TEXT),
                                 bundle.getLong(KEY_TIMESTAMP),
-                                senderPerson);
+                                senderPerson,
+                                bundle.getBoolean(KEY_REMOTE_INPUT_HISTORY, false));
                         if (bundle.containsKey(KEY_DATA_MIME_TYPE) &&
                                 bundle.containsKey(KEY_DATA_URI)) {
                             message.setData(bundle.getString(KEY_DATA_MIME_TYPE),
diff --git a/core/java/com/android/internal/widget/MessagingGroup.java b/core/java/com/android/internal/widget/MessagingGroup.java
index 07d78fe..627cf47 100644
--- a/core/java/com/android/internal/widget/MessagingGroup.java
+++ b/core/java/com/android/internal/widget/MessagingGroup.java
@@ -22,6 +22,8 @@
 import android.annotation.StyleRes;
 import android.app.Person;
 import android.content.Context;
+import android.content.res.ColorStateList;
+import android.graphics.Color;
 import android.graphics.Point;
 import android.graphics.Rect;
 import android.graphics.drawable.Icon;
@@ -29,6 +31,7 @@
 import android.util.AttributeSet;
 import android.util.DisplayMetrics;
 import android.util.Pools;
+import android.util.TypedValue;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.view.ViewGroup;
@@ -36,6 +39,7 @@
 import android.view.ViewTreeObserver;
 import android.widget.ImageView;
 import android.widget.LinearLayout;
+import android.widget.ProgressBar;
 import android.widget.RemoteViews;
 
 import com.android.internal.R;
@@ -58,6 +62,7 @@
     private CharSequence mAvatarName = "";
     private Icon mAvatarIcon;
     private int mTextColor;
+    private int mSendingTextColor;
     private List<MessagingMessage> mMessages;
     private ArrayList<MessagingMessage> mAddedMessages = new ArrayList<>();
     private boolean mFirstLayout;
@@ -69,6 +74,8 @@
     private MessagingImageMessage mIsolatedMessage;
     private boolean mTransformingImages;
     private Point mDisplaySize = new Point();
+    private ProgressBar mSendingSpinner;
+    private View mSendingSpinnerContainer;
 
     public MessagingGroup(@NonNull Context context) {
         super(context);
@@ -96,6 +103,8 @@
         mSenderName.addOnLayoutChangeListener(MessagingLayout.MESSAGING_PROPERTY_ANIMATOR);
         mAvatarView = findViewById(R.id.message_icon);
         mImageContainer = findViewById(R.id.messaging_group_icon_container);
+        mSendingSpinner = findViewById(R.id.messaging_group_sending_progress);
+        mSendingSpinnerContainer = findViewById(R.id.messaging_group_sending_progress_container);
         DisplayMetrics displayMetrics = getResources().getDisplayMetrics();
         mDisplaySize.x = displayMetrics.widthPixels;
         mDisplaySize.y = displayMetrics.heightPixels;
@@ -139,17 +148,37 @@
         mAvatarView.setVisibility(VISIBLE);
         mSenderName.setVisibility(VISIBLE);
         mTextColor = getNormalTextColor();
+        mSendingTextColor = calculateSendingTextColor();
+    }
+
+    public void setSending(boolean sending) {
+        int visibility = sending ? View.VISIBLE : View.GONE;
+        if (mSendingSpinnerContainer.getVisibility() != visibility) {
+            mSendingSpinnerContainer.setVisibility(visibility);
+            updateMessageColor();
+        }
     }
 
     private int getNormalTextColor() {
         return mContext.getColor(R.color.notification_secondary_text_color_light);
     }
 
+    private int calculateSendingTextColor() {
+        TypedValue alphaValue = new TypedValue();
+        mContext.getResources().getValue(
+                R.dimen.notification_secondary_text_disabled_alpha, alphaValue, true);
+        float alpha = alphaValue.getFloat();
+        return Color.valueOf(
+                Color.red(mTextColor),
+                Color.green(mTextColor),
+                Color.blue(mTextColor),
+                alpha).toArgb();
+    }
+
     public void setAvatar(Icon icon) {
         mAvatarIcon = icon;
         mAvatarView.setImageIcon(icon);
         mAvatarSymbol = "";
-        mLayoutColor = 0;
         mAvatarName = "";
     }
 
@@ -321,13 +350,26 @@
                 || layoutColor != mLayoutColor) {
             setAvatar(cachedIcon);
             mAvatarSymbol = avatarSymbol;
-            mLayoutColor = layoutColor;
+            setLayoutColor(layoutColor);
             mAvatarName = avatarName;
         }
     }
 
     public void setLayoutColor(int layoutColor) {
-        mLayoutColor = layoutColor;
+        if (layoutColor != mLayoutColor){
+            mLayoutColor = layoutColor;
+            mSendingSpinner.setIndeterminateTintList(ColorStateList.valueOf(mLayoutColor));
+        }
+    }
+
+    private void updateMessageColor() {
+        if (mMessages != null) {
+            int color = mSendingSpinnerContainer.getVisibility() == View.VISIBLE
+                    ? mSendingTextColor : mTextColor;
+            for (MessagingMessage message : mMessages) {
+                message.setColor(message.getMessage().isRemoteInputHistory() ? color : mTextColor);
+            }
+        }
     }
 
     public void setMessages(List<MessagingMessage> group) {
@@ -336,7 +378,6 @@
         MessagingImageMessage isolatedMessage = null;
         for (int messageIndex = 0; messageIndex < group.size(); messageIndex++) {
             MessagingMessage message = group.get(messageIndex);
-            message.setColor(mTextColor);
             if (message.getGroup() != this) {
                 message.setMessagingGroup(this);
                 mAddedMessages.add(message);
@@ -376,6 +417,7 @@
         }
         mIsolatedMessage = isolatedMessage;
         mMessages = group;
+        updateMessageColor();
     }
 
     /**
diff --git a/core/java/com/android/internal/widget/MessagingLayout.java b/core/java/com/android/internal/widget/MessagingLayout.java
index f8236c7..df20e63 100644
--- a/core/java/com/android/internal/widget/MessagingLayout.java
+++ b/core/java/com/android/internal/widget/MessagingLayout.java
@@ -149,7 +149,9 @@
         }
         addRemoteInputHistoryToMessages(newMessages,
                 extras.getCharSequenceArray(Notification.EXTRA_REMOTE_INPUT_HISTORY));
-        bind(newMessages, newHistoricMessages);
+        boolean showSpinner =
+                extras.getBoolean(Notification.EXTRA_SHOW_REMOTE_INPUT_SPINNER, false);
+        bind(newMessages, newHistoricMessages, showSpinner);
     }
 
     private void addRemoteInputHistoryToMessages(
@@ -161,17 +163,18 @@
         for (int i = remoteInputHistory.length - 1; i >= 0; i--) {
             CharSequence message = remoteInputHistory[i];
             newMessages.add(new Notification.MessagingStyle.Message(
-                    message, 0, (Person) null));
+                    message, 0, (Person) null, true /* remoteHistory */));
         }
     }
 
     private void bind(List<Notification.MessagingStyle.Message> newMessages,
-            List<Notification.MessagingStyle.Message> newHistoricMessages) {
+            List<Notification.MessagingStyle.Message> newHistoricMessages,
+            boolean showSpinner) {
 
         List<MessagingMessage> historicMessages = createMessages(newHistoricMessages,
                 true /* isHistoric */);
         List<MessagingMessage> messages = createMessages(newMessages, false /* isHistoric */);
-        addMessagesToGroups(historicMessages, messages);
+        addMessagesToGroups(historicMessages, messages, showSpinner);
 
         // Let's remove the remaining messages
         mMessages.forEach(REMOVE_MESSAGE);
@@ -308,7 +311,7 @@
     }
 
     private void addMessagesToGroups(List<MessagingMessage> historicMessages,
-            List<MessagingMessage> messages) {
+            List<MessagingMessage> messages, boolean showSpinner) {
         // Let's first find our groups!
         List<List<MessagingMessage>> groups = new ArrayList<>();
         List<Person> senders = new ArrayList<>();
@@ -317,11 +320,11 @@
         findGroups(historicMessages, messages, groups, senders);
 
         // Let's now create the views and reorder them accordingly
-        createGroupViews(groups, senders);
+        createGroupViews(groups, senders, showSpinner);
     }
 
     private void createGroupViews(List<List<MessagingMessage>> groups,
-            List<Person> senders) {
+            List<Person> senders, boolean showSpinner) {
         mGroups.clear();
         for (int groupIndex = 0; groupIndex < groups.size(); groupIndex++) {
             List<MessagingMessage> group = groups.get(groupIndex);
@@ -346,6 +349,7 @@
                 nameOverride = mNameReplacement;
             }
             newGroup.setSender(sender, nameOverride);
+            newGroup.setSending(groupIndex == (groups.size() - 1) && showSpinner);
             mGroups.add(newGroup);
 
             if (mMessagingLinearLayout.indexOfChild(newGroup) != groupIndex) {
diff --git a/core/java/com/android/internal/widget/MessagingMessage.java b/core/java/com/android/internal/widget/MessagingMessage.java
index a2cc7cf..f0b6068 100644
--- a/core/java/com/android/internal/widget/MessagingMessage.java
+++ b/core/java/com/android/internal/widget/MessagingMessage.java
@@ -82,6 +82,9 @@
         if (!Objects.equals(message.getDataUri(), ownMessage.getDataUri())) {
             return false;
         }
+        if (message.isRemoteInputHistory() != ownMessage.isRemoteInputHistory()) {
+            return false;
+        }
         return true;
     }
 
diff --git a/core/res/res/layout/notification_material_reply_text.xml b/core/res/res/layout/notification_material_reply_text.xml
index 84603b0..dbf2dd0 100644
--- a/core/res/res/layout/notification_material_reply_text.xml
+++ b/core/res/res/layout/notification_material_reply_text.xml
@@ -30,6 +30,7 @@
             android:id="@+id/action_divider"
             android:layout_marginTop="@dimen/notification_content_margin"
             android:layout_marginBottom="@dimen/notification_content_margin"
+            android:layout_marginEnd="@dimen/notification_content_margin_end"
             android:background="@drawable/notification_template_divider" />
 
     <TextView
@@ -50,12 +51,31 @@
             android:textAppearance="@style/TextAppearance.Material.Notification.Reply"
             android:singleLine="true" />
 
-    <TextView
-            android:id="@+id/notification_material_reply_text_1"
+    <LinearLayout
+            android:id="@+id/notification_material_reply_text_1_container"
             android:layout_width="match_parent"
             android:layout_height="wrap_content"
+            android:orientation="horizontal"
+            android:layout_marginStart="@dimen/notification_content_margin_start"
+            android:layout_marginEnd="@dimen/notification_content_margin_end">
+        <TextView
+                android:id="@+id/notification_material_reply_text_1"
+                android:layout_width="match_parent"
+                android:layout_height="wrap_content"
+                android:layout_weight="1"
+                android:layout_marginEnd="@dimen/notification_content_margin_end"
+                android:layout_gravity="center"
+                android:textAppearance="@style/TextAppearance.Material.Notification.Reply"
+                android:singleLine="true" />
+        <ProgressBar
+            android:id="@+id/notification_material_reply_progress"
+            android:layout_height="@dimen/messaging_group_sending_progress_size"
+            android:layout_width="@dimen/messaging_group_sending_progress_size"
+            android:layout_marginStart="@dimen/notification_content_margin_start"
             android:layout_marginEnd="@dimen/notification_content_margin_end"
-            android:textAppearance="@style/TextAppearance.Material.Notification.Reply"
-            android:singleLine="true" />
+            android:layout_gravity="center"
+            android:indeterminate="true"
+            style="?android:attr/progressBarStyleSmall" />
+    </LinearLayout>
 
 </LinearLayout>
diff --git a/core/res/res/layout/notification_template_messaging_group.xml b/core/res/res/layout/notification_template_messaging_group.xml
index bd1030e..d2fd467 100644
--- a/core/res/res/layout/notification_template_messaging_group.xml
+++ b/core/res/res/layout/notification_template_messaging_group.xml
@@ -50,4 +50,19 @@
         android:layout_height="@dimen/messaging_avatar_size"
         android:layout_marginStart="12dp"
         android:visibility="gone"/>
+    <FrameLayout
+        android:id="@+id/messaging_group_sending_progress_container"
+        android:layout_width="@dimen/messaging_group_sending_progress_size"
+        android:layout_height="@dimen/messaging_avatar_size"
+        android:layout_marginStart="12dp"
+        android:layout_gravity="top"
+        android:visibility="gone">
+        <ProgressBar
+            android:id="@+id/messaging_group_sending_progress"
+            android:layout_height="@dimen/messaging_group_sending_progress_size"
+            android:layout_width="@dimen/messaging_group_sending_progress_size"
+            android:layout_gravity="center"
+            android:indeterminate="true"
+            style="?android:attr/progressBarStyleSmall" />
+    </FrameLayout>
 </com.android.internal.widget.MessagingGroup>
diff --git a/core/res/res/values/colors.xml b/core/res/res/values/colors.xml
index da9fed0..095a632 100644
--- a/core/res/res/values/colors.xml
+++ b/core/res/res/values/colors.xml
@@ -137,6 +137,7 @@
     <color name="notification_primary_text_color_light">@color/primary_text_default_material_light</color>
     <color name="notification_primary_text_color_dark">@color/primary_text_default_material_dark</color>
     <color name="notification_secondary_text_color_light">@color/primary_text_default_material_light</color>
+    <item name="notification_secondary_text_disabled_alpha" format="float" type="dimen">0.30</item>
     <color name="notification_secondary_text_color_dark">@color/primary_text_default_material_dark</color>
     <color name="notification_default_color_dark">@color/primary_text_default_material_light</color>
     <color name="notification_default_color_light">#a3202124</color>
diff --git a/core/res/res/values/dimens.xml b/core/res/res/values/dimens.xml
index a135b28..3e0a8a6 100644
--- a/core/res/res/values/dimens.xml
+++ b/core/res/res/values/dimens.xml
@@ -635,6 +635,8 @@
 
     <dimen name="messaging_avatar_size">@dimen/notification_right_icon_size</dimen>
 
+    <dimen name="messaging_group_sending_progress_size">24dp</dimen>
+
     <!-- Max width/height of the autofill data set picker as a fraction of the screen width/height -->
     <dimen name="autofill_dataset_picker_max_width">90%</dimen>
     <dimen name="autofill_dataset_picker_max_height">90%</dimen>
diff --git a/core/res/res/values/symbols.xml b/core/res/res/values/symbols.xml
index 590f988..55e69a3 100644
--- a/core/res/res/values/symbols.xml
+++ b/core/res/res/values/symbols.xml
@@ -2668,8 +2668,10 @@
 
   <java-symbol type="id" name="notification_material_reply_container" />
   <java-symbol type="id" name="notification_material_reply_text_1" />
+  <java-symbol type="id" name="notification_material_reply_text_1_container" />
   <java-symbol type="id" name="notification_material_reply_text_2" />
   <java-symbol type="id" name="notification_material_reply_text_3" />
+  <java-symbol type="id" name="notification_material_reply_progress" />
 
   <java-symbol type="string" name="notification_hidden_text" />
   <java-symbol type="id" name="app_name_text" />
@@ -2999,6 +3001,7 @@
   <java-symbol type="color" name="notification_secondary_text_color_dark" />
   <java-symbol type="color" name="notification_default_color_light" />
   <java-symbol type="color" name="notification_default_color_dark" />
+  <java-symbol type="dimen" name="notification_secondary_text_disabled_alpha" />
 
   <java-symbol type="string" name="app_category_game" />
   <java-symbol type="string" name="app_category_audio" />
@@ -3254,11 +3257,14 @@
   <java-symbol type="id" name="clip_children_tag" />
   <java-symbol type="drawable" name="ic_reply_notification_large" />
   <java-symbol type="dimen" name="messaging_avatar_size" />
+  <java-symbol type="dimen" name="messaging_group_sending_progress_size" />
   <java-symbol type="dimen" name="messaging_image_rounding" />
   <java-symbol type="dimen" name="messaging_image_min_size" />
   <java-symbol type="dimen" name="messaging_image_max_height" />
   <java-symbol type="dimen" name="messaging_image_extra_spacing" />
   <java-symbol type="id" name="messaging_group_icon_container" />
+  <java-symbol type="id" name="messaging_group_sending_progress" />
+  <java-symbol type="id" name="messaging_group_sending_progress_container" />
 
   <java-symbol type="integer" name="config_stableDeviceDisplayWidth" />
   <java-symbol type="integer" name="config_stableDeviceDisplayHeight" />
diff --git a/packages/SystemUI/src/com/android/systemui/Dependency.java b/packages/SystemUI/src/com/android/systemui/Dependency.java
index 5fce0a6..7d19784 100644
--- a/packages/SystemUI/src/com/android/systemui/Dependency.java
+++ b/packages/SystemUI/src/com/android/systemui/Dependency.java
@@ -21,6 +21,7 @@
 import android.os.HandlerThread;
 import android.os.Looper;
 import android.os.Process;
+import android.os.ServiceManager;
 import android.util.ArrayMap;
 import android.view.IWindowManager;
 import android.view.WindowManagerGlobal;
@@ -28,6 +29,7 @@
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.app.ColorDisplayController;
 import com.android.internal.logging.MetricsLogger;
+import com.android.internal.statusbar.IStatusBarService;
 import com.android.internal.util.Preconditions;
 import com.android.settingslib.bluetooth.LocalBluetoothManager;
 import com.android.systemui.assist.AssistManager;
@@ -320,6 +322,9 @@
 
         mProviders.put(VibratorHelper.class, () -> new VibratorHelper(mContext));
 
+        mProviders.put(IStatusBarService.class, () -> IStatusBarService.Stub.asInterface(
+                ServiceManager.getService(Context.STATUS_BAR_SERVICE)));
+
         // Put all dependencies above here so the factory can override them if it wants.
         SystemUIFactory.getInstance().injectDependencies(mProviders, mContext);
     }
diff --git a/packages/SystemUI/src/com/android/systemui/SystemUIFactory.java b/packages/SystemUI/src/com/android/systemui/SystemUIFactory.java
index ac4da73..721890f 100644
--- a/packages/SystemUI/src/com/android/systemui/SystemUIFactory.java
+++ b/packages/SystemUI/src/com/android/systemui/SystemUIFactory.java
@@ -20,7 +20,6 @@
 import android.content.Context;
 import android.util.ArrayMap;
 import android.util.Log;
-import android.view.View;
 import android.view.ViewGroup;
 
 import com.android.internal.logging.MetricsLogger;
@@ -43,12 +42,11 @@
 import com.android.systemui.statusbar.NotificationRemoteInputManager;
 import com.android.systemui.statusbar.NotificationViewHierarchyManager;
 import com.android.systemui.statusbar.ScrimView;
-import com.android.systemui.statusbar.SmartReplyLogger;
+import com.android.systemui.statusbar.SmartReplyController;
 import com.android.systemui.statusbar.notification.VisualStabilityManager;
 import com.android.systemui.statusbar.phone.DozeParameters;
 import com.android.systemui.statusbar.phone.KeyguardBouncer;
 import com.android.systemui.statusbar.phone.KeyguardDismissUtil;
-import com.android.systemui.statusbar.phone.LightBarController;
 import com.android.systemui.statusbar.phone.LockIcon;
 import com.android.systemui.statusbar.phone.LockscreenWallpaper;
 import com.android.systemui.statusbar.phone.NotificationGroupManager;
@@ -58,7 +56,6 @@
 import com.android.systemui.statusbar.phone.StatusBar;
 import com.android.systemui.statusbar.phone.StatusBarIconController;
 import com.android.systemui.statusbar.phone.StatusBarKeyguardViewManager;
-import com.android.systemui.statusbar.policy.DeviceProvisionedController;
 import com.android.systemui.statusbar.policy.SmartReplyConstants;
 
 import java.util.function.Consumer;
@@ -150,6 +147,6 @@
                 () -> new NotificationViewHierarchyManager(context));
         providers.put(NotificationEntryManager.class, () -> new NotificationEntryManager(context));
         providers.put(KeyguardDismissUtil.class, KeyguardDismissUtil::new);
-        providers.put(SmartReplyLogger.class, () -> new SmartReplyLogger(context));
+        providers.put(SmartReplyController.class, () -> new SmartReplyController());
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationContentView.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationContentView.java
index e527be1..5477468 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationContentView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationContentView.java
@@ -81,7 +81,7 @@
 
     private SmartReplyConstants mSmartReplyConstants;
     private SmartReplyView mExpandedSmartReplyView;
-    private SmartReplyLogger mSmartReplyLogger;
+    private SmartReplyController mSmartReplyController;
 
     private NotificationViewWrapper mContractedWrapper;
     private NotificationViewWrapper mExpandedWrapper;
@@ -154,7 +154,7 @@
         super(context, attrs);
         mHybridGroupManager = new HybridGroupManager(getContext(), this);
         mSmartReplyConstants = Dependency.get(SmartReplyConstants.class);
-        mSmartReplyLogger = Dependency.get(SmartReplyLogger.class);
+        mSmartReplyController = Dependency.get(SmartReplyController.class);
         initView();
     }
 
@@ -1359,7 +1359,7 @@
                     applySmartReplyView(mExpandedChild, remoteInput, pendingIntent, entry);
             if (mExpandedSmartReplyView != null && remoteInput != null
                     && remoteInput.getChoices() != null && remoteInput.getChoices().length > 0) {
-                mSmartReplyLogger.smartRepliesAdded(entry, remoteInput.getChoices().length);
+                mSmartReplyController.smartRepliesAdded(entry, remoteInput.getChoices().length);
             }
         }
     }
@@ -1377,6 +1377,13 @@
             smartReplyContainer.setVisibility(View.GONE);
             return null;
         }
+        // If we are showing the spinner we don't want to add the buttons.
+        boolean showingSpinner = entry.notification.getNotification()
+                .extras.getBoolean(Notification.EXTRA_SHOW_REMOTE_INPUT_SPINNER, false);
+        if (showingSpinner) {
+            smartReplyContainer.setVisibility(View.GONE);
+            return null;
+        }
         SmartReplyView smartReplyView = null;
         if (smartReplyContainer.getChildCount() == 0) {
             smartReplyView = SmartReplyView.inflate(mContext, smartReplyContainer);
@@ -1389,7 +1396,7 @@
         }
         if (smartReplyView != null) {
             smartReplyView.setRepliesFromRemoteInput(remoteInput, pendingIntent,
-                    mSmartReplyLogger, entry);
+                    mSmartReplyController, entry, smartReplyContainer);
             smartReplyContainer.setVisibility(View.VISIBLE);
         }
         return smartReplyView;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationEntryManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationEntryManager.java
index 3a79e70..7681530 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationEntryManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationEntryManager.java
@@ -474,37 +474,12 @@
         if (FORCE_REMOTE_INPUT_HISTORY
                 && shouldKeepForRemoteInput(entry)
                 && entry.row != null && !entry.row.isDismissed()) {
-            StatusBarNotification sbn = entry.notification;
-
-            Notification.Builder b = Notification.Builder
-                    .recoverBuilder(mContext, sbn.getNotification().clone());
-            CharSequence[] oldHistory = sbn.getNotification().extras
-                    .getCharSequenceArray(Notification.EXTRA_REMOTE_INPUT_HISTORY);
-            CharSequence[] newHistory;
-            if (oldHistory == null) {
-                newHistory = new CharSequence[1];
-            } else {
-                newHistory = new CharSequence[oldHistory.length + 1];
-                System.arraycopy(oldHistory, 0, newHistory, 1, oldHistory.length);
-            }
             CharSequence remoteInputText = entry.remoteInputText;
             if (TextUtils.isEmpty(remoteInputText)) {
                 remoteInputText = entry.remoteInputTextWhenReset;
             }
-            newHistory[0] = String.valueOf(remoteInputText);
-            b.setRemoteInputHistory(newHistory);
-
-            Notification newNotification = b.build();
-
-            // Undo any compatibility view inflation
-            newNotification.contentView = sbn.getNotification().contentView;
-            newNotification.bigContentView = sbn.getNotification().bigContentView;
-            newNotification.headsUpContentView = sbn.getNotification().headsUpContentView;
-
-            StatusBarNotification newSbn = new StatusBarNotification(sbn.getPackageName(),
-                    sbn.getOpPkg(),
-                    sbn.getId(), sbn.getTag(), sbn.getUid(), sbn.getInitialPid(),
-                    newNotification, sbn.getUser(), sbn.getOverrideGroupKey(), sbn.getPostTime());
+            StatusBarNotification newSbn = rebuildNotificationWithRemoteInput(entry,
+                    remoteInputText, false /* showSpinner */);
             boolean updated = false;
             entry.onRemoteInputInserted();
             try {
@@ -554,6 +529,39 @@
         mCallback.onNotificationRemoved(key, old);
     }
 
+    public StatusBarNotification rebuildNotificationWithRemoteInput(NotificationData.Entry entry,
+            CharSequence remoteInputText, boolean showSpinner) {
+        StatusBarNotification sbn = entry.notification;
+
+        Notification.Builder b = Notification.Builder
+                .recoverBuilder(mContext, sbn.getNotification().clone());
+        CharSequence[] oldHistory = sbn.getNotification().extras
+                .getCharSequenceArray(Notification.EXTRA_REMOTE_INPUT_HISTORY);
+        CharSequence[] newHistory;
+        if (oldHistory == null) {
+            newHistory = new CharSequence[1];
+        } else {
+            newHistory = new CharSequence[oldHistory.length + 1];
+            System.arraycopy(oldHistory, 0, newHistory, 1, oldHistory.length);
+        }
+        newHistory[0] = String.valueOf(remoteInputText);
+        b.setRemoteInputHistory(newHistory);
+        b.setShowRemoteInputSpinner(showSpinner);
+
+        Notification newNotification = b.build();
+
+        // Undo any compatibility view inflation
+        newNotification.contentView = sbn.getNotification().contentView;
+        newNotification.bigContentView = sbn.getNotification().bigContentView;
+        newNotification.headsUpContentView = sbn.getNotification().headsUpContentView;
+
+        StatusBarNotification newSbn = new StatusBarNotification(sbn.getPackageName(),
+                sbn.getOpPkg(),
+                sbn.getId(), sbn.getTag(), sbn.getUid(), sbn.getInitialPid(),
+                newNotification, sbn.getUser(), sbn.getOverrideGroupKey(), sbn.getPostTime());
+        return newSbn;
+    }
+
     private boolean shouldKeepForRemoteInput(NotificationData.Entry entry) {
         if (entry == null) {
             return false;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyLogger.java b/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java
similarity index 64%
rename from packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyLogger.java
rename to packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java
index 75dd77d..67da68c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyLogger.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java
@@ -15,24 +15,32 @@
  */
 package com.android.systemui.statusbar;
 
-import android.content.Context;
 import android.os.RemoteException;
-import android.os.ServiceManager;
+import android.service.notification.StatusBarNotification;
+
 import com.android.internal.statusbar.IStatusBarService;
+import com.android.systemui.Dependency;
+
 
 /**
- * Handles reporting when smart replies are added to a notification
+ * Handles when smart replies are added to a notification
  * and clicked upon.
  */
-public class SmartReplyLogger {
-    protected IStatusBarService mBarService;
+public class SmartReplyController {
+    private IStatusBarService mBarService;
+    private NotificationEntryManager mNotificationEntryManager;
 
-    public SmartReplyLogger(Context context) {
-        mBarService = IStatusBarService.Stub.asInterface(
-                ServiceManager.getService(Context.STATUS_BAR_SERVICE));
+    public SmartReplyController() {
+        mBarService = Dependency.get(IStatusBarService.class);
+        mNotificationEntryManager = Dependency.get(NotificationEntryManager.class);
     }
 
-    public void smartReplySent(NotificationData.Entry entry, int replyIndex) {
+    public void smartReplySent(NotificationData.Entry entry, int replyIndex, CharSequence reply) {
+        StatusBarNotification newSbn =
+                mNotificationEntryManager.rebuildNotificationWithRemoteInput(entry, reply,
+                        true /* showSpinner */);
+        mNotificationEntryManager.updateNotification(newSbn, null /* ranking */);
+
         try {
             mBarService.onNotificationSmartReplySent(entry.notification.getKey(),
                     replyIndex);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/SmartReplyView.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/SmartReplyView.java
index 1431682..b4fa2e8 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/SmartReplyView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/SmartReplyView.java
@@ -26,7 +26,7 @@
 import com.android.systemui.Dependency;
 import com.android.systemui.R;
 import com.android.systemui.statusbar.NotificationData;
-import com.android.systemui.statusbar.SmartReplyLogger;
+import com.android.systemui.statusbar.SmartReplyController;
 import com.android.systemui.statusbar.notification.NotificationUtils;
 import com.android.systemui.statusbar.phone.KeyguardDismissUtil;
 
@@ -73,6 +73,8 @@
 
     private PriorityQueue<Button> mCandidateButtonQueueForSqueezing;
 
+    private View mSmartReplyContainer;
+
     public SmartReplyView(Context context, AttributeSet attrs) {
         super(context, attrs);
         mConstants = Dependency.get(SmartReplyConstants.class);
@@ -133,7 +135,9 @@
     }
 
     public void setRepliesFromRemoteInput(RemoteInput remoteInput, PendingIntent pendingIntent,
-            SmartReplyLogger smartReplyLogger, NotificationData.Entry entry) {
+            SmartReplyController smartReplyController, NotificationData.Entry entry,
+            View smartReplyContainer) {
+        mSmartReplyContainer = smartReplyContainer;
         removeAllViews();
         if (remoteInput != null && pendingIntent != null) {
             CharSequence[] choices = remoteInput.getChoices();
@@ -141,7 +145,7 @@
                 for (int i = 0; i < choices.length; ++i) {
                     Button replyButton = inflateReplyButton(
                             getContext(), this, i, choices[i], remoteInput, pendingIntent,
-                            smartReplyLogger, entry);
+                            smartReplyController, entry);
                     addView(replyButton);
                 }
             }
@@ -157,7 +161,7 @@
     @VisibleForTesting
     Button inflateReplyButton(Context context, ViewGroup root, int replyIndex,
             CharSequence choice, RemoteInput remoteInput, PendingIntent pendingIntent,
-            SmartReplyLogger smartReplyLogger, NotificationData.Entry entry) {
+            SmartReplyController smartReplyController, NotificationData.Entry entry) {
         Button b = (Button) LayoutInflater.from(context).inflate(
                 R.layout.smart_reply_button, root, false);
         b.setText(choice);
@@ -173,7 +177,8 @@
             } catch (PendingIntent.CanceledException e) {
                 Log.w(TAG, "Unable to send smart reply", e);
             }
-            smartReplyLogger.smartReplySent(entry, replyIndex);
+            smartReplyController.smartReplySent(entry, replyIndex, b.getText());
+            mSmartReplyContainer.setVisibility(View.GONE);
             return false; // do not defer
         };
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationEntryManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationEntryManagerTest.java
index 7cfd7a3..8172626 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationEntryManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationEntryManagerTest.java
@@ -367,4 +367,44 @@
         mEntryManager.tagForeground(mEntry.notification);
         Assert.assertEquals(0, mEntry.mActiveAppOps.size());
     }
+
+    @Test
+    public void testRebuildWithRemoteInput_noExistingInputNoSpinner() {
+        StatusBarNotification newSbn =
+                mEntryManager.rebuildNotificationWithRemoteInput(mEntry, "A Reply", false);
+        CharSequence[] messages = newSbn.getNotification().extras
+                .getCharSequenceArray(Notification.EXTRA_REMOTE_INPUT_HISTORY);
+        Assert.assertEquals(1, messages.length);
+        Assert.assertEquals("A Reply", messages[0]);
+        Assert.assertFalse(newSbn.getNotification().extras
+                .getBoolean(Notification.EXTRA_SHOW_REMOTE_INPUT_SPINNER, false));
+    }
+
+    @Test
+    public void testRebuildWithRemoteInput_noExistingInputWithSpinner() {
+        StatusBarNotification newSbn =
+                mEntryManager.rebuildNotificationWithRemoteInput(mEntry, "A Reply", true);
+        CharSequence[] messages = newSbn.getNotification().extras
+                .getCharSequenceArray(Notification.EXTRA_REMOTE_INPUT_HISTORY);
+        Assert.assertEquals(1, messages.length);
+        Assert.assertEquals("A Reply", messages[0]);
+        Assert.assertTrue(newSbn.getNotification().extras
+                .getBoolean(Notification.EXTRA_SHOW_REMOTE_INPUT_SPINNER, false));
+    }
+
+    @Test
+    public void testRebuildWithRemoteInput_withExistingInput() {
+        // Setup a notification entry with 1 remote input.
+        StatusBarNotification newSbn =
+                mEntryManager.rebuildNotificationWithRemoteInput(mEntry, "A Reply", false);
+        NotificationData.Entry entry = new NotificationData.Entry(newSbn);
+
+        // Try rebuilding to add another reply.
+        newSbn = mEntryManager.rebuildNotificationWithRemoteInput(entry, "Reply 2", true);
+        CharSequence[] messages = newSbn.getNotification().extras
+                .getCharSequenceArray(Notification.EXTRA_REMOTE_INPUT_HISTORY);
+        Assert.assertEquals(2, messages.length);
+        Assert.assertEquals("Reply 2", messages[0]);
+        Assert.assertEquals("A Reply", messages[1]);
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java
new file mode 100644
index 0000000..84dceae
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java
@@ -0,0 +1,121 @@
+/*
+ * Copyright (C) 2018 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.systemui.statusbar;
+
+import static org.mockito.ArgumentMatchers.argThat;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.isNull;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.app.Notification;
+import android.os.RemoteException;
+import android.service.notification.StatusBarNotification;
+import android.support.test.filters.SmallTest;
+import android.testing.AndroidTestingRunner;
+import android.testing.TestableLooper;
+
+import com.android.internal.statusbar.IStatusBarService;
+import com.android.systemui.R;
+import com.android.systemui.SysuiTestCase;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+@RunWith(AndroidTestingRunner.class)
+@TestableLooper.RunWithLooper
+@SmallTest
+public class SmartReplyControllerTest extends SysuiTestCase {
+    private static final String TEST_NOTIFICATION_KEY = "akey";
+    private static final String TEST_CHOICE_TEXT = "A Reply";
+    private static final int TEST_CHOICE_INDEX = 2;
+    private static final int TEST_CHOICE_COUNT = 4;
+
+    private Notification mNotification;
+    private NotificationData.Entry mEntry;
+
+    @Mock
+    private NotificationEntryManager mNotificationEntryManager;
+    @Mock
+    private IStatusBarService mIStatusBarService;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+
+        mDependency.injectTestDependency(NotificationEntryManager.class,
+                mNotificationEntryManager);
+        mDependency.injectTestDependency(IStatusBarService.class, mIStatusBarService);
+
+        mNotification = new Notification.Builder(mContext, "")
+                .setSmallIcon(R.drawable.ic_person)
+                .setContentTitle("Title")
+                .setContentText("Text").build();
+        StatusBarNotification sbn = mock(StatusBarNotification.class);
+        when(sbn.getNotification()).thenReturn(mNotification);
+        when(sbn.getKey()).thenReturn(TEST_NOTIFICATION_KEY);
+        mEntry = new NotificationData.Entry(sbn);
+    }
+
+    @Test
+    public void testSendSmartReply_updatesRemoteInput() throws RemoteException {
+        StatusBarNotification sbn = mock(StatusBarNotification.class);
+        when(sbn.getKey()).thenReturn(TEST_NOTIFICATION_KEY);
+        when(mNotificationEntryManager.rebuildNotificationWithRemoteInput(
+                argThat(entry -> entry.notification.getKey().equals(TEST_NOTIFICATION_KEY)),
+                eq(TEST_CHOICE_TEXT), eq(true))).thenReturn(sbn);
+
+        SmartReplyController controller = new SmartReplyController();
+        controller.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT);
+
+        // Sending smart reply should make calls to NotificationEntryManager
+        // to update the notification with reply and spinner.
+        verify(mNotificationEntryManager).rebuildNotificationWithRemoteInput(
+                argThat(entry -> entry.notification.getKey().equals(TEST_NOTIFICATION_KEY)),
+                eq(TEST_CHOICE_TEXT), eq(true));
+        verify(mNotificationEntryManager).updateNotification(
+                argThat(sbn2 -> sbn2.getKey().equals(TEST_NOTIFICATION_KEY)), isNull());
+    }
+
+    @Test
+    public void testSendSmartReply_logsToStatusBar() throws RemoteException {
+        StatusBarNotification sbn = mock(StatusBarNotification.class);
+        when(sbn.getKey()).thenReturn(TEST_NOTIFICATION_KEY);
+        when(mNotificationEntryManager.rebuildNotificationWithRemoteInput(
+                argThat(entry -> entry.notification.getKey().equals(TEST_NOTIFICATION_KEY)),
+                eq(TEST_CHOICE_TEXT), eq(true))).thenReturn(sbn);
+
+        SmartReplyController controller = new SmartReplyController();
+        controller.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT);
+
+        // Check we log the result to the status bar service.
+        verify(mIStatusBarService).onNotificationSmartReplySent(TEST_NOTIFICATION_KEY,
+                TEST_CHOICE_INDEX);
+    }
+
+    @Test
+    public void testShowSmartReply_logsToStatusBar() throws RemoteException {
+        SmartReplyController controller = new SmartReplyController();
+        controller.smartRepliesAdded(mEntry, TEST_CHOICE_COUNT);
+
+        // Check we log the result to the status bar service.
+        verify(mIStatusBarService).onNotificationSmartRepliesAdded(TEST_NOTIFICATION_KEY,
+                TEST_CHOICE_COUNT);
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SmartReplyViewTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SmartReplyViewTest.java
index 2bb8106..99c06e6 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SmartReplyViewTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SmartReplyViewTest.java
@@ -23,9 +23,10 @@
 import static junit.framework.Assert.fail;
 
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
+import android.app.Notification;
 import android.app.PendingIntent;
 import android.app.RemoteInput;
 import android.content.Intent;
@@ -44,7 +45,7 @@
 import com.android.systemui.R;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.statusbar.NotificationData;
-import com.android.systemui.statusbar.SmartReplyLogger;
+import com.android.systemui.statusbar.SmartReplyController;
 import com.android.systemui.statusbar.phone.KeyguardDismissUtil;
 
 import java.util.concurrent.atomic.AtomicReference;
@@ -64,19 +65,22 @@
     private static final String TEST_ACTION = "com.android.SMART_REPLY_VIEW_ACTION";
 
     private static final String[] TEST_CHOICES = new String[]{"Hello", "What's up?", "I'm here"};
+    private static final String TEST_NOTIFICATION_KEY = "akey";
 
     private static final int WIDTH_SPEC = MeasureSpec.makeMeasureSpec(500, MeasureSpec.EXACTLY);
     private static final int HEIGHT_SPEC = MeasureSpec.makeMeasureSpec(400, MeasureSpec.AT_MOST);
 
     private BlockingQueueIntentReceiver mReceiver;
     private SmartReplyView mView;
+    private View mContainer;
 
     private int mSingleLinePaddingHorizontal;
     private int mDoubleLinePaddingHorizontal;
     private int mSpacing;
 
-    @Mock private SmartReplyLogger mLogger;
+    @Mock private SmartReplyController mLogger;
     private NotificationData.Entry mEntry;
+    private Notification mNotification;
 
     @Before
     public void setUp() {
@@ -86,6 +90,7 @@
         mDependency.get(KeyguardDismissUtil.class).setDismissHandler(
             (action, cancelAction, afterKeyguardGone) -> action.onDismiss());
 
+        mContainer = new View(mContext, null);
         mView = SmartReplyView.inflate(mContext, null);
 
 
@@ -96,9 +101,14 @@
                 R.dimen.smart_reply_button_padding_horizontal_double_line);
         mSpacing = res.getDimensionPixelSize(R.dimen.smart_reply_button_spacing);
 
-        StatusBarNotification notification = mock(StatusBarNotification.class);
-        when(notification.getKey()).thenReturn("akey");
-        mEntry = new NotificationData.Entry(notification);
+        mNotification = new Notification.Builder(mContext, "")
+                .setSmallIcon(R.drawable.ic_person)
+                .setContentTitle("Title")
+                .setContentText("Text").build();
+        StatusBarNotification sbn = mock(StatusBarNotification.class);
+        when(sbn.getNotification()).thenReturn(mNotification);
+        when(sbn.getKey()).thenReturn(TEST_NOTIFICATION_KEY);
+        mEntry = new NotificationData.Entry(sbn);
     }
 
     @After
@@ -155,10 +165,18 @@
     }
 
     @Test
-    public void testSendSmartReply_LoggerCall() {
+    public void testSendSmartReply_controllerCalled() {
         setRepliesFromRemoteInput(TEST_CHOICES);
         mView.getChildAt(2).performClick();
-        verify(mLogger).smartReplySent(mEntry, 2);
+        verify(mLogger).smartReplySent(mEntry, 2, TEST_CHOICES[2]);
+    }
+
+    @Test
+    public void testSendSmartReply_hidesContainer() {
+        mContainer.setVisibility(View.VISIBLE);
+        setRepliesFromRemoteInput(TEST_CHOICES);
+        mView.getChildAt(0).performClick();
+        assertEquals(View.GONE, mContainer.getVisibility());
     }
 
     @Test
@@ -340,7 +358,7 @@
         PendingIntent pendingIntent = PendingIntent.getBroadcast(mContext, 0,
                 new Intent(TEST_ACTION), 0);
         RemoteInput input = new RemoteInput.Builder(TEST_RESULT_KEY).setChoices(choices).build();
-        mView.setRepliesFromRemoteInput(input, pendingIntent, mLogger, mEntry);
+        mView.setRepliesFromRemoteInput(input, pendingIntent, mLogger, mEntry, mContainer);
     }
 
     /** Builds a {@link ViewGroup} whose measures and layout mirror a {@link SmartReplyView}. */