Populate buckets to ensure correct corner rounds on notifications

Fixes: 204127880
Bug: 200956118
Bug: 199388061
Test: manual
Test: atest SystemUITests:ShadeListBuilderTest SystemUITests:NotifSpecBuilderTest

Merged-In: I36f6f2ac2a503daa01594248a1e6cffcfbd5c348
Change-Id: I36f6f2ac2a503daa01594248a1e6cffcfbd5c348
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationLockscreenUserManagerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationLockscreenUserManagerImpl.java
index db7d5c1..856052e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationLockscreenUserManagerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationLockscreenUserManagerImpl.java
@@ -19,8 +19,8 @@
 import static android.app.admin.DevicePolicyManager.ACTION_DEVICE_POLICY_MANAGER_STATE_CHANGED;
 
 import static com.android.systemui.DejankUtils.whitelistIpcs;
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_MEDIA_CONTROLS;
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_SILENT;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_MEDIA_CONTROLS;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_SILENT;
 
 import android.app.ActivityManager;
 import android.app.KeyguardManager;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntry.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntry.java
index 94ee868..66d019e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntry.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntry.java
@@ -31,7 +31,7 @@
 import static android.app.NotificationManager.Policy.SUPPRESSED_EFFECT_STATUS_BAR;
 
 import static com.android.systemui.statusbar.notification.collection.NotifCollection.REASON_NOT_CANCELED;
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_ALERTING;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_ALERTING;
 
 import static java.util.Objects.requireNonNull;
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/ShadeListBuilder.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/ShadeListBuilder.java
index 122fb1c..3730524 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/ShadeListBuilder.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/ShadeListBuilder.java
@@ -52,6 +52,7 @@
 import com.android.systemui.statusbar.notification.collection.listbuilder.pluggable.NotifStabilityManager;
 import com.android.systemui.statusbar.notification.collection.listbuilder.pluggable.Pluggable;
 import com.android.systemui.statusbar.notification.collection.notifcollection.CollectionReadyForBuildListener;
+import com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt;
 import com.android.systemui.util.Assert;
 import com.android.systemui.util.time.SystemClock;
 
@@ -738,7 +739,7 @@
             if (entry instanceof GroupEntry) {
                 GroupEntry parent = (GroupEntry) entry;
                 for (NotificationEntry child : parent.getChildren()) {
-                    child.getAttachState().setSection(section);
+                    setEntrySection(child, section);
                 }
                 parent.sortChildren(sChildComparator);
             }
@@ -957,10 +958,18 @@
             }
         }
 
-        entry.getAttachState().setSection(finalSection);
+        setEntrySection(entry, finalSection);
         return finalSection;
     }
 
