Revert^2 "Simplify PriorityNestedScrollConnection"

This reverts commit 7f60633a61a30bf859f63e8c7e9eac26ddc67d70.

Reason for revert: A fix will be applied on top of this revert

Change-Id: I405ad66f57d65bfeec98cc2ab4f98896c2f616da

7 files changed, 165 insertions, 112 deletions

diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationScrimNestedScrollConnection.kt b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationScrimNestedScrollConnection.kt
index 8b9e9274b448..779934d8d211 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationScrimNestedScrollConnection.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationScrimNestedScrollConnection.kt
@@ -18,6 +18,8 @@ package com.android.systemui.notifications.ui.composable
 
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.ui.input.nestedscroll.NestedScrollConnection
+import androidx.compose.ui.util.fastCoerceAtLeast
+import androidx.compose.ui.util.fastCoerceAtMost
 import com.android.compose.nestedscroll.PriorityNestedScrollConnection
 
 /**
@@ -44,7 +46,7 @@ fun NotificationScrimNestedScrollConnection(
         orientation = Orientation.Vertical,
         // scrolling up and inner content is taller than the scrim, so scrim needs to
         // expand; content can scroll once scrim is at the minScrimOffset.
-        canStartPreScroll = { offsetAvailable, offsetBeforeStart ->
+        canStartPreScroll = { offsetAvailable, offsetBeforeStart, _ ->
             offsetAvailable < 0 &&
                 offsetBeforeStart == 0f &&
                 contentHeight() > minVisibleScrimHeight() &&
@@ -52,25 +54,21 @@ fun NotificationScrimNestedScrollConnection(
         },
         // scrolling down and content is done scrolling to top. After that, the scrim
         // needs to collapse; collapse the scrim until it is at the maxScrimOffset.
-        canStartPostScroll = { offsetAvailable, _ ->
+        canStartPostScroll = { offsetAvailable, _, _ ->
             offsetAvailable > 0 && (scrimOffset() < maxScrimOffset || isCurrentGestureOverscroll())
         },
         canStartPostFling = { false },
-        canContinueScroll = {
-            val currentHeight = scrimOffset()
-            minScrimOffset() < currentHeight && currentHeight < maxScrimOffset
-        },
-        canScrollOnFling = true,
+        canStopOnPreFling = { false },
         onStart = { offsetAvailable -> onStart(offsetAvailable) },
-        onScroll = { offsetAvailable ->
+        onScroll = { offsetAvailable, _ ->
             val currentHeight = scrimOffset()
             val amountConsumed =
                 if (offsetAvailable > 0) {
                     val amountLeft = maxScrimOffset - currentHeight
-                    offsetAvailable.coerceAtMost(amountLeft)
+                    offsetAvailable.fastCoerceAtMost(amountLeft)
                 } else {
                     val amountLeft = minScrimOffset() - currentHeight
-                    offsetAvailable.coerceAtLeast(amountLeft)
+                    offsetAvailable.fastCoerceAtLeast(amountLeft)
                 }
             snapScrimOffset(currentHeight + amountConsumed)
             amountConsumed
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationStackNestedScrollConnection.kt b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationStackNestedScrollConnection.kt
index a706585deebc..42dd1be51d44 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationStackNestedScrollConnection.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationStackNestedScrollConnection.kt
@@ -28,6 +28,7 @@ import androidx.compose.ui.platform.LocalConfiguration
 import androidx.compose.ui.platform.LocalDensity
 import androidx.compose.ui.unit.IntOffset
 import androidx.compose.ui.unit.dp
+import androidx.compose.ui.util.fastCoerceAtLeast
 import com.android.compose.nestedscroll.PriorityNestedScrollConnection
 import kotlin.math.max
 import kotlin.math.roundToInt
@@ -86,17 +87,20 @@ fun NotificationStackNestedScrollConnection(
 ): PriorityNestedScrollConnection {
     return PriorityNestedScrollConnection(
         orientation = Orientation.Vertical,
-        canStartPreScroll = { _, _ -> false },
-        canStartPostScroll = { offsetAvailable, offsetBeforeStart ->
+        canStartPreScroll = { _, _, _ -> false },
+        canStartPostScroll = { offsetAvailable, offsetBeforeStart, _ ->
             offsetAvailable < 0f && offsetBeforeStart < 0f && !canScrollForward()
         },
         canStartPostFling = { velocityAvailable -> velocityAvailable < 0f && !canScrollForward() },
-        canContinueScroll = { stackOffset() > 0f },
-        canScrollOnFling = true,
+        canStopOnPreFling = { false },
         onStart = { offsetAvailable -> onStart(offsetAvailable) },
-        onScroll = { offsetAvailable ->
-            onScroll(offsetAvailable)
-            offsetAvailable
+        onScroll = { offsetAvailable, _ ->
+            val minOffset = 0f
+            val consumed = offsetAvailable.fastCoerceAtLeast(minOffset - stackOffset())
+            if (consumed != 0f) {
+                onScroll(consumed)
+            }
+            consumed
         },
         onStop = { velocityAvailable ->
             onStop(velocityAvailable)
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
index e5ef79b37b88..03fa971f332d 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
@@ -612,7 +612,7 @@ internal class NestedScrollHandlerImpl(
 
         return PriorityNestedScrollConnection(
             orientation = orientation,
-            canStartPreScroll = { offsetAvailable, offsetBeforeStart ->
+            canStartPreScroll = { offsetAvailable, offsetBeforeStart, _ ->
                 canChangeScene =
                     if (isExternalOverscrollGesture()) false else offsetBeforeStart == 0f
 
@@ -644,7 +644,7 @@ internal class NestedScrollHandlerImpl(
                 isIntercepting = true
                 true
             },
-            canStartPostScroll = { offsetAvailable, offsetBeforeStart ->
+            canStartPostScroll = { offsetAvailable, offsetBeforeStart, _ ->
                 val behavior: NestedScrollBehavior =
                     when {
                         offsetAvailable > 0f -> topOrLeftBehavior
@@ -709,8 +709,7 @@ internal class NestedScrollHandlerImpl(
 
                 canStart
             },
-            canContinueScroll = { true },
-            canScrollOnFling = false,
+            canStopOnPreFling = { true },
             onStart = { offsetAvailable ->
                 val pointersInfo = pointersInfo()
                 dragController =
@@ -720,7 +719,7 @@ internal class NestedScrollHandlerImpl(
                         overSlop = if (isIntercepting) 0f else offsetAvailable,
                     )
             },
-            onScroll = { offsetAvailable ->
+            onScroll = { offsetAvailable, _ ->
                 val controller = dragController ?: error("Should be called after onStart")
 
                 val pointersInfo = pointersInfoOwner.pointersInfo()
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/LargeTopAppBarNestedScrollConnection.kt b/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/LargeTopAppBarNestedScrollConnection.kt
index 4ae323517b26..b05f02224327 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/LargeTopAppBarNestedScrollConnection.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/LargeTopAppBarNestedScrollConnection.kt
@@ -18,6 +18,8 @@ package com.android.compose.nestedscroll
 
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.ui.input.nestedscroll.NestedScrollConnection
+import androidx.compose.ui.util.fastCoerceAtLeast
+import androidx.compose.ui.util.fastCoerceAtMost
 
 /**
  * A [NestedScrollConnection] that listens for all vertical scroll events and responds in the
@@ -43,30 +45,26 @@ fun LargeTopAppBarNestedScrollConnection(
         orientation = Orientation.Vertical,
         // When swiping up, the LargeTopAppBar will shrink (to [minHeight]) and the content will
         // expand. Then, you can then scroll down the content.
-        canStartPreScroll = { offsetAvailable, offsetBeforeStart ->
+        canStartPreScroll = { offsetAvailable, offsetBeforeStart, _ ->
             offsetAvailable < 0 && offsetBeforeStart == 0f && height() > minHeight()
         },
         // When swiping down, the content will scroll up until it reaches the top. Then, the
         // LargeTopAppBar will expand until it reaches its [maxHeight].
-        canStartPostScroll = { offsetAvailable, _ ->
+        canStartPostScroll = { offsetAvailable, _, _ ->
             offsetAvailable > 0 && height() < maxHeight()
         },
         canStartPostFling = { false },
-        canContinueScroll = {
-            val currentHeight = height()
-            minHeight() < currentHeight && currentHeight < maxHeight()
-        },
-        canScrollOnFling = true,
+        canStopOnPreFling = { false },
         onStart = { /* do nothing */ },
