Merge "PSS Task Switcher: Improve unit tests - Use real repositories" into udc-qpr-dev

7 files changed, 211 insertions, 183 deletions

diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/data/repository/FakeMediaProjectionManager.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/data/repository/FakeMediaProjectionManager.kt
new file mode 100644
index 000000000000..45f0a8c62125
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/data/repository/FakeMediaProjectionManager.kt
@@ -0,0 +1,71 @@
+/*
+ * Copyright (C) 2023 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.mediaprojection.taskswitcher.data.repository
+
+import android.media.projection.MediaProjectionInfo
+import android.media.projection.MediaProjectionManager
+import android.os.Binder
+import android.os.IBinder
+import android.os.UserHandle
+import android.view.ContentRecordingSession
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.whenever
+
+class FakeMediaProjectionManager {
+
+    val mediaProjectionManager = mock<MediaProjectionManager>()
+
+    private val callbacks = mutableListOf<MediaProjectionManager.Callback>()
+
+    init {
+        whenever(mediaProjectionManager.addCallback(any(), any())).thenAnswer {
+            callbacks += it.arguments[0] as MediaProjectionManager.Callback
+            return@thenAnswer Unit
+        }
+        whenever(mediaProjectionManager.removeCallback(any())).thenAnswer {
+            callbacks -= it.arguments[0] as MediaProjectionManager.Callback
+            return@thenAnswer Unit
+        }
+    }
+
+    fun dispatchOnStart(info: MediaProjectionInfo = DEFAULT_INFO) {
+        callbacks.forEach { it.onStart(info) }
+    }
+
+    fun dispatchOnStop(info: MediaProjectionInfo = DEFAULT_INFO) {
+        callbacks.forEach { it.onStop(info) }
+    }
+
+    fun dispatchOnSessionSet(
+        info: MediaProjectionInfo = DEFAULT_INFO,
+        session: ContentRecordingSession?
+    ) {
+        callbacks.forEach { it.onRecordingSessionSet(info, session) }
+    }
+
+    companion object {
+        fun createDisplaySession(): ContentRecordingSession =
+            ContentRecordingSession.createDisplaySession(/* displayToMirror = */ 123)
+        fun createSingleTaskSession(token: IBinder = Binder()): ContentRecordingSession =
+            ContentRecordingSession.createTaskSession(token)
+
+        private const val DEFAULT_PACKAGE_NAME = "com.media.projection.test"
+        private val DEFAULT_USER_HANDLE = UserHandle.getUserHandleForUid(UserHandle.myUserId())
+        private val DEFAULT_INFO = MediaProjectionInfo(DEFAULT_PACKAGE_NAME, DEFAULT_USER_HANDLE)
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/data/repository/FakeMediaProjectionRepository.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/data/repository/FakeMediaProjectionRepository.kt
deleted file mode 100644
index c59fd60cca9b..000000000000
--- a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/data/repository/FakeMediaProjectionRepository.kt
+++ /dev/null
@@ -1,42 +0,0 @@
-/*
- * Copyright (C) 2023 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.mediaprojection.taskswitcher.data.repository
-
-import android.app.TaskInfo
-import com.android.systemui.mediaprojection.taskswitcher.data.model.MediaProjectionState
-import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.asStateFlow
-
-class FakeMediaProjectionRepository : MediaProjectionRepository {
-
-    private val state = MutableStateFlow<MediaProjectionState>(MediaProjectionState.NotProjecting)
-
-    fun switchProjectedTask(newTask: TaskInfo) {
-        state.value = MediaProjectionState.SingleTask(newTask)
-    }
-
-    override val mediaProjectionState: Flow<MediaProjectionState> = state.asStateFlow()
-
-    fun projectEntireScreen() {
-        state.value = MediaProjectionState.EntireScreen
-    }
-
-    fun stopProjecting() {
-        state.value = MediaProjectionState.NotProjecting
-    }
-}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/data/repository/FakeTasksRepository.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/data/repository/FakeTasksRepository.kt
deleted file mode 100644
index 593e3893fb2a..000000000000
--- a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/data/repository/FakeTasksRepository.kt
+++ /dev/null
@@ -1,68 +0,0 @@
-/*
- * Copyright (C) 2023 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.mediaprojection.taskswitcher.data.repository
-
-import android.app.ActivityManager.RunningTaskInfo
-import android.content.Intent
-import android.os.IBinder
-import android.window.IWindowContainerToken
-import android.window.WindowContainerToken
-import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.asStateFlow
-
-class FakeTasksRepository : TasksRepository {
-
-    private val _foregroundTask = MutableStateFlow(DEFAULT_TASK)
-
-    override val foregroundTask: Flow<RunningTaskInfo> = _foregroundTask.asStateFlow()
-
-    private val runningTasks = mutableListOf(DEFAULT_TASK)
-
-    override suspend fun findRunningTaskFromWindowContainerToken(
-        windowContainerToken: IBinder
-    ): RunningTaskInfo? = runningTasks.firstOrNull { it.token.asBinder() == windowContainerToken }
-
-    fun addRunningTask(task: RunningTaskInfo) {
-        runningTasks.add(task)
-    }
-
-    fun moveTaskToForeground(task: RunningTaskInfo) {
-        _foregroundTask.value = task
-    }
-
-    companion object {
-        val DEFAULT_TASK = createTask(taskId = -1)
-        val LAUNCHER_INTENT: Intent = Intent(Intent.ACTION_MAIN).addCategory(Intent.CATEGORY_HOME)
-
-        fun createTask(
-            taskId: Int,
-            token: WindowContainerToken = createToken(),
-            baseIntent: Intent = Intent()
-        ) =
-            RunningTaskInfo().apply {
-                this.taskId = taskId
-                this.token = token
-                this.baseIntent = baseIntent
-            }
-
-        fun createToken(): WindowContainerToken {
-            val realToken = object : IWindowContainerToken.Stub() {}
-            return WindowContainerToken(realToken)
-        }
-    }
-}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/data/repository/MediaProjectionManagerRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/data/repository/MediaProjectionManagerRepositoryTest.kt
index 2b074655bb02..3a74c7255cbb 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/data/repository/MediaProjectionManagerRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/data/repository/MediaProjectionManagerRepositoryTest.kt
@@ -16,27 +16,22 @@
 
 package com.android.systemui.mediaprojection.taskswitcher.data.repository
 
-import android.media.projection.MediaProjectionInfo
-import android.media.projection.MediaProjectionManager
 import android.os.Binder
 import android.os.Handler
-import android.os.UserHandle
 import android.testing.AndroidTestingRunner
 import android.view.ContentRecordingSession
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.mediaprojection.taskswitcher.data.model.MediaProjectionState
-import com.android.systemui.util.mockito.any
-import com.android.systemui.util.mockito.mock
-import com.android.systemui.util.mockito.whenever
+import com.android.systemui.mediaprojection.taskswitcher.data.repository.FakeActivityTaskManager.Companion.createTask
+import com.android.systemui.mediaprojection.taskswitcher.data.repository.FakeActivityTaskManager.Companion.createToken
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
-import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 
@@ -45,29 +40,26 @@ import org.junit.runner.RunWith
 @SmallTest
 class MediaProjectionManagerRepositoryTest : SysuiTestCase() {
 
-    private val mediaProjectionManager = mock<MediaProjectionManager>()
-
     private val dispatcher = StandardTestDispatcher()
     private val testScope = TestScope(dispatcher)
-    private val tasksRepo = FakeTasksRepository()
 
-    private lateinit var callback: MediaProjectionManager.Callback
-    private lateinit var repo: MediaProjectionManagerRepository
+    private val fakeMediaProjectionManager = FakeMediaProjectionManager()
+    private val fakeActivityTaskManager = FakeActivityTaskManager()
 
-    @Before
-    fun setUp() {
-        whenever(mediaProjectionManager.addCallback(any(), any())).thenAnswer {
-            callback = it.arguments[0] as MediaProjectionManager.Callback
-            return@thenAnswer Unit
-        }
-        repo =
-            MediaProjectionManagerRepository(
-                mediaProjectionManager = mediaProjectionManager,
-                handler = Handler.getMain(),
-                applicationScope = testScope.backgroundScope,
-                tasksRepository = tasksRepo
-            )
-    }
+    private val tasksRepo =
+        ActivityTaskManagerTasksRepository(
+            activityTaskManager = fakeActivityTaskManager.activityTaskManager,
+            applicationScope = testScope.backgroundScope,
+            backgroundDispatcher = dispatcher
+        )
+
+    private val repo =
+        MediaProjectionManagerRepository(
+            mediaProjectionManager = fakeMediaProjectionManager.mediaProjectionManager,
+            handler = Handler.getMain(),
+            applicationScope = testScope.backgroundScope,
+            tasksRepository = tasksRepo
+        )
 
     @Test
     fun mediaProjectionState_onStart_emitsNotProjecting() =
@@ -75,7 +67,7 @@ class MediaProjectionManagerRepositoryTest : SysuiTestCase() {
             val state by collectLastValue(repo.mediaProjectionState)
             runCurrent()
 
-            callback.onStart(TEST_MEDIA_INFO)
+            fakeMediaProjectionManager.dispatchOnStart()
 
             assertThat(state).isEqualTo(MediaProjectionState.NotProjecting)
         }
@@ -86,7 +78,7 @@ class MediaProjectionManagerRepositoryTest : SysuiTestCase() {
             val state by collectLastValue(repo.mediaProjectionState)
             runCurrent()
 
-            callback.onStop(TEST_MEDIA_INFO)
+            fakeMediaProjectionManager.dispatchOnStop()
 
             assertThat(state).isEqualTo(MediaProjectionState.NotProjecting)
         }
@@ -97,7 +89,7 @@ class MediaProjectionManagerRepositoryTest : SysuiTestCase() {
             val state by collectLastValue(repo.mediaProjectionState)
             runCurrent()
 
-            callback.onRecordingSessionSet(TEST_MEDIA_INFO, /* session= */ null)
+            fakeMediaProjectionManager.dispatchOnSessionSet(session = null)
 
             assertThat(state).isEqualTo(MediaProjectionState.NotProjecting)
         }
