Merge changes I027cc96c,I44bb8191 into tm-qpr-dev

* changes:
  Add a TestableAlertDialog
  Pass Secure.LOCKSCREEN_ALLOW_TRIVIAL_CONTROLS

core/java/android/service/controls/ControlsProviderService.java

@@ -68,6 +68,18 @@ public abstract class ControlsProviderService extends Service {
     public static final String META_DATA_PANEL_ACTIVITY =
             "android.service.controls.META_DATA_PANEL_ACTIVITY";
 
+    /**
+     * Boolean extra containing the value of
+     * {@link android.provider.Settings.Secure#LOCKSCREEN_ALLOW_TRIVIAL_CONTROLS}.
+     *
+     * This is passed with the intent when the panel specified by {@link #META_DATA_PANEL_ACTIVITY}
+     * is launched.
+     *
+     * @hide
+     */
+    public static final String EXTRA_LOCKSCREEN_ALLOW_TRIVIAL_CONTROLS =
+            "android.service.controls.EXTRA_LOCKSCREEN_ALLOW_TRIVIAL_CONTROLS";
+
     /**
      * @hide
      */

packages/SystemUI/src/com/android/systemui/controls/ui/ControlsUiControllerImpl.kt

@@ -28,6 +28,7 @@ import android.content.Intent
 import android.graphics.drawable.Drawable
 import android.graphics.drawable.LayerDrawable
 import android.service.controls.Control
+import android.service.controls.ControlsProviderService
 import android.util.Log
 import android.view.ContextThemeWrapper
 import android.view.LayoutInflater
@@ -48,6 +49,7 @@ import com.android.systemui.Dumpable
 import com.android.systemui.R
 import com.android.systemui.controls.ControlsMetricsLogger
 import com.android.systemui.controls.ControlsServiceInfo
+import com.android.systemui.controls.ControlsSettingsRepository
 import com.android.systemui.controls.CustomIconCache
 import com.android.systemui.controls.controller.ControlsController
 import com.android.systemui.controls.controller.StructureInfo
@@ -96,6 +98,7 @@ class ControlsUiControllerImpl @Inject constructor (
         private val userFileManager: UserFileManager,
         private val userTracker: UserTracker,
         private val taskViewFactory: Optional<TaskViewFactory>,
+        private val controlsSettingsRepository: ControlsSettingsRepository,
         dumpManager: DumpManager
 ) : ControlsUiController, Dumpable {
 
@@ -354,7 +357,6 @@ class ControlsUiControllerImpl @Inject constructor (
                 } else {
                     items[0]
                 }
-
         maybeUpdateSelectedItem(selectionItem)
 
         createControlsSpaceFrame()
@@ -374,11 +376,20 @@ class ControlsUiControllerImpl @Inject constructor (
     }
 
     private fun createPanelView(componentName: ComponentName) {
-        val pendingIntent = PendingIntent.getActivity(
+        val setting = controlsSettingsRepository
+                .allowActionOnTrivialControlsInLockscreen.value
+        val pendingIntent = PendingIntent.getActivityAsUser(
                 context,
                 0,
-                Intent().setComponent(componentName),
+                Intent()
-                PendingIntent.FLAG_IMMUTABLE
+                        .setComponent(componentName)
+                        .putExtra(
+                                ControlsProviderService.EXTRA_LOCKSCREEN_ALLOW_TRIVIAL_CONTROLS,
+                                setting
+                        ),
+                PendingIntent.FLAG_IMMUTABLE or PendingIntent.FLAG_UPDATE_CURRENT,
+                null,
+                userTracker.userHandle
         )
 
         parent.requireViewById<View>(R.id.controls_scroll_view).visibility = View.GONE
@@ -698,6 +709,8 @@ class ControlsUiControllerImpl @Inject constructor (
             println("hidden: $hidden")
             println("selectedItem: $selectedItem")
             println("lastSelections: $lastSelections")
+            println("setting: ${controlsSettingsRepository
+                    .allowActionOnTrivialControlsInLockscreen.value}")
         }
     }
 }

packages/SystemUI/tests/src/com/android/systemui/controls/ui/ControlsUiControllerImplTest.kt

@@ -16,15 +16,26 @@
 
 package com.android.systemui.controls.ui
 
+import android.app.PendingIntent
 import android.content.ComponentName
 import android.content.Context
+import android.content.pm.ApplicationInfo
+import android.content.pm.ServiceInfo
+import android.os.UserHandle
+import android.service.controls.ControlsProviderService
 import android.testing.AndroidTestingRunner
 import android.testing.TestableLooper
+import android.util.AttributeSet
+import android.view.LayoutInflater
+import android.view.View
 import android.widget.FrameLayout
 import androidx.test.filters.SmallTest
+import com.android.systemui.R
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.controls.ControlsMetricsLogger
+import com.android.systemui.controls.ControlsServiceInfo
 import com.android.systemui.controls.CustomIconCache
+import com.android.systemui.controls.FakeControlsSettingsRepository
 import com.android.systemui.controls.controller.ControlsController
 import com.android.systemui.controls.controller.StructureInfo
 import com.android.systemui.controls.management.ControlsListingController
@@ -38,19 +49,26 @@ import com.android.systemui.statusbar.policy.KeyguardStateController
 import com.android.systemui.util.FakeSharedPreferences
 import com.android.systemui.util.concurrency.FakeExecutor
 import com.android.systemui.util.mockito.any
+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.android.systemui.util.time.FakeSystemClock
+import com.android.wm.shell.TaskView
 import com.android.wm.shell.TaskViewFactory
 import com.google.common.truth.Truth.assertThat
 import dagger.Lazy
 import java.util.Optional
+import java.util.function.Consumer
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mock
 import org.mockito.Mockito.anyInt
 import org.mockito.Mockito.anyString
-import org.mockito.Mockito.mock
+import org.mockito.Mockito.clearInvocations
 import org.mockito.Mockito.never
+import org.mockito.Mockito.spy
 import org.mockito.Mockito.verify
 import org.mockito.Mockito.`when`
 import org.mockito.MockitoAnnotations
@@ -70,9 +88,9 @@ class ControlsUiControllerImplTest : SysuiTestCase() {
     @Mock lateinit var userFileManager: UserFileManager
     @Mock lateinit var userTracker: UserTracker
     @Mock lateinit var taskViewFactory: TaskViewFactory
-    @Mock lateinit var activityContext: Context
     @Mock lateinit var dumpManager: DumpManager
     val sharedPreferences = FakeSharedPreferences()
+    lateinit var controlsSettingsRepository: FakeControlsSettingsRepository
 
     var uiExecutor = FakeExecutor(FakeSystemClock())
     var bgExecutor = FakeExecutor(FakeSystemClock())
@@ -83,6 +101,17 @@ class ControlsUiControllerImplTest : SysuiTestCase() {
     fun setup() {
         MockitoAnnotations.initMocks(this)
 
+        controlsSettingsRepository = FakeControlsSettingsRepository()
+
+        // This way, it won't be cloned every time `LayoutInflater.fromContext` is called, but we
+        // need to clone it once so we don't modify the original one.
+        mContext.addMockSystemService(
+            Context.LAYOUT_INFLATER_SERVICE,
+            mContext.baseContext
+                .getSystemService(LayoutInflater::class.java)!!
+                .cloneInContext(mContext)
+        )
+
         parent = FrameLayout(mContext)
 
         underTest =
@@ -100,6 +129,7 @@ class ControlsUiControllerImplTest : SysuiTestCase() {
                 userFileManager,
                 userTracker,
                 Optional.of(taskViewFactory),
+                controlsSettingsRepository,
                 dumpManager
             )
         `when`(
@@ -113,11 +143,12 @@ class ControlsUiControllerImplTest : SysuiTestCase() {
         `when`(userFileManager.getSharedPreferences(anyString(), anyInt(), anyInt()))
             .thenReturn(sharedPreferences)
         `when`(userTracker.userId).thenReturn(0)
+        `when`(userTracker.userHandle).thenReturn(UserHandle.of(0))
     }
 
     @Test
     fun testGetPreferredStructure() {
-        val structureInfo = mock(StructureInfo::class.java)
+        val structureInfo = mock<StructureInfo>()
         underTest.getPreferredSelectedItem(listOf(structureInfo))
         verify(userFileManager)
             .getSharedPreferences(
@@ -189,14 +220,195 @@ class ControlsUiControllerImplTest : SysuiTestCase() {
     @Test
     fun testPanelDoesNotRefreshControls() {
         val panel = SelectedItem.PanelItem("App name", ComponentName("pkg", "cls"))
+        setUpPanel(panel)
+
+        underTest.show(parent, {}, context)
+        verify(controlsController, never()).refreshStatus(any(), any())
+    }
+
+    @Test
+    fun testPanelCallsTaskViewFactoryCreate() {
+        mockLayoutInflater()
+        val panel = SelectedItem.PanelItem("App name", ComponentName("pkg", "cls"))
+        val serviceInfo = setUpPanel(panel)
+
+        underTest.show(parent, {}, context)
+
+        val captor = argumentCaptor<ControlsListingController.ControlsListingCallback>()
+
+        verify(controlsListingController).addCallback(capture(captor))
+
+        captor.value.onServicesUpdated(listOf(serviceInfo))
+        FakeExecutor.exhaustExecutors(uiExecutor, bgExecutor)
+
+        verify(taskViewFactory).create(eq(context), eq(uiExecutor), any())
+    }
+
+    @Test
+    fun testPanelControllerStartActivityWithCorrectArguments() {
+        mockLayoutInflater()
+        controlsSettingsRepository.setAllowActionOnTrivialControlsInLockscreen(true)
+
+        val panel = SelectedItem.PanelItem("App name", ComponentName("pkg", "cls"))
+        val serviceInfo = setUpPanel(panel)
+
+        underTest.show(parent, {}, context)
+
+        val captor = argumentCaptor<ControlsListingController.ControlsListingCallback>()
+
+        verify(controlsListingController).addCallback(capture(captor))
+
+        captor.value.onServicesUpdated(listOf(serviceInfo))
+        FakeExecutor.exhaustExecutors(uiExecutor, bgExecutor)
+
+        val pendingIntent = verifyPanelCreatedAndStartTaskView()
+
+        with(pendingIntent) {
+            assertThat(isActivity).isTrue()
+            assertThat(intent.component).isEqualTo(serviceInfo.panelActivity)
+            assertThat(
+                    intent.getBooleanExtra(
+                        ControlsProviderService.EXTRA_LOCKSCREEN_ALLOW_TRIVIAL_CONTROLS,
+                        false
+                    )
+                )
+                .isTrue()
+        }
+    }
+
+    @Test
+    fun testPendingIntentExtrasAreModified() {
+        mockLayoutInflater()
+        controlsSettingsRepository.setAllowActionOnTrivialControlsInLockscreen(true)
+
+        val panel = SelectedItem.PanelItem("App name", ComponentName("pkg", "cls"))
+        val serviceInfo = setUpPanel(panel)
+
+        underTest.show(parent, {}, context)
+
+        val captor = argumentCaptor<ControlsListingController.ControlsListingCallback>()
+
+        verify(controlsListingController).addCallback(capture(captor))
+
+        captor.value.onServicesUpdated(listOf(serviceInfo))
+        FakeExecutor.exhaustExecutors(uiExecutor, bgExecutor)
+
+        val pendingIntent = verifyPanelCreatedAndStartTaskView()
+        assertThat(
+                pendingIntent.intent.getBooleanExtra(
+                    ControlsProviderService.EXTRA_LOCKSCREEN_ALLOW_TRIVIAL_CONTROLS,
+                    false
+                )
+            )
+            .isTrue()
+
+        underTest.hide()
+
+        clearInvocations(controlsListingController, taskViewFactory)
+        controlsSettingsRepository.setAllowActionOnTrivialControlsInLockscreen(false)
+        underTest.show(parent, {}, context)
+
+        verify(controlsListingController).addCallback(capture(captor))
+        captor.value.onServicesUpdated(listOf(serviceInfo))
+        FakeExecutor.exhaustExecutors(uiExecutor, bgExecutor)
+
+        val newPendingIntent = verifyPanelCreatedAndStartTaskView()
+        assertThat(
+                newPendingIntent.intent.getBooleanExtra(
+                    ControlsProviderService.EXTRA_LOCKSCREEN_ALLOW_TRIVIAL_CONTROLS,
+                    false
+                )
+            )
+            .isFalse()
+    }
+
+    private fun setUpPanel(panel: SelectedItem.PanelItem): ControlsServiceInfo {
+        val activity = ComponentName("pkg", "activity")
         sharedPreferences
             .edit()
             .putString("controls_component", panel.componentName.flattenToString())
             .putString("controls_structure", panel.appName.toString())
             .putBoolean("controls_is_panel", true)
             .commit()
+        return ControlsServiceInfo(panel.componentName, panel.appName, activity)
+    }
 
-        underTest.show(parent, {}, activityContext)
+    private fun verifyPanelCreatedAndStartTaskView(): PendingIntent {
-        verify(controlsController, never()).refreshStatus(any(), any())
+        val taskViewConsumerCaptor = argumentCaptor<Consumer<TaskView>>()
+        verify(taskViewFactory).create(eq(context), eq(uiExecutor), capture(taskViewConsumerCaptor))
+
+        val taskView: TaskView = mock {
+            `when`(this.post(any())).thenAnswer {
+                uiExecutor.execute(it.arguments[0] as Runnable)
+                true
+            }
+        }
+        // calls PanelTaskViewController#launchTaskView
+        taskViewConsumerCaptor.value.accept(taskView)
+        val listenerCaptor = argumentCaptor<TaskView.Listener>()
+        verify(taskView).setListener(any(), capture(listenerCaptor))
+        listenerCaptor.value.onInitialized()
+        FakeExecutor.exhaustExecutors(uiExecutor, bgExecutor)
+
+        val pendingIntentCaptor = argumentCaptor<PendingIntent>()
+        verify(taskView).startActivity(capture(pendingIntentCaptor), any(), any(), any())
+        return pendingIntentCaptor.value
+    }
+
+    private fun ControlsServiceInfo(
+        componentName: ComponentName,
+        label: CharSequence,
+        panelComponentName: ComponentName? = null
+    ): ControlsServiceInfo {
+        val serviceInfo =
+            ServiceInfo().apply {
+                applicationInfo = ApplicationInfo()
+                packageName = componentName.packageName
+                name = componentName.className
+            }
+        return spy(ControlsServiceInfo(mContext, serviceInfo)).apply {
+            `when`(loadLabel()).thenReturn(label)
+            `when`(loadIcon()).thenReturn(mock())
+            `when`(panelActivity).thenReturn(panelComponentName)
+        }
+    }
+
+    private fun mockLayoutInflater() {
+        LayoutInflater.from(context)
+            .setPrivateFactory(
+                object : LayoutInflater.Factory2 {
+                    override fun onCreateView(
+                        view: View?,
+                        name: String,
+                        context: Context,
+                        attrs: AttributeSet
+                    ): View? {
+                        return onCreateView(name, context, attrs)
+                    }
+
+                    override fun onCreateView(
+                        name: String,
+                        context: Context,
+                        attrs: AttributeSet
+                    ): View? {
+                        if (FrameLayout::class.java.simpleName.equals(name)) {
+                            val mock: FrameLayout = mock {
+                                `when`(this.context).thenReturn(context)
+                                `when`(this.id).thenReturn(R.id.controls_panel)
+                                `when`(this.requireViewById<View>(any())).thenCallRealMethod()
+                                `when`(this.findViewById<View>(R.id.controls_panel))
+                                    .thenReturn(this)
+                                `when`(this.post(any())).thenAnswer {
+                                    uiExecutor.execute(it.arguments[0] as Runnable)
+                                    true
+                                }
+                            }
+                            return mock
+                        } else {
+                            return null
+                        }
+                    }
+                }
+            )
     }
 }

packages/SystemUI/tests/src/com/android/systemui/util/TestableAlertDialogTest.kt

@@ -0,0 +1,333 @@
+/*
+ * Copyright (C) 2022 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.util
+
+import android.content.DialogInterface
+import android.content.DialogInterface.BUTTON_NEGATIVE
+import android.content.DialogInterface.BUTTON_NEUTRAL
+import android.content.DialogInterface.BUTTON_POSITIVE
+import android.testing.AndroidTestingRunner
+import android.testing.TestableLooper
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.mock
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mockito.anyInt
+import org.mockito.Mockito.inOrder
+import org.mockito.Mockito.never
+import org.mockito.Mockito.verify
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+@TestableLooper.RunWithLooper
+class TestableAlertDialogTest : SysuiTestCase() {
+
+    @Test
+    fun dialogNotShowingWhenCreated() {
+        val dialog = TestableAlertDialog(context)
+
+        assertThat(dialog.isShowing).isFalse()
+    }
+
+    @Test
+    fun dialogShownDoesntCrash() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.show()
+    }
+
+    @Test
+    fun dialogShowing() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.show()
+
+        assertThat(dialog.isShowing).isTrue()
+    }
+
+    @Test
+    fun showListenerCalled() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnShowListener = mock()
+        dialog.setOnShowListener(listener)
+
+        dialog.show()
+
+        verify(listener).onShow(dialog)
+    }
+
+    @Test
+    fun showListenerRemoved() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnShowListener = mock()
+        dialog.setOnShowListener(listener)
+        dialog.setOnShowListener(null)
+
+        dialog.show()
+
+        verify(listener, never()).onShow(any())
+    }
+
+    @Test
+    fun dialogHiddenNotShowing() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.show()
+        dialog.hide()
+
+        assertThat(dialog.isShowing).isFalse()
+    }
+
+    @Test
+    fun dialogDismissNotShowing() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.show()
+        dialog.dismiss()
+
+        assertThat(dialog.isShowing).isFalse()
+    }
+
+    @Test
+    fun dismissListenerCalled_ifShowing() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnDismissListener = mock()
+        dialog.setOnDismissListener(listener)
+
+        dialog.show()
+        dialog.dismiss()
+
+        verify(listener).onDismiss(dialog)
+    }
+
+    @Test
+    fun dismissListenerNotCalled_ifNotShowing() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnDismissListener = mock()
+        dialog.setOnDismissListener(listener)
+
+        dialog.dismiss()
+
+        verify(listener, never()).onDismiss(any())
+    }
+
+    @Test
+    fun dismissListenerRemoved() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnDismissListener = mock()
+        dialog.setOnDismissListener(listener)
+        dialog.setOnDismissListener(null)
+
+        dialog.show()
+        dialog.dismiss()
+
+        verify(listener, never()).onDismiss(any())
+    }
+
+    @Test
+    fun cancelListenerCalled_showing() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnCancelListener = mock()
+        dialog.setOnCancelListener(listener)
+
+        dialog.show()
+        dialog.cancel()
+
+        verify(listener).onCancel(dialog)
+    }
+
+    @Test
+    fun cancelListenerCalled_notShowing() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnCancelListener = mock()
+        dialog.setOnCancelListener(listener)
+
+        dialog.cancel()
+
+        verify(listener).onCancel(dialog)
+    }
+
+    @Test
+    fun dismissCalledOnCancel_showing() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnDismissListener = mock()
+        dialog.setOnDismissListener(listener)
+
+        dialog.show()
+        dialog.cancel()
+
+        verify(listener).onDismiss(dialog)
+    }
+
+    @Test
+    fun dialogCancelNotShowing() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.show()
+        dialog.cancel()
+
+        assertThat(dialog.isShowing).isFalse()
+    }
+
+    @Test
+    fun cancelListenerRemoved() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnCancelListener = mock()
+        dialog.setOnCancelListener(listener)
+        dialog.setOnCancelListener(null)
+
+        dialog.show()
+        dialog.cancel()
+
+        verify(listener, never()).onCancel(any())
+    }
+
+    @Test
+    fun positiveButtonClick() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnClickListener = mock()
+        dialog.setButton(BUTTON_POSITIVE, "", listener)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_POSITIVE)
+
+        verify(listener).onClick(dialog, BUTTON_POSITIVE)
+    }
+
+    @Test
+    fun positiveButtonListener_noCalledWhenClickOtherButtons() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnClickListener = mock()
+        dialog.setButton(BUTTON_POSITIVE, "", listener)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_NEUTRAL)
+        dialog.clickButton(BUTTON_NEGATIVE)
+
+        verify(listener, never()).onClick(any(), anyInt())
+    }
+
+    @Test
+    fun negativeButtonClick() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnClickListener = mock()
+        dialog.setButton(BUTTON_NEGATIVE, "", listener)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_NEGATIVE)
+
+        verify(listener).onClick(dialog, DialogInterface.BUTTON_NEGATIVE)
+    }
+
+    @Test
+    fun negativeButtonListener_noCalledWhenClickOtherButtons() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnClickListener = mock()
+        dialog.setButton(BUTTON_NEGATIVE, "", listener)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_NEUTRAL)
+        dialog.clickButton(BUTTON_POSITIVE)
+
+        verify(listener, never()).onClick(any(), anyInt())
+    }
+
+    @Test
+    fun neutralButtonClick() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnClickListener = mock()
+        dialog.setButton(BUTTON_NEUTRAL, "", listener)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_NEUTRAL)
+
+        verify(listener).onClick(dialog, BUTTON_NEUTRAL)
+    }
+
+    @Test
+    fun neutralButtonListener_noCalledWhenClickOtherButtons() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnClickListener = mock()
+        dialog.setButton(BUTTON_NEUTRAL, "", listener)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_POSITIVE)
+        dialog.clickButton(BUTTON_NEGATIVE)
+
+        verify(listener, never()).onClick(any(), anyInt())
+    }
+
+    @Test
+    fun sameClickListenerCalledCorrectly() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnClickListener = mock()
+        dialog.setButton(BUTTON_POSITIVE, "", listener)
+        dialog.setButton(BUTTON_NEUTRAL, "", listener)
+        dialog.setButton(BUTTON_NEGATIVE, "", listener)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_POSITIVE)
+        dialog.clickButton(BUTTON_NEGATIVE)
+        dialog.clickButton(BUTTON_NEUTRAL)
+
+        val inOrder = inOrder(listener)
+        inOrder.verify(listener).onClick(dialog, BUTTON_POSITIVE)
+        inOrder.verify(listener).onClick(dialog, BUTTON_NEGATIVE)
+        inOrder.verify(listener).onClick(dialog, BUTTON_NEUTRAL)
+    }
+
+    @Test(expected = IllegalArgumentException::class)
+    fun clickBadButton() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.clickButton(10000)
+    }
+
+    @Test
+    fun clickButtonDismisses_positive() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_POSITIVE)
+
+        assertThat(dialog.isShowing).isFalse()
+    }
+
+    @Test
+    fun clickButtonDismisses_negative() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_NEGATIVE)
+
+        assertThat(dialog.isShowing).isFalse()
+    }
+
+    @Test
+    fun clickButtonDismisses_neutral() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_NEUTRAL)
+
+        assertThat(dialog.isShowing).isFalse()
+    }
+}

packages/SystemUI/tests/utils/src/com/android/systemui/util/TestableAlertDialog.kt

@@ -0,0 +1,141 @@
+/*
+ * Copyright (C) 2022 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.util
+
+import android.app.AlertDialog
+import android.content.Context
+import android.content.DialogInterface
+import java.lang.IllegalArgumentException
+
+/**
+ * [AlertDialog] that is easier to test. Due to [AlertDialog] being a class and not an interface,
+ * there are some things that cannot be avoided, like the creation of a [Handler] on the main thread
+ * (and therefore needing a prepared [Looper] in the test).
+ *
+ * It bypasses calls to show, clicks on buttons, cancel and dismiss so it all can happen bounded in
+ * the test. It tries to be as close in behavior as a real [AlertDialog].
+ *
+ * It will only call [onCreate] as part of its lifecycle, but not any of the other lifecycle methods
+ * in [Dialog].
+ *
+ * In order to test clicking on buttons, use [clickButton] instead of calling [View.callOnClick] on
+ * the view returned by [getButton] to bypass the internal [Handler].
+ */
+class TestableAlertDialog(context: Context) : AlertDialog(context) {
+
+    private var _onDismissListener: DialogInterface.OnDismissListener? = null
+    private var _onCancelListener: DialogInterface.OnCancelListener? = null
+    private var _positiveButtonClickListener: DialogInterface.OnClickListener? = null
+    private var _negativeButtonClickListener: DialogInterface.OnClickListener? = null
+    private var _neutralButtonClickListener: DialogInterface.OnClickListener? = null
+    private var _onShowListener: DialogInterface.OnShowListener? = null
+    private var _dismissOverride: Runnable? = null
+
+    private var showing = false
+    private var visible = false
+    private var created = false
+
+    override fun show() {
+        if (!created) {
+            created = true
+            onCreate(null)
+        }
+        if (isShowing) return
+        showing = true
+        visible = true
+        _onShowListener?.onShow(this)
+    }
+
+    override fun hide() {
+        visible = false
+    }
+
+    override fun isShowing(): Boolean {
+        return visible && showing
+    }
+
+    override fun dismiss() {
+        if (!showing) {
+            return
+        }
+        if (_dismissOverride != null) {
+            _dismissOverride?.run()
+            return
+        }
+        _onDismissListener?.onDismiss(this)
+        showing = false
+    }
+
+    override fun cancel() {
+        _onCancelListener?.onCancel(this)
+        dismiss()
+    }
+
+    override fun setOnDismissListener(listener: DialogInterface.OnDismissListener?) {
+        _onDismissListener = listener
+    }
+
+    override fun setOnCancelListener(listener: DialogInterface.OnCancelListener?) {
+        _onCancelListener = listener
+    }
+
+    override fun setOnShowListener(listener: DialogInterface.OnShowListener?) {
+        _onShowListener = listener
+    }
+
+    override fun takeCancelAndDismissListeners(
+        msg: String?,
+        cancel: DialogInterface.OnCancelListener?,
+        dismiss: DialogInterface.OnDismissListener?
+    ): Boolean {
+        _onCancelListener = cancel
+        _onDismissListener = dismiss
+        return true
+    }
+
+    override fun setButton(
+        whichButton: Int,
+        text: CharSequence?,
+        listener: DialogInterface.OnClickListener?
+    ) {
+        super.setButton(whichButton, text, listener)
+        when (whichButton) {
+            DialogInterface.BUTTON_POSITIVE -> _positiveButtonClickListener = listener
+            DialogInterface.BUTTON_NEGATIVE -> _negativeButtonClickListener = listener
+            DialogInterface.BUTTON_NEUTRAL -> _neutralButtonClickListener = listener
+            else -> Unit
+        }
+    }
+
+    /**
+     * Click one of the buttons in the [AlertDialog] and call the corresponding listener.
+     *
+     * Button ids are from [DialogInterface].
+     */
+    fun clickButton(whichButton: Int) {
+        val listener =
+            when (whichButton) {
+                DialogInterface.BUTTON_POSITIVE -> _positiveButtonClickListener
+                DialogInterface.BUTTON_NEGATIVE -> _negativeButtonClickListener
+                DialogInterface.BUTTON_NEUTRAL -> _neutralButtonClickListener
+                else -> throw IllegalArgumentException("Wrong button $whichButton")
+            }
+        listener?.onClick(this, whichButton)
+        dismiss()
+    }
+}