diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt
index 250dc7b924..d40f8ab389 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt
@@ -78,7 +78,7 @@ class SplitAnimationControllerTest {
     private val mockSplitSourceDrawable: Drawable = mock()
     private val mockSplitSourceView: View = mock()
 
-    private val stateManager: StateManager<*> = mock()
+    private val stateManager: StateManager<*, *> = mock()
     private val depthController: DepthController = mock()
     private val transitionInfo: TransitionInfo = mock()
     private val transaction: Transaction = mock()
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
index aa08ca4c05..bab84ef28b 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
@@ -31,6 +31,7 @@ import com.android.launcher3.logging.StatsLogManager.StatsLogger
 import com.android.launcher3.model.data.ItemInfo
 import com.android.launcher3.statehandlers.DepthController
 import com.android.launcher3.statemanager.StateManager
+import com.android.launcher3.statemanager.StatefulActivity
 import com.android.launcher3.util.ComponentKey
 import com.android.launcher3.util.SplitConfigurationOptions
 import com.android.quickstep.RecentsModel
@@ -61,7 +62,7 @@ class SplitSelectStateControllerTest {
     private val depthController: DepthController = mock()
     private val statsLogManager: StatsLogManager = mock()
     private val statsLogger: StatsLogger = mock()
-    private val stateManager: StateManager<LauncherState> = mock()
+    private val stateManager: StateManager<LauncherState, StatefulActivity<LauncherState>> = mock()
     private val handler: Handler = mock()
     private val context: RecentsViewContainer = mock()
     private val recentsModel: RecentsModel = mock()
diff --git a/quickstep/tests/src/com/android/launcher3/taskbar/FallbackTaskbarUIControllerTest.kt b/quickstep/tests/src/com/android/launcher3/taskbar/FallbackTaskbarUIControllerTest.kt
index f3115c6a52..04012c027d 100644
--- a/quickstep/tests/src/com/android/launcher3/taskbar/FallbackTaskbarUIControllerTest.kt
+++ b/quickstep/tests/src/com/android/launcher3/taskbar/FallbackTaskbarUIControllerTest.kt
@@ -37,7 +37,7 @@ class FallbackTaskbarUIControllerTest : TaskbarBaseTestCase() {
     lateinit var stateListener: StateManager.StateListener<RecentsState>
 
     private val recentsActivity: RecentsActivity = mock()
-    private val stateManager: StateManager<RecentsState> = mock()
+    private val stateManager: StateManager<RecentsState, RecentsActivity> = mock()
 
     @Before
     override fun setup() {