Merge "Simplify keyguard notification count logic" into tm-dev

packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java

@@ -417,7 +417,6 @@ public class NotificationStackScrollLayout extends ViewGroup implements Dumpable
     private NotificationShelf mShelf;
     private int mMaxDisplayedNotifications = -1;
     private float mKeyguardBottomPadding = -1;
-    private float mKeyguardNotificationAvailableSpace = -1;
     @VisibleForTesting int mStatusBarHeight;
     private int mMinInteractionHeight;
     private final Rect mClipRect = new Rect();
@@ -775,9 +774,11 @@ public class NotificationStackScrollLayout extends ViewGroup implements Dumpable
         y = (int) mMaxLayoutHeight;
         drawDebugInfo(canvas, y, Color.MAGENTA, /* label= */ "mMaxLayoutHeight = " + y);
 
+        // The space between mTopPadding and mKeyguardBottomPadding determines the available space
+        // for notifications on keyguard.
         if (mKeyguardBottomPadding >= 0) {
             y = getHeight() - (int) mKeyguardBottomPadding;
-            drawDebugInfo(canvas, y, Color.GRAY,
+            drawDebugInfo(canvas, y, Color.RED,
                     /* label= */ "getHeight() - mKeyguardBottomPadding = " + y);
         }
 
@@ -789,7 +790,7 @@ public class NotificationStackScrollLayout extends ViewGroup implements Dumpable
         drawDebugInfo(canvas, y, Color.CYAN, /* label= */ "mAmbientState.getStackY() = " + y);
 
         y = (int) (mAmbientState.getStackY() + mAmbientState.getStackHeight());
-        drawDebugInfo(canvas, y, Color.BLUE,
+        drawDebugInfo(canvas, y, Color.LTGRAY,
                 /* label= */ "mAmbientState.getStackY() + mAmbientState.getStackHeight() = " + y);
 
         y = (int) mAmbientState.getStackY() + mContentHeight;
@@ -800,10 +801,6 @@ public class NotificationStackScrollLayout extends ViewGroup implements Dumpable
         drawDebugInfo(canvas, y, Color.YELLOW,
                 /* label= */ "mAmbientState.getStackY() + mIntrinsicContentHeight = " + y);
 
-        y = (int) (mAmbientState.getStackY() + mKeyguardNotificationAvailableSpace);
-        drawDebugInfo(canvas, y, Color.RED, /* label= */
-                "mAmbientState.getStackY() + mKeyguardNotificationAvailableSpace = " + y);
-
         drawDebugInfo(canvas, mRoundedRectClippingBottom, Color.DKGRAY,
                 /* label= */ "mRoundedRectClippingBottom) = " + y);
     }
@@ -2267,10 +2264,11 @@ public class NotificationStackScrollLayout extends ViewGroup implements Dumpable
     @ShadeViewRefactor(RefactorComponent.STATE_RESOLVER)
     private void updateContentHeight() {
         final float scrimTopPadding = mAmbientState.isOnKeyguard() ? 0 : mMinimumPaddings;
+        final int shelfIntrinsicHeight = mShelf != null ? mShelf.getIntrinsicHeight() : 0;
         final int height =
                 (int) scrimTopPadding + (int) mNotificationStackSizeCalculator.computeHeight(
                         /* notificationStackScrollLayout= */ this, mMaxDisplayedNotifications,
-                        mShelf != null ? mShelf.getIntrinsicHeight() : 0);
+                        shelfIntrinsicHeight);
         mIntrinsicContentHeight = height;
 
         // The topPadding can be bigger than the regular padding when qs is expanded, in that
@@ -4914,15 +4912,6 @@ public class NotificationStackScrollLayout extends ViewGroup implements Dumpable
         mKeyguardBottomPadding = keyguardBottomPadding;
     }
 
