Merge "Promote FeatureFlag best practices in notifications code & tests" into udc-qpr-dev am: febba3b65f am: ebee1d214a

Original change: https://googleplex-android-review.googlesource.com/c/platform/frameworks/base/+/23984923

Change-Id: I88980b1b7fc1d21289c176f9ba2c666a50b1c110
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/packages/SystemUI/src/com/android/systemui/flags/ViewRefactorFlag.kt b/packages/SystemUI/src/com/android/systemui/flags/ViewRefactorFlag.kt
new file mode 100644
index 0000000..eaecda5
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/flags/ViewRefactorFlag.kt
@@ -0,0 +1,99 @@
+/*
+ * Copyright (C) 2023 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.flags
+
+import android.util.Log
+import com.android.systemui.Dependency
+
+/**
+ * This class promotes best practices for flag guarding System UI view refactors.
+ * * [isEnabled] allows changing an implementation.
+ * * [assertDisabled] allows authors to flag code as being "dead" when the flag gets enabled and
+ *   ensure that it is not being invoked accidentally in the post-flag refactor.
+ * * [expectEnabled] allows authors to guard new code with a "safe" alternative when invoked on
+ *   flag-disabled builds, but with a check that should crash eng builds or tests when the
+ *   expectation is violated.
+ *
+ * The constructors prefer that you provide a [FeatureFlags] instance, but does not require it,
+ * falling back to [Dependency.get]. This fallback should ONLY be used to flag-guard code changes
+ * inside views where injecting flag values after initialization can be error-prone.
+ */
+class ViewRefactorFlag
+private constructor(
+    private val injectedFlags: FeatureFlags?,
+    private val flag: BooleanFlag,
+    private val readFlagValue: (FeatureFlags) -> Boolean
+) {
+    @JvmOverloads
+    constructor(
+        flags: FeatureFlags? = null,
+        flag: UnreleasedFlag
+    ) : this(flags, flag, { it.isEnabled(flag) })
+
+    @JvmOverloads
+    constructor(
+        flags: FeatureFlags? = null,
+        flag: ReleasedFlag
+    ) : this(flags, flag, { it.isEnabled(flag) })
+
+    /** Whether the flag is enabled. Called to switch between an old behavior and a new behavior. */
+    val isEnabled by lazy {
+        @Suppress("DEPRECATION")
+        val featureFlags = injectedFlags ?: Dependency.get(FeatureFlags::class.java)
+        readFlagValue(featureFlags)
+    }
+
+    /**
+     * Called to ensure code is only run when the flag is disabled. This will throw an exception if
+     * the flag is enabled to ensure that the refactor author catches issues in testing.
+     *
+     * Example usage:
+     * ```
+     * public void setController(NotificationShelfController notificationShelfController) {
+     *     mShelfRefactor.assertDisabled();
+     *     mController = notificationShelfController;
+     * }
+     * ````
+     */
+    fun assertDisabled() = check(!isEnabled) { "Code path not supported when $flag is enabled." }
+
+    /**
+     * Called to ensure code is only run when the flag is enabled. This protects users from the
+     * unintended behaviors caused by accidentally running new logic, while also crashing on an eng
+     * build to ensure that the refactor author catches issues in testing.
+     *
+     * Example usage:
+     * ```
+     * public void setShelfIcons(NotificationIconContainer icons) {
+     *     if (mShelfRefactor.expectEnabled()) {
+     *         mShelfIcons = icons;
+     *     }
+     * }
+     * ```
+     */
+    fun expectEnabled(): Boolean {
+        if (!isEnabled) {
+            val message = "Code path not supported when $flag is disabled."
+            Log.wtf(TAG, message, Exception(message))
+        }
+        return isEnabled
+    }
+
+    private companion object {
+        private const val TAG = "ViewRefactorFlag"
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/LegacyNotificationShelfControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/LegacyNotificationShelfControllerImpl.java
index 4ec5f46..7a989cf 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/LegacyNotificationShelfControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/LegacyNotificationShelfControllerImpl.java
@@ -19,7 +19,6 @@
 import android.view.View;
 
 import com.android.systemui.flags.FeatureFlags;
-import com.android.systemui.flags.Flags;
 import com.android.systemui.statusbar.notification.row.ActivatableNotificationViewController;
 import com.android.systemui.statusbar.notification.row.dagger.NotificationRowScope;
 import com.android.systemui.statusbar.notification.stack.AmbientState;
@@ -52,7 +51,6 @@
         mActivatableNotificationViewController = activatableNotificationViewController;
         mKeyguardBypassController = keyguardBypassController;
         mStatusBarStateController = statusBarStateController;
-        mView.setSensitiveRevealAnimEnabled(featureFlags.isEnabled(Flags.SENSITIVE_REVEAL_ANIM));
         mOnAttachStateChangeListener = new View.OnAttachStateChangeListener() {
             @Override
             public void onViewAttachedToWindow(View v) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java
index 25a1dc6..3f37c60 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java
@@ -24,7 +24,6 @@
 import android.graphics.Rect;
 import android.util.AttributeSet;
 import android.util.IndentingPrintWriter;
-import android.util.Log;
 import android.util.MathUtils;
 import android.view.View;
 import android.view.ViewGroup;
@@ -40,6 +39,8 @@
 import com.android.internal.policy.SystemBarUtils;
 import com.android.systemui.R;
 import com.android.systemui.animation.ShadeInterpolation;
+import com.android.systemui.flags.Flags;
+import com.android.systemui.flags.ViewRefactorFlag;
 import com.android.systemui.plugins.statusbar.StatusBarStateController.StateListener;
 import com.android.systemui.shade.transition.LargeScreenShadeInterpolator;
 import com.android.systemui.statusbar.notification.NotificationUtils;
@@ -95,8 +96,10 @@
     private float mCornerAnimationDistance;
     private NotificationShelfController mController;
     private float mActualWidth = -1;
-    private boolean mSensitiveRevealAnimEnabled;
-    private boolean mShelfRefactorFlagEnabled;
+    private final ViewRefactorFlag mSensitiveRevealAnim =
+            new ViewRefactorFlag(Flags.SENSITIVE_REVEAL_ANIM);
+    private final ViewRefactorFlag mShelfRefactor =
+            new ViewRefactorFlag(Flags.NOTIFICATION_SHELF_REFACTOR);
     private boolean mCanModifyColorOfNotifications;
     private boolean mCanInteract;
     private NotificationStackScrollLayout mHostLayout;
@@ -130,7 +133,7 @@
 
     public void bind(AmbientState ambientState,
                      NotificationStackScrollLayoutController hostLayoutController) {
-        assertRefactorFlagDisabled();
+        mShelfRefactor.assertDisabled();
         mAmbientState = ambientState;
         mHostLayoutController = hostLayoutController;
         hostLayoutController.setOnNotificationRemovedListener((child, isTransferInProgress) -> {
@@ -140,7 +143,7 @@
 
     public void bind(AmbientState ambientState, NotificationStackScrollLayout hostLayout,
             NotificationRoundnessManager roundnessManager) {
-        if (!checkRefactorFlagEnabled()) return;
+        if (!mShelfRefactor.expectEnabled()) return;
         mAmbientState = ambientState;
         mHostLayout = hostLayout;
         mRoundnessManager = roundnessManager;
@@ -268,7 +271,7 @@
         }
 
         final float stackEnd = ambientState.getStackY() + ambientState.getStackHeight();
-        if (mSensitiveRevealAnimEnabled && viewState.hidden) {
+        if (mSensitiveRevealAnim.isEnabled() && viewState.hidden) {
             // if the shelf is hidden, position it at the end of the stack (plus the clip
             // padding), such that when it appears animated, it will smoothly move in from the
             // bottom, without jump cutting any notifications
@@ -279,7 +282,7 @@
     }
 
     private int getSpeedBumpIndex() {
-        if (mShelfRefactorFlagEnabled) {
+        if (mShelfRefactor.isEnabled()) {
             return mHostLayout.getSpeedBumpIndex();
         } else {
             return mHostLayoutController.getSpeedBumpIndex();
@@ -413,7 +416,7 @@
                     expandingAnimated, isLastChild, shelfClipStart);
 
             // TODO(b/172289889) scale mPaddingBetweenElements with expansion amount
-            if ((!mSensitiveRevealAnimEnabled && ((isLastChild && !child.isInShelf())
+            if ((!mSensitiveRevealAnim.isEnabled() && ((isLastChild && !child.isInShelf())
                     || backgroundForceHidden)) || aboveShelf) {
                 notificationClipEnd = shelfStart + getIntrinsicHeight();
             } else {
@@ -462,7 +465,7 @@
                 // if the shelf is visible, but if the shelf is hidden, it causes incorrect curling.
                 // notificationClipEnd handles the discrepancy between a visible and hidden shelf,
                 // so we use that when on the keyguard (and while animating away) to reduce curling.
-                final float keyguardSafeShelfStart = !mSensitiveRevealAnimEnabled
+                final float keyguardSafeShelfStart = !mSensitiveRevealAnim.isEnabled()
                         && mAmbientState.isOnKeyguard() ? notificationClipEnd : shelfStart;
                 updateCornerRoundnessOnScroll(anv, viewStart, keyguardSafeShelfStart);
             }
@@ -504,7 +507,7 @@
     }
 
     private ExpandableView getHostLayoutChildAt(int index) {
-        if (mShelfRefactorFlagEnabled) {
+        if (mShelfRefactor.isEnabled()) {
             return (ExpandableView) mHostLayout.getChildAt(index);
         } else {
             return mHostLayoutController.getChildAt(index);
@@ -512,7 +515,7 @@
     }
 
     private int getHostLayoutChildCount() {
-        if (mShelfRefactorFlagEnabled) {
+        if (mShelfRefactor.isEnabled()) {
             return mHostLayout.getChildCount();
         } else {
             return mHostLayoutController.getChildCount();
@@ -520,7 +523,7 @@
     }
 
     private boolean canModifyColorOfNotifications() {
-        if (mShelfRefactorFlagEnabled) {
+        if (mShelfRefactor.isEnabled()) {
             return mCanModifyColorOfNotifications && mAmbientState.isShadeExpanded();
         } else {
             return mController.canModifyColorOfNotifications();
@@ -583,7 +586,7 @@
     }
 
     private boolean isViewAffectedBySwipe(ExpandableView expandableView) {
-        if (!mShelfRefactorFlagEnabled) {
+        if (!mShelfRefactor.isEnabled()) {
             return mHostLayoutController.isViewAffectedBySwipe(expandableView);
         } else {
             return mRoundnessManager.isViewAffectedBySwipe(expandableView);
@@ -607,7 +610,7 @@
     }
 
     private View getHostLayoutTransientView(int index) {
-        if (mShelfRefactorFlagEnabled) {
+        if (mShelfRefactor.isEnabled()) {
             return mHostLayout.getTransientView(index);
         } else {
             return mHostLayoutController.getTransientView(index);
@@ -615,7 +618,7 @@
     }
 
     private int getHostLayoutTransientViewCount() {
-        if (mShelfRefactorFlagEnabled) {
+        if (mShelfRefactor.isEnabled()) {
             return mHostLayout.getTransientViewCount();
         } else {
             return mHostLayoutController.getTransientViewCount();
@@ -961,7 +964,7 @@
 
     @Override
     public void onStateChanged(int newState) {
-        assertRefactorFlagDisabled();
+        mShelfRefactor.assertDisabled();
         mStatusBarState = newState;
         updateInteractiveness();
     }
@@ -975,7 +978,7 @@
     }
 
     private boolean canInteract() {
-        if (mShelfRefactorFlagEnabled) {
+        if (mShelfRefactor.isEnabled()) {
             return mCanInteract;
         } else {
             return mStatusBarState == StatusBarState.KEYGUARD;
@@ -1018,32 +1021,18 @@
         return false;
     }
 
-    private void assertRefactorFlagDisabled() {
-        if (mShelfRefactorFlagEnabled) {
-            NotificationShelfController.throwIllegalFlagStateError(false);
-        }
-    }
-
-    private boolean checkRefactorFlagEnabled() {
-        if (!mShelfRefactorFlagEnabled) {
-            Log.wtf(TAG,
-                    "Code path not supported when Flags.NOTIFICATION_SHELF_REFACTOR is disabled.");
-        }
-        return mShelfRefactorFlagEnabled;
-    }
-
     public void setController(NotificationShelfController notificationShelfController) {
-        assertRefactorFlagDisabled();
+        mShelfRefactor.assertDisabled();
         mController = notificationShelfController;
     }
 
     public void setCanModifyColorOfNotifications(boolean canModifyColorOfNotifications) {
-        if (!checkRefactorFlagEnabled()) return;
+        if (!mShelfRefactor.expectEnabled()) return;
         mCanModifyColorOfNotifications = canModifyColorOfNotifications;
     }
 
     public void setCanInteract(boolean canInteract) {
-        if (!checkRefactorFlagEnabled()) return;
+        if (!mShelfRefactor.expectEnabled()) return;
         mCanInteract = canInteract;
         updateInteractiveness();
     }
@@ -1053,27 +1042,15 @@
     }
 
     private int getIndexOfViewInHostLayout(ExpandableView child) {
-        if (mShelfRefactorFlagEnabled) {
+        if (mShelfRefactor.isEnabled()) {
             return mHostLayout.indexOfChild(child);
         } else {
             return mHostLayoutController.indexOfChild(child);
         }
     }
 
-    /**
-     * Set whether the sensitive reveal animation feature flag is enabled
-     * @param enabled true if enabled
-     */
-    public void setSensitiveRevealAnimEnabled(boolean enabled) {
-        mSensitiveRevealAnimEnabled = enabled;
-    }
-
-    public void setRefactorFlagEnabled(boolean enabled) {
-        mShelfRefactorFlagEnabled = enabled;
-    }
-
     public void requestRoundnessResetFor(ExpandableView child) {
-        if (!checkRefactorFlagEnabled()) return;
+        if (!mShelfRefactor.expectEnabled()) return;
         child.requestRoundnessReset(SHELF_SCROLL);
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelfController.kt b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelfController.kt
index 1619dda..8a3e217 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelfController.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelfController.kt
@@ -16,11 +16,8 @@
 
 package com.android.systemui.statusbar
 
-import android.util.Log
 import android.view.View
 import android.view.View.OnClickListener
-import com.android.systemui.flags.FeatureFlags
-import com.android.systemui.flags.Flags
 import com.android.systemui.statusbar.notification.row.ExpandableView
 import com.android.systemui.statusbar.notification.stack.AmbientState
 import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayout
@@ -49,29 +46,4 @@
 
     /** @see View.setOnClickListener */
     fun setOnClickListener(listener: OnClickListener)
-
-    companion object {
-        @JvmStatic
-        fun assertRefactorFlagDisabled(featureFlags: FeatureFlags) {
-            if (featureFlags.isEnabled(Flags.NOTIFICATION_SHELF_REFACTOR)) {
-                throwIllegalFlagStateError(expected = false)
-            }
-        }
-
-        @JvmStatic
-        fun checkRefactorFlagEnabled(featureFlags: FeatureFlags): Boolean =
-            featureFlags.isEnabled(Flags.NOTIFICATION_SHELF_REFACTOR).also { enabled ->
-                if (!enabled) {
-                    Log.wtf("NotifShelf", getErrorMessage(expected = true))
-                }
-            }
-
-        @JvmStatic
-        fun throwIllegalFlagStateError(expected: Boolean): Nothing =
-            error(getErrorMessage(expected))
-
-        private fun getErrorMessage(expected: Boolean): String =
-            "Code path not supported when Flags.NOTIFICATION_SHELF_REFACTOR is " +
-                if (expected) "disabled" else "enabled"
-    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/Roundable.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/Roundable.kt
index 1cf9c1e..1c5aa3c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/Roundable.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/Roundable.kt
@@ -3,10 +3,10 @@
 import android.util.FloatProperty
 import android.view.View
 import androidx.annotation.FloatRange
-import com.android.systemui.Dependency
 import com.android.systemui.R
 import com.android.systemui.flags.FeatureFlags
 import com.android.systemui.flags.Flags
+import com.android.systemui.flags.ViewRefactorFlag
 import com.android.systemui.statusbar.notification.stack.AnimationProperties
 import com.android.systemui.statusbar.notification.stack.StackStateAnimator
 import kotlin.math.abs
@@ -46,14 +46,14 @@
     @JvmDefault
     val topCornerRadius: Float
         get() =
-            if (roundableState.newHeadsUpAnimFlagEnabled) roundableState.topCornerRadius
+            if (roundableState.newHeadsUpAnim.isEnabled) roundableState.topCornerRadius
             else topRoundness * maxRadius
 
     /** Current bottom corner in pixel, based on [bottomRoundness] and [maxRadius] */
     @JvmDefault
     val bottomCornerRadius: Float
         get() =
-            if (roundableState.newHeadsUpAnimFlagEnabled) roundableState.bottomCornerRadius
+            if (roundableState.newHeadsUpAnim.isEnabled) roundableState.bottomCornerRadius
             else bottomRoundness * maxRadius
 
     /** Get and update the current radii */
@@ -335,13 +335,12 @@
     internal val targetView: View,
     private val roundable: Roundable,
     maxRadius: Float,
-    private val featureFlags: FeatureFlags = Dependency.get(FeatureFlags::class.java)
+    featureFlags: FeatureFlags? = null
 ) {
     internal var maxRadius = maxRadius
         private set
 
-    internal val newHeadsUpAnimFlagEnabled
-        get() = featureFlags.isEnabled(Flags.IMPROVED_HUN_ANIMATIONS)
+    internal val newHeadsUpAnim = ViewRefactorFlag(featureFlags, Flags.IMPROVED_HUN_ANIMATIONS)
 
     /** Animatable for top roundness */
     private val topAnimatable = topAnimatable(roundable)
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/inflation/NotificationRowBinderImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/inflation/NotificationRowBinderImpl.java
index 1896080..9ecf50e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/inflation/NotificationRowBinderImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/inflation/NotificationRowBinderImpl.java
@@ -16,7 +16,6 @@
 
 package com.android.systemui.statusbar.notification.collection.inflation;
 
-import static com.android.systemui.flags.Flags.NOTIFICATION_INLINE_REPLY_ANIMATION;
 import static com.android.systemui.statusbar.notification.row.NotificationRowContentBinder.FLAG_CONTENT_VIEW_CONTRACTED;
 import static com.android.systemui.statusbar.notification.row.NotificationRowContentBinder.FLAG_CONTENT_VIEW_EXPANDED;
 import static com.android.systemui.statusbar.notification.row.NotificationRowContentBinder.FLAG_CONTENT_VIEW_PUBLIC;
@@ -176,8 +175,6 @@
         entry.setRow(row);
         mNotifBindPipeline.manageRow(entry, row);
         mPresenter.onBindRow(row);
-        row.setInlineReplyAnimationFlagEnabled(
-                mFeatureFlags.isEnabled(NOTIFICATION_INLINE_REPLY_ANIMATION));
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ActivatableNotificationView.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ActivatableNotificationView.java
index 908c11a..36a8e98 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ActivatableNotificationView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ActivatableNotificationView.java
@@ -566,7 +566,7 @@
 
     @Override
     public float getTopCornerRadius() {
-        if (isNewHeadsUpAnimFlagEnabled()) {
+        if (mImprovedHunAnimation.isEnabled()) {
             return super.getTopCornerRadius();
         }
 
@@ -576,7 +576,7 @@
 
     @Override
     public float getBottomCornerRadius() {
-        if (isNewHeadsUpAnimFlagEnabled()) {
+        if (mImprovedHunAnimation.isEnabled()) {
             return super.getBottomCornerRadius();
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
index b34c281..42b99a1 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
@@ -77,6 +77,7 @@
 import com.android.systemui.classifier.FalsingCollector;
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.flags.Flags;
+import com.android.systemui.flags.ViewRefactorFlag;
 import com.android.systemui.plugins.FalsingManager;
 import com.android.systemui.plugins.PluginListener;
 import com.android.systemui.plugins.statusbar.NotificationMenuRowPlugin;
@@ -275,7 +276,8 @@
     private OnExpandClickListener mOnExpandClickListener;
     private View.OnClickListener mOnFeedbackClickListener;
     private Path mExpandingClipPath;
-    private boolean mIsInlineReplyAnimationFlagEnabled = false;
+    private final ViewRefactorFlag mInlineReplyAnimation =
+            new ViewRefactorFlag(Flags.NOTIFICATION_INLINE_REPLY_ANIMATION);
 
     // Listener will be called when receiving a long click event.
     // Use #setLongPressPosition to optionally assign positional data with the long press.
@@ -3121,10 +3123,6 @@
         return showingLayout != null && showingLayout.requireRowToHaveOverlappingRendering();
     }
 
-    public void setInlineReplyAnimationFlagEnabled(boolean isEnabled) {
-        mIsInlineReplyAnimationFlagEnabled = isEnabled;
-    }
-
     @Override
     public void setActualHeight(int height, boolean notifyListeners) {
         boolean changed = height != getActualHeight();
@@ -3144,7 +3142,7 @@
         }
         int contentHeight = Math.max(getMinHeight(), height);
         for (NotificationContentView l : mLayouts) {
-            if (mIsInlineReplyAnimationFlagEnabled) {
+            if (mInlineReplyAnimation.isEnabled()) {
                 l.setContentHeight(height);
             } else {
                 l.setContentHeight(contentHeight);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableOutlineView.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableOutlineView.java
index 7f23c1b..c8f13a6 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableOutlineView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableOutlineView.java
@@ -28,10 +28,9 @@
 import android.view.View;
 import android.view.ViewOutlineProvider;
 
-import com.android.systemui.Dependency;
 import com.android.systemui.R;
-import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.flags.Flags;
+import com.android.systemui.flags.ViewRefactorFlag;
 import com.android.systemui.statusbar.notification.RoundableState;
 import com.android.systemui.statusbar.notification.stack.NotificationChildrenContainer;
 import com.android.systemui.util.DumpUtilsKt;
@@ -50,7 +49,8 @@
     private float mOutlineAlpha = -1f;
     private boolean mAlwaysRoundBothCorners;
     private Path mTmpPath = new Path();
-    private final FeatureFlags mFeatureFlags;
+    protected final ViewRefactorFlag mImprovedHunAnimation =
+            new ViewRefactorFlag(Flags.IMPROVED_HUN_ANIMATIONS);
 
     /**
      * {@code false} if the children views of the {@link ExpandableOutlineView} are translated when
@@ -126,7 +126,7 @@
             return EMPTY_PATH;
         }
         float bottomRadius = mAlwaysRoundBothCorners ? getMaxRadius() : getBottomCornerRadius();
-        if (!isNewHeadsUpAnimFlagEnabled() && (topRadius + bottomRadius > height)) {
+        if (!mImprovedHunAnimation.isEnabled() && (topRadius + bottomRadius > height)) {
             float overShoot = topRadius + bottomRadius - height;
             float currentTopRoundness = getTopRoundness();
             float currentBottomRoundness = getBottomRoundness();
@@ -167,7 +167,6 @@
         super(context, attrs);
         setOutlineProvider(mProvider);
         initDimens();
-        mFeatureFlags = Dependency.get(FeatureFlags.class);
     }
 
     @Override
@@ -376,8 +375,4 @@
         });
     }
 
-    // TODO(b/290365128) replace with ViewRefactorFlag
-    protected boolean isNewHeadsUpAnimFlagEnabled() {
-        return mFeatureFlags.isEnabled(Flags.IMPROVED_HUN_ANIMATIONS);
-    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/shelf/ui/viewbinder/NotificationShelfViewBinder.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/shelf/ui/viewbinder/NotificationShelfViewBinder.kt
index 23a58d2..22a87a7 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/shelf/ui/viewbinder/NotificationShelfViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/shelf/ui/viewbinder/NotificationShelfViewBinder.kt
@@ -21,7 +21,6 @@
 import androidx.lifecycle.repeatOnLifecycle
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.flags.FeatureFlags
-import com.android.systemui.flags.Flags
 import com.android.systemui.lifecycle.repeatWhenAttached
 import com.android.systemui.plugins.FalsingManager
 import com.android.systemui.statusbar.LegacyNotificationShelfControllerImpl
@@ -66,7 +65,7 @@
     override fun setOnClickListener(listener: View.OnClickListener) = unsupported
 
     private val unsupported: Nothing
-        get() = NotificationShelfController.throwIllegalFlagStateError(expected = true)
+        get() = error("Code path not supported when NOTIFICATION_SHELF_REFACTOR is disabled")
 }
 
 /** Binds a [NotificationShelf] to its [view model][NotificationShelfViewModel]. */
@@ -80,8 +79,6 @@
     ) {
         ActivatableNotificationViewBinder.bind(viewModel, shelf, falsingManager)
         shelf.apply {
-            setRefactorFlagEnabled(true)
-            setSensitiveRevealAnimEnabled(featureFlags.isEnabled(Flags.SENSITIVE_REVEAL_ANIM))
             // TODO(278765923): Replace with eventual NotificationIconContainerViewBinder#bind()
             notificationIconAreaController.setShelfIcons(shelfIcons)
             repeatWhenAttached {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/AmbientState.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/AmbientState.java
index b0f3f59..95e74f2 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/AmbientState.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/AmbientState.java
@@ -29,7 +29,6 @@
 import com.android.systemui.R;
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.dump.DumpManager;
-import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.shade.transition.LargeScreenShadeInterpolator;
 import com.android.systemui.statusbar.NotificationShelf;
 import com.android.systemui.statusbar.StatusBarState;
@@ -57,7 +56,6 @@
     private final SectionProvider mSectionProvider;
     private final BypassController mBypassController;
     private final LargeScreenShadeInterpolator mLargeScreenShadeInterpolator;
-    private final FeatureFlags mFeatureFlags;
     /**
      *  Used to read bouncer states.
      */
@@ -261,13 +259,12 @@
             @NonNull SectionProvider sectionProvider,
             @NonNull BypassController bypassController,
             @Nullable StatusBarKeyguardViewManager statusBarKeyguardViewManager,
-            @NonNull LargeScreenShadeInterpolator largeScreenShadeInterpolator,
-            @NonNull FeatureFlags featureFlags) {
+            @NonNull LargeScreenShadeInterpolator largeScreenShadeInterpolator
+    ) {
         mSectionProvider = sectionProvider;
         mBypassController = bypassController;
         mStatusBarKeyguardViewManager = statusBarKeyguardViewManager;
         mLargeScreenShadeInterpolator = largeScreenShadeInterpolator;
-        mFeatureFlags = featureFlags;
         reload(context);
         dumpManager.registerDumpable(this);
     }
@@ -753,10 +750,6 @@
         return mLargeScreenShadeInterpolator;
     }
 
-    public FeatureFlags getFeatureFlags() {
-        return mFeatureFlags;
-    }
-
     @Override
     public void dump(PrintWriter pw, String[] args) {
         pw.println("mTopPadding=" + mTopPadding);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
index c1ceb3c..d71bc2f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
@@ -89,6 +89,7 @@
 import com.android.systemui.R;
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.flags.Flags;
+import com.android.systemui.flags.ViewRefactorFlag;
 import com.android.systemui.plugins.ActivityStarter;
 import com.android.systemui.plugins.statusbar.NotificationSwipeActionHelper;
 import com.android.systemui.shade.ShadeController;
@@ -198,7 +199,8 @@
     private Set<Integer> mDebugTextUsedYPositions;
     private final boolean mDebugRemoveAnimation;
     private final boolean mSensitiveRevealAnimEndabled;
-    private boolean mAnimatedInsets;
+    private final ViewRefactorFlag mAnimatedInsets;
+    private final ViewRefactorFlag mShelfRefactor;
 
     private int mContentHeight;
     private float mIntrinsicContentHeight;
@@ -621,7 +623,9 @@
         mDebugLines = featureFlags.isEnabled(Flags.NSSL_DEBUG_LINES);
         mDebugRemoveAnimation = featureFlags.isEnabled(Flags.NSSL_DEBUG_REMOVE_ANIMATION);
         mSensitiveRevealAnimEndabled = featureFlags.isEnabled(Flags.SENSITIVE_REVEAL_ANIM);
-        setAnimatedInsetsEnabled(featureFlags.isEnabled(Flags.ANIMATED_NOTIFICATION_SHADE_INSETS));
+        mAnimatedInsets =
+                new ViewRefactorFlag(featureFlags, Flags.ANIMATED_NOTIFICATION_SHADE_INSETS);
+        mShelfRefactor = new ViewRefactorFlag(featureFlags, Flags.NOTIFICATION_SHELF_REFACTOR);
         mSectionsManager = Dependency.get(NotificationSectionsManager.class);
         mScreenOffAnimationController =
                 Dependency.get(ScreenOffAnimationController.class);
@@ -660,7 +664,7 @@
         mGroupMembershipManager = Dependency.get(GroupMembershipManager.class);
         mGroupExpansionManager = Dependency.get(GroupExpansionManager.class);
         setImportantForAccessibility(IMPORTANT_FOR_ACCESSIBILITY_YES);
-        if (mAnimatedInsets) {
+        if (mAnimatedInsets.isEnabled()) {
             setWindowInsetsAnimationCallback(mInsetsCallback);
         }
     }
@@ -730,11 +734,6 @@
     }
 
     @VisibleForTesting
-    void setAnimatedInsetsEnabled(boolean enabled) {
-        mAnimatedInsets = enabled;
-    }
-
-    @VisibleForTesting
     public void updateFooter() {
         if (mFooterView == null) {
             return;
@@ -1773,7 +1772,7 @@
             return;
         }
         mForcedScroll = v;
-        if (mAnimatedInsets) {
+        if (mAnimatedInsets.isEnabled()) {
             updateForcedScroll();
         } else {
             scrollTo(v);
@@ -1822,7 +1821,7 @@
 
     @Override
     public WindowInsets onApplyWindowInsets(WindowInsets insets) {
-        if (!mAnimatedInsets) {
+        if (!mAnimatedInsets.isEnabled()) {
             mBottomInset = insets.getInsets(WindowInsets.Type.ime()).bottom;
         }
         mWaterfallTopInset = 0;
@@ -1830,11 +1829,11 @@
         if (cutout != null) {
             mWaterfallTopInset = cutout.getWaterfallInsets().top;
         }
-        if (mAnimatedInsets && !mIsInsetAnimationRunning) {
+        if (mAnimatedInsets.isEnabled() && !mIsInsetAnimationRunning) {
             // update bottom inset e.g. after rotation
             updateBottomInset(insets);
         }
-        if (!mAnimatedInsets) {
+        if (!mAnimatedInsets.isEnabled()) {
             int range = getScrollRange();
             if (mOwnScrollY > range) {
                 // HACK: We're repeatedly getting staggered insets here while the IME is
@@ -2714,7 +2713,7 @@
      * @param listener callback for notification removed
      */
     public void setOnNotificationRemovedListener(OnNotificationRemovedListener listener) {
-        NotificationShelfController.assertRefactorFlagDisabled(mAmbientState.getFeatureFlags());
+        mShelfRefactor.assertDisabled();
         mOnNotificationRemovedListener = listener;
     }
 
@@ -2727,7 +2726,7 @@
         if (!mChildTransferInProgress) {
             onViewRemovedInternal(expandableView, this);
         }
-        if (mAmbientState.getFeatureFlags().isEnabled(Flags.NOTIFICATION_SHELF_REFACTOR)) {
+        if (mShelfRefactor.isEnabled()) {
             mShelf.requestRoundnessResetFor(expandableView);
         } else {
             if (mOnNotificationRemovedListener != null) {
@@ -4943,18 +4942,12 @@
 
     @Nullable
     public ExpandableView getShelf() {
-        if (NotificationShelfController.checkRefactorFlagEnabled(mAmbientState.getFeatureFlags())) {
-            return mShelf;
-        } else {
-            return null;
-        }
+        if (!mShelfRefactor.expectEnabled()) return null;
+        return mShelf;
     }
 
     public void setShelf(NotificationShelf shelf) {
-        if (!NotificationShelfController.checkRefactorFlagEnabled(
-                mAmbientState.getFeatureFlags())) {
-            return;
-        }
+        if (!mShelfRefactor.expectEnabled()) return;
         int index = -1;
         if (mShelf != null) {
             index = indexOfChild(mShelf);
@@ -4968,7 +4961,7 @@
     }
 
     public void setShelfController(NotificationShelfController notificationShelfController) {
-        NotificationShelfController.assertRefactorFlagDisabled(mAmbientState.getFeatureFlags());
+        mShelfRefactor.assertDisabled();
         int index = -1;
         if (mShelf != null) {
             index = indexOfChild(mShelf);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
index ef7375a..4668aa4 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
@@ -65,6 +65,7 @@
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.flags.Flags;
+import com.android.systemui.flags.ViewRefactorFlag;
 import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository;
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor;
 import com.android.systemui.keyguard.shared.model.KeyguardState;
@@ -205,6 +206,7 @@
     private boolean mIsInTransitionToAod = false;
 
     private final FeatureFlags mFeatureFlags;
+    private final ViewRefactorFlag mShelfRefactor;
     private final NotificationTargetsHelper mNotificationTargetsHelper;
     private final SecureSettings mSecureSettings;
     private final NotificationDismissibilityProvider mDismissibilityProvider;
@@ -718,6 +720,7 @@
         mShadeController = shadeController;
         mNotifIconAreaController = notifIconAreaController;
         mFeatureFlags = featureFlags;
+        mShelfRefactor = new ViewRefactorFlag(featureFlags, Flags.NOTIFICATION_SHELF_REFACTOR);
         mNotificationTargetsHelper = notificationTargetsHelper;
         mSecureSettings = secureSettings;
         mDismissibilityProvider = dismissibilityProvider;
@@ -1432,7 +1435,7 @@
     }
 
     public void setShelfController(NotificationShelfController notificationShelfController) {
-        NotificationShelfController.assertRefactorFlagDisabled(mFeatureFlags);
+        mShelfRefactor.assertDisabled();
         mView.setShelfController(notificationShelfController);
     }
 
@@ -1645,12 +1648,12 @@
     }
 
     public void setShelf(NotificationShelf shelf) {
-        if (!NotificationShelfController.checkRefactorFlagEnabled(mFeatureFlags)) return;
+        if (!mShelfRefactor.expectEnabled()) return;
         mView.setShelf(shelf);
     }
 
     public int getShelfHeight() {
-        if (!NotificationShelfController.checkRefactorFlagEnabled(mFeatureFlags)) {
+        if (!mShelfRefactor.expectEnabled()) {
             return 0;
         }
         ExpandableView shelf = mView.getShelf();
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationIconAreaController.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationIconAreaController.java
index e18c9d8..0bf0f4b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationIconAreaController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationIconAreaController.java
@@ -24,6 +24,8 @@
 import com.android.systemui.demomode.DemoMode;
 import com.android.systemui.demomode.DemoModeController;
 import com.android.systemui.flags.FeatureFlags;
+import com.android.systemui.flags.Flags;
+import com.android.systemui.flags.ViewRefactorFlag;
 import com.android.systemui.plugins.DarkIconDispatcher;
 import com.android.systemui.plugins.DarkIconDispatcher.DarkReceiver;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
@@ -88,7 +90,7 @@
     private final ArrayList<Rect> mTintAreas = new ArrayList<>();
     private final Context mContext;
 
-    private final FeatureFlags mFeatureFlags;
+    private final ViewRefactorFlag mShelfRefactor;
 
     private int mAodIconAppearTranslation;
 
@@ -120,12 +122,13 @@
             Optional<Bubbles> bubblesOptional,
             DemoModeController demoModeController,
             DarkIconDispatcher darkIconDispatcher,
-            FeatureFlags featureFlags, StatusBarWindowController statusBarWindowController,
+            FeatureFlags featureFlags,
+            StatusBarWindowController statusBarWindowController,
             ScreenOffAnimationController screenOffAnimationController) {
         mContrastColorUtil = ContrastColorUtil.getInstance(context);
         mContext = context;
         mStatusBarStateController = statusBarStateController;
-        mFeatureFlags = featureFlags;
+        mShelfRefactor = new ViewRefactorFlag(featureFlags, Flags.NOTIFICATION_SHELF_REFACTOR);
         mStatusBarStateController.addCallback(this);
         mMediaManager = notificationMediaManager;
         mDozeParameters = dozeParameters;
@@ -179,12 +182,12 @@
     }
 
     public void setupShelf(NotificationShelfController notificationShelfController) {
-        NotificationShelfController.assertRefactorFlagDisabled(mFeatureFlags);
+        mShelfRefactor.assertDisabled();
         mShelfIcons = notificationShelfController.getShelfIcons();
     }
 
     public void setShelfIcons(NotificationIconContainer icons) {
-        if (NotificationShelfController.checkRefactorFlagEnabled(mFeatureFlags)) {
+        if (mShelfRefactor.expectEnabled()) {
             mShelfIcons = icons;
         }
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/ExpandHelperTest.java b/packages/SystemUI/tests/src/com/android/systemui/ExpandHelperTest.java
index 6decb88..5867a40c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/ExpandHelperTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/ExpandHelperTest.java
@@ -30,6 +30,8 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.keyguard.KeyguardUpdateMonitor;
+import com.android.systemui.flags.FakeFeatureFlags;
+import com.android.systemui.flags.Flags;
 import com.android.systemui.statusbar.NotificationMediaManager;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.row.NotificationTestHelper;
@@ -43,12 +45,14 @@
 @RunWithLooper
 public class ExpandHelperTest extends SysuiTestCase {
 
+    private final FakeFeatureFlags mFeatureFlags = new FakeFeatureFlags();
     private ExpandableNotificationRow mRow;
     private ExpandHelper mExpandHelper;
     private ExpandHelper.Callback mCallback;
 
     @Before
     public void setUp() throws Exception {
+        mFeatureFlags.setDefault(Flags.NOTIFICATION_INLINE_REPLY_ANIMATION);
         mDependency.injectMockDependency(KeyguardUpdateMonitor.class);
         mDependency.injectMockDependency(NotificationMediaManager.class);
         allowTestableLooperAsMainThread();
@@ -56,7 +60,8 @@
         NotificationTestHelper helper = new NotificationTestHelper(
                 mContext,
                 mDependency,
-                TestableLooper.get(this));
+                TestableLooper.get(this),
+                mFeatureFlags);
         mRow = helper.createRow();
         mCallback = mock(ExpandHelper.Callback.class);
         mExpandHelper = new ExpandHelper(context, mCallback, 10, 100);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowTest.java
index 608778e..1dc8453 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowTest.java
@@ -27,6 +27,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
@@ -87,6 +88,7 @@
 @RunWithLooper
 public class ExpandableNotificationRowTest extends SysuiTestCase {
 
+    private final FakeFeatureFlags mFeatureFlags = new FakeFeatureFlags();
     private NotificationTestHelper mNotificationTestHelper;
     @Rule public MockitoRule mockito = MockitoJUnit.rule();
 
@@ -96,12 +98,10 @@
         mNotificationTestHelper = new NotificationTestHelper(
                 mContext,
                 mDependency,
-                TestableLooper.get(this));
+                TestableLooper.get(this),
+                mFeatureFlags);
         mNotificationTestHelper.setDefaultInflationFlags(FLAG_CONTENT_VIEW_ALL);
-
-        FakeFeatureFlags fakeFeatureFlags = new FakeFeatureFlags();
-        fakeFeatureFlags.set(Flags.SENSITIVE_REVEAL_ANIM, false);
-        mNotificationTestHelper.setFeatureFlags(fakeFeatureFlags);
+        mFeatureFlags.setDefault(Flags.SENSITIVE_REVEAL_ANIM);
     }
 
     @Test
@@ -183,6 +183,14 @@
     }
 
     @Test
+    public void testSetSensitiveOnNotifRowNotifiesOfHeightChange_withOtherFlagValue()
+            throws Exception {
+        FakeFeatureFlags flags = mFeatureFlags;
+        flags.set(Flags.SENSITIVE_REVEAL_ANIM, !flags.isEnabled(Flags.SENSITIVE_REVEAL_ANIM));
+        testSetSensitiveOnNotifRowNotifiesOfHeightChange();
+    }
+
+    @Test
     public void testSetSensitiveOnNotifRowNotifiesOfHeightChange() throws Exception {
         // GIVEN a sensitive notification row that's currently redacted
         ExpandableNotificationRow row = mNotificationTestHelper.createRow();
@@ -199,10 +207,19 @@
         // WHEN the row is set to no longer be sensitive
         row.setSensitive(false, true);
 
+        boolean expectAnimation = mFeatureFlags.isEnabled(Flags.SENSITIVE_REVEAL_ANIM);
         // VERIFY that the height change listener is invoked
         assertThat(row.getShowingLayout()).isSameInstanceAs(row.getPrivateLayout());
         assertThat(row.getIntrinsicHeight()).isGreaterThan(0);
-        verify(listener).onHeightChanged(eq(row), eq(false));
+        verify(listener).onHeightChanged(eq(row), eq(expectAnimation));
+    }
+
+    @Test
+    public void testSetSensitiveOnGroupRowNotifiesOfHeightChange_withOtherFlagValue()
+            throws Exception {
+        FakeFeatureFlags flags = mFeatureFlags;
+        flags.set(Flags.SENSITIVE_REVEAL_ANIM, !flags.isEnabled(Flags.SENSITIVE_REVEAL_ANIM));
+        testSetSensitiveOnGroupRowNotifiesOfHeightChange();
     }
 
     @Test
@@ -222,10 +239,19 @@
         // WHEN the row is set to no longer be sensitive
         group.setSensitive(false, true);
 
+        boolean expectAnimation = mFeatureFlags.isEnabled(Flags.SENSITIVE_REVEAL_ANIM);
         // VERIFY that the height change listener is invoked
         assertThat(group.getShowingLayout()).isSameInstanceAs(group.getPrivateLayout());
         assertThat(group.getIntrinsicHeight()).isGreaterThan(0);
-        verify(listener).onHeightChanged(eq(group), eq(false));
+        verify(listener).onHeightChanged(eq(group), eq(expectAnimation));
+    }
+
+    @Test
+    public void testSetSensitiveOnPublicRowDoesNotNotifyOfHeightChange_withOtherFlagValue()
+            throws Exception {
+        FakeFeatureFlags flags = mFeatureFlags;
+        flags.set(Flags.SENSITIVE_REVEAL_ANIM, !flags.isEnabled(Flags.SENSITIVE_REVEAL_ANIM));
+        testSetSensitiveOnPublicRowDoesNotNotifyOfHeightChange();
     }
 
     @Test
@@ -254,7 +280,7 @@
         assertThat(publicRow.getIntrinsicHeight()).isGreaterThan(0);
         assertThat(publicRow.getPrivateLayout().getMinHeight())
                 .isEqualTo(publicRow.getPublicLayout().getMinHeight());
-        verify(listener, never()).onHeightChanged(eq(publicRow), eq(false));
+        verify(listener, never()).onHeightChanged(eq(publicRow), anyBoolean());
     }
 
     private void measureAndLayout(ExpandableNotificationRow row) {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java
index 1a644d3..d21029d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java
@@ -48,12 +48,16 @@
 import android.view.LayoutInflater;
 import android.widget.RemoteViews;
 
+import androidx.annotation.NonNull;
+
 import com.android.internal.logging.MetricsLogger;
 import com.android.internal.statusbar.IStatusBarService;
 import com.android.systemui.TestableDependency;
 import com.android.systemui.classifier.FalsingCollectorFake;
 import com.android.systemui.classifier.FalsingManagerFake;
+import com.android.systemui.flags.FakeFeatureFlags;
 import com.android.systemui.flags.FeatureFlags;
+import com.android.systemui.flags.Flags;
 import com.android.systemui.media.controls.util.MediaFeatureFlag;
 import com.android.systemui.media.dialog.MediaOutputDialogFactory;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
@@ -90,6 +94,7 @@
 
 import org.mockito.ArgumentCaptor;
 
+import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executor;
@@ -130,14 +135,24 @@
     private final NotificationDismissibilityProvider mDismissibilityProvider;
     public final Runnable mFutureDismissalRunnable;
     private @InflationFlag int mDefaultInflationFlags;
-    private FeatureFlags mFeatureFlags;
+    private final FakeFeatureFlags mFeatureFlags;
 
     public NotificationTestHelper(
             Context context,
             TestableDependency dependency,
             TestableLooper testLooper) {
+        this(context, dependency, testLooper, new FakeFeatureFlags());
+    }
+
+    public NotificationTestHelper(
+            Context context,
+            TestableDependency dependency,
+            TestableLooper testLooper,
+            @NonNull FakeFeatureFlags featureFlags) {
         mContext = context;
         mTestLooper = testLooper;
+        mFeatureFlags = Objects.requireNonNull(featureFlags);
+        dependency.injectTestDependency(FeatureFlags.class, mFeatureFlags);
         dependency.injectMockDependency(NotificationMediaManager.class);
         dependency.injectMockDependency(NotificationShadeWindowController.class);
         dependency.injectMockDependency(MediaOutputDialogFactory.class);
@@ -183,17 +198,12 @@
         mFutureDismissalRunnable = mock(Runnable.class);
         when(mOnUserInteractionCallback.registerFutureDismissal(any(), anyInt()))
                 .thenReturn(mFutureDismissalRunnable);
-        mFeatureFlags = mock(FeatureFlags.class);
     }
 
     public void setDefaultInflationFlags(@InflationFlag int defaultInflationFlags) {
         mDefaultInflationFlags = defaultInflationFlags;
     }
 
-    public void setFeatureFlags(FeatureFlags featureFlags) {
-        mFeatureFlags = featureFlags;
-    }
-
     public ExpandableNotificationRowLogger getMockLogger() {
         return mMockLogger;
     }
@@ -527,6 +537,10 @@
             @InflationFlag int extraInflationFlags,
             int importance)
             throws Exception {
+        // NOTE: This flag is read when the ExpandableNotificationRow is inflated, so it needs to be
+        //  set, but we do not want to override an existing value that is needed by a specific test.
+        mFeatureFlags.setDefault(Flags.IMPROVED_HUN_ANIMATIONS);
+
         LayoutInflater inflater = (LayoutInflater) mContext.getSystemService(
                 mContext.LAYOUT_INFLATER_SERVICE);
         mRow = (ExpandableNotificationRow) inflater.inflate(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/AmbientStateTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/AmbientStateTest.kt
index 09382ec..3d75288 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/AmbientStateTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/AmbientStateTest.kt
@@ -20,7 +20,6 @@
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.dump.DumpManager
-import com.android.systemui.flags.FeatureFlags
 import com.android.systemui.shade.transition.LargeScreenShadeInterpolator
 import com.android.systemui.statusbar.StatusBarState
 import com.android.systemui.statusbar.phone.StatusBarKeyguardViewManager
@@ -42,7 +41,6 @@
     private val bypassController = StackScrollAlgorithm.BypassController { false }
     private val statusBarKeyguardViewManager = mock<StatusBarKeyguardViewManager>()
     private val largeScreenShadeInterpolator = mock<LargeScreenShadeInterpolator>()
-    private val featureFlags = mock<FeatureFlags>()
 
     private lateinit var sut: AmbientState
 
@@ -55,8 +53,7 @@
                 sectionProvider,
                 bypassController,
                 statusBarKeyguardViewManager,
-                largeScreenShadeInterpolator,
-                featureFlags
+                largeScreenShadeInterpolator
             )
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSectionsManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSectionsManagerTest.java
index f38881c..4b145d8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSectionsManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSectionsManagerTest.java
@@ -30,7 +30,6 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.systemui.SysuiTestCase;
-import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.media.controls.ui.KeyguardMediaController;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.statusbar.StatusBarState;
@@ -65,7 +64,6 @@
     @Mock private SectionHeaderController mPeopleHeaderController;
     @Mock private SectionHeaderController mAlertingHeaderController;
     @Mock private SectionHeaderController mSilentHeaderController;
-    @Mock private FeatureFlags mFeatureFlag;
 
     private NotificationSectionsManager mSectionsManager;
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt
index 8d751e3..1dc0ab0 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt
@@ -9,7 +9,9 @@
 import com.android.systemui.R
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.ShadeInterpolation
+import com.android.systemui.flags.FakeFeatureFlags
 import com.android.systemui.flags.FeatureFlags
+import com.android.systemui.flags.Flags
 import com.android.systemui.shade.transition.LargeScreenShadeInterpolator
 import com.android.systemui.statusbar.NotificationShelf
 import com.android.systemui.statusbar.StatusBarIconView
@@ -20,6 +22,7 @@
 import junit.framework.Assert.assertEquals
 import junit.framework.Assert.assertFalse
 import junit.framework.Assert.assertTrue
+import org.junit.Assume.assumeTrue
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -34,22 +37,32 @@
 @SmallTest
 @RunWith(AndroidTestingRunner::class)
 @RunWithLooper
-class NotificationShelfTest : SysuiTestCase() {
+open class NotificationShelfTest : SysuiTestCase() {
+
+    open val useShelfRefactor: Boolean = false
+    open val useSensitiveReveal: Boolean = false
+    private val flags = FakeFeatureFlags()
 
     @Mock
     private lateinit var largeScreenShadeInterpolator: LargeScreenShadeInterpolator
     @Mock
-    private lateinit var flags: FeatureFlags
-    @Mock
     private lateinit var ambientState: AmbientState
     @Mock
     private lateinit var hostLayoutController: NotificationStackScrollLayoutController
+    @Mock
+    private lateinit var hostLayout: NotificationStackScrollLayout
+    @Mock
+    private lateinit var roundnessManager: NotificationRoundnessManager
 
     private lateinit var shelf: NotificationShelf
 
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
+        mDependency.injectTestDependency(FeatureFlags::class.java, flags)
+        flags.set(Flags.NOTIFICATION_SHELF_REFACTOR, useShelfRefactor)
+        flags.set(Flags.SENSITIVE_REVEAL_ANIM, useSensitiveReveal)
+        flags.setDefault(Flags.IMPROVED_HUN_ANIMATIONS)
         val root = FrameLayout(context)
         shelf = LayoutInflater.from(root.context)
                 .inflate(/* resource = */ R.layout.status_bar_notification_shelf,
@@ -57,10 +70,13 @@
                     /* attachToRoot = */false) as NotificationShelf
 
         whenever(ambientState.largeScreenShadeInterpolator).thenReturn(largeScreenShadeInterpolator)
-        whenever(ambientState.featureFlags).thenReturn(flags)
         whenever(ambientState.isSmallScreen).thenReturn(true)
 
-        shelf.bind(ambientState, /* hostLayoutController */ hostLayoutController)
+        if (useShelfRefactor) {
+            shelf.bind(ambientState, hostLayout, roundnessManager)
+        } else {
+            shelf.bind(ambientState, hostLayoutController)
+        }
         shelf.layout(/* left */ 0, /* top */ 0, /* right */ 30, /* bottom */5)
     }
 
@@ -345,7 +361,7 @@
     @Test
     fun updateState_withNullLastVisibleBackgroundChild_hideShelf() {
         // GIVEN
-        shelf.setSensitiveRevealAnimEnabled(true)
+        assumeTrue(useSensitiveReveal)
         whenever(ambientState.stackY).thenReturn(100f)
         whenever(ambientState.stackHeight).thenReturn(100f)
         val paddingBetweenElements =
@@ -372,7 +388,7 @@
     @Test
     fun updateState_withNullFirstViewInShelf_hideShelf() {
         // GIVEN
-        shelf.setSensitiveRevealAnimEnabled(true)
+        assumeTrue(useSensitiveReveal)
         whenever(ambientState.stackY).thenReturn(100f)
         whenever(ambientState.stackHeight).thenReturn(100f)
         val paddingBetweenElements =
@@ -399,7 +415,7 @@
     @Test
     fun updateState_withCollapsedShade_hideShelf() {
         // GIVEN
-        shelf.setSensitiveRevealAnimEnabled(true)
+        assumeTrue(useSensitiveReveal)
         whenever(ambientState.stackY).thenReturn(100f)
         whenever(ambientState.stackHeight).thenReturn(100f)
         val paddingBetweenElements =
@@ -426,7 +442,7 @@
     @Test
     fun updateState_withHiddenSectionBeforeShelf_hideShelf() {
         // GIVEN
-        shelf.setSensitiveRevealAnimEnabled(true)
+        assumeTrue(useSensitiveReveal)
         whenever(ambientState.stackY).thenReturn(100f)
         whenever(ambientState.stackHeight).thenReturn(100f)
         val paddingBetweenElements =
@@ -486,3 +502,25 @@
         assertEquals(expectedAlpha, shelf.viewState.alpha)
     }
 }
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+@RunWithLooper
+class NotificationShelfWithRefactorTest : NotificationShelfTest() {
+    override val useShelfRefactor: Boolean = true
+}
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+@RunWithLooper
+class NotificationShelfWithSensitiveRevealTest : NotificationShelfTest() {
+    override val useSensitiveReveal: Boolean = true
+}
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+@RunWithLooper
+class NotificationShelfWithBothFlagsTest : NotificationShelfTest() {
+    override val useShelfRefactor: Boolean = true
+    override val useSensitiveReveal: Boolean = true
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutControllerTest.java
index ee8325e..07eadf7c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutControllerTest.java
@@ -54,7 +54,6 @@
 import com.android.systemui.classifier.FalsingManagerFake;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.flags.FakeFeatureFlags;
-import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.flags.Flags;
 import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository;
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor;
@@ -119,6 +118,7 @@
 @RunWith(AndroidTestingRunner.class)
 public class NotificationStackScrollLayoutControllerTest extends SysuiTestCase {
 
+    private final FakeFeatureFlags mFeatureFlags = new FakeFeatureFlags();
     @Mock private NotificationGutsManager mNotificationGutsManager;
     @Mock private NotificationsController mNotificationsController;
     @Mock private NotificationVisibilityProvider mVisibilityProvider;
@@ -157,7 +157,6 @@
     @Mock private StackStateLogger mStackLogger;
     @Mock private NotificationStackScrollLogger mLogger;
     @Mock private NotificationStackSizeCalculator mNotificationStackSizeCalculator;
-    private FakeFeatureFlags mFeatureFlags = new FakeFeatureFlags();
     @Mock private NotificationTargetsHelper mNotificationTargetsHelper;
     @Mock private SecureSettings mSecureSettings;
     @Mock private NotificationIconAreaController mIconAreaController;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java
index 8ad271b..72fcdec 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java
@@ -68,7 +68,9 @@
 import com.android.systemui.R;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.dump.DumpManager;
+import com.android.systemui.flags.FakeFeatureFlags;
 import com.android.systemui.flags.FeatureFlags;
+import com.android.systemui.flags.Flags;
 import com.android.systemui.shade.ShadeController;
 import com.android.systemui.shade.transition.LargeScreenShadeInterpolator;
 import com.android.systemui.statusbar.EmptyShadeView;
@@ -106,6 +108,7 @@
 @TestableLooper.RunWithLooper
 public class NotificationStackScrollLayoutTest extends SysuiTestCase {
 
+    private final FakeFeatureFlags mFeatureFlags = new FakeFeatureFlags();
     private NotificationStackScrollLayout mStackScroller;  // Normally test this
     private NotificationStackScrollLayout mStackScrollerInternal;  // See explanation below
     private AmbientState mAmbientState;
@@ -129,7 +132,6 @@
     @Mock private NotificationStackSizeCalculator mNotificationStackSizeCalculator;
     @Mock private StatusBarKeyguardViewManager mStatusBarKeyguardViewManager;
     @Mock private LargeScreenShadeInterpolator mLargeScreenShadeInterpolator;
-    @Mock private FeatureFlags mFeatureFlags;
 
     @Before
     public void setUp() throws Exception {
@@ -143,11 +145,25 @@
                 mNotificationSectionsManager,
                 mBypassController,
                 mStatusBarKeyguardViewManager,
-                mLargeScreenShadeInterpolator,
-                mFeatureFlags
+                mLargeScreenShadeInterpolator
         ));
 
+        // Register the debug flags we use
+        assertFalse(Flags.NSSL_DEBUG_LINES.getDefault());
+        assertFalse(Flags.NSSL_DEBUG_REMOVE_ANIMATION.getDefault());
+        mFeatureFlags.set(Flags.NSSL_DEBUG_LINES, false);
+        mFeatureFlags.set(Flags.NSSL_DEBUG_REMOVE_ANIMATION, false);
+
+        // Register the feature flags we use
+        // TODO: Ideally we wouldn't need to set these unless a test actually reads them,
+        //  and then we would test both configurations, but currently they are all read
+        //  in the constructor.
+        mFeatureFlags.setDefault(Flags.SENSITIVE_REVEAL_ANIM);
+        mFeatureFlags.setDefault(Flags.ANIMATED_NOTIFICATION_SHADE_INSETS);
+        mFeatureFlags.setDefault(Flags.NOTIFICATION_SHELF_REFACTOR);
+
         // Inject dependencies before initializing the layout
+        mDependency.injectTestDependency(FeatureFlags.class, mFeatureFlags);
         mDependency.injectTestDependency(SysuiStatusBarStateController.class, mBarState);
         mDependency.injectMockDependency(ShadeController.class);
         mDependency.injectTestDependency(
@@ -176,13 +192,18 @@
         mStackScrollerInternal.initView(getContext(), mNotificationSwipeHelper,
                 mNotificationStackSizeCalculator);
         mStackScroller = spy(mStackScrollerInternal);
-        mStackScroller.setShelfController(notificationShelfController);
+        if (!mFeatureFlags.isEnabled(Flags.NOTIFICATION_SHELF_REFACTOR)) {
+            mStackScroller.setShelfController(notificationShelfController);
+        }
         mStackScroller.setNotificationsController(mNotificationsController);
         mStackScroller.setEmptyShadeView(mEmptyShadeView);
         when(mStackScrollLayoutController.isHistoryEnabled()).thenReturn(true);
         when(mStackScrollLayoutController.getNotificationRoundnessManager())
                 .thenReturn(mNotificationRoundnessManager);
         mStackScroller.setController(mStackScrollLayoutController);
+        if (mFeatureFlags.isEnabled(Flags.NOTIFICATION_SHELF_REFACTOR)) {
+            mStackScroller.setShelf(mNotificationShelf);
+        }
 
         doNothing().when(mGroupExpansionManager).collapseGroups();
         doNothing().when(mExpandHelper).cancelImmediately();
@@ -899,7 +920,6 @@
     @Test
     public void testWindowInsetAnimationProgress_updatesBottomInset() {
         int bottomImeInset = 100;
-        mStackScrollerInternal.setAnimatedInsetsEnabled(true);
         WindowInsets windowInsets = new WindowInsets.Builder()
                 .setInsets(ime(), Insets.of(0, 0, 0, bottomImeInset)).build();
         ArrayList<WindowInsetsAnimation> windowInsetsAnimations = new ArrayList<>();
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSwipeHelperTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSwipeHelperTest.java
index df65c09..85a2bdd 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSwipeHelperTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSwipeHelperTest.java
@@ -47,6 +47,7 @@
 
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.classifier.FalsingManagerFake;
+import com.android.systemui.flags.FakeFeatureFlags;
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.plugins.statusbar.NotificationMenuRowPlugin;
 import com.android.systemui.plugins.statusbar.NotificationSwipeActionHelper.SnoozeOption;
@@ -83,7 +84,7 @@
     private Handler mHandler;
     private ExpandableNotificationRow mNotificationRow;
     private Runnable mFalsingCheck;
-    private FeatureFlags mFeatureFlags;
+    private final FeatureFlags mFeatureFlags = new FakeFeatureFlags();
 
     private static final int FAKE_ROW_WIDTH = 20;
     private static final int FAKE_ROW_HEIGHT = 20;
@@ -96,7 +97,6 @@
         mCallback = mock(NotificationSwipeHelper.NotificationCallback.class);
         mListener = mock(NotificationMenuRowPlugin.OnMenuEventListener.class);
         mNotificationRoundnessManager = mock(NotificationRoundnessManager.class);
-        mFeatureFlags = mock(FeatureFlags.class);
         mSwipeHelper = spy(new NotificationSwipeHelper(
                 mContext.getResources(),
                 ViewConfiguration.get(mContext),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationTargetsHelperTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationTargetsHelperTest.kt
index 45725ce..e30947c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationTargetsHelperTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationTargetsHelperTest.kt
@@ -18,6 +18,7 @@
 @RunWith(AndroidTestingRunner::class)
 @RunWithLooper
 class NotificationTargetsHelperTest : SysuiTestCase() {
+    private val featureFlags = FakeFeatureFlags()
     lateinit var notificationTestHelper: NotificationTestHelper
     private val sectionsManager: NotificationSectionsManager = mock()
     private val stackScrollLayout: NotificationStackScrollLayout = mock()
@@ -26,10 +27,10 @@
     fun setUp() {
         allowTestableLooperAsMainThread()
         notificationTestHelper =
-            NotificationTestHelper(mContext, mDependency, TestableLooper.get(this))
+            NotificationTestHelper(mContext, mDependency, TestableLooper.get(this), featureFlags)
     }
 
-    private fun notificationTargetsHelper() = NotificationTargetsHelper(FakeFeatureFlags())
+    private fun notificationTargetsHelper() = NotificationTargetsHelper(featureFlags)
 
     @Test
     fun targetsForFirstNotificationInGroup() {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithmTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithmTest.kt
index 4c97d20..987861d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithmTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithmTest.kt
@@ -8,7 +8,6 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.ShadeInterpolation.getContentAlpha
 import com.android.systemui.dump.DumpManager
-import com.android.systemui.flags.FeatureFlags
 import com.android.systemui.shade.transition.LargeScreenShadeInterpolator
 import com.android.systemui.statusbar.EmptyShadeView
 import com.android.systemui.statusbar.NotificationShelf
@@ -45,7 +44,6 @@
     private val dumpManager = mock<DumpManager>()
     private val mStatusBarKeyguardViewManager = mock<StatusBarKeyguardViewManager>()
     private val notificationShelf = mock<NotificationShelf>()
-    private val featureFlags = mock<FeatureFlags>()
     private val emptyShadeView = EmptyShadeView(context, /* attrs= */ null).apply {
         layout(/* l= */ 0, /* t= */ 0, /* r= */ 100, /* b= */ 100)
     }
@@ -56,7 +54,6 @@
             /* bypassController */ { false },
             mStatusBarKeyguardViewManager,
             largeScreenShadeInterpolator,
-            featureFlags,
         )
 
     private val testableResources = mContext.getOrCreateTestableResources()
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarterTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarterTest.java
index 9c7f619..33144f2 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarterTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarterTest.java
@@ -64,7 +64,7 @@
 import com.android.systemui.animation.ActivityLaunchAnimator;
 import com.android.systemui.assist.AssistManager;
 import com.android.systemui.classifier.FalsingCollectorFake;
-import com.android.systemui.flags.FeatureFlags;
+import com.android.systemui.flags.FakeFeatureFlags;
 import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository;
 import com.android.systemui.plugins.ActivityStarter;
 import com.android.systemui.plugins.ActivityStarter.OnDismissAction;
@@ -117,6 +117,7 @@
 public class StatusBarNotificationActivityStarterTest extends SysuiTestCase {
 
     private static final int DISPLAY_ID = 0;
+    private final FakeFeatureFlags mFeatureFlags = new FakeFeatureFlags();
 
     @Mock
     private AssistManager mAssistManager;
@@ -256,7 +257,7 @@
                         notificationAnimationProvider,
                         mock(LaunchFullScreenIntentProvider.class),
                         mPowerInteractor,
-                        mock(FeatureFlags.class),
+                        mFeatureFlags,
                         mUserTracker
                 );
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/RemoteInputViewTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/RemoteInputViewTest.java
index 391c8ca..7c285b8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/RemoteInputViewTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/RemoteInputViewTest.java
@@ -102,6 +102,8 @@
             "com.android.sysuitest.dummynotificationsender";
     private static final int DUMMY_MESSAGE_APP_ID = Process.LAST_APPLICATION_UID - 1;
 
+    private final FakeFeatureFlags mFeatureFlags = new FakeFeatureFlags();
+
     @Mock private RemoteInputController mController;
     @Mock private ShortcutManager mShortcutManager;
     @Mock private RemoteInputQuickSettingsDisabler mRemoteInputQuickSettingsDisabler;
@@ -453,8 +455,7 @@
     private RemoteInputViewController bindController(
             RemoteInputView view,
             NotificationEntry entry) {
-        FakeFeatureFlags fakeFeatureFlags = new FakeFeatureFlags();
-        fakeFeatureFlags.set(Flags.NOTIFICATION_INLINE_REPLY_ANIMATION, true);
+        mFeatureFlags.set(Flags.NOTIFICATION_INLINE_REPLY_ANIMATION, true);
         RemoteInputViewControllerImpl viewController = new RemoteInputViewControllerImpl(
                 view,
                 entry,
@@ -462,7 +463,7 @@
                 mController,
                 mShortcutManager,
                 mUiEventLoggerFake,
-                fakeFeatureFlags
+                mFeatureFlags
                 );
         viewController.bind();
         return viewController;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
index 4839eeb..17bb73b 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
@@ -92,7 +92,7 @@
 import com.android.systemui.biometrics.AuthController;
 import com.android.systemui.colorextraction.SysuiColorExtractor;
 import com.android.systemui.dump.DumpManager;
-import com.android.systemui.flags.FeatureFlags;
+import com.android.systemui.flags.FakeFeatureFlags;
 import com.android.systemui.keyguard.KeyguardViewMediator;
 import com.android.systemui.model.SysUiState;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
@@ -305,6 +305,7 @@
     private TestableLooper mTestableLooper;
 
     private FakeDisplayTracker mDisplayTracker = new FakeDisplayTracker(mContext);
+    private final FakeFeatureFlags mFeatureFlags = new FakeFeatureFlags();
 
     private UserHandle mUser0;
 
@@ -423,7 +424,7 @@
                 mCommonNotifCollection,
                 mNotifPipeline,
                 mSysUiState,
-                mock(FeatureFlags.class),
+                mFeatureFlags,
                 mNotifPipelineFlags,
                 syncExecutor);
         mBubblesManager.addNotifCallback(mNotifCallback);
@@ -432,7 +433,8 @@
         mNotificationTestHelper = new NotificationTestHelper(
                 mContext,
                 mDependency,
-                TestableLooper.get(this));
+                TestableLooper.get(this),
+                mFeatureFlags);
         mRow = mNotificationTestHelper.createBubble(mDeleteIntent);
         mRow2 = mNotificationTestHelper.createBubble(mDeleteIntent);
         mNonBubbleNotifRow = mNotificationTestHelper.createRow();
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/flags/FakeFeatureFlags.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/flags/FakeFeatureFlags.kt
index b94f816e..36fa7e6 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/flags/FakeFeatureFlags.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/flags/FakeFeatureFlags.kt
@@ -62,6 +62,32 @@
         }
     }
 
+    /**
+     * Set the given flag's default value if no other value has been set.
+     *
+     * REMINDER: You should always test your code with your flag in both configurations, so
+     *  generally you should be setting a particular value.  This method should be reserved for
+     *  situations where the flag needs to be read (e.g. in the class constructor), but its
+     *  value shouldn't affect the actual test cases. In those cases, it's mildly safer to use
+     *  this method than to hard-code `false` or `true` because then at least if you're wrong,
+     *  and the flag value *does* matter, you'll notice when the flag is flipped and tests
+     *  start failing.
+     */
+    fun setDefault(flag: BooleanFlag) = booleanFlags.putIfAbsent(flag.id, flag.default)
+
+    /**
+     * Set the given flag's default value if no other value has been set.
+     *
+     * REMINDER: You should always test your code with your flag in both configurations, so
+     *  generally you should be setting a particular value.  This method should be reserved for
+     *  situations where the flag needs to be read (e.g. in the class constructor), but its
+     *  value shouldn't affect the actual test cases. In those cases, it's mildly safer to use
+     *  this method than to hard-code `false` or `true` because then at least if you're wrong,
+     *  and the flag value *does* matter, you'll notice when the flag is flipped and tests
+     *  start failing.
+     */
+    fun setDefault(flag: SysPropBooleanFlag) = booleanFlags.putIfAbsent(flag.id, flag.default)
+
     private fun notifyFlagChanged(flag: Flag<*>) {
         flagListeners[flag.id]?.let { listeners ->
             listeners.forEach { listener ->