Merge "Update model when adding favorites" into rvc-dev

2 files changed, 35 insertions, 0 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/controls/management/AllModel.kt b/packages/SystemUI/src/com/android/systemui/controls/management/AllModel.kt
index 3fd583fa7481..11181e56838e 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/management/AllModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/management/AllModel.kt
@@ -65,11 +65,18 @@ class AllModel(
     override val elements: List<ElementWrapper> = createWrappers(controls)
 
     override fun changeFavoriteStatus(controlId: String, favorite: Boolean) {
+        val toChange = elements.firstOrNull {
+            it is ControlWrapper && it.controlStatus.control.controlId == controlId
+        } as ControlWrapper?
+        if (favorite == toChange?.controlStatus?.favorite) return
         if (favorite) {
             favoriteIds.add(controlId)
         } else {
             favoriteIds.remove(controlId)
         }
+        toChange?.let {
+            it.controlStatus.favorite = favorite
+        }
     }
 
     private fun createWrappers(list: List<ControlStatus>): List<ElementWrapper> {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/controls/management/AllModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/controls/management/AllModelTest.kt
index 9adab5d2369f..5e0d28f6f795 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/controls/management/AllModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/controls/management/AllModelTest.kt
@@ -25,6 +25,7 @@ import com.android.systemui.SysuiTestCase
 import com.android.systemui.controls.ControlStatus
 import com.android.systemui.controls.controller.ControlInfo
 import org.junit.Assert.assertEquals
+import org.junit.Assert.assertFalse
 import org.junit.Assert.assertTrue
 import org.junit.Before
 import org.junit.Test
@@ -156,12 +157,26 @@ class AllModelTest : SysuiTestCase() {
     }
 
     @Test
+    fun testAddFavorite_changesModelFlag() {
+        val indexToAdd = 6
+        val id = "$idPrefix$indexToAdd"
+        model.changeFavoriteStatus(id, true)
+        assertTrue(
+                (model.elements.first {
+                    it is ControlWrapper && it.controlStatus.control.controlId == id
+                } as ControlWrapper)
+                        .controlStatus.favorite
+        )
+    }
+
+    @Test
     fun testAddFavorite_alreadyThere() {
         val indexToAdd = 7
         model.changeFavoriteStatus("$idPrefix$indexToAdd", true)
 
         val expectedFavorites = favoritesIndices.map(controls::get).map(ControlStatus::control)
 
+        assertEquals(expectedFavorites.size, model.favorites.size)
         model.favorites.zip(expectedFavorites).forEach {
             assertTrue(sameControl(it.first, it.second))
         }
@@ -182,6 +197,19 @@ class AllModelTest : SysuiTestCase() {
     }
 
     @Test
+    fun testRemoveFavorite_changesModelFlag() {
+        val indexToRemove = 3
+        val id = "$idPrefix$indexToRemove"
+        model.changeFavoriteStatus(id, false)
+        assertFalse(
+                (model.elements.first {
+                    it is ControlWrapper && it.controlStatus.control.controlId == id
+                } as ControlWrapper)
+                        .controlStatus.favorite
+        )
+    }
+
+    @Test
     fun testRemoveFavorite_notThere() {
         val indexToRemove = 4
         model.changeFavoriteStatus("$idPrefix$indexToRemove", false)