Added data refresh after new shortcut is added/deleted

This causes shortcut custom shortcut repository to get updated leading
to an update in shortcuts being displayed in the UI.

Fix: 379718710
Test: CustomShortcutCategoriesRepositoryTest
Test: CustomInputGesturesRepositoryTest
Flag: com.android.systemui.keyboard_shortcut_helper_shortcut_customizer
Change-Id: Ie6c11981875038ccc8ca28662b1db0c5a2dd8e4a

7 files changed, 371 insertions, 90 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
new file mode 100644
index 000000000000..e659ef274980
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/repository/CustomInputGesturesRepositoryTest.kt
@@ -0,0 +1,118 @@
+/*
+ * Copyright (C) 2024 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.keyboard.shortcut.data.repository
+
+import android.content.Context
+import android.content.Context.INPUT_SERVICE
+import android.hardware.input.InputGestureData
+import android.hardware.input.InputManager.CUSTOM_INPUT_GESTURE_RESULT_SUCCESS
+import android.hardware.input.fakeInputManager
+import android.platform.test.annotations.EnableFlags
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.hardware.input.Flags.FLAG_ENABLE_CUSTOMIZABLE_INPUT_GESTURES
+import com.android.hardware.input.Flags.FLAG_USE_KEY_GESTURE_EVENT_HANDLER
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.keyboard.shortcut.customInputGesturesRepository
+import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.allAppsInputGestureData
+import com.android.systemui.kosmos.testScope
+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.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
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+@EnableFlags(FLAG_ENABLE_CUSTOMIZABLE_INPUT_GESTURES, FLAG_USE_KEY_GESTURE_EVENT_HANDLER)
+class CustomInputGesturesRepositoryTest : SysuiTestCase() {
+
+    private val mockUserContext: Context = mock()
+    private val kosmos = testKosmos().also {
+        it.userTracker = FakeUserTracker(onCreateCurrentUserContext = { mockUserContext })
+    }
+
+    private val inputManager = kosmos.fakeInputManager.inputManager
+    private val testScope = kosmos.testScope
+    private val customInputGesturesRepository = kosmos.customInputGesturesRepository
+
+    @Before
+    fun setup(){
+        whenever(mockUserContext.getSystemService(INPUT_SERVICE)).thenReturn(inputManager)
+    }
+
+    @Test
+    fun customInputGestures_initialValueReturnsDataFromAPI() {
+        testScope.runTest {
+            val customInputGestures = listOf(allAppsInputGestureData)
+            whenever(
+                inputManager.getCustomInputGestures(/* filter= */ InputGestureData.Filter.KEY)
+            ).then { return@then customInputGestures }
+
+            val inputGestures by collectLastValue(customInputGesturesRepository.customInputGestures)
+
+            assertThat(inputGestures).containsExactly(allAppsInputGestureData)
+        }
+    }
+
+    @Test
+    fun customInputGestures_isUpdatedToMostRecentDataAfterNewGestureIsAdded() {
+        testScope.runTest {
+            var customInputGestures = listOf<InputGestureData>()
+            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
+                return@then CUSTOM_INPUT_GESTURE_RESULT_SUCCESS
+            }
+
+            val inputGestures by collectLastValue(customInputGesturesRepository.customInputGestures)
+            assertThat(inputGestures).isEmpty()
+
+            customInputGesturesRepository.addCustomInputGesture(allAppsInputGestureData)
+            assertThat(inputGestures).containsExactly(allAppsInputGestureData)
+        }
+    }
+
+    @Test
+    fun retrieveCustomInputGestures_retrievesMostRecentData() {
+        testScope.runTest {
+            var customInputGestures = listOf<InputGestureData>()
+            whenever(
+                inputManager.getCustomInputGestures(/* filter= */ InputGestureData.Filter.KEY)
+            ).then { return@then customInputGestures }
+
+
+            assertThat(customInputGesturesRepository.retrieveCustomInputGestures()).isEmpty()
+
+            customInputGestures = listOf(allAppsInputGestureData)
+
+            assertThat(customInputGesturesRepository.retrieveCustomInputGestures())
+                .containsExactly(allAppsInputGestureData)
+        }
+    }
+
+}
\ 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 0d32b7fb1b3e..72916a35814f 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
@@ -18,6 +18,8 @@ package com.android.systemui.keyboard.shortcut.data.repository
 
 import android.content.Context
 import android.content.Context.INPUT_SERVICE
