Add ScreenshotActionsController between provider and viewModel

The asynchronicity of various processes in screenshots/action providers
means that we can show actions that are associated with the wrong
screenshot (see bug for details).

This change removes direct access from the ScreenshotActionProvider to
the ScreenshotViewModel, and instead abstracts out the action providers
to a ScreenshotActionsController that only updates the ViewModel if the
actions are associated with the current "live" screenshot.

Bug: 329659738
Bug: 341062985
Fix: 341062985
Flag: com.android.systemui.screenshot_shelf_ui2
Test: manual via taking several screenshots in succession
Test: atest ScreenshotActionsControllerTest
Merged-In: Ie18396172239b30218372229630da3d3d716406b
Change-Id: Ie18396172239b30218372229630da3d3d716406b

packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotActionsController.kt

+/*
+ * Copyright (C) 2024 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.screenshot
+
+import android.app.assist.AssistContent
+import com.android.systemui.screenshot.ui.viewmodel.ActionButtonAppearance
+import com.android.systemui.screenshot.ui.viewmodel.ScreenshotViewModel
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
+import java.util.UUID
+
+/**
+ * Responsible for obtaining the actions for each screenshot and sending them to the view model.
+ * Ensures that only actions from screenshots that are currently being shown are added to the view
+ * model.
+ */
+class ScreenshotActionsController
+@AssistedInject
+constructor(
+    private val viewModel: ScreenshotViewModel,
+    private val actionsProviderFactory: ScreenshotActionsProvider.Factory,
+    @Assisted val actionExecutor: ActionExecutor
+) {
+    private val actionProviders: MutableMap<UUID, ScreenshotActionsProvider> = mutableMapOf()
+    private var currentScreenshotId: UUID? = null
+
+    fun setCurrentScreenshot(screenshot: ScreenshotData): UUID {
+        val screenshotId = UUID.randomUUID()
+        currentScreenshotId = screenshotId
+        actionProviders[screenshotId] =
+            actionsProviderFactory.create(
+                screenshotId,
+                screenshot,
+                actionExecutor,
+                ActionsCallback(screenshotId),
+            )
+        return screenshotId
+    }
+
+    fun endScreenshotSession() {
+        currentScreenshotId = null
+    }
+
+    fun onAssistContent(screenshotId: UUID, assistContent: AssistContent?) {
+        actionProviders[screenshotId]?.onAssistContent(assistContent)
+    }
+
+    fun onScrollChipReady(screenshotId: UUID, onClick: Runnable) {
+        if (screenshotId == currentScreenshotId) {
+            actionProviders[screenshotId]?.onScrollChipReady(onClick)
+        }
+    }
+
+    fun onScrollChipInvalidated() {
+        for (provider in actionProviders.values) {
+            provider.onScrollChipInvalidated()
+        }
+    }
+
+    fun setCompletedScreenshot(screenshotId: UUID, result: ScreenshotSavedResult) {
+        if (screenshotId == currentScreenshotId) {
+            actionProviders[screenshotId]?.setCompletedScreenshot(result)
+        }
+    }
+
+    @AssistedFactory
+    interface Factory {
+        fun getController(actionExecutor: ActionExecutor): ScreenshotActionsController
+    }
+
+    inner class ActionsCallback(private val screenshotId: UUID) {
+        fun providePreviewAction(onClick: () -> Unit) {
+            if (screenshotId == currentScreenshotId) {
+                viewModel.setPreviewAction(onClick)
+            }
+        }
+
+        fun provideActionButton(
+            appearance: ActionButtonAppearance,
+            showDuringEntrance: Boolean,
+            onClick: () -> Unit
+        ): Int {
+            if (screenshotId == currentScreenshotId) {
+                return viewModel.addAction(appearance, showDuringEntrance, onClick)
+            }
+            return 0
+        }
+
+        fun updateActionButtonAppearance(buttonId: Int, appearance: ActionButtonAppearance) {
+            if (screenshotId == currentScreenshotId) {
+                viewModel.updateActionAppearance(buttonId, appearance)
+            }
+        }
+
+        fun updateActionButtonVisibility(buttonId: Int, visible: Boolean) {
+            if (screenshotId == currentScreenshotId) {
+                viewModel.setActionVisibility(buttonId, visible)
+            }
+        }
+    }
+}

packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotActionsProvider.kt

