From 639a826dd40a3de9d2d5d0a554cfca79aeebfda6 Mon Sep 17 00:00:00 2001
From: Anton Potapov <apotapov@google.com>
Date: Wed, 15 Mar 2023 17:25:23 +0000
Subject: [PATCH] Bind controls edit and add flows

It's possible to navigate from editing controls to adding new ones and
back when the flag is active. This affects both setting up a new app and
amending an existing one.

Protected by Flags.CONTROLS_MANAGEMENT_NEW_FLOWS

Test: Manual + added autotests
Fixes: 265180342
Change-Id: I471e7581982453e126f515272b7199a40e257643
---
 .../res/layout/controls_management.xml        |  23 +++
 packages/SystemUI/res/values/strings.xml      |   9 +
 .../controller/ControlsControllerImpl.kt      |   2 -
 .../systemui/controls/management/AllModel.kt  |   9 +-
 .../management/ControlsEditingActivity.kt     |  53 +++++-
 .../management/ControlsFavoritingActivity.kt  |  89 ++++++++--
 .../controls/management/ControlsModel.kt      |   7 +-
 .../ControlsProviderSelectorActivity.kt       |   5 +-
 .../controls/management/AllModelTest.kt       |  25 +++
 .../management/ControlsEditingActivityTest.kt | 142 +++++++++++++--
 .../ControlsFavoritingActivityTest.kt         | 164 ++++++++++++++++--
 .../controls/management/StartActivityData.kt  |   6 +
 12 files changed, 473 insertions(+), 61 deletions(-)
 create mode 100644 packages/SystemUI/tests/src/com/android/systemui/controls/management/StartActivityData.kt

diff --git a/packages/SystemUI/res/layout/controls_management.xml b/packages/SystemUI/res/layout/controls_management.xml
index b9e711e54b3b0..d8967d4706ce2 100644
--- a/packages/SystemUI/res/layout/controls_management.xml
+++ b/packages/SystemUI/res/layout/controls_management.xml
@@ -76,6 +76,29 @@
                 app:layout_constraintBottom_toBottomOf="parent"
                 app:layout_constraintStart_toStartOf="parent"/>
 
+            <Button
+                android:id="@+id/rearrange"
+                android:visibility="gone"
+                android:layout_width="wrap_content"
+                android:layout_height="match_parent"
+                android:gravity="center_vertical"
+                style="@*android:style/Widget.DeviceDefault.Button.Borderless.Colored"
+                app:layout_constraintTop_toTopOf="parent"
+                app:layout_constraintBottom_toBottomOf="parent"
+                app:layout_constraintStart_toStartOf="parent"/>
+
+            <Button
+                android:id="@+id/addControls"
+                android:visibility="gone"
+                android:layout_width="wrap_content"
+                android:layout_height="match_parent"
+                android:gravity="center_vertical"
+                android:text="@string/controls_favorite_add_controls"
+                style="@*android:style/Widget.DeviceDefault.Button.Borderless.Colored"
+                app:layout_constraintTop_toTopOf="parent"
+                app:layout_constraintBottom_toBottomOf="parent"
+                app:layout_constraintStart_toStartOf="parent"/>
+
             <Button
                 android:id="@+id/done"
                 android:layout_width="wrap_content"
diff --git a/packages/SystemUI/res/values/strings.xml b/packages/SystemUI/res/values/strings.xml
index 324ba02a7a46c..1dd12eef4639b 100644
--- a/packages/SystemUI/res/values/strings.xml
+++ b/packages/SystemUI/res/values/strings.xml
@@ -2471,6 +2471,15 @@
     <!-- Controls management favorites screen. See other apps button [CHAR LIMIT=30] -->
     <string name="controls_favorite_see_other_apps">See other apps</string>
 
+    <!-- Controls management favorites screen. Rearrange controls button [CHAR LIMIT=30]-->
+    <string name="controls_favorite_rearrange_button">Rearrange</string>
+
+    <!-- Controls management edit screen. Add controls button [CHAR LIMIT=30]-->
+    <string name="controls_favorite_add_controls">Add controls</string>
+
+    <!-- Controls management edit screen. Return to editing button [CHAR LIMIT=30]-->
+    <string name="controls_favorite_back_to_editing">Back to editing</string>
+
     <!-- Controls management controls screen error on load message [CHAR LIMIT=NONE] -->
     <string name="controls_favorite_load_error">Controls could not be loaded. Check the <xliff:g id="app" example="System UI">%s</xliff:g> app to make sure that the app settings haven\u2019t changed.</string>
     <!-- Controls management controls screen no controls found on load message [CHAR LIMIT=NONE] -->
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 ac1150e0fcb19..e8c97bf772716 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsControllerImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/controller/ControlsControllerImpl.kt
@@ -31,7 +31,6 @@ import android.service.controls.actions.ControlAction
 import android.util.ArrayMap
 import android.util.Log
 import com.android.internal.annotations.VisibleForTesting
-import com.android.internal.notification.NotificationAccessConfirmationActivityContract.EXTRA_USER_ID
 import com.android.systemui.Dumpable
 import com.android.systemui.backup.BackupHelper
 import com.android.systemui.controls.ControlStatus
@@ -44,7 +43,6 @@ import com.android.systemui.controls.ui.SelectedItem
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.dump.DumpManager
-import com.android.systemui.people.widget.PeopleSpaceWidgetProvider.EXTRA_USER_HANDLE
 import com.android.systemui.settings.UserFileManager
 import com.android.systemui.settings.UserTracker
 import com.android.systemui.statusbar.policy.DeviceControlsControllerImpl.Companion.PREFS_CONTROLS_FILE
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 00a406e4dbc04..be428a84da2fa 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/management/AllModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/management/AllModel.kt
@@ -75,9 +75,12 @@ class AllModel(
         } else {
             favoriteIds.remove(controlId)
         }