+    private void setEntrySection(ListEntry entry, NotifSection finalSection) {
+        entry.getAttachState().setSection(finalSection);
+        NotificationEntry representativeEntry = entry.getRepresentativeEntry();
+        if (representativeEntry != null && finalSection != null) {
+            representativeEntry.setBucket(finalSection.getBucket());
+        }
+    }
+
     @NonNull
     private NotifSection findSection(ListEntry entry) {
         for (int i = 0; i < mNotifSections.size(); i++) {
@@ -1039,13 +1048,13 @@
         void onRenderList(@NonNull List<ListEntry> entries);
     }
 
-    private static final NotifSectioner DEFAULT_SECTIONER =
-            new NotifSectioner("UnknownSection") {
-                @Override
-                public boolean isInSection(ListEntry entry) {
-                    return true;
-                }
-            };
+    private static final NotifSectioner DEFAULT_SECTIONER = new NotifSectioner("UnknownSection",
+            NotificationPriorityBucketKt.BUCKET_UNKNOWN) {
+        @Override
+        public boolean isInSection(ListEntry entry) {
+            return true;
+        }
+    };
 
     private static final int MIN_CHILDREN_FOR_GROUP = 2;
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/AppOpsCoordinator.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/AppOpsCoordinator.java
index 3a87f68..301b185 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/AppOpsCoordinator.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/AppOpsCoordinator.java
@@ -30,6 +30,7 @@
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.notification.collection.listbuilder.pluggable.NotifFilter;
 import com.android.systemui.statusbar.notification.collection.listbuilder.pluggable.NotifSectioner;
+import com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt;
 import com.android.systemui.util.concurrency.DelayableExecutor;
 
 import javax.inject.Inject;
@@ -102,7 +103,8 @@
     /**
      * Puts foreground service notifications into its own section.
      */
-    private final NotifSectioner mNotifSectioner = new NotifSectioner("ForegroundService") {
+    private final NotifSectioner mNotifSectioner = new NotifSectioner("ForegroundService",
+            NotificationPriorityBucketKt.BUCKET_FOREGROUND_SERVICE) {
         @Override
         public boolean isInSection(ListEntry entry) {
             NotificationEntry notificationEntry = entry.getRepresentativeEntry();
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/ConversationCoordinator.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/ConversationCoordinator.kt
index efec94f..c385836 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/ConversationCoordinator.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/ConversationCoordinator.kt
@@ -26,6 +26,7 @@
 import com.android.systemui.statusbar.notification.dagger.PeopleHeader
 import com.android.systemui.statusbar.notification.people.PeopleNotificationIdentifier
 import com.android.systemui.statusbar.notification.people.PeopleNotificationIdentifier.Companion.TYPE_NON_PERSON
+import com.android.systemui.statusbar.notification.stack.BUCKET_PEOPLE
 import javax.inject.Inject
 
 /**
@@ -45,7 +46,7 @@
         }
     }
 
-    val sectioner = object : NotifSectioner("People") {
+    val sectioner = object : NotifSectioner("People", BUCKET_PEOPLE) {
         override fun isInSection(entry: ListEntry): Boolean =
                 isConversation(entry.representativeEntry!!)
         override fun getHeaderNodeController() =
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinator.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinator.java
index 1bde312..4d36251 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinator.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinator.java
@@ -34,6 +34,7 @@
 import com.android.systemui.statusbar.notification.dagger.IncomingHeader;
 import com.android.systemui.statusbar.notification.interruption.HeadsUpViewBinder;
 import com.android.systemui.statusbar.notification.interruption.NotificationInterruptStateProvider;
+import com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt;
 import com.android.systemui.statusbar.policy.HeadsUpManager;
 import com.android.systemui.statusbar.policy.OnHeadsUpChangedListener;
 
@@ -196,7 +197,8 @@
         }
     };
 
-    private final NotifSectioner mNotifSectioner = new NotifSectioner("HeadsUp") {
+    private final NotifSectioner mNotifSectioner = new NotifSectioner("HeadsUp",
+            NotificationPriorityBucketKt.BUCKET_HEADS_UP) {
         @Override
         public boolean isInSection(ListEntry entry) {
             return isCurrentlyShowingHun(entry);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/RankingCoordinator.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/RankingCoordinator.java
index c62214c..1a6a63a 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/RankingCoordinator.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/RankingCoordinator.java
@@ -31,6 +31,7 @@
 import com.android.systemui.statusbar.notification.collection.render.SectionHeaderController;
 import com.android.systemui.statusbar.notification.dagger.AlertingHeader;
 import com.android.systemui.statusbar.notification.dagger.SilentHeader;
+import com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt;
 
 import java.util.List;
 
@@ -82,7 +83,8 @@
         return mSilentNotifSectioner;
     }
 
-    private final NotifSectioner mAlertingNotifSectioner = new NotifSectioner("Alerting") {
+    private final NotifSectioner mAlertingNotifSectioner = new NotifSectioner("Alerting",
+            NotificationPriorityBucketKt.BUCKET_ALERTING) {
         @Override
         public boolean isInSection(ListEntry entry) {
             return mHighPriorityProvider.isHighPriority(entry);
@@ -99,7 +101,8 @@
         }
     };
 
-    private final NotifSectioner mSilentNotifSectioner = new NotifSectioner("Silent") {
+    private final NotifSectioner mSilentNotifSectioner = new NotifSectioner("Silent",
+            NotificationPriorityBucketKt.BUCKET_SILENT) {
         @Override
         public boolean isInSection(ListEntry entry) {
             return !mHighPriorityProvider.isHighPriority(entry);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/listbuilder/NotifSection.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/listbuilder/NotifSection.kt
index c9fc992..6424e37 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/listbuilder/NotifSection.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/listbuilder/NotifSection.kt
@@ -18,14 +18,17 @@
 
 import com.android.systemui.statusbar.notification.collection.listbuilder.pluggable.NotifSectioner
 import com.android.systemui.statusbar.notification.collection.render.NodeController
+import com.android.systemui.statusbar.notification.stack.PriorityBucket
 
 data class NotifSection(
     val sectioner: NotifSectioner,
     val index: Int
 ) {
     val label: String
-        get() = "Section($index, \"${sectioner.name}\")"
+        get() = "Section($index, $bucket, \"${sectioner.name}\")"
 
     val headerController: NodeController?
         get() = sectioner.headerNodeController
+
+    @PriorityBucket val bucket: Int = sectioner.bucket
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/listbuilder/pluggable/NotifSectioner.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/listbuilder/pluggable/NotifSectioner.java
index 58afca8..ef9ee11 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/listbuilder/pluggable/NotifSectioner.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/listbuilder/pluggable/NotifSectioner.java
@@ -22,15 +22,28 @@
 import com.android.systemui.statusbar.notification.collection.ShadeListBuilder;
 import com.android.systemui.statusbar.notification.collection.render.NodeController;
 import com.android.systemui.statusbar.notification.collection.render.NodeSpec;
+import com.android.systemui.statusbar.notification.stack.PriorityBucket;
 
 import java.util.List;
 
 /**
- * Pluggable for participating in notif sectioning. See {@link ShadeListBuilder#setSections}.
+ * Pluggable for participating in notif sectioning. See {@link ShadeListBuilder#setSectioners}.
  */
 public abstract class NotifSectioner extends Pluggable<NotifSectioner> {
-    protected NotifSectioner(String name) {
+    @PriorityBucket
+    private final int mBucket;
+
+    protected NotifSectioner(String name, @PriorityBucket int bucket) {
         super(name);
+        mBucket = bucket;
+    }
+
+    /**
+     * @return the "bucket" value to apply to entries in this section
+     */
+    @PriorityBucket
+    public final int getBucket() {
+        return mBucket;
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationPanelLogger.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationPanelLogger.java
index c147023..9faef1b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationPanelLogger.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationPanelLogger.java
@@ -16,12 +16,12 @@
 
 package com.android.systemui.statusbar.notification.logging;
 
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_ALERTING;
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_FOREGROUND_SERVICE;
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_HEADS_UP;
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_MEDIA_CONTROLS;
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_PEOPLE;
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_SILENT;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_ALERTING;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_FOREGROUND_SERVICE;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_HEADS_UP;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_MEDIA_CONTROLS;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_PEOPLE;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_SILENT;
 
 import android.annotation.Nullable;
 import android.service.notification.StatusBarNotification;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationPriorityBucket.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationPriorityBucket.kt
new file mode 100644
index 0000000..31f4857
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationPriorityBucket.kt
@@ -0,0 +1,25 @@
+package com.android.systemui.statusbar.notification.stack
+
+import android.annotation.IntDef
+
+/**
+ * For now, declare the available notification buckets (sections) here so that other
+ * presentation code can decide what to do based on an entry's buckets
+ */
+@Retention(AnnotationRetention.SOURCE)
+@IntDef(
+        prefix = ["BUCKET_"],
+        value = [
+            BUCKET_UNKNOWN, BUCKET_MEDIA_CONTROLS, BUCKET_HEADS_UP, BUCKET_FOREGROUND_SERVICE,
+            BUCKET_PEOPLE, BUCKET_ALERTING, BUCKET_SILENT
+        ]
+)
+annotation class PriorityBucket
+
+const val BUCKET_UNKNOWN = 0
+const val BUCKET_MEDIA_CONTROLS = 1
+const val BUCKET_HEADS_UP = 2
+const val BUCKET_FOREGROUND_SERVICE = 3
+const val BUCKET_PEOPLE = 4
+const val BUCKET_ALERTING = 5
+const val BUCKET_SILENT = 6
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationSection.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationSection.java
index ab39de0..bc172ce 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationSection.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationSection.java
@@ -16,7 +16,7 @@
 
 package com.android.systemui.statusbar.notification.stack;
 
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_MEDIA_CONTROLS;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_MEDIA_CONTROLS;
 
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationSectionsManager.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationSectionsManager.kt
index 6ee14b5..5f157a7 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationSectionsManager.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationSectionsManager.kt
@@ -16,7 +16,6 @@
 package com.android.systemui.statusbar.notification.stack
 
 import android.annotation.ColorInt
-import android.annotation.IntDef
 import android.annotation.LayoutRes
 import android.util.Log
 import android.view.LayoutInflater
@@ -448,25 +447,3 @@
         private const val DEBUG = false
     }
 }
-
-/**
- * For now, declare the available notification buckets (sections) here so that other
- * presentation code can decide what to do based on an entry's buckets
- */
-@Retention(AnnotationRetention.SOURCE)
-@IntDef(
-        prefix = ["BUCKET_"],
-        value = [
-            BUCKET_UNKNOWN, BUCKET_MEDIA_CONTROLS, BUCKET_HEADS_UP, BUCKET_FOREGROUND_SERVICE,
-            BUCKET_PEOPLE, BUCKET_ALERTING, BUCKET_SILENT
-        ]
-)
-annotation class PriorityBucket
-
-const val BUCKET_UNKNOWN = 0
-const val BUCKET_MEDIA_CONTROLS = 1
-const val BUCKET_HEADS_UP = 2
-const val BUCKET_FOREGROUND_SERVICE = 3
-const val BUCKET_PEOPLE = 4
-const val BUCKET_ALERTING = 5
-const val BUCKET_SILENT = 6
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 1f76425..6aa5f07 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
@@ -17,7 +17,7 @@
 package com.android.systemui.statusbar.notification.stack;
 
 import static com.android.internal.jank.InteractionJankMonitor.CUJ_NOTIFICATION_SHADE_SCROLL_FLING;
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_SILENT;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_SILENT;
 import static com.android.systemui.statusbar.notification.stack.StackStateAnimator.ANIMATION_DURATION_SWIPE;
 import static com.android.systemui.util.Utils.shouldUseSplitNotificationShade;
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationLockscreenUserManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationLockscreenUserManagerTest.java
index ea21aa9..23cca72 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationLockscreenUserManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationLockscreenUserManagerTest.java
@@ -20,10 +20,10 @@
 import static android.app.NotificationManager.IMPORTANCE_LOW;
 import static android.content.Intent.ACTION_USER_SWITCHED;
 
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_ALERTING;
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_MEDIA_CONTROLS;
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_PEOPLE;
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_SILENT;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_ALERTING;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_MEDIA_CONTROLS;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_PEOPLE;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_SILENT;
 
 import static junit.framework.Assert.assertFalse;
 import static junit.framework.Assert.assertTrue;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/ShadeListBuilderTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/ShadeListBuilderTest.java
index c862f97..e9e6718 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/ShadeListBuilderTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/ShadeListBuilderTest.java
@@ -1660,7 +1660,7 @@
         private final String mPackage;
 
         PackageSectioner(String pkg) {
-            super("PackageSection_" + pkg);
+            super("PackageSection_" + pkg, 0);
             mPackage = pkg;
         }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/NodeSpecBuilderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/NodeSpecBuilderTest.kt
index 2e676bb..ed48452 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/NodeSpecBuilderTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/NodeSpecBuilderTest.kt
@@ -26,6 +26,10 @@
 import com.android.systemui.statusbar.notification.collection.getAttachState
 import com.android.systemui.statusbar.notification.collection.listbuilder.NotifSection
 import com.android.systemui.statusbar.notification.collection.listbuilder.pluggable.NotifSectioner
+import com.android.systemui.statusbar.notification.stack.BUCKET_ALERTING
+import com.android.systemui.statusbar.notification.stack.BUCKET_PEOPLE
+import com.android.systemui.statusbar.notification.stack.BUCKET_SILENT
+import com.android.systemui.statusbar.notification.stack.PriorityBucket
 import com.android.systemui.util.mockito.any
 import org.junit.Before
 import org.junit.Test
@@ -45,11 +49,15 @@
     private var headerController1: NodeController = buildFakeController("header1")
     private var headerController2: NodeController = buildFakeController("header2")
 
-    private val section0 = buildSection(0, headerController0)
-    private val section0NoHeader = buildSection(0, null)
-    private val section1 = buildSection(1, headerController1)
-    private val section1NoHeader = buildSection(1, null)
-    private val section2 = buildSection(2, headerController2)
+    private val section0Bucket = BUCKET_PEOPLE
+    private val section1Bucket = BUCKET_ALERTING
+    private val section2Bucket = BUCKET_SILENT
+
+    private val section0 = buildSection(0, section0Bucket, headerController0)
+    private val section0NoHeader = buildSection(0, section0Bucket, null)
+    private val section1 = buildSection(1, section1Bucket, headerController1)
+    private val section1NoHeader = buildSection(1, section1Bucket, null)
+    private val section2 = buildSection(2, section2Bucket, headerController2)
 
     private val fakeViewBarn = FakeViewBarn()
 
@@ -297,8 +305,12 @@
     return controller
 }
 
-private fun buildSection(index: Int, nodeController: NodeController?): NotifSection {
-    return NotifSection(object : NotifSectioner("Section $index") {
+private fun buildSection(
+    index: Int,
+    @PriorityBucket bucket: Int,
+    nodeController: NodeController?
+): NotifSection {
+    return NotifSection(object : NotifSectioner("Section $index (bucket=$bucket)", bucket) {
 
         override fun isInSection(entry: ListEntry?): Boolean {
             throw NotImplementedError("This should never be called")
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 c1d2ea8..f11f8c4 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
@@ -18,11 +18,11 @@
 
 import static android.view.ViewGroup.LayoutParams.WRAP_CONTENT;
 
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_ALERTING;
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_FOREGROUND_SERVICE;
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_HEADS_UP;
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_PEOPLE;
-import static com.android.systemui.statusbar.notification.stack.NotificationSectionsManagerKt.BUCKET_SILENT;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_ALERTING;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_FOREGROUND_SERVICE;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_HEADS_UP;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_PEOPLE;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_SILENT;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -64,7 +64,6 @@
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.Answers;
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
@@ -608,7 +607,7 @@
         }
     }
 
-    private View mockNotification(int bucket, boolean isGone) {
+    private View mockNotification(@PriorityBucket int bucket, boolean isGone) {
         ExpandableNotificationRow notifRow =
                 mock(ExpandableNotificationRow.class, RETURNS_DEEP_STUBS);
         when(notifRow.getVisibility()).thenReturn(View.VISIBLE);