+import android.hardware.input.InputGestureData
+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.fakeInputManager
 import android.platform.test.annotations.DisableFlags
@@ -34,15 +36,19 @@ import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.keyboard.shared.model.ShortcutCustomizationRequestResult
 import com.android.systemui.keyboard.shortcut.customShortcutCategoriesRepository
 import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.ALL_SUPPORTED_MODIFIERS
+import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.allAppsShortcutAddRequest
 import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.allAppsInputGestureData
 import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.allCustomizableInputGesturesWithSimpleShortcutCombinations
+import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.allAppsShortcutCategory
 import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.customizableInputGestureWithUnknownKeyGestureType
 import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.expectedShortcutCategoriesWithSimpleShortcutCombination
 import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.goHomeInputGestureData
-import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.standardAddCustomShortcutRequestInfo
-import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.standardDeleteCustomShortcutRequestInfo
+import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.allAppsShortcutDeleteRequest
 import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.standardKeyCombination
 import com.android.systemui.keyboard.shortcut.shared.model.KeyCombination
+import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCustomizationRequestInfo
+import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCustomizationRequestInfo.Add
+import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCustomizationRequestInfo.Delete
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutKey
 import com.android.systemui.keyboard.shortcut.shortcutHelperTestHelper
 import com.android.systemui.kosmos.testScope
