Remove calls to `new SystemUIDialog(context)`

This removes low-hanging fruit, but replacing calls to
`new SystemUIDialog(context)` with `Provider<SystemUIDialog>#get`.

Flag: NA
Bug: 219008720
Test: atest SystemUITests
Change-Id: I9364b35d84f4d5afb084547a14bc04686d1a9ba3

21 files changed, 260 insertions, 134 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/controls/management/PanelConfirmationDialogFactory.kt b/packages/SystemUI/src/com/android/systemui/controls/management/PanelConfirmationDialogFactory.kt
index 5df26b3176ff..a6b432019486 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/management/PanelConfirmationDialogFactory.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/management/PanelConfirmationDialogFactory.kt
@@ -28,27 +28,26 @@ import javax.inject.Inject
 /**
  * Factory to create dialogs for consenting to show app panels for specific apps.
  *
- * [internalDialogFactory] is for facilitating testing.
+ * [dialogFactory] is for facilitating testing.
  */
-class PanelConfirmationDialogFactory(
-    private val internalDialogFactory: (Context) -> SystemUIDialog
+class PanelConfirmationDialogFactory @Inject constructor(
+        private val dialogFactory: SystemUIDialog.Factory
 ) {
-    @Inject constructor() : this({ SystemUIDialog(it) })
 
     /**
      * Creates a dialog to show to the user. [response] will be true if an only if the user responds
      * affirmatively.
      */
     fun createConfirmationDialog(
-        context: Context,
-        appName: CharSequence,
-        response: Consumer<Boolean>
+            context: Context,
+            appName: CharSequence,
+            response: Consumer<Boolean>
     ): Dialog {
         val listener =
             DialogInterface.OnClickListener { _, which ->
                 response.accept(which == DialogInterface.BUTTON_POSITIVE)
             }
-        return internalDialogFactory(context).apply {
+        return dialogFactory.create(context).apply {
             setTitle(this.context.getString(R.string.controls_panel_authorization_title, appName))
             setMessage(this.context.getString(R.string.controls_panel_authorization, appName))
             setCanceledOnTouchOutside(true)
diff --git a/packages/SystemUI/src/com/android/systemui/controls/ui/ControlsDialogsFactory.kt b/packages/SystemUI/src/com/android/systemui/controls/ui/ControlsDialogsFactory.kt
index 2ad6014fd7cd..e42a4a6af0de 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/ui/ControlsDialogsFactory.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/ui/ControlsDialogsFactory.kt
@@ -25,20 +25,21 @@ import com.android.systemui.statusbar.phone.SystemUIDialog
 import java.util.function.Consumer
 import javax.inject.Inject
 
-class ControlsDialogsFactory(private val internalDialogFactory: (Context) -> SystemUIDialog) {
+class ControlsDialogsFactory @Inject constructor(
+        private val dialogFactory: SystemUIDialog.Factory
+) {
 
-    @Inject constructor() : this({ SystemUIDialog(it) })
 
     fun createRemoveAppDialog(
-        context: Context,
-        appName: CharSequence,
-        response: Consumer<Boolean>
+            context: Context,
+            appName: CharSequence,
+            response: Consumer<Boolean>
     ): Dialog {
         val listener =
             DialogInterface.OnClickListener { _, which ->
                 response.accept(which == DialogInterface.BUTTON_POSITIVE)
             }
-        return internalDialogFactory(context).apply {
+        return dialogFactory.create(context).apply {
             setTitle(context.getString(R.string.controls_panel_remove_app_authorization, appName))
             setCanceledOnTouchOutside(true)
             setOnCancelListener { response.accept(false) }
diff --git a/packages/SystemUI/src/com/android/systemui/power/PowerNotificationWarnings.java b/packages/SystemUI/src/com/android/systemui/power/PowerNotificationWarnings.java
index a6c623391bb0..7e06f5a21113 100644
--- a/packages/SystemUI/src/com/android/systemui/power/PowerNotificationWarnings.java
+++ b/packages/SystemUI/src/com/android/systemui/power/PowerNotificationWarnings.java
@@ -87,6 +87,7 @@ import java.util.Locale;
 import java.util.Objects;
 
 import javax.inject.Inject;
+import javax.inject.Provider;
 
 /**
  */
@@ -149,6 +150,7 @@ public class PowerNotificationWarnings implements PowerUI.WarningsUI {
     public static final String EXTRA_CONFIRM_ONLY = "extra_confirm_only";
 
     private final Context mContext;
+    private final SystemUIDialog.Factory mSystemUIDialogFactory;
     private final NotificationManager mNoMan;
     private final PowerManager mPowerMan;
     private final KeyguardManager mKeyguard;
@@ -186,11 +188,17 @@ public class PowerNotificationWarnings implements PowerUI.WarningsUI {
     /**
      */
     @Inject
-    public PowerNotificationWarnings(Context context, ActivityStarter activityStarter,
-            BroadcastSender broadcastSender, Lazy<BatteryController> batteryControllerLazy,
-            DialogLaunchAnimator dialogLaunchAnimator, UiEventLogger uiEventLogger,
-            GlobalSettings globalSettings, UserTracker userTracker) {
+    public PowerNotificationWarnings(
+            Context context,
+            ActivityStarter activityStarter,
+            BroadcastSender broadcastSender,
+            Lazy<BatteryController> batteryControllerLazy,
+            DialogLaunchAnimator dialogLaunchAnimator,
+            UiEventLogger uiEventLogger,
+            UserTracker userTracker,
+            SystemUIDialog.Factory systemUIDialogFactory) {
         mContext = context;
+        mSystemUIDialogFactory = systemUIDialogFactory;
         mNoMan = mContext.getSystemService(NotificationManager.class);
         mPowerMan = (PowerManager) context.getSystemService(Context.POWER_SERVICE);
         mKeyguard = mContext.getSystemService(KeyguardManager.class);
@@ -444,7 +452,7 @@ public class PowerNotificationWarnings implements PowerUI.WarningsUI {
 
     private void showHighTemperatureDialog() {
         if (mHighTempDialog != null) return;
-        final SystemUIDialog d = new SystemUIDialog(mContext);
+        final SystemUIDialog d = mSystemUIDialogFactory.create();
         d.setIconAttribute(android.R.attr.alertDialogIcon);
         d.setTitle(R.string.high_temp_title);
         d.setMessage(R.string.high_temp_dialog_message);
@@ -479,7 +487,7 @@ public class PowerNotificationWarnings implements PowerUI.WarningsUI {
 
     private void showThermalShutdownDialog() {
         if (mThermalShutdownDialog != null) return;
-        final SystemUIDialog d = new SystemUIDialog(mContext);
+        final SystemUIDialog d = mSystemUIDialogFactory.create();
         d.setIconAttribute(android.R.attr.alertDialogIcon);
         d.setTitle(R.string.thermal_shutdown_title);
         d.setMessage(R.string.thermal_shutdown_dialog_message);
@@ -643,7 +651,7 @@ public class PowerNotificationWarnings implements PowerUI.WarningsUI {
 
     private void showStartSaverConfirmation(Bundle extras) {
         if (mSaverConfirmation != null || mUseExtraSaverConfirmation) return;
-        final SystemUIDialog d = new SystemUIDialog(mContext);
+        final SystemUIDialog d = mSystemUIDialogFactory.create();
         final boolean confirmOnly = extras.getBoolean(BatterySaverUtils.EXTRA_CONFIRM_TEXT_ONLY);
         final int batterySaverTriggerMode =
                 extras.getInt(BatterySaverUtils.EXTRA_POWER_SAVE_MODE_TRIGGER,
diff --git a/packages/SystemUI/src/com/android/systemui/qs/FgsManagerController.kt b/packages/SystemUI/src/com/android/systemui/qs/FgsManagerController.kt
index 6f35cfbfb4a5..b5def41fb3c7 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/FgsManagerController.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/FgsManagerController.kt
@@ -148,7 +148,8 @@ class FgsManagerControllerImpl @Inject constructor(
     private val deviceConfigProxy: DeviceConfigProxy,
     private val dialogLaunchAnimator: DialogLaunchAnimator,
     private val broadcastDispatcher: BroadcastDispatcher,
-    private val dumpManager: DumpManager
+    private val dumpManager: DumpManager,
+    private val systemUIDialogFactory: SystemUIDialog.Factory,
 ) : Dumpable, FgsManagerController {
 
     companion object {
@@ -375,7 +376,7 @@ class FgsManagerControllerImpl @Inject constructor(
     override fun showDialog(expandable: Expandable?) {
         synchronized(lock) {
             if (dialog == null) {
-                val dialog = SystemUIDialog(context)
+                val dialog = systemUIDialogFactory.create()
                 dialog.setTitle(R.string.fgs_manager_dialog_title)
                 dialog.setMessage(R.string.fgs_manager_dialog_message)
 
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/DataSaverTile.java b/packages/SystemUI/src/com/android/systemui/qs/tiles/DataSaverTile.java
index ccf7afbe7016..c9b002209fa8 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/DataSaverTile.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/DataSaverTile.java
@@ -55,6 +55,7 @@ public class DataSaverTile extends QSTileImpl<BooleanState> implements
 
     private final DataSaverController mDataSaverController;
     private final DialogLaunchAnimator mDialogLaunchAnimator;
+    private final SystemUIDialog.Factory mSystemUIDialogFactory;
 
     @Inject
     public DataSaverTile(
@@ -68,12 +69,14 @@ public class DataSaverTile extends QSTileImpl<BooleanState> implements
             ActivityStarter activityStarter,
             QSLogger qsLogger,
             DataSaverController dataSaverController,
-            DialogLaunchAnimator dialogLaunchAnimator
+            DialogLaunchAnimator dialogLaunchAnimator,
+            SystemUIDialog.Factory systemUIDialogFactory
     ) {
         super(host, uiEventLogger, backgroundLooper, mainHandler, falsingManager, metricsLogger,
                 statusBarStateController, activityStarter, qsLogger);
         mDataSaverController = dataSaverController;
         mDialogLaunchAnimator = dialogLaunchAnimator;
+        mSystemUIDialogFactory = systemUIDialogFactory;
         mDataSaverController.observe(getLifecycle(), this);
     }
 
@@ -98,7 +101,7 @@ public class DataSaverTile extends QSTileImpl<BooleanState> implements
         // Show a dialog to confirm first. Dialogs shown by the DialogLaunchAnimator must be created
         // and shown on the main thread, so we post it to the UI handler.
         mUiHandler.post(() -> {
-            SystemUIDialog dialog = new SystemUIDialog(mContext);
+            SystemUIDialog dialog = mSystemUIDialogFactory.create();
             dialog.setTitle(com.android.internal.R.string.data_saver_enable_title);
             dialog.setMessage(com.android.internal.R.string.data_saver_description);
             dialog.setPositiveButton(com.android.internal.R.string.data_saver_enable_button,
diff --git a/packages/SystemUI/src/com/android/systemui/qs/user/UserSwitchDialogController.kt b/packages/SystemUI/src/com/android/systemui/qs/user/UserSwitchDialogController.kt
index acd7510a6c2a..41cd221186fe 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/user/UserSwitchDialogController.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/user/UserSwitchDialogController.kt
@@ -23,7 +23,6 @@ import android.content.DialogInterface.BUTTON_NEUTRAL
 import android.content.Intent
 import android.provider.Settings
 import android.view.LayoutInflater
-import androidx.annotation.VisibleForTesting
 import com.android.internal.jank.InteractionJankMonitor
 import com.android.internal.logging.UiEventLogger
 import com.android.systemui.res.R
@@ -44,31 +43,15 @@ import javax.inject.Provider
  * Controller for [UserDialog].
  */
 @SysUISingleton
-class UserSwitchDialogController @VisibleForTesting constructor(
-    private val userDetailViewAdapterProvider: Provider<UserDetailView.Adapter>,
-    private val activityStarter: ActivityStarter,
-    private val falsingManager: FalsingManager,
-    private val dialogLaunchAnimator: DialogLaunchAnimator,
-    private val uiEventLogger: UiEventLogger,
-    private val dialogFactory: (Context) -> SystemUIDialog
+class UserSwitchDialogController @Inject constructor(
+        private val userDetailViewAdapterProvider: Provider<UserDetailView.Adapter>,
+        private val activityStarter: ActivityStarter,
+        private val falsingManager: FalsingManager,
+        private val dialogLaunchAnimator: DialogLaunchAnimator,
+        private val uiEventLogger: UiEventLogger,
+        private val dialogFactory: SystemUIDialog.Factory
 ) {
 
-    @Inject
-    constructor(
-        userDetailViewAdapterProvider: Provider<UserDetailView.Adapter>,
-        activityStarter: ActivityStarter,
-        falsingManager: FalsingManager,
-        dialogLaunchAnimator: DialogLaunchAnimator,
-        uiEventLogger: UiEventLogger
-    ) : this(
-        userDetailViewAdapterProvider,
-        activityStarter,
-        falsingManager,
-        dialogLaunchAnimator,
-        uiEventLogger,
-        { SystemUIDialog(it) }
-    )
-
     companion object {
         private const val INTERACTION_JANK_TAG = "switch_user"
         private val USER_SETTINGS_INTENT = Intent(Settings.ACTION_USER_SETTINGS)
@@ -81,7 +64,7 @@ class UserSwitchDialogController @VisibleForTesting constructor(
      * [userDetailViewAdapterProvider] and show it as launched from [expandable].
      */
     fun showDialog(context: Context, expandable: Expandable) {
-        with(dialogFactory(context)) {
+        with(dialogFactory.create()) {
             setShowForAllUsers(true)
             setCanceledOnTouchOutside(true)
 
diff --git a/packages/SystemUI/src/com/android/systemui/reardisplay/RearDisplayDialogController.java b/packages/SystemUI/src/com/android/systemui/reardisplay/RearDisplayDialogController.java
index f07162377358..9076182def70 100644
--- a/packages/SystemUI/src/com/android/systemui/reardisplay/RearDisplayDialogController.java
+++ b/packages/SystemUI/src/com/android/systemui/reardisplay/RearDisplayDialogController.java
@@ -21,8 +21,10 @@ import android.annotation.SuppressLint;
 import android.annotation.TestApi;
 import android.content.Context;
 import android.content.res.Configuration;
+import android.content.res.Resources;
 import android.hardware.devicestate.DeviceStateManager;
 import android.hardware.devicestate.DeviceStateManagerGlobal;
+import android.view.LayoutInflater;
 import android.view.View;
 import android.view.ViewGroup.LayoutParams;
 import android.widget.LinearLayout;
@@ -72,20 +74,27 @@ public class RearDisplayDialogController implements
     private DeviceStateManager.DeviceStateCallback mDeviceStateManagerCallback =
             new DeviceStateManagerCallback();
 
-    private final Context mContext;
     private final CommandQueue mCommandQueue;
     private final Executor mExecutor;
+    private final Resources mResources;
+    private final LayoutInflater mLayoutInflater;
+    private final SystemUIDialog.Factory mSystemUIDialogFactory;
 
-    @VisibleForTesting
-    SystemUIDialog mRearDisplayEducationDialog;
+    private SystemUIDialog mRearDisplayEducationDialog;
     @Nullable LinearLayout mDialogViewContainer;
 
     @Inject
-    public RearDisplayDialogController(Context context, CommandQueue commandQueue,
-            @Main Executor executor) {
-        mContext = context;
+    public RearDisplayDialogController(
+            CommandQueue commandQueue,
+            @Main Executor executor,
+            @Main Resources resources,
+            LayoutInflater layoutInflater,
+            SystemUIDialog.Factory systemUIDialogFactory) {
         mCommandQueue = commandQueue;
         mExecutor = executor;
+        mResources = resources;
+        mLayoutInflater = layoutInflater;
+        mSystemUIDialogFactory = systemUIDialogFactory;
     }
 
     @Override
@@ -104,8 +113,7 @@ public class RearDisplayDialogController implements
         if (mRearDisplayEducationDialog != null && mRearDisplayEducationDialog.isShowing()
                 && mDialogViewContainer != null) {
             // Refresh the dialog view when configuration is changed.
-            Context dialogContext = mRearDisplayEducationDialog.getContext();
-            View dialogView = createDialogView(dialogContext);
+            View dialogView = createDialogView(mRearDisplayEducationDialog.getContext());
             mDialogViewContainer.removeAllViews();
             mDialogViewContainer.addView(dialogView);
         }
@@ -114,9 +122,7 @@ public class RearDisplayDialogController implements
     private void createAndShowDialog() {
         mServiceNotified = false;
         Context dialogContext = mRearDisplayEducationDialog.getContext();
-
         View dialogView = createDialogView(dialogContext);
-
         mDialogViewContainer = new LinearLayout(dialogContext);
         mDialogViewContainer.setLayoutParams(
                 new LinearLayout.LayoutParams(
@@ -133,11 +139,11 @@ public class RearDisplayDialogController implements
 
     private View createDialogView(Context context) {
         View dialogView;
+        LayoutInflater inflater = mLayoutInflater.cloneInContext(context);
         if (mStartedFolded) {
-            dialogView = View.inflate(context,
-                    R.layout.activity_rear_display_education, null);
+            dialogView = inflater.inflate(R.layout.activity_rear_display_education, null);
         } else {
-            dialogView = View.inflate(context,
+            dialogView = inflater.inflate(
                     R.layout.activity_rear_display_education_opened, null);
         }
         LottieAnimationView animationView = dialogView.findViewById(
@@ -172,9 +178,9 @@ public class RearDisplayDialogController implements
      * Ensures we're not using old values from when the dialog may have been shown previously.
      */
     private void initializeValues(int startingBaseState) {
-        mRearDisplayEducationDialog = new SystemUIDialog(mContext);
+        mRearDisplayEducationDialog = mSystemUIDialogFactory.create();
         if (mFoldedStates == null) {
-            mFoldedStates = mContext.getResources().getIntArray(
+            mFoldedStates = mResources.getIntArray(
                     com.android.internal.R.array.config_foldedDeviceStates);
         }
         mStartedFolded = isFoldedState(startingBaseState);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/SystemUIDialog.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/SystemUIDialog.java
index af6da3fb6e51..3394eacddbd8 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/SystemUIDialog.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/SystemUIDialog.java
@@ -149,6 +149,14 @@ public class SystemUIDialog extends AlertDialog implements ViewRootImpl.ConfigCh
             return create(new DialogDelegate<>(){}, mContext);
         }
 
+        /** Creates a new instance of {@link SystemUIDialog} with no customized behavior.
+         *
+         * When you just need a dialog created with a specific {@link Context}, call this.
+         */
+        public SystemUIDialog create(Context context) {
+            return create(new DialogDelegate<>(){}, context);
+        }
+
         /**
          * Creates a new instance of {@link SystemUIDialog} with {@code delegate} as the {@link
          * Delegate}.
diff --git a/packages/SystemUI/src/com/android/systemui/tuner/TunerServiceImpl.java b/packages/SystemUI/src/com/android/systemui/tuner/TunerServiceImpl.java
index 8087a8755a6e..550a65c01bfc 100644
--- a/packages/SystemUI/src/com/android/systemui/tuner/TunerServiceImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/tuner/TunerServiceImpl.java
@@ -48,6 +48,8 @@ import com.android.systemui.statusbar.phone.StatusBarIconController;
 import com.android.systemui.statusbar.phone.SystemUIDialog;
 import com.android.systemui.util.leak.LeakDetector;
 
+import dagger.Lazy;
+
 import java.util.HashSet;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
@@ -87,6 +89,7 @@ public class TunerServiceImpl extends TunerService {
     // Set of all tunables, used for leak detection.
     private final HashSet<Tunable> mTunables = LeakDetector.ENABLED ? new HashSet<>() : null;
     private final Context mContext;
+    private final Lazy<SystemUIDialog.Factory> mSystemUIDialogFactoryLazy;
     private final LeakDetector mLeakDetector;
     private final DemoModeController mDemoModeController;
 
@@ -104,9 +107,11 @@ public class TunerServiceImpl extends TunerService {
             @Main Handler mainHandler,
             LeakDetector leakDetector,
             DemoModeController demoModeController,
-            UserTracker userTracker) {
+            UserTracker userTracker,
+            Lazy<SystemUIDialog.Factory> systemUIDialogFactoryLazy) {
         super(context);
         mContext = context;
+        mSystemUIDialogFactoryLazy = systemUIDialogFactoryLazy;
         mContentResolver = mContext.getContentResolver();
         mLeakDetector = leakDetector;
         mDemoModeController = demoModeController;
@@ -301,7 +306,7 @@ public class TunerServiceImpl extends TunerService {
 
     @Override
     public void showResetRequest(Runnable onDisabled) {
-        SystemUIDialog dialog = new SystemUIDialog(mContext);
+        SystemUIDialog dialog = mSystemUIDialogFactoryLazy.get().create();
         dialog.setShowForAllUsers(true);
         dialog.setMessage(R.string.remove_from_settings_prompt);
         dialog.setButton(DialogInterface.BUTTON_NEGATIVE, mContext.getString(R.string.cancel),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/accessibility/fontscaling/FontScalingDialogDelegateTest.kt b/packages/SystemUI/tests/src/com/android/systemui/accessibility/fontscaling/FontScalingDialogDelegateTest.kt
index bfb5485e47b7..c52571188256 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/accessibility/fontscaling/FontScalingDialogDelegateTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/accessibility/fontscaling/FontScalingDialogDelegateTest.kt
@@ -120,7 +120,7 @@ class FontScalingDialogDelegateTest : SysuiTestCase() {
             fontScalingDialogDelegate
         )
 
-        whenever(dialogFactory.create(any())).thenReturn(dialog)
+        whenever(dialogFactory.create(any(), any())).thenReturn(dialog)
     }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bluetooth/BroadcastDialogDelegateTest.java b/packages/SystemUI/tests/src/com/android/systemui/bluetooth/BroadcastDialogDelegateTest.java
index 4022d4388ab1..3ff43c6a3787 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bluetooth/BroadcastDialogDelegateTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/bluetooth/BroadcastDialogDelegateTest.java
@@ -28,8 +28,6 @@ import static org.mockito.Mockito.when;
 
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
-import android.view.LayoutInflater;
-import android.view.View;
 import android.widget.Button;
 import android.widget.TextView;
 
@@ -95,7 +93,7 @@ public class BroadcastDialogDelegateTest extends SysuiTestCase {
 
         mFeatureFlags.set(Flags.WM_ENABLE_PREDICTIVE_BACK_QS_DIALOG_ANIM, true);
         when(mSysUiState.setFlag(anyInt(), anyBoolean())).thenReturn(mSysUiState);
-        when(mSystemUIDialogFactory.create(any())).thenReturn(mDialog);
+        when(mSystemUIDialogFactory.create(any(), any())).thenReturn(mDialog);
 
         mBroadcastDialogDelegate = new BroadcastDialogDelegate(
                 mContext,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/contrast/ContrastDialogDelegateTest.kt b/packages/SystemUI/tests/src/com/android/systemui/contrast/ContrastDialogDelegateTest.kt
index 65f68f9df3e1..35ac2ae4ed44 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/contrast/ContrastDialogDelegateTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/contrast/ContrastDialogDelegateTest.kt
@@ -31,6 +31,7 @@ import com.android.systemui.flags.FakeFeatureFlags
 import com.android.systemui.flags.FeatureFlags
 import com.android.systemui.model.SysUiState
 import com.android.systemui.settings.UserTracker
+import com.android.systemui.statusbar.phone.DialogDelegate
 import com.android.systemui.statusbar.phone.SystemUIDialog
 import com.android.systemui.util.concurrency.FakeExecutor
 import com.android.systemui.util.mockito.any
@@ -69,7 +70,8 @@ class ContrastDialogDelegateTest : SysuiTestCase() {
         mDependency.injectTestDependency(SysUiState::class.java, sysuiState)
         mDependency.injectMockDependency(DialogLaunchAnimator::class.java)
         whenever(sysuiState.setFlag(any(), any())).thenReturn(sysuiState)
-        whenever(sysuiDialogFactory.create(any())).thenReturn(sysuiDialog)
+        whenever(sysuiDialogFactory.create(any(SystemUIDialog.Delegate::class.java)))
+            .thenReturn(sysuiDialog)
         whenever(sysuiDialog.layoutInflater).thenReturn(LayoutInflater.from(mContext))
 
         whenever(mockUserTracker.userId).thenReturn(context.userId)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/controls/management/PanelConfirmationDialogFactoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/controls/management/PanelConfirmationDialogFactoryTest.kt
index 4e8f86615522..7f0ea9a7a6d0 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/controls/management/PanelConfirmationDialogFactoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/controls/management/PanelConfirmationDialogFactoryTest.kt
@@ -17,34 +17,48 @@
 
 package com.android.systemui.controls.management
 
+import android.content.Context
 import android.content.DialogInterface
 import android.testing.AndroidTestingRunner
 import androidx.test.filters.SmallTest
-import com.android.systemui.res.R
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.res.R
 import com.android.systemui.statusbar.phone.SystemUIDialog
+import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.argumentCaptor
 import com.android.systemui.util.mockito.capture
-import com.android.systemui.util.mockito.eq
-import com.android.systemui.util.mockito.mock
 import com.google.common.truth.Truth.assertThat
+import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.ArgumentCaptor
+import org.mockito.ArgumentMatchers.eq
+import org.mockito.Mock
 import org.mockito.Mockito.verify
-import org.mockito.Mockito.`when`
+import org.mockito.MockitoAnnotations
+import org.mockito.Mockito.`when` as whenever
 
 @SmallTest
 @RunWith(AndroidTestingRunner::class)
 class PanelConfirmationDialogFactoryTest : SysuiTestCase() {
 
+    @Mock private lateinit var mockDialog : SystemUIDialog
+    @Mock private lateinit var mockDialogFactory : SystemUIDialog.Factory
+    private lateinit var factory : PanelConfirmationDialogFactory
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+
+        whenever(mockDialogFactory.create(any(Context::class.java))).thenReturn(mockDialog)
+        whenever(mockDialog.context).thenReturn(mContext)
+        factory = PanelConfirmationDialogFactory(mockDialogFactory)
+    }
+
     @Test
     fun testDialogHasCorrectInfo() {
-        val mockDialog: SystemUIDialog = mock() { `when`(context).thenReturn(mContext) }
-        val factory = PanelConfirmationDialogFactory { mockDialog }
         val appName = "appName"
 
-        factory.createConfirmationDialog(context, appName) {}
+        factory.createConfirmationDialog(mContext, appName) {}
 
         verify(mockDialog).setCanceledOnTouchOutside(true)
         verify(mockDialog)
@@ -55,12 +69,9 @@ class PanelConfirmationDialogFactoryTest : SysuiTestCase() {
 
     @Test
     fun testDialogPositiveButton() {
-        val mockDialog: SystemUIDialog = mock() { `when`(context).thenReturn(mContext) }
-        val factory = PanelConfirmationDialogFactory { mockDialog }
-
         var response: Boolean? = null
 
-        factory.createConfirmationDialog(context, "") { response = it }
+        factory.createConfirmationDialog(mContext,"") { response = it }
 
         val captor: ArgumentCaptor<DialogInterface.OnClickListener> = argumentCaptor()
         verify(mockDialog).setPositiveButton(eq(R.string.controls_dialog_ok), capture(captor))
@@ -72,12 +83,9 @@ class PanelConfirmationDialogFactoryTest : SysuiTestCase() {
 
     @Test
     fun testDialogNeutralButton() {
-        val mockDialog: SystemUIDialog = mock() { `when`(context).thenReturn(mContext) }
-        val factory = PanelConfirmationDialogFactory { mockDialog }
-
         var response: Boolean? = null
 
-        factory.createConfirmationDialog(context, "") { response = it }
+        factory.createConfirmationDialog(mContext, "") { response = it }
 
         val captor: ArgumentCaptor<DialogInterface.OnClickListener> = argumentCaptor()
         verify(mockDialog).setNeutralButton(eq(R.string.cancel), capture(captor))
@@ -89,12 +97,9 @@ class PanelConfirmationDialogFactoryTest : SysuiTestCase() {
 
     @Test
     fun testDialogCancel() {
-        val mockDialog: SystemUIDialog = mock() { `when`(context).thenReturn(mContext) }
-        val factory = PanelConfirmationDialogFactory { mockDialog }
-
         var response: Boolean? = null
 
-        factory.createConfirmationDialog(context, "") { response = it }
+        factory.createConfirmationDialog(mContext, "") { response = it }
 
         val captor: ArgumentCaptor<DialogInterface.OnCancelListener> = argumentCaptor()
         verify(mockDialog).setOnCancelListener(capture(captor))
diff --git a/packages/SystemUI/tests/src/com/android/systemui/controls/ui/ControlsDialogsFactoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/controls/ui/ControlsDialogsFactoryTest.kt
index 8eebceebe874..38c6a0e236ed 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/controls/ui/ControlsDialogsFactoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/controls/ui/ControlsDialogsFactoryTest.kt
@@ -17,17 +17,23 @@
 
 package com.android.systemui.controls.ui
 
+import android.content.Context
 import android.testing.AndroidTestingRunner
 import androidx.test.filters.SmallTest
 import com.android.systemui.res.R
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.statusbar.phone.SystemUIDialog
 import com.android.systemui.util.FakeSystemUIDialogController
+import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.any
 import org.mockito.Mockito.eq
 import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
 
 @SmallTest
 @RunWith(AndroidTestingRunner::class)
@@ -37,18 +43,24 @@ class ControlsDialogsFactoryTest : SysuiTestCase() {
         const val APP_NAME = "Test App"
     }
 
-    private val fakeDialogController = FakeSystemUIDialogController()
+    @Mock
+    private lateinit var mockDialogFactory : SystemUIDialog.Factory
+
+    private val fakeDialogController = FakeSystemUIDialogController(mContext)
 
     private lateinit var underTest: ControlsDialogsFactory
 
     @Before
     fun setup() {
-        underTest = ControlsDialogsFactory { fakeDialogController.dialog }
+        MockitoAnnotations.initMocks(this)
+        whenever(mockDialogFactory.create(any(Context::class.java)))
+            .thenReturn(fakeDialogController.dialog)
+        underTest = ControlsDialogsFactory(mockDialogFactory)
     }
 
     @Test
     fun testCreatesRemoveAppDialog() {
-        val dialog = underTest.createRemoveAppDialog(context, APP_NAME) {}
+        val dialog = underTest.createRemoveAppDialog(mContext, APP_NAME) {}
 
         verify(dialog)
             .setTitle(
@@ -60,7 +72,7 @@ class ControlsDialogsFactoryTest : SysuiTestCase() {
     @Test
     fun testPositiveClickRemoveAppDialogWorks() {
         var dialogResult: Boolean? = null
-        underTest.createRemoveAppDialog(context, APP_NAME) { dialogResult = it }
+        underTest.createRemoveAppDialog(mContext, APP_NAME) { dialogResult = it }
 
         fakeDialogController.clickPositive()
 
@@ -70,7 +82,7 @@ class ControlsDialogsFactoryTest : SysuiTestCase() {
     @Test
     fun testNeutralClickRemoveAppDialogWorks() {
         var dialogResult: Boolean? = null
-        underTest.createRemoveAppDialog(context, APP_NAME) { dialogResult = it }
+        underTest.createRemoveAppDialog(mContext, APP_NAME) { dialogResult = it }
 
         fakeDialogController.clickNeutral()
 
@@ -80,7 +92,7 @@ class ControlsDialogsFactoryTest : SysuiTestCase() {
     @Test
     fun testCancelRemoveAppDialogWorks() {
         var dialogResult: Boolean? = null
-        underTest.createRemoveAppDialog(context, APP_NAME) { dialogResult = it }
+        underTest.createRemoveAppDialog(mContext, APP_NAME) { dialogResult = it }
 
         fakeDialogController.cancel()
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/controls/ui/ControlsUiControllerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/controls/ui/ControlsUiControllerImplTest.kt
index 11bd9cb240a5..36ae0c740c48 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/controls/ui/ControlsUiControllerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/controls/ui/ControlsUiControllerImplTest.kt
@@ -51,6 +51,7 @@ import com.android.systemui.flags.FeatureFlags
 import com.android.systemui.plugins.ActivityStarter
 import com.android.systemui.res.R
 import com.android.systemui.settings.UserTracker
+import com.android.systemui.statusbar.phone.SystemUIDialog
 import com.android.systemui.statusbar.policy.KeyguardStateController
 import com.android.systemui.util.FakeSystemUIDialogController
 import com.android.systemui.util.concurrency.FakeExecutor
@@ -97,9 +98,10 @@ class ControlsUiControllerImplTest : SysuiTestCase() {
     @Mock lateinit var authorizedPanelsRepository: AuthorizedPanelsRepository
     @Mock lateinit var featureFlags: FeatureFlags
     @Mock lateinit var packageManager: PackageManager
+    @Mock lateinit var systemUIDialogFactory: SystemUIDialog.Factory
 
     private val preferredPanelRepository = FakeSelectedComponentRepository()
-    private val fakeDialogController = FakeSystemUIDialogController()
+    private lateinit var fakeDialogController: FakeSystemUIDialogController
     private val uiExecutor = FakeExecutor(FakeSystemClock())
     private val bgExecutor = FakeExecutor(FakeSystemClock())
 
@@ -114,6 +116,9 @@ class ControlsUiControllerImplTest : SysuiTestCase() {
     fun setup() {
         MockitoAnnotations.initMocks(this)
 
+        fakeDialogController = FakeSystemUIDialogController(mContext)
+        whenever(systemUIDialogFactory.create(any(Context::class.java)))
+            .thenReturn(fakeDialogController.dialog)
         controlsSettingsRepository = FakeControlsSettingsRepository()
 
         // This way, it won't be cloned every time `LayoutInflater.fromContext` is called, but we
@@ -146,10 +151,7 @@ class ControlsUiControllerImplTest : SysuiTestCase() {
                 authorizedPanelsRepository,
                 preferredPanelRepository,
                 featureFlags,
-                ControlsDialogsFactory {
-                    isRemoveAppDialogCreated = true
-                    fakeDialogController.dialog
-                },
+                ControlsDialogsFactory(systemUIDialogFactory),
                 dumpManager,
             )
         `when`(userTracker.userId).thenReturn(0)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/power/PowerNotificationWarningsTest.java b/packages/SystemUI/tests/src/com/android/systemui/power/PowerNotificationWarningsTest.java
index 6248bb1009dc..1a303b08b396 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/power/PowerNotificationWarningsTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/power/PowerNotificationWarningsTest.java
@@ -55,6 +55,7 @@ import com.android.systemui.animation.DialogLaunchAnimator;
 import com.android.systemui.broadcast.BroadcastSender;
 import com.android.systemui.plugins.ActivityStarter;
 import com.android.systemui.settings.UserTracker;
+import com.android.systemui.statusbar.phone.SystemUIDialog;
 import com.android.systemui.statusbar.policy.BatteryController;
 import com.android.systemui.util.NotificationChannels;
 import com.android.systemui.util.settings.FakeGlobalSettings;
@@ -77,7 +78,6 @@ public class PowerNotificationWarningsTest extends SysuiTestCase {
     public static final String FORMATTED_45M = "0h 45m";
     public static final String FORMATTED_HOUR = "1h 0m";
     private final NotificationManager mMockNotificationManager = mock(NotificationManager.class);
-    private final GlobalSettings mGlobalSettings = new FakeGlobalSettings();
     private PowerNotificationWarnings mPowerNotificationWarnings;
 
     @Mock
@@ -90,6 +90,10 @@ public class PowerNotificationWarningsTest extends SysuiTestCase {
     private UserTracker mUserTracker;
     @Mock
     private View mView;
+    @Mock
+    private SystemUIDialog.Factory mSystemUIDialogFactory;
+    @Mock
+    private SystemUIDialog mSystemUIDialog;
 
     private BroadcastReceiver mReceiver;
 
@@ -113,9 +117,16 @@ public class PowerNotificationWarningsTest extends SysuiTestCase {
         when(mUserTracker.getUserId()).thenReturn(ActivityManager.getCurrentUser());
         when(mUserTracker.getUserHandle()).thenReturn(
                 UserHandle.of(ActivityManager.getCurrentUser()));
-        mPowerNotificationWarnings = new PowerNotificationWarnings(wrapper, starter,
-                broadcastSender, () -> mBatteryController, mDialogLaunchAnimator, mUiEventLogger,
-                mGlobalSettings, mUserTracker);
+        when(mSystemUIDialogFactory.create()).thenReturn(mSystemUIDialog);
+        mPowerNotificationWarnings = new PowerNotificationWarnings(
+                wrapper,
+                starter,
+                broadcastSender,
+                () -> mBatteryController,
+                mDialogLaunchAnimator,
+                mUiEventLogger,
+                mUserTracker,
+                mSystemUIDialogFactory);
         BatteryStateSnapshot snapshot = new BatteryStateSnapshot(100, false, false, 1,
                 BatteryManager.BATTERY_HEALTH_GOOD, 5, 15);
         mPowerNotificationWarnings.updateSnapshot(snapshot);
@@ -251,7 +262,7 @@ public class PowerNotificationWarningsTest extends SysuiTestCase {
 
         verify(mDialogLaunchAnimator, never()).showFromView(any(), any());
 
-        assertThat(mPowerNotificationWarnings.getSaverConfirmationDialog().isShowing()).isTrue();
+        verify(mPowerNotificationWarnings.getSaverConfirmationDialog()).show();
         mPowerNotificationWarnings.getSaverConfirmationDialog().dismiss();
     }
 
@@ -266,7 +277,7 @@ public class PowerNotificationWarningsTest extends SysuiTestCase {
 
         verify(mDialogLaunchAnimator, never()).showFromView(any(), any());
 
-        assertThat(mPowerNotificationWarnings.getSaverConfirmationDialog().isShowing()).isTrue();
+        verify(mPowerNotificationWarnings.getSaverConfirmationDialog()).show();
         mPowerNotificationWarnings.getSaverConfirmationDialog().dismiss();
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/FgsManagerControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/qs/FgsManagerControllerTest.java
index f5a3becc7017..698868d67071 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/FgsManagerControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/FgsManagerControllerTest.java
@@ -27,6 +27,7 @@ import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import android.app.IActivityManager;
 import android.app.IForegroundServiceObserver;
@@ -53,6 +54,7 @@ import com.android.systemui.animation.DialogLaunchAnimator;
 import com.android.systemui.broadcast.BroadcastDispatcher;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.settings.UserTracker;
+import com.android.systemui.statusbar.phone.SystemUIDialog;
 import com.android.systemui.util.DeviceConfigProxyFake;
 import com.android.systemui.util.concurrency.FakeExecutor;
 import com.android.systemui.util.time.FakeSystemClock;
@@ -95,6 +97,10 @@ public class FgsManagerControllerTest extends SysuiTestCase {
     BroadcastDispatcher mBroadcastDispatcher;
     @Mock
     DumpManager mDumpManager;
+    @Mock
+    SystemUIDialog.Factory mSystemUIDialogFactory;
+    @Mock
+    SystemUIDialog mSystemUIDialog;
 
     private FgsManagerController mFmc;
 
@@ -114,6 +120,7 @@ public class FgsManagerControllerTest extends SysuiTestCase {
         mSystemClock = new FakeSystemClock();
         mMainExecutor = new FakeExecutor(mSystemClock);
         mBackgroundExecutor = new FakeExecutor(mSystemClock);
+        when(mSystemUIDialogFactory.create()).thenReturn(mSystemUIDialog);
 
         mUserProfiles = new ArrayList<>();
         Mockito.doReturn(mUserProfiles).when(mUserTracker).getUserProfiles();
@@ -325,7 +332,8 @@ public class FgsManagerControllerTest extends SysuiTestCase {
                 mDeviceConfigProxyFake,
                 mDialogLaunchAnimator,
                 mBroadcastDispatcher,
-                mDumpManager
+                mDumpManager,
+                mSystemUIDialogFactory
         );
         fmc.init();
         Assert.assertTrue(fmc.getIncludesUserVisibleJobs());
@@ -351,7 +359,8 @@ public class FgsManagerControllerTest extends SysuiTestCase {
                 mDeviceConfigProxyFake,
                 mDialogLaunchAnimator,
                 mBroadcastDispatcher,
-                mDumpManager
+                mDumpManager,
+                mSystemUIDialogFactory
         );
         fmc.init();
         Assert.assertFalse(fmc.getIncludesUserVisibleJobs());
@@ -457,7 +466,8 @@ public class FgsManagerControllerTest extends SysuiTestCase {
                 mDeviceConfigProxyFake,
                 mDialogLaunchAnimator,
                 mBroadcastDispatcher,
-                mDumpManager
+                mDumpManager,
+                mSystemUIDialogFactory
         );
         result.init();
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/DataSaverTileTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/DataSaverTileTest.kt
index 51e95be3611b..c109a1e95f66 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/DataSaverTileTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/DataSaverTileTest.kt
@@ -32,7 +32,9 @@ import com.android.systemui.qs.QSHost
 import com.android.systemui.qs.QsEventLogger
 import com.android.systemui.qs.logging.QSLogger
 import com.android.systemui.qs.tileimpl.QSTileImpl
+import com.android.systemui.statusbar.phone.SystemUIDialog
 import com.android.systemui.statusbar.policy.DataSaverController
+import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
 import org.junit.Before
@@ -49,8 +51,6 @@ class DataSaverTileTest : SysuiTestCase() {
 
     @Mock private lateinit var mHost: QSHost
     @Mock private lateinit var mMetricsLogger: MetricsLogger
-    @Mock private lateinit var mStatusBarStateController: StatusBarStateController
-    @Mock private lateinit var mActivityStarter: ActivityStarter
     @Mock private lateinit var mQsLogger: QSLogger
     private val falsingManager = FalsingManagerFake()
     @Mock private lateinit var statusBarStateController: StatusBarStateController
@@ -58,6 +58,8 @@ class DataSaverTileTest : SysuiTestCase() {
     @Mock private lateinit var dataSaverController: DataSaverController
     @Mock private lateinit var dialogLaunchAnimator: DialogLaunchAnimator
     @Mock private lateinit var uiEventLogger: QsEventLogger
+    @Mock private lateinit var systemUIDialogFactory: SystemUIDialog.Factory
+    @Mock private lateinit var systemUIDialog: SystemUIDialog
 
     private lateinit var testableLooper: TestableLooper
     private lateinit var tile: DataSaverTile
@@ -67,7 +69,8 @@ class DataSaverTileTest : SysuiTestCase() {
         MockitoAnnotations.initMocks(this)
         testableLooper = TestableLooper.get(this)
 
-        Mockito.`when`(mHost.context).thenReturn(mContext)
+        whenever(mHost.context).thenReturn(mContext)
+        whenever(systemUIDialogFactory.create()).thenReturn(systemUIDialog)
 
         tile =
             DataSaverTile(
@@ -81,7 +84,8 @@ class DataSaverTileTest : SysuiTestCase() {
                 activityStarter,
                 mQsLogger,
                 dataSaverController,
-                dialogLaunchAnimator
+                dialogLaunchAnimator,
+                systemUIDialogFactory
             )
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/user/UserSwitchDialogControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/user/UserSwitchDialogControllerTest.kt
index 0a34810f4d3f..945490f1983d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/user/UserSwitchDialogControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/user/UserSwitchDialogControllerTest.kt
@@ -36,6 +36,7 @@ import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.capture
 import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.whenever
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -56,6 +57,8 @@ import org.mockito.MockitoAnnotations
 class UserSwitchDialogControllerTest : SysuiTestCase() {
 
     @Mock
+    private lateinit var dialogFactory: SystemUIDialog.Factory
+    @Mock
     private lateinit var dialog: SystemUIDialog
     @Mock
     private lateinit var falsingManager: FalsingManager
@@ -80,7 +83,8 @@ class UserSwitchDialogControllerTest : SysuiTestCase() {
     fun setUp() {
         MockitoAnnotations.initMocks(this)
 
-        `when`(dialog.context).thenReturn(mContext)
+        whenever(dialog.context).thenReturn(mContext)
+        whenever(dialogFactory.create()).thenReturn(dialog)
 
         controller = UserSwitchDialogController(
             { userDetailViewAdapter },
@@ -88,7 +92,7 @@ class UserSwitchDialogControllerTest : SysuiTestCase() {
             falsingManager,
             dialogLaunchAnimator,
             uiEventLogger,
-            { dialog }
+            dialogFactory
         )
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/reardisplay/RearDisplayDialogControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/reardisplay/RearDisplayDialogControllerTest.java
index 273ce85f89f5..35bf7753358e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/reardisplay/RearDisplayDialogControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/reardisplay/RearDisplayDialogControllerTest.java
@@ -18,25 +18,42 @@ package com.android.systemui.reardisplay;
 
 import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertNotSame;
-import static junit.framework.Assert.assertTrue;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import android.content.res.Configuration;
+import android.content.res.Resources;
 import android.hardware.devicestate.DeviceStateManager;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
+import android.view.LayoutInflater;
+import android.view.View;
 import android.widget.TextView;
 
 import androidx.test.filters.SmallTest;
 
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.res.R;
+import com.android.systemui.flags.FakeFeatureFlags;
+import com.android.systemui.flags.Flags;
+import com.android.systemui.model.SysUiState;
+import com.android.systemui.res.R;
 import com.android.systemui.statusbar.CommandQueue;
+import com.android.systemui.statusbar.phone.SystemUIDialog;
 import com.android.systemui.util.concurrency.FakeExecutor;
 import com.android.systemui.util.time.FakeSystemClock;
 
+import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
@@ -45,24 +62,49 @@ public class RearDisplayDialogControllerTest extends SysuiTestCase {
 
     @Mock
     private CommandQueue mCommandQueue;
+    @Mock
+    private SystemUIDialog.Factory mSystemUIDialogFactory;
+    @Mock
+    private SystemUIDialog mSystemUIDialog;
+    private final FakeFeatureFlags mFeatureFlags = new FakeFeatureFlags();
+    @Mock
+    private SysUiState mSysUiState;
+    @Mock
+    private Resources mResources;
 
-    private FakeExecutor mFakeExecutor = new FakeExecutor(new FakeSystemClock());
+    LayoutInflater mLayoutInflater = LayoutInflater.from(mContext);
 
+    private final FakeExecutor mFakeExecutor = new FakeExecutor(new FakeSystemClock());
 
     private static final int CLOSED_BASE_STATE = 0;
     private static final int OPEN_BASE_STATE = 1;
 
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+
+        mFeatureFlags.set(Flags.WM_ENABLE_PREDICTIVE_BACK_QS_DIALOG_ANIM, true);
+        when(mSysUiState.setFlag(anyInt(), anyBoolean())).thenReturn(mSysUiState);
+        when(mSystemUIDialogFactory.create()).thenReturn(mSystemUIDialog);
+        when(mSystemUIDialog.getContext()).thenReturn(mContext);
+    }
     @Test
     public void testClosedDialogIsShown() {
-        RearDisplayDialogController controller = new RearDisplayDialogController(mContext,
-                mCommandQueue, mFakeExecutor);
+        RearDisplayDialogController controller = new RearDisplayDialogController(
+                mCommandQueue,
+                mFakeExecutor,
+                mResources,
+                mLayoutInflater,
+                mSystemUIDialogFactory);
         controller.setDeviceStateManagerCallback(new TestDeviceStateManagerCallback());
         controller.setFoldedStates(new int[]{0});
         controller.setAnimationRepeatCount(0);
 
         controller.showRearDisplayDialog(CLOSED_BASE_STATE);
-        assertTrue(controller.mRearDisplayEducationDialog.isShowing());
-        TextView deviceClosedTitleTextView = controller.mRearDisplayEducationDialog.findViewById(
+        verify(mSystemUIDialog).show();
+
+        View container = getDialogViewContainer();
+        TextView deviceClosedTitleTextView = container.findViewById(
                 R.id.rear_display_title_text_view);
         assertEquals(deviceClosedTitleTextView.getText().toString(),
                 getContext().getResources().getString(
@@ -71,20 +113,28 @@ public class RearDisplayDialogControllerTest extends SysuiTestCase {
 
     @Test
     public void testClosedDialogIsRefreshedOnConfigurationChange() {
-        RearDisplayDialogController controller = new RearDisplayDialogController(mContext,
-                mCommandQueue, mFakeExecutor);
+        RearDisplayDialogController controller = new RearDisplayDialogController(
+                mCommandQueue,
+                mFakeExecutor,
+                mResources,
+                mLayoutInflater,
+                mSystemUIDialogFactory);
         controller.setDeviceStateManagerCallback(new TestDeviceStateManagerCallback());
         controller.setFoldedStates(new int[]{0});
         controller.setAnimationRepeatCount(0);
 
         controller.showRearDisplayDialog(CLOSED_BASE_STATE);
-        assertTrue(controller.mRearDisplayEducationDialog.isShowing());
-        TextView deviceClosedTitleTextView = controller.mRearDisplayEducationDialog.findViewById(
+        verify(mSystemUIDialog).show();
+        View container = getDialogViewContainer();
+        TextView deviceClosedTitleTextView = container.findViewById(
                 R.id.rear_display_title_text_view);
 
+        reset(mSystemUIDialog);
+        when(mSystemUIDialog.isShowing()).thenReturn(true);
+        when(mSystemUIDialog.getContext()).thenReturn(mContext);
+
         controller.onConfigChanged(new Configuration());
-        assertTrue(controller.mRearDisplayEducationDialog.isShowing());
-        TextView deviceClosedTitleTextView2 = controller.mRearDisplayEducationDialog.findViewById(
+        TextView deviceClosedTitleTextView2 = container.findViewById(
                 R.id.rear_display_title_text_view);
 
         assertNotSame(deviceClosedTitleTextView, deviceClosedTitleTextView2);
@@ -92,22 +142,33 @@ public class RearDisplayDialogControllerTest extends SysuiTestCase {
 
     @Test
     public void testOpenDialogIsShown() {
-        RearDisplayDialogController controller = new RearDisplayDialogController(mContext,
-                mCommandQueue, mFakeExecutor);
+        RearDisplayDialogController controller = new RearDisplayDialogController(
+                mCommandQueue,
+                mFakeExecutor,
+                mResources,
+                mLayoutInflater,
+                mSystemUIDialogFactory);
         controller.setDeviceStateManagerCallback(new TestDeviceStateManagerCallback());
         controller.setFoldedStates(new int[]{0});
         controller.setAnimationRepeatCount(0);
 
         controller.showRearDisplayDialog(OPEN_BASE_STATE);
 
-        assertTrue(controller.mRearDisplayEducationDialog.isShowing());
-        TextView deviceClosedTitleTextView = controller.mRearDisplayEducationDialog.findViewById(
+        verify(mSystemUIDialog).show();
+        View container = getDialogViewContainer();
+        TextView deviceClosedTitleTextView = container.findViewById(
                 R.id.rear_display_title_text_view);
         assertEquals(deviceClosedTitleTextView.getText().toString(),
                 getContext().getResources().getString(
                         R.string.rear_display_unfolded_bottom_sheet_title));
     }
 
+    private View getDialogViewContainer() {
+        ArgumentCaptor<View> viewCaptor = ArgumentCaptor.forClass(View.class);
+        verify(mSystemUIDialog).setView(viewCaptor.capture());
+
+        return viewCaptor.getValue();
+    }
     /**
      * Empty device state manager callbacks, so we can verify that the correct
      * dialogs are being created regardless of device state of the test device.
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/util/FakeSystemUIDialogController.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/util/FakeSystemUIDialogController.kt
index 0c9ce0f145f1..697b5087a865 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/util/FakeSystemUIDialogController.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/util/FakeSystemUIDialogController.kt
@@ -17,6 +17,7 @@
 
 package com.android.systemui.util
 
+import android.content.Context
 import android.content.DialogInterface
 import com.android.systemui.statusbar.phone.SystemUIDialog
 import com.android.systemui.util.mockito.any
@@ -27,13 +28,15 @@ import org.mockito.Mockito.doAnswer
 import org.mockito.Mockito.verify
 import org.mockito.stubbing.Stubber
 
-class FakeSystemUIDialogController {
+class FakeSystemUIDialogController(context: Context) {
 
     val dialog: SystemUIDialog = mock()
 
+
     private val clickListeners: MutableMap<Int, DialogInterface.OnClickListener> = mutableMapOf()
 
     init {
+        whenever(dialog.context).thenReturn(context)
         saveListener(DialogInterface.BUTTON_POSITIVE)
             .whenever(dialog)
             .setPositiveButton(any(), any())