-    /**
-     * For debugging only. Enables to draw a line related to the available size for notifications in
-     * keyguard.
-     */
-    public void setKeyguardAvailableSpaceForDebug(float keyguardNotificationAvailableSpace) {
-        mKeyguardNotificationAvailableSpace = keyguardNotificationAvailableSpace;
-    }
-
-
     @ShadeViewRefactor(RefactorComponent.SHADE_VIEW)
     public void setShouldShowShelfOnly(boolean shouldShowShelfOnly) {
         mShouldShowShelfOnly = shouldShowShelfOnly;

packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java

@@ -1305,12 +1305,6 @@ public class NotificationStackScrollLayoutController {
         mView.setKeyguardBottomPadding(keyguardBottomPadding);
     }
 
-    /** For debugging only. */
-    public void mKeyguardNotificationAvailableSpaceForDebug(
-            float keyguardNotificationAvailableSpace) {
-        mView.setKeyguardAvailableSpaceForDebug(keyguardNotificationAvailableSpace);
-    }
-
     public RemoteInputController.Delegate createDelegate() {
         return new RemoteInputController.Delegate() {
             public void setRemoteInputActive(NotificationEntry entry,

packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackSizeCalculator.kt

@@ -29,6 +29,7 @@ import com.android.systemui.statusbar.notification.row.ExpandableView
 import com.android.systemui.util.children
 import javax.inject.Inject
 import kotlin.math.max
+import kotlin.math.min
 import kotlin.properties.Delegates.notNull
 
 private const val TAG = "NotificationStackSizeCalculator"
@@ -51,9 +52,7 @@ constructor(
      */
     private var maxKeyguardNotifications by notNull<Int>()
 
-    /**
+    /** Minimum space between two notifications, see [calculateGapAndDividerHeight]. */
-     * Minimum space between two notifications, see [calculateGapAndDividerHeight].
-     */
     private var dividerHeight by notNull<Int>()
 
     init {
@@ -61,55 +60,34 @@ constructor(
     }
 
     /**
-     * Given the [availableSpace] constraint, calculates how many notification to show.
+     * Given the [totalAvailableSpace] constraint, calculates how many notification to show.
      *
      * This number is only valid in keyguard.
      *
-     * @param availableSpace space for notifications. This doesn't include the space for the shelf.
+     * @param totalAvailableSpace space for notifications. This includes the space for the shelf.
      */
     fun computeMaxKeyguardNotifications(
         stack: NotificationStackScrollLayout,
-        availableSpace: Float,
+        totalAvailableSpace: Float,
-        shelfHeight: Float
+        shelfIntrinsicHeight: Float
     ): Int {
+        val stackHeightSequence = computeHeightPerNotificationLimit(stack, shelfIntrinsicHeight)
+
+        var maxNotifications =
+            stackHeightSequence.lastIndexWhile { stackHeight -> stackHeight <= totalAvailableSpace }
+
+        if (onLockscreen()) {
+            maxNotifications = min(maxKeyguardNotifications, maxNotifications)
+        }
+
+        // Could be < 0 if the space available is less than the shelf size. Returns 0 in this case.
+        maxNotifications = max(0, maxNotifications)
         log {
             "computeMaxKeyguardNotifications(" +
-                "availableSpace=$availableSpace shelfHeight=$shelfHeight)"
+                "availableSpace=$totalAvailableSpace" +
+                " shelfHeight=$shelfIntrinsicHeight) -> $maxNotifications"
         }
-
+        return maxNotifications
-        val children: Sequence<ExpandableView> = stack.childrenSequence
-        var remainingSpace: Float = availableSpace
-        var count = 0
-        var previous: ExpandableView? = null
-        val onLockscreen = true
-        val showableRows = children.filter { it.isShowable(onLockscreen) }
-        val showableRowsCount = showableRows.count()
-        log { "\tshowableRowsCount=$showableRowsCount "}
-
-        showableRows.forEachIndexed { i, current ->
-            val spaceNeeded = current.spaceNeeded(count, previous, stack, onLockscreen)
-            val spaceAfter = remainingSpace - spaceNeeded
-            previous = current
-            log { "\ti=$i spaceNeeded=$spaceNeeded remainingSpace=$remainingSpace " +
-                    "spaceAfter=$spaceAfter" }
-
-            if (remainingSpace - spaceNeeded >= 0 && count < maxKeyguardNotifications) {
-                count += 1
-                remainingSpace -= spaceNeeded
-            } else if (remainingSpace - spaceNeeded > -shelfHeight && i == showableRowsCount - 1) {
-                log { "Show all notifications. Shelf not needed." }
-                // If this is the last one, and it fits using the space shelf would use, then we can
-                // display it, as the shelf will not be needed (as all notifications are shown).
-                return count + 1
-            } else {
-                log {
-                    "No more fit. Returning $count. Space used: ${availableSpace - remainingSpace}"
-                }
-                return count
-            }
-        }
-        log { "All fit. Returning $count" }
-        return count
     }
 
     /**
@@ -119,47 +97,60 @@ constructor(
      * @param stack stack containing notifications as children.
      * @param maxNotifications Maximum number of notifications. When reached, the others will go
      * into the shelf.
-     * @param shelfHeight height of the shelf. It might be zero.
+     * @param shelfIntrinsicHeight height of the shelf, without any padding. It might be zero.
      *
      * @return height of the stack, including shelf height, if needed.
      */
     fun computeHeight(
         stack: NotificationStackScrollLayout,
         maxNotifications: Int,
-        shelfHeight: Float
+        shelfIntrinsicHeight: Float
     ): Float {
-        val children: Sequence<ExpandableView> = stack.childrenSequence
+        val heightPerMaxNotifications =
-        val maxNotificationsArg = infiniteIfNegative(maxNotifications)
+            computeHeightPerNotificationLimit(stack, shelfIntrinsicHeight)
+        val height =
+            heightPerMaxNotifications.elementAtOrElse(maxNotifications) {
+                heightPerMaxNotifications.last() // Height with all notifications visible.
+            }
+        log { "computeHeight(maxNotifications=$maxNotifications) -> $height" }
+        return height
+    }
+
+    /** The ith result in the sequence is the height with ith max notifications. */
+    private fun computeHeightPerNotificationLimit(
+        stack: NotificationStackScrollLayout,
+        shelfIntrinsicHeight: Float
+    ): Sequence<Float> = sequence {
+        val children = stack.showableChildren().toList()
         var height = 0f
         var previous: ExpandableView? = null
-        var count = 0
         val onLockscreen = onLockscreen()
 
-        log { "computeHeight(maxNotification=$maxNotifications, shelf=$shelfHeight" }
+        yield(dividerHeight + shelfIntrinsicHeight) // Only shelf.
-        children.filter { it.isShowable(onLockscreen) }.forEach { current ->
+
-            if (count < maxNotificationsArg) {
+        children.forEachIndexed { i, currentNotification ->
-                val spaceNeeded = current.spaceNeeded(count, previous, stack, onLockscreen)
+            height += currentNotification.spaceNeeded(i, previous, stack, onLockscreen)
-                log { "\ti=$count spaceNeeded=$spaceNeeded" }
+            previous = currentNotification
-                height += spaceNeeded
+
-                count += 1
+            val shelfHeight =
+                if (i == children.lastIndex) {
+                    0f // No shelf needed.
                 } else {
-                height += current.calculateGapAndDividerHeight(stack, previous, count)
+                    val spaceBeforeShelf =
-                height += shelfHeight
+                        calculateGapAndDividerHeight(
-                log { "returning height with shelf -> $height" }
+                            stack, previous = currentNotification, current = children[i + 1], i)
-                return height
+                    spaceBeforeShelf + shelfIntrinsicHeight
                 }
-            previous = current
+
+            yield(height + shelfHeight)
         }
-        log { "Returning height without shelf -> $height" }
-        return height
     }
 
     fun updateResources() {
         maxKeyguardNotifications =
             infiniteIfNegative(resources.getInteger(R.integer.keyguard_max_notification_count))
 
-        dividerHeight =
+        dividerHeight = max(1, resources.getDimensionPixelSize(R.dimen.notification_divider_height))
-            max(1, resources.getDimensionPixelSize(R.dimen.notification_divider_height))
     }
 
     private val NotificationStackScrollLayout.childrenSequence: Sequence<ExpandableView>
@@ -180,7 +171,7 @@ constructor(
             } else {
                 intrinsicHeight.toFloat()
             }
-        size += calculateGapAndDividerHeight(stack, previousView, visibleIndex)
+        size += calculateGapAndDividerHeight(stack, previousView, current = this, visibleIndex)
         return size
     }
 
@@ -200,18 +191,22 @@ constructor(
         return true
     }
 
-    private fun ExpandableView.calculateGapAndDividerHeight(
+    private fun calculateGapAndDividerHeight(
         stack: NotificationStackScrollLayout,
         previous: ExpandableView?,
+        current: ExpandableView?,
         visibleIndex: Int
     ): Float {
-        var height = stack.calculateGapHeight(previous, /* current= */ this, visibleIndex)
+        var height = stack.calculateGapHeight(previous, current, visibleIndex)
         if (visibleIndex != 0) {
             height += dividerHeight
         }
         return height
     }
 
+    private fun NotificationStackScrollLayout.showableChildren() =
+        this.childrenSequence.filter { it.isShowable(onLockscreen()) }
+
     /**
      * Can a view be shown on the lockscreen when calculating the number of allowed notifications to
      * show?
@@ -240,4 +235,8 @@ constructor(
         } else {
             v
         }
+
+    /** Returns the last index where [predicate] returns true, or -1 if it was always false. */
+    private fun <T> Sequence<T>.lastIndexWhile(predicate: (T) -> Boolean): Int =
+        takeWhile(predicate).count() - 1
 }

packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationPanelViewController.java

@@ -317,8 +317,6 @@ public class NotificationPanelViewController extends PanelViewController {
     private boolean mShouldUseSplitNotificationShade;
     // The bottom padding reserved for elements of the keyguard measuring notifications
     private float mKeyguardNotificationBottomPadding;
-    // Space available for notifications.
-    private float mKeyguardNotificationAvailableSpace;
     // Current max allowed keyguard notifications determined by measuring the panel
     private int mMaxAllowedKeyguardNotifications;
 
@@ -1245,8 +1243,6 @@ public class NotificationPanelViewController extends PanelViewController {
                     mMaxAllowedKeyguardNotifications);
             mNotificationStackScrollLayoutController.setKeyguardBottomPaddingForDebug(
                     mKeyguardNotificationBottomPadding);
-            mNotificationStackScrollLayoutController.mKeyguardNotificationAvailableSpaceForDebug(
-                    mKeyguardNotificationAvailableSpace);
         } else {
             // no max when not on the keyguard
             mNotificationStackScrollLayoutController.setMaxDisplayedNotifications(-1);
@@ -1468,13 +1464,11 @@ public class NotificationPanelViewController extends PanelViewController {
      * @return the maximum keyguard notifications that can fit on the screen
      */
     private int computeMaxKeyguardNotifications() {
-        int notificationPadding = Math.max(
-                1, mResources.getDimensionPixelSize(R.dimen.notification_divider_height));
         float topPadding = mNotificationStackScrollLayoutController.getTopPadding();
-        float shelfHeight =
+        float shelfIntrinsicHeight =
                 mNotificationShelfController.getVisibility() == View.GONE
                         ? 0
-                        : mNotificationShelfController.getIntrinsicHeight() + notificationPadding;
+                        : mNotificationShelfController.getIntrinsicHeight();
 
         // Padding to add to the bottom of the stack to keep a minimum distance from the top of
         // the lock icon.
@@ -1493,13 +1487,11 @@ public class NotificationPanelViewController extends PanelViewController {
         float availableSpace =
                 mNotificationStackScrollLayoutController.getHeight()
                         - topPadding
-                        - shelfHeight
                         - bottomPadding;
-        mKeyguardNotificationAvailableSpace = availableSpace;
 
         return mNotificationStackSizeCalculator.computeMaxKeyguardNotifications(
                 mNotificationStackScrollLayoutController.getView(), availableSpace,
-                shelfHeight);
+                shelfIntrinsicHeight);
     }
 
     private void updateClock() {

packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackSizeCalculatorTest.kt

@@ -16,6 +16,7 @@
 
 package com.android.systemui.statusbar.notification.stack
 
+import android.annotation.DimenRes
 import android.service.notification.StatusBarNotification
 import android.testing.AndroidTestingRunner
 import android.view.View.VISIBLE
@@ -27,6 +28,7 @@ import com.android.systemui.statusbar.notification.collection.NotificationEntry
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
 import com.android.systemui.statusbar.notification.row.ExpandableView
 import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.mockito.nullable
 import com.google.common.truth.Truth.assertThat
 import org.junit.Before
@@ -34,8 +36,8 @@ import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mock
 import org.mockito.Mockito.mock
-import org.mockito.MockitoAnnotations
 import org.mockito.Mockito.`when` as whenever
+import org.mockito.MockitoAnnotations
 
 @SmallTest
 @RunWith(AndroidTestingRunner::class)
@@ -49,17 +51,15 @@ class NotificationStackSizeCalculatorTest : SysuiTestCase() {
 
     private lateinit var sizeCalculator: NotificationStackSizeCalculator
 
+    private val gapHeight = px(R.dimen.notification_section_divider_height)
+    private val dividerHeight = px(R.dimen.notification_divider_height)
+    private val shelfHeight = px(R.dimen.notification_shelf_height)
+    private val rowHeight = px(R.dimen.notification_max_height)
+
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
 
-        whenever(stackLayout.calculateGapHeight(nullable(), nullable(), any()))
-            .thenReturn(GAP_HEIGHT)
-        with(testableResources) {
-            addOverride(R.integer.keyguard_max_notification_count, -1)
-            addOverride(R.dimen.notification_divider_height, DIVIDER_HEIGHT.toInt())
-        }
-
         sizeCalculator =
             NotificationStackSizeCalculator(
                 statusBarStateController = sysuiStatusBarStateController,
@@ -68,7 +68,7 @@ class NotificationStackSizeCalculatorTest : SysuiTestCase() {
 
     @Test
     fun computeMaxKeyguardNotifications_zeroSpace_returnZero() {
-        val rows = listOf(createMockRow(height = ROW_HEIGHT))
+        val rows = listOf(createMockRow(height = rowHeight))
 
         val maxNotifications =
             computeMaxKeyguardNotifications(rows, availableSpace = 0f, shelfHeight = 0f)
@@ -87,105 +87,78 @@ class NotificationStackSizeCalculatorTest : SysuiTestCase() {
     }
 
     @Test
-    fun computeMaxKeyguardNotifications_spaceForOne_returnsOne() {
+    fun computeMaxKeyguardNotifications_spaceForOneAndShelf_returnsOne() {
-        val rowHeight = ROW_HEIGHT
+        setGapHeight(gapHeight)
-        val totalSpaceForEachRow = GAP_HEIGHT + rowHeight
+        val shelfHeight = rowHeight / 2 // Shelf absence won't leave room for another row.
-        val shelfHeight =
+        val availableSpace =
-            totalSpaceForEachRow / 2 // In this way shelf absence will not leave room for another.
+            listOf(rowHeight + dividerHeight, gapHeight + dividerHeight + shelfHeight).sum()
-        val spaceForOne = totalSpaceForEachRow
+        val rows = listOf(createMockRow(rowHeight), createMockRow(rowHeight))
-        val rows =
-            listOf(
-                createMockRow(rowHeight),
-                createMockRow(rowHeight))
 
-        val maxNotifications =
+        val maxNotifications = computeMaxKeyguardNotifications(rows, availableSpace, shelfHeight)
-            computeMaxKeyguardNotifications(
-                rows, availableSpace = spaceForOne, shelfHeight = shelfHeight)
-
-        assertThat(maxNotifications).isEqualTo(1)
-    }
-
-    @Test
-    fun computeMaxKeyguardNotifications_spaceForOne_shelfUsableForLastNotification_returnsTwo() {
-        val rowHeight = ROW_HEIGHT
-        val totalSpaceForEachRow = GAP_HEIGHT + rowHeight
-        val shelfHeight = totalSpaceForEachRow + DIVIDER_HEIGHT
-        val spaceForOne = totalSpaceForEachRow
-        val rows =
-            listOf(
-                createMockRow(rowHeight),
-                createMockRow(rowHeight))
-
-        val maxNotifications =
-            computeMaxKeyguardNotifications(
-                rows, availableSpace = spaceForOne, shelfHeight = shelfHeight)
 
         assertThat(maxNotifications).isEqualTo(1)
     }
 
     @Test
     fun computeMaxKeyguardNotifications_spaceForTwo_returnsTwo() {
-        val rowHeight = ROW_HEIGHT
+        setGapHeight(gapHeight)
-        val totalSpaceForEachRow = GAP_HEIGHT + rowHeight
+        val shelfHeight = shelfHeight + dividerHeight
-        val spaceForTwo = totalSpaceForEachRow * 2 + DIVIDER_HEIGHT
+        val availableSpace =
-        val rows =
             listOf(
-                createMockRow(rowHeight),
+                    rowHeight + dividerHeight,
-                createMockRow(rowHeight),
+                    gapHeight + rowHeight + dividerHeight,
-                createMockRow(rowHeight))
+                    gapHeight + dividerHeight + shelfHeight)
+                .sum()
+        val rows =
+            listOf(createMockRow(rowHeight), createMockRow(rowHeight), createMockRow(rowHeight))
 
-        val maxNotifications = computeMaxKeyguardNotifications(rows, spaceForTwo, shelfHeight = 0f)
+        val maxNotifications = computeMaxKeyguardNotifications(rows, availableSpace, shelfHeight)
 
         assertThat(maxNotifications).isEqualTo(2)
     }
 
     @Test
     fun computeHeight_returnsAtMostSpaceAvailable_withGapBeforeShelf() {
-        val rowHeight = ROW_HEIGHT
+        setGapHeight(gapHeight)
-        val shelfHeight = SHELF_HEIGHT
+        val shelfHeight = shelfHeight
-        val totalSpaceForEachRow = GAP_HEIGHT + rowHeight + DIVIDER_HEIGHT
+        val availableSpace =
-        val availableSpace = totalSpaceForEachRow * 2
+            listOf(
+                    rowHeight + dividerHeight,
+                    gapHeight + rowHeight + dividerHeight,
+                    gapHeight + dividerHeight + shelfHeight)
+                .sum()
 
         // All rows in separate sections (default setup).
         val rows =
-            listOf(
+            listOf(createMockRow(rowHeight), createMockRow(rowHeight), createMockRow(rowHeight))
-                createMockRow(rowHeight),
-                createMockRow(rowHeight),
-                createMockRow(rowHeight))
 
         val maxNotifications = computeMaxKeyguardNotifications(rows, availableSpace, shelfHeight)
         assertThat(maxNotifications).isEqualTo(2)
 
-        val height = sizeCalculator.computeHeight(stackLayout, maxNotifications, SHELF_HEIGHT)
+        val height = sizeCalculator.computeHeight(stackLayout, maxNotifications, this.shelfHeight)
-        assertThat(height).isAtMost(availableSpace + GAP_HEIGHT + SHELF_HEIGHT)
+        assertThat(height).isAtMost(availableSpace)
     }
 
     @Test
-    fun computeHeight_returnsAtMostSpaceAvailable_noGapBeforeShelf() {
+    fun computeHeight_noGapBeforeShelf_returnsAtMostSpaceAvailable() {
-        val rowHeight = ROW_HEIGHT
-        val shelfHeight = SHELF_HEIGHT
-        val totalSpaceForEachRow = GAP_HEIGHT + rowHeight + DIVIDER_HEIGHT
-        val availableSpace = totalSpaceForEachRow * 1
-
         // Both rows are in the same section.
-        whenever(stackLayout.calculateGapHeight(nullable(), nullable(), any()))
+        setGapHeight(0f)
-                .thenReturn(0f)
+        val rowHeight = rowHeight
-        val rows =
+        val shelfHeight = shelfHeight
-                listOf(
+        val availableSpace = listOf(rowHeight + dividerHeight, dividerHeight + shelfHeight).sum()
-                        createMockRow(rowHeight),
+        val rows = listOf(createMockRow(rowHeight), createMockRow(rowHeight))
-                        createMockRow(rowHeight))
 
         val maxNotifications = computeMaxKeyguardNotifications(rows, availableSpace, shelfHeight)
         assertThat(maxNotifications).isEqualTo(1)
 
-        val height = sizeCalculator.computeHeight(stackLayout, maxNotifications, SHELF_HEIGHT)
+        val height = sizeCalculator.computeHeight(stackLayout, maxNotifications, this.shelfHeight)
-        assertThat(height).isAtMost(availableSpace + SHELF_HEIGHT)
+        assertThat(height).isAtMost(availableSpace)
     }
 
     private fun computeMaxKeyguardNotifications(
         rows: List<ExpandableView>,
         availableSpace: Float,
-        shelfHeight: Float = SHELF_HEIGHT
+        shelfHeight: Float = this.shelfHeight
     ): Int {
         setupChildren(rows)
         return sizeCalculator.computeMaxKeyguardNotifications(
@@ -204,9 +177,9 @@ class NotificationStackSizeCalculatorTest : SysuiTestCase() {
         (1..number).map { createMockRow() }.toList()
 
     private fun createMockRow(
-        height: Float = ROW_HEIGHT,
+        height: Float = rowHeight,
         isRemoved: Boolean = false,
-        visibility: Int = VISIBLE,
+        visibility: Int = VISIBLE
     ): ExpandableNotificationRow {
         val row = mock(ExpandableNotificationRow::class.java)
         val entry = mock(NotificationEntry::class.java)
@@ -220,11 +193,12 @@ class NotificationStackSizeCalculatorTest : SysuiTestCase() {
         return row
     }
 
-    /** Default dimensions for tests that don't overwrite them. */
+    private fun setGapHeight(height: Float) {
-    companion object {
+        whenever(stackLayout.calculateGapHeight(nullable(), nullable(), any())).thenReturn(height)
-        const val GAP_HEIGHT = 12f
+        whenever(stackLayout.calculateGapHeight(nullable(), nullable(), /* visibleIndex= */ eq(0)))
-        const val DIVIDER_HEIGHT = 3f
+            .thenReturn(0f)
-        const val SHELF_HEIGHT = 14f
-        const val ROW_HEIGHT = SHELF_HEIGHT * 3
     }
+
+    private fun px(@DimenRes id: Int): Float =
+        testableResources.resources.getDimensionPixelSize(id).toFloat()
 }