-        if (changed && !modified) {
-            modified = true
-            controlsModelCallback.onFirstChange()
+        if (changed) {
+            if (!modified) {
+                modified = true
+                controlsModelCallback.onFirstChange()
+            }
+            controlsModelCallback.onChange()
         }
         toChange?.let {
             it.controlStatus.favorite = favorite
diff --git a/packages/SystemUI/src/com/android/systemui/controls/management/ControlsEditingActivity.kt b/packages/SystemUI/src/com/android/systemui/controls/management/ControlsEditingActivity.kt
index 7df08651d5ab6..d629e3ea365e7 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/management/ControlsEditingActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/management/ControlsEditingActivity.kt
@@ -27,6 +27,7 @@ import android.view.ViewGroup
 import android.view.ViewStub
 import android.widget.Button
 import android.widget.TextView
+import android.widget.Toast
 import android.window.OnBackInvokedCallback
 import android.window.OnBackInvokedDispatcher
 import androidx.activity.ComponentActivity
@@ -38,8 +39,9 @@ import com.android.systemui.controls.CustomIconCache
 import com.android.systemui.controls.controller.ControlsControllerImpl
 import com.android.systemui.controls.controller.StructureInfo
 import com.android.systemui.controls.ui.ControlsActivity
-import com.android.systemui.controls.ui.ControlsUiController
 import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.flags.FeatureFlags
+import com.android.systemui.flags.Flags
 import com.android.systemui.settings.UserTracker
 import java.util.concurrent.Executor
 import javax.inject.Inject
@@ -48,17 +50,19 @@ import javax.inject.Inject
  * Activity for rearranging and removing controls for a given structure
  */
 open class ControlsEditingActivity @Inject constructor(
+    featureFlags: FeatureFlags,
     @Main private val mainExecutor: Executor,
     private val controller: ControlsControllerImpl,
     private val userTracker: UserTracker,
     private val customIconCache: CustomIconCache,
-    private val uiController: ControlsUiController
 ) : ComponentActivity() {
 
     companion object {
         private const val DEBUG = false
         private const val TAG = "ControlsEditingActivity"
         const val EXTRA_STRUCTURE = ControlsFavoritingActivity.EXTRA_STRUCTURE
+        const val EXTRA_APP = ControlsFavoritingActivity.EXTRA_APP
+        const val EXTRA_FROM_FAVORITING = "extra_from_favoriting"
         private val SUBTITLE_ID = R.string.controls_favorite_rearrange
         private val EMPTY_TEXT_ID = R.string.controls_favorite_removed
     }
@@ -68,7 +72,12 @@ open class ControlsEditingActivity @Inject constructor(
     private lateinit var model: FavoritesModel
     private lateinit var subtitle: TextView
     private lateinit var saveButton: View
+    private lateinit var addControls: View
 
+    private var isFromFavoriting: Boolean = false
+
+    private val isNewFlowEnabled: Boolean =
+        featureFlags.isEnabled(Flags.CONTROLS_MANAGEMENT_NEW_FLOWS)
     private val userTrackerCallback: UserTracker.Callback = object : UserTracker.Callback {
         private val startingUser = controller.currentUserId
 
@@ -93,7 +102,7 @@ open class ControlsEditingActivity @Inject constructor(
         intent.getParcelableExtra<ComponentName>(Intent.EXTRA_COMPONENT_NAME)?.let {
             component = it
         } ?: run(this::finish)
-
+        isFromFavoriting = intent.getBooleanExtra(EXTRA_FROM_FAVORITING, false)
         intent.getCharSequenceExtra(EXTRA_STRUCTURE)?.let {
             structure = it
         } ?: run(this::finish)
@@ -165,8 +174,42 @@ open class ControlsEditingActivity @Inject constructor(
     }
 
     private fun bindButtons() {
+        addControls = requireViewById<Button>(R.id.addControls).apply {
+            isEnabled = true
+            visibility = if (isNewFlowEnabled) View.VISIBLE else View.GONE
+            setOnClickListener {
+                if (saveButton.isEnabled) {
+                    // The user has made changes
+                    Toast.makeText(
+                        applicationContext,
+                        R.string.controls_favorite_toast_no_changes,
+                        Toast.LENGTH_SHORT
+                    ).show()
+                }
+                if (isFromFavoriting) {
+                    animateExitAndFinish()
+                } else {
+                    startActivity(Intent(context, ControlsFavoritingActivity::class.java).also {
+                        it.putExtra(ControlsFavoritingActivity.EXTRA_STRUCTURE, structure)
+                        it.putExtra(Intent.EXTRA_COMPONENT_NAME, component)
+                        it.putExtra(
+                            ControlsFavoritingActivity.EXTRA_APP,
+                            intent.getCharSequenceExtra(EXTRA_APP),
+                        )
+                        it.putExtra(
+                            ControlsFavoritingActivity.EXTRA_SOURCE,
+                            ControlsFavoritingActivity.EXTRA_SOURCE_VALUE_FROM_EDITING,
+                        )
+                    },
+                                  ActivityOptions.makeSceneTransitionAnimation(
+                                      this@ControlsEditingActivity
+                                  ).toBundle(),
+                    )
+                }
+            }
+        }
         saveButton = requireViewById<Button>(R.id.done).apply {
-            isEnabled = false
+            isEnabled = isFromFavoriting
             setText(R.string.save)
             setOnClickListener {
                 saveFavorites()
@@ -194,6 +237,8 @@ open class ControlsEditingActivity @Inject constructor(
             }
         }
 
+        override fun onChange() = Unit
+
         override fun onFirstChange() {
             saveButton.isEnabled = true
         }
diff --git a/packages/SystemUI/src/com/android/systemui/controls/management/ControlsFavoritingActivity.kt b/packages/SystemUI/src/com/android/systemui/controls/management/ControlsFavoritingActivity.kt
index 3e97d3132bc78..d3ffc95853355 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/management/ControlsFavoritingActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/management/ControlsFavoritingActivity.kt
@@ -37,6 +37,7 @@ import android.widget.Toast
 import android.window.OnBackInvokedCallback
 import android.window.OnBackInvokedDispatcher
 import androidx.activity.ComponentActivity
+import androidx.annotation.VisibleForTesting
 import androidx.viewpager2.widget.ViewPager2
 import com.android.systemui.Prefs
 import com.android.systemui.R
@@ -45,20 +46,20 @@ import com.android.systemui.controls.TooltipManager
 import com.android.systemui.controls.controller.ControlsControllerImpl
 import com.android.systemui.controls.controller.StructureInfo
 import com.android.systemui.controls.ui.ControlsActivity
-import com.android.systemui.controls.ui.ControlsUiController
 import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.flags.FeatureFlags
+import com.android.systemui.flags.Flags
 import com.android.systemui.settings.UserTracker
 import java.text.Collator
 import java.util.concurrent.Executor
-import java.util.function.Consumer
 import javax.inject.Inject
 
 open class ControlsFavoritingActivity @Inject constructor(
+    featureFlags: FeatureFlags,
     @Main private val executor: Executor,
     private val controller: ControlsControllerImpl,
     private val listingController: ControlsListingController,
     private val userTracker: UserTracker,
-    private val uiController: ControlsUiController
 ) : ComponentActivity() {
 
     companion object {
@@ -71,7 +72,10 @@ open class ControlsFavoritingActivity @Inject constructor(
         // If provided, show this structure page first
         const val EXTRA_STRUCTURE = "extra_structure"
         const val EXTRA_SINGLE_STRUCTURE = "extra_single_structure"
-        const val EXTRA_FROM_PROVIDER_SELECTOR = "extra_from_provider_selector"
+        const val EXTRA_SOURCE = "extra_source"
+        const val EXTRA_SOURCE_UNDEFINED: Byte = 0
+        const val EXTRA_SOURCE_VALUE_FROM_PROVIDER_SELECTOR: Byte = 1
+        const val EXTRA_SOURCE_VALUE_FROM_EDITING: Byte = 2
         private const val TOOLTIP_PREFS_KEY = Prefs.Key.CONTROLS_STRUCTURE_SWIPE_TOOLTIP_COUNT
         private const val TOOLTIP_MAX_SHOWN = 2
     }
@@ -79,7 +83,7 @@ open class ControlsFavoritingActivity @Inject constructor(
     private var component: ComponentName? = null
     private var appName: CharSequence? = null
     private var structureExtra: CharSequence? = null
-    private var fromProviderSelector = false
+    private var openSource = EXTRA_SOURCE_UNDEFINED
 
     private lateinit var structurePager: ViewPager2
     private lateinit var statusText: TextView
@@ -89,12 +93,19 @@ open class ControlsFavoritingActivity @Inject constructor(
     private var mTooltipManager: TooltipManager? = null
     private lateinit var doneButton: View
     private lateinit var otherAppsButton: View
+    private lateinit var rearrangeButton: Button
     private var listOfStructures = emptyList<StructureContainer>()
 
     private lateinit var comparator: Comparator<StructureContainer>
     private var cancelLoadRunnable: Runnable? = null
     private var isPagerLoaded = false
 
+    private val fromProviderSelector: Boolean
+        get() = openSource == EXTRA_SOURCE_VALUE_FROM_PROVIDER_SELECTOR
+    private val fromEditing: Boolean
+        get() = openSource == EXTRA_SOURCE_VALUE_FROM_EDITING
+    private val isNewFlowEnabled: Boolean =
+        featureFlags.isEnabled(Flags.CONTROLS_MANAGEMENT_NEW_FLOWS)
     private val userTrackerCallback: UserTracker.Callback = object : UserTracker.Callback {
         private val startingUser = controller.currentUserId
 
@@ -117,14 +128,20 @@ open class ControlsFavoritingActivity @Inject constructor(
 
         override fun onServicesUpdated(serviceInfos: List<ControlsServiceInfo>) {
             if (serviceInfos.size > 1) {
-                otherAppsButton.post {
-                    otherAppsButton.visibility = View.VISIBLE
+                val newVisibility = if (isNewFlowEnabled) View.GONE else View.VISIBLE
+                if (otherAppsButton.visibility != newVisibility) {
+                    otherAppsButton.post {
+                        otherAppsButton.visibility = newVisibility
+                    }
                 }
             }
         }
     }
 
     override fun onBackPressed() {
+        if (fromEditing) {
+            animateExitAndFinish()
+        }
         if (!fromProviderSelector) {
             openControlsOrigin()
         }
@@ -139,7 +156,7 @@ open class ControlsFavoritingActivity @Inject constructor(
         appName = intent.getCharSequenceExtra(EXTRA_APP)
         structureExtra = intent.getCharSequenceExtra(EXTRA_STRUCTURE)
         component = intent.getParcelableExtra<ComponentName>(Intent.EXTRA_COMPONENT_NAME)
-        fromProviderSelector = intent.getBooleanExtra(EXTRA_FROM_PROVIDER_SELECTOR, false)
+        openSource = intent.getByteExtra(EXTRA_SOURCE, EXTRA_SOURCE_UNDEFINED)
 
         bindViews()
     }
@@ -148,14 +165,19 @@ open class ControlsFavoritingActivity @Inject constructor(
         override fun onFirstChange() {
             doneButton.isEnabled = true
         }
+
+        override fun onChange() {
+            val structure: StructureContainer = listOfStructures[structurePager.currentItem]
+            rearrangeButton.isEnabled = structure.model.favorites.isNotEmpty()
+        }
     }
 
     private fun loadControls() {
-        component?.let {
+        component?.let { componentName ->
             statusText.text = resources.getText(com.android.internal.R.string.loading)
             val emptyZoneString = resources.getText(
                     R.string.controls_favorite_other_zone_header)
-            controller.loadForComponent(it, Consumer { data ->
+            controller.loadForComponent(componentName, { data ->
                 val allControls = data.allControls
                 val favoriteKeys = data.favoritesIds
                 val error = data.errorOnLoad
@@ -213,7 +235,7 @@ open class ControlsFavoritingActivity @Inject constructor(
                         ControlsAnimations.enterAnimation(structurePager).start()
                     }
                 }
-            }, Consumer { runnable -> cancelLoadRunnable = runnable })
+            }, { runnable -> cancelLoadRunnable = runnable })
         }
     }
 
@@ -299,7 +321,8 @@ open class ControlsFavoritingActivity @Inject constructor(
         bindButtons()
     }
 
-    private fun animateExitAndFinish() {
+    @VisibleForTesting
+    internal open fun animateExitAndFinish() {
         val rootView = requireViewById<ViewGroup>(R.id.controls_management_root)
         ControlsAnimations.exitAnimation(
                 rootView,
@@ -312,6 +335,32 @@ open class ControlsFavoritingActivity @Inject constructor(
     }
 
     private fun bindButtons() {
+        rearrangeButton = requireViewById<Button>(R.id.rearrange).apply {
+            text = if (fromEditing) {
+                getString(R.string.controls_favorite_back_to_editing)
+            } else {
+                getString(R.string.controls_favorite_rearrange_button)
+            }
+            isEnabled = false
+            visibility = if (isNewFlowEnabled) View.VISIBLE else View.GONE
+            setOnClickListener {
+                if (component == null) return@setOnClickListener
+                saveFavorites()
+                startActivity(
+                    Intent(context, ControlsEditingActivity::class.java).also {
+                        it.putExtra(Intent.EXTRA_COMPONENT_NAME, component)
+                        it.putExtra(ControlsEditingActivity.EXTRA_APP, appName)
+                        it.putExtra(ControlsEditingActivity.EXTRA_FROM_FAVORITING, true)
+                        it.putExtra(
+                            ControlsEditingActivity.EXTRA_STRUCTURE,
+                            listOfStructures[structurePager.currentItem].structureName,
+                        )
+                    },
+                    ActivityOptions
+                        .makeSceneTransitionAnimation(this@ControlsFavoritingActivity).toBundle()
+                )
+            }
+        }
         otherAppsButton = requireViewById<Button>(R.id.other_apps).apply {
             setOnClickListener {
                 if (doneButton.isEnabled) {
@@ -335,18 +384,22 @@ open class ControlsFavoritingActivity @Inject constructor(
             isEnabled = false
             setOnClickListener {
                 if (component == null) return@setOnClickListener
-                listOfStructures.forEach {
-                    val favoritesForStorage = it.model.favorites
-                    controller.replaceFavoritesForStructure(
-                        StructureInfo(component!!, it.structureName, favoritesForStorage)
-                    )
-                }
+                saveFavorites()
                 animateExitAndFinish()
                 openControlsOrigin()
             }
         }
     }
 
+    private fun saveFavorites() {
+        listOfStructures.forEach {
+            val favoritesForStorage = it.model.favorites
+            controller.replaceFavoritesForStructure(
+                StructureInfo(component!!, it.structureName, favoritesForStorage)
+            )
+        }
+    }
+
     private fun openControlsOrigin() {
         startActivity(
             Intent(applicationContext, ControlsActivity::class.java),
diff --git a/packages/SystemUI/src/com/android/systemui/controls/management/ControlsModel.kt b/packages/SystemUI/src/com/android/systemui/controls/management/ControlsModel.kt
index d65481a8a8302..3455e6dbc6561 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/management/ControlsModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/management/ControlsModel.kt
@@ -71,6 +71,11 @@ interface ControlsModel {
          * Use to notify that the model has changed for the first time
          */
         fun onFirstChange()
+
+        /**
+         * Use to notify that the model has changed
+         */
+        fun onChange()
     }
 
     /**
@@ -132,7 +137,7 @@ data class ControlInfoWrapper(
         controlInfo: ControlInfo,
         favorite: Boolean,
         customIconGetter: (ComponentName, String) -> Icon?
-    ): this(component, controlInfo, favorite) {
+    ) : this(component, controlInfo, favorite) {
         this.customIconGetter = customIconGetter
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/controls/management/ControlsProviderSelectorActivity.kt b/packages/SystemUI/src/com/android/systemui/controls/management/ControlsProviderSelectorActivity.kt
index 3808e73ca085a..92aff0624bdc8 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/management/ControlsProviderSelectorActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/management/ControlsProviderSelectorActivity.kt
@@ -211,7 +211,10 @@ open class ControlsProviderSelectorActivity @Inject constructor(
                     putExtra(ControlsFavoritingActivity.EXTRA_APP,
                             listingController.getAppLabel(it))
                     putExtra(Intent.EXTRA_COMPONENT_NAME, it)
-                    putExtra(ControlsFavoritingActivity.EXTRA_FROM_PROVIDER_SELECTOR, true)
+                    putExtra(
+                        ControlsFavoritingActivity.EXTRA_SOURCE,
+                        ControlsFavoritingActivity.EXTRA_SOURCE_VALUE_FROM_PROVIDER_SELECTOR,
+                    )
                 }
                 startActivity(intent, ActivityOptions.makeSceneTransitionAnimation(this).toBundle())
                 animateExitAndFinish()
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 236384b095142..4ea96169e8a58 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
@@ -32,6 +32,7 @@ import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mock
 import org.mockito.Mockito.never
+import org.mockito.Mockito.times
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
 
@@ -161,6 +162,7 @@ class AllModelTest : SysuiTestCase() {
         }
 
         verify(controlsModelCallback).onFirstChange()
+        verify(controlsModelCallback).onChange()
     }
 
     @Test
@@ -176,6 +178,7 @@ class AllModelTest : SysuiTestCase() {
         )
 
         verify(controlsModelCallback).onFirstChange()
+        verify(controlsModelCallback).onChange()
     }
 
     @Test
@@ -191,6 +194,7 @@ class AllModelTest : SysuiTestCase() {
         }
 
         verify(controlsModelCallback, never()).onFirstChange()
+        verify(controlsModelCallback, never()).onChange()
     }
 
     @Test
@@ -207,6 +211,7 @@ class AllModelTest : SysuiTestCase() {
         }
 
         verify(controlsModelCallback).onFirstChange()
+        verify(controlsModelCallback).onChange()
     }
 
     @Test
@@ -222,6 +227,7 @@ class AllModelTest : SysuiTestCase() {
         )
 
         verify(controlsModelCallback).onFirstChange()
+        verify(controlsModelCallback).onChange()
     }
 
     @Test
@@ -236,5 +242,24 @@ class AllModelTest : SysuiTestCase() {
         }
 
         verify(controlsModelCallback, never()).onFirstChange()
+        verify(controlsModelCallback, never()).onChange()
+    }
+
+    @Test
+    fun testAddSecondChange_callbacks() {
+        model.changeFavoriteStatus("${idPrefix}4", true)
+        model.changeFavoriteStatus("${idPrefix}5", true)
+
+        verify(controlsModelCallback).onFirstChange()
+        verify(controlsModelCallback, times(2)).onChange()
+    }
+
+    @Test
+    fun testRemoveSecondChange_callbacks() {
+        model.changeFavoriteStatus("${idPrefix}1", false)
+        model.changeFavoriteStatus("${idPrefix}3", false)
+
+        verify(controlsModelCallback).onFirstChange()
+        verify(controlsModelCallback, times(2)).onChange()
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/controls/management/ControlsEditingActivityTest.kt b/packages/SystemUI/tests/src/com/android/systemui/controls/management/ControlsEditingActivityTest.kt
index 3b6f7d19e93eb..42106756b473d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/controls/management/ControlsEditingActivityTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/controls/management/ControlsEditingActivityTest.kt
@@ -2,27 +2,33 @@ package com.android.systemui.controls.management
 
 import android.content.ComponentName
 import android.content.Intent
+import android.os.Bundle
 import android.testing.AndroidTestingRunner
 import android.testing.TestableLooper
+import android.view.View
+import android.widget.Button
 import android.window.OnBackInvokedCallback
 import android.window.OnBackInvokedDispatcher
 import androidx.test.filters.SmallTest
 import androidx.test.rule.ActivityTestRule
 import androidx.test.runner.intercepting.SingleActivityFactory
+import com.android.systemui.R
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.controls.CustomIconCache
 import com.android.systemui.controls.controller.ControlsControllerImpl
-import com.android.systemui.controls.ui.ControlsUiController
+import com.android.systemui.flags.FakeFeatureFlags
+import com.android.systemui.flags.Flags
 import com.android.systemui.settings.UserTracker
 import com.android.systemui.util.concurrency.FakeExecutor
 import com.android.systemui.util.time.FakeSystemClock
+import com.google.common.truth.Truth.assertThat
 import java.util.concurrent.CountDownLatch
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.ArgumentCaptor
-import org.mockito.ArgumentMatchers
+import org.mockito.ArgumentMatchers.eq
 import org.mockito.Captor
 import org.mockito.Mock
 import org.mockito.Mockito.verify
@@ -32,7 +38,15 @@ import org.mockito.MockitoAnnotations
 @RunWith(AndroidTestingRunner::class)
 @TestableLooper.RunWithLooper
 class ControlsEditingActivityTest : SysuiTestCase() {
+
+    private companion object {
+        val TEST_COMPONENT = ComponentName("TestPackageName", "TestClassName")
+        val TEST_STRUCTURE: CharSequence = "TestStructure"
+        val TEST_APP: CharSequence = "TestApp"
+    }
+
     private val uiExecutor = FakeExecutor(FakeSystemClock())
+    private val featureFlags = FakeFeatureFlags()
 
     @Mock lateinit var controller: ControlsControllerImpl
 
@@ -40,9 +54,6 @@ class ControlsEditingActivityTest : SysuiTestCase() {
 
     @Mock lateinit var customIconCache: CustomIconCache
 
-    @Mock lateinit var uiController: ControlsUiController
-
-    private lateinit var controlsEditingActivity: ControlsEditingActivity_Factory
     private var latch: CountDownLatch = CountDownLatch(1)
 
     @Mock private lateinit var mockDispatcher: OnBackInvokedDispatcher
@@ -58,11 +69,11 @@ class ControlsEditingActivityTest : SysuiTestCase() {
                 ) {
                 override fun create(intent: Intent?): TestableControlsEditingActivity {
                     return TestableControlsEditingActivity(
+                        featureFlags,
                         uiExecutor,
                         controller,
                         userTracker,
                         customIconCache,
-                        uiController,
                         mockDispatcher,
                         latch
                     )
@@ -75,19 +86,17 @@ class ControlsEditingActivityTest : SysuiTestCase() {
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
-        val intent = Intent()
-        intent.putExtra(ControlsEditingActivity.EXTRA_STRUCTURE, "TestTitle")
-        val cname = ComponentName("TestPackageName", "TestClassName")
-        intent.putExtra(Intent.EXTRA_COMPONENT_NAME, cname)
-        activityRule.launchActivity(intent)
+
+        featureFlags.set(Flags.CONTROLS_MANAGEMENT_NEW_FLOWS, false)
     }
 
     @Test
     fun testBackCallbackRegistrationAndUnregistration() {
+        launchActivity()
         // 1. ensure that launching the activity results in it registering a callback
         verify(mockDispatcher)
             .registerOnBackInvokedCallback(
-                ArgumentMatchers.eq(OnBackInvokedDispatcher.PRIORITY_DEFAULT),
+                eq(OnBackInvokedDispatcher.PRIORITY_DEFAULT),
                 captureCallback.capture()
             )
         activityRule.finishActivity()
@@ -96,15 +105,102 @@ class ControlsEditingActivityTest : SysuiTestCase() {
         verify(mockDispatcher).unregisterOnBackInvokedCallback(captureCallback.value)
     }
 
-    public class TestableControlsEditingActivity(
-        private val executor: FakeExecutor,
-        private val controller: ControlsControllerImpl,
-        private val userTracker: UserTracker,
-        private val customIconCache: CustomIconCache,
-        private val uiController: ControlsUiController,
+    @Test
+    fun testNewFlowDisabled_addControlsButton_gone() {
+        with(launchActivity()) {
+            val addControlsButton = requireViewById<Button>(R.id.addControls)
+            assertThat(addControlsButton.visibility).isEqualTo(View.GONE)
+        }
+    }
+
+    @Test
+    fun testNewFlowEnabled_addControlsButton_visible() {
+        featureFlags.set(Flags.CONTROLS_MANAGEMENT_NEW_FLOWS, true)
+        with(launchActivity()) {
+            val addControlsButton = requireViewById<Button>(R.id.addControls)
+            assertThat(addControlsButton.visibility).isEqualTo(View.VISIBLE)
+            assertThat(addControlsButton.isEnabled).isTrue()
+        }
+    }
+
+    @Test
+    fun testNotLaunchFromFavoriting_saveButton_disabled() {
+        featureFlags.set(Flags.CONTROLS_MANAGEMENT_NEW_FLOWS, true)
+        with(launchActivity(isFromFavoriting = false)) {
+            val saveButton = requireViewById<Button>(R.id.done)
+            assertThat(saveButton.isEnabled).isFalse()
+        }
+    }
+
+    @Test
+    fun testLaunchFromFavoriting_saveButton_enabled() {
+        featureFlags.set(Flags.CONTROLS_MANAGEMENT_NEW_FLOWS, true)
+        with(launchActivity(isFromFavoriting = true)) {
+            val saveButton = requireViewById<Button>(R.id.done)
+            assertThat(saveButton.isEnabled).isTrue()
+        }
+    }
+
+    @Test
+    fun testNotFromFavoriting_addControlsPressed_launchesFavouriting() {
+        featureFlags.set(Flags.CONTROLS_MANAGEMENT_NEW_FLOWS, true)
+        with(launchActivity(isFromFavoriting = false)) {
+            val addControls = requireViewById<Button>(R.id.addControls)
+
+            activityRule.runOnUiThread { addControls.performClick() }
+
+            with(startActivityData!!.intent) {
+                assertThat(component)
+                    .isEqualTo(ComponentName(context, ControlsFavoritingActivity::class.java))
+                assertThat(getCharSequenceExtra(ControlsFavoritingActivity.EXTRA_STRUCTURE))
+                    .isEqualTo(TEST_STRUCTURE)
+                assertThat(
+                        getParcelableExtra(Intent.EXTRA_COMPONENT_NAME, ComponentName::class.java)
+                    )
+                    .isEqualTo(TEST_COMPONENT)
+                assertThat(getCharSequenceExtra(ControlsFavoritingActivity.EXTRA_APP))
+                    .isEqualTo(TEST_APP)
+                assertThat(getByteExtra(ControlsFavoritingActivity.EXTRA_SOURCE, -1))
+                    .isEqualTo(ControlsFavoritingActivity.EXTRA_SOURCE_VALUE_FROM_EDITING)
+            }
+        }
+    }
+
+    private fun launchActivity(
+        componentName: ComponentName = TEST_COMPONENT,
+        structure: CharSequence = TEST_STRUCTURE,
+        isFromFavoriting: Boolean = false,
+        app: CharSequence = TEST_APP,
+    ): TestableControlsEditingActivity =
+        activityRule.launchActivity(
+            Intent().apply {
+                putExtra(ControlsEditingActivity.EXTRA_FROM_FAVORITING, isFromFavoriting)
+                putExtra(ControlsEditingActivity.EXTRA_STRUCTURE, structure)
+                putExtra(Intent.EXTRA_COMPONENT_NAME, componentName)
+                putExtra(ControlsEditingActivity.EXTRA_APP, app)
+            }
+        )
+
+    class TestableControlsEditingActivity(
+        featureFlags: FakeFeatureFlags,
+        executor: FakeExecutor,
+        controller: ControlsControllerImpl,
+        userTracker: UserTracker,
+        customIconCache: CustomIconCache,
         private val mockDispatcher: OnBackInvokedDispatcher,
         private val latch: CountDownLatch
-    ) : ControlsEditingActivity(executor, controller, userTracker, customIconCache, uiController) {
+    ) :
+        ControlsEditingActivity(
+            featureFlags,
+            executor,
+            controller,
+            userTracker,
+            customIconCache,
+        ) {
+
+        var startActivityData: StartActivityData? = null
+            private set
+
         override fun getOnBackInvokedDispatcher(): OnBackInvokedDispatcher {
             return mockDispatcher
         }
@@ -114,5 +210,13 @@ class ControlsEditingActivityTest : SysuiTestCase() {
             // ensures that test runner thread does not proceed until ui thread is done
             latch.countDown()
         }
+
+        override fun startActivity(intent: Intent) {
+            startActivityData = StartActivityData(intent, null)
+        }
+
+        override fun startActivity(intent: Intent, options: Bundle?) {
+            startActivityData = StartActivityData(intent, options)
+        }
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/controls/management/ControlsFavoritingActivityTest.kt b/packages/SystemUI/tests/src/com/android/systemui/controls/management/ControlsFavoritingActivityTest.kt
index 365523233353a..68846168d17b6 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/controls/management/ControlsFavoritingActivityTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/controls/management/ControlsFavoritingActivityTest.kt
@@ -1,30 +1,49 @@
 package com.android.systemui.controls.management
 
+import android.content.ComponentName
 import android.content.Intent
+import android.os.Bundle
+import android.service.controls.Control
 import android.testing.AndroidTestingRunner
 import android.testing.TestableLooper
+import android.view.View
+import android.widget.Button
 import android.window.OnBackInvokedCallback
 import android.window.OnBackInvokedDispatcher
 import androidx.test.filters.FlakyTest
 import androidx.test.filters.SmallTest
 import androidx.test.rule.ActivityTestRule
 import androidx.test.runner.intercepting.SingleActivityFactory
+import com.android.systemui.R
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.controls.ControlStatus
+import com.android.systemui.controls.ControlsServiceInfo
+import com.android.systemui.controls.controller.ControlsController
 import com.android.systemui.controls.controller.ControlsControllerImpl
-import com.android.systemui.controls.ui.ControlsUiController
+import com.android.systemui.controls.controller.createLoadDataObject
 import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.flags.FakeFeatureFlags
+import com.android.systemui.flags.FeatureFlags
+import com.android.systemui.flags.Flags
 import com.android.systemui.settings.UserTracker
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.capture
+import com.android.systemui.util.mockito.whenever
+import com.google.common.truth.Truth.assertThat
 import com.google.common.util.concurrent.MoreExecutors
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.Executor
+import java.util.function.Consumer
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.Answers
 import org.mockito.ArgumentCaptor
-import org.mockito.ArgumentMatchers
 import org.mockito.Captor
 import org.mockito.Mock
+import org.mockito.Mockito.eq
+import org.mockito.Mockito.mock
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
 
@@ -32,7 +51,19 @@ import org.mockito.MockitoAnnotations
 @RunWith(AndroidTestingRunner::class)
 @TestableLooper.RunWithLooper
 class ControlsFavoritingActivityTest : SysuiTestCase() {
+
+    private companion object {
+        val TEST_COMPONENT = ComponentName("TestPackageName", "TestClassName")
+        val TEST_CONTROL =
+            mock(Control::class.java, Answers.RETURNS_MOCKS)!!.apply {
+                whenever(structure).thenReturn(TEST_STRUCTURE)
+            }
+        val TEST_STRUCTURE: CharSequence = "TestStructure"
+        val TEST_APP: CharSequence = "TestApp"
+    }
+
     @Main private val executor: Executor = MoreExecutors.directExecutor()
+    private val featureFlags = FakeFeatureFlags()
 
     @Mock lateinit var controller: ControlsControllerImpl
 
@@ -40,13 +71,15 @@ class ControlsFavoritingActivityTest : SysuiTestCase() {
 
     @Mock lateinit var userTracker: UserTracker
 
-    @Mock lateinit var uiController: ControlsUiController
-
-    private lateinit var controlsFavoritingActivity: ControlsFavoritingActivity_Factory
     private var latch: CountDownLatch = CountDownLatch(1)
 
     @Mock private lateinit var mockDispatcher: OnBackInvokedDispatcher
     @Captor private lateinit var captureCallback: ArgumentCaptor<OnBackInvokedCallback>
+    @Captor
+    private lateinit var listingCallback:
+        ArgumentCaptor<ControlsListingController.ControlsListingCallback>
+    @Captor
+    private lateinit var controlsCallback: ArgumentCaptor<Consumer<ControlsController.LoadData>>
 
     @Rule
     @JvmField
@@ -58,11 +91,11 @@ class ControlsFavoritingActivityTest : SysuiTestCase() {
                 ) {
                 override fun create(intent: Intent?): TestableControlsFavoritingActivity {
                     return TestableControlsFavoritingActivity(
+                        featureFlags,
                         executor,
                         controller,
                         listingController,
                         userTracker,
-                        uiController,
                         mockDispatcher,
                         latch
                     )
@@ -75,19 +108,18 @@ class ControlsFavoritingActivityTest : SysuiTestCase() {
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
-        val intent = Intent()
-        intent.putExtra(ControlsFavoritingActivity.EXTRA_FROM_PROVIDER_SELECTOR, true)
-        activityRule.launchActivity(intent)
+        featureFlags.set(Flags.CONTROLS_MANAGEMENT_NEW_FLOWS, false)
     }
 
     // b/259549854 to root-cause and fix
     @FlakyTest
     @Test
     fun testBackCallbackRegistrationAndUnregistration() {
+        launchActivity()
         // 1. ensure that launching the activity results in it registering a callback
         verify(mockDispatcher)
             .registerOnBackInvokedCallback(
-                ArgumentMatchers.eq(OnBackInvokedDispatcher.PRIORITY_DEFAULT),
+                eq(OnBackInvokedDispatcher.PRIORITY_DEFAULT),
                 captureCallback.capture()
             )
         activityRule.finishActivity()
@@ -96,22 +128,116 @@ class ControlsFavoritingActivityTest : SysuiTestCase() {
         verify(mockDispatcher).unregisterOnBackInvokedCallback(captureCallback.value)
     }
 
-    public class TestableControlsFavoritingActivity(
+    @Test
+    fun testNewFlowEnabled_buttons() {
+        featureFlags.set(Flags.CONTROLS_MANAGEMENT_NEW_FLOWS, true)
+        with(launchActivity()) {
+            verify(listingController).addCallback(listingCallback.capture())
+            listingCallback.value.onServicesUpdated(
+                listOf(mock(ControlsServiceInfo::class.java), mock(ControlsServiceInfo::class.java))
+            )
+
+            val rearrangeButton = requireViewById<Button>(R.id.rearrange)
+            assertThat(rearrangeButton.visibility).isEqualTo(View.VISIBLE)
+            assertThat(rearrangeButton.isEnabled).isFalse()
+            assertThat(requireViewById<Button>(R.id.other_apps).visibility).isEqualTo(View.GONE)
+        }
+    }
+
+    @Test
+    fun testNewFlowDisabled_buttons() {
+        with(launchActivity()) {
+            verify(listingController).addCallback(listingCallback.capture())
+            activityRule.runOnUiThread {
+                listingCallback.value.onServicesUpdated(
+                    listOf(
+                        mock(ControlsServiceInfo::class.java),
+                        mock(ControlsServiceInfo::class.java)
+                    )
+                )
+            }
+
+            val rearrangeButton = requireViewById<Button>(R.id.rearrange)
+            assertThat(rearrangeButton.visibility).isEqualTo(View.GONE)
+            assertThat(rearrangeButton.isEnabled).isFalse()
+            assertThat(requireViewById<Button>(R.id.other_apps).visibility).isEqualTo(View.VISIBLE)
+        }
+    }
+
+    @Test
+    fun testNewFlowEnabled_rearrangePressed_savesAndlaunchesActivity() {
+        featureFlags.set(Flags.CONTROLS_MANAGEMENT_NEW_FLOWS, true)
+        with(launchActivity()) {
+            verify(listingController).addCallback(capture(listingCallback))
+            listingCallback.value.onServicesUpdated(
+                listOf(mock(ControlsServiceInfo::class.java), mock(ControlsServiceInfo::class.java))
+            )
+            verify(controller).loadForComponent(any(), capture(controlsCallback), any())
+            activityRule.runOnUiThread {
+                controlsCallback.value.accept(
+                    createLoadDataObject(
+                        listOf(ControlStatus(TEST_CONTROL, TEST_COMPONENT, true)),
+                        emptyList(),
+                    )
+                )
+                requireViewById<Button>(R.id.rearrange).performClick()
+            }
+
+            verify(controller).replaceFavoritesForStructure(any())
+            with(startActivityData!!.intent) {
+                assertThat(component)
+                    .isEqualTo(ComponentName(context, ControlsEditingActivity::class.java))
+                assertThat(
+                        getParcelableExtra(Intent.EXTRA_COMPONENT_NAME, ComponentName::class.java)
+                    )
+                    .isEqualTo(TEST_COMPONENT)
+                assertThat(getCharSequenceExtra(ControlsEditingActivity.EXTRA_APP))
+                    .isEqualTo(TEST_APP)
+                assertThat(getBooleanExtra(ControlsEditingActivity.EXTRA_FROM_FAVORITING, false))
+                    .isTrue()
+                assertThat(getCharSequenceExtra(ControlsEditingActivity.EXTRA_STRUCTURE))
+                    .isEqualTo("")
+            }
+        }
+    }
+
+    private fun launchActivity(
+        componentName: ComponentName = TEST_COMPONENT,
+        structure: CharSequence = TEST_STRUCTURE,
+        app: CharSequence = TEST_APP,
+        source: Byte = ControlsFavoritingActivity.EXTRA_SOURCE_VALUE_FROM_PROVIDER_SELECTOR,
+    ): TestableControlsFavoritingActivity =
+        activityRule.launchActivity(
+            Intent().apply {
+                putExtra(Intent.EXTRA_COMPONENT_NAME, componentName)
+                putExtra(ControlsFavoritingActivity.EXTRA_STRUCTURE, structure)
+                putExtra(ControlsFavoritingActivity.EXTRA_APP, app)
+                putExtra(ControlsFavoritingActivity.EXTRA_SOURCE, source)
+            }
+        )
+
+    class TestableControlsFavoritingActivity(
+        featureFlags: FeatureFlags,
         executor: Executor,
         controller: ControlsControllerImpl,
         listingController: ControlsListingController,
         userTracker: UserTracker,
-        uiController: ControlsUiController,
         private val mockDispatcher: OnBackInvokedDispatcher,
         private val latch: CountDownLatch
     ) :
         ControlsFavoritingActivity(
+            featureFlags,
             executor,
             controller,
             listingController,
             userTracker,
-            uiController
         ) {
+
+        var triedToFinish = false
+
+        var startActivityData: StartActivityData? = null
+            private set
+
         override fun getOnBackInvokedDispatcher(): OnBackInvokedDispatcher {
             return mockDispatcher
         }
@@ -121,5 +247,17 @@ class ControlsFavoritingActivityTest : SysuiTestCase() {
             // ensures that test runner thread does not proceed until ui thread is done
             latch.countDown()
         }
+
+        override fun startActivity(intent: Intent) {
+            startActivityData = StartActivityData(intent, null)
+        }
+
+        override fun startActivity(intent: Intent, options: Bundle?) {
+            startActivityData = StartActivityData(intent, options)
+        }
+
+        override fun animateExitAndFinish() {
+            triedToFinish = true
+        }
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/controls/management/StartActivityData.kt b/packages/SystemUI/tests/src/com/android/systemui/controls/management/StartActivityData.kt
new file mode 100644
index 0000000000000..977e3ba899f6b
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/controls/management/StartActivityData.kt
@@ -0,0 +1,6 @@
+package com.android.systemui.controls.management
+
+import android.content.Intent
+import android.os.Bundle
+
+data class StartActivityData(val intent: Intent, val options: Bundle?)