@@ -108,8 +100,9 @@ class MediaProjectionManagerRepositoryTest : SysuiTestCase() {
             val state by collectLastValue(repo.mediaProjectionState)
             runCurrent()
 
-            val session = ContentRecordingSession.createDisplaySession(/* displayToMirror= */ 123)
-            callback.onRecordingSessionSet(TEST_MEDIA_INFO, session)
+            fakeMediaProjectionManager.dispatchOnSessionSet(
+                session = ContentRecordingSession.createDisplaySession(/* displayToMirror= */ 123)
+            )
 
             assertThat(state).isEqualTo(MediaProjectionState.EntireScreen)
         }
@@ -120,9 +113,10 @@ class MediaProjectionManagerRepositoryTest : SysuiTestCase() {
             val state by collectLastValue(repo.mediaProjectionState)
             runCurrent()
 
-            val session =
-                ContentRecordingSession.createTaskSession(/* taskWindowContainerToken= */ null)
-            callback.onRecordingSessionSet(TEST_MEDIA_INFO, session)
+            fakeMediaProjectionManager.dispatchOnSessionSet(
+                session =
+                    ContentRecordingSession.createTaskSession(/* taskWindowContainerToken= */ null)
+            )
 
             assertThat(state).isEqualTo(MediaProjectionState.EntireScreen)
         }
