Merge "new implementation of IME shortcuts retrieval" into main

8 files changed, 455 insertions, 25 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCategoriesRepository.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCategoriesRepository.kt
index 04bde26fdd88..c00bd6ff93b0 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCategoriesRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCategoriesRepository.kt
@@ -16,11 +16,23 @@
 
 package com.android.systemui.keyboard.shortcut.data.repository
 
+import android.view.KeyEvent
+import android.view.KeyboardShortcutGroup
+import android.view.KeyboardShortcutInfo
+import android.view.WindowManager
+import android.view.WindowManager.KeyboardShortcutsReceiver
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.keyboard.shortcut.data.source.MultitaskingShortcutsSource
 import com.android.systemui.keyboard.shortcut.data.source.SystemShortcutsSource
+import com.android.systemui.keyboard.shortcut.shared.model.Shortcut
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCategory
+import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCategoryType
+import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCommand
+import com.android.systemui.keyboard.shortcut.shared.model.ShortcutHelperState.Active
+import com.android.systemui.keyboard.shortcut.shared.model.shortcutCategory
 import javax.inject.Inject
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.suspendCancellableCoroutine
 
 @SysUISingleton
 class ShortcutHelperCategoriesRepository
@@ -28,11 +40,77 @@ class ShortcutHelperCategoriesRepository
 constructor(
     private val systemShortcutsSource: SystemShortcutsSource,
     private val multitaskingShortcutsSource: MultitaskingShortcutsSource,
+    private val windowManager: WindowManager,
+    shortcutHelperStateRepository: ShortcutHelperStateRepository
 ) {
 
-    fun systemShortcutsCategory(): ShortcutCategory =
-        systemShortcutsSource.systemShortcutsCategory()
+    val systemShortcutsCategory =
+        shortcutHelperStateRepository.state.map {
+            if (it is Active) systemShortcutsSource.systemShortcutsCategory() else null
+        }
 
-    fun multitaskingShortcutsCategory(): ShortcutCategory =
-        multitaskingShortcutsSource.multitaskingShortcutCategory()
+    val multitaskingShortcutsCategory =
+        shortcutHelperStateRepository.state.map {
+            if (it is Active) multitaskingShortcutsSource.multitaskingShortcutCategory() else null
+        }
+
+    val imeShortcutsCategory =
+        shortcutHelperStateRepository.state.map {
+            if (it is Active) retrieveImeShortcuts(it.deviceId) else null
+        }
+
+    private suspend fun retrieveImeShortcuts(deviceId: Int): ShortcutCategory {
+        return suspendCancellableCoroutine { continuation ->
+            val shortcutsReceiver = KeyboardShortcutsReceiver { shortcutGroups ->
+                continuation.resumeWith(Result.success(toShortcutCategory(shortcutGroups)))
+            }
+            windowManager.requestImeKeyboardShortcuts(shortcutsReceiver, deviceId)
+        }
+    }
+
+    private fun toShortcutCategory(shortcutGroups: List<KeyboardShortcutGroup>) =
+        shortcutCategory(ShortcutCategoryType.IME) {
+            shortcutGroups.map { shortcutGroup ->
+                subCategory(shortcutGroup.label.toString(), toShortcuts(shortcutGroup.items))
+            }
+        }
+
+    private fun toShortcuts(infoList: List<KeyboardShortcutInfo>) =
+        infoList.mapNotNull { toShortcut(it) }
+
+    private fun toShortcut(shortcutInfo: KeyboardShortcutInfo): Shortcut? {
+        val shortcutCommand = toShortcutCommand(shortcutInfo)
+        return if (shortcutCommand == null) null
+        else Shortcut(label = shortcutInfo.label!!.toString(), commands = listOf(shortcutCommand))
+    }
+
+    private fun toShortcutCommand(info: KeyboardShortcutInfo): ShortcutCommand? {
+        val keyCodes = mutableListOf<Int>()
+        var remainingModifiers = info.modifiers
+        SUPPORTED_MODIFIERS.forEach { supportedModifier ->
+            if ((supportedModifier and remainingModifiers) != 0) {
+                keyCodes += supportedModifier
+                // "Remove" the modifier from the remaining modifiers
+                remainingModifiers = remainingModifiers and supportedModifier.inv()
+            }
+        }
+        if (remainingModifiers != 0) {
+            // There is a remaining modifier we don't support
+            return null
+        }
+        keyCodes += info.keycode
+        return ShortcutCommand(keyCodes)
+    }
+
+    companion object {
+        private val SUPPORTED_MODIFIERS =
+            listOf(
+                KeyEvent.META_META_ON,
+                KeyEvent.META_CTRL_ON,
+                KeyEvent.META_ALT_ON,
+                KeyEvent.META_SHIFT_ON,
+                KeyEvent.META_SYM_ON,
+                KeyEvent.META_FUNCTION_ON
+            )
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/domain/interactor/ShortcutHelperCategoriesInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/domain/interactor/ShortcutHelperCategoriesInteractor.kt
index 883407c5f6a4..57d4b4a02fce 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/domain/interactor/ShortcutHelperCategoriesInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/domain/interactor/ShortcutHelperCategoriesInteractor.kt
@@ -18,30 +18,56 @@ package com.android.systemui.keyboard.shortcut.domain.interactor
 
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.keyboard.shortcut.data.repository.ShortcutHelperCategoriesRepository
-import com.android.systemui.keyboard.shortcut.data.repository.ShortcutHelperStateRepository
+import com.android.systemui.keyboard.shortcut.shared.model.Shortcut
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCategory
-import com.android.systemui.keyboard.shortcut.shared.model.ShortcutHelperState
+import com.android.systemui.keyboard.shortcut.shared.model.ShortcutSubCategory
 import javax.inject.Inject
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.map
 
 @SysUISingleton
 class ShortcutHelperCategoriesInteractor
 @Inject
 constructor(
-    stateRepository: ShortcutHelperStateRepository,
     categoriesRepository: ShortcutHelperCategoriesRepository,
 ) {
 
+    private val systemsShortcutCategory = categoriesRepository.systemShortcutsCategory
+    private val multitaskingShortcutsCategory = categoriesRepository.multitaskingShortcutsCategory
+    private val imeShortcutsCategory =
+        categoriesRepository.imeShortcutsCategory.map { groupSubCategoriesInCategory(it) }
+
     val shortcutCategories: Flow<List<ShortcutCategory>> =
-        stateRepository.state.map { state ->
-            when (state) {
-                is ShortcutHelperState.Active ->
-                    listOf(
-                        categoriesRepository.systemShortcutsCategory(),
-                        categoriesRepository.multitaskingShortcutsCategory()
-                    )
-                is ShortcutHelperState.Inactive -> emptyList()
-            }
+        combine(systemsShortcutCategory, multitaskingShortcutsCategory, imeShortcutsCategory) {
+            shortcutCategories ->
+            shortcutCategories.filterNotNull()
+        }
+
+    private fun groupSubCategoriesInCategory(
+        shortcutCategory: ShortcutCategory?
+    ): ShortcutCategory? {
+        if (shortcutCategory == null) {
+            return null
         }
+        val subCategoriesWithGroupedShortcuts =
+            shortcutCategory.subCategories.map {
+                ShortcutSubCategory(
+                    label = it.label,
+                    shortcuts = groupShortcutsInSubcategory(it.shortcuts)
+                )
+            }
+        return ShortcutCategory(
+            type = shortcutCategory.type,
+            subCategories = subCategoriesWithGroupedShortcuts
+        )
+    }
+
+    private fun groupShortcutsInSubcategory(shortcuts: List<Shortcut>) =
+        shortcuts
+            .groupBy { it.label }
+            .entries
+            .map { (commonLabel, groupedShortcuts) ->
+                Shortcut(label = commonLabel, commands = groupedShortcuts.flatMap { it.commands })
+            }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/domain/interactor/ShortcutHelperStateInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/domain/interactor/ShortcutHelperStateInteractor.kt
index 3d707f70538e..299628ee19f6 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/domain/interactor/ShortcutHelperStateInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/domain/interactor/ShortcutHelperStateInteractor.kt
@@ -26,6 +26,7 @@ import com.android.systemui.shared.system.QuickStepContract
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.launch
 
 @SysUISingleton
@@ -38,7 +39,7 @@ constructor(
     private val repository: ShortcutHelperStateRepository
 ) {
 
-    val state: Flow<ShortcutHelperState> = repository.state
+    val state: Flow<ShortcutHelperState> = repository.state.asStateFlow()
 
     fun onViewClosed() {
         repository.hide()
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/shared/model/ShortcutCategory.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/shared/model/ShortcutCategory.kt
index c5e8d2c12fda..3ac7fa8f8ece 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/shared/model/ShortcutCategory.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/shared/model/ShortcutCategory.kt
@@ -19,6 +19,7 @@ package com.android.systemui.keyboard.shortcut.shared.model
 enum class ShortcutCategoryType {
     SYSTEM,
     MULTI_TASKING,
+    IME
 }
 
 data class ShortcutCategory(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCategoriesRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCategoriesRepositoryTest.kt
new file mode 100644
index 000000000000..6985439f0552
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCategoriesRepositoryTest.kt
@@ -0,0 +1,130 @@
+/*
+ * 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.view.KeyEvent
+import android.view.KeyboardShortcutGroup
+import android.view.KeyboardShortcutInfo
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.keyboard.shortcut.shared.model.Shortcut
+import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCategoryType
+import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCommand
+import com.android.systemui.keyboard.shortcut.shared.model.shortcutCategory
+import com.android.systemui.keyboard.shortcut.shortcutHelperCategoriesRepository
+import com.android.systemui.keyboard.shortcut.shortcutHelperTestHelper
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class ShortcutHelperCategoriesRepositoryTest : SysuiTestCase() {
+    @OptIn(ExperimentalCoroutinesApi::class)
+    private val kosmos = testKosmos().also { it.testDispatcher = UnconfinedTestDispatcher() }
+    private val repo = kosmos.shortcutHelperCategoriesRepository
+    private val helper = kosmos.shortcutHelperTestHelper
+    private val testScope = kosmos.testScope
+
+    @Test
+    fun stateActive_imeShortcuts_shortcutInfoCorrectlyConverted() =
+        testScope.runTest {
+            helper.setImeShortcuts(imeShortcutsGroupWithPreviousLanguageSwitchShortcut)
+            val imeShortcutCategory by collectLastValue(repo.imeShortcutsCategory)
+
+            helper.showFromActivity()
+
+            assertThat(imeShortcutCategory)
+                .isEqualTo(expectedImeShortcutCategoryWithPreviousLanguageSwitchShortcut)
+        }
+
+    @Test
+    fun stateActive_imeShortcuts_discardUnsupportedShortcutInfoModifiers() =
+        testScope.runTest {
+            helper.setImeShortcuts(imeShortcutsGroupWithUnsupportedShortcutModifiers)
+            val imeShortcutCategory by collectLastValue(repo.imeShortcutsCategory)
+
+            helper.showFromActivity()
+
+            assertThat(imeShortcutCategory)
+                .isEqualTo(expectedImeShortcutCategoryWithDiscardedUnsupportedShortcuts)
+        }
+
+    private val switchToPreviousLanguageCommand =
+        ShortcutCommand(
+            listOf(KeyEvent.META_CTRL_ON, KeyEvent.META_SHIFT_ON, KeyEvent.KEYCODE_SPACE)
+        )
+
+    private val expectedImeShortcutCategoryWithDiscardedUnsupportedShortcuts =
+        shortcutCategory(ShortcutCategoryType.IME) { subCategory("input", emptyList()) }
+
+    private val switchToPreviousLanguageKeyboardShortcutInfo =
+        KeyboardShortcutInfo(
+            /* label = */ "switch to previous language",
+            /* keycode = */ switchToPreviousLanguageCommand.keyCodes[2],
+            /* modifiers = */ switchToPreviousLanguageCommand.keyCodes[0] or
+                switchToPreviousLanguageCommand.keyCodes[1],
+        )
+
+    private val expectedImeShortcutCategoryWithPreviousLanguageSwitchShortcut =
+        shortcutCategory(ShortcutCategoryType.IME) {
+            subCategory(
+                "input",
+                listOf(
+                    Shortcut(
+                        switchToPreviousLanguageKeyboardShortcutInfo.label!!.toString(),
+                        listOf(switchToPreviousLanguageCommand)
+                    )
+                )
+            )
+        }
+
+    private val imeShortcutsGroupWithPreviousLanguageSwitchShortcut =
+        listOf(
+            KeyboardShortcutGroup(
+                "input",
+                listOf(
+                    switchToPreviousLanguageKeyboardShortcutInfo,
+                )
+            )
+        )
+
+    private val shortcutInfoWithUnsupportedModifier =
+        KeyboardShortcutInfo(
+            /* label = */ "unsupported shortcut",
+            /* keycode = */ KeyEvent.KEYCODE_SPACE,
+            /* modifiers = */ 32
+        )
+
+    private val imeShortcutsGroupWithUnsupportedShortcutModifiers =
+        listOf(
+            KeyboardShortcutGroup(
+                "input",
+                listOf(
+                    shortcutInfoWithUnsupportedModifier,
+                )
+            )
+        )
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyboard/shortcut/domain/interactor/ShortcutHelperCategoriesInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyboard/shortcut/domain/interactor/ShortcutHelperCategoriesInteractorTest.kt
index 9c9e48e9200e..5c7ce3ef1009 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyboard/shortcut/domain/interactor/ShortcutHelperCategoriesInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyboard/shortcut/domain/interactor/ShortcutHelperCategoriesInteractorTest.kt
@@ -16,10 +16,17 @@
 
 package com.android.systemui.keyboard.shortcut.domain.interactor
 
+import android.view.KeyEvent
+import android.view.KeyboardShortcutGroup
+import android.view.KeyboardShortcutInfo
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCategory
+import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCategoryType
+import com.android.systemui.keyboard.shortcut.shared.model.ShortcutSubCategory
+import com.android.systemui.keyboard.shortcut.shared.model.shortcut
 import com.android.systemui.keyboard.shortcut.shortcutHelperCategoriesInteractor
 import com.android.systemui.keyboard.shortcut.shortcutHelperMultiTaskingShortcutsSource
 import com.android.systemui.keyboard.shortcut.shortcutHelperSystemShortcutsSource
@@ -57,6 +64,7 @@ class ShortcutHelperCategoriesInteractorTest : SysuiTestCase() {
     @Test
     fun categories_stateActive_emitsAllCategoriesInOrder() =
         testScope.runTest {
+            helper.setImeShortcuts(imeShortcutGroups)
             val categories by collectLastValue(interactor.shortcutCategories)
 
             helper.showFromActivity()
@@ -64,7 +72,8 @@ class ShortcutHelperCategoriesInteractorTest : SysuiTestCase() {
             assertThat(categories)
                 .containsExactly(
                     systemShortcutsSource.systemShortcutsCategory(),
-                    multitaskingShortcutsSource.multitaskingShortcutCategory()
+                    multitaskingShortcutsSource.multitaskingShortcutCategory(),
+                    imeShortcutCategory
                 )
                 .inOrder()
         }
@@ -78,4 +87,165 @@ class ShortcutHelperCategoriesInteractorTest : SysuiTestCase() {
 
             assertThat(categories).isEmpty()
         }
+
+    fun categories_stateActive_emitsGroupedShortcuts() =
+        testScope.runTest {
+            helper.setImeShortcuts(imeShortcutsGroupsWithDuplicateLabels)
+            val categories by collectLastValue(interactor.shortcutCategories)
+
+            helper.showFromActivity()
+
+            assertThat(categories)
+                .containsExactly(
+                    systemShortcutsSource.systemShortcutsCategory(),
+                    multitaskingShortcutsSource.multitaskingShortcutCategory(),
+                    expectedGroupedShortcutCategories
+                )
+        }
+
+    private val switchToNextLanguageShortcut =
+        shortcut(label = "switch to next language") {
+            command(KeyEvent.META_CTRL_ON, KeyEvent.KEYCODE_SPACE)
+        }
+
+    private val switchToNextLanguageKeyboardShortcutInfo =
+        KeyboardShortcutInfo(
+            /* label = */ switchToNextLanguageShortcut.label,
+            /* keycode = */ switchToNextLanguageShortcut.commands[0].keyCodes[1],
+            /* modifiers = */ switchToNextLanguageShortcut.commands[0].keyCodes[0],
+        )
+
+    private val switchToNextLanguageShortcutAlternative =
+        shortcut("switch to next language") {
+            command(KeyEvent.META_CTRL_ON, KeyEvent.KEYCODE_SPACE)
+        }
+
+    private val switchToNextLanguageKeyboardShortcutInfoAlternative =
+        KeyboardShortcutInfo(
+            /* label = */ switchToNextLanguageShortcutAlternative.label,
+            /* keycode = */ switchToNextLanguageShortcutAlternative.commands[0].keyCodes[1],
+            /* modifiers = */ switchToNextLanguageShortcutAlternative.commands[0].keyCodes[0],
+        )
+
+    private val switchToPreviousLanguageShortcut =
+        shortcut("switch to previous language") {
+            command(
+                KeyEvent.META_SHIFT_ON,
+                KeyEvent.KEYCODE_SPACE,
+            )
+        }
+
+    private val switchToPreviousLanguageKeyboardShortcutInfo =
+        KeyboardShortcutInfo(
+            /* label = */ switchToPreviousLanguageShortcut.label,
+            /* keycode = */ switchToPreviousLanguageShortcut.commands[0].keyCodes[1],
+            /* modifiers = */ switchToPreviousLanguageShortcut.commands[0].keyCodes[0],
+        )
+
+    private val switchToPreviousLanguageShortcutAlternative =
+        shortcut("switch to previous language") {
+            command(
+                KeyEvent.META_SHIFT_ON,
+                KeyEvent.KEYCODE_SPACE,
+            )
+        }
+
+    private val switchToPreviousLanguageKeyboardShortcutInfoAlternative =
+        KeyboardShortcutInfo(
+            /* label = */ switchToPreviousLanguageShortcutAlternative.label,
+            /* keycode = */ switchToPreviousLanguageShortcutAlternative.commands[0].keyCodes[1],
+            /* modifiers = */ switchToPreviousLanguageShortcutAlternative.commands[0].keyCodes[0],
+        )
+
+    private val showOnscreenKeyboardShortcut =
+        shortcut(label = "Show on-screen keyboard") {
+            command(KeyEvent.META_ALT_ON, KeyEvent.KEYCODE_K)
+        }
+
+    private val showOnScreenKeyboardShortcutInfo =
+        KeyboardShortcutInfo(
+            /* label = */ showOnscreenKeyboardShortcut.label,
+            /* keycode = */ showOnscreenKeyboardShortcut.commands[0].keyCodes[1],
+            /* modifiers = */ showOnscreenKeyboardShortcut.commands[0].keyCodes[0],
+        )
+
+    private val accessClipboardShortcut =
+        shortcut(label = "Access clipboard") { command(KeyEvent.META_ALT_ON, KeyEvent.KEYCODE_V) }
+
+    private val accessClipboardShortcutInfo =
+        KeyboardShortcutInfo(
+            /* label = */ accessClipboardShortcut.label,
+            /* keycode = */ accessClipboardShortcut.commands[0].keyCodes[1],
+            /* modifiers = */ accessClipboardShortcut.commands[0].keyCodes[0],
+        )
+
+    private val imeShortcutGroups =
+        listOf(
+            KeyboardShortcutGroup(
+                /* label = */ "input",
+                /* shortcutInfoList = */ listOf(
+                    switchToNextLanguageKeyboardShortcutInfo,
+                    switchToPreviousLanguageKeyboardShortcutInfo
+                )
+            )
+        )
+
+    private val imeShortcutCategory =
+        ShortcutCategory(
+            type = ShortcutCategoryType.IME,
+            subCategories =
+                listOf(
+                    ShortcutSubCategory(
+                        imeShortcutGroups[0].label.toString(),
+                        listOf(switchToNextLanguageShortcut, switchToPreviousLanguageShortcut)
+                    )
+                )
+        )
+
+    private val imeShortcutsGroupsWithDuplicateLabels =
+        listOf(
+            KeyboardShortcutGroup(
+                "input",
+                listOf(
+                    switchToNextLanguageKeyboardShortcutInfo,
+                    switchToNextLanguageKeyboardShortcutInfoAlternative,
+                    switchToPreviousLanguageKeyboardShortcutInfo,
+                    switchToPreviousLanguageKeyboardShortcutInfoAlternative
+                )
+            ),
+            KeyboardShortcutGroup(
+                "Gboard",
+                listOf(
+                    showOnScreenKeyboardShortcutInfo,
+                    accessClipboardShortcutInfo,
+                )
+            )
+        )
+
+    private val expectedGroupedShortcutCategories =
+        ShortcutCategory(
+            type = ShortcutCategoryType.IME,
+            subCategories =
+                listOf(
+                    ShortcutSubCategory(
+                        imeShortcutsGroupsWithDuplicateLabels[0].label.toString(),
+                        listOf(
+                            switchToNextLanguageShortcut.copy(
+                                commands =
+                                    switchToNextLanguageShortcut.commands +
+                                        switchToNextLanguageShortcutAlternative.commands
+                            ),
+                            switchToPreviousLanguageShortcut.copy(
+                                commands =
+                                    switchToPreviousLanguageShortcut.commands +
+                                        switchToPreviousLanguageShortcutAlternative.commands
+                            )
+                        ),
+                    ),
+                    ShortcutSubCategory(
+                        imeShortcutsGroupsWithDuplicateLabels[1].label.toString(),
+                        listOf(showOnscreenKeyboardShortcut, accessClipboardShortcut),
+                    )
+                )
+        )
 }
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 f51036f3b54b..e00f9806f6a2 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
@@ -19,6 +19,7 @@ package com.android.systemui.keyboard.shortcut
 import android.content.applicationContext
 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.ShortcutHelperCategoriesRepository
 import com.android.systemui.keyboard.shortcut.data.repository.ShortcutHelperStateRepository
@@ -59,6 +60,8 @@ val Kosmos.shortcutHelperCategoriesRepository by
         ShortcutHelperCategoriesRepository(
             shortcutHelperSystemShortcutsSource,
             shortcutHelperMultiTaskingShortcutsSource,
+            windowManager,
+            shortcutHelperStateRepository
         )
     }
 
@@ -68,7 +71,8 @@ val Kosmos.shortcutHelperTestHelper by
             shortcutHelperStateRepository,
             applicationContext,
             broadcastDispatcher,
-            fakeCommandQueue
+            fakeCommandQueue,
+            windowManager
         )
     }
 
@@ -83,12 +87,7 @@ val Kosmos.shortcutHelperStateInteractor by
     }
 
 val Kosmos.shortcutHelperCategoriesInteractor by
-    Kosmos.Fixture {
-        ShortcutHelperCategoriesInteractor(
-            shortcutHelperStateRepository,
-            shortcutHelperCategoriesRepository
-        )
-    }
+    Kosmos.Fixture { ShortcutHelperCategoriesInteractor(shortcutHelperCategoriesRepository) }
 
 val Kosmos.shortcutHelperViewModel by
     Kosmos.Fixture { ShortcutHelperViewModel(testDispatcher, shortcutHelperStateInteractor) }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperTestHelper.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperTestHelper.kt
index 36608ff57c93..40510db24f47 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperTestHelper.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperTestHelper.kt
@@ -18,20 +18,45 @@ package com.android.systemui.keyboard.shortcut.data.repository
 
 import android.content.Context
 import android.content.Intent
+import android.view.KeyboardShortcutGroup
+import android.view.WindowManager
+import android.view.WindowManager.KeyboardShortcutsReceiver
 import com.android.systemui.broadcast.FakeBroadcastDispatcher
 import com.android.systemui.keyguard.data.repository.FakeCommandQueue
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.whenever
 
 class ShortcutHelperTestHelper(
     repo: ShortcutHelperStateRepository,
     private val context: Context,
     private val fakeBroadcastDispatcher: FakeBroadcastDispatcher,
     private val fakeCommandQueue: FakeCommandQueue,
+    windowManager: WindowManager
 ) {
 
+    companion object {
+        const val DEFAULT_DEVICE_ID = 123
+    }
+
+    private var imeShortcuts: List<KeyboardShortcutGroup> = emptyList()
+
     init {
+        whenever(windowManager.requestImeKeyboardShortcuts(any(), any())).thenAnswer {
+            val keyboardShortcutReceiver = it.getArgument<KeyboardShortcutsReceiver>(0)
+            keyboardShortcutReceiver.onKeyboardShortcutsReceived(imeShortcuts)
+            return@thenAnswer Unit
+        }
         repo.start()
     }
 
+    /**
+     * Use this method to set what ime shortcuts should be returned from windowManager in tests. By
+     * default windowManager.requestImeKeyboardShortcuts will return emptyList. See init block.
+     */
+    fun setImeShortcuts(imeShortcuts: List<KeyboardShortcutGroup>) {
+        this.imeShortcuts = imeShortcuts
+    }
+
     fun hideThroughCloseSystemDialogs() {
         fakeBroadcastDispatcher.sendIntentToMatchingReceiversOnly(
             context,