From 92e34b22732be0adfef8407267e89cdd1208c19f Mon Sep 17 00:00:00 2001
From: Lucas Silva <lusilva@google.com>
Date: Mon, 10 Feb 2025 11:17:37 -0500
Subject: 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
---
 .../systemui/communal/ui/compose/ContentListState.kt  |  5 +++--
 .../view/viewmodel/CommunalEditModeViewModelTest.kt   | 19 +++++++++++++++++++
 .../communal/ui/viewmodel/BaseCommunalViewModel.kt    |  2 +-
 .../ui/viewmodel/CommunalEditModeViewModel.kt         |  5 ++++-
 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,
             )
@@ -213,6 +214,24 @@ class CommunalEditModeViewModelTest : SysuiTestCase() {
             assertThat(appWidgetId).isEqualTo(1)
         }
 
+    @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 {
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)
     }
-- 
cgit v1.2.3-59-g8ed1b