Controls UI - Structure switching

Create a popupdialog similiar to a spinner widget. Move 'add controls'
into there as a permanent item. Support multiple structures per app,
but also default empty structures to just use the app name.

Bug: 148207527
Test: visual
Change-Id: I77671bd40859dfb749a90064b654a0bd14526622

15 files changed, 466 insertions, 199 deletions

diff --git a/packages/SystemUI/res/drawable/controls_list_divider.xml b/packages/SystemUI/res/drawable/controls_list_divider.xml
new file mode 100644
index 000000000000..f8211d5297f3
--- /dev/null
+++ b/packages/SystemUI/res/drawable/controls_list_divider.xml
@@ -0,0 +1,23 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2020 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+
+<shape xmlns:android="http://schemas.android.com/apk/res/android"
+       android:tint="@color/control_secondary_text">
+    <solid android:color="#33000000" />
+    <size
+        android:height="1dp"
+        android:width="1dp" />
+</shape>
diff --git a/packages/SystemUI/res/drawable/ic_more_vert.xml b/packages/SystemUI/res/drawable/ic_more_vert.xml
deleted file mode 100644
index 1309fa875b55..000000000000
--- a/packages/SystemUI/res/drawable/ic_more_vert.xml
+++ /dev/null
@@ -1,24 +0,0 @@
-<!--
-    Copyright (C) 2020 The Android Open Source Project
-
-    Licensed under the Apache License, Version 2.0 (the "License");
-    you may not use this file except in compliance with the License.
-    You may obtain a copy of the License at
-
-         http://www.apache.org/licenses/LICENSE-2.0
-
-    Unless required by applicable law or agreed to in writing, software
-    distributed under the License is distributed on an "AS IS" BASIS,
-    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-    See the License for the specific language governing permissions and
-    limitations under the License.
--->
-<vector xmlns:android="http://schemas.android.com/apk/res/android"
-    android:width="24dp"
-    android:height="24dp"
-    android:viewportWidth="24"
-    android:viewportHeight="24">
-  <path
-      android:fillColor="#FF000000"
-      android:pathData="M12,8c1.1,0 2,-0.9 2,-2s-0.9,-2 -2,-2 -2,0.9 -2,2 0.9,2 2,2zM12,10c-1.1,0 -2,0.9 -2,2s0.9,2 2,2 2,-0.9 2,-2 -0.9,-2 -2,-2zM12,16c-1.1,0 -2,0.9 -2,2s0.9,2 2,2 2,-0.9 2,-2 -0.9,-2 -2,-2z"/>
-</vector>
diff --git a/packages/SystemUI/res/layout/controls_spinner_item.xml b/packages/SystemUI/res/layout/controls_spinner_item.xml
new file mode 100644
index 000000000000..6b880545b296
--- /dev/null
+++ b/packages/SystemUI/res/layout/controls_spinner_item.xml
@@ -0,0 +1,50 @@
+<!--
+  ~ Copyright (C) 2019 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+<LinearLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:orientation="horizontal"
+    android:layout_width="match_parent"
+    android:layout_height="wrap_content"
+    android:paddingTop="12dp"
+    android:paddingBottom="12dp">
+
+    <Space
+        android:layout_weight="1"
+        android:layout_width="0dp"
+        android:layout_height="1dp" />
+    <ImageView
+        android:id="@+id/app_icon"
+        android:layout_gravity="center"
+        android:layout_width="34dp"
+        android:layout_height="24dp" 
+        android:layout_marginEnd="10dp" />
+
+    <TextView
+        android:id="@+id/controls_spinner_item"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:singleLine="true"
+        android:layout_gravity="center"
+        android:textSize="25sp"
+        android:textColor="@color/control_secondary_text"
+        android:fontFamily="@*android:string/config_headlineFontFamily" />
+
+    <Space
+        android:layout_weight="1"
+        android:layout_width="0dp"
+        android:layout_height="1dp" />
+</LinearLayout>
+    
diff --git a/packages/SystemUI/res/layout/controls_with_favorites.xml b/packages/SystemUI/res/layout/controls_with_favorites.xml
index 2cd9505b8fe4..77bcc3575fad 100644
--- a/packages/SystemUI/res/layout/controls_with_favorites.xml
+++ b/packages/SystemUI/res/layout/controls_with_favorites.xml
@@ -14,44 +14,47 @@
   ~ limitations under the License.
   -->
 <merge
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    xmlns:app="http://schemas.android.com/apk/res-auto">
+    xmlns:android="http://schemas.android.com/apk/res/android">
 
-  <androidx.constraintlayout.widget.ConstraintLayout
+  <LinearLayout
+      android:id="@+id/controls_header"
+      android:orientation="horizontal"
       android:layout_width="match_parent"
       android:layout_height="wrap_content"
-      android:paddingBottom="20dp">
+      android:paddingTop="12dp">
 
-    <TextView
-        android:text="@string/quick_controls_title"
-        android:layout_width="match_parent"
-        android:layout_height="wrap_content"
-        android:singleLine="true"
-        android:gravity="center"
-        android:textSize="25sp"
-        android:textColor="@*android:color/foreground_material_dark"
-        android:fontFamily="@*android:string/config_headlineFontFamily"
-        app:layout_constraintBottom_toBottomOf="parent"
-        app:layout_constraintEnd_toEndOf="parent"
-        app:layout_constraintStart_toStartOf="parent"
-        app:layout_constraintTop_toTopOf="parent" />
+    <Space
+        android:layout_weight="1"
+        android:layout_width="0dp"
+        android:layout_height="1dp" />
 
     <ImageView
-        android:id="@+id/controls_more"
-        android:src="@drawable/ic_more_vert"
-        android:layout_width="34dp"
+        android:id="@+id/app_icon"
+        android:layout_gravity="center"
+        android:layout_width="24dp"
         android:layout_height="24dp" 
-        android:layout_marginEnd="10dp"
-        android:tint="@*android:color/foreground_material_dark"
-        app:layout_constraintEnd_toEndOf="parent"
-        app:layout_constraintBottom_toBottomOf="parent"
-        app:layout_constraintTop_toTopOf="parent" />
+        android:layout_marginEnd="10dp" />
+
+    <TextView
+        style="@style/Control.Spinner.Header"
+        android:clickable="false"
+        android:id="@+id/app_or_structure_spinner"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:singleLine="true"
+        android:layout_gravity="center"
+        android:ellipsize="end" />
 
