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

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> {

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
@@ -155,6 +156,19 @@ 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
@@ -162,6 +176,7 @@ class AllModelTest : SysuiTestCase() {
 
         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))
         }
@@ -181,6 +196,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