diff --git a/quickstep/tests/src/com/android/launcher3/taskbar/FallbackTaskbarUIControllerTest.kt b/quickstep/tests/src/com/android/launcher3/taskbar/FallbackTaskbarUIControllerTest.kt
index 8c13fe3..f3115c6 100644
--- a/quickstep/tests/src/com/android/launcher3/taskbar/FallbackTaskbarUIControllerTest.kt
+++ b/quickstep/tests/src/com/android/launcher3/taskbar/FallbackTaskbarUIControllerTest.kt
@@ -24,11 +24,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.ArgumentCaptor
-import org.mockito.Mock
-import org.mockito.Mockito.times
-import org.mockito.Mockito.verify
-import org.mockito.Mockito.`when` as whenever
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.times
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
 
 @RunWith(AndroidJUnit4::class)
 class FallbackTaskbarUIControllerTest : TaskbarBaseTestCase() {
@@ -36,8 +36,8 @@
     lateinit var fallbackTaskbarUIController: FallbackTaskbarUIController
     lateinit var stateListener: StateManager.StateListener<RecentsState>
 
-    @Mock lateinit var recentsActivity: RecentsActivity
-    @Mock lateinit var stateManager: StateManager<RecentsState>
+    private val recentsActivity: RecentsActivity = mock()
+    private val stateManager: StateManager<RecentsState> = mock()
 
     @Before
     override fun setup() {
@@ -46,10 +46,10 @@
         fallbackTaskbarUIController = FallbackTaskbarUIController(recentsActivity)
 
         // Capture registered state listener to send events to in our tests
-        val captor = ArgumentCaptor.forClass(StateManager.StateListener::class.java)
+        val captor = argumentCaptor<StateManager.StateListener<RecentsState>>()
         fallbackTaskbarUIController.init(taskbarControllers)
         verify(stateManager).addStateListener(captor.capture())
-        stateListener = captor.value as StateManager.StateListener<RecentsState>
+        stateListener = captor.lastValue
     }
 
     @Test
diff --git a/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarKeyguardControllerTest.kt b/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarKeyguardControllerTest.kt
index 148e36c..ed88c29 100644
--- a/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarKeyguardControllerTest.kt
+++ b/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarKeyguardControllerTest.kt
@@ -23,17 +23,17 @@
 import com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING_OCCLUDED
 import org.junit.Before
 import org.junit.Test
-import org.mockito.Mock
-import org.mockito.Mockito.anyBoolean
-import org.mockito.Mockito.never
-import org.mockito.Mockito.times
-import org.mockito.Mockito.verify
-import org.mockito.Mockito.`when` as whenever
+import org.mockito.kotlin.any
+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
 
 class TaskbarKeyguardControllerTest : TaskbarBaseTestCase() {
 
-    @Mock lateinit var baseDragLayer: TaskbarDragLayer
-    @Mock lateinit var keyguardManager: KeyguardManager
+    private val baseDragLayer: TaskbarDragLayer = mock()
+    private val keyguardManager: KeyguardManager = mock()
 
     @Before
     override fun setup() {
@@ -50,7 +50,7 @@
     @Test
     fun uninterestingFlags_noActions() {
         setFlags(0)
-        verify(navbarButtonsViewController, never()).setKeyguardVisible(anyBoolean(), anyBoolean())
+        verify(navbarButtonsViewController, never()).setKeyguardVisible(any(), any())
     }
 
     @Test
diff --git a/quickstep/tests/src/com/android/launcher3/taskbar/navbutton/NavButtonLayoutFactoryTest.kt b/quickstep/tests/src/com/android/launcher3/taskbar/navbutton/NavButtonLayoutFactoryTest.kt
index 16bfe70..37fcf43 100644
--- a/quickstep/tests/src/com/android/launcher3/taskbar/navbutton/NavButtonLayoutFactoryTest.kt
+++ b/quickstep/tests/src/com/android/launcher3/taskbar/navbutton/NavButtonLayoutFactoryTest.kt
@@ -19,32 +19,29 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.Mockito.`when` as whenever
-import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
 @RunWith(AndroidJUnit4::class)
 class NavButtonLayoutFactoryTest {
 
-    @Mock lateinit var mockDeviceProfile: DeviceProfile
-    @Mock lateinit var mockParentButtonContainer: FrameLayout
-    @Mock lateinit var mockNavLayout: LinearLayout
-    @Mock lateinit var mockStartContextualLayout: ViewGroup
-    @Mock lateinit var mockEndContextualLayout: ViewGroup
-    @Mock lateinit var mockResources: Resources
-    @Mock lateinit var mockBackButton: ImageView
-    @Mock lateinit var mockRecentsButton: ImageView
-    @Mock lateinit var mockHomeButton: ImageView
-    @Mock lateinit var mockImeSwitcher: ImageView
-    @Mock lateinit var mockRotationButton: RotationButton
-    @Mock lateinit var mockA11yButton: ImageView
+    private val mockDeviceProfile: DeviceProfile = mock()
+    private val mockParentButtonContainer: FrameLayout = mock()
+    private val mockNavLayout: LinearLayout = mock()
+    private val mockStartContextualLayout: ViewGroup = mock()
+    private val mockEndContextualLayout: ViewGroup = mock()
+    private val mockResources: Resources = mock()
+    private val mockBackButton: ImageView = mock()
+    private val mockRecentsButton: ImageView = mock()
+    private val mockHomeButton: ImageView = mock()
+    private val mockImeSwitcher: ImageView = mock()
+    private val mockRotationButton: RotationButton = mock()
+    private val mockA11yButton: ImageView = mock()
 
     private var surfaceRotation = Surface.ROTATION_0
 
     @Before
     fun setup() {
-        MockitoAnnotations.initMocks(this)
-
         // Init end nav buttons
         whenever(mockNavLayout.childCount).thenReturn(3)
         whenever(mockNavLayout.findViewById<View>(R.id.back)).thenReturn(mockBackButton)
@@ -155,13 +152,13 @@
         mockDeviceProfile.isTaskbarPresent = false
         setDeviceProfileLandscape()
         val layoutter: NavButtonLayoutFactory.NavButtonLayoutter =
-                getLayoutter(
-                        isKidsMode = false,
-                        isInSetup = false,
-                        isThreeButtonNav = true,
-                        phoneMode = true,
-                        surfaceRotation = ROTATION_270
-                )
+            getLayoutter(
+                isKidsMode = false,
+                isInSetup = false,
+                isThreeButtonNav = true,
+                phoneMode = true,
+                surfaceRotation = ROTATION_270
+            )
         assert(layoutter is PhoneSeascapeNavLayoutter)
     }
 
diff --git a/quickstep/tests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt b/quickstep/tests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt
index 7e07b81..50803fe 100644
--- a/quickstep/tests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt
@@ -32,39 +32,36 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.MockitoAnnotations
-import org.mockito.Mockito.`when` as whenever
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
 @RunWith(AndroidJUnit4::class)
 class SplitAnimationControllerTest {
 
     private val taskId = 9
 
-    @Mock lateinit var mockSplitSelectStateController: SplitSelectStateController
+    private val mockSplitSelectStateController: SplitSelectStateController = mock()
     // TaskView
-    @Mock lateinit var mockTaskView: TaskView
-    @Mock lateinit var mockThumbnailView: TaskThumbnailView
-    @Mock lateinit var mockBitmap: Bitmap
-    @Mock lateinit var mockIconView: IconView
-    @Mock lateinit var mockTaskViewDrawable: Drawable
+    private val mockTaskView: TaskView = mock()
+    private val mockThumbnailView: TaskThumbnailView = mock()
+    private val mockBitmap: Bitmap = mock()
+    private val mockIconView: IconView = mock()
+    private val mockTaskViewDrawable: Drawable = mock()
     // GroupedTaskView
-    @Mock lateinit var mockGroupedTaskView: GroupedTaskView
-    @Mock lateinit var mockTask: Task
-    @Mock lateinit var mockTaskKey: Task.TaskKey
-    @Mock lateinit var mockTaskIdAttributeContainer: TaskIdAttributeContainer
+    private val mockGroupedTaskView: GroupedTaskView = mock()
+    private val mockTask: Task = mock()
+    private val mockTaskKey: Task.TaskKey = mock()
+    private val mockTaskIdAttributeContainer: TaskIdAttributeContainer = mock()
 
     // SplitSelectSource
-    @Mock lateinit var splitSelectSource: SplitConfigurationOptions.SplitSelectSource
-    @Mock lateinit var mockSplitSourceDrawable: Drawable
-    @Mock lateinit var mockSplitSourceView: View
+    private val splitSelectSource: SplitConfigurationOptions.SplitSelectSource = mock()
+    private val mockSplitSourceDrawable: Drawable = mock()
+    private val mockSplitSourceView: View = mock()
 
     lateinit var splitAnimationController: SplitAnimationController
 
     @Before
     fun setup() {
-        MockitoAnnotations.initMocks(this)
-
         whenever(mockTaskView.thumbnail).thenReturn(mockThumbnailView)
         whenever(mockThumbnailView.thumbnail).thenReturn(mockBitmap)
         whenever(mockTaskView.iconView).thenReturn(mockIconView)
@@ -85,12 +82,14 @@
         // Missing taskView icon
         whenever(mockIconView.drawable).thenReturn(null)
 
-        val splitAnimInitProps : SplitAnimationController.Companion.SplitAnimInitProps =
-                splitAnimationController.getFirstAnimInitViews(
-                        { mockTaskView }, { splitSelectSource })
+        val splitAnimInitProps: SplitAnimationController.Companion.SplitAnimInitProps =
+            splitAnimationController.getFirstAnimInitViews({ mockTaskView }, { splitSelectSource })
 
-        assertEquals("Did not fallback to use splitSource icon drawable",
-                mockSplitSourceDrawable, splitAnimInitProps.iconDrawable)
+        assertEquals(
+            "Did not fallback to use splitSource icon drawable",
+            mockSplitSourceDrawable,
+            splitAnimInitProps.iconDrawable
+        )
     }
 
     @Test
@@ -99,12 +98,14 @@
         whenever(mockSplitSelectStateController.isAnimateCurrentTaskDismissal).thenReturn(true)
         whenever(mockSplitSelectStateController.isDismissingFromSplitPair).thenReturn(false)
 
-        val splitAnimInitProps : SplitAnimationController.Companion.SplitAnimInitProps =
-                splitAnimationController.getFirstAnimInitViews(
-                        { mockTaskView }, { splitSelectSource })
+        val splitAnimInitProps: SplitAnimationController.Companion.SplitAnimInitProps =
+            splitAnimationController.getFirstAnimInitViews({ mockTaskView }, { splitSelectSource })
 
-        assertEquals("Did not use taskView icon drawable", mockTaskViewDrawable,
-                splitAnimInitProps.iconDrawable)
+        assertEquals(
+            "Did not use taskView icon drawable",
+            mockTaskViewDrawable,
+            splitAnimInitProps.iconDrawable
+        )
     }
 
     @Test
@@ -116,12 +117,14 @@
         // Set split source to null
         whenever(splitSelectSource.drawable).thenReturn(null)
 
-        val splitAnimInitProps : SplitAnimationController.Companion.SplitAnimInitProps =
-                splitAnimationController.getFirstAnimInitViews(
-                        { mockTaskView }, { splitSelectSource })
+        val splitAnimInitProps: SplitAnimationController.Companion.SplitAnimInitProps =
+            splitAnimationController.getFirstAnimInitViews({ mockTaskView }, { splitSelectSource })
 
-        assertEquals("Did not use taskView icon drawable", mockTaskViewDrawable,
-                splitAnimInitProps.iconDrawable)
+        assertEquals(
+            "Did not use taskView icon drawable",
+            mockTaskViewDrawable,
+            splitAnimInitProps.iconDrawable
+        )
     }
 
     @Test
@@ -130,12 +133,14 @@
         whenever(mockSplitSelectStateController.isAnimateCurrentTaskDismissal).thenReturn(false)
         whenever(mockSplitSelectStateController.isDismissingFromSplitPair).thenReturn(false)
 
-        val splitAnimInitProps : SplitAnimationController.Companion.SplitAnimInitProps =
-                splitAnimationController.getFirstAnimInitViews(
-                        { mockTaskView }, { splitSelectSource })
+        val splitAnimInitProps: SplitAnimationController.Companion.SplitAnimInitProps =
+            splitAnimationController.getFirstAnimInitViews({ mockTaskView }, { splitSelectSource })
 
-        assertEquals("Did not use splitSource icon drawable", mockSplitSourceDrawable,
-                splitAnimInitProps.iconDrawable)
+        assertEquals(
+            "Did not use splitSource icon drawable",
+            mockSplitSourceDrawable,
+            splitAnimInitProps.iconDrawable
+        )
     }
 
     @Test
@@ -154,12 +159,17 @@
         whenever(mockTaskKey.getId()).thenReturn(taskId)
         whenever(mockSplitSelectStateController.initialTaskId).thenReturn(taskId)
         whenever(mockGroupedTaskView.taskIdAttributeContainers)
-                .thenReturn(Array(1) { mockTaskIdAttributeContainer })
-        val splitAnimInitProps : SplitAnimationController.Companion.SplitAnimInitProps =
-                splitAnimationController.getFirstAnimInitViews(
-                        { mockGroupedTaskView }, { splitSelectSource })
+            .thenReturn(Array(1) { mockTaskIdAttributeContainer })
+        val splitAnimInitProps: SplitAnimationController.Companion.SplitAnimInitProps =
+            splitAnimationController.getFirstAnimInitViews(
+                { mockGroupedTaskView },
+                { splitSelectSource }
+            )
 
-        assertEquals("Did not use splitSource icon drawable", mockSplitSourceDrawable,
-                splitAnimInitProps.iconDrawable)
+        assertEquals(
+            "Did not use splitSource icon drawable",
+            mockSplitSourceDrawable,
+            splitAnimInitProps.iconDrawable
+        )
     }
-}
\ No newline at end of file
+}
diff --git a/quickstep/tests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt b/quickstep/tests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
index f198741..f292f9a 100644
--- a/quickstep/tests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
@@ -33,11 +33,11 @@
 import com.android.launcher3.statemanager.StatefulActivity
 import com.android.launcher3.util.ComponentKey
 import com.android.launcher3.util.SplitConfigurationOptions