@@ -55,6 +61,7 @@ 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.anyOrNull
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.whenever
@@ -187,11 +194,11 @@ class CustomShortcutCategoriesRepositoryTest : SysuiTestCase() {
     @Test
     fun shortcutBeingCustomized_updatedOnCustomizationRequested() {
         testScope.runTest {
-            repo.onCustomizationRequested(standardAddCustomShortcutRequestInfo)
+            repo.onCustomizationRequested(allAppsShortcutAddRequest)
 
             val shortcutBeingCustomized = repo.getShortcutBeingCustomized()
 
-            assertThat(shortcutBeingCustomized).isEqualTo(standardAddCustomShortcutRequestInfo)
+            assertThat(shortcutBeingCustomized).isEqualTo(allAppsShortcutAddRequest)
         }
     }
 
@@ -211,7 +218,7 @@ class CustomShortcutCategoriesRepositoryTest : SysuiTestCase() {
     fun buildInputGestureDataForShortcutBeingCustomized_noKeyCombinationSelected_returnsNull() {
         testScope.runTest {
             helper.toggle(deviceId = 123)
-            repo.onCustomizationRequested(standardAddCustomShortcutRequestInfo)
+            repo.onCustomizationRequested(allAppsShortcutAddRequest)
 
             val inputGestureData = repo.buildInputGestureDataForShortcutBeingCustomized()
 
@@ -223,7 +230,7 @@ class CustomShortcutCategoriesRepositoryTest : SysuiTestCase() {
     fun buildInputGestureDataForShortcutBeingCustomized_successfullyBuildInputGestureData() {
         testScope.runTest {
             helper.toggle(deviceId = 123)
-            repo.onCustomizationRequested(standardAddCustomShortcutRequestInfo)
+            repo.onCustomizationRequested(allAppsShortcutAddRequest)
             repo.updateUserKeyCombination(standardKeyCombination)
             val inputGestureData = repo.buildInputGestureDataForShortcutBeingCustomized()
 
@@ -242,13 +249,78 @@ class CustomShortcutCategoriesRepositoryTest : SysuiTestCase() {
                 .thenReturn(listOf(allAppsInputGestureData, goHomeInputGestureData))
             whenever(inputManager.removeCustomInputGesture(allAppsInputGestureData))
                 .thenReturn(CUSTOM_INPUT_GESTURE_RESULT_SUCCESS)
+            helper.toggle(deviceId = 123)
+
+            val result = customizeShortcut(allAppsShortcutDeleteRequest)
+            assertThat(result).isEqualTo(ShortcutCustomizationRequestResult.SUCCESS)
+        }
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_CUSTOMIZABLE_INPUT_GESTURES, FLAG_USE_KEY_GESTURE_EVENT_HANDLER)
+    fun categories_isUpdatedAfterCustomShortcutIsDeleted() {
+        testScope.runTest {
+            // TODO(b/380445594) refactor tests and move these stubbing to ShortcutHelperTestHelper
+            var customInputGestures = listOf(allAppsInputGestureData)
+            whenever(inputManager.getCustomInputGestures(anyOrNull())).then {
+                return@then customInputGestures
+            }
+            whenever(inputManager.removeCustomInputGesture(any())).then {
+                val inputGestureToRemove = it.getArgument<InputGestureData>(0)
+                val containsGesture = customInputGestures.contains(inputGestureToRemove)
+                customInputGestures = customInputGestures - inputGestureToRemove
+                return@then if (containsGesture) CUSTOM_INPUT_GESTURE_RESULT_SUCCESS
+                else CUSTOM_INPUT_GESTURE_RESULT_ERROR_DOES_NOT_EXIST
+            }
+            val categories by collectLastValue(repo.categories)
+            helper.toggle(deviceId = 123)
 
+            customizeShortcut(customizationRequest = allAppsShortcutDeleteRequest)
+            assertThat(categories).isEmpty()
+        }
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_CUSTOMIZABLE_INPUT_GESTURES, FLAG_USE_KEY_GESTURE_EVENT_HANDLER)
+    fun categories_isUpdatedAfterCustomShortcutIsAdded() {
+        testScope.runTest {
+            // TODO(b/380445594) refactor tests and move these stubbings to ShortcutHelperTestHelper
+            var customInputGestures = listOf<InputGestureData>()
+            whenever(inputManager.getCustomInputGestures(anyOrNull())).then {
+                return@then customInputGestures
+            }
+            whenever(inputManager.addCustomInputGesture(any())).then {
+                val inputGestureToAdd = it.getArgument<InputGestureData>(0)
+                customInputGestures = customInputGestures + inputGestureToAdd
+                return@then CUSTOM_INPUT_GESTURE_RESULT_SUCCESS
+            }
+            val categories by collectLastValue(repo.categories)
             helper.toggle(deviceId = 123)
-            repo.onCustomizationRequested(standardDeleteCustomShortcutRequestInfo)
 
-            val result = repo.deleteShortcutCurrentlyBeingCustomized()
+            customizeShortcut(allAppsShortcutAddRequest, standardKeyCombination)
+            assertThat(categories).containsExactly(allAppsShortcutCategory)
+        }
+    }
 
-            assertThat(result).isEqualTo(ShortcutCustomizationRequestResult.SUCCESS)
+    private suspend fun customizeShortcut(
+        customizationRequest: ShortcutCustomizationRequestInfo,
+        keyCombination: KeyCombination? = null
+    ): ShortcutCustomizationRequestResult{
+        repo.onCustomizationRequested(customizationRequest)
+        repo.updateUserKeyCombination(keyCombination)
+
+        return when (customizationRequest) {
+            is Add -> {
+                repo.confirmAndSetShortcutCurrentlyBeingCustomized()
+            }
+
+            is Delete -> {
+                repo.deleteShortcutCurrentlyBeingCustomized()
+            }
+
+            else -> {
+                ShortcutCustomizationRequestResult.ERROR_OTHER
+            }
         }
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/source/TestShortcuts.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/source/TestShortcuts.kt
index 6d22b4903920..7855d4219788 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/source/TestShortcuts.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/source/TestShortcuts.kt
@@ -35,6 +35,7 @@ import android.view.KeyEvent.META_SHIFT_RIGHT_ON
 import android.view.KeyEvent.META_SYM_ON
 import android.view.KeyboardShortcutGroup
 import android.view.KeyboardShortcutInfo
+import com.android.systemui.keyboard.shortcut.data.repository.ShortcutHelperKeys
 import com.android.systemui.keyboard.shortcut.shared.model.KeyCombination
 import com.android.systemui.keyboard.shortcut.shared.model.Shortcut
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCategory
@@ -47,6 +48,7 @@ import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCustomization
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutKey
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutSubCategory
 import com.android.systemui.keyboard.shortcut.shared.model.shortcut
+import com.android.systemui.keyboard.shortcut.shared.model.shortcutCategory
 import com.android.systemui.keyboard.shortcut.ui.model.ShortcutCustomizationUiState
 import com.android.systemui.res.R
 
@@ -558,14 +560,14 @@ object TestShortcuts {
             ),
         )
 
-    val standardAddCustomShortcutRequestInfo =
+    val allAppsShortcutAddRequest =
         ShortcutCustomizationRequestInfo.Add(
             label = "Open apps list",
             categoryType = System,
             subCategoryLabel = "System controls",
         )
 
-    val standardDeleteCustomShortcutRequestInfo =
+    val allAppsShortcutDeleteRequest =
         ShortcutCustomizationRequestInfo.Delete(
             label = "Open apps list",
             categoryType = System,
@@ -610,6 +612,20 @@ object TestShortcuts {
             )
             .build()
 
+    val allAppsShortcutCategory =
+        shortcutCategory(System) {
+            subCategory("System controls") {
+                shortcut("Open apps list") {
+                    command {
+                        isCustom(true)
+                        key(ShortcutHelperKeys.metaModifierIconResId)
+                        key("Shift")
+                        key("A")
+                    }
+                }
+            }
+        }
+
     val expectedStandardDeleteShortcutUiState =
         ShortcutCustomizationUiState.DeleteShortcutDialog(isDialogShowing = false)
 
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 d0ce34c2d68d..d3d1a3506725 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
@@ -34,9 +34,9 @@ import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.goHomeIn
 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
-import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.standardAddCustomShortcutRequestInfo
+import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.allAppsShortcutAddRequest
 import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.standardAddShortcutRequest
-import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.standardDeleteCustomShortcutRequestInfo
+import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts.allAppsShortcutDeleteRequest
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutKey
 import com.android.systemui.keyboard.shortcut.shortcutCustomizationViewModelFactory
 import com.android.systemui.keyboard.shortcut.shortcutHelperTestHelper
@@ -99,7 +99,7 @@ class ShortcutCustomizationViewModelTest : SysuiTestCase() {
     @Test
     fun uiState_correctlyUpdatedWhenDeleteShortcutCustomizationIsRequested() {
         testScope.runTest {
-            viewModel.onShortcutCustomizationRequested(standardDeleteCustomShortcutRequestInfo)
+            viewModel.onShortcutCustomizationRequested(allAppsShortcutDeleteRequest)
             val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
 
             assertThat(uiState).isEqualTo(expectedStandardDeleteShortcutUiState)
@@ -122,7 +122,7 @@ class ShortcutCustomizationViewModelTest : SysuiTestCase() {
     fun uiState_consumedOnDeleteDialogShown() {
         testScope.runTest {
             val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
-            viewModel.onShortcutCustomizationRequested(standardDeleteCustomShortcutRequestInfo)
+            viewModel.onShortcutCustomizationRequested(allAppsShortcutDeleteRequest)
             viewModel.onDialogShown()
 
             assertThat(
@@ -170,7 +170,7 @@ class ShortcutCustomizationViewModelTest : SysuiTestCase() {
     fun uiState_errorMessage_isEmptyByDefault() {
         testScope.runTest {
             val uiState by collectLastValue(viewModel.shortcutCustomizationUiState)
-            viewModel.onShortcutCustomizationRequested(standardAddCustomShortcutRequestInfo)
+            viewModel.onShortcutCustomizationRequested(allAppsShortcutAddRequest)
             viewModel.onDialogShown()
 
             assertThat((uiState as ShortcutCustomizationUiState.AddShortcutDialog).errorMessage)
@@ -298,7 +298,7 @@ class ShortcutCustomizationViewModelTest : SysuiTestCase() {
     }
 
     private suspend fun openAddShortcutDialogAndSetShortcut() {
-        viewModel.onShortcutCustomizationRequested(standardAddCustomShortcutRequestInfo)
+        viewModel.onShortcutCustomizationRequested(allAppsShortcutAddRequest)
         viewModel.onDialogShown()
 
         viewModel.onKeyPressed(keyDownEventWithActionKeyPressed)
@@ -308,7 +308,7 @@ class ShortcutCustomizationViewModelTest : SysuiTestCase() {
     }
 
     private suspend fun openDeleteShortcutDialogAndDeleteShortcut() {
-        viewModel.onShortcutCustomizationRequested(standardDeleteCustomShortcutRequestInfo)
+        viewModel.onShortcutCustomizationRequested(allAppsShortcutDeleteRequest)
         viewModel.onDialogShown()
 
         viewModel.deleteShortcutCurrentlyBeingCustomized()
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/CustomInputGesturesRepository.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/CustomInputGesturesRepository.kt
new file mode 100644
index 000000000000..9ffdafc549c7
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/CustomInputGesturesRepository.kt
@@ -0,0 +1,118 @@
+/*
+ * Copyright (C) 2024 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.keyboard.shortcut.data.repository
+
+import android.content.Context
+import android.content.Context.INPUT_SERVICE
+import android.hardware.input.InputGestureData
+import android.hardware.input.InputManager
+import android.hardware.input.InputManager.CUSTOM_INPUT_GESTURE_RESULT_ERROR_ALREADY_EXISTS
+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.InputSettings
+import android.util.Log
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.keyboard.shared.model.ShortcutCustomizationRequestResult
+import com.android.systemui.keyboard.shared.model.ShortcutCustomizationRequestResult.ERROR_OTHER
+import com.android.systemui.keyboard.shared.model.ShortcutCustomizationRequestResult.ERROR_RESERVED_COMBINATION
+import com.android.systemui.keyboard.shared.model.ShortcutCustomizationRequestResult.SUCCESS
+import com.android.systemui.settings.UserTracker
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.onStart
+import kotlinx.coroutines.withContext
+import javax.inject.Inject
+import kotlin.coroutines.CoroutineContext
+
+class CustomInputGesturesRepository
+@Inject
+constructor(private val userTracker: UserTracker,
+    @Background private val bgCoroutineContext: CoroutineContext)
+{
+
+    private val userContext: Context
+        get() = userTracker.createCurrentUserContext(userTracker.userContext)
+
+    // Input manager created with user context to provide correct user id when requesting custom
+    // shortcut
+    private val inputManager: InputManager
+        get() = userContext.getSystemService(INPUT_SERVICE) as InputManager
+
+    private val _customInputGesture = MutableStateFlow<List<InputGestureData>>(emptyList())
+
+    val customInputGestures =
+        _customInputGesture.onStart { refreshCustomInputGestures() }
+
+    private fun refreshCustomInputGestures() {
+        _customInputGesture.value = retrieveCustomInputGestures()
+    }
+
+    fun retrieveCustomInputGestures(): List<InputGestureData> {
+        return if (InputSettings.isCustomizableInputGesturesFeatureFlagEnabled()) {
+            inputManager.getCustomInputGestures(/* filter= */ InputGestureData.Filter.KEY)
+        } else emptyList()
+    }
+
+    suspend fun addCustomInputGesture(inputGesture: InputGestureData): ShortcutCustomizationRequestResult {
+        return withContext(bgCoroutineContext) {
+            when (val result = inputManager.addCustomInputGesture(inputGesture)) {
+                CUSTOM_INPUT_GESTURE_RESULT_SUCCESS -> {
+                    refreshCustomInputGestures()
+                    SUCCESS
+                }
+                CUSTOM_INPUT_GESTURE_RESULT_ERROR_ALREADY_EXISTS ->
+                    ERROR_RESERVED_COMBINATION
+
+                CUSTOM_INPUT_GESTURE_RESULT_ERROR_RESERVED_GESTURE ->
+                    ERROR_RESERVED_COMBINATION
+
+                else -> {
+                    Log.w(
+                        TAG,
+                        "Attempted to add inputGesture: $inputGesture " +
+                                "but ran into an error with code: $result",
+                    )
+                    ERROR_OTHER
+                }
+            }
+        }
+    }
+
+    suspend fun deleteCustomInputGesture(inputGesture: InputGestureData): ShortcutCustomizationRequestResult {
+        return withContext(bgCoroutineContext){
+            when (
+                val result = inputManager.removeCustomInputGesture(inputGesture)
+            ) {
+                CUSTOM_INPUT_GESTURE_RESULT_SUCCESS -> {
+                    refreshCustomInputGestures()
+                    SUCCESS
+                }
+                else -> {
+                    Log.w(
+                        TAG,
+                        "Attempted to delete inputGesture: $inputGesture " +
+                                "but ran into an error with code: $result",
+                    )
+                    ERROR_OTHER
+                }
+            }
+        }
+    }
+
+    private companion object {
+        private const val TAG = "CustomInputGesturesRepository"
+    }
+}
\ No newline at end of file
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/CustomShortcutCategoriesRepository.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/CustomShortcutCategoriesRepository.kt
index 321fd57d3e8b..d1bd51c23d45 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/CustomShortcutCategoriesRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/CustomShortcutCategoriesRepository.kt
@@ -17,17 +17,11 @@
 package com.android.systemui.keyboard.shortcut.data.repository
 
 import android.content.Context
-import android.content.Context.INPUT_SERVICE
 import android.hardware.input.InputGestureData
 import android.hardware.input.InputGestureData.Builder
 import android.hardware.input.InputGestureData.KeyTrigger
 import android.hardware.input.InputGestureData.createKeyTrigger
 import android.hardware.input.InputManager
-import android.hardware.input.InputManager.CUSTOM_INPUT_GESTURE_RESULT_ERROR_ALREADY_EXISTS
-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.InputSettings
-import android.hardware.input.KeyGestureEvent
 import android.hardware.input.KeyGestureEvent.KeyGestureType
 import android.util.Log
 import androidx.annotation.VisibleForTesting
@@ -44,9 +38,6 @@ import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCategoryType
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCustomizationRequestInfo
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutHelperState.Active
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutKey
-import com.android.systemui.settings.UserTracker
-import javax.inject.Inject
-import kotlin.coroutines.CoroutineContext
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -55,27 +46,22 @@ import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.withContext
+import javax.inject.Inject
+import kotlin.coroutines.CoroutineContext
 
 @SysUISingleton
 class CustomShortcutCategoriesRepository
 @Inject
 constructor(
     stateRepository: ShortcutHelperStateRepository,
-    private val userTracker: UserTracker,
     @Background private val backgroundScope: CoroutineScope,
     @Background private val bgCoroutineContext: CoroutineContext,
     private val shortcutCategoriesUtils: ShortcutCategoriesUtils,
     private val context: Context,
     private val inputGestureMaps: InputGestureMaps,
-) : ShortcutCategoriesRepository {
-
-    private val userContext: Context
-        get() = userTracker.createCurrentUserContext(userTracker.userContext)
-
-    // Input manager created with user context to provide correct user id when requesting custom
-    // shortcut
+    private val customInputGesturesRepository: CustomInputGesturesRepository,
     private val inputManager: InputManager
-        get() = userContext.getSystemService(INPUT_SERVICE) as InputManager
+) : ShortcutCategoriesRepository {
 
     private val _selectedKeyCombination = MutableStateFlow<KeyCombination?>(null)
     private val _shortcutBeingCustomized = mutableStateOf<ShortcutCustomizationRequestInfo?>(null)
@@ -125,14 +111,12 @@ constructor(
             )
 
     override val categories: Flow<List<ShortcutCategory>> =
-        activeInputDevice
-            .map { inputDevice ->
+        combine(activeInputDevice, customInputGesturesRepository.customInputGestures)
+        { inputDevice, inputGestures ->
                 if (inputDevice == null) {
                     emptyList()
                 } else {
-                    val customInputGesturesForUser: List<InputGestureData> =
-                        getCustomInputGestures()
-                    val sources = toInternalGroupSources(customInputGesturesForUser)
+                    val sources = toInternalGroupSources(inputGestures)
                     val supportedKeyCodes =
                         shortcutCategoriesUtils.fetchSupportedKeyCodes(
                             inputDevice.id,
@@ -181,56 +165,23 @@ constructor(
 
     private fun retrieveInputGestureDataForShortcutBeingDeleted(): InputGestureData? {
         val keyGestureType = getKeyGestureTypeFromShortcutBeingDeletedLabel()
-        return getCustomInputGestures().firstOrNull { it.action.keyGestureType() == keyGestureType }
+        return customInputGesturesRepository.retrieveCustomInputGestures()
+            .firstOrNull { it.action.keyGestureType() == keyGestureType }
     }
 
-    suspend fun confirmAndSetShortcutCurrentlyBeingCustomized():
-        ShortcutCustomizationRequestResult {
-        return withContext(bgCoroutineContext) {
-            val inputGestureData =
-                buildInputGestureDataForShortcutBeingCustomized()
-                    ?: return@withContext ShortcutCustomizationRequestResult.ERROR_OTHER
-
-            return@withContext when (inputManager.addCustomInputGesture(inputGestureData)) {
-                CUSTOM_INPUT_GESTURE_RESULT_SUCCESS -> ShortcutCustomizationRequestResult.SUCCESS
-                CUSTOM_INPUT_GESTURE_RESULT_ERROR_ALREADY_EXISTS ->
-                    ShortcutCustomizationRequestResult.ERROR_RESERVED_COMBINATION
+    suspend fun confirmAndSetShortcutCurrentlyBeingCustomized(): ShortcutCustomizationRequestResult {
+        val inputGestureData =
+            buildInputGestureDataForShortcutBeingCustomized()
+                ?: return ShortcutCustomizationRequestResult.ERROR_OTHER
 
-                CUSTOM_INPUT_GESTURE_RESULT_ERROR_RESERVED_GESTURE ->
-                    ShortcutCustomizationRequestResult.ERROR_RESERVED_COMBINATION
-
-                else -> ShortcutCustomizationRequestResult.ERROR_OTHER
-            }
-        }
-    }
-
-    suspend fun deleteShortcutCurrentlyBeingCustomized():
-        ShortcutCustomizationRequestResult {
-        return withContext(bgCoroutineContext) {
-            val inputGestureData =
-                retrieveInputGestureDataForShortcutBeingDeleted()
-                    ?: return@withContext ShortcutCustomizationRequestResult.ERROR_OTHER
-            return@withContext when (
-                val result = inputManager.removeCustomInputGesture(inputGestureData)
-            ) {
-                CUSTOM_INPUT_GESTURE_RESULT_SUCCESS -> ShortcutCustomizationRequestResult.SUCCESS
-                else -> {
-                    Log.w(
-                        TAG,
-                        "Attempted to delete shortcut being customized " +
-                            "${_shortcutBeingCustomized.value} but ran into an error. InputGestureData" +
-                            " = $inputGestureData, error code: $result",
-                    )
-                    ShortcutCustomizationRequestResult.ERROR_OTHER
-                }
-            }
-        }
+        return customInputGesturesRepository.addCustomInputGesture(inputGestureData)
     }
 
-    private fun getCustomInputGestures(): List<InputGestureData> {
-        return if (InputSettings.isCustomizableInputGesturesFeatureFlagEnabled()) {
-            inputManager.getCustomInputGestures(/* filter= */ InputGestureData.Filter.KEY)
-        } else emptyList()
+    suspend fun deleteShortcutCurrentlyBeingCustomized(): ShortcutCustomizationRequestResult {
+        val inputGestureData =
+            retrieveInputGestureDataForShortcutBeingDeleted()
+                ?: return ShortcutCustomizationRequestResult.ERROR_OTHER
+        return customInputGesturesRepository.deleteCustomInputGesture(inputGestureData)
     }
 
     private fun Builder.addKeyGestureTypeFromShortcutLabel(): Builder {
@@ -347,7 +298,7 @@ constructor(
     }
 
     private fun fetchGroupLabelByGestureType(
-        @KeyGestureEvent.KeyGestureType keyGestureType: Int
+        @KeyGestureType keyGestureType: Int
     ): String? {
         inputGestureMaps.gestureToInternalKeyboardShortcutGroupLabelResIdMap[keyGestureType]?.let {
             return context.getString(it)
@@ -355,7 +306,7 @@ constructor(
     }
 
     private fun fetchShortcutInfoLabelByGestureType(
-        @KeyGestureEvent.KeyGestureType keyGestureType: Int
+        @KeyGestureType keyGestureType: Int
     ): String? {
         inputGestureMaps.gestureToInternalKeyboardShortcutInfoLabelResIdMap[keyGestureType]?.let {
             return context.getString(it)
@@ -363,7 +314,7 @@ constructor(
     }
 
     private fun fetchShortcutCategoryTypeByGestureType(
-        @KeyGestureEvent.KeyGestureType keyGestureType: Int
+        @KeyGestureType keyGestureType: Int
     ): ShortcutCategoryType? {
         return inputGestureMaps.gestureToShortcutCategoryTypeMap[keyGestureType]
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/KeyboardShortcutHelperKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/KeyboardShortcutHelperKosmos.kt
index e7672ffabc79..2bff0c66889f 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/KeyboardShortcutHelperKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/KeyboardShortcutHelperKosmos.kt
@@ -22,6 +22,7 @@ import android.content.res.mainResources
 import android.hardware.input.fakeInputManager
 import android.view.windowManager
 import com.android.systemui.broadcast.broadcastDispatcher
+import com.android.systemui.keyboard.shortcut.data.repository.CustomInputGesturesRepository
 import com.android.systemui.keyboard.shortcut.data.repository.CustomShortcutCategoriesRepository
 import com.android.systemui.keyboard.shortcut.data.repository.DefaultShortcutCategoriesRepository
 import com.android.systemui.keyboard.shortcut.data.repository.InputGestureMaps
@@ -106,16 +107,21 @@ val Kosmos.defaultShortcutCategoriesRepository by
 
 val Kosmos.inputGestureMaps by Kosmos.Fixture { InputGestureMaps(applicationContext) }
 
+val Kosmos.customInputGesturesRepository by Kosmos.Fixture {
+    CustomInputGesturesRepository(userTracker, testDispatcher)
+}
+
 val Kosmos.customShortcutCategoriesRepository by
     Kosmos.Fixture {
         CustomShortcutCategoriesRepository(
             shortcutHelperStateRepository,
-            userTracker,
             applicationCoroutineScope,
             testDispatcher,
             shortcutCategoriesUtils,
             applicationContext,
             inputGestureMaps,
+            customInputGesturesRepository,
+            fakeInputManager.inputManager
         )
     }