Restart home controls activity when dream is redirecting wakes

When the dream is redirecting calls to wakeUp, the dream is not
terminated but the embedded activity is. This results in a blank dream.
The fix here will restart the embedded activity in this case, after a
short delay to allow the wakeup animation to play first.

Fixes: 341957747
Test: atest HomeControlsDreamServiceTest
Flag: android.service.controls.flags.home_panel_dream
Change-Id: I495b20235affa49e31b14bc0e50acee007f2813d

6 files changed, 167 insertions, 54 deletions

diff --git a/core/java/android/service/dreams/DreamService.java b/core/java/android/service/dreams/DreamService.java
index 71066ac7ac39..3f9c819cd62f 100644
--- a/core/java/android/service/dreams/DreamService.java
+++ b/core/java/android/service/dreams/DreamService.java
@@ -1276,13 +1276,22 @@ public class DreamService extends Service implements Window.Callback {
         });
     }
 
+    /**
+     * Whether or not wake requests will be redirected.
+     *
+     * @hide
+     */
+    public boolean getRedirectWake() {
+        return mOverlayConnection != null && mRedirectWake;
+    }
+
     private void wakeUp(boolean fromSystem) {
         if (mDebug) {
             Slog.v(mTag, "wakeUp(): fromSystem=" + fromSystem + ", mWaking=" + mWaking
                     + ", mFinished=" + mFinished);
         }
 
-        if (!fromSystem && mOverlayConnection != null && mRedirectWake) {
+        if (!fromSystem && getRedirectWake()) {
             mOverlayConnection.addConsumer(overlay -> {
                 try {
                     overlay.onWakeRequested();
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/homecontrols/HomeControlsDreamServiceTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/homecontrols/HomeControlsDreamServiceTest.kt
index 723f6a2bfff4..9300db9a24c8 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/homecontrols/HomeControlsDreamServiceTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/homecontrols/HomeControlsDreamServiceTest.kt
@@ -16,29 +16,40 @@
 package com.android.systemui.dreams.homecontrols
 
 import android.app.Activity
+import android.content.Intent
+import android.service.controls.ControlsProviderService.CONTROLS_SURFACE_ACTIVITY_PANEL
+import android.service.controls.ControlsProviderService.CONTROLS_SURFACE_DREAM
+import android.service.controls.ControlsProviderService.EXTRA_CONTROLS_SURFACE
+import android.window.TaskFragmentInfo
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.controls.settings.FakeControlsSettingsRepository
 import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.kosmos.testScope
-import com.android.systemui.log.core.FakeLogBuffer.Factory.Companion.create
 import com.android.systemui.log.logcatLogBuffer
 import com.android.systemui.testKosmos
-import com.android.systemui.util.mockito.any
-import com.android.systemui.util.mockito.whenever
 import com.android.systemui.util.wakelock.WakeLockFake
 import com.google.common.truth.Truth.assertThat
 import java.util.Optional
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.advanceUntilIdle
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.Mockito.never
-import org.mockito.Mockito.verify
-import org.mockito.MockitoAnnotations
-
+import org.mockito.kotlin.any
+import org.mockito.kotlin.argThat
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.doReturn
+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
+import org.mockito.kotlin.whenever
+
+@OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class HomeControlsDreamServiceTest : SysuiTestCase() {
@@ -46,31 +57,38 @@ class HomeControlsDreamServiceTest : SysuiTestCase() {
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
 
-    private lateinit var fakeWakeLockBuilder: WakeLockFake.Builder
-    private lateinit var fakeWakeLock: WakeLockFake
-
-    @Mock private lateinit var taskFragmentComponentFactory: TaskFragmentComponent.Factory
-    @Mock private lateinit var taskFragmentComponent: TaskFragmentComponent
-    @Mock private lateinit var activity: Activity
+    private val fakeWakeLock = WakeLockFake()
+    private val fakeWakeLockBuilder by lazy {
+        WakeLockFake.Builder(context).apply { setWakeLock(fakeWakeLock) }
+    }
+
+    private val taskFragmentComponent = mock<TaskFragmentComponent>()
+    private val activity = mock<Activity>()
+    private val onCreateCallback = argumentCaptor<(TaskFragmentInfo) -> Unit>()
+    private val onInfoChangedCallback = argumentCaptor<(TaskFragmentInfo) -> Unit>()
+    private val hideCallback = argumentCaptor<() -> Unit>()
+    private val dreamServiceDelegate =
+        mock<DreamServiceDelegate> { on { getActivity(any()) } doReturn activity }
+
+    private val taskFragmentComponentFactory =
+        mock<TaskFragmentComponent.Factory> {
+            on {
+                create(
+                    activity = eq(activity),
+                    onCreateCallback = onCreateCallback.capture(),
+                    onInfoChangedCallback = onInfoChangedCallback.capture(),
+                    hide = hideCallback.capture(),
+                )
+            } doReturn taskFragmentComponent
+        }
 
-    private lateinit var underTest: HomeControlsDreamService
+    private val underTest: HomeControlsDreamService by lazy { buildService() }
 
     @Before
-    fun setup() =
-        with(kosmos) {
-            MockitoAnnotations.initMocks(this@HomeControlsDreamServiceTest)
-            whenever(taskFragmentComponentFactory.create(any(), any(), any(), any()))
-                .thenReturn(taskFragmentComponent)
-
-            fakeWakeLock = WakeLockFake()
-            fakeWakeLockBuilder = WakeLockFake.Builder(context)
-            fakeWakeLockBuilder.setWakeLock(fakeWakeLock)
-
-            whenever(controlsComponent.getControlsListingController())
-                .thenReturn(Optional.of(controlsListingController))
-
-            underTest = buildService { activity }
-        }
+    fun setup() {
+        whenever(kosmos.controlsComponent.getControlsListingController())
+            .thenReturn(Optional.of(kosmos.controlsListingController))
+    }
 
     @Test
     fun testOnAttachedToWindowCreatesTaskFragmentComponent() =
@@ -90,9 +108,12 @@ class HomeControlsDreamServiceTest : SysuiTestCase() {
     @Test
     fun testNotCreatingTaskFragmentComponentWhenActivityIsNull() =
         testScope.runTest {
-            underTest = buildService { null }
+            val serviceWithNullActivity =
+                buildService(
+                    mock<DreamServiceDelegate> { on { getActivity(underTest) } doReturn null }
+                )
 
-            underTest.onAttachedToWindow()
+            serviceWithNullActivity.onAttachedToWindow()
             verify(taskFragmentComponentFactory, never()).create(any(), any(), any(), any())
         }
 
@@ -102,6 +123,7 @@ class HomeControlsDreamServiceTest : SysuiTestCase() {
             underTest.onAttachedToWindow()
             assertThat(fakeWakeLock.isHeld).isTrue()
         }
+
     @Test
     fun testDetachWindow_wakeLockCanBeReleased() =
         testScope.runTest {
@@ -112,14 +134,60 @@ class HomeControlsDreamServiceTest : SysuiTestCase() {
             assertThat(fakeWakeLock.isHeld).isFalse()
         }
 
-    private fun buildService(activityProvider: DreamActivityProvider): HomeControlsDreamService =
+    @Test
+    fun testFinishesDreamWithoutRestartingActivityWhenNotRedirectingWakes() =
+        testScope.runTest {
+            whenever(dreamServiceDelegate.redirectWake(any())).thenReturn(false)
+            underTest.onAttachedToWindow()
+            onCreateCallback.firstValue.invoke(mock<TaskFragmentInfo>())
+            verify(taskFragmentComponent, times(1)).startActivityInTaskFragment(intentMatcher())
+
+            // Task fragment becomes empty
+            onInfoChangedCallback.firstValue.invoke(
+                mock<TaskFragmentInfo> { on { isEmpty } doReturn true }
+            )
+            advanceUntilIdle()
+            // Dream is finished and activity is not restarted
+            verify(taskFragmentComponent, times(1)).startActivityInTaskFragment(intentMatcher())
+            verify(dreamServiceDelegate, never()).wakeUp(any())
+            verify(dreamServiceDelegate).finish(any())
+        }
+
+    @Test
+    fun testRestartsActivityWhenRedirectingWakes() =
+        testScope.runTest {
+            whenever(dreamServiceDelegate.redirectWake(any())).thenReturn(true)
+            underTest.onAttachedToWindow()
+            onCreateCallback.firstValue.invoke(mock<TaskFragmentInfo>())
+            verify(taskFragmentComponent, times(1)).startActivityInTaskFragment(intentMatcher())
+
+            // Task fragment becomes empty
+            onInfoChangedCallback.firstValue.invoke(
+                mock<TaskFragmentInfo> { on { isEmpty } doReturn true }
+            )
+            advanceUntilIdle()
+            // Activity is restarted instead of finishing the dream.
+            verify(taskFragmentComponent, times(2)).startActivityInTaskFragment(intentMatcher())
+            verify(dreamServiceDelegate).wakeUp(any())
+            verify(dreamServiceDelegate, never()).finish(any())
+        }
+
+    private fun intentMatcher() =
+        argThat<Intent> {
+            getIntExtra(EXTRA_CONTROLS_SURFACE, CONTROLS_SURFACE_ACTIVITY_PANEL) ==
+                CONTROLS_SURFACE_DREAM
+        }
+
+    private fun buildService(
+        activityProvider: DreamServiceDelegate = dreamServiceDelegate
+    ): HomeControlsDreamService =
         with(kosmos) {
             return HomeControlsDreamService(
                 controlsSettingsRepository = FakeControlsSettingsRepository(),
                 taskFragmentFactory = taskFragmentComponentFactory,
                 homeControlsComponentInteractor = homeControlsComponentInteractor,
-                fakeWakeLockBuilder,
-                dreamActivityProvider = activityProvider,
+                wakeLockBuilder = fakeWakeLockBuilder,
+                dreamServiceDelegate = activityProvider,
                 bgDispatcher = testDispatcher,
                 logBuffer = logcatLogBuffer("HomeControlsDreamServiceTest")
             )
diff --git a/packages/SystemUI/src/com/android/systemui/dreams/dagger/DreamModule.java b/packages/SystemUI/src/com/android/systemui/dreams/dagger/DreamModule.java
index b0d134f5f15f..f6ac7a579140 100644
--- a/packages/SystemUI/src/com/android/systemui/dreams/dagger/DreamModule.java
+++ b/packages/SystemUI/src/com/android/systemui/dreams/dagger/DreamModule.java
@@ -33,8 +33,8 @@ import com.android.systemui.dreams.DreamOverlayNotificationCountProvider;
 import com.android.systemui.dreams.DreamOverlayService;
 import com.android.systemui.dreams.SystemDialogsCloser;
 import com.android.systemui.dreams.complication.dagger.ComplicationComponent;
-import com.android.systemui.dreams.homecontrols.DreamActivityProvider;
-import com.android.systemui.dreams.homecontrols.DreamActivityProviderImpl;
+import com.android.systemui.dreams.homecontrols.DreamServiceDelegate;
+import com.android.systemui.dreams.homecontrols.DreamServiceDelegateImpl;
 import com.android.systemui.dreams.homecontrols.HomeControlsDreamService;
 import com.android.systemui.qs.QsEventLogger;
 import com.android.systemui.qs.pipeline.shared.TileSpec;
@@ -202,8 +202,8 @@ public interface DreamModule {
     }
 
 
-    /** Provides activity for dream service */
+    /** Provides delegate to allow for testing of dream service */
     @Binds
-    DreamActivityProvider bindActivityProvider(DreamActivityProviderImpl impl);
+    DreamServiceDelegate bindDreamDelegate(DreamServiceDelegateImpl impl);
 
 }
diff --git a/packages/SystemUI/src/com/android/systemui/dreams/homecontrols/DreamActivityProvider.kt b/packages/SystemUI/src/com/android/systemui/dreams/homecontrols/DreamServiceDelegate.kt
index b35b7f5debb3..2cfb02eadd19 100644
--- a/packages/SystemUI/src/com/android/systemui/dreams/homecontrols/DreamActivityProvider.kt
+++ b/packages/SystemUI/src/com/android/systemui/dreams/homecontrols/DreamServiceDelegate.kt
@@ -18,10 +18,17 @@ package com.android.systemui.dreams.homecontrols
 import android.app.Activity
 import android.service.dreams.DreamService
 
-fun interface DreamActivityProvider {
-    /**
-     * Provides abstraction for getting the activity associated with a dream service, so that the
-     * activity can be mocked in tests.
-     */
+/** Provides abstraction for [DreamService] methods, so they can be mocked in tests. */
+interface DreamServiceDelegate {
+    /** Wrapper for [DreamService.getActivity] which can be mocked in tests. */
     fun getActivity(dreamService: DreamService): Activity?
+
+    /** Wrapper for [DreamService.wakeUp] which can be mocked in tests. */
+    fun wakeUp(dreamService: DreamService)
+
+    /** Wrapper for [DreamService.finish] which can be mocked in tests. */
+    fun finish(dreamService: DreamService)
+
+    /** Wrapper for [DreamService.getRedirectWake] which can be mocked in tests. */
+    fun redirectWake(dreamService: DreamService): Boolean
 }
diff --git a/packages/SystemUI/src/com/android/systemui/dreams/homecontrols/DreamActivityProviderImpl.kt b/packages/SystemUI/src/com/android/systemui/dreams/homecontrols/DreamServiceDelegateImpl.kt
index 0854e939645b..7dc5434c595e 100644
--- a/packages/SystemUI/src/com/android/systemui/dreams/homecontrols/DreamActivityProviderImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/dreams/homecontrols/DreamServiceDelegateImpl.kt
@@ -19,8 +19,20 @@ import android.app.Activity
 import android.service.dreams.DreamService
 import javax.inject.Inject
 
-class DreamActivityProviderImpl @Inject constructor() : DreamActivityProvider {
+class DreamServiceDelegateImpl @Inject constructor() : DreamServiceDelegate {
     override fun getActivity(dreamService: DreamService): Activity {
         return dreamService.activity
     }
+
+    override fun finish(dreamService: DreamService) {
+        dreamService.finish()
+    }
+
+    override fun wakeUp(dreamService: DreamService) {
+        dreamService.wakeUp()
+    }
+
+    override fun redirectWake(dreamService: DreamService): Boolean {
+        return dreamService.redirectWake
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/dreams/homecontrols/HomeControlsDreamService.kt b/packages/SystemUI/src/com/android/systemui/dreams/homecontrols/HomeControlsDreamService.kt
index 76187c614b5d..77c54ec1eac3 100644
--- a/packages/SystemUI/src/com/android/systemui/dreams/homecontrols/HomeControlsDreamService.kt
+++ b/packages/SystemUI/src/com/android/systemui/dreams/homecontrols/HomeControlsDreamService.kt
@@ -31,6 +31,7 @@ import com.android.systemui.log.dagger.DreamLog
 import com.android.systemui.util.wakelock.WakeLock
 import com.android.systemui.util.wakelock.WakeLock.Builder.NO_TIMEOUT
 import javax.inject.Inject
+import kotlin.time.Duration.Companion.milliseconds
 import kotlin.time.Duration.Companion.seconds
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
@@ -46,7 +47,7 @@ constructor(
     private val taskFragmentFactory: TaskFragmentComponent.Factory,
     private val homeControlsComponentInteractor: HomeControlsComponentInteractor,
     private val wakeLockBuilder: WakeLock.Builder,
-    private val dreamActivityProvider: DreamActivityProvider,
+    private val dreamServiceDelegate: DreamServiceDelegate,
     @Background private val bgDispatcher: CoroutineDispatcher,
     @DreamLog logBuffer: LogBuffer
 ) : DreamService() {
@@ -65,7 +66,7 @@ constructor(
 
     override fun onAttachedToWindow() {
         super.onAttachedToWindow()
-        val activity = dreamActivityProvider.getActivity(this)
+        val activity = dreamServiceDelegate.getActivity(this)
         if (activity == null) {
             finish()
             return
@@ -79,9 +80,9 @@ constructor(
             taskFragmentFactory
                 .create(
                     activity = activity,
-                    onCreateCallback = this::onTaskFragmentCreated,
+                    onCreateCallback = { launchActivity() },
                     onInfoChangedCallback = this::onTaskFragmentInfoChanged,
-                    hide = { endDream() }
+                    hide = { endDream(false) }
                 )
                 .apply { createTaskFragment() }
 
@@ -91,16 +92,24 @@ constructor(
     private fun onTaskFragmentInfoChanged(taskFragmentInfo: TaskFragmentInfo) {
         if (taskFragmentInfo.isEmpty) {
             logger.d("Finishing dream due to TaskFragment being empty")
-            endDream()
+            endDream(true)
         }
     }
 
-    private fun endDream() {
+    private fun endDream(handleRedirect: Boolean) {
         homeControlsComponentInteractor.onDreamEndUnexpectedly()
-        finish()
+        if (handleRedirect && dreamServiceDelegate.redirectWake(this)) {
+            dreamServiceDelegate.wakeUp(this)
+            serviceScope.launch {
+                delay(ACTIVITY_RESTART_DELAY)
+                launchActivity()
+            }
+        } else {
+            dreamServiceDelegate.finish(this)
+        }
     }
 
-    private fun onTaskFragmentCreated(taskFragmentInfo: TaskFragmentInfo) {
+    private fun launchActivity() {
         val setting = controlsSettingsRepository.allowActionOnTrivialControlsInLockscreen.value
         val componentName = homeControlsComponentInteractor.panelComponent.value
         logger.d("Starting embedding $componentName")
@@ -134,6 +143,14 @@ constructor(
          * complete.
          */
         val CANCELLATION_DELAY_AFTER_DETACHED = 5.seconds
+
+        /**
+         * Defines the delay after wakeup where we should attempt to restart the embedded activity.
+         * When a wakeup is redirected, the dream service may keep running. In this case, we should
+         * restart the activity if it finished. This delays ensures the activity is only restarted
+         * after the wakeup transition has played.
+         */
+        val ACTIVITY_RESTART_DELAY = 334.milliseconds
         const val TAG = "HomeControlsDreamService"
     }
 }