Added Tests for Early custom shortcut verification

Test: ShortcutCustomizationViewModelTest
CustomShortcutCategoriesRepositoryTest CustomInputGesturesRepositoryTest
Flag: com.android.systemui.keyboard_shortcut_helper_shortcut_customizer
Fix: 381063978

Change-Id: Iba9538f0889c113dd364b6557e8cd8af259de6c0

4 files changed, 226 insertions, 94 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/repository/CustomInputGesturesRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/repository/CustomInputGesturesRepositoryTest.kt
index 698fac107a1d..4d81cb0ce726 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/repository/CustomInputGesturesRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/repository/CustomInputGesturesRepositoryTest.kt
@@ -19,6 +19,7 @@ package com.android.systemui.keyboard.shortcut.data.repository
 import android.content.Context
 import android.content.Context.INPUT_SERVICE
 import android.content.Intent
+import android.hardware.input.FakeInputManager
 import android.hardware.input.InputGestureData
 import android.hardware.input.InputManager
 import android.hardware.input.InputManager.CUSTOM_INPUT_GESTURE_RESULT_SUCCESS
@@ -57,13 +58,14 @@ class CustomInputGesturesRepositoryTest : SysuiTestCase() {
     private val secondaryUserContext: Context = mock()
     private var activeUserContext: Context = primaryUserContext
 
-    private val kosmos = testKosmos().also {
-        it.userTracker = FakeUserTracker(onCreateCurrentUserContext = { activeUserContext })
-    }
+    private val kosmos =
+        testKosmos().also {
+            it.userTracker = FakeUserTracker(onCreateCurrentUserContext = { activeUserContext })
+        }
 
     private val inputManager = kosmos.fakeInputManager.inputManager
     private val broadcastDispatcher = kosmos.broadcastDispatcher
-    private val inputManagerForSecondaryUser: InputManager = mock()
+    private val inputManagerForSecondaryUser: InputManager = FakeInputManager().inputManager
     private val testScope = kosmos.testScope
     private val testHelper = kosmos.shortcutHelperTestHelper
     private val customInputGesturesRepository = kosmos.customInputGesturesRepository
@@ -94,9 +96,10 @@ class CustomInputGesturesRepositoryTest : SysuiTestCase() {
     fun customInputGestures_initialValueReturnsDataFromAPI() {
         testScope.runTest {
             val customInputGestures = listOf(allAppsInputGestureData)
-            whenever(
-                inputManager.getCustomInputGestures(/* filter= */ InputGestureData.Filter.KEY)
-            ).then { return@then customInputGestures }
+            whenever(inputManager.getCustomInputGestures(/* filter= */ InputGestureData.Filter.KEY))
+                .then {
+                    return@then customInputGestures
+                }
 
             val inputGestures by collectLastValue(customInputGesturesRepository.customInputGestures)
 
@@ -108,9 +111,10 @@ class CustomInputGesturesRepositoryTest : SysuiTestCase() {
     fun customInputGestures_isUpdatedToMostRecentDataAfterNewGestureIsAdded() {
         testScope.runTest {
             var customInputGestures = listOf<InputGestureData>()
-            whenever(
-                inputManager.getCustomInputGestures(/* filter= */ InputGestureData.Filter.KEY)
-            ).then { return@then customInputGestures }
+            whenever(inputManager.getCustomInputGestures(/* filter= */ InputGestureData.Filter.KEY))
+                .then {
+                    return@then customInputGestures
+                }
             whenever(inputManager.addCustomInputGesture(any())).then { invocation ->
                 val inputGesture = invocation.getArgument<InputGestureData>(0)
                 customInputGestures = customInputGestures + inputGesture
@@ -129,10 +133,10 @@ class CustomInputGesturesRepositoryTest : SysuiTestCase() {
     fun retrieveCustomInputGestures_retrievesMostRecentData() {
         testScope.runTest {
             var customInputGestures = listOf<InputGestureData>()
-            whenever(
-                inputManager.getCustomInputGestures(/* filter= */ InputGestureData.Filter.KEY)
-            ).then { return@then customInputGestures }
-
+            whenever(inputManager.getCustomInputGestures(/* filter= */ InputGestureData.Filter.KEY))
+                .then {
+                    return@then customInputGestures
+                }
 
             assertThat(customInputGesturesRepository.retrieveCustomInputGestures()).isEmpty()
 
@@ -143,24 +147,38 @@ class CustomInputGesturesRepositoryTest : SysuiTestCase() {
         }
     }
 
+    @Test
+    fun getInputGestureByTrigger_returnsInputGestureFromInputManager() =
+        testScope.runTest {
+            inputManager.addCustomInputGesture(allAppsInputGestureData)
+
+            val inputGestureData =
+                customInputGesturesRepository.getInputGestureByTrigger(
+                    allAppsInputGestureData.trigger
+                )
+
+            assertThat(inputGestureData).isEqualTo(allAppsInputGestureData)
+        }
+
     private fun setCustomInputGesturesForPrimaryUser(vararg inputGesture: InputGestureData) {
-        whenever(
-            inputManager.getCustomInputGestures(/* filter= */ InputGestureData.Filter.KEY)
-        ).thenReturn(inputGesture.toList())
+        whenever(inputManager.getCustomInputGestures(/* filter= */ InputGestureData.Filter.KEY))
+            .thenReturn(inputGesture.toList())
     }
 
     private fun setCustomInputGesturesForSecondaryUser(vararg inputGesture: InputGestureData) {
         whenever(
-            inputManagerForSecondaryUser.getCustomInputGestures(/* filter= */ InputGestureData.Filter.KEY)
-        ).thenReturn(inputGesture.toList())
+                inputManagerForSecondaryUser.getCustomInputGestures(
+                    /* filter= */ InputGestureData.Filter.KEY
+                )
+            )
+            .thenReturn(inputGesture.toList())
     }
 
     private fun switchToSecondaryUser() {
         activeUserContext = secondaryUserContext
         broadcastDispatcher.sendIntentToMatchingReceiversOnly(
             context,
-            Intent(Intent.ACTION_USER_SWITCHED)
+            Intent(Intent.ACTION_USER_SWITCHED),
         )
     }
-
-}
\ No newline at end of file
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/repository/CustomShortcutCategoriesRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/repository/CustomShortcutCategoriesRepositoryTest.kt
index 4cfb26e6555b..522572dcffb7 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/repository/CustomShortcutCategoriesRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/repository/CustomShortcutCategoriesRepositoryTest.kt
@@ -24,6 +24,7 @@ import android.hardware.input.InputGestureData
 import android.hardware.input.InputGestureData.createKeyTrigger
 import android.hardware.input.InputManager.CUSTOM_INPUT_GESTURE_RESULT_ERROR_DOES_NOT_EXIST
 import android.hardware.input.InputManager.CUSTOM_INPUT_GESTURE_RESULT_SUCCESS
+import android.hardware.input.KeyGestureEvent.KEY_GESTURE_TYPE_HOME
 import android.hardware.input.KeyGestureEvent.KEY_GESTURE_TYPE_LAUNCH_APPLICATION
 import android.hardware.input.fakeInputManager
 import android.platform.test.annotations.DisableFlags
@@ -336,28 +337,6 @@ class CustomShortcutCategoriesRepositoryTest : SysuiTestCase() {
         }
     }
 
-    private suspend fun customizeShortcut(
-        customizationRequest: ShortcutCustomizationRequestInfo,
-        keyCombination: KeyCombination? = null,
-    ): ShortcutCustomizationRequestResult {
-        repo.onCustomizationRequested(customizationRequest)
-        repo.updateUserKeyCombination(keyCombination)
-
-        return when (customizationRequest) {
-            is SingleShortcutCustomization.Add -> {
-                repo.confirmAndSetShortcutCurrentlyBeingCustomized()
-            }
-
-            is SingleShortcutCustomization.Delete -> {
-                repo.deleteShortcutCurrentlyBeingCustomized()
-            }
-
-            else -> {
-                ShortcutCustomizationRequestResult.ERROR_OTHER
-            }
-        }
-    }
-
     @Test
     @EnableFlags(FLAG_ENABLE_CUSTOMIZABLE_INPUT_GESTURES, FLAG_USE_KEY_GESTURE_EVENT_HANDLER)
     fun categories_isUpdatedAfterCustomShortcutsAreReset() {
@@ -387,10 +366,66 @@ class CustomShortcutCategoriesRepositoryTest : SysuiTestCase() {
         }
     }
 
+    @Test
+    fun selectedKeyCombinationIsAvailable_whenTriggerIsNotRegisteredInInputManager() =
+        testScope.runTest {
+            helper.toggle(deviceId = 123)
+            repo.onCustomizationRequested(allAppsShortcutAddRequest)
+            repo.updateUserKeyCombination(standardKeyCombination)
+
+            assertThat(repo.isSelectedKeyCombinationAvailable()).isTrue()
+        }
+
+    @Test
+    fun selectedKeyCombinationIsNotAvailable_whenTriggerIsRegisteredInInputManager() =
+        testScope.runTest {
+            inputManager.addCustomInputGesture(buildInputGestureWithStandardKeyCombination())
+
+            helper.toggle(deviceId = 123)
+            repo.onCustomizationRequested(allAppsShortcutAddRequest)
+            repo.updateUserKeyCombination(standardKeyCombination)
+
+            assertThat(repo.isSelectedKeyCombinationAvailable()).isFalse()
+        }
+
     private fun setApiAppLaunchBookmarks(appLaunchBookmarks: List<InputGestureData>) {
         whenever(inputManager.appLaunchBookmarks).thenReturn(appLaunchBookmarks)
     }
 
+    private suspend fun customizeShortcut(
+        customizationRequest: ShortcutCustomizationRequestInfo,
+        keyCombination: KeyCombination? = null,
+    ): ShortcutCustomizationRequestResult {
+        repo.onCustomizationRequested(customizationRequest)
+        repo.updateUserKeyCombination(keyCombination)
+
+        return when (customizationRequest) {
+            is SingleShortcutCustomization.Add -> {
+                repo.confirmAndSetShortcutCurrentlyBeingCustomized()
+            }
+
+            is SingleShortcutCustomization.Delete -> {
+                repo.deleteShortcutCurrentlyBeingCustomized()
+            }
+
+            else -> {
+                ShortcutCustomizationRequestResult.ERROR_OTHER
+            }
+        }
+    }
+
+    private fun buildInputGestureWithStandardKeyCombination() =
+        InputGestureData.Builder()
+            .setKeyGestureType(KEY_GESTURE_TYPE_HOME)
+            .setTrigger(
+                createKeyTrigger(
+                    /* keycode= */ standardKeyCombination.keyCode!!,
+                    /* modifierState= */ standardKeyCombination.modifiers and
+                        ALL_SUPPORTED_MODIFIERS,
+                )
+            )
+            .build()
+
     private fun simpleInputGestureDataForAppLaunchShortcut(
         keyCode: Int = KEYCODE_A,
         modifiers: Int = META_CTRL_ON or META_ALT_ON,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutCustomizationViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutCustomizationViewModelTest.kt
index 11efe67cb2ec..6eef5eb09812 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutCustomizationViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutCustomizationViewModelTest.kt
@@ -18,11 +18,15 @@ package com.android.systemui.keyboard.shortcut.ui.viewmodel
 
 import android.content.Context
 import android.content.Context.INPUT_SERVICE
-import android.hardware.input.InputManager.CUSTOM_INPUT_GESTURE_RESULT_ERROR_ALREADY_EXISTS
+import android.hardware.input.InputGestureData
+import android.hardware.input.InputGestureData.createKeyTrigger
 import android.hardware.input.InputManager.CUSTOM_INPUT_GESTURE_RESULT_ERROR_OTHER
 import android.hardware.input.InputManager.CUSTOM_INPUT_GESTURE_RESULT_ERROR_RESERVED_GESTURE
-import android.hardware.input.InputManager.CUSTOM_INPUT_GESTURE_RESULT_SUCCESS
+import android.hardware.input.KeyGestureEvent.KEY_GESTURE_TYPE_HOME
 import android.hardware.input.fakeInputManager
+import android.view.KeyEvent.KEYCODE_A
+import android.view.KeyEvent.META_CTRL_ON
+import android.view.KeyEvent.META_META_ON
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
@@ -30,7 +34,6 @@ import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.allAppsInputGestureData
 import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.allAppsShortcutAddRequest
 import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.allAppsShortcutDeleteRequest
-import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.goHomeInputGestureData
 import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.keyDownEventWithActionKeyPressed
 import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.keyDownEventWithoutActionKeyPressed
 import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.keyUpEventWithActionKeyPressed
@@ -44,16 +47,17 @@ import com.android.systemui.keyboard.shortcut.ui.model.ShortcutCustomizationUiSt
 import com.android.systemui.keyboard.shortcut.ui.model.ShortcutCustomizationUiState.DeleteShortcutDialog
 import com.android.systemui.keyboard.shortcut.ui.model.ShortcutCustomizationUiState.ResetShortcutDialog
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.kosmos.useUnconfinedTestDispatcher
 import com.android.systemui.res.R
 import com.android.systemui.settings.FakeUserTracker
 import com.android.systemui.settings.userTracker
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.kotlin.any
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.whenever
 
@@ -63,7 +67,7 @@ class ShortcutCustomizationViewModelTest : SysuiTestCase() {
 
     private val mockUserContext: Context = mock()
     private val kosmos =
-        testKosmos().also {
+        testKosmos().useUnconfinedTestDispatcher().also {
             it.userTracker = FakeUserTracker(onCreateCurrentUserContext = { mockUserContext })
         }
     private val testScope = kosmos.testScope
@@ -75,6 +79,7 @@ class ShortcutCustomizationViewModelTest : SysuiTestCase() {
     fun setup() {
         helper.showFromActivity()
         whenever(mockUserContext.getSystemService(INPUT_SERVICE)).thenReturn(inputManager)
+        testScope.backgroundScope.launch { viewModel.activate() }
     }
 
     @Test
@@ -146,8 +151,6 @@ class ShortcutCustomizationViewModelTest : SysuiTestCase() {
     fun uiState_becomeInactiveAfterSuccessfullySettingShortcut() {
         testScope.runTest {
             val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
-            whenever(inputManager.addCustomInputGesture(any()))
-                .thenReturn(CUSTOM_INPUT_GESTURE_RESULT_SUCCESS)
 
             openAddShortcutDialogAndSetShortcut()
 
@@ -166,11 +169,38 @@ class ShortcutCustomizationViewModelTest : SysuiTestCase() {
     }
 
     @Test
-    fun uiState_errorMessage_isKeyCombinationInUse_whenKeyCombinationAlreadyExists() {
+    fun uiState_errorMessage_onKeyPressed_isKeyCombinationInUse_whenKeyCombinationAlreadyExists() {
         testScope.runTest {
+            inputManager.addCustomInputGesture(buildSimpleInputGestureWithMetaCtrlATrigger())
+            val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
+
+            openAddShortcutDialogAndPressKeyCombination()
+
+            assertThat((uiState as AddShortcutDialog).errorMessage)
+                .isEqualTo(
+                    context.getString(
+                        R.string.shortcut_customizer_key_combination_in_use_error_message
+                    )
+                )
+        }
+    }
+
+    @Test
+    fun uiState_errorMessage_onKeyPressed_isEmpty_whenKeyCombinationIsAvailable() {
+        testScope.runTest {
+            val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
+
+            openAddShortcutDialogAndPressKeyCombination()
+
+            assertThat((uiState as AddShortcutDialog).errorMessage).isEmpty()
+        }
+    }
+
+    @Test
+    fun uiState_errorMessage_onSetShortcut_isKeyCombinationInUse_whenKeyCombinationAlreadyExists() {
+        testScope.runTest {
+            inputManager.addCustomInputGesture(buildSimpleInputGestureWithMetaCtrlATrigger())
             val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
-            whenever(inputManager.addCustomInputGesture(any()))
-                .thenReturn(CUSTOM_INPUT_GESTURE_RESULT_ERROR_ALREADY_EXISTS)
 
             openAddShortcutDialogAndSetShortcut()
 
@@ -184,11 +214,12 @@ class ShortcutCustomizationViewModelTest : SysuiTestCase() {
     }
 
     @Test
-    fun uiState_errorMessage_isKeyCombinationInUse_whenKeyCombinationIsReserved() {
+    fun uiState_errorMessage_onSetShortcut_isKeyCombinationInUse_whenKeyCombinationIsReserved() {
         testScope.runTest {
+            inputManager.addCustomInputGesture(buildSimpleInputGestureWithMetaCtrlATrigger())
+            kosmos.fakeInputManager.addCustomInputGestureErrorCode =
+                CUSTOM_INPUT_GESTURE_RESULT_ERROR_RESERVED_GESTURE
             val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
-            whenever(inputManager.addCustomInputGesture(any()))
-                .thenReturn(CUSTOM_INPUT_GESTURE_RESULT_ERROR_RESERVED_GESTURE)
 
             openAddShortcutDialogAndSetShortcut()
 
@@ -202,11 +233,12 @@ class ShortcutCustomizationViewModelTest : SysuiTestCase() {
     }
 
     @Test
-    fun uiState_errorMessage_isGenericError_whenErrorIsUnknown() {
+    fun uiState_errorMessage_onSetShortcut_isGenericError_whenErrorIsUnknown() {
         testScope.runTest {
+            inputManager.addCustomInputGesture(buildSimpleInputGestureWithMetaCtrlATrigger())
+            kosmos.fakeInputManager.addCustomInputGestureErrorCode =
+                CUSTOM_INPUT_GESTURE_RESULT_ERROR_OTHER
             val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
-            whenever(inputManager.addCustomInputGesture(any()))
-                .thenReturn(CUSTOM_INPUT_GESTURE_RESULT_ERROR_OTHER)
 
             openAddShortcutDialogAndSetShortcut()
 
@@ -219,10 +251,7 @@ class ShortcutCustomizationViewModelTest : SysuiTestCase() {
     fun uiState_becomesInactiveAfterSuccessfullyDeletingShortcut() {
         testScope.runTest {
             val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
-            whenever(inputManager.getCustomInputGestures(any()))
-                .thenReturn(listOf(goHomeInputGestureData, allAppsInputGestureData))
-            whenever(inputManager.removeCustomInputGesture(any()))
-                .thenReturn(CUSTOM_INPUT_GESTURE_RESULT_SUCCESS)
+            inputManager.addCustomInputGesture(allAppsInputGestureData)
 
             openDeleteShortcutDialogAndDeleteShortcut()
 
@@ -234,7 +263,6 @@ class ShortcutCustomizationViewModelTest : SysuiTestCase() {
     fun uiState_becomesInactiveAfterSuccessfullyResettingShortcuts() {
         testScope.runTest {
             val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
-            whenever(inputManager.getCustomInputGestures(any())).thenReturn(emptyList())
 
             openResetShortcutDialogAndResetAllCustomShortcuts()
 
@@ -305,29 +333,34 @@ class ShortcutCustomizationViewModelTest : SysuiTestCase() {
             viewModel.onShortcutKeyCombinationSelected(keyDownEventWithActionKeyPressed)
             viewModel.onShortcutKeyCombinationSelected(keyUpEventWithActionKeyPressed)
             viewModel.clearSelectedKeyCombination()
-            
             assertThat((uiState as AddShortcutDialog).pressedKeys).isEmpty()
         }
     }
 
     private suspend fun openAddShortcutDialogAndSetShortcut() {
-        viewModel.onShortcutCustomizationRequested(allAppsShortcutAddRequest)
+        openAddShortcutDialogAndPressKeyCombination()
+        viewModel.onSetShortcut()
+    }
 
+    private fun openAddShortcutDialogAndPressKeyCombination() {
+        viewModel.onShortcutCustomizationRequested(allAppsShortcutAddRequest)
         viewModel.onShortcutKeyCombinationSelected(keyDownEventWithActionKeyPressed)
         viewModel.onShortcutKeyCombinationSelected(keyUpEventWithActionKeyPressed)
-
-        viewModel.onSetShortcut()
     }
 
     private suspend fun openDeleteShortcutDialogAndDeleteShortcut() {
         viewModel.onShortcutCustomizationRequested(allAppsShortcutDeleteRequest)
-
         viewModel.deleteShortcutCurrentlyBeingCustomized()
     }
 
     private suspend fun openResetShortcutDialogAndResetAllCustomShortcuts() {
         viewModel.onShortcutCustomizationRequested(ShortcutCustomizationRequestInfo.Reset)
-
         viewModel.resetAllCustomShortcuts()
     }
+
+    private fun buildSimpleInputGestureWithMetaCtrlATrigger() =
+        InputGestureData.Builder()
+            .setKeyGestureType(KEY_GESTURE_TYPE_HOME)
+            .setTrigger(createKeyTrigger(KEYCODE_A, META_CTRL_ON or META_META_ON))
+            .build()
 }
diff --git a/packages/SystemUI/tests/utils/src/android/hardware/input/FakeInputManager.kt b/packages/SystemUI/tests/utils/src/android/hardware/input/FakeInputManager.kt
index de4bbecaaf0e..42c509eeaa0b 100644
--- a/packages/SystemUI/tests/utils/src/android/hardware/input/FakeInputManager.kt
+++ b/packages/SystemUI/tests/utils/src/android/hardware/input/FakeInputManager.kt
@@ -16,16 +16,19 @@
 
 package android.hardware.input
 
+import android.hardware.input.InputGestureData.Trigger
+import android.hardware.input.InputManager.CUSTOM_INPUT_GESTURE_RESULT_ERROR_ALREADY_EXISTS
+import android.hardware.input.InputManager.CUSTOM_INPUT_GESTURE_RESULT_ERROR_DOES_NOT_EXIST
+import android.hardware.input.InputManager.CUSTOM_INPUT_GESTURE_RESULT_SUCCESS
 import android.hardware.input.InputManager.InputDeviceListener
 import android.view.InputDevice
 import android.view.KeyCharacterMap
 import android.view.KeyCharacterMap.VIRTUAL_KEYBOARD
 import android.view.KeyEvent
-import com.android.systemui.util.mockito.any
-import com.android.systemui.util.mockito.mock
-import com.android.systemui.util.mockito.whenever
 import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.invocation.InvocationOnMock
+import org.mockito.kotlin.any
+import org.mockito.kotlin.mock
 
 class FakeInputManager {
 
@@ -49,36 +52,79 @@ class FakeInputManager {
         )
 
     private var inputDeviceListener: InputDeviceListener? = null
+    private val customInputGestures: MutableMap<Trigger, InputGestureData> = mutableMapOf()
+    var addCustomInputGestureErrorCode = CUSTOM_INPUT_GESTURE_RESULT_ERROR_ALREADY_EXISTS
+
+    val inputManager: InputManager = mock {
+        on { getCustomInputGestures(any()) }.then { customInputGestures.values.toList() }
+
+        on { addCustomInputGesture(any()) }
+            .then {
+                val inputGestureData = it.getArgument<InputGestureData>(0)
+                val trigger = inputGestureData.trigger
+
+                if (customInputGestures.containsKey(trigger)) {
+                    addCustomInputGestureErrorCode
+                } else {
+                    customInputGestures[trigger] = inputGestureData
+                    CUSTOM_INPUT_GESTURE_RESULT_SUCCESS
+                }
+            }
+
+        on { removeCustomInputGesture(any()) }
+            .then {
+                val inputGestureData = it.getArgument<InputGestureData>(0)
+                val trigger = inputGestureData.trigger
+
+                if (customInputGestures.containsKey(trigger)) {
+                    customInputGestures.remove(trigger)
+                    CUSTOM_INPUT_GESTURE_RESULT_SUCCESS
+                } else {
+                    CUSTOM_INPUT_GESTURE_RESULT_ERROR_DOES_NOT_EXIST
+                }
+            }
+
+        on { removeAllCustomInputGestures(any()) }.then { customInputGestures.clear() }
 
-    val inputManager =
-        mock<InputManager> {
-            whenever(getInputDevice(anyInt())).thenAnswer { invocation ->
+        on { getInputGesture(any()) }
+            .then {
+                val trigger = it.getArgument<Trigger>(0)
+                customInputGestures[trigger]
+            }
+
+        on { getInputDevice(anyInt()) }
+            .thenAnswer { invocation ->
                 val deviceId = invocation.arguments[0] as Int
                 return@thenAnswer devices[deviceId]
             }
-            whenever(inputDeviceIds).thenAnswer {
+        on { inputDeviceIds }
+            .thenAnswer {
                 return@thenAnswer devices.keys.toIntArray()
             }
 
-            fun setDeviceEnabled(invocation: InvocationOnMock, enabled: Boolean) {
-                val deviceId = invocation.arguments[0] as Int
-                val device = devices[deviceId] ?: return
-                devices[deviceId] = device.copy(enabled = enabled)
-            }
+        fun setDeviceEnabled(invocation: InvocationOnMock, enabled: Boolean) {
+            val deviceId = invocation.arguments[0] as Int
+            val device = devices[deviceId] ?: return
+            devices[deviceId] = device.copy(enabled = enabled)
+        }
 
-            whenever(disableInputDevice(anyInt())).thenAnswer { invocation ->
-                setDeviceEnabled(invocation, enabled = false)
-            }
-            whenever(enableInputDevice(anyInt())).thenAnswer { invocation ->
-                setDeviceEnabled(invocation, enabled = true)
-            }
-            whenever(deviceHasKeys(any(), any())).thenAnswer { invocation ->
+        on { disableInputDevice(anyInt()) }
+            .thenAnswer { invocation -> setDeviceEnabled(invocation, enabled = false) }
+        on { enableInputDevice(anyInt()) }
+            .thenAnswer { invocation -> setDeviceEnabled(invocation, enabled = true) }
+        on { deviceHasKeys(any(), any()) }
+            .thenAnswer { invocation ->
                 val deviceId = invocation.arguments[0] as Int
                 val keyCodes = invocation.arguments[1] as IntArray
                 val supportedKeyCodes = supportedKeyCodesByDeviceId[deviceId]!!
                 return@thenAnswer keyCodes.map { supportedKeyCodes.contains(it) }.toBooleanArray()
             }
-        }
+    }
+
+    fun resetCustomInputGestures() {
+        customInputGestures.clear()
+        addCustomInputGestureErrorCode = CUSTOM_INPUT_GESTURE_RESULT_ERROR_ALREADY_EXISTS
+    }
 
     fun addPhysicalKeyboardIfNotPresent(deviceId: Int, enabled: Boolean = true) {
         if (devices.containsKey(deviceId)) {
@@ -97,7 +143,7 @@ class FakeInputManager {
         vendorId: Int = 0,
         productId: Int = 0,
         isFullKeyboard: Boolean = true,
-        enabled: Boolean = true
+        enabled: Boolean = true,
     ) {
         check(id > 0) { "Physical keyboard ids have to be > 0" }
         addKeyboard(id, vendorId, productId, isFullKeyboard, enabled)
@@ -113,7 +159,7 @@ class FakeInputManager {
         vendorId: Int = 0,
         productId: Int = 0,
         isFullKeyboard: Boolean = true,
-        enabled: Boolean = true
+        enabled: Boolean = true,
     ) {
         val keyboardType =
             if (isFullKeyboard) InputDevice.KEYBOARD_TYPE_ALPHABETIC
@@ -152,7 +198,7 @@ class FakeInputManager {
         id: Int = getId(),
         type: Int = keyboardType,
         sources: Int = getSources(),
-        enabled: Boolean = isEnabled
+        enabled: Boolean = isEnabled,
     ) =
         InputDevice.Builder()
             .setId(id)