Move SmartReplyView inflation off the UI thread.

To avoid blocking the UI thread we here move the inflation of
SmartReplyView and its smart suggestion buttons off the UI thread.

We move these inflations into NotificationContentInflater where they are
performed before the inflations of other notification content views -
reusing the existing AsyncTask to run these inflations on a background
thread.
More specifically: for both the expanded state and the heads-up state of
a notification we
    1. inflate a SmartReplyView,
    2. inflate the buttons within that view (reply buttons and action
    buttons), and
    3. create a SmartRepliesAndActions object
from the AsyncTask (on a background thread). We then pass these objects
back to the UI thread where we attach them to the view hierarchy or use
them to make certain decisions.

In terms of thread-safety the background thread reads constants from
SmartReplyConstants so that class must be thread-safe (at least so that
the background thread doesn't store a local copy of the constants).
We also pass a SmartReplyController reference on the background thread,
that reference is only used on the UI thread (in onClick listeners).

Bug:119801785
Test: post message notification containing smart replies and actions.
Post both heads-up and notification shade notifications. Click on
replies and actions. Ensure the smart-buttons are updated when a
notification is updated.

Change-Id: I4698a6895cf65c46461cc429b2b6eb5905acb629

5 files changed, 251 insertions, 72 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentInflater.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentInflater.java
index b65c4a5f71d8..b4dd1144e761 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentInflater.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentInflater.java
@@ -34,12 +34,17 @@ import android.widget.RemoteViews;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.widget.ImageMessageConsumer;
+import com.android.systemui.Dependency;
 import com.android.systemui.statusbar.InflationTask;
+import com.android.systemui.statusbar.SmartReplyController;
 import com.android.systemui.statusbar.notification.InflationException;
 import com.android.systemui.statusbar.notification.MediaNotificationProcessor;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.notification.row.wrapper.NotificationViewWrapper;
 import com.android.systemui.statusbar.phone.StatusBar;
+import com.android.systemui.statusbar.policy.HeadsUpManager;
+import com.android.systemui.statusbar.policy.InflatedSmartReplies;
+import com.android.systemui.statusbar.policy.SmartReplyConstants;
 import com.android.systemui.util.Assert;
 
 import java.lang.annotation.Retention;
@@ -278,6 +283,8 @@ public class NotificationContentInflater {
         InflationProgress result = createRemoteViews(reInflateFlags, builder, mIsLowPriority,
                 mIsChildInGroup, mUsesIncreasedHeight, mUsesIncreasedHeadsUpHeight,
                 mRedactAmbient, packageContext);
+        result = inflateSmartReplyViews(result, reInflateFlags, mRow.getEntry(),
+                mRow.getContext(), mRow.getHeadsUpManager());
         apply(
                 inflateSynchronously,
                 result,
@@ -306,6 +313,7 @@ public class NotificationContentInflater {
                 if (mRow.getPrivateLayout().isContentViewInactive(VISIBLE_TYPE_HEADSUP)) {
                     mRow.getPrivateLayout().setHeadsUpChild(null);
                     mCachedContentViews.remove(FLAG_CONTENT_VIEW_HEADS_UP);
+                    mRow.getPrivateLayout().setHeadsUpInflatedSmartReplies(null);
                 }
                 break;
             case FLAG_CONTENT_VIEW_AMBIENT:
@@ -336,12 +344,33 @@ public class NotificationContentInflater {
         }
     }
 
+    private static InflationProgress inflateSmartReplyViews(InflationProgress result,
+            @InflationFlag int reInflateFlags, NotificationEntry entry, Context context,
+            HeadsUpManager headsUpManager) {
+        SmartReplyConstants smartReplyConstants = Dependency.get(SmartReplyConstants.class);
+        SmartReplyController smartReplyController = Dependency.get(SmartReplyController.class);
+        if ((reInflateFlags & FLAG_CONTENT_VIEW_EXPANDED) != 0 && result.newExpandedView != null) {
+            result.expandedInflatedSmartReplies =
+                    InflatedSmartReplies.inflate(
+                            context, entry, smartReplyConstants, smartReplyController,
+                            headsUpManager);
+        }
+        if ((reInflateFlags & FLAG_CONTENT_VIEW_HEADS_UP) != 0 && result.newHeadsUpView != null) {
+            result.headsUpInflatedSmartReplies =
+                    InflatedSmartReplies.inflate(
+                            context, entry, smartReplyConstants, smartReplyController,
+                            headsUpManager);
+        }
+        return result;
+    }
+
     private static InflationProgress createRemoteViews(@InflationFlag int reInflateFlags,
             Notification.Builder builder, boolean isLowPriority, boolean isChildInGroup,
             boolean usesIncreasedHeight, boolean usesIncreasedHeadsUpHeight, boolean redactAmbient,
             Context packageContext) {
         InflationProgress result = new InflationProgress();
         isLowPriority = isLowPriority && !isChildInGroup;
+
         if ((reInflateFlags & FLAG_CONTENT_VIEW_CONTRACTED) != 0) {
             result.newContentView = createContentView(builder, isLowPriority, usesIncreasedHeight);
         }
@@ -661,6 +690,12 @@ public class NotificationContentInflater {
                 } else if (result.newExpandedView == null) {
                     privateLayout.setExpandedChild(null);
                 }
+                if (result.newExpandedView != null) {
+                    privateLayout.setExpandedInflatedSmartReplies(
+                            result.expandedInflatedSmartReplies);
+                } else {
+                    privateLayout.setExpandedInflatedSmartReplies(null);
+                }
                 cachedContentViews.put(FLAG_CONTENT_VIEW_EXPANDED, result.newExpandedView);
                 row.setExpandable(result.newExpandedView != null);
             }
@@ -671,6 +706,12 @@ public class NotificationContentInflater {
                 } else if (result.newHeadsUpView == null) {
                     privateLayout.setHeadsUpChild(null);
                 }
+                if (result.newHeadsUpView != null) {
+                    privateLayout.setHeadsUpInflatedSmartReplies(
+                            result.headsUpInflatedSmartReplies);
+                } else {
+                    privateLayout.setHeadsUpInflatedSmartReplies(null);
+                }
                 cachedContentViews.put(FLAG_CONTENT_VIEW_HEADS_UP, result.newHeadsUpView);
             }
 
@@ -846,9 +887,12 @@ public class NotificationContentInflater {
                             packageContext);
                     processor.processNotification(notification, recoveredBuilder);
                 }
-                return createRemoteViews(mReInflateFlags, recoveredBuilder, mIsLowPriority,
+                InflationProgress inflationProgress = createRemoteViews(mReInflateFlags,
+                        recoveredBuilder, mIsLowPriority,
                         mIsChildInGroup, mUsesIncreasedHeight, mUsesIncreasedHeadsUpHeight,
                         mRedactAmbient, packageContext);
+                return inflateSmartReplyViews(inflationProgress, mReInflateFlags, mRow.getEntry(),
+                        mRow.getContext(), mRow.getHeadsUpManager());
             } catch (Exception e) {
                 mError = e;
                 return null;
@@ -927,6 +971,9 @@ public class NotificationContentInflater {
         private View inflatedPublicView;
         private CharSequence headsUpStatusBarText;
         private CharSequence headsUpStatusBarTextPublic;
+
+        private InflatedSmartReplies expandedInflatedSmartReplies;
+        private InflatedSmartReplies headsUpInflatedSmartReplies;
     }
 
     @VisibleForTesting
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentView.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentView.java
index dae14debd683..646617c72128 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentView.java
@@ -93,6 +93,8 @@ public class NotificationContentView extends FrameLayout {
     private SmartReplyView mExpandedSmartReplyView;
     private SmartReplyView mHeadsUpSmartReplyView;
     private SmartReplyController mSmartReplyController;
+    private InflatedSmartReplies mExpandedInflatedSmartReplies;
+    private InflatedSmartReplies mHeadsUpInflatedSmartReplies;
 
     private NotificationViewWrapper mContractedWrapper;
     private NotificationViewWrapper mExpandedWrapper;
@@ -1316,8 +1318,22 @@ public class NotificationContentView extends FrameLayout {
             return;
         }
 
-        SmartRepliesAndActions smartRepliesAndActions =
-                InflatedSmartReplies.chooseSmartRepliesAndActions(mSmartReplyConstants, entry);
+        applyRemoteInput(entry, InflatedSmartReplies.hasFreeformRemoteInput(entry));
+
+        if (mExpandedInflatedSmartReplies == null && mHeadsUpInflatedSmartReplies == null) {
+            if (DEBUG) {
+                Log.d(TAG, "Both expanded, and heads-up InflatedSmartReplies are null, "
+                        + "don't add smart replies.");
+            }
+            return;
+        }
+        // The inflated smart-reply objects for the expanded view and the heads-up view both contain
+        // the same SmartRepliesAndActions to avoid discrepancies between the two views. We here
+        // reuse that object for our local SmartRepliesAndActions to avoid discrepancies between
+        // this class and the InflatedSmartReplies classes.
+        SmartRepliesAndActions smartRepliesAndActions = mExpandedInflatedSmartReplies != null
+                ? mExpandedInflatedSmartReplies.getSmartRepliesAndActions()
+                : mHeadsUpInflatedSmartReplies.getSmartRepliesAndActions();
         if (DEBUG) {
             Log.d(TAG, String.format("Adding suggestions for %s, %d actions, and %d replies.",
                     entry.notification.getKey(),
@@ -1326,7 +1342,6 @@ public class NotificationContentView extends FrameLayout {
                     smartRepliesAndActions.smartReplies == null ? 0 :
                             smartRepliesAndActions.smartReplies.choices.length));
         }
-        applyRemoteInput(entry, InflatedSmartReplies.hasFreeformRemoteInput(entry));
         applySmartReplyView(smartRepliesAndActions, entry);
     }
 
@@ -1429,11 +1444,12 @@ public class NotificationContentView extends FrameLayout {
         return null;
     }
 
-    private void applySmartReplyView(SmartRepliesAndActions smartRepliesAndActions,
+    private void applySmartReplyView(
+            SmartRepliesAndActions smartRepliesAndActions,
             NotificationEntry entry) {
         if (mExpandedChild != null) {
-            mExpandedSmartReplyView =
-                    applySmartReplyView(mExpandedChild, smartRepliesAndActions, entry);
+            mExpandedSmartReplyView = applySmartReplyView(mExpandedChild, smartRepliesAndActions,
+                    entry, mExpandedInflatedSmartReplies);
             if (mExpandedSmartReplyView != null) {
                 if (smartRepliesAndActions.smartReplies != null
                         || smartRepliesAndActions.smartActions != null) {
@@ -1455,18 +1471,21 @@ public class NotificationContentView extends FrameLayout {
             }
         }
         if (mHeadsUpChild != null && mSmartReplyConstants.getShowInHeadsUp()) {
-            mHeadsUpSmartReplyView =
-                    applySmartReplyView(mHeadsUpChild, smartRepliesAndActions, entry);
+            mHeadsUpSmartReplyView = applySmartReplyView(mHeadsUpChild, smartRepliesAndActions,
+                    entry, mHeadsUpInflatedSmartReplies);
         }
     }
 
+    @Nullable
     private SmartReplyView applySmartReplyView(View view,
-            SmartRepliesAndActions smartRepliesAndActions, NotificationEntry entry) {
+            SmartRepliesAndActions smartRepliesAndActions,
+            NotificationEntry entry, InflatedSmartReplies inflatedSmartReplyView) {
         View smartReplyContainerCandidate = view.findViewById(
                 com.android.internal.R.id.smart_reply_container);
         if (!(smartReplyContainerCandidate instanceof LinearLayout)) {
             return null;
         }
+
         LinearLayout smartReplyContainer = (LinearLayout) smartReplyContainerCandidate;
         if (!InflatedSmartReplies.shouldShowSmartReplyView(entry, smartRepliesAndActions)) {
             smartReplyContainer.setVisibility(View.GONE);
@@ -1474,31 +1493,57 @@ public class NotificationContentView extends FrameLayout {
         }
 
         SmartReplyView smartReplyView = null;
-        if (smartReplyContainer.getChildCount() == 0) {
-            smartReplyView = SmartReplyView.inflate(mContext, smartReplyContainer);
+        if (smartReplyContainer.getChildCount() == 1
+                && smartReplyContainer.getChildAt(0) instanceof SmartReplyView) {
+            // If we already have a SmartReplyView - replace it with the newly inflated one. The
+            // newly inflated one is connected to the new inflated smart reply/action buttons.
+            smartReplyContainer.removeAllViews();
+        }
+        if (smartReplyContainer.getChildCount() == 0
+                && inflatedSmartReplyView != null
+                && inflatedSmartReplyView.getSmartReplyView() != null) {
+            smartReplyView = inflatedSmartReplyView.getSmartReplyView();
             smartReplyContainer.addView(smartReplyView);
-        } else if (smartReplyContainer.getChildCount() == 1) {
-            View child = smartReplyContainer.getChildAt(0);
-            if (child instanceof SmartReplyView) {
-                smartReplyView = (SmartReplyView) child;
-            }
         }
         if (smartReplyView != null) {
             smartReplyView.resetSmartSuggestions(smartReplyContainer);
-            if (smartRepliesAndActions.smartReplies != null) {
-                smartReplyView.addRepliesFromRemoteInput(
-                        smartRepliesAndActions.smartReplies, mSmartReplyController, entry);
-            }
-            if (smartRepliesAndActions.smartActions != null) {
-                smartReplyView.addSmartActions(
-                        smartRepliesAndActions.smartActions, mSmartReplyController, entry,
-                        mContainingNotification.getHeadsUpManager());
-            }
+            smartReplyView.addPreInflatedButtons(
+                    inflatedSmartReplyView.getSmartSuggestionButtons());
             smartReplyContainer.setVisibility(View.VISIBLE);
         }
         return smartReplyView;
     }
 
+    /**
+     * Set pre-inflated views necessary to display smart replies and actions in the expanded
+     * notification state.
+     *
+     * @param inflatedSmartReplies the pre-inflated state to add to this view. If null the existing
+     * {@link SmartReplyView} related to the expanded notification state is cleared.
+     */
+    public void setExpandedInflatedSmartReplies(
+            @Nullable InflatedSmartReplies inflatedSmartReplies) {
+        mExpandedInflatedSmartReplies = inflatedSmartReplies;
+        if (inflatedSmartReplies == null) {
+            mExpandedSmartReplyView = null;
+        }
+    }
+
+    /**
+     * Set pre-inflated views necessary to display smart replies and actions in the heads-up
+     * notification state.
+     *
+     * @param inflatedSmartReplies the pre-inflated state to add to this view. If null the existing
+     * {@link SmartReplyView} related to the heads-up notification state is cleared.
+     */
+    public void setHeadsUpInflatedSmartReplies(
+            @Nullable InflatedSmartReplies inflatedSmartReplies) {
+        mHeadsUpInflatedSmartReplies = inflatedSmartReplies;
+        if (inflatedSmartReplies == null) {
+            mHeadsUpSmartReplyView = null;
+        }
+    }
+
     public void closeRemoteInput() {
         if (mHeadsUpRemoteInput != null) {
             mHeadsUpRemoteInput.close();
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/InflatedSmartReplies.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/InflatedSmartReplies.java
index 754bfb2b30aa..d8ea1f6eef5f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/InflatedSmartReplies.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/InflatedSmartReplies.java
@@ -20,13 +20,17 @@ import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.app.Notification;
 import android.app.RemoteInput;
+import android.content.Context;
 import android.os.Build;
 import android.util.Log;
 import android.util.Pair;
+import android.widget.Button;
 
 import com.android.internal.util.ArrayUtils;
+import com.android.systemui.statusbar.SmartReplyController;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 
+import java.util.ArrayList;
 import java.util.List;
 
 /**
@@ -36,8 +40,63 @@ import java.util.List;
 public class InflatedSmartReplies {
     private static final String TAG = "InflatedSmartReplies";
     private static final boolean DEBUG = Log.isLoggable(TAG, Log.DEBUG);
+    @Nullable private final SmartReplyView mSmartReplyView;
+    @Nullable private final List<Button> mSmartSuggestionButtons;
+    @NonNull private final SmartRepliesAndActions mSmartRepliesAndActions;
 
-    private InflatedSmartReplies() { }
+    private InflatedSmartReplies(
+            @Nullable SmartReplyView smartReplyView,
+            @Nullable List<Button> smartSuggestionButtons,
+            @NonNull SmartRepliesAndActions smartRepliesAndActions) {
+        mSmartReplyView = smartReplyView;
+        mSmartSuggestionButtons = smartSuggestionButtons;
+        mSmartRepliesAndActions = smartRepliesAndActions;
+    }
+
+    @Nullable public SmartReplyView getSmartReplyView() {
+        return mSmartReplyView;
+    }
+
+    @Nullable public List<Button> getSmartSuggestionButtons() {
+        return mSmartSuggestionButtons;
+    }
+
+    @NonNull public SmartRepliesAndActions getSmartRepliesAndActions() {
+        return mSmartRepliesAndActions;
+    }
+
+    /**
+     * Inflate a SmartReplyView and its smart suggestions.
+     */
+    public static InflatedSmartReplies inflate(
+            Context context,
+            NotificationEntry entry,
+            SmartReplyConstants smartReplyConstants,
+            SmartReplyController smartReplyController,
+            HeadsUpManager headsUpManager) {
+        SmartRepliesAndActions smartRepliesAndActions =
+                chooseSmartRepliesAndActions(smartReplyConstants, entry);
+        if (!shouldShowSmartReplyView(entry, smartRepliesAndActions)) {
+            return new InflatedSmartReplies(null /* smartReplyView */,
+                    null /* smartSuggestionButtons */, smartRepliesAndActions);
+        }
+
+        SmartReplyView smartReplyView = SmartReplyView.inflate(context);
+
+        List<Button> suggestionButtons = new ArrayList<>();
+        if (smartRepliesAndActions.smartReplies != null) {
+            suggestionButtons.addAll(smartReplyView.inflateRepliesFromRemoteInput(
+                    smartRepliesAndActions.smartReplies, smartReplyController, entry));
+        }
+        if (smartRepliesAndActions.smartActions != null) {
+            suggestionButtons.addAll(
+                    smartReplyView.inflateSmartActions(smartRepliesAndActions.smartActions,
+                            smartReplyController, entry, headsUpManager));
+        }
+
+        return new InflatedSmartReplies(smartReplyView, suggestionButtons,
+                smartRepliesAndActions);
+    }
 
     /**
      * Returns whether we should show the smart reply view and its smart suggestions.
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 45d215ef309c..ed5487f74356 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/SmartReplyView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/SmartReplyView.java
@@ -196,66 +196,81 @@ public class SmartReplyView extends ViewGroup {
     }
 
     /**
+     * Add buttons to the {@link SmartReplyView} - these buttons must have been preinflated using
+     * one of the methods in this class.
+     */
+    public void addPreInflatedButtons(List<Button> smartSuggestionButtons) {
+        for (Button button : smartSuggestionButtons) {
+            addView(button);
+        }
+        reallocateCandidateButtonQueueForSqueezing();
+    }
+
+    /**
      * Add smart replies to this view, using the provided {@link RemoteInput} and
      * {@link PendingIntent} to respond when the user taps a smart reply. Only the replies that fit
      * into the notification are shown.
      */
-    public void addRepliesFromRemoteInput(
-            SmartReplies smartReplies,
+    public List<Button> inflateRepliesFromRemoteInput(
+            @NonNull SmartReplies smartReplies,
             SmartReplyController smartReplyController, NotificationEntry entry) {
+        List<Button> buttons = new ArrayList<>();
+
         if (smartReplies.remoteInput != null && smartReplies.pendingIntent != null) {
             if (smartReplies.choices != null) {
                 for (int i = 0; i < smartReplies.choices.length; ++i) {
-                    Button replyButton = inflateReplyButton(
-                            getContext(), this, i, smartReplies, smartReplyController, entry);
-                    addView(replyButton);
+                    buttons.add(inflateReplyButton(
+                            this, getContext(), i, smartReplies, smartReplyController, entry));
                 }
                 this.mSmartRepliesGeneratedByAssistant = smartReplies.fromAssistant;
             }
         }
-        reallocateCandidateButtonQueueForSqueezing();
+        return buttons;
     }
 
     /**
      * Add smart actions to be shown next to smart replies. Only the actions that fit into the
      * notification are shown.
      */
-    public void addSmartActions(SmartActions smartActions,
+    public List<Button> inflateSmartActions(@NonNull SmartActions smartActions,
             SmartReplyController smartReplyController, NotificationEntry entry,
             HeadsUpManager headsUpManager) {
+        List<Button> buttons = new ArrayList<>();
         int numSmartActions = smartActions.actions.size();
         for (int n = 0; n < numSmartActions; n++) {
             Notification.Action action = smartActions.actions.get(n);
             if (action.actionIntent != null) {
-                Button actionButton = inflateActionButton(
-                        getContext(), this, n, smartActions, smartReplyController, entry,
-                        headsUpManager);
-                addView(actionButton);
+                buttons.add(inflateActionButton(
+                        this, getContext(), n, smartActions, smartReplyController, entry,
+                        headsUpManager));
             }
         }
-        reallocateCandidateButtonQueueForSqueezing();
+        return buttons;
     }
 
-    public static SmartReplyView inflate(Context context, ViewGroup root) {
-        return (SmartReplyView)
-                LayoutInflater.from(context).inflate(R.layout.smart_reply_view, root, false);
+    /**
+     * Inflate an instance of this class.
+     */
+    public static SmartReplyView inflate(Context context) {
+        return (SmartReplyView) LayoutInflater.from(context).inflate(
+                R.layout.smart_reply_view, null /* root */);
     }
 
     @VisibleForTesting
-    Button inflateReplyButton(Context context, ViewGroup root, int replyIndex,
-            SmartReplies smartReplies, SmartReplyController smartReplyController,
+    static Button inflateReplyButton(SmartReplyView smartReplyView, Context context,
+            int replyIndex, SmartReplies smartReplies, SmartReplyController smartReplyController,
             NotificationEntry entry) {
         Button b = (Button) LayoutInflater.from(context).inflate(
-                R.layout.smart_reply_button, root, false);
+                R.layout.smart_reply_button, smartReplyView, false);
         CharSequence choice = smartReplies.choices[replyIndex];
         b.setText(choice);
 
         OnDismissAction action = () -> {
-            if (mConstants.getEffectiveEditChoicesBeforeSending(
+            if (smartReplyView.mConstants.getEffectiveEditChoicesBeforeSending(
                     smartReplies.remoteInput.getEditChoicesBeforeSending())) {
                 EditedSuggestionInfo editedSuggestionInfo =
                         new EditedSuggestionInfo(choice, replyIndex);
-                mRemoteInputManager.activateRemoteInput(b,
+                smartReplyView.mRemoteInputManager.activateRemoteInput(b,
                         new RemoteInput[] { smartReplies.remoteInput }, smartReplies.remoteInput,
                         smartReplies.pendingIntent, editedSuggestionInfo);
                 return false;
@@ -276,33 +291,41 @@ public class SmartReplyView extends ViewGroup {
             } catch (PendingIntent.CanceledException e) {
                 Log.w(TAG, "Unable to send smart reply", e);
             }
-            mSmartReplyContainer.setVisibility(View.GONE);
+            // Note that as inflateReplyButton is called mSmartReplyContainer is null, but when the
+            // reply Button is added to the SmartReplyView mSmartReplyContainer will be set. So, it
+            // will not be possible for a user to trigger this on-click-listener without
+            // mSmartReplyContainer being set.
+            smartReplyView.mSmartReplyContainer.setVisibility(View.GONE);
             return false; // do not defer
         };
 
         b.setOnClickListener(view -> {
-            mKeyguardDismissUtil.executeWhenUnlocked(action);
+            smartReplyView.mKeyguardDismissUtil.executeWhenUnlocked(action);
         });
 
         b.setAccessibilityDelegate(new AccessibilityDelegate() {
             public void onInitializeAccessibilityNodeInfo(View host, AccessibilityNodeInfo info) {
                 super.onInitializeAccessibilityNodeInfo(host, info);
-                String label = getResources().getString(R.string.accessibility_send_smart_reply);
+                String label = smartReplyView.getResources().getString(
+                        R.string.accessibility_send_smart_reply);
                 info.addAction(new AccessibilityAction(AccessibilityNodeInfo.ACTION_CLICK, label));
             }
         });
 
-        setColors(b, mCurrentBackgroundColor, mDefaultStrokeColor, mDefaultTextColor, mRippleColor);
+        SmartReplyView.setButtonColors(b, smartReplyView.mCurrentBackgroundColor,
+                smartReplyView.mDefaultStrokeColor, smartReplyView.mDefaultTextColor,
+                smartReplyView.mRippleColor, smartReplyView.mStrokeWidth);
         return b;
     }
 
     @VisibleForTesting
-    Button inflateActionButton(Context context, ViewGroup root, int actionIndex,
-            SmartActions smartActions, SmartReplyController smartReplyController,
-            NotificationEntry entry, HeadsUpManager headsUpManager) {
+    static Button inflateActionButton(SmartReplyView smartReplyView, Context context,
+            int actionIndex, SmartActions smartActions,
+            SmartReplyController smartReplyController, NotificationEntry entry,
+            HeadsUpManager headsUpManager) {
         Notification.Action action = smartActions.actions.get(actionIndex);
         Button button = (Button) LayoutInflater.from(context).inflate(
-                R.layout.smart_action_button, root, false);
+                R.layout.smart_action_button, smartReplyView, false);
         button.setText(action.title);
 
         Drawable iconDrawable = action.getIcon().loadDrawable(context);
@@ -313,7 +336,7 @@ public class SmartReplyView extends ViewGroup {
         button.setCompoundDrawables(iconDrawable, null, null, null);
 
         button.setOnClickListener(view ->
-                getActivityStarter().startPendingIntentDismissingKeyguard(
+                smartReplyView.getActivityStarter().startPendingIntentDismissingKeyguard(
                         action.actionIntent,
                         () -> {
                             smartReplyController.smartActionClicked(
@@ -803,12 +826,13 @@ public class SmartReplyView extends ViewGroup {
         int childCount = getChildCount();
         for (int i = 0; i < childCount; i++) {
             final Button child = (Button) getChildAt(i);
-            setColors(child, backgroundColor, strokeColor, textColor, rippleColor);
+            setButtonColors(child, backgroundColor, strokeColor, textColor, rippleColor,
+                    mStrokeWidth);
         }
     }
 
-    private void setColors(Button button, int backgroundColor, int strokeColor, int textColor,
-            int rippleColor) {
+    private static void setButtonColors(Button button, int backgroundColor, int strokeColor,
+            int textColor, int rippleColor, int strokeWidth) {
         Drawable drawable = button.getBackground();
         if (drawable instanceof RippleDrawable) {
             // Mutate in case other notifications are using this drawable.
@@ -821,7 +845,7 @@ public class SmartReplyView extends ViewGroup {
                 if (background instanceof GradientDrawable) {
                     GradientDrawable gradientDrawable = (GradientDrawable) background;
                     gradientDrawable.setColor(backgroundColor);
-                    gradientDrawable.setStroke(mStrokeWidth, strokeColor);
+                    gradientDrawable.setStroke(strokeWidth, strokeColor);
                 }
             }
             button.setBackground(drawable);
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 6793ecaabdd7..60b0d61558e8 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
@@ -113,7 +113,7 @@ public class SmartReplyViewTest extends SysuiTestCase {
         mDependency.injectTestDependency(SmartReplyConstants.class, mConstants);
 
         mContainer = new View(mContext, null);
-        mView = SmartReplyView.inflate(mContext, null);
+        mView = SmartReplyView.inflate(mContext);
 
         // Any number of replies are fine.
         when(mConstants.getMinNumSystemGeneratedReplies()).thenReturn(0);
@@ -402,17 +402,18 @@ public class SmartReplyViewTest extends SysuiTestCase {
     }
 
     private void setSmartReplies(CharSequence[] choices) {
-        setSmartReplies(choices, false /* fromAssistant */);
+        mView.resetSmartSuggestions(mContainer);
+        List<Button> replyButtons = inflateSmartReplies(choices, false /* fromAssistant */);
+        mView.addPreInflatedButtons(replyButtons);
     }
 
-    private void setSmartReplies(CharSequence[] choices, boolean fromAssistant) {
+    private List<Button> inflateSmartReplies(CharSequence[] choices, boolean fromAssistant) {
         PendingIntent pendingIntent = PendingIntent.getBroadcast(mContext, 0,
                 new Intent(TEST_ACTION), 0);
         RemoteInput input = new RemoteInput.Builder(TEST_RESULT_KEY).setChoices(choices).build();
         SmartReplyView.SmartReplies smartReplies =
                 new SmartReplyView.SmartReplies(choices, input, pendingIntent, fromAssistant);
-        mView.resetSmartSuggestions(mContainer);
-        mView.addRepliesFromRemoteInput(smartReplies, mLogger, mEntry);
+        return mView.inflateRepliesFromRemoteInput(smartReplies, mLogger, mEntry);
     }
 
     private Notification.Action createAction(String actionTitle) {
@@ -431,11 +432,12 @@ public class SmartReplyViewTest extends SysuiTestCase {
 
     private void setSmartActions(String[] actionTitles) {
         mView.resetSmartSuggestions(mContainer);
-        mView.addSmartActions(
+        List<Button> actions = mView.inflateSmartActions(
                 new SmartReplyView.SmartActions(createActions(actionTitles), false),
                 mLogger,
                 mEntry,
                 mHeadsUpManager);
+        mView.addPreInflatedButtons(actions);
     }
 
     private void setSmartRepliesAndActions(CharSequence[] choices, String[] actionTitles) {
@@ -444,12 +446,14 @@ public class SmartReplyViewTest extends SysuiTestCase {
 
     private void setSmartRepliesAndActions(
             CharSequence[] choices, String[] actionTitles, boolean fromAssistant) {
-        setSmartReplies(choices, fromAssistant);
-        mView.addSmartActions(
+        mView.resetSmartSuggestions(mContainer);
+        List<Button> smartSuggestions = inflateSmartReplies(choices, fromAssistant);
+        smartSuggestions.addAll(mView.inflateSmartActions(
                 new SmartReplyView.SmartActions(createActions(actionTitles), fromAssistant),
                 mLogger,
                 mEntry,
-                mHeadsUpManager);
+                mHeadsUpManager));
+        mView.addPreInflatedButtons(smartSuggestions);
     }
 
     private ViewGroup buildExpectedView(CharSequence[] choices, int lineCount) {
@@ -485,8 +489,8 @@ public class SmartReplyViewTest extends SysuiTestCase {
         SmartReplyView.SmartReplies smartReplies =
                 new SmartReplyView.SmartReplies(choices, null, null, false);
         for (int i = 0; i < choices.length; ++i) {
-            Button current = mView.inflateReplyButton(mContext, mView, i, smartReplies,
-                    null, null);
+            Button current = SmartReplyView.inflateReplyButton(mView, mContext, i, smartReplies,
+                    null /* SmartReplyController */, null /* NotificationEntry */);
             current.setPadding(paddingHorizontal, current.getPaddingTop(), paddingHorizontal,
                     current.getPaddingBottom());
             if (previous != null) {
@@ -752,7 +756,7 @@ public class SmartReplyViewTest extends SysuiTestCase {
     }
 
     private Button inflateActionButton(Notification.Action action) {
-        return mView.inflateActionButton(getContext(), mView, 0,
+        return SmartReplyView.inflateActionButton(mView, getContext(), 0,
                 new SmartReplyView.SmartActions(Collections.singletonList(action), false),
                 mLogger, mEntry, mHeadsUpManager);
     }