-  </androidx.constraintlayout.widget.ConstraintLayout>
+    <Space
+        android:layout_weight="1"
+        android:layout_width="0dp"
+        android:layout_height="1dp" />
+  </LinearLayout>
 
   <LinearLayout
       android:id="@+id/global_actions_controls_list"
       android:layout_width="match_parent"
       android:layout_height="wrap_content"
-      android:orientation="vertical" />
+      android:orientation="vertical"
+      android:paddingTop="20dp" />
 </merge>
diff --git a/packages/SystemUI/res/values/dimens.xml b/packages/SystemUI/res/values/dimens.xml
index aefe4a20a496..016b48bb48a1 100644
--- a/packages/SystemUI/res/values/dimens.xml
+++ b/packages/SystemUI/res/values/dimens.xml
@@ -1217,6 +1217,7 @@
 
     <!-- Home Controls -->
     <dimen name="control_spacing">4dp</dimen>
+    <dimen name="control_list_divider">1dp</dimen>
     <dimen name="control_corner_radius">15dp</dimen>
     <dimen name="control_height">100dp</dimen>
     <dimen name="control_padding">15dp</dimen>
diff --git a/packages/SystemUI/res/values/styles.xml b/packages/SystemUI/res/values/styles.xml
index d9b1452b6bb1..83c507095534 100644
--- a/packages/SystemUI/res/values/styles.xml
+++ b/packages/SystemUI/res/values/styles.xml
@@ -656,6 +656,12 @@
         <item name="android:fontFamily">@*android:string/config_bodyFontFamily</item>
     </style>
 
+    <style name="Control.Spinner.Header" parent="@*android:style/Widget.DeviceDefault.Spinner.DropDown">
+        <item name="android:textSize">25sp</item>
+        <item name="android:textColor">@color/control_primary_text</item>
+        <item name="android:fontFamily">@*android:string/config_headlineFontFamily</item>
+    </style>
+
     <style name="TextAppearance.Control.Status">
         <item name="android:textSize">12sp</item>
         <item name="android:textColor">@color/control_primary_text</item>
@@ -669,5 +675,8 @@
         <item name="android:textSize">12sp</item>
         <item name="android:textColor">@color/control_secondary_text</item>
     </style>
+    <style name="Control.ListPopupWindow" parent="@android:style/Widget.ListPopupWindow">
+        <item name="android:overlapAnchor">true</item>
+    </style>
     
 </resources>