@@ -29,10 +29,10 @@ import com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_EDIT_TAPPED
 import com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_PREVIEW_TAPPED
 import com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_SHARE_TAPPED
 import com.android.systemui.screenshot.ui.viewmodel.ActionButtonAppearance
-import com.android.systemui.screenshot.ui.viewmodel.ScreenshotViewModel
 import dagger.assisted.Assisted
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
+import java.util.UUID
 
 /**
  * Provides actions for screenshots. This class can be overridden by a vendor-specific SysUI
@@ -51,9 +51,10 @@ interface ScreenshotActionsProvider {
 
     interface Factory {
         fun create(
+            requestId: UUID,
             request: ScreenshotData,
-            requestId: String,
             actionExecutor: ActionExecutor,
+            actionsCallback: ScreenshotActionsController.ActionsCallback,
         ): ScreenshotActionsProvider
     }
 }
@@ -62,11 +63,11 @@ class DefaultScreenshotActionsProvider
 @AssistedInject
 constructor(
     private val context: Context,
-    private val viewModel: ScreenshotViewModel,
     private val uiEventLogger: UiEventLogger,
+    @Assisted val requestId: UUID,
     @Assisted val request: ScreenshotData,
-    @Assisted val requestId: String,
     @Assisted val actionExecutor: ActionExecutor,
+    @Assisted val actionsCallback: ScreenshotActionsController.ActionsCallback,
 ) : ScreenshotActionsProvider {
     private var addedScrollChip = false
     private var onScrollClick: Runnable? = null
@@ -74,7 +75,7 @@ constructor(
     private var result: ScreenshotSavedResult? = null
 
     init {
-        viewModel.setPreviewAction {
+        actionsCallback.providePreviewAction {
             debugLog(LogConfig.DEBUG_ACTIONS) { "Preview tapped" }
             uiEventLogger.log(SCREENSHOT_PREVIEW_TAPPED, 0, request.packageNameString)
             onDeferrableActionTapped { result ->
@@ -85,40 +86,41 @@ constructor(
                 )
             }
         }
-        viewModel.addAction(
+
+        actionsCallback.provideActionButton(
             ActionButtonAppearance(
-                AppCompatResources.getDrawable(context, R.drawable.ic_screenshot_edit),
-                context.resources.getString(R.string.screenshot_edit_label),
-                context.resources.getString(R.string.screenshot_edit_description),
+                AppCompatResources.getDrawable(context, R.drawable.ic_screenshot_share),
+                context.resources.getString(R.string.screenshot_share_label),
+                context.resources.getString(R.string.screenshot_share_description),
             ),
             showDuringEntrance = true,
         ) {
-            debugLog(LogConfig.DEBUG_ACTIONS) { "Edit tapped" }
-            uiEventLogger.log(SCREENSHOT_EDIT_TAPPED, 0, request.packageNameString)
+            debugLog(LogConfig.DEBUG_ACTIONS) { "Share tapped" }
+            uiEventLogger.log(SCREENSHOT_SHARE_TAPPED, 0, request.packageNameString)
             onDeferrableActionTapped { result ->
                 actionExecutor.startSharedTransition(
-                    createEdit(result.uri, context),
+                    createShareWithSubject(result.uri, result.subject),
                     result.user,
-                    true
+                    false
                 )
             }
         }
 
-        viewModel.addAction(
+        actionsCallback.provideActionButton(
             ActionButtonAppearance(
-                AppCompatResources.getDrawable(context, R.drawable.ic_screenshot_share),
-                context.resources.getString(R.string.screenshot_share_label),
-                context.resources.getString(R.string.screenshot_share_description),
+                AppCompatResources.getDrawable(context, R.drawable.ic_screenshot_edit),
+                context.resources.getString(R.string.screenshot_edit_label),
+                context.resources.getString(R.string.screenshot_edit_description),
             ),
             showDuringEntrance = true,
         ) {
-            debugLog(LogConfig.DEBUG_ACTIONS) { "Share tapped" }
-            uiEventLogger.log(SCREENSHOT_SHARE_TAPPED, 0, request.packageNameString)
+            debugLog(LogConfig.DEBUG_ACTIONS) { "Edit tapped" }
+            uiEventLogger.log(SCREENSHOT_EDIT_TAPPED, 0, request.packageNameString)
             onDeferrableActionTapped { result ->
                 actionExecutor.startSharedTransition(
-                    createShareWithSubject(result.uri, result.subject),
+                    createEdit(result.uri, context),
                     result.user,
-                    false
+                    true
                 )
             }
         }
@@ -127,7 +129,7 @@ constructor(
     override fun onScrollChipReady(onClick: Runnable) {
         onScrollClick = onClick
         if (!addedScrollChip) {
-            viewModel.addAction(
+            actionsCallback.provideActionButton(
                 ActionButtonAppearance(
                     AppCompatResources.getDrawable(context, R.drawable.ic_screenshot_scroll),
                     context.resources.getString(R.string.screenshot_scroll_label),
@@ -161,9 +163,10 @@ constructor(
     @AssistedFactory
     interface Factory : ScreenshotActionsProvider.Factory {
         override fun create(
+            requestId: UUID,
             request: ScreenshotData,
-            requestId: String,
             actionExecutor: ActionExecutor,
+            actionsCallback: ScreenshotActionsController.ActionsCallback,
         ): DefaultScreenshotActionsProvider
     }
 

packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java

@@ -192,7 +192,6 @@ public class ScreenshotController {
     private final WindowContext mContext;
     private final FeatureFlags mFlags;
     private final ScreenshotViewProxy mViewProxy;
-    private final ScreenshotActionsProvider.Factory mActionsProviderFactory;
     private final ScreenshotNotificationsController mNotificationsController;
     private final ScreenshotSmartActions mScreenshotSmartActions;
     private final UiEventLogger mUiEventLogger;
@@ -202,7 +201,7 @@ public class ScreenshotController {
     private final ExecutorService mBgExecutor;
     private final BroadcastSender mBroadcastSender;
     private final BroadcastDispatcher mBroadcastDispatcher;
-    private final ActionExecutor mActionExecutor;
+    private final ScreenshotActionsController mActionsController;
 
     private final WindowManager mWindowManager;
     private final WindowManager.LayoutParams mWindowLayoutParams;
@@ -217,6 +216,8 @@ public class ScreenshotController {
     private final ActionIntentExecutor mActionIntentExecutor;
     private final UserManager mUserManager;
     private final AssistContentRequester mAssistContentRequester;
+    private final ActionExecutor mActionExecutor;
+
 
     private final MessageContainerController mMessageContainerController;
     private Bitmap mScreenBitmap;
@@ -225,7 +226,6 @@ public class ScreenshotController {
     private boolean mBlockAttach;
     private Animator mScreenshotAnimation;
     private RequestCallback mCurrentRequestCallback;
-    private ScreenshotActionsProvider mActionsProvider;
     private String mPackageName = "";
     private final BroadcastReceiver mCopyBroadcastReceiver;
 
@@ -252,7 +252,6 @@ public class ScreenshotController {
             WindowManager windowManager,
             FeatureFlags flags,
             ScreenshotViewProxy.Factory viewProxyFactory,
-            ScreenshotActionsProvider.Factory actionsProviderFactory,
             ScreenshotSmartActions screenshotSmartActions,
             ScreenshotNotificationsController.Factory screenshotNotificationsControllerFactory,
             UiEventLogger uiEventLogger,
@@ -264,6 +263,7 @@ public class ScreenshotController {
             BroadcastSender broadcastSender,
             BroadcastDispatcher broadcastDispatcher,
             ScreenshotNotificationSmartActionsProvider screenshotNotificationSmartActionsProvider,
+            ScreenshotActionsController.Factory screenshotActionsControllerFactory,
             ActionIntentExecutor actionIntentExecutor,
             ActionExecutor.Factory actionExecutorFactory,
             UserManager userManager,
@@ -275,7 +275,6 @@ public class ScreenshotController {
     ) {
         mScreenshotSmartActions = screenshotSmartActions;
         mWindowManager = windowManager;
-        mActionsProviderFactory = actionsProviderFactory;
         mNotificationsController = screenshotNotificationsControllerFactory.create(
                 display.getDisplayId());
         mUiEventLogger = uiEventLogger;
@@ -323,6 +322,8 @@ public class ScreenshotController {
                     finishDismiss();
                     return Unit.INSTANCE;
                 });
+        mActionsController = screenshotActionsControllerFactory.getController(mActionExecutor);
+
 
         // Sound is only reproduced from the controller of the default display.
         if (display.getDisplayId() == Display.DEFAULT_DISPLAY) {
@@ -399,20 +400,21 @@ public class ScreenshotController {
             return;
         }
 
+        final UUID requestId;
         if (screenshotShelfUi2()) {
-            final UUID requestId = UUID.randomUUID();
-            final String screenshotId = String.format("Screenshot_%s", requestId);
-            mActionsProvider = mActionsProviderFactory.create(
-                    screenshot, screenshotId, mActionExecutor);
+            requestId = mActionsController.setCurrentScreenshot(screenshot);
             saveScreenshotInBackground(screenshot, requestId, finisher);
 
             if (screenshot.getTaskId() >= 0) {
-                mAssistContentRequester.requestAssistContent(screenshot.getTaskId(),
-                        assistContent -> mActionsProvider.onAssistContent(assistContent));
+                mAssistContentRequester.requestAssistContent(
+                        screenshot.getTaskId(),
+                        assistContent ->
+                                mActionsController.onAssistContent(requestId, assistContent));
             } else {
-                mActionsProvider.onAssistContent(null);
+                mActionsController.onAssistContent(requestId, null);
             }
         } else {
+            requestId = UUID.randomUUID(); // passed through but unused for legacy UI
             saveScreenshotInWorkerThread(screenshot.getUserHandle(), finisher,
                     this::showUiOnActionsReady, this::showUiOnQuickShareActionReady);
         }
@@ -421,7 +423,7 @@ public class ScreenshotController {
         setWindowFocusable(true);
         mViewProxy.requestFocus();
 
-        enqueueScrollCaptureRequest(screenshot.getUserHandle());
+        enqueueScrollCaptureRequest(requestId, screenshot.getUserHandle());
 
         attachWindow();
 
@@ -574,11 +576,11 @@ public class ScreenshotController {
         mWindow.setContentView(mViewProxy.getView());
     }
 
-    private void enqueueScrollCaptureRequest(UserHandle owner) {
+    private void enqueueScrollCaptureRequest(UUID requestId, UserHandle owner) {
         // Wait until this window is attached to request because it is
         // the reference used to locate the target window (below).
         withWindowAttached(() -> {
-            requestScrollCapture(owner);
+            requestScrollCapture(requestId, owner);
             mWindow.peekDecorView().getViewRootImpl().setActivityConfigCallback(
                     new ViewRootImpl.ActivityConfigCallback() {
                         @Override
@@ -588,14 +590,14 @@ public class ScreenshotController {
                                 // Hide the scroll chip until we know it's available in this
                                 // orientation
                                 if (screenshotShelfUi2()) {
-                                    mActionsProvider.onScrollChipInvalidated();
+                                    mActionsController.onScrollChipInvalidated();
                                 } else {
                                     mViewProxy.hideScrollChip();
                                 }
                                 // Delay scroll capture eval a bit to allow the underlying activity
                                 // to set up in the new orientation.
                                 mScreenshotHandler.postDelayed(
-                                        () -> requestScrollCapture(owner), 150);
+                                        () -> requestScrollCapture(requestId, owner), 150);
                                 mViewProxy.updateInsets(
                                         mWindowManager.getCurrentWindowMetrics().getWindowInsets());
                                 // Screenshot animation calculations won't be valid anymore,
@@ -617,7 +619,7 @@ public class ScreenshotController {
         });
     }
 
-    private void requestScrollCapture(UserHandle owner) {
+    private void requestScrollCapture(UUID requestId, UserHandle owner) {
         mScrollCaptureExecutor.requestScrollCapture(
                 mDisplay.getDisplayId(),
                 mWindow.getDecorView().getWindowToken(),
@@ -625,10 +627,8 @@ public class ScreenshotController {
                     mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_LONG_SCREENSHOT_IMPRESSION,
                             0, response.getPackageName());
                     if (screenshotShelfUi2()) {
-                        if (mActionsProvider != null) {
-                            mActionsProvider.onScrollChipReady(
-                                    () -> onScrollButtonClicked(owner, response));
-                        }
+                        mActionsController.onScrollChipReady(requestId,
+                                () -> onScrollButtonClicked(owner, response));
                     } else {
                         mViewProxy.showScrollChip(response.getPackageName(),
                                 () -> onScrollButtonClicked(owner, response));
@@ -814,6 +814,7 @@ public class ScreenshotController {
     /** Reset screenshot view and then call onCompleteRunnable */
     private void finishDismiss() {
         Log.d(TAG, "finishDismiss");
+        mActionsController.endScreenshotSession();
         mScrollCaptureExecutor.close();
         if (mCurrentRequestCallback != null) {
             mCurrentRequestCallback.onFinish();
@@ -834,9 +835,8 @@ public class ScreenshotController {
                 ImageExporter.Result result = future.get();
                 Log.d(TAG, "Saved screenshot: " + result);
                 logScreenshotResultStatus(result.uri, screenshot.getUserHandle());
-                mScreenshotHandler.resetTimeout();
                 if (result.uri != null) {
-                    mActionsProvider.setCompletedScreenshot(new ScreenshotSavedResult(
+                    mActionsController.setCompletedScreenshot(requestId, new ScreenshotSavedResult(
                             result.uri, screenshot.getUserOrDefault(), result.timestamp));
                 }
                 if (DEBUG_CALLBACK) {

packages/SystemUI/tests/src/com/android/systemui/screenshot/DefaultScreenshotActionsProviderTest.kt

@@ -21,41 +21,38 @@ import android.net.Uri
 import android.os.Process
 import android.os.UserHandle
 import android.testing.AndroidTestingRunner
-import android.view.accessibility.AccessibilityManager
 import androidx.test.filters.SmallTest
 import com.android.internal.logging.UiEventLogger
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.screenshot.ui.viewmodel.ScreenshotViewModel
-import com.android.systemui.util.mockito.argumentCaptor
-import com.android.systemui.util.mockito.capture
-import com.android.systemui.util.mockito.eq
-import com.android.systemui.util.mockito.mock
 import com.google.common.truth.Truth.assertThat
+import java.util.UUID
 import kotlin.test.Test
 import kotlinx.coroutines.test.runTest
-import org.junit.Assert.assertNotNull
 import org.junit.Before
 import org.junit.runner.RunWith
 import org.mockito.Mockito.verifyNoMoreInteractions
+import org.mockito.kotlin.any
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
 import org.mockito.kotlin.never
+import org.mockito.kotlin.times
 import org.mockito.kotlin.verify
 
 @RunWith(AndroidTestingRunner::class)
 @SmallTest
 class DefaultScreenshotActionsProviderTest : SysuiTestCase() {
     private val actionExecutor = mock<ActionExecutor>()
-    private val accessibilityManager = mock<AccessibilityManager>()
     private val uiEventLogger = mock<UiEventLogger>()
+    private val actionsCallback = mock<ScreenshotActionsController.ActionsCallback>()
 
     private val request = ScreenshotData.forTesting()
     private val validResult = ScreenshotSavedResult(Uri.EMPTY, Process.myUserHandle(), 0)
 
-    private lateinit var viewModel: ScreenshotViewModel
     private lateinit var actionsProvider: ScreenshotActionsProvider
 
     @Before
     fun setUp() {
-        viewModel = ScreenshotViewModel(accessibilityManager)
         request.userHandle = UserHandle.OWNER
     }
 
@@ -63,8 +60,9 @@ class DefaultScreenshotActionsProviderTest : SysuiTestCase() {
     fun previewActionAccessed_beforeScreenshotCompleted_doesNothing() {
         actionsProvider = createActionsProvider()
 
-        assertNotNull(viewModel.previewAction.value)
-        viewModel.previewAction.value!!.invoke()
+        val previewActionCaptor = argumentCaptor<() -> Unit>()
+        verify(actionsCallback).providePreviewAction(previewActionCaptor.capture())
+        previewActionCaptor.firstValue.invoke()
         verifyNoMoreInteractions(actionExecutor)
     }
 
@@ -72,13 +70,13 @@ class DefaultScreenshotActionsProviderTest : SysuiTestCase() {
     fun actionButtonsAccessed_beforeScreenshotCompleted_doesNothing() {
         actionsProvider = createActionsProvider()
 
-        assertThat(viewModel.actions.value.size).isEqualTo(2)
-        val firstAction = viewModel.actions.value[0]
-        assertThat(firstAction.onClicked).isNotNull()
-        val secondAction = viewModel.actions.value[1]
-        assertThat(secondAction.onClicked).isNotNull()
-        firstAction.onClicked!!.invoke()
-        secondAction.onClicked!!.invoke()
+        val actionButtonCaptor = argumentCaptor<() -> Unit>()
+        verify(actionsCallback, times(2))
+            .provideActionButton(any(), any(), actionButtonCaptor.capture())
+        val firstAction = actionButtonCaptor.firstValue
+        val secondAction = actionButtonCaptor.secondValue
+        firstAction.invoke()
+        secondAction.invoke()
         verifyNoMoreInteractions(actionExecutor)
     }
 
@@ -87,29 +85,39 @@ class DefaultScreenshotActionsProviderTest : SysuiTestCase() {
         actionsProvider = createActionsProvider()
 
         actionsProvider.setCompletedScreenshot(validResult)
-        viewModel.actions.value[0].onClicked!!.invoke()
 
-        verify(uiEventLogger).log(eq(ScreenshotEvent.SCREENSHOT_EDIT_TAPPED), eq(0), eq(""))
+        val actionButtonCaptor = argumentCaptor<() -> Unit>()
+        verify(actionsCallback, times(2))
+            .provideActionButton(any(), any(), actionButtonCaptor.capture())
+        actionButtonCaptor.firstValue.invoke()
+
+        verify(uiEventLogger).log(eq(ScreenshotEvent.SCREENSHOT_SHARE_TAPPED), eq(0), eq(""))
         val intentCaptor = argumentCaptor<Intent>()
         verify(actionExecutor)
-            .startSharedTransition(capture(intentCaptor), eq(Process.myUserHandle()), eq(true))
-        assertThat(intentCaptor.value.action).isEqualTo(Intent.ACTION_EDIT)
+            .startSharedTransition(intentCaptor.capture(), eq(Process.myUserHandle()), eq(false))
+        assertThat(intentCaptor.firstValue.action).isEqualTo(Intent.ACTION_CHOOSER)
     }
 
     @Test
     fun actionAccessed_whilePending_launchesMostRecentAction() = runTest {
         actionsProvider = createActionsProvider()
 
-        viewModel.actions.value[0].onClicked!!.invoke()
-        viewModel.previewAction.value!!.invoke()
-        viewModel.actions.value[1].onClicked!!.invoke()
+        val previewActionCaptor = argumentCaptor<() -> Unit>()
+        verify(actionsCallback).providePreviewAction(previewActionCaptor.capture())
+        val actionButtonCaptor = argumentCaptor<() -> Unit>()
+        verify(actionsCallback, times(2))
+            .provideActionButton(any(), any(), actionButtonCaptor.capture())
+
+        actionButtonCaptor.firstValue.invoke()
+        previewActionCaptor.firstValue.invoke()
+        actionButtonCaptor.secondValue.invoke()
         actionsProvider.setCompletedScreenshot(validResult)
 
-        verify(uiEventLogger).log(eq(ScreenshotEvent.SCREENSHOT_SHARE_TAPPED), eq(0), eq(""))
+        verify(uiEventLogger).log(eq(ScreenshotEvent.SCREENSHOT_EDIT_TAPPED), eq(0), eq(""))
         val intentCaptor = argumentCaptor<Intent>()
         verify(actionExecutor)
-            .startSharedTransition(capture(intentCaptor), eq(Process.myUserHandle()), eq(false))
-        assertThat(intentCaptor.value.action).isEqualTo(Intent.ACTION_CHOOSER)
+            .startSharedTransition(intentCaptor.capture(), eq(Process.myUserHandle()), eq(true))
+        assertThat(intentCaptor.firstValue.action).isEqualTo(Intent.ACTION_EDIT)
     }
 
     @Test
@@ -117,9 +125,12 @@ class DefaultScreenshotActionsProviderTest : SysuiTestCase() {
         actionsProvider = createActionsProvider()
 
         val onScrollClick = mock<Runnable>()
-        val numActions = viewModel.actions.value.size
         actionsProvider.onScrollChipReady(onScrollClick)
-        viewModel.actions.value[numActions].onClicked!!.invoke()
+        val actionButtonCaptor = argumentCaptor<() -> Unit>()
+        // share, edit, scroll
+        verify(actionsCallback, times(3))
+            .provideActionButton(any(), any(), actionButtonCaptor.capture())
+        actionButtonCaptor.thirdValue.invoke()
 
         verify(onScrollClick).run()
     }
@@ -129,10 +140,13 @@ class DefaultScreenshotActionsProviderTest : SysuiTestCase() {
         actionsProvider = createActionsProvider()
 
         val onScrollClick = mock<Runnable>()
-        val numActions = viewModel.actions.value.size
         actionsProvider.onScrollChipReady(onScrollClick)
+        val actionButtonCaptor = argumentCaptor<() -> Unit>()
         actionsProvider.onScrollChipInvalidated()
-        viewModel.actions.value[numActions].onClicked!!.invoke()
+        // share, edit, scroll
+        verify(actionsCallback, times(3))
+            .provideActionButton(any(), any(), actionButtonCaptor.capture())
+        actionButtonCaptor.thirdValue.invoke()
 
         verify(onScrollClick, never()).run()
     }
@@ -143,11 +157,15 @@ class DefaultScreenshotActionsProviderTest : SysuiTestCase() {
 
         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()
+        val actionButtonCaptor = argumentCaptor<() -> Unit>()
+        // share, edit, scroll
+        verify(actionsCallback, times(3))
+            .provideActionButton(any(), any(), actionButtonCaptor.capture())
+        actionButtonCaptor.thirdValue.invoke()
 
         verify(onScrollClick2).run()
         verify(onScrollClick, never()).run()
@@ -156,11 +174,11 @@ class DefaultScreenshotActionsProviderTest : SysuiTestCase() {
     private fun createActionsProvider(): ScreenshotActionsProvider {
         return DefaultScreenshotActionsProvider(
             context,
-            viewModel,
             uiEventLogger,
+            UUID.randomUUID(),
             request,
-            "testid",
-            actionExecutor
+            actionExecutor,
+            actionsCallback,
         )
     }
 }

packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotActionsControllerTest.kt

+/*
+ * Copyright (C) 2024 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.screenshot
+
+import android.testing.AndroidTestingRunner
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.screenshot.ui.viewmodel.ScreenshotViewModel
+import java.util.UUID
+import kotlin.test.Test
+import org.junit.Before
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.never
+import org.mockito.kotlin.verify
+
+@RunWith(AndroidTestingRunner::class)
+@SmallTest
+class ScreenshotActionsControllerTest : SysuiTestCase() {
+    private val screenshotData = mock<ScreenshotData>()
+    private val actionExecutor = mock<ActionExecutor>()
+    private val viewModel = mock<ScreenshotViewModel>()
+    private val onClick = mock<() -> Unit>()
+
+    private lateinit var actionsController: ScreenshotActionsController
+    private lateinit var fakeActionsProvider1: FakeActionsProvider
+    private lateinit var fakeActionsProvider2: FakeActionsProvider
+    private val actionsProviderFactory =
+        object : ScreenshotActionsProvider.Factory {
+            var isFirstCall = true
+            override fun create(
+                requestId: UUID,
+                request: ScreenshotData,
+                actionExecutor: ActionExecutor,
+                actionsCallback: ScreenshotActionsController.ActionsCallback
+            ): ScreenshotActionsProvider {
+                return if (isFirstCall) {
+                    isFirstCall = false
+                    fakeActionsProvider1 = FakeActionsProvider(actionsCallback)
+                    fakeActionsProvider1
+                } else {
+                    fakeActionsProvider2 = FakeActionsProvider(actionsCallback)
+                    fakeActionsProvider2
+                }
+            }
+        }
+
+    @Before
+    fun setUp() {
+        actionsController =
+            ScreenshotActionsController(viewModel, actionsProviderFactory, actionExecutor)
+    }
+
+    @Test
+    fun setPreview_onCurrentScreenshot_updatesViewModel() {
+        actionsController.setCurrentScreenshot(screenshotData)
+        fakeActionsProvider1.callPreview(onClick)
+
+        verify(viewModel).setPreviewAction(onClick)
+    }
+
+    @Test
+    fun setPreview_onNonCurrentScreenshot_doesNotUpdateViewModel() {
+        actionsController.setCurrentScreenshot(screenshotData)
+        actionsController.setCurrentScreenshot(screenshotData)
+        fakeActionsProvider1.callPreview(onClick)
+
+        verify(viewModel, never()).setPreviewAction(any())
+    }
+
+    class FakeActionsProvider(
+        private val actionsCallback: ScreenshotActionsController.ActionsCallback
+    ) : ScreenshotActionsProvider {
+
+        fun callPreview(onClick: () -> Unit) {
+            actionsCallback.providePreviewAction(onClick)
+        }
+
+        override fun onScrollChipReady(onClick: Runnable) {}
+
+        override fun onScrollChipInvalidated() {}
+
+        override fun setCompletedScreenshot(result: ScreenshotSavedResult) {}
+    }
+}