Clear selected key when widget removed

If a widget is removed, clear the selected key if the widget being
removed was selected.

Fixes: 395072939
Test: atest CommunalEditModeViewModelTest
Flag: EXEMPT bugfix
Change-Id: I2d626a441fed098118e24580493250bb7bb65283

4 files changed, 27 insertions, 4 deletions

diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/ContentListState.kt b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/ContentListState.kt
index 16002bc709fd..8ad96a5bcb37 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/ContentListState.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/ContentListState.kt
@@ -57,7 +57,8 @@ class ContentListState
 internal constructor(
     communalContent: List<CommunalContentModel>,
     private val onAddWidget: (componentName: ComponentName, user: UserHandle, rank: Int) -> Unit,
-    private val onDeleteWidget: (id: Int, componentName: ComponentName, rank: Int) -> Unit,
+    private val onDeleteWidget:
+        (id: Int, key: String, componentName: ComponentName, rank: Int) -> Unit,
     private val onReorderWidgets: (widgetIdToRankMap: Map<Int, Int>) -> Unit,
     private val onResizeWidget:
         (
@@ -81,7 +82,7 @@ internal constructor(
         if (list[indexToRemove].isWidgetContent()) {
             val widget = list[indexToRemove] as CommunalContentModel.WidgetContent
             list.apply { removeAt(indexToRemove) }
-            onDeleteWidget(widget.appWidgetId, widget.componentName, widget.rank)
+            onDeleteWidget(widget.appWidgetId, widget.key, widget.componentName, widget.rank)
         }
     }
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt
index 522650bcde3c..f5f5dd8f81f9 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt
@@ -201,6 +201,7 @@ class CommunalEditModeViewModelTest : SysuiTestCase() {
 
             underTest.onDeleteWidget(
                 id = 0,
+                key = "key_0",
                 componentName = ComponentName("test_package", "test_class"),
                 rank = 30,
             )
@@ -214,6 +215,24 @@ class CommunalEditModeViewModelTest : SysuiTestCase() {
         }
 
     @Test
+    fun deleteWidget_clearsSelectedKey() =
+        kosmos.runTest {
+            val selectedKey by collectLastValue(underTest.selectedKey)
+            underTest.setSelectedKey("test_key")
+            assertThat(selectedKey).isEqualTo("test_key")
+
+            // Selected key is deleted.
+            underTest.onDeleteWidget(
+                id = 0,
+                key = "test_key",
+                componentName = ComponentName("test_package", "test_class"),
+                rank = 30,
+            )
+
+            assertThat(selectedKey).isNull()
+        }
+
+    @Test
     fun reorderWidget_uiEventLogging_start() =
         kosmos.runTest {
             underTest.onReorderWidgetStart(CommunalContentModel.KEY.widget(123))
diff --git a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/BaseCommunalViewModel.kt b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/BaseCommunalViewModel.kt
index 49003a735fbd..f4b532155ee7 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/BaseCommunalViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/BaseCommunalViewModel.kt
@@ -153,7 +153,7 @@ abstract class BaseCommunalViewModel(
     ) {}
 
     /** Called as the UI requests deleting a widget. */
-    open fun onDeleteWidget(id: Int, componentName: ComponentName, rank: Int) {}
+    open fun onDeleteWidget(id: Int, key: String, componentName: ComponentName, rank: Int) {}
 
     /** Called as the UI detects a tap event on the widget. */
     open fun onTapWidget(componentName: ComponentName, rank: Int) {}
diff --git a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt
index 8aba11190623..59beb1ed6c10 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt
@@ -141,7 +141,10 @@ constructor(
         metricsLogger.logAddWidget(componentName.flattenToString(), rank)
     }
 
-    override fun onDeleteWidget(id: Int, componentName: ComponentName, rank: Int) {
+    override fun onDeleteWidget(id: Int, key: String, componentName: ComponentName, rank: Int) {
+        if (selectedKey.value == key) {
+            setSelectedKey(null)
+        }
         communalInteractor.deleteWidget(id)
         metricsLogger.logRemoveWidget(componentName.flattenToString(), rank)
     }