diff --git a/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsBindingControllerImpl.kt b/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsBindingControllerImpl.kt
index 1d937ad75c99..8f02c252beb1 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsBindingControllerImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsBindingControllerImpl.kt
@@ -162,11 +162,9 @@ open class ControlsBindingControllerImpl @Inject constructor(
 
     override fun onComponentRemoved(componentName: ComponentName) {
         backgroundExecutor.execute {
-            synchronized(componentMap) {
-                val removed = componentMap.remove(Key(componentName, currentUser))
-                removed?.let {
-                    it.unbindService()
-                    tokenMap.remove(it.token)
+            currentProvider?.let {
+                if (it.componentName == componentName) {
+                    unbind()
                 }
             }
         }
@@ -182,16 +180,10 @@ open class ControlsBindingControllerImpl @Inject constructor(
 
     private abstract inner class CallbackRunnable(val token: IBinder) : Runnable {
         protected val provider: ControlsProviderLifecycleManager? = currentProvider
-    }
 
-    private inner class OnLoadRunnable(
-        token: IBinder,
-        val list: List<Control>,
-        val callback: ControlsBindingController.LoadCallback
-    ) : CallbackRunnable(token) {
         override fun run() {
             if (provider == null) {
-                Log.e(TAG, "No provider found for token:$token")
+                Log.e(TAG, "No current provider set")
                 return
             }
             if (provider.user != currentUser) {
@@ -202,8 +194,21 @@ open class ControlsBindingControllerImpl @Inject constructor(
                 Log.e(TAG, "Provider for token:$token does not exist anymore")
                 return
             }
+
+            doRun()
+        }
+
+        abstract fun doRun()
+    }
+
+    private inner class OnLoadRunnable(
+        token: IBinder,
+        val list: List<Control>,
+        val callback: ControlsBindingController.LoadCallback
+    ) : CallbackRunnable(token) {
+        override fun doRun() {
             callback.accept(list)
-            provider.unbindService()
+            provider?.unbindService()
         }
     }
 
@@ -211,14 +216,11 @@ open class ControlsBindingControllerImpl @Inject constructor(
         token: IBinder,
         val control: Control
     ) : CallbackRunnable(token) {
-        override fun run() {
+        override fun doRun() {
             if (!refreshing.get()) {
                 Log.d(TAG, "onRefresh outside of window from:${provider?.componentName}")
             }
-            if (provider?.user != currentUser) {
-                Log.e(TAG, "User ${provider?.user} is not current user")
-                return
-            }
+
             provider?.let {
                 lazyController.get().refreshStatus(it.componentName, control)
             }
@@ -229,7 +231,7 @@ open class ControlsBindingControllerImpl @Inject constructor(
         token: IBinder,
         val subscription: IControlsSubscription
     ) : CallbackRunnable(token) {
-        override fun run() {
+        override fun doRun() {
             if (!refreshing.get()) {
                 Log.d(TAG, "onRefresh outside of window from '${provider?.componentName}'")
             }
@@ -242,7 +244,7 @@ open class ControlsBindingControllerImpl @Inject constructor(
     private inner class OnCompleteRunnable(
         token: IBinder
     ) : CallbackRunnable(token) {
-        override fun run() {
+        override fun doRun() {
             provider?.let {
                 Log.i(TAG, "onComplete receive from '${it.componentName}'")
             }
@@ -253,7 +255,7 @@ open class ControlsBindingControllerImpl @Inject constructor(
         token: IBinder,
         val error: String
     ) : CallbackRunnable(token) {
-        override fun run() {
+        override fun doRun() {
             provider?.let {
                 Log.e(TAG, "onError receive from '${it.componentName}': $error")
             }
@@ -265,11 +267,7 @@ open class ControlsBindingControllerImpl @Inject constructor(
         val controlId: String,
         @ControlAction.ResponseResult val response: Int
     ) : CallbackRunnable(token) {
-        override fun run() {
-            if (provider?.user != currentUser) {
-                Log.e(TAG, "User ${provider?.user} is not current user")
-                return
-            }
+        override fun doRun() {
             provider?.let {
                 lazyController.get().onActionResponse(it.componentName, controlId, response)
             }
@@ -281,7 +279,7 @@ open class ControlsBindingControllerImpl @Inject constructor(
         val error: String,
         val callback: ControlsBindingController.LoadCallback
     ) : CallbackRunnable(token) {
-        override fun run() {
+        override fun doRun() {
             callback.error(error)
             provider?.let {
                 Log.e(TAG, "onError receive from '${it.componentName}': $error")
diff --git a/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsControllerImpl.kt b/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsControllerImpl.kt
index 22016a775028..dedd341a46be 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsControllerImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsControllerImpl.kt
@@ -145,21 +145,23 @@ class ControlsControllerImpl @Inject constructor (
      * If some component has been removed, the new set of favorites will also be saved.
      */
     private val listingCallback = object : ControlsListingController.ControlsListingCallback {
-        override fun onServicesUpdated(candidates: List<ControlsServiceInfo>) {
+        override fun onServicesUpdated(serviceInfos: List<ControlsServiceInfo>) {
             executor.execute {
-                val candidateComponents = candidates.map(ControlsServiceInfo::componentName)
-                synchronized(currentFavorites) {
-                    val components = currentFavorites.keys.toSet() // create a copy
-                    components.forEach {
-                        if (it !in candidateComponents) {
-                            currentFavorites.remove(it)
-                            bindingController.onComponentRemoved(it)
-                        }
-                    }
-                    // Check if something has been removed, if so, store the new list
-                    if (components.size > currentFavorites.size) {
-                        persistenceWrapper.storeFavorites(favoritesAsListLocked())
-                    }
+                val serviceInfoSet = serviceInfos.map(ControlsServiceInfo::componentName).toSet()
+                val favoriteComponentSet = Favorites.getAllStructures().map {
+                    it.componentName
+                }.toSet()
+
+                var changed = false
+                favoriteComponentSet.subtract(serviceInfoSet).forEach {
+                    changed = true
+                    Favorites.removeStructures(it)
+                    bindingController.onComponentRemoved(it)
+                }
+
+                // Check if something has been removed, if so, store the new list
+                if (changed) {
+                    persistenceWrapper.storeFavorites(Favorites.getAllStructures())
                 }
             }
         }
@@ -183,6 +185,7 @@ class ControlsControllerImpl @Inject constructor (
 
         if (shouldLoad) {
             Favorites.load(persistenceWrapper.readFavorites())
+            listingController.addCallback(listingCallback)
         }
     }
 
@@ -220,39 +223,42 @@ class ControlsControllerImpl @Inject constructor (
                 componentName,
                 object : ControlsBindingController.LoadCallback {
                     override fun accept(controls: List<Control>) {
-                        val favoritesForComponentKeys = Favorites
-                            .getControlsForComponent(componentName).map { it.controlId }
+                        executor.execute {
+                            val favoritesForComponentKeys = Favorites
+                                .getControlsForComponent(componentName).map { it.controlId }
+
+                            val changed = Favorites.updateControls(componentName, controls)
+                            if (changed) {
+                                persistenceWrapper.storeFavorites(Favorites.getAllStructures())
+                            }
+                            val removed = findRemoved(favoritesForComponentKeys.toSet(), controls)
+                            val controlsWithFavorite = controls.map {
+                                ControlStatus(it, it.controlId in favoritesForComponentKeys)
+                            }
+                            val loadData = createLoadDataObject(
+                                Favorites.getControlsForComponent(componentName)
+                                    .filter { it.controlId in removed }
+                                    .map { createRemovedStatus(componentName, it) } +
+                                controlsWithFavorite,
+                                favoritesForComponentKeys
+                            )
 
-                        val changed = Favorites.updateControls(componentName, controls)
-                        if (changed) {
-                            persistenceWrapper.storeFavorites(Favorites.getAllStructures())
-                        }
-                        val removed = findRemovedLocked(favoritesForComponentKeys.toSet(),
-                            controls)
-                        val controlsWithFavorite = controls.map {
-                            ControlStatus(it, it.controlId in favoritesForComponentKeys)
+                            dataCallback.accept(loadData)
                         }
-                        val loadData = createLoadDataObject(
-                            Favorites.getControlsForComponent(componentName)
-                                .filter { it.controlId in removed }
-                                .map { createRemovedStatus(componentName, it) } +
-                            controlsWithFavorite,
-                            favoritesForComponentKeys
-                        )
-
-                        dataCallback.accept(loadData)
                     }
 
                     override fun error(message: String) {
-                        Favorites.getControlsForComponent(componentName).let { controls ->
-                            val keys = controls.map { it.controlId }
-                            val loadData = createLoadDataObject(
-                                controls.map { createRemovedStatus(componentName, it, false) },
-                                keys,
-                                true
-                            )
-                            dataCallback.accept(loadData)
+                        val loadData = Favorites.getControlsForComponent(componentName).let {
+                            controls ->
+                                val keys = controls.map { it.controlId }
+                                createLoadDataObject(
+                                    controls.map { createRemovedStatus(componentName, it, false) },
+                                    keys,
+                                    true
+                                )
                         }
+
+                        dataCallback.accept(loadData)
                     }
                 }
         )
@@ -278,7 +284,7 @@ class ControlsControllerImpl @Inject constructor (
         return ControlStatus(control, true, setRemoved)
     }
 
-    private fun findRemovedLocked(favoriteKeys: Set<String>, list: List<Control>): Set<String> {
+    private fun findRemoved(favoriteKeys: Set<String>, list: List<Control>): Set<String> {
         val controlsKeys = list.map { it.controlId }
         return favoriteKeys.minus(controlsKeys)
     }
@@ -296,8 +302,10 @@ class ControlsControllerImpl @Inject constructor (
 
     override fun replaceFavoritesForStructure(structureInfo: StructureInfo) {
         if (!confirmAvailability()) return
-        Favorites.replaceControls(structureInfo)
-        persistenceWrapper.storeFavorites(Favorites.getAllStructures())
+        executor.execute {
+            Favorites.replaceControls(structureInfo)
+            persistenceWrapper.storeFavorites(Favorites.getAllStructures())
+        }
     }
 
     override fun refreshStatus(componentName: ComponentName, control: Control) {
@@ -358,6 +366,8 @@ class ControlsControllerImpl @Inject constructor (
 /**
  * Relies on immutable data for thread safety. When necessary to update favMap, use reassignment to
  * replace it, which will not disrupt any ongoing map traversal.
+ *
+ * Update/replace calls should use thread isolation to avoid race conditions.
  */
 private object Favorites {
     private var favMap = mapOf<ComponentName, List<StructureInfo>>()
@@ -416,11 +426,17 @@ private object Favorites {
 
         val newFavMap = favMap.toMutableMap()
         newFavMap.put(componentName, structures)
-        favMap = newFavMap.toMap()
+        favMap = newFavMap
 
         return true
     }
 
+    fun removeStructures(componentName: ComponentName) {
+        val newFavMap = favMap.toMutableMap()
+        newFavMap.remove(componentName)
+        favMap = newFavMap
+    }
+
     fun replaceControls(updatedStructure: StructureInfo) {
         val newFavMap = favMap.toMutableMap()
         val structures = mutableListOf<StructureInfo>()
diff --git a/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsFavoritePersistenceWrapper.kt b/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsFavoritePersistenceWrapper.kt
index fe3f4f8ef07d..4bea6ef3d7b9 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsFavoritePersistenceWrapper.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsFavoritePersistenceWrapper.kt
@@ -52,6 +52,10 @@ class ControlsFavoritePersistenceWrapper(
         private const val TAG_ID = "id"
         private const val TAG_TITLE = "title"
         private const val TAG_TYPE = "type"
+        private const val TAG_VERSION = "version"
+
+        // must increment with every change to the XML structure
+        private const val VERSION = 1
     }
 
     /**
@@ -83,6 +87,10 @@ class ControlsFavoritePersistenceWrapper(
                     setOutput(writer, "utf-8")
                     setFeature("http://xmlpull.org/v1/doc/features.html#indent-output", true)
                     startDocument(null, true)
+                    startTag(null, TAG_VERSION)
+                    text("$VERSION")
+                    endTag(null, TAG_VERSION)
+
                     startTag(null, TAG_STRUCTURES)
                     structures.forEach { s ->
                         startTag(null, TAG_STRUCTURE)
diff --git a/packages/SystemUI/src/com/android/systemui/controls/management/AppAdapter.kt b/packages/SystemUI/src/com/android/systemui/controls/management/AppAdapter.kt
index 855e4a501f0a..8b3454a2bc7c 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/management/AppAdapter.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/management/AppAdapter.kt
@@ -26,15 +26,13 @@ import android.widget.TextView
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
 import androidx.recyclerview.widget.RecyclerView
-import com.android.settingslib.applications.DefaultAppInfo
-import com.android.settingslib.widget.CandidateInfo
 import com.android.systemui.R
 import com.android.systemui.controls.ControlsServiceInfo
 import java.text.Collator
 import java.util.concurrent.Executor
 
 /**
- * Adapter for binding [CandidateInfo] related to [ControlsProviderService].
+ * Adapter for binding [ControlsServiceInfo] related to [ControlsProviderService].
  *
  * This class handles subscribing and keeping track of the list of valid applications for
  * displaying.
@@ -56,16 +54,16 @@ class AppAdapter(
     private val resources: Resources
 ) : RecyclerView.Adapter<AppAdapter.Holder>() {
 
-    private var listOfServices = emptyList<CandidateInfo>()
+    private var listOfServices = emptyList<ControlsServiceInfo>()
 
     private val callback = object : ControlsListingController.ControlsListingCallback {
-        override fun onServicesUpdated(candidates: List<ControlsServiceInfo>) {
+        override fun onServicesUpdated(serviceInfos: List<ControlsServiceInfo>) {
             backgroundExecutor.execute {
                 val collator = Collator.getInstance(resources.configuration.locales[0])
-                val localeComparator = compareBy<CandidateInfo, CharSequence>(collator) {
+                val localeComparator = compareBy<ControlsServiceInfo, CharSequence>(collator) {
                     it.loadLabel()
                 }
-                listOfServices = candidates.sortedWith(localeComparator)
+                listOfServices = serviceInfos.sortedWith(localeComparator)
                 uiExecutor.execute(::notifyDataSetChanged)
             }
         }
@@ -101,11 +99,10 @@ class AppAdapter(
          * Bind data to the view
          * @param data Information about the [ControlsProviderService] to bind to the data
          */
-        fun bindData(data: CandidateInfo) {
+        fun bindData(data: ControlsServiceInfo) {
             icon.setImageDrawable(data.loadIcon())
             title.text = data.loadLabel()
-            favorites.text = favRenderer.renderFavoritesForComponent(
-                    (data as DefaultAppInfo).componentName)
+            favorites.text = favRenderer.renderFavoritesForComponent(data.componentName)
         }
     }
 }
@@ -123,4 +120,4 @@ class FavoritesRenderer(
             return ""
         }
     }
-}
\ No newline at end of file
+}
diff --git a/packages/SystemUI/src/com/android/systemui/controls/management/ControlsListingController.kt b/packages/SystemUI/src/com/android/systemui/controls/management/ControlsListingController.kt
index efbe7c0bdb38..647daccca8bd 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/management/ControlsListingController.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/management/ControlsListingController.kt
@@ -45,6 +45,6 @@ interface ControlsListingController :
 
     @FunctionalInterface
     interface ControlsListingCallback {
-        fun onServicesUpdated(candidates: List<ControlsServiceInfo>)
+        fun onServicesUpdated(serviceInfos: List<ControlsServiceInfo>)
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/controls/management/ControlsListingControllerImpl.kt b/packages/SystemUI/src/com/android/systemui/controls/management/ControlsListingControllerImpl.kt
index 8db694eef064..9b108cf3e34b 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/management/ControlsListingControllerImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/management/ControlsListingControllerImpl.kt
@@ -24,7 +24,6 @@ import android.os.UserHandle
 import android.service.controls.ControlsProviderService
 import android.util.Log
 import com.android.internal.annotations.VisibleForTesting
-import com.android.settingslib.applications.DefaultAppInfo
 import com.android.settingslib.applications.ServiceListing
 import com.android.settingslib.widget.CandidateInfo
 import com.android.systemui.controls.ControlsServiceInfo
@@ -157,7 +156,7 @@ class ControlsListingControllerImpl @VisibleForTesting constructor(
      * @return a label as returned by [CandidateInfo.loadLabel] or `null`.
      */
     override fun getAppLabel(name: ComponentName): CharSequence? {
-        return getCurrentServices().firstOrNull { (it as? DefaultAppInfo)?.componentName == name }
+        return getCurrentServices().firstOrNull { it.componentName == name }
                 ?.loadLabel()
     }
-}
\ No newline at end of file
+}
diff --git a/packages/SystemUI/src/com/android/systemui/controls/ui/ControlsUiControllerImpl.kt b/packages/SystemUI/src/com/android/systemui/controls/ui/ControlsUiControllerImpl.kt
index fd95b9c80fdf..71fc01717baa 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/ui/ControlsUiControllerImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/ui/ControlsUiControllerImpl.kt
@@ -22,19 +22,25 @@ import android.content.ComponentName
 import android.content.Context
 import android.content.Intent
 import android.content.ServiceConnection
+import android.content.SharedPreferences
 import android.graphics.drawable.Drawable
 import android.os.IBinder
 import android.service.controls.Control
 import android.service.controls.TokenProvider
 import android.util.Log
+import android.view.ContextThemeWrapper
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
+import android.view.WindowManager
+import android.widget.AdapterView
+import android.widget.ArrayAdapter
 import android.widget.ImageView
 import android.widget.LinearLayout
+import android.widget.ListPopupWindow
 import android.widget.Space
-
-import com.android.settingslib.widget.CandidateInfo
+import android.widget.TextView
+import com.android.systemui.controls.ControlsServiceInfo
 import com.android.systemui.controls.controller.ControlInfo
 import com.android.systemui.controls.controller.ControlsController
 import com.android.systemui.controls.controller.StructureInfo
@@ -43,7 +49,6 @@ import com.android.systemui.controls.management.ControlsProviderSelectorActivity
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.R
-import com.android.systemui.controls.ControlsServiceInfo
 import com.android.systemui.util.concurrency.DelayableExecutor
 
 import dagger.Lazy
@@ -123,12 +128,17 @@ class ControlsUiControllerImpl @Inject constructor (
     val context: Context,
     @Main val uiExecutor: DelayableExecutor,
     @Background val bgExecutor: DelayableExecutor,
-    val controlsListingController: Lazy<ControlsListingController>
+    val controlsListingController: Lazy<ControlsListingController>,
+    @Main val sharedPreferences: SharedPreferences
 ) : ControlsUiController {
 
     companion object {
+        private const val PREF_COMPONENT = "controls_component"
+        private const val PREF_STRUCTURE = "controls_structure"
+
+        private val EMPTY_COMPONENT = ComponentName("", "")
         private val EMPTY_STRUCTURE = StructureInfo(
-            ComponentName("", ""),
+            EMPTY_COMPONENT,
             "",
             mutableListOf<ControlInfo>()
         )
@@ -139,45 +149,66 @@ class ControlsUiControllerImpl @Inject constructor (
     private val controlsById = mutableMapOf<ControlKey, ControlWithState>()
     private val controlViewsById = mutableMapOf<ControlKey, ControlViewHolder>()
     private lateinit var parent: ViewGroup
+    private lateinit var lastItems: List<SelectionItem>
+    private var popup: ListPopupWindow? = null
+
+    private val addControlsItem = SelectionItem(
+        context.resources.getString(R.string.controls_providers_title),
+        "",
+        context.getDrawable(R.drawable.ic_add),
+        EMPTY_COMPONENT
+    )
 
     override val available: Boolean
         get() = controlsController.get().available
 
-    private val listingCallback = object : ControlsListingController.ControlsListingCallback {
-        override fun onServicesUpdated(candidates: List<ControlsServiceInfo>) {
-            bgExecutor.execute {
-                val collator = Collator.getInstance(context.resources.configuration.locales[0])
-                val localeComparator = compareBy<CandidateInfo, CharSequence>(collator) {
-                    it.loadLabel()
+    private lateinit var listingCallback: ControlsListingController.ControlsListingCallback
+
+    private fun createCallback(
+        onResult: (List<SelectionItem>) -> Unit
+    ): ControlsListingController.ControlsListingCallback {
+        return object : ControlsListingController.ControlsListingCallback {
+            override fun onServicesUpdated(serviceInfos: List<ControlsServiceInfo>) {
+                bgExecutor.execute {
+                    val collator = Collator.getInstance(context.resources.configuration.locales[0])
+                    val localeComparator = compareBy<ControlsServiceInfo, CharSequence>(collator) {
+                        it.loadLabel()
+                    }
+
+                    val mList = serviceInfos.toMutableList()
+                    mList.sortWith(localeComparator)
+                    lastItems = mList.map {
+                        SelectionItem(it.loadLabel(), "", it.loadIcon(), it.componentName)
+                    }
+                    uiExecutor.execute {
+                        onResult(lastItems)
+                    }
                 }
-
-                val mList = candidates.toMutableList()
-                mList.sortWith(localeComparator)
-                loadInitialSetupViewIcons(mList.map { it.loadLabel() to it.loadIcon() })
             }
         }
     }
 
     override fun show(parent: ViewGroup) {
         Log.d(ControlsUiController.TAG, "show()")
-
         this.parent = parent
 
         allStructures = controlsController.get().getFavorites()
+        selectedStructure = loadPreference(allStructures)
 
-        if (allStructures.isEmpty()) {
-            showInitialSetupView()
+        if (selectedStructure.controls.isEmpty() && allStructures.size <= 1) {
+            // only show initial view if there are really no favorites across any structure
+            listingCallback = createCallback(::showInitialSetupView)
         } else {
-            selectedStructure = allStructures.get(0)
             selectedStructure.controls.map {
                 ControlWithState(selectedStructure.componentName, it, null)
             }.associateByTo(controlsById) {
                 ControlKey(selectedStructure.componentName, it.ci.controlId)
             }
-
-            showControlsView()
+            listingCallback = createCallback(::showControlsView)
         }
 
+        controlsListingController.get().addCallback(listingCallback)
+
         // Temp code to pass auth
         tokenProviderConnection = TokenProviderConnection(controlsController.get(), context,
                 selectedStructure)
@@ -191,29 +222,21 @@ class ControlsUiControllerImpl @Inject constructor (
         }
     }
 
-    private fun showInitialSetupView() {
+    private fun showInitialSetupView(items: List<SelectionItem>) {
+        parent.removeAllViews()
+
         val inflater = LayoutInflater.from(context)
         inflater.inflate(R.layout.controls_no_favorites, parent, true)
 
         val viewGroup = parent.requireViewById(R.id.controls_no_favorites_group) as ViewGroup
         viewGroup.setOnClickListener(launchSelectorActivityListener(context))
 
-        controlsListingController.get().addCallback(listingCallback)
-    }
-
-    private fun loadInitialSetupViewIcons(icons: List<Pair<CharSequence, Drawable>>) {
-        uiExecutor.execute {
-            val viewGroup = parent.requireViewById(R.id.controls_icon_row) as ViewGroup
-            viewGroup.removeAllViews()
-
-            val inflater = LayoutInflater.from(context)
-            icons.forEach {
-                val imageView = inflater.inflate(R.layout.controls_icon, viewGroup, false)
-                        as ImageView
-                imageView.setContentDescription(it.first)
-                imageView.setImageDrawable(it.second)
-                viewGroup.addView(imageView)
-            }
+        val iconRowGroup = parent.requireViewById(R.id.controls_icon_row) as ViewGroup
+        items.forEach {
+            val imageView = inflater.inflate(R.layout.controls_icon, viewGroup, false) as ImageView
+            imageView.setContentDescription(it.getTitle())
+            imageView.setImageDrawable(it.icon)
+            iconRowGroup.addView(imageView)
         }
     }
 
@@ -229,14 +252,16 @@ class ControlsUiControllerImpl @Inject constructor (
         }
     }
 
-    private fun showControlsView() {
+    private fun showControlsView(items: List<SelectionItem>) {
+        parent.removeAllViews()
+        controlViewsById.clear()
+
         val inflater = LayoutInflater.from(context)
         inflater.inflate(R.layout.controls_with_favorites, parent, true)
 
         val listView = parent.requireViewById(R.id.global_actions_controls_list) as ViewGroup
         var lastRow: ViewGroup = createRow(inflater, listView)
         selectedStructure.controls.forEach {
-            Log.d(ControlsUiController.TAG, "favorited control id: " + it.controlId)
             if (lastRow.getChildCount() == 2) {
                 lastRow = createRow(inflater, listView)
             }
@@ -249,16 +274,109 @@ class ControlsUiControllerImpl @Inject constructor (
             controlViewsById.put(key, cvh)
         }
 
+        // add spacer if necessary to keep control size consistent
         if ((selectedStructure.controls.size % 2) == 1) {
             lastRow.addView(Space(context), LinearLayout.LayoutParams(0, 0, 1f))
         }
 
-        val moreImageView = parent.requireViewById(R.id.controls_more) as View
-        moreImageView.setOnClickListener(launchSelectorActivityListener(context))
+        val itemsByComponent = items.associateBy { it.componentName }
+        var adapter = ItemAdapter(context, R.layout.controls_spinner_item).apply {
+            val listItems = allStructures.mapNotNull {
+                itemsByComponent.get(it.componentName)?.copy(structure = it.structure)
+            }
+
+            addAll(listItems + addControlsItem)
+        }
+
+        /*
+         * Default spinner widget does not work with the window type required
+         * for this dialog. Use a textView with the ListPopupWindow to achieve
+         * a similar effect
+         */
+        parent.requireViewById<TextView>(R.id.app_or_structure_spinner).apply {
+            setText((adapter.findSelectionItem(selectedStructure) ?: adapter.getItem(0)).getTitle())
+        }
+        val anchor = parent.requireViewById<ViewGroup>(R.id.controls_header)
+        anchor.setOnClickListener(object : View.OnClickListener {
+            override fun onClick(v: View) {
+                popup = ListPopupWindow(
+                    ContextThemeWrapper(context, R.style.Control_ListPopupWindow))
+                popup?.apply {
+                    setWindowLayoutType(WindowManager.LayoutParams.TYPE_VOLUME_OVERLAY)
+                    setAnchorView(anchor)
+                    setAdapter(adapter)
+                    setModal(true)
+                    setOnItemClickListener(object : AdapterView.OnItemClickListener {
+                        override fun onItemClick(
+                            parent: AdapterView<*>,
+                            view: View,
+                            pos: Int,
+                            id: Long
+                        ) {
+                            val listItem = parent.getItemAtPosition(pos) as SelectionItem
+                            this@ControlsUiControllerImpl.switchAppOrStructure(listItem)
+                            dismiss()
+                        }
+                    })
+                    // need to call show() first in order to construct the listView
+                    show()
+                    getListView()?.apply {
+                        setDividerHeight(
+                            context.resources.getDimensionPixelSize(R.dimen.control_list_divider))
+                        setDivider(
+                            context.resources.getDrawable(R.drawable.controls_list_divider))
+                    }
+                    show()
+                }
+            }
+        })
+
+        parent.requireViewById<ImageView>(R.id.app_icon).apply {
+            setContentDescription("My Home")
+            setImageDrawable(items[0].icon)
+        }
+    }
+
+    private fun loadPreference(structures: List<StructureInfo>): StructureInfo {
+        if (structures.isEmpty()) return EMPTY_STRUCTURE
+
+        val component = sharedPreferences.getString(PREF_COMPONENT, null)?.let {
+            ComponentName.unflattenFromString(it)
+        } ?: EMPTY_COMPONENT
+        val structure = sharedPreferences.getString(PREF_STRUCTURE, "")
+
+        return structures.firstOrNull {
+            component == it.componentName && structure == it.structure
+        } ?: structures.get(0)
+    }
+
+    private fun updatePreferences(si: StructureInfo) {
+        sharedPreferences.edit()
+            .putString(PREF_COMPONENT, si.componentName.flattenToString())
+            .putString(PREF_STRUCTURE, si.structure.toString())
+            .commit()
+    }
+
+    private fun switchAppOrStructure(item: SelectionItem) {
+        if (item == addControlsItem) {
+            launchSelectorActivityListener(context)(parent)
+        } else {
+            val newSelection = allStructures.first {
+                it.structure == item.structure && it.componentName == item.componentName
+            }
+
+            if (newSelection != selectedStructure) {
+                selectedStructure = newSelection
+                updatePreferences(selectedStructure)
+                showControlsView(lastItems)
+            }
+        }
     }
 
     override fun hide() {
         Log.d(ControlsUiController.TAG, "hide()")
+        popup?.dismiss()
+
         controlsController.get().unsubscribe()
         context.unbindService(tokenProviderConnection)
         tokenProviderConnection = null
@@ -298,3 +416,46 @@ class ControlsUiControllerImpl @Inject constructor (
         return row
     }
 }
+
+private data class SelectionItem(
+    val appName: CharSequence,
+    val structure: CharSequence,
+    val icon: Drawable,
+    val componentName: ComponentName
+) {
+    fun getTitle() = if (structure.isEmpty()) { appName } else { structure }
+}
+
+private class ItemAdapter(
+    val parentContext: Context,
+    val resource: Int
+) : ArrayAdapter<SelectionItem>(parentContext, resource) {
+
+    val layoutInflater = LayoutInflater.from(context)
+
+    override fun getView(position: Int, convertView: View?, parent: ViewGroup): View {
+        val item = getItem(position)
+        val view = convertView ?: layoutInflater.inflate(resource, parent, false)
+        view.requireViewById<TextView>(R.id.controls_spinner_item).apply {
+            setText(item.getTitle())
+        }
+        view.requireViewById<ImageView>(R.id.app_icon).apply {
+            setContentDescription(item.getTitle())
+            setImageDrawable(item.icon)
+        }
+        return view
+    }
+
+    fun findSelectionItem(si: StructureInfo): SelectionItem? {
+        var i = 0
+        while (i < getCount()) {
+            val item = getItem(i)
+            if (item.componentName == si.componentName &&
+                item.structure == si.structure) {
+                return item
+            }
+            i++
+        }
+        return null
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ControlsBindingControllerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ControlsBindingControllerImplTest.kt
index ea58d4764beb..eceb1ddaaf06 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ControlsBindingControllerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ControlsBindingControllerImplTest.kt
@@ -39,6 +39,7 @@ import org.mockito.Mockito
 import org.mockito.Mockito.`when`
 import org.mockito.Mockito.mock
 import org.mockito.Mockito.never
+import org.mockito.Mockito.times
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
 
@@ -176,20 +177,13 @@ class ControlsBindingControllerImplTest : SysuiTestCase() {
 
     @Test
     fun testComponentRemoved_existingIsUnbound() {
-        controller.bindServices(listOf(
-                TEST_COMPONENT_NAME_1,
-                TEST_COMPONENT_NAME_2,
-                TEST_COMPONENT_NAME_3
-        ))
+        controller.bindService(TEST_COMPONENT_NAME_1)
 
-        controller.onComponentRemoved(TEST_COMPONENT_NAME_2)
+        controller.onComponentRemoved(TEST_COMPONENT_NAME_1)
 
         executor.runAllReady()
 
-        providers.forEach {
-            verify(it, if (it.componentName == TEST_COMPONENT_NAME_2) times(1) else never())
-                    .unbindService()
-        }
+        verify(providers[0], times(1)).unbindService()
     }
 }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ControlsControllerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ControlsControllerImplTest.kt
index d75131fe7de8..45ea3c96b819 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ControlsControllerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/controls/controller/ControlsControllerImplTest.kt
@@ -47,12 +47,14 @@ import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.ArgumentCaptor
+import org.mockito.ArgumentMatchers
 import org.mockito.Captor
 import org.mockito.Mock
 import org.mockito.Mockito
 import org.mockito.Mockito.`when`
 import org.mockito.Mockito.inOrder
 import org.mockito.Mockito.mock
+import org.mockito.Mockito.never
 import org.mockito.Mockito.reset
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
@@ -149,6 +151,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
         assertTrue(controller.available)
         verify(broadcastDispatcher).registerReceiver(
                 capture(broadcastReceiverCaptor), any(), any(), eq(UserHandle.ALL))
+
         verify(listingController).addCallback(capture(listingCallbackCaptor))
     }
 
@@ -210,6 +213,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
     fun testSubscribeFavorites() {
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO)
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO_2)
+        delayableExecutor.runAllReady()
 
         controller.subscribeToFavorites(TEST_STRUCTURE_INFO)
 
@@ -241,6 +245,8 @@ class ControlsControllerImplTest : SysuiTestCase() {
 
         controlLoadCallbackCaptor.value.accept(listOf(control))
 
+        delayableExecutor.runAllReady()
+
         assertTrue(loaded)
     }
 
@@ -251,6 +257,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
         val control2 = builderFromInfo(TEST_CONTROL_INFO_2).build()
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO)
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO_2)
+        delayableExecutor.runAllReady()
 
         controller.loadForComponent(TEST_COMPONENT, Consumer { data ->
             val controls = data.allControls
@@ -272,6 +279,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
                 capture(controlLoadCallbackCaptor))
 
         controlLoadCallbackCaptor.value.accept(listOf(control, control2))
+        delayableExecutor.runAllReady()
 
         assertTrue(loaded)
     }
@@ -280,6 +288,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
     fun testLoadForComponent_removed() {
         var loaded = false
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO)
+        delayableExecutor.runAllReady()
 
         controller.loadForComponent(TEST_COMPONENT, Consumer { data ->
             val controls = data.allControls
@@ -300,6 +309,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
                 capture(controlLoadCallbackCaptor))
 
         controlLoadCallbackCaptor.value.accept(emptyList())
+        delayableExecutor.runAllReady()
 
         assertTrue(loaded)
     }
@@ -308,6 +318,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
     fun testErrorOnLoad_notRemoved() {
         var loaded = false
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO)
+        delayableExecutor.runAllReady()
 
         controller.loadForComponent(TEST_COMPONENT, Consumer { data ->
             val controls = data.allControls
@@ -335,6 +346,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
     @Test
     fun testFavoriteInformationModifiedOnLoad() {
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO)
+        delayableExecutor.runAllReady()
 
         val newControlInfo = TEST_CONTROL_INFO.copy(controlTitle = TEST_CONTROL_TITLE_2)
         val control = builderFromInfo(newControlInfo).build()
@@ -345,6 +357,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
                 capture(controlLoadCallbackCaptor))
 
         controlLoadCallbackCaptor.value.accept(listOf(control))
+        delayableExecutor.runAllReady()
 
         val favorites = controller.getFavorites().flatMap { it.controls }
         assertEquals(1, favorites.size)
@@ -370,6 +383,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
     @Test
     fun testSwitchUsers() {
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO)
+        delayableExecutor.runAllReady()
 
         reset(persistenceWrapper)
         val intent = Intent(Intent.ACTION_USER_SWITCHED).apply {
@@ -401,6 +415,8 @@ class ControlsControllerImplTest : SysuiTestCase() {
     @Test
     fun testDisableFeature_clearFavorites() {
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO)
+        delayableExecutor.runAllReady()
+
         assertFalse(controller.getFavorites().isEmpty())
 
         Settings.Secure.putIntForUser(mContext.contentResolver,
@@ -412,6 +428,8 @@ class ControlsControllerImplTest : SysuiTestCase() {
     @Test
     fun testDisableFeature_noChangeForNotCurrentUser() {
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO)
+        delayableExecutor.runAllReady()
+
         Settings.Secure.putIntForUser(mContext.contentResolver,
                 ControlsControllerImpl.CONTROLS_AVAILABLE, 0, otherUser)
         controller.settingObserver.onChange(false, ControlsControllerImpl.URI, otherUser)
@@ -440,6 +458,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
     @Test
     fun testCountFavoritesForComponent_singleComponent() {
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO)
+        delayableExecutor.runAllReady()
 
         assertEquals(1, controller.countFavoritesForComponent(TEST_COMPONENT))
         assertEquals(0, controller.countFavoritesForComponent(TEST_COMPONENT_2))
@@ -449,6 +468,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
     fun testCountFavoritesForComponent_multipleComponents() {
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO)
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO_2)
+        delayableExecutor.runAllReady()
 
         assertEquals(1, controller.countFavoritesForComponent(TEST_COMPONENT))
         assertEquals(1, controller.countFavoritesForComponent(TEST_COMPONENT_2))
@@ -457,6 +477,8 @@ class ControlsControllerImplTest : SysuiTestCase() {
     @Test
     fun testGetFavoritesForComponent() {
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO)
+        delayableExecutor.runAllReady()
+
         assertEquals(listOf(TEST_STRUCTURE_INFO),
             controller.getFavoritesForComponent(TEST_COMPONENT))
     }
@@ -464,6 +486,8 @@ class ControlsControllerImplTest : SysuiTestCase() {
     @Test
     fun testGetFavoritesForComponent_otherComponent() {
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO_2)
+        delayableExecutor.runAllReady()
+
         assertTrue(controller.getFavoritesForComponent(TEST_COMPONENT).isEmpty())
     }
 
@@ -477,6 +501,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
                 "Home",
                 listOf(TEST_CONTROL_INFO, controlInfo)
         ))
+        delayableExecutor.runAllReady()
 
         assertEquals(listOf(TEST_CONTROL_INFO, controlInfo),
             controller.getFavoritesForComponent(TEST_COMPONENT).flatMap { it.controls })
@@ -487,6 +512,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
                 "Home",
                 listOf(controlInfo, TEST_CONTROL_INFO)
         ))
+        delayableExecutor.runAllReady()
 
         assertEquals(listOf(controlInfo, TEST_CONTROL_INFO),
             controller.getFavoritesForComponent(TEST_COMPONENT).flatMap { it.controls })
@@ -495,6 +521,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
     @Test
     fun testReplaceFavoritesForStructure_noFavorites() {
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO)
+        delayableExecutor.runAllReady()
 
         assertEquals(1, controller.countFavoritesForComponent(TEST_COMPONENT))
         assertEquals(listOf(TEST_STRUCTURE_INFO),
@@ -505,6 +532,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
     fun testReplaceFavoritesForStructure_differentComponentsAreFilteredOut() {
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO)
         controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO_2)
+        delayableExecutor.runAllReady()
 
         assertEquals(1, controller.countFavoritesForComponent(TEST_COMPONENT))
         assertEquals(listOf(TEST_CONTROL_INFO),
@@ -530,6 +558,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
                 "Home",
                 listOf(TEST_CONTROL_INFO)
         ))
+        delayableExecutor.runAllReady()
 
         assertEquals(1, controller.countFavoritesForComponent(newComponent))
         assertEquals(listOf(TEST_CONTROL_INFO), controller
@@ -543,6 +572,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
         val listOrder1 = listOf(TEST_CONTROL_INFO, controlInfo)
         val structure1 = StructureInfo(TEST_COMPONENT, "Home", listOrder1)
         controller.replaceFavoritesForStructure(structure1)
+        delayableExecutor.runAllReady()
 
         assertEquals(2, controller.countFavoritesForComponent(TEST_COMPONENT))
         assertEquals(listOrder1, controller.getFavoritesForComponent(TEST_COMPONENT)
@@ -552,6 +582,7 @@ class ControlsControllerImplTest : SysuiTestCase() {
         val structure2 = StructureInfo(TEST_COMPONENT, "Home", listOrder2)
 
         controller.replaceFavoritesForStructure(structure2)
+        delayableExecutor.runAllReady()
 
         assertEquals(2, controller.countFavoritesForComponent(TEST_COMPONENT))
         assertEquals(listOrder2, controller.getFavoritesForComponent(TEST_COMPONENT)
@@ -560,7 +591,8 @@ class ControlsControllerImplTest : SysuiTestCase() {
 
     @Test
     fun testPackageRemoved_noFavorites_noRemovals() {
-        controller.changeFavoriteStatus(TEST_CONTROL_INFO, true)
+        controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO)
+        delayableExecutor.runAllReady()
 
         val serviceInfo = mock(ServiceInfo::class.java)
         `when`(serviceInfo.componentName).thenReturn(TEST_COMPONENT)
@@ -569,21 +601,21 @@ class ControlsControllerImplTest : SysuiTestCase() {
         // Don't want to check what happens before this call
         reset(persistenceWrapper)
         listingCallbackCaptor.value.onServicesUpdated(listOf(info))
-
         delayableExecutor.runAllReady()
 
         verify(bindingController, never()).onComponentRemoved(any())
 
-        assertEquals(1, controller.getFavoriteControls().size)
-        assertEquals(TEST_CONTROL_INFO, controller.getFavoriteControls()[0])
+        assertEquals(1, controller.getFavorites().size)
+        assertEquals(TEST_STRUCTURE_INFO, controller.getFavorites()[0])
 
         verify(persistenceWrapper, never()).storeFavorites(ArgumentMatchers.anyList())
     }
 
     @Test
     fun testPackageRemoved_hasFavorites() {
-        controller.changeFavoriteStatus(TEST_CONTROL_INFO, true)
-        controller.changeFavoriteStatus(TEST_CONTROL_INFO_2, true)
+        controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO)
+        controller.replaceFavoritesForStructure(TEST_STRUCTURE_INFO_2)
+        delayableExecutor.runAllReady()
 
         val serviceInfo = mock(ServiceInfo::class.java)
         `when`(serviceInfo.componentName).thenReturn(TEST_COMPONENT)
@@ -591,14 +623,14 @@ class ControlsControllerImplTest : SysuiTestCase() {
 
         // Don't want to check what happens before this call
         reset(persistenceWrapper)
-        listingCallbackCaptor.value.onServicesUpdated(listOf(info))
 
+        listingCallbackCaptor.value.onServicesUpdated(listOf(info))
         delayableExecutor.runAllReady()
 
         verify(bindingController).onComponentRemoved(TEST_COMPONENT_2)
 
-        assertEquals(1, controller.getFavoriteControls().size)
-        assertEquals(TEST_CONTROL_INFO, controller.getFavoriteControls()[0])
+        assertEquals(1, controller.getFavorites().size)
+        assertEquals(TEST_STRUCTURE_INFO, controller.getFavorites()[0])
 
         verify(persistenceWrapper).storeFavorites(ArgumentMatchers.anyList())
     }