Keep selection of widget at the end of widget dragging

Previously, the Glanceable Hub clears widget selection as soon as a
widget dragging starts. This change makes it so that a drag action also
selects that widget being dragged, and does not reset it after the
dragging completes.

Test: atest CommunalEditModeViewModelTest
Test: verified selection remains at the end of a widget dragging;
      see bug for before and after video
Fix: 376492600
Flag: com.android.systemui.communal_hub
Change-Id: I51970469cb73c77e909142b154cbbc67f4ae4dad

4 files changed, 14 insertions, 13 deletions

diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/GridDragDropState.kt b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/GridDragDropState.kt
index ef5e90bd7aad..7a500805809d 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/GridDragDropState.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/GridDragDropState.kt
@@ -94,7 +94,7 @@ internal constructor(
     private val scope: CoroutineScope,
     private val updateDragPositionForRemove: (draggingBoundingBox: IntRect) -> Boolean,
 ) {
-    var draggingItemKey by mutableStateOf<Any?>(null)
+    var draggingItemKey by mutableStateOf<String?>(null)
         private set
 
     var isDraggingToRemove by mutableStateOf(false)
@@ -138,7 +138,7 @@ internal constructor(
             // before content padding from the initial pointer position
             .firstItemAtOffset(normalizedOffset - contentOffset)
             ?.apply {
-                draggingItemKey = key
+                draggingItemKey = key as String
                 draggingItemInitialOffset = this.offset.toOffset()
                 return true
             }
@@ -284,7 +284,9 @@ fun Modifier.dragContainer(
                             contentOffset,
                         )
                     ) {
-                        viewModel.onReorderWidgetStart()
+                        // draggingItemKey is guaranteed to be non-null here because it is set in
+                        // onDragStart()
+                        viewModel.onReorderWidgetStart(dragDropState.draggingItemKey!!)
                     }
                 },
                 onDragEnd = {
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 5bbd3ffc625a..18cc8bf5f0d3 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
@@ -172,16 +172,16 @@ class CommunalEditModeViewModelTest : SysuiTestCase() {
         }
 
     @Test
-    fun selectedKey_onReorderWidgets_isCleared() =
+    fun selectedKey_onReorderWidgets_isSet() =
         testScope.runTest {
             val selectedKey by collectLastValue(underTest.selectedKey)
 
+            underTest.setSelectedKey(null)
+            assertThat(selectedKey).isNull()
+
             val key = CommunalContentModel.KEY.widget(123)
-            underTest.setSelectedKey(key)
+            underTest.onReorderWidgetStart(key)
             assertThat(selectedKey).isEqualTo(key)
-
-            underTest.onReorderWidgetStart()
-            assertThat(selectedKey).isNull()
         }
 
     @Test
@@ -234,7 +234,7 @@ class CommunalEditModeViewModelTest : SysuiTestCase() {
 
     @Test
     fun reorderWidget_uiEventLogging_start() {
-        underTest.onReorderWidgetStart()
+        underTest.onReorderWidgetStart(CommunalContentModel.KEY.widget(123))
         verify(uiEventLogger).log(CommunalUiEvent.COMMUNAL_HUB_REORDER_WIDGET_START)
     }
 
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 5ecf2e6b2551..a339af3694e7 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
@@ -195,7 +195,7 @@ abstract class BaseCommunalViewModel(
     open fun onDismissCtaTile() {}
 
     /** Called as the user starts dragging a widget to reorder. */
-    open fun onReorderWidgetStart() {}
+    open fun onReorderWidgetStart(draggingItemKey: String) {}
 
     /** Called as the user finishes dragging a widget to reorder. */
     open fun onReorderWidgetEnd() {}
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 736ed5c7d336..52bf0004cbe4 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
@@ -164,9 +164,8 @@ constructor(
         )
     }
 
-    override fun onReorderWidgetStart() {
-        // Clear selection status
-        setSelectedKey(null)
+    override fun onReorderWidgetStart(draggingItemKey: String) {
+        setSelectedKey(draggingItemKey)
         _reorderingWidgets.value = true
         uiEventLogger.log(CommunalUiEvent.COMMUNAL_HUB_REORDER_WIDGET_START)
     }