Merge changes I09c162e3,I4af5b80f into main

* changes:
  Shortcut Helper - Add "back", "home", and "recents" keys in system shortcuts
  Shortcut Helper - Filter shortcut commands containing unsupported keys
diff --git a/packages/SystemUI/res/drawable/ic_arrow_back_2.xml b/packages/SystemUI/res/drawable/ic_arrow_back_2.xml
new file mode 100644
index 0000000..8522d38
--- /dev/null
+++ b/packages/SystemUI/res/drawable/ic_arrow_back_2.xml
@@ -0,0 +1,24 @@
+<!--
+  ~ 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.
+  -->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportHeight="960"
+    android:viewportWidth="960">
+    <path
+        android:fillColor="#000"
+        android:pathData="M640,760L200,480L640,200L640,760ZM560,480L560,480L560,480ZM560,614L560,346L350,480L560,614Z" />
+</vector>
\ No newline at end of file
diff --git a/packages/SystemUI/res/drawable/ic_check_box_outline_blank.xml b/packages/SystemUI/res/drawable/ic_check_box_outline_blank.xml
new file mode 100644
index 0000000..f413d900
--- /dev/null
+++ b/packages/SystemUI/res/drawable/ic_check_box_outline_blank.xml
@@ -0,0 +1,24 @@
+<!--
+  ~ 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.
+  -->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="960"
+    android:viewportHeight="960">
+    <path
+        android:fillColor="#000"
+        android:pathData="M200,840Q167,840 143.5,816.5Q120,793 120,760L120,200Q120,167 143.5,143.5Q167,120 200,120L760,120Q793,120 816.5,143.5Q840,167 840,200L840,760Q840,793 816.5,816.5Q793,840 760,840L200,840ZM200,760L760,760Q760,760 760,760Q760,760 760,760L760,200Q760,200 760,200Q760,200 760,200L200,200Q200,200 200,200Q200,200 200,200L200,760Q200,760 200,760Q200,760 200,760Z"/>
+</vector>
diff --git a/packages/SystemUI/res/drawable/ic_radio_button_unchecked.xml b/packages/SystemUI/res/drawable/ic_radio_button_unchecked.xml
new file mode 100644
index 0000000..5bf914b
--- /dev/null
+++ b/packages/SystemUI/res/drawable/ic_radio_button_unchecked.xml
@@ -0,0 +1,24 @@
+<!--
+  ~ 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.
+  -->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="960"
+    android:viewportHeight="960">
+    <path
+        android:fillColor="#000"
+        android:pathData="M480,880Q397,880 324,848.5Q251,817 197,763Q143,709 111.5,636Q80,563 80,480Q80,397 111.5,324Q143,251 197,197Q251,143 324,111.5Q397,80 480,80Q563,80 636,111.5Q709,143 763,197Q817,251 848.5,324Q880,397 880,480Q880,563 848.5,636Q817,709 763,763Q709,817 636,848.5Q563,880 480,880ZM480,800Q614,800 707,707Q800,614 800,480Q800,346 707,253Q614,160 480,160Q346,160 253,253Q160,346 160,480Q160,614 253,707Q346,800 480,800ZM480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Z"/>
+</vector>
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 495e8f3..85bd0b0 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
@@ -71,6 +71,35 @@
     stateRepository: ShortcutHelperStateRepository
 ) {
 
+    private val sources =
+        listOf(
+            InternalGroupsSource(
+                source = systemShortcutsSource,
+                isTrusted = true,
+                typeProvider = { System }
+            ),
+            InternalGroupsSource(
+                source = multitaskingShortcutsSource,
+                isTrusted = true,
+                typeProvider = { MultiTasking }
+            ),
+            InternalGroupsSource(
+                source = appCategoriesShortcutsSource,
+                isTrusted = true,
+                typeProvider = { AppCategories }
+            ),
+            InternalGroupsSource(
+                source = inputShortcutsSource,
+                isTrusted = false,
+                typeProvider = { InputMethodEditor }
+            ),
+            InternalGroupsSource(
+                source = currentAppShortcutsSource,
+                isTrusted = false,
+                typeProvider = { groups -> getCurrentAppShortcutCategoryType(groups) }
+            ),
+        )
+
     private val activeInputDevice =
         stateRepository.state.map {
             if (it is Active) {
@@ -82,17 +111,20 @@
 
     val categories: Flow<List<ShortcutCategory>> =
         activeInputDevice
-            .map {
-                if (it == null) {
+            .map { inputDevice ->
+                if (inputDevice == null) {
                     return@map emptyList()
                 }
-                return@map listOfNotNull(
-                    fetchSystemShortcuts(it),
-                    fetchMultiTaskingShortcuts(it),
-                    fetchAppCategoriesShortcuts(it),
-                    fetchImeShortcuts(it),
-                    fetchCurrentAppShortcuts(it),
-                )
+                val groupsFromAllSources = sources.map { it.source.shortcutGroups(inputDevice.id) }
+                val supportedKeyCodes = fetchSupportedKeyCodes(inputDevice.id, groupsFromAllSources)
+                return@map sources.mapIndexedNotNull { index, internalGroupsSource ->
+                    fetchShortcutCategory(
+                        internalGroupsSource,
+                        groupsFromAllSources[index],
+                        inputDevice,
+                        supportedKeyCodes,
+                    )
+                }
             }
             .stateIn(
                 scope = backgroundScope,
@@ -100,49 +132,22 @@
                 initialValue = emptyList(),
             )
 
-    private suspend fun fetchSystemShortcuts(inputDevice: InputDevice) =
-        toShortcutCategory(
-            inputDevice.keyCharacterMap,
-            System,
-            systemShortcutsSource.shortcutGroups(inputDevice.id),
-            keepIcons = true,
-        )
-
-    private suspend fun fetchMultiTaskingShortcuts(inputDevice: InputDevice) =
-        toShortcutCategory(
-            inputDevice.keyCharacterMap,
-            MultiTasking,
-            multitaskingShortcutsSource.shortcutGroups(inputDevice.id),
-            keepIcons = true,
-        )
-
-    private suspend fun fetchAppCategoriesShortcuts(inputDevice: InputDevice) =
-        toShortcutCategory(
-            inputDevice.keyCharacterMap,
-            AppCategories,
-            appCategoriesShortcutsSource.shortcutGroups(inputDevice.id),
-            keepIcons = true,
-        )
-
-    private suspend fun fetchImeShortcuts(inputDevice: InputDevice) =
-        toShortcutCategory(
-            inputDevice.keyCharacterMap,
-            InputMethodEditor,
-            inputShortcutsSource.shortcutGroups(inputDevice.id),
-            keepIcons = false,
-        )
-
-    private suspend fun fetchCurrentAppShortcuts(inputDevice: InputDevice): ShortcutCategory? {
-        val shortcutGroups = currentAppShortcutsSource.shortcutGroups(inputDevice.id)
-        val categoryType = getCurrentAppShortcutCategoryType(shortcutGroups)
-        return if (categoryType == null) {
+    private fun fetchShortcutCategory(
+        internalGroupsSource: InternalGroupsSource,
+        groups: List<KeyboardShortcutGroup>,
+        inputDevice: InputDevice,
+        supportedKeyCodes: Set<Int>,
+    ): ShortcutCategory? {
+        val type = internalGroupsSource.typeProvider(groups)
+        return if (type == null) {
             null
         } else {
             toShortcutCategory(
                 inputDevice.keyCharacterMap,
-                categoryType,
-                shortcutGroups,
-                keepIcons = false
+                type,
+                groups,
+                internalGroupsSource.isTrusted,
+                supportedKeyCodes,
             )
         }
     }
@@ -162,13 +167,19 @@
         type: ShortcutCategoryType,
         shortcutGroups: List<KeyboardShortcutGroup>,
         keepIcons: Boolean,
+        supportedKeyCodes: Set<Int>,
     ): ShortcutCategory? {
         val subCategories =
             shortcutGroups
                 .map { shortcutGroup ->
                     ShortcutSubCategory(
                         shortcutGroup.label.toString(),
-                        toShortcuts(keyCharacterMap, shortcutGroup.items, keepIcons)
+                        toShortcuts(
+                            keyCharacterMap,
+                            shortcutGroup.items,
+                            keepIcons,
+                            supportedKeyCodes,
+                        )
                     )
                 }
                 .filter { it.shortcuts.isNotEmpty() }
@@ -184,7 +195,15 @@
         keyCharacterMap: KeyCharacterMap,
         infoList: List<KeyboardShortcutInfo>,
         keepIcons: Boolean,
-    ) = infoList.mapNotNull { toShortcut(keyCharacterMap, it, keepIcons) }
+        supportedKeyCodes: Set<Int>,
+    ) =
+        infoList
+            .filter {
+                // Allow KEYCODE_UNKNOWN (0) because shortcuts can have just modifiers and no
+                // keycode, or they could have a baseCharacter instead of a keycode.
+                it.keycode == KeyEvent.KEYCODE_UNKNOWN || supportedKeyCodes.contains(it.keycode)
+            }
+            .mapNotNull { toShortcut(keyCharacterMap, it, keepIcons) }
 
     private fun toShortcut(
         keyCharacterMap: KeyCharacterMap,
@@ -268,6 +287,29 @@
         return null
     }
 
+    private suspend fun fetchSupportedKeyCodes(
+        deviceId: Int,
+        groupsFromAllSources: List<List<KeyboardShortcutGroup>>
+    ): Set<Int> =
+        withContext(backgroundDispatcher) {
+            val allUsedKeyCodes =
+                groupsFromAllSources
+                    .flatMap { groups -> groups.flatMap { group -> group.items } }
+                    .map { info -> info.keycode }
+                    .distinct()
+            val keyCodesSupported =
+                inputManager.deviceHasKeys(deviceId, allUsedKeyCodes.toIntArray())
+            return@withContext allUsedKeyCodes
+                .filterIndexed { index, _ -> keyCodesSupported[index] }
+                .toSet()
+        }
+
+    private class InternalGroupsSource(
+        val source: KeyboardShortcutGroupsSource,
+        val isTrusted: Boolean,
+        val typeProvider: (groups: List<KeyboardShortcutGroup>) -> ShortcutCategoryType?,
+    )
+
     companion object {
         private const val TAG = "SHCategoriesRepo"
 
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperKeys.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperKeys.kt
index cbe6fc7..8db16fa 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperKeys.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperKeys.kt
@@ -98,6 +98,7 @@
 import android.view.KeyEvent.KEYCODE_PAGE_DOWN
 import android.view.KeyEvent.KEYCODE_PAGE_UP
 import android.view.KeyEvent.KEYCODE_PERIOD
+import android.view.KeyEvent.KEYCODE_RECENT_APPS
 import android.view.KeyEvent.KEYCODE_SCROLL_LOCK
 import android.view.KeyEvent.KEYCODE_SHIFT_LEFT
 import android.view.KeyEvent.KEYCODE_SHIFT_RIGHT
@@ -118,6 +119,9 @@
     val keyIcons =
         mapOf(
             META_META_ON to R.drawable.ic_ksh_key_meta,
+            KEYCODE_BACK to R.drawable.ic_arrow_back_2,
+            KEYCODE_HOME to R.drawable.ic_radio_button_unchecked,
+            KEYCODE_RECENT_APPS to R.drawable.ic_check_box_outline_blank,
         )
 
     val specialKeyLabels =
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/source/SystemShortcutsSource.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/source/SystemShortcutsSource.kt
index e55e339..7c0c75e 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/source/SystemShortcutsSource.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/source/SystemShortcutsSource.kt
@@ -18,14 +18,17 @@
 
 import android.content.res.Resources
 import android.view.KeyEvent.KEYCODE_A
+import android.view.KeyEvent.KEYCODE_BACK
 import android.view.KeyEvent.KEYCODE_DEL
 import android.view.KeyEvent.KEYCODE_DPAD_LEFT
 import android.view.KeyEvent.KEYCODE_ENTER
 import android.view.KeyEvent.KEYCODE_ESCAPE
 import android.view.KeyEvent.KEYCODE_H
+import android.view.KeyEvent.KEYCODE_HOME
 import android.view.KeyEvent.KEYCODE_I
 import android.view.KeyEvent.KEYCODE_L
 import android.view.KeyEvent.KEYCODE_N
+import android.view.KeyEvent.KEYCODE_RECENT_APPS
 import android.view.KeyEvent.KEYCODE_S
 import android.view.KeyEvent.KEYCODE_SLASH
 import android.view.KeyEvent.KEYCODE_TAB
@@ -60,24 +63,36 @@
                 command(META_META_ON)
             },
             // Access home screen:
+            //  - Home button
             //  - Meta + H
             //  - Meta + Enter
             shortcutInfo(resources.getString(R.string.group_system_access_home_screen)) {
+                command(modifiers = 0, KEYCODE_HOME)
+            },
+            shortcutInfo(resources.getString(R.string.group_system_access_home_screen)) {
                 command(META_META_ON, KEYCODE_H)
             },
             shortcutInfo(resources.getString(R.string.group_system_access_home_screen)) {
                 command(META_META_ON, KEYCODE_ENTER)
             },
             // Overview of open apps:
+            //  - Recent apps button
             //  - Meta + Tab
             shortcutInfo(resources.getString(R.string.group_system_overview_open_apps)) {
+                command(modifiers = 0, KEYCODE_RECENT_APPS)
+            },
+            shortcutInfo(resources.getString(R.string.group_system_overview_open_apps)) {
                 command(META_META_ON, KEYCODE_TAB)
             },
             // Back: go back to previous state (back button)
+            //  - Back button
             //  - Meta + Escape OR
             //  - Meta + Backspace OR
             //  - Meta + Left arrow
             shortcutInfo(resources.getString(R.string.group_system_go_back)) {
+                command(modifiers = 0, KEYCODE_BACK)
+            },
+            shortcutInfo(resources.getString(R.string.group_system_go_back)) {
                 command(META_META_ON, KEYCODE_ESCAPE)
             },
             shortcutInfo(resources.getString(R.string.group_system_go_back)) {
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
index 14837f2..6e883c2 100644
--- 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
@@ -16,12 +16,28 @@
 
 package com.android.systemui.keyboard.shortcut.data.repository
 
+import android.hardware.input.fakeInputManager
+import android.view.KeyEvent.KEYCODE_A
+import android.view.KeyEvent.KEYCODE_B
+import android.view.KeyEvent.KEYCODE_C
+import android.view.KeyEvent.KEYCODE_D
+import android.view.KeyEvent.KEYCODE_E
+import android.view.KeyEvent.KEYCODE_F
+import android.view.KeyEvent.KEYCODE_G
+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.data.source.FakeKeyboardShortcutGroupsSource
 import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts
+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.ShortcutKey
+import com.android.systemui.keyboard.shortcut.shared.model.ShortcutSubCategory
 import com.android.systemui.keyboard.shortcut.shortcutHelperAppCategoriesShortcutsSource
 import com.android.systemui.keyboard.shortcut.shortcutHelperCategoriesRepository
 import com.android.systemui.keyboard.shortcut.shortcutHelperCurrentAppShortcutsSource
@@ -47,13 +63,14 @@
 
     private val fakeSystemSource = FakeKeyboardShortcutGroupsSource()
     private val fakeMultiTaskingSource = FakeKeyboardShortcutGroupsSource()
+    private val fakeAppCategoriesSource = FakeKeyboardShortcutGroupsSource()
 
     private val kosmos =
         testKosmos().also {
             it.testDispatcher = UnconfinedTestDispatcher()
             it.shortcutHelperSystemShortcutsSource = fakeSystemSource
             it.shortcutHelperMultiTaskingShortcutsSource = fakeMultiTaskingSource
-            it.shortcutHelperAppCategoriesShortcutsSource = FakeKeyboardShortcutGroupsSource()
+            it.shortcutHelperAppCategoriesShortcutsSource = fakeAppCategoriesSource
             it.shortcutHelperInputShortcutsSource = FakeKeyboardShortcutGroupsSource()
             it.shortcutHelperCurrentAppShortcutsSource = FakeKeyboardShortcutGroupsSource()
         }
@@ -61,6 +78,7 @@
     private val repo = kosmos.shortcutHelperCategoriesRepository
     private val helper = kosmos.shortcutHelperTestHelper
     private val testScope = kosmos.testScope
+    private val fakeInputManager = kosmos.fakeInputManager
 
     @Before
     fun setUp() {
@@ -87,4 +105,74 @@
             // though fetching shortcuts again would have returned a new result.
             assertThat(secondCategories).isEqualTo(firstCategories)
         }
+
+    @Test
+    fun categories_filtersShortcutsWithUnsupportedKeyCodes() =
+        testScope.runTest {
+            fakeSystemSource.setGroups(
+                listOf(
+                    simpleGroup(
+                        simpleShortcutInfo(KEYCODE_A),
+                        simpleShortcutInfo(KEYCODE_B),
+                    ),
+                    simpleGroup(
+                        simpleShortcutInfo(KEYCODE_C),
+                    ),
+                )
+            )
+            fakeMultiTaskingSource.setGroups(
+                listOf(
+                    simpleGroup(
+                        simpleShortcutInfo(KEYCODE_D),
+                    ),
+                    simpleGroup(
+                        simpleShortcutInfo(KEYCODE_E),
+                        simpleShortcutInfo(KEYCODE_F),
+                    ),
+                )
+            )
+            fakeAppCategoriesSource.setGroups(listOf(simpleGroup(simpleShortcutInfo(KEYCODE_G))))
+
+            fakeInputManager.removeKeysFromKeyboard(deviceId = 123, KEYCODE_A, KEYCODE_D, KEYCODE_G)
+            helper.toggle(deviceId = 123)
+
+            val categories by collectLastValue(repo.categories)
+            assertThat(categories)
+                .containsExactly(
+                    ShortcutCategory(
+                        ShortcutCategoryType.System,
+                        listOf(
+                            simpleSubCategory(simpleShortcut("B")),
+                            simpleSubCategory(simpleShortcut("C")),
+                        )
+                    ),
+                    ShortcutCategory(
+                        ShortcutCategoryType.MultiTasking,
+                        listOf(
+                            simpleSubCategory(
+                                simpleShortcut("E"),
+                                simpleShortcut("F"),
+                            ),
+                        )
+                    ),
+                )
+        }
+
+    private fun simpleSubCategory(vararg shortcuts: Shortcut) =
+        ShortcutSubCategory(simpleGroupLabel, shortcuts.asList())
+
+    private fun simpleShortcut(vararg keys: String) =
+        Shortcut(
+            label = simpleShortcutLabel,
+            commands = listOf(ShortcutCommand(keys.map { ShortcutKey.Text(it) }))
+        )
+
+    private fun simpleGroup(vararg shortcuts: KeyboardShortcutInfo) =
+        KeyboardShortcutGroup(simpleGroupLabel, shortcuts.asList())
+
+    private fun simpleShortcutInfo(keyCode: Int = 0) =
+        KeyboardShortcutInfo(simpleShortcutLabel, keyCode, /* modifiers= */ 0)
+
+    private val simpleShortcutLabel = "shortcut label"
+    private val simpleGroupLabel = "group label"
 }
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 c4f93d1..6e7c05c 100644
--- a/packages/SystemUI/tests/utils/src/android/hardware/input/FakeInputManager.kt
+++ b/packages/SystemUI/tests/utils/src/android/hardware/input/FakeInputManager.kt
@@ -19,6 +19,8 @@
 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
@@ -38,6 +40,12 @@
             .build()
 
     private val devices = mutableMapOf<Int, InputDevice>(VIRTUAL_KEYBOARD to virtualKeyboard)
+    private val allKeyCodes = (0..KeyEvent.MAX_KEYCODE)
+    private val supportedKeyCodesByDeviceId =
+        mutableMapOf(
+            // Mark all keys supported by default
+            VIRTUAL_KEYBOARD to allKeyCodes.toMutableSet()
+        )
 
     val inputManager =
         mock<InputManager> {
@@ -61,13 +69,31 @@
             whenever(enableInputDevice(anyInt())).thenAnswer { invocation ->
                 setDeviceEnabled(invocation, enabled = true)
             }
+            whenever(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 addPhysicalKeyboardIfNotPresent(deviceId: Int, enabled: Boolean = true) {
+        if (devices.containsKey(deviceId)) {
+            return
+        }
+        addPhysicalKeyboard(deviceId, enabled)
+    }
+
     fun addPhysicalKeyboard(id: Int, enabled: Boolean = true) {
         check(id > 0) { "Physical keyboard ids have to be > 0" }
         addKeyboard(id, enabled)
     }
 
+    fun removeKeysFromKeyboard(deviceId: Int, vararg keyCodes: Int) {
+        addPhysicalKeyboardIfNotPresent(deviceId)
+        supportedKeyCodesByDeviceId[deviceId]!!.removeAll(keyCodes.asList())
+    }
+
     private fun addKeyboard(id: Int, enabled: Boolean = true) {
         devices[id] =
             InputDevice.Builder()
@@ -77,6 +103,7 @@
                 .setEnabled(enabled)
                 .setKeyCharacterMap(keyCharacterMap)
                 .build()
+        supportedKeyCodesByDeviceId[id] = allKeyCodes.toMutableSet()
     }
 
     private fun InputDevice.copy(
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 6ca5cd8..8b45662 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
@@ -95,7 +95,7 @@
     }
 
     fun toggle(deviceId: Int) {
-        fakeInputManager.addPhysicalKeyboard(deviceId)
+        fakeInputManager.addPhysicalKeyboardIfNotPresent(deviceId)
         fakeCommandQueue.doForEachCallback { it.toggleKeyboardShortcutsMenu(deviceId) }
     }