-        onScroll = { offsetAvailable ->
+        onScroll = { offsetAvailable, _ ->
             val currentHeight = height()
             val amountConsumed =
                 if (offsetAvailable > 0) {
                     val amountLeft = maxHeight() - currentHeight
-                    offsetAvailable.coerceAtMost(amountLeft)
+                    offsetAvailable.fastCoerceAtMost(amountLeft)
                 } else {
                     val amountLeft = minHeight() - currentHeight
-                    offsetAvailable.coerceAtLeast(amountLeft)
+                    offsetAvailable.fastCoerceAtLeast(amountLeft)
                 }
             onHeightChanged(currentHeight + amountConsumed)
             amountConsumed
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/PriorityNestedScrollConnection.kt b/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/PriorityNestedScrollConnection.kt
index a3641e6635e7..e3cddc1a8f1d 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/PriorityNestedScrollConnection.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/nestedscroll/PriorityNestedScrollConnection.kt
@@ -27,25 +27,39 @@ import kotlin.math.sign
 internal typealias SuspendedValue<T> = suspend () -> T
 
 /**
- * This [NestedScrollConnection] waits for a child to scroll ([onPreScroll] or [onPostScroll]), and
- * then decides (via [canStartPreScroll] or [canStartPostScroll]) if it should take over scrolling.
- * If it does, it will scroll before its children, until [canContinueScroll] allows it.
+ * A [NestedScrollConnection] that intercepts scroll events in priority mode.
  *
- * Note: Call [reset] before destroying this object to make sure you always get a call to [onStop]
- * after [onStart].
+ * Priority mode allows this connection to take control over scroll events within a nested scroll
+ * hierarchy. When in priority mode, this connection consumes scroll events before its children,
+ * enabling custom scrolling behaviors like sticky headers.
  *
+ * @param orientation The orientation of the scroll.
+ * @param canStartPreScroll lambda that returns true if the connection can start consuming scroll
+ *   events in pre-scroll mode.
+ * @param canStartPostScroll lambda that returns true if the connection can start consuming scroll
+ *   events in post-scroll mode.
+ * @param canStartPostFling lambda that returns true if the connection can start consuming scroll
+ *   events in post-fling mode.
+ * @param canStopOnPreFling lambda that returns true if the connection can stop consuming scroll
+ *   events in pre-fling (i.e. as soon as the user lifts their fingers).
+ * @param onStart lambda that is called when the connection starts consuming scroll events.
+ * @param onScroll lambda that is called when the connection consumes a scroll event and returns the
+ *   consumed amount.
+ * @param onStop lambda that is called when the connection stops consuming scroll events and returns
+ *   the consumed velocity.
  * @sample LargeTopAppBarNestedScrollConnection
  * @sample com.android.compose.animation.scene.NestedScrollHandlerImpl.nestedScrollConnection
  */
 class PriorityNestedScrollConnection(
     orientation: Orientation,
-    private val canStartPreScroll: (offsetAvailable: Float, offsetBeforeStart: Float) -> Boolean,
-    private val canStartPostScroll: (offsetAvailable: Float, offsetBeforeStart: Float) -> Boolean,
+    private val canStartPreScroll:
+        (offsetAvailable: Float, offsetBeforeStart: Float, source: NestedScrollSource) -> Boolean,
+    private val canStartPostScroll:
+        (offsetAvailable: Float, offsetBeforeStart: Float, source: NestedScrollSource) -> Boolean,
     private val canStartPostFling: (velocityAvailable: Float) -> Boolean,
-    private val canContinueScroll: (source: NestedScrollSource) -> Boolean,
-    private val canScrollOnFling: Boolean,
+    private val canStopOnPreFling: () -> Boolean,
     private val onStart: (offsetAvailable: Float) -> Unit,
-    private val onScroll: (offsetAvailable: Float) -> Float,
+    private val onScroll: (offsetAvailable: Float, source: NestedScrollSource) -> Float,
     private val onStop: (velocityAvailable: Float) -> SuspendedValue<Float>,
 ) : NestedScrollConnection, SpaceVectorConverter by SpaceVectorConverter(orientation) {
 
@@ -64,62 +78,48 @@ class PriorityNestedScrollConnection(
         // the beginning or from the last fling gesture.
         val offsetBeforeStart = offsetScrolledBeforePriorityMode - availableFloat
 
-        if (
-            isPriorityMode ||
-                (source == NestedScrollSource.SideEffect && !canScrollOnFling) ||
-                !canStartPostScroll(availableFloat, offsetBeforeStart)
-        ) {
+        if (isPriorityMode || !canStartPostScroll(availableFloat, offsetBeforeStart, source)) {
             // The priority mode cannot start so we won't consume the available offset.
             return Offset.Zero
         }
 
-        return onPriorityStart(availableFloat).toOffset()
+        return start(availableFloat, source).toOffset()
     }
 
     override fun onPreScroll(available: Offset, source: NestedScrollSource): Offset {
         if (!isPriorityMode) {
-            if (source == NestedScrollSource.UserInput || canScrollOnFling) {
-                val availableFloat = available.toFloat()
-                if (canStartPreScroll(availableFloat, offsetScrolledBeforePriorityMode)) {
-                    return onPriorityStart(availableFloat).toOffset()
-                }
-                // We want to track the amount of offset consumed before entering priority mode
-                offsetScrolledBeforePriorityMode += availableFloat
+            val availableFloat = available.toFloat()
+            if (canStartPreScroll(availableFloat, offsetScrolledBeforePriorityMode, source)) {
+                return start(availableFloat, source).toOffset()
             }
-
-            return Offset.Zero
-        }
-
-        val availableFloat = available.toFloat()
-        if (!canContinueScroll(source)) {
-            // Step 3a: We have lost priority and we no longer need to intercept scroll events.
-            onPriorityStop(velocity = 0f)
-
-            // We've just reset offsetScrolledBeforePriorityMode to 0f
             // We want to track the amount of offset consumed before entering priority mode
             offsetScrolledBeforePriorityMode += availableFloat
-
             return Offset.Zero
         }
 
-        // Step 2: We have the priority and can consume the scroll events.
-        return onScroll(availableFloat).toOffset()
+        return scroll(available.toFloat(), source).toOffset()
     }
 
     override suspend fun onPreFling(available: Velocity): Velocity {
-        if (isPriorityMode && canScrollOnFling) {
-            // We don't want to consume the velocity, we prefer to continue receiving scroll events.
+        if (!isPriorityMode) {
+            resetOffsetTracker()
             return Velocity.Zero
         }
-        // Step 3b: The finger is lifted, we can stop intercepting scroll events and use the speed
-        // of the fling gesture.
-        return onPriorityStop(velocity = available.toFloat()).invoke().toVelocity()
+
+        if (canStopOnPreFling()) {
+            // Step 3b: The finger is lifted, we can stop intercepting scroll events and use the
+            // velocity of the fling gesture.
+            return stop(velocityAvailable = available.toFloat()).toVelocity()
+        }
+
+        // We don't want to consume the velocity, we prefer to continue receiving scroll events.
+        return Velocity.Zero
     }
 
     override suspend fun onPostFling(consumed: Velocity, available: Velocity): Velocity {
         val availableFloat = available.toFloat()
         if (isPriorityMode) {
-            return onPriorityStop(velocity = availableFloat).invoke().toVelocity()
+            return stop(velocityAvailable = availableFloat).toVelocity()
         }
 
         if (!canStartPostFling(availableFloat)) {
@@ -131,10 +131,14 @@ class PriorityNestedScrollConnection(
         // TODO(b/291053278): Remove canStartPostFling() and instead make it possible to define the
         // overscroll behavior on the Scene level.
         val smallOffset = availableFloat.sign
-        onPriorityStart(availableOffset = smallOffset)
+        start(
+            availableOffset = smallOffset,
+            source = NestedScrollSource.SideEffect,
+            skipScroll = true,
+        )
 
         // This is the last event of a scroll gesture.
-        return onPriorityStop(availableFloat).invoke().toVelocity()
+        return stop(availableFloat).toVelocity()
     }
 
     /**
@@ -143,13 +147,25 @@ class PriorityNestedScrollConnection(
      * TODO(b/303224944) This method should be removed.
      */
     fun reset() {
-        // Step 3c: To ensure that an onStop is always called for every onStart.
-        onPriorityStop(velocity = 0f)
+        if (isPriorityMode) {
+            // Step 3c: To ensure that an onStop is always called for every onStart.
+            cancel()
+        } else {
+            resetOffsetTracker()
+        }
     }
 
-    private fun onPriorityStart(availableOffset: Float): Float {
-        if (isPriorityMode) {
-            error("This should never happen, onPriorityStart() was called when isPriorityMode")
+    private fun shouldStop(consumed: Float): Boolean {
+        return consumed == 0f
+    }
+
+    private fun start(
+        availableOffset: Float,
+        source: NestedScrollSource,
+        skipScroll: Boolean = false,
+    ): Float {
+        check(!isPriorityMode) {
+            "This should never happen, start() was called when isPriorityMode"
         }
 
         // Step 1: It's our turn! We start capturing scroll events when one of our children has an
@@ -160,19 +176,41 @@ class PriorityNestedScrollConnection(
         // lifted (step 3b), or this object has been destroyed (step 3c).
         onStart(availableOffset)
 
-        return onScroll(availableOffset)
+        return if (skipScroll) 0f else scroll(availableOffset, source)
     }
 
-    private fun onPriorityStop(velocity: Float): SuspendedValue<Float> {
-        // We can restart tracking the consumed offsets from scratch.
-        offsetScrolledBeforePriorityMode = 0f
+    private fun scroll(offsetAvailable: Float, source: NestedScrollSource): Float {
+        // Step 2: We have the priority and can consume the scroll events.
+        val consumedByScroll = onScroll(offsetAvailable, source)
 
-        if (!isPriorityMode) {
-            return { 0f }
+        if (shouldStop(consumedByScroll)) {
+            // Step 3a: We have lost priority and we no longer need to intercept scroll events.
+            cancel()
+
+            // We've just reset offsetScrolledBeforePriorityMode to 0f
+            // We want to track the amount of offset consumed before entering priority mode
+            offsetScrolledBeforePriorityMode += offsetAvailable - consumedByScroll
         }
 
+        return consumedByScroll
+    }
+
+    /** Reset the tracking of consumed offsets before entering in priority mode. */
+    private fun resetOffsetTracker() {
+        offsetScrolledBeforePriorityMode = 0f
+    }
+
+    private suspend fun stop(velocityAvailable: Float): Float {
+        check(isPriorityMode) { "This should never happen, stop() was called before start()" }
         isPriorityMode = false
+        resetOffsetTracker()
+        return onStop(velocityAvailable).invoke()
+    }
 
-        return onStop(velocity)
+    private fun cancel() {
+        check(isPriorityMode) { "This should never happen, cancel() was called before start()" }
+        isPriorityMode = false
+        resetOffsetTracker()
+        onStop(0f)
     }
 }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/nestedscroll/PriorityNestedScrollConnectionTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/nestedscroll/PriorityNestedScrollConnectionTest.kt
index badc43bd3e0f..9363ecbf63dc 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/nestedscroll/PriorityNestedScrollConnectionTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/nestedscroll/PriorityNestedScrollConnectionTest.kt
@@ -34,29 +34,28 @@ class PriorityNestedScrollConnectionTest {
     private var canStartPreScroll = false
     private var canStartPostScroll = false
     private var canStartPostFling = false
-    private var canContinueScroll = false
+    private var canStopOnPreFling = true
     private var isStarted = false
     private var lastScroll: Float? = null
-    private var returnOnScroll = 0f
+    private var consumeScroll = true
     private var lastStop: Float? = null
-    private var returnOnStop = 0f
+    private var consumeStop = true
 
     private val scrollConnection =
         PriorityNestedScrollConnection(
             orientation = Orientation.Vertical,
-            canStartPreScroll = { _, _ -> canStartPreScroll },
-            canStartPostScroll = { _, _ -> canStartPostScroll },
+            canStartPreScroll = { _, _, _ -> canStartPreScroll },
+            canStartPostScroll = { _, _, _ -> canStartPostScroll },
             canStartPostFling = { canStartPostFling },
-            canContinueScroll = { canContinueScroll },
-            canScrollOnFling = false,
+            canStopOnPreFling = { canStopOnPreFling },
             onStart = { isStarted = true },
-            onScroll = {
-                lastScroll = it
-                returnOnScroll
+            onScroll = { offsetAvailable, _ ->
+                lastScroll = offsetAvailable
+                if (consumeScroll) offsetAvailable else 0f
             },
             onStop = {
                 lastStop = it
-                { returnOnStop }
+                { if (consumeStop) it else 0f }
             },
         )
 
@@ -85,7 +84,7 @@ class PriorityNestedScrollConnectionTest {
         canStartPostScroll = true
         scrollConnection.onPostScroll(
             consumed = Offset.Zero,
-            available = Offset.Zero,
+            available = Offset(1f, 1f),
             source = UserInput,
         )
     }
@@ -136,21 +135,21 @@ class PriorityNestedScrollConnectionTest {
     @Test
     fun step2_onPriorityMode_shouldContinueIfAllowed() {
         startPriorityModePostScroll()
-        canContinueScroll = true
 
-        scrollConnection.onPreScroll(available = Offset(1f, 1f), source = UserInput)
+        val scroll1 = scrollConnection.onPreScroll(available = Offset(0f, 1f), source = UserInput)
         assertThat(lastScroll).isEqualTo(1f)
+        assertThat(scroll1.y).isEqualTo(1f)
 
-        canContinueScroll = false
-        scrollConnection.onPreScroll(available = Offset(2f, 2f), source = UserInput)
-        assertThat(lastScroll).isNotEqualTo(2f)
-        assertThat(lastScroll).isEqualTo(1f)
+        consumeScroll = false
+        val scroll2 = scrollConnection.onPreScroll(available = Offset(0f, 2f), source = UserInput)
+        assertThat(lastScroll).isEqualTo(2f)
+        assertThat(scroll2.y).isEqualTo(0f)
     }
 
     @Test
     fun step3a_onPriorityMode_shouldStopIfCannotContinue() {
         startPriorityModePostScroll()
-        canContinueScroll = false
+        consumeScroll = false
 
         scrollConnection.onPreScroll(available = Offset.Zero, source = UserInput)
 
@@ -160,17 +159,27 @@ class PriorityNestedScrollConnectionTest {
     @Test
     fun step3b_onPriorityMode_shouldStopOnFling() = runTest {
         startPriorityModePostScroll()
-        canContinueScroll = true
 
         scrollConnection.onPreFling(available = Velocity.Zero)
 
-        assertThat(lastStop).isNotNull()
+        assertThat(lastStop).isEqualTo(0f)
+    }
+
+    @Test
+    fun ifCannotStopOnPreFling_shouldStopOnPostFling() = runTest {
+        startPriorityModePostScroll()
+        canStopOnPreFling = false
+
+        scrollConnection.onPreFling(available = Velocity.Zero)
+        assertThat(lastStop).isNull()
+
+        scrollConnection.onPostFling(consumed = Velocity.Zero, available = Velocity.Zero)
+        assertThat(lastStop).isEqualTo(0f)
     }
 
     @Test
     fun step3c_onPriorityMode_shouldStopOnReset() {
         startPriorityModePostScroll()
-        canContinueScroll = true
 
         scrollConnection.reset()
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationScrimNestedScrollConnectionTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationScrimNestedScrollConnectionTest.kt
index 1f1680e0c273..75479ad35725 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationScrimNestedScrollConnectionTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationScrimNestedScrollConnectionTest.kt
@@ -31,6 +31,7 @@ import org.junit.runner.RunWith
 @RunWith(AndroidJUnit4::class)
 class NotificationScrimNestedScrollConnectionTest : SysuiTestCase() {
     private var isStarted = false
+    private var wasStarted = false
     private var scrimOffset = 0f
     private var contentHeight = 0f
     private var isCurrentGestureOverscroll = false
@@ -46,7 +47,10 @@ class NotificationScrimNestedScrollConnectionTest : SysuiTestCase() {
             minVisibleScrimHeight = { MIN_VISIBLE_SCRIM_HEIGHT },
             isCurrentGestureOverscroll = { isCurrentGestureOverscroll },
             onStart = { isStarted = true },
-            onStop = { isStarted = false },
+            onStop = {
+                wasStarted = true
+                isStarted = false
+            },
         )
 
     @Test
@@ -165,6 +169,7 @@ class NotificationScrimNestedScrollConnectionTest : SysuiTestCase() {
             )
 
         assertThat(offsetConsumed).isEqualTo(Offset.Zero)
+        assertThat(wasStarted).isEqualTo(false)
         assertThat(isStarted).isEqualTo(false)
     }
 
@@ -181,7 +186,9 @@ class NotificationScrimNestedScrollConnectionTest : SysuiTestCase() {
             )
 
         assertThat(offsetConsumed).isEqualTo(Offset.Zero)
-        assertThat(isStarted).isEqualTo(true)
+        // Returning 0 offset will immediately stop the connection
+        assertThat(wasStarted).isEqualTo(true)
+        assertThat(isStarted).isEqualTo(false)
     }
 
     @Test