@@ -134,8 +128,9 @@ class MediaProjectionManagerRepositoryTest : SysuiTestCase() {
             runCurrent()
 
             val taskWindowContainerToken = Binder()
-            val session = ContentRecordingSession.createTaskSession(taskWindowContainerToken)
-            callback.onRecordingSessionSet(TEST_MEDIA_INFO, session)
+            fakeMediaProjectionManager.dispatchOnSessionSet(
+                session = ContentRecordingSession.createTaskSession(taskWindowContainerToken)
+            )
 
             assertThat(state).isEqualTo(MediaProjectionState.EntireScreen)
         }
@@ -143,20 +138,16 @@ class MediaProjectionManagerRepositoryTest : SysuiTestCase() {
     @Test
     fun mediaProjectionState_sessionSet_taskWithToken_matchingRunningTask_emitsSingleTask() =
         testScope.runTest {
-            val token = FakeTasksRepository.createToken()
-            val task = FakeTasksRepository.createTask(taskId = 1, token = token)
-            tasksRepo.addRunningTask(task)
+            val token = createToken()
+            val task = createTask(taskId = 1, token = token)
+            fakeActivityTaskManager.addRunningTasks(task)
             val state by collectLastValue(repo.mediaProjectionState)
             runCurrent()
 
-            val session = ContentRecordingSession.createTaskSession(token.asBinder())
-            callback.onRecordingSessionSet(TEST_MEDIA_INFO, session)
+            fakeMediaProjectionManager.dispatchOnSessionSet(
+                session = ContentRecordingSession.createTaskSession(token.asBinder())
+            )
 
             assertThat(state).isEqualTo(MediaProjectionState.SingleTask(task))
         }
-
-    companion object {
-        val TEST_MEDIA_INFO =
-            MediaProjectionInfo(/* packageName= */ "com.test.package", UserHandle.CURRENT)
-    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/domain/interactor/TaskSwitchInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/domain/interactor/TaskSwitchInteractorTest.kt
index 112950b860e8..b2ebe1bcbc8b 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/domain/interactor/TaskSwitchInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/domain/interactor/TaskSwitchInteractorTest.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.mediaprojection.taskswitcher.domain.interactor
 
 import android.content.Intent
+import android.os.Handler
 import android.testing.AndroidTestingRunner
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
@@ -24,7 +25,9 @@ import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.mediaprojection.taskswitcher.data.repository.ActivityTaskManagerTasksRepository
 import com.android.systemui.mediaprojection.taskswitcher.data.repository.FakeActivityTaskManager
 import com.android.systemui.mediaprojection.taskswitcher.data.repository.FakeActivityTaskManager.Companion.createTask
-import com.android.systemui.mediaprojection.taskswitcher.data.repository.FakeMediaProjectionRepository
+import com.android.systemui.mediaprojection.taskswitcher.data.repository.FakeMediaProjectionManager
+import com.android.systemui.mediaprojection.taskswitcher.data.repository.FakeMediaProjectionManager.Companion.createSingleTaskSession
+import com.android.systemui.mediaprojection.taskswitcher.data.repository.MediaProjectionManagerRepository
 import com.android.systemui.mediaprojection.taskswitcher.domain.model.TaskSwitchState
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -43,7 +46,8 @@ class TaskSwitchInteractorTest : SysuiTestCase() {
     private val testScope = TestScope(dispatcher)
 
     private val fakeActivityTaskManager = FakeActivityTaskManager()
-    private val mediaRepo = FakeMediaProjectionRepository()
+    private val fakeMediaProjectionManager = FakeMediaProjectionManager()
+
     private val tasksRepo =
         ActivityTaskManagerTasksRepository(
             activityTaskManager = fakeActivityTaskManager.activityTaskManager,
@@ -51,15 +55,26 @@ class TaskSwitchInteractorTest : SysuiTestCase() {
             backgroundDispatcher = dispatcher
         )
 
+    private val mediaRepo =
+        MediaProjectionManagerRepository(
+            mediaProjectionManager = fakeMediaProjectionManager.mediaProjectionManager,
+            handler = Handler.getMain(),
+            applicationScope = testScope.backgroundScope,
+            tasksRepository = tasksRepo,
+        )
+
     private val interactor = TaskSwitchInteractor(mediaRepo, tasksRepo)
 
     @Test
     fun taskSwitchChanges_notProjecting_foregroundTaskChange_emitsNotProjectingTask() =
         testScope.runTest {
-            mediaRepo.stopProjecting()
+            val backgroundTask = createTask(taskId = 0)
+            val foregroundTask = createTask(taskId = 1)
             val taskSwitchState by collectLastValue(interactor.taskSwitchChanges)
 
-            fakeActivityTaskManager.moveTaskToForeground(createTask(taskId = 1))
+            fakeActivityTaskManager.addRunningTasks(backgroundTask, foregroundTask)
+            fakeMediaProjectionManager.dispatchOnStop()
+            fakeActivityTaskManager.moveTaskToForeground(foregroundTask)
 
             assertThat(taskSwitchState).isEqualTo(TaskSwitchState.NotProjectingTask)
         }
@@ -67,10 +82,15 @@ class TaskSwitchInteractorTest : SysuiTestCase() {
     @Test
     fun taskSwitchChanges_projectingScreen_foregroundTaskChange_emitsNotProjectingTask() =
         testScope.runTest {
-            mediaRepo.projectEntireScreen()
+            val backgroundTask = createTask(taskId = 0)
+            val foregroundTask = createTask(taskId = 1)
             val taskSwitchState by collectLastValue(interactor.taskSwitchChanges)
 
-            fakeActivityTaskManager.moveTaskToForeground(createTask(taskId = 1))
+            fakeActivityTaskManager.addRunningTasks(backgroundTask, foregroundTask)
+            fakeMediaProjectionManager.dispatchOnSessionSet(
+                session = FakeMediaProjectionManager.createDisplaySession()
+            )
+            fakeActivityTaskManager.moveTaskToForeground(foregroundTask)
 
             assertThat(taskSwitchState).isEqualTo(TaskSwitchState.NotProjectingTask)
         }
@@ -80,9 +100,12 @@ class TaskSwitchInteractorTest : SysuiTestCase() {
         testScope.runTest {
             val projectedTask = createTask(taskId = 0)
             val foregroundTask = createTask(taskId = 1)
-            mediaRepo.switchProjectedTask(projectedTask)
             val taskSwitchState by collectLastValue(interactor.taskSwitchChanges)
 
+            fakeActivityTaskManager.addRunningTasks(projectedTask, foregroundTask)
+            fakeMediaProjectionManager.dispatchOnSessionSet(
+                session = createSingleTaskSession(token = projectedTask.token.asBinder())
+            )
             fakeActivityTaskManager.moveTaskToForeground(foregroundTask)
 
             assertThat(taskSwitchState)
@@ -99,9 +122,12 @@ class TaskSwitchInteractorTest : SysuiTestCase() {
         testScope.runTest {
             val projectedTask = createTask(taskId = 0)
             val foregroundTask = createTask(taskId = 1, baseIntent = LAUNCHER_INTENT)
-            mediaRepo.switchProjectedTask(projectedTask)
             val taskSwitchState by collectLastValue(interactor.taskSwitchChanges)
 
+            fakeActivityTaskManager.addRunningTasks(projectedTask, foregroundTask)
+            fakeMediaProjectionManager.dispatchOnSessionSet(
+                session = createSingleTaskSession(projectedTask.token.asBinder())
+            )
             fakeActivityTaskManager.moveTaskToForeground(foregroundTask)
 
             assertThat(taskSwitchState).isEqualTo(TaskSwitchState.TaskUnchanged)
@@ -111,11 +137,13 @@ class TaskSwitchInteractorTest : SysuiTestCase() {
     fun taskSwitchChanges_projectingTask_foregroundTaskSame_emitsTaskUnchanged() =
         testScope.runTest {
             val projectedTask = createTask(taskId = 0)
-            val foregroundTask = createTask(taskId = 0)
-            mediaRepo.switchProjectedTask(projectedTask)
             val taskSwitchState by collectLastValue(interactor.taskSwitchChanges)
 
-            fakeActivityTaskManager.moveTaskToForeground(foregroundTask)
+            fakeActivityTaskManager.addRunningTasks(projectedTask)
+            fakeMediaProjectionManager.dispatchOnSessionSet(
+                session = createSingleTaskSession(projectedTask.token.asBinder())
+            )
+            fakeActivityTaskManager.moveTaskToForeground(projectedTask)
 
             assertThat(taskSwitchState).isEqualTo(TaskSwitchState.TaskUnchanged)
         }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/ui/TaskSwitcherNotificationCoordinatorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/ui/TaskSwitcherNotificationCoordinatorTest.kt
index cfbbf768e1fe..b396cafb28d5 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/ui/TaskSwitcherNotificationCoordinatorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/ui/TaskSwitcherNotificationCoordinatorTest.kt
@@ -18,13 +18,15 @@ package com.android.systemui.mediaprojection.taskswitcher.ui
 
 import android.app.Notification
 import android.app.NotificationManager
+import android.os.Handler
 import android.testing.AndroidTestingRunner
 import androidx.test.filters.SmallTest
 import com.android.systemui.R
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.mediaprojection.taskswitcher.data.repository.ActivityTaskManagerTasksRepository
 import com.android.systemui.mediaprojection.taskswitcher.data.repository.FakeActivityTaskManager
-import com.android.systemui.mediaprojection.taskswitcher.data.repository.FakeMediaProjectionRepository
+import com.android.systemui.mediaprojection.taskswitcher.data.repository.FakeMediaProjectionManager
+import com.android.systemui.mediaprojection.taskswitcher.data.repository.MediaProjectionManagerRepository
 import com.android.systemui.mediaprojection.taskswitcher.domain.interactor.TaskSwitchInteractor
 import com.android.systemui.mediaprojection.taskswitcher.ui.viewmodel.TaskSwitcherNotificationViewModel
 import com.android.systemui.util.mockito.any
@@ -51,14 +53,25 @@ class TaskSwitcherNotificationCoordinatorTest : SysuiTestCase() {
 
     private val dispatcher = UnconfinedTestDispatcher()
     private val testScope = TestScope(dispatcher)
+
     private val fakeActivityTaskManager = FakeActivityTaskManager()
-    private val mediaRepo = FakeMediaProjectionRepository()
+    private val fakeMediaProjectionManager = FakeMediaProjectionManager()
+
     private val tasksRepo =
         ActivityTaskManagerTasksRepository(
             activityTaskManager = fakeActivityTaskManager.activityTaskManager,
             applicationScope = testScope.backgroundScope,
             backgroundDispatcher = dispatcher
         )
+
+    private val mediaRepo =
+        MediaProjectionManagerRepository(
+            mediaProjectionManager = fakeMediaProjectionManager.mediaProjectionManager,
+            handler = Handler.getMain(),
+            applicationScope = testScope.backgroundScope,
+            tasksRepository = tasksRepo,
+        )
+
     private val interactor = TaskSwitchInteractor(mediaRepo, tasksRepo)
     private val viewModel = TaskSwitcherNotificationViewModel(interactor)
 
@@ -90,7 +103,7 @@ class TaskSwitcherNotificationCoordinatorTest : SysuiTestCase() {
     @Test
     fun hideNotification() {
         testScope.runTest {
-            mediaRepo.stopProjecting()
+            fakeMediaProjectionManager.dispatchOnStop()
 
             verify(notificationManager).cancel(any())
         }
@@ -99,7 +112,7 @@ class TaskSwitcherNotificationCoordinatorTest : SysuiTestCase() {
     @Test
     fun notificationIdIsConsistent() {
         testScope.runTest {
-            mediaRepo.stopProjecting()
+            fakeMediaProjectionManager.dispatchOnStop()
             val idCancel = argumentCaptor<Int>()
             verify(notificationManager).cancel(idCancel.capture())
 
@@ -114,7 +127,11 @@ class TaskSwitcherNotificationCoordinatorTest : SysuiTestCase() {
     private fun switchTask() {
         val projectedTask = FakeActivityTaskManager.createTask(taskId = 1)
         val foregroundTask = FakeActivityTaskManager.createTask(taskId = 2)
-        mediaRepo.switchProjectedTask(projectedTask)
+        fakeActivityTaskManager.addRunningTasks(projectedTask, foregroundTask)
+        fakeMediaProjectionManager.dispatchOnSessionSet(
+            session =
+                FakeMediaProjectionManager.createSingleTaskSession(projectedTask.token.asBinder())
+        )
         fakeActivityTaskManager.moveTaskToForeground(foregroundTask)
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/ui/viewmodel/TaskSwitcherNotificationViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/ui/viewmodel/TaskSwitcherNotificationViewModelTest.kt
index ea44fb3b1f6e..7d38de4bc2d3 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/ui/viewmodel/TaskSwitcherNotificationViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/taskswitcher/ui/viewmodel/TaskSwitcherNotificationViewModelTest.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.mediaprojection.taskswitcher.ui.viewmodel
 
 import android.content.Intent
+import android.os.Handler
 import android.testing.AndroidTestingRunner
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
@@ -24,7 +25,10 @@ import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.mediaprojection.taskswitcher.data.repository.ActivityTaskManagerTasksRepository
 import com.android.systemui.mediaprojection.taskswitcher.data.repository.FakeActivityTaskManager
 import com.android.systemui.mediaprojection.taskswitcher.data.repository.FakeActivityTaskManager.Companion.createTask
-import com.android.systemui.mediaprojection.taskswitcher.data.repository.FakeMediaProjectionRepository
+import com.android.systemui.mediaprojection.taskswitcher.data.repository.FakeMediaProjectionManager
+import com.android.systemui.mediaprojection.taskswitcher.data.repository.FakeMediaProjectionManager.Companion.createDisplaySession
+import com.android.systemui.mediaprojection.taskswitcher.data.repository.FakeMediaProjectionManager.Companion.createSingleTaskSession
+import com.android.systemui.mediaprojection.taskswitcher.data.repository.MediaProjectionManagerRepository
 import com.android.systemui.mediaprojection.taskswitcher.domain.interactor.TaskSwitchInteractor
 import com.android.systemui.mediaprojection.taskswitcher.ui.model.TaskSwitcherNotificationUiState
 import com.google.common.truth.Truth.assertThat
@@ -44,13 +48,23 @@ class TaskSwitcherNotificationViewModelTest : SysuiTestCase() {
     private val testScope = TestScope(dispatcher)
 
     private val fakeActivityTaskManager = FakeActivityTaskManager()
-    private val mediaRepo = FakeMediaProjectionRepository()
+    private val fakeMediaProjectionManager = FakeMediaProjectionManager()
+
     private val tasksRepo =
         ActivityTaskManagerTasksRepository(
             activityTaskManager = fakeActivityTaskManager.activityTaskManager,
             applicationScope = testScope.backgroundScope,
             backgroundDispatcher = dispatcher
         )
+
+    private val mediaRepo =
+        MediaProjectionManagerRepository(
+            mediaProjectionManager = fakeMediaProjectionManager.mediaProjectionManager,
+            handler = Handler.getMain(),
+            applicationScope = testScope.backgroundScope,
+            tasksRepository = tasksRepo,
+        )
+
     private val interactor = TaskSwitchInteractor(mediaRepo, tasksRepo)
 
     private val viewModel = TaskSwitcherNotificationViewModel(interactor)
@@ -58,19 +72,23 @@ class TaskSwitcherNotificationViewModelTest : SysuiTestCase() {
     @Test
     fun uiState_notProjecting_emitsNotShowing() =
         testScope.runTest {
-            mediaRepo.stopProjecting()
             val uiState by collectLastValue(viewModel.uiState)
 
+            fakeMediaProjectionManager.dispatchOnStop()
+
             assertThat(uiState).isEqualTo(TaskSwitcherNotificationUiState.NotShowing)
         }
 
     @Test
     fun uiState_notProjecting_foregroundTaskChanged_emitsNotShowing() =
         testScope.runTest {
-            mediaRepo.stopProjecting()
+            val backgroundTask = createTask(taskId = 0)
+            val foregroundTask = createTask(taskId = 1)
             val uiState by collectLastValue(viewModel.uiState)
 
-            mediaRepo.switchProjectedTask(createTask(taskId = 1))
+            fakeActivityTaskManager.addRunningTasks(backgroundTask, foregroundTask)
+            fakeMediaProjectionManager.dispatchOnStop()
+            fakeActivityTaskManager.moveTaskToForeground(foregroundTask)
 
             assertThat(uiState).isEqualTo(TaskSwitcherNotificationUiState.NotShowing)
         }
@@ -78,19 +96,23 @@ class TaskSwitcherNotificationViewModelTest : SysuiTestCase() {
     @Test
     fun uiState_projectingEntireScreen_emitsNotShowing() =
         testScope.runTest {
-            mediaRepo.projectEntireScreen()
             val uiState by collectLastValue(viewModel.uiState)
 
+            fakeMediaProjectionManager.dispatchOnSessionSet(session = createDisplaySession())
+
             assertThat(uiState).isEqualTo(TaskSwitcherNotificationUiState.NotShowing)
         }
 
     @Test
     fun uiState_projectingEntireScreen_foregroundTaskChanged_emitsNotShowing() =
         testScope.runTest {
-            mediaRepo.projectEntireScreen()
+            val backgroundTask = createTask(taskId = 0)
+            val foregroundTask = createTask(taskId = 1)
             val uiState by collectLastValue(viewModel.uiState)
 
-            mediaRepo.switchProjectedTask(createTask(taskId = 1))
+            fakeActivityTaskManager.addRunningTasks(backgroundTask, foregroundTask)
+            fakeMediaProjectionManager.dispatchOnSessionSet(session = createDisplaySession())
+            fakeActivityTaskManager.moveTaskToForeground(foregroundTask)
 
             assertThat(uiState).isEqualTo(TaskSwitcherNotificationUiState.NotShowing)
         }
@@ -100,9 +122,12 @@ class TaskSwitcherNotificationViewModelTest : SysuiTestCase() {
         testScope.runTest {
             val projectedTask = createTask(taskId = 1)
             val foregroundTask = createTask(taskId = 2)
-            mediaRepo.switchProjectedTask(projectedTask)
             val uiState by collectLastValue(viewModel.uiState)
 
+            fakeActivityTaskManager.addRunningTasks(projectedTask, foregroundTask)
+            fakeMediaProjectionManager.dispatchOnSessionSet(
+                session = createSingleTaskSession(projectedTask.token.asBinder())
+            )
             fakeActivityTaskManager.moveTaskToForeground(foregroundTask)
 
             assertThat(uiState)
@@ -113,9 +138,12 @@ class TaskSwitcherNotificationViewModelTest : SysuiTestCase() {
     fun uiState_projectingTask_foregroundTaskChanged_same_emitsNotShowing() =
         testScope.runTest {
             val projectedTask = createTask(taskId = 1)
-            mediaRepo.switchProjectedTask(projectedTask)
             val uiState by collectLastValue(viewModel.uiState)
 
+            fakeActivityTaskManager.addRunningTasks(projectedTask)
+            fakeMediaProjectionManager.dispatchOnSessionSet(
+                session = createSingleTaskSession(projectedTask.token.asBinder())
+            )
             fakeActivityTaskManager.moveTaskToForeground(projectedTask)
 
             assertThat(uiState).isEqualTo(TaskSwitcherNotificationUiState.NotShowing)
@@ -126,9 +154,12 @@ class TaskSwitcherNotificationViewModelTest : SysuiTestCase() {
         testScope.runTest {
             val projectedTask = createTask(taskId = 1)
             val foregroundTask = createTask(taskId = 2, baseIntent = LAUNCHER_INTENT)
-            mediaRepo.switchProjectedTask(projectedTask)
             val uiState by collectLastValue(viewModel.uiState)
 
+            fakeActivityTaskManager.addRunningTasks(projectedTask, foregroundTask)
+            fakeMediaProjectionManager.dispatchOnSessionSet(
+                session = createSingleTaskSession(projectedTask.token.asBinder())
+            )
             fakeActivityTaskManager.moveTaskToForeground(foregroundTask)
 
             assertThat(uiState).isEqualTo(TaskSwitcherNotificationUiState.NotShowing)