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 b65c4a5f71d8a..b4dd1144e7619 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 dae14debd6839..646617c72128e 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 754bfb2b30aad..d8ea1f6eef5fa 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 45d215ef309c7..ed5487f743561 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/SmartReplyView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/SmartReplyView.java
@@ -195,67 +195,82 @@ public class SmartReplyView extends ViewGroup {
         mCurrentBackgroundColor = mDefaultBackgroundColor;
     }
 
+    /**
+     * 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 6793ecaabdd72..60b0d61558e81 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);
     }