Merge "Merge "Update scroll chip on rotation" into 24D1-dev am: e577313649" into main
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotActionsProvider.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotActionsProvider.kt
index ef1d87d..a1dd415 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotActionsProvider.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotActionsProvider.kt
@@ -40,6 +40,7 @@
  */
 interface ScreenshotActionsProvider {
     fun onScrollChipReady(onClick: Runnable)
+    fun onScrollChipInvalidated()
     fun setCompletedScreenshot(result: ScreenshotSavedResult)
 
     /**
@@ -67,6 +68,8 @@
     @Assisted val requestId: String,
     @Assisted val actionExecutor: ActionExecutor,
 ) : ScreenshotActionsProvider {
+    private var addedScrollChip = false
+    private var onScrollClick: Runnable? = null
     private var pendingAction: ((ScreenshotSavedResult) -> Unit)? = null
     private var result: ScreenshotSavedResult? = null
 
@@ -122,18 +125,26 @@
     }
 
     override fun onScrollChipReady(onClick: Runnable) {
-        viewModel.addAction(
-            ActionButtonAppearance(
-                AppCompatResources.getDrawable(context, R.drawable.ic_screenshot_scroll),
-                context.resources.getString(R.string.screenshot_scroll_label),
-                context.resources.getString(R.string.screenshot_scroll_label),
-            ),
-            showDuringEntrance = true,
-        ) {
-            onClick.run()
+        onScrollClick = onClick
+        if (!addedScrollChip) {
+            viewModel.addAction(
+                ActionButtonAppearance(
+                    AppCompatResources.getDrawable(context, R.drawable.ic_screenshot_scroll),
+                    context.resources.getString(R.string.screenshot_scroll_label),
+                    context.resources.getString(R.string.screenshot_scroll_label),
+                ),
+                showDuringEntrance = true,
+            ) {
+                onScrollClick?.run()
+            }
+            addedScrollChip = true
         }
     }
 
+    override fun onScrollChipInvalidated() {
+        onScrollClick = null
+    }
+
     override fun setCompletedScreenshot(result: ScreenshotSavedResult) {
         if (this.result != null) {
             Log.e(TAG, "Got a second completed screenshot for existing request!")
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
index bd90de2..2f026ae 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
@@ -586,7 +586,11 @@
                             if (mConfigChanges.applyNewConfig(mContext.getResources())) {
                                 // Hide the scroll chip until we know it's available in this
                                 // orientation
-                                mViewProxy.hideScrollChip();
+                                if (screenshotShelfUi2()) {
+                                    mActionsProvider.onScrollChipInvalidated();
+                                } else {
+                                    mViewProxy.hideScrollChip();
+                                }
                                 // Delay scroll capture eval a bit to allow the underlying activity
                                 // to set up in the new orientation.
                                 mScreenshotHandler.postDelayed(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/DefaultScreenshotActionsProviderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/DefaultScreenshotActionsProviderTest.kt
index 853e50a..896c3bf 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/DefaultScreenshotActionsProviderTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/DefaultScreenshotActionsProviderTest.kt
@@ -37,6 +37,7 @@
 import org.junit.Before
 import org.junit.runner.RunWith
 import org.mockito.Mockito.verifyNoMoreInteractions
+import org.mockito.kotlin.never
 import org.mockito.kotlin.verify
 
 @RunWith(AndroidTestingRunner::class)
@@ -111,6 +112,47 @@
         assertThat(intentCaptor.value.action).isEqualTo(Intent.ACTION_CHOOSER)
     }
 
+    @Test
+    fun scrollChipClicked_callsOnClick() = runTest {
+        actionsProvider = createActionsProvider()
+
+        val onScrollClick = mock<Runnable>()
+        val numActions = viewModel.actions.value.size
+        actionsProvider.onScrollChipReady(onScrollClick)
+        viewModel.actions.value[numActions].onClicked!!.invoke()
+
+        verify(onScrollClick).run()
+    }
+
+    @Test
+    fun scrollChipClicked_afterInvalidate_doesNothing() = runTest {
+        actionsProvider = createActionsProvider()
+
+        val onScrollClick = mock<Runnable>()
+        val numActions = viewModel.actions.value.size
+        actionsProvider.onScrollChipReady(onScrollClick)
+        actionsProvider.onScrollChipInvalidated()
+        viewModel.actions.value[numActions].onClicked!!.invoke()
+
+        verify(onScrollClick, never()).run()
+    }
+
+    @Test
+    fun scrollChipClicked_afterUpdate_runsNewAction() = runTest {
+        actionsProvider = createActionsProvider()
+
+        val onScrollClick = mock<Runnable>()
+        val onScrollClick2 = mock<Runnable>()
+        val numActions = viewModel.actions.value.size
+        actionsProvider.onScrollChipReady(onScrollClick)
+        actionsProvider.onScrollChipInvalidated()
+        actionsProvider.onScrollChipReady(onScrollClick2)
+        viewModel.actions.value[numActions].onClicked!!.invoke()
+
+        verify(onScrollClick2).run()
+        verify(onScrollClick, never()).run()
+    }
+
     private fun createActionsProvider(): ScreenshotActionsProvider {
         return DefaultScreenshotActionsProvider(
             context,