-import com.android.launcher3.util.withArgCaptor
 import com.android.quickstep.RecentsModel
 import com.android.quickstep.SystemUiProxy
 import com.android.systemui.shared.recents.model.Task
 import com.android.wm.shell.common.split.SplitScreenConstants.SNAP_TO_50_50
+import java.util.function.Consumer
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertFalse
 import org.junit.Assert.assertNull
@@ -45,23 +45,22 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.Mockito.verify
-import org.mockito.Mockito.`when`
-import org.mockito.MockitoAnnotations
-import java.util.function.Consumer
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
 
 @RunWith(AndroidJUnit4::class)
 class SplitSelectStateControllerTest {
 
-    @Mock lateinit var systemUiProxy: SystemUiProxy
-    @Mock lateinit var depthController: DepthController
-    @Mock lateinit var statsLogManager: StatsLogManager
-    @Mock lateinit var stateManager: StateManager<LauncherState>
-    @Mock lateinit var handler: Handler
-    @Mock lateinit var context: StatefulActivity<*>
-    @Mock lateinit var recentsModel: RecentsModel
-    @Mock lateinit var pendingIntent: PendingIntent
+    private val systemUiProxy: SystemUiProxy = mock()
+    private val depthController: DepthController = mock()
+    private val statsLogManager: StatsLogManager = mock()
+    private val stateManager: StateManager<LauncherState> = mock()
+    private val handler: Handler = mock()
+    private val context: StatefulActivity<*> = mock()
+    private val recentsModel: RecentsModel = mock()
+    private val pendingIntent: PendingIntent = mock()
 
     lateinit var splitSelectStateController: SplitSelectStateController
 
@@ -69,11 +68,12 @@
     private val nonPrimaryUserHandle = UserHandle(ActivityManager.RunningTaskInfo().userId + 10)
 
     private var taskIdCounter = 0
-    private fun getUniqueId(): Int { return ++taskIdCounter }
+    private fun getUniqueId(): Int {
+        return ++taskIdCounter
+    }
 
     @Before
     fun setup() {
-        MockitoAnnotations.initMocks(this)
         splitSelectStateController =
             SplitSelectStateController(
                 context,
@@ -111,13 +111,14 @@
 
         // Capture callback from recentsModel#getTasks()
         val consumer =
-            withArgCaptor<Consumer<ArrayList<GroupTask>>> {
-                splitSelectStateController.findLastActiveTasksAndRunCallback(
-                    listOf(nonMatchingComponent),
-                    taskConsumer
-                )
-                verify(recentsModel).getTasks(capture())
-            }
+            argumentCaptor<Consumer<ArrayList<GroupTask>>> {
+                    splitSelectStateController.findLastActiveTasksAndRunCallback(
+                        listOf(nonMatchingComponent),
+                        taskConsumer
+                    )
+                    verify(recentsModel).getTasks(capture())
+                }
+                .lastValue
 
         // Send our mocked tasks
         consumer.accept(tasks)
@@ -162,13 +163,14 @@
 
         // Capture callback from recentsModel#getTasks()
         val consumer =
-            withArgCaptor<Consumer<ArrayList<GroupTask>>> {
-                splitSelectStateController.findLastActiveTasksAndRunCallback(
-                    listOf(matchingComponent),
-                    taskConsumer
-                )
-                verify(recentsModel).getTasks(capture())
-            }
+            argumentCaptor<Consumer<ArrayList<GroupTask>>> {
+                    splitSelectStateController.findLastActiveTasksAndRunCallback(
+                        listOf(matchingComponent),
+                        taskConsumer
+                    )
+                    verify(recentsModel).getTasks(capture())
+                }
+                .lastValue
 
         // Send our mocked tasks
         consumer.accept(tasks)
@@ -201,13 +203,14 @@
 
         // Capture callback from recentsModel#getTasks()
         val consumer =
-            withArgCaptor<Consumer<ArrayList<GroupTask>>> {
-                splitSelectStateController.findLastActiveTasksAndRunCallback(
-                    listOf(nonPrimaryUserComponent),
-                    taskConsumer
-                )
-                verify(recentsModel).getTasks(capture())
-            }
+            argumentCaptor<Consumer<ArrayList<GroupTask>>> {
+                    splitSelectStateController.findLastActiveTasksAndRunCallback(
+                        listOf(nonPrimaryUserComponent),
+                        taskConsumer
+                    )
+                    verify(recentsModel).getTasks(capture())
+                }
+                .lastValue
 
         // Send our mocked tasks
         consumer.accept(tasks)
@@ -255,13 +258,14 @@
 
         // Capture callback from recentsModel#getTasks()
         val consumer =
-            withArgCaptor<Consumer<ArrayList<GroupTask>>> {
-                splitSelectStateController.findLastActiveTasksAndRunCallback(
-                    listOf(nonPrimaryUserComponent),
-                    taskConsumer
-                )
-                verify(recentsModel).getTasks(capture())
-            }
+            argumentCaptor<Consumer<ArrayList<GroupTask>>> {
+                    splitSelectStateController.findLastActiveTasksAndRunCallback(
+                        listOf(nonPrimaryUserComponent),
+                        taskConsumer
+                    )
+                    verify(recentsModel).getTasks(capture())
+                }
+                .lastValue
 
         // Send our mocked tasks
         consumer.accept(tasks)
@@ -306,13 +310,14 @@
 
         // Capture callback from recentsModel#getTasks()
         val consumer =
-            withArgCaptor<Consumer<ArrayList<GroupTask>>> {
-                splitSelectStateController.findLastActiveTasksAndRunCallback(
-                    listOf(matchingComponent),
-                    taskConsumer
-                )
-                verify(recentsModel).getTasks(capture())
-            }
+            argumentCaptor<Consumer<ArrayList<GroupTask>>> {
+                    splitSelectStateController.findLastActiveTasksAndRunCallback(
+                        listOf(matchingComponent),
+                        taskConsumer
+                    )
+                    verify(recentsModel).getTasks(capture())
+                }
+                .lastValue
 
         // Send our mocked tasks
         consumer.accept(tasks)
@@ -327,10 +332,7 @@
             ComponentKey(ComponentName(matchingPackage, matchingClass), primaryUserHandle)
 
         val groupTask1 =
-            generateGroupTask(
-                ComponentName("hotdog", "pie"),
-                ComponentName("pumpkin", "pie")
-            )
+            generateGroupTask(ComponentName("hotdog", "pie"), ComponentName("pumpkin", "pie"))
         val groupTask2 =
             generateGroupTask(
                 ComponentName("pomegranate", "juice"),
@@ -361,13 +363,14 @@
 
         // Capture callback from recentsModel#getTasks()
         val consumer =
-            withArgCaptor<Consumer<ArrayList<GroupTask>>> {
-                splitSelectStateController.findLastActiveTasksAndRunCallback(
-                    listOf(nonMatchingComponent, matchingComponent),
-                    taskConsumer
-                )
-                verify(recentsModel).getTasks(capture())
-            }
+            argumentCaptor<Consumer<ArrayList<GroupTask>>> {
+                    splitSelectStateController.findLastActiveTasksAndRunCallback(
+                        listOf(nonMatchingComponent, matchingComponent),
+                        taskConsumer
+                    )
+                    verify(recentsModel).getTasks(capture())
+                }
+                .lastValue
 
         // Send our mocked tasks
         consumer.accept(tasks)
@@ -381,10 +384,7 @@
             ComponentKey(ComponentName(matchingPackage, matchingClass), primaryUserHandle)
 
         val groupTask1 =
-            generateGroupTask(
-                ComponentName("hotdog", "pie"),
-                ComponentName("pumpkin", "pie")
-            )
+            generateGroupTask(ComponentName("hotdog", "pie"), ComponentName("pumpkin", "pie"))
         val groupTask2 =
             generateGroupTask(
                 ComponentName("pomegranate", "juice"),
@@ -415,13 +415,14 @@
 
         // Capture callback from recentsModel#getTasks()
         val consumer =
-            withArgCaptor<Consumer<ArrayList<GroupTask>>> {
-                splitSelectStateController.findLastActiveTasksAndRunCallback(
-                    listOf(matchingComponent, matchingComponent),
-                    taskConsumer
-                )
-                verify(recentsModel).getTasks(capture())
-            }
+            argumentCaptor<Consumer<ArrayList<GroupTask>>> {
+                    splitSelectStateController.findLastActiveTasksAndRunCallback(
+                        listOf(matchingComponent, matchingComponent),
+                        taskConsumer
+                    )
+                    verify(recentsModel).getTasks(capture())
+                }
+                .lastValue
 
         // Send our mocked tasks
         consumer.accept(tasks)
@@ -479,13 +480,14 @@
 
         // Capture callback from recentsModel#getTasks()
         val consumer =
-            withArgCaptor<Consumer<ArrayList<GroupTask>>> {
-                splitSelectStateController.findLastActiveTasksAndRunCallback(
-                    listOf(matchingComponent, matchingComponent),
-                    taskConsumer
-                )
-                verify(recentsModel).getTasks(capture())
-            }
+            argumentCaptor<Consumer<ArrayList<GroupTask>>> {
+                    splitSelectStateController.findLastActiveTasksAndRunCallback(
+                        listOf(matchingComponent, matchingComponent),
+                        taskConsumer
+                    )
+                    verify(recentsModel).getTasks(capture())
+                }
+                .lastValue
 
         // Send our mocked tasks
         consumer.accept(tasks)
@@ -531,7 +533,7 @@
     @Test
     fun secondPendingIntentSet() {
         val itemInfo = ItemInfo()
-        `when`(pendingIntent.creatorUserHandle).thenReturn(primaryUserHandle)
+        whenever(pendingIntent.creatorUserHandle).thenReturn(primaryUserHandle)
         splitSelectStateController.setInitialTaskSelect(null, 0, itemInfo, null, 1)
         splitSelectStateController.setSecondTask(pendingIntent)
         assertTrue(splitSelectStateController.isBothSplitAppsConfirmed)
diff --git a/quickstep/tests/src/com/android/quickstep/util/unfold/PreemptiveUnfoldTransitionProgressProviderTest.kt b/quickstep/tests/src/com/android/quickstep/util/unfold/PreemptiveUnfoldTransitionProgressProviderTest.kt
index f73be72..6a418a4 100644
--- a/quickstep/tests/src/com/android/quickstep/util/unfold/PreemptiveUnfoldTransitionProgressProviderTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/util/unfold/PreemptiveUnfoldTransitionProgressProviderTest.kt
@@ -21,19 +21,17 @@
 import android.testing.TestableLooper.RunWithLooper
 import android.util.Log
 import androidx.test.filters.SmallTest
-import com.android.launcher3.util.any
-import com.android.launcher3.util.mock
 import com.android.systemui.unfold.UnfoldTransitionProgressProvider
 import com.android.systemui.unfold.UnfoldTransitionProgressProvider.TransitionProgressListener
 import org.junit.After
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mockito.anyBoolean
-import org.mockito.Mockito.anyFloat
-import org.mockito.Mockito.inOrder
-import org.mockito.Mockito.never
-import org.mockito.Mockito.verify
+import org.mockito.kotlin.any
+import org.mockito.kotlin.inOrder
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.never
+import org.mockito.kotlin.verify
 
 @SmallTest
 @RunWith(AndroidTestingRunner::class)
@@ -74,7 +72,7 @@
         provider.preemptivelyStartTransition(initialProgress = null)
 
         verify(listener).onTransitionStarted()
-        verify(listener, never()).onTransitionProgress(anyFloat())
+        verify(listener, never()).onTransitionProgress(any())
     }
 
     @Test
@@ -90,7 +88,7 @@
         provider.preemptivelyStartTransition()
         provider.cancelPreemptiveStart()
 
-        with(inOrder(listener)) {
+        inOrder(listener) {
             verify(listener).onTransitionStarted()
             verify(listener).onTransitionFinished()
         }
@@ -111,7 +109,7 @@
         source.onTransitionStarted()
         source.onTransitionFinished()
 
-        with(inOrder(listener)) {
+        inOrder(listener) {
             verify(listener).onTransitionStarted()
             verify(listener).onTransitionFinished()
         }
@@ -152,7 +150,7 @@
         provider.preemptivelyStartTransition()
         source.onTransitionFinished()
 
-        with(inOrder(listener)) {
+        inOrder(listener) {
             verify(listener).onTransitionStarted()
             verify(listener).onTransitionFinished()
         }
@@ -165,7 +163,7 @@
         testableLooper.moveTimeForward(PREEMPTIVE_UNFOLD_TIMEOUT_MS + 1)
         testableLooper.processAllMessages()
 
-        with(inOrder(listener)) {
+        inOrder(listener) {
             verify(listener).onTransitionStarted()
             verify(listener).onTransitionFinished()
         }
@@ -178,7 +176,7 @@
         testableLooper.moveTimeForward(PREEMPTIVE_UNFOLD_TIMEOUT_MS + 1)
         testableLooper.processAllMessages()
 
-        verify(testWtfHandler).onTerribleFailure(any(), any(), anyBoolean())
+        verify(testWtfHandler).onTerribleFailure(any(), any(), any())
     }
 
     @Test
@@ -225,7 +223,7 @@
 
         source.onTransitionFinished()
 
-        with(inOrder(listener)) {
+        inOrder(listener) {
             verify(listener).onTransitionStarted()
             verify(listener).onTransitionFinished()
         }
diff --git a/tests/Android.bp b/tests/Android.bp
index 654edad..582c7fd 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -87,6 +87,7 @@
         "androidx.test.espresso.contrib",
         "androidx.test.espresso.intents",
         "androidx.test.uiautomator_uiautomator",
+        "mockito-kotlin2",
         "mockito-target-extended-minus-junit4",
         "launcher_log_protos_lite",
         "truth-prebuilt",
diff --git a/tests/src/com/android/launcher3/AbstractDeviceProfileTest.kt b/tests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
index a52ba9e..84fa988 100644
--- a/tests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
+++ b/tests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
@@ -38,10 +38,10 @@
 import org.junit.After
 import org.junit.Before
 import org.junit.Rule
-import org.mockito.ArgumentMatchers
-import org.mockito.Mockito.mock
-import org.mockito.Mockito.spy
-import org.mockito.Mockito.`when` as whenever
+import org.mockito.kotlin.any
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.spy
+import org.mockito.kotlin.whenever
 
 /**
  * This is an abstract class for DeviceProfile tests that create an InvariantDeviceProfile based on
@@ -52,8 +52,8 @@
 abstract class AbstractDeviceProfileTest {
     protected var context: Context? = null
     protected open val runningContext: Context = ApplicationProvider.getApplicationContext()
-    private var displayController: DisplayController = mock(DisplayController::class.java)
-    private var windowManagerProxy: WindowManagerProxy = mock(WindowManagerProxy::class.java)
+    private val displayController: DisplayController = mock()
+    private val windowManagerProxy: WindowManagerProxy = mock()
     private lateinit var originalDisplayController: DisplayController
     private lateinit var originalWindowManagerProxy: WindowManagerProxy
 
@@ -288,11 +288,10 @@
     ) {
         val windowsBounds = perDisplayBoundsCache[displayInfo]!!
         val realBounds = windowsBounds[rotation]
-        whenever(windowManagerProxy.getDisplayInfo(ArgumentMatchers.any())).thenReturn(displayInfo)
-        whenever(windowManagerProxy.getRealBounds(ArgumentMatchers.any(), ArgumentMatchers.any()))
-            .thenReturn(realBounds)
-        whenever(windowManagerProxy.getRotation(ArgumentMatchers.any())).thenReturn(rotation)
-        whenever(windowManagerProxy.getNavigationMode(ArgumentMatchers.any()))
+        whenever(windowManagerProxy.getDisplayInfo(any())).thenReturn(displayInfo)
+        whenever(windowManagerProxy.getRealBounds(any(), any())).thenReturn(realBounds)
+        whenever(windowManagerProxy.getRotation(any())).thenReturn(rotation)
+        whenever(windowManagerProxy.getNavigationMode(any()))
             .thenReturn(
                 if (isGestureMode) NavigationMode.NO_BUTTON else NavigationMode.THREE_BUTTONS
             )
diff --git a/tests/src/com/android/launcher3/FakeInvariantDeviceProfileTest.kt b/tests/src/com/android/launcher3/FakeInvariantDeviceProfileTest.kt
index 42338bf..a421006 100644
--- a/tests/src/com/android/launcher3/FakeInvariantDeviceProfileTest.kt
+++ b/tests/src/com/android/launcher3/FakeInvariantDeviceProfileTest.kt
@@ -27,9 +27,9 @@
 import java.io.PrintWriter
 import java.io.StringWriter
 import org.junit.Before
-import org.mockito.ArgumentMatchers.any
-import org.mockito.Mockito.mock
-import org.mockito.Mockito.`when` as whenever
+import org.mockito.kotlin.any
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
 /**
  * This is an abstract class for DeviceProfile tests that don't need the real Context and mock an
@@ -41,7 +41,7 @@
 
     protected var context: Context? = null
     protected var inv: InvariantDeviceProfile? = null
-    protected var info: Info = mock(Info::class.java)
+    protected val info: Info = mock()
     protected var windowBounds: WindowBounds? = null
     protected var isMultiWindowMode: Boolean = false
     protected var transposeLayoutWithOrientation: Boolean = false
diff --git a/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.kt b/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.kt
index 1155227..78c61d5 100644
--- a/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.kt
+++ b/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.kt
@@ -24,20 +24,19 @@
 import com.android.launcher3.util.Executors
 import com.android.launcher3.util.IntArray
 import com.android.launcher3.util.TestUtil.runOnExecutorSync
-import com.android.launcher3.util.any
-import com.android.launcher3.util.eq
-import com.android.launcher3.util.same
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
 import org.mockito.Mockito.times
-import org.mockito.Mockito.verify
-import org.mockito.Mockito.verifyZeroInteractions
-import org.mockito.Mockito.`when` as whenever
-import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.any
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.same
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.verifyZeroInteractions
+import org.mockito.kotlin.whenever
 
 /** Tests for [AddWorkspaceItemsTask] */
 @SmallTest
@@ -46,12 +45,11 @@
 
     private lateinit var mDataModelCallbacks: MyCallbacks
 
-    @Mock private lateinit var mWorkspaceItemSpaceFinder: WorkspaceItemSpaceFinder
+    private val mWorkspaceItemSpaceFinder: WorkspaceItemSpaceFinder = mock()
 
     @Before
     override fun setup() {
         super.setup()
-        MockitoAnnotations.initMocks(this)
         mDataModelCallbacks = MyCallbacks()
         Executors.MAIN_EXECUTOR.submit { mModelHelper.model.addCallbacks(mDataModelCallbacks) }
             .get()
diff --git a/tests/src/com/android/launcher3/util/DisplayControllerTest.kt b/tests/src/com/android/launcher3/util/DisplayControllerTest.kt
index a94dd2e..8670d40 100644
--- a/tests/src/com/android/launcher3/util/DisplayControllerTest.kt
+++ b/tests/src/com/android/launcher3/util/DisplayControllerTest.kt
@@ -42,11 +42,13 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.Mockito.doNothing
-import org.mockito.Mockito.verify
-import org.mockito.Mockito.`when` as whenever
-import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.any
+import org.mockito.kotlin.anyOrNull
+import org.mockito.kotlin.doNothing
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
 import org.mockito.stubbing.Answer
 
 /** Unit tests for {@link DisplayController} */
@@ -56,13 +58,13 @@
 
     private val appContext: Context = ApplicationProvider.getApplicationContext()
 
-    @Mock private lateinit var context: SandboxContext
-    @Mock private lateinit var windowManagerProxy: WindowManagerProxy
-    @Mock private lateinit var launcherPrefs: LauncherPrefs
-    @Mock private lateinit var displayManager: DisplayManager
-    @Mock private lateinit var display: Display
-    @Mock private lateinit var resources: Resources
-    @Mock private lateinit var displayInfoChangeListener: DisplayInfoChangeListener
+    private val context: SandboxContext = mock()
+    private val windowManagerProxy: WindowManagerProxy = mock()
+    private val launcherPrefs: LauncherPrefs = mock()
+    private val displayManager: DisplayManager = mock()
+    private val display: Display = mock()
+    private val resources: Resources = mock()
+    private val displayInfoChangeListener: DisplayInfoChangeListener = mock()
 
     private lateinit var displayController: DisplayController
 
@@ -88,7 +90,6 @@
 
     @Before
     fun setUp() {
-        MockitoAnnotations.initMocks(this)
         whenever(context.getObject(eq(WindowManagerProxy.INSTANCE))).thenReturn(windowManagerProxy)
         whenever(context.getObject(eq(LauncherPrefs.INSTANCE))).thenReturn(launcherPrefs)
         whenever(launcherPrefs.get(TASKBAR_PINNING)).thenReturn(false)
@@ -112,10 +113,10 @@
 
         whenever(windowManagerProxy.getNavigationMode(any())).thenReturn(NavigationMode.NO_BUTTON)
         // Mock context
-        whenever(context.createWindowContext(any(), any(), nullable())).thenReturn(context)
+        whenever(context.createWindowContext(any(), any(), anyOrNull())).thenReturn(context)
         whenever(context.getSystemService(eq(DisplayManager::class.java)))
             .thenReturn(displayManager)
-        doNothing().`when`(context).registerComponentCallbacks(any())
+        doNothing().whenever(context).registerComponentCallbacks(any())
 
         // Mock display
         whenever(display.rotation).thenReturn(displayInfo.rotation)
diff --git a/tests/src/com/android/launcher3/util/KotlinMockitoHelpers.kt b/tests/src/com/android/launcher3/util/KotlinMockitoHelpers.kt
deleted file mode 100644
index c9c9616..0000000
--- a/tests/src/com/android/launcher3/util/KotlinMockitoHelpers.kt
+++ /dev/null
@@ -1,119 +0,0 @@
-/*
- * 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.launcher3.util
-
-/**
- * Kotlin versions of popular mockito methods that can return null in situations when Kotlin expects
- * a non-null value. Kotlin will throw an IllegalStateException when this takes place ("x must not
- * be null"). To fix this, we can use methods that modify the return type to be nullable. This
- * causes Kotlin to skip the null checks.
- */
-import org.mockito.ArgumentCaptor
-import org.mockito.Mockito
-
-/**
- * Returns Mockito.eq() as nullable type to avoid java.lang.IllegalStateException when null is
- * returned.
- *
- * Generic T is nullable because implicitly bounded by Any?.
- */
-fun <T> eq(obj: T): T = Mockito.eq<T>(obj)
-
-/**
- * Returns Mockito.same() as nullable type to avoid java.lang.IllegalStateException when null is
- * returned.
- *
- * Generic T is nullable because implicitly bounded by Any?.
- */
-fun <T> same(obj: T): T = Mockito.same<T>(obj)
-
-/**
- * Returns Mockito.any() as nullable type to avoid java.lang.IllegalStateException when null is
- * returned.
- *
- * Generic T is nullable because implicitly bounded by Any?.
- */
-fun <T> any(type: Class<T>): T = Mockito.any<T>(type)
-
-inline fun <reified T> any(): T = any(T::class.java)
-
-/** Kotlin type-inferred version of Mockito.nullable() */
-inline fun <reified T> nullable(): T? = Mockito.nullable(T::class.java)
-
-/**
- * Returns ArgumentCaptor.capture() as nullable type to avoid java.lang.IllegalStateException when
- * null is returned.
- *
- * Generic T is nullable because implicitly bounded by Any?.
- */
-fun <T> capture(argumentCaptor: ArgumentCaptor<T>): T = argumentCaptor.capture()
-
-/**
- * Helper function for creating an argumentCaptor in kotlin.
- *
- * Generic T is nullable because implicitly bounded by Any?.
- */
-inline fun <reified T : Any> argumentCaptor(): ArgumentCaptor<T> =
-    ArgumentCaptor.forClass(T::class.java)
-
-/**
- * Helper function for creating new mocks, without the need to pass in a [Class] instance.
- *
- * Generic T is nullable because implicitly bounded by Any?.
- */
-inline fun <reified T : Any> mock(): T = Mockito.mock(T::class.java)
-
-/**
- * A kotlin implemented wrapper of [ArgumentCaptor] which prevents the following exception when
- * kotlin tests are mocking kotlin objects and the methods take non-null parameters:
- * ```
- *     java.lang.NullPointerException: capture() must not be null
- * ```
- */
-class KotlinArgumentCaptor<T> constructor(clazz: Class<T>) {
-    private val wrapped: ArgumentCaptor<T> = ArgumentCaptor.forClass(clazz)
-    fun capture(): T = wrapped.capture()
-    val value: T
-        get() = wrapped.value
-}
-
-/**
- * Helper function for creating an argumentCaptor in kotlin.
- *
- * Generic T is nullable because implicitly bounded by Any?.
- */
-inline fun <reified T : Any> kotlinArgumentCaptor(): KotlinArgumentCaptor<T> =
-    KotlinArgumentCaptor(T::class.java)
-
-/**
- * Helper function for creating and using a single-use ArgumentCaptor in kotlin.
- *
- * ```
- *    val captor = argumentCaptor<Foo>()
- *    verify(...).someMethod(captor.capture())
- *    val captured = captor.value
- * ```
- *
- * becomes:
- * ```
- *    val captured = withArgCaptor<Foo> { verify(...).someMethod(capture()) }
- * ```
- *
- * NOTE: this uses the KotlinArgumentCaptor to avoid the NullPointerException.
- */
-inline fun <reified T : Any> withArgCaptor(block: KotlinArgumentCaptor<T>.() -> Unit): T =
-    kotlinArgumentCaptor<T>().apply { block() }.value
diff --git a/tests/src/com/android/launcher3/util/LockedUserStateTest.kt b/tests/src/com/android/launcher3/util/LockedUserStateTest.kt
index 92ab2cb..2c4a54f 100644
--- a/tests/src/com/android/launcher3/util/LockedUserStateTest.kt
+++ b/tests/src/com/android/launcher3/util/LockedUserStateTest.kt
@@ -26,29 +26,27 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.Mockito.verify
-import org.mockito.Mockito.verifyZeroInteractions
-import org.mockito.Mockito.`when`
-import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.verifyZeroInteractions
+import org.mockito.kotlin.whenever
 
 /** Unit tests for {@link LockedUserState} */
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class LockedUserStateTest {
 
-    @Mock lateinit var userManager: UserManager
-    @Mock lateinit var context: Context
+    private val userManager: UserManager = mock()
+    private val context: Context = mock()
 
     @Before
     fun setup() {
-        MockitoAnnotations.initMocks(this)
-        `when`(context.getSystemService(UserManager::class.java)).thenReturn(userManager)
+        whenever(context.getSystemService(UserManager::class.java)).thenReturn(userManager)
     }
 
     @Test
     fun runOnUserUnlocked_runs_action_immediately_if_already_unlocked() {
-        `when`(userManager.isUserUnlocked(Process.myUserHandle())).thenReturn(true)
+        whenever(userManager.isUserUnlocked(Process.myUserHandle())).thenReturn(true)
         val action: Runnable = mock()
         LockedUserState(context).runOnUserUnlocked(action)
         verify(action).run()
@@ -56,7 +54,7 @@
 
     @Test
     fun runOnUserUnlocked_waits_to_run_action_until_user_is_unlocked() {
-        `when`(userManager.isUserUnlocked(Process.myUserHandle())).thenReturn(false)
+        whenever(userManager.isUserUnlocked(Process.myUserHandle())).thenReturn(false)
         val action: Runnable = mock()
         val state = LockedUserState(context)
         state.runOnUserUnlocked(action)
@@ -67,13 +65,13 @@
 
     @Test
     fun isUserUnlocked_returns_true_when_user_is_unlocked() {
-        `when`(userManager.isUserUnlocked(Process.myUserHandle())).thenReturn(true)
+        whenever(userManager.isUserUnlocked(Process.myUserHandle())).thenReturn(true)
         assertThat(LockedUserState(context).isUserUnlocked).isTrue()
     }
 
     @Test
     fun isUserUnlocked_returns_false_when_user_is_locked() {
-        `when`(userManager.isUserUnlocked(Process.myUserHandle())).thenReturn(false)
+        whenever(userManager.isUserUnlocked(Process.myUserHandle())).thenReturn(false)
         assertThat(LockedUserState(context).isUserUnlocked).isFalse()
     }
 }