diff --git a/AndroidManifest.xml b/AndroidManifest.xml
index 8cfd9b5..b28da4f 100644
--- a/AndroidManifest.xml
+++ b/AndroidManifest.xml
@@ -1427,6 +1427,19 @@
                        android:value="true" />
         </activity>
 
+        <activity
+            android:name=".Settings$ManageAdaptiveNotificationsActivity"
+            android:exported="true">
+            <intent-filter>
+                <action android:name="android.settings.MANAGE_ADAPTIVE_NOTIFICATIONS"/>
+                <category android:name="android.intent.category.DEFAULT"/>
+            </intent-filter>
+            <meta-data android:name="com.android.settings.FRAGMENT_CLASS"
+                android:value="com.android.settings.notification.PoliteNotificationsPreferenceFragment" />
+            <meta-data android:name="com.android.settings.HIGHLIGHT_MENU_KEY"
+                android:value="@string/menu_key_notifications"/>
+        </activity>
+
         <activity android:name="Settings$WallpaperSettingsActivity"
                   android:label="@string/wallpaper_settings_fragment_title"
                   android:icon="@drawable/ic_wallpaper"
diff --git a/proguard.flags b/proguard.flags
index 492404c..d1c97f3 100644
--- a/proguard.flags
+++ b/proguard.flags
@@ -65,3 +65,6 @@
 -keep class androidx.window.extensions.** { *; }
 -dontwarn androidx.window.extensions.**
 -keep class androidx.window.** { *; }
+
+# Keep the com.android.settings.media_drm.FakeFeatureFlagsImpl
+-keep class com.android.settings.media_drm.FakeFeatureFlagsImpl { *; }
diff --git a/res/values/strings.xml b/res/values/strings.xml
index a490856..d610e1f 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -12383,7 +12383,7 @@
     <!-- Toast on failure to reformat data to ext4 -->
     <string name="format_ext4_failure_toast">Failed to reformat and wipe the data partition to ext4.</string>
     <!-- Dialog to OEM unlock the device before using 16K developer option -->
-    <string name="confirm_oem_unlock_for_16k_title">Bootloader Unlock Required for 16KB Mode</string>
+    <string name="confirm_oem_unlock_for_16k_title">Bootloader Unlock Required</string>
     <string name="confirm_oem_unlock_for_16k_text">This device needs to have the bootloader unlocked before using the 16KB developer option.
         Software integrity cannot be guaranteed in this mode, and any data stored on the phone while the bootloader is unlocked may be at risk.
         All user data and settings will be wiped when activating 16KB mode. Once the bootloader is unlocked, activating the 16KB option will require two reboots.
@@ -12412,10 +12412,10 @@
     <!-- persistent notification 16k page agnostic mode text -->
     <string name="page_agnostic_16k_pages_text_short">You are in the 16KB mode of the page-agnostic mode. Software integrity cannot be guaranteed in this mode,
         and any data stored on the phone while the bootloader is unlocked may be at risk. Some features will be disabled in these modes, so some applications may not work.
-        In order to re-enter the production mode, you must, switch back to 4K mode and then lock the bootloader of the device. Tap to read more.</string>
+        In order to re-enter the production mode, you must, switch back to 4KB mode and then lock the bootloader of the device. Tap to read more.</string>
     <string name="page_agnostic_16k_pages_text">You are in the 16KB mode of the page-agnostic mode. Software integrity cannot be guaranteed in this mode,
         and any data stored on the phone while the bootloader is unlocked may be at risk. Some features will be disabled in these modes, so some applications may not work.
-        In order to re-enter the production mode, you must, switch back to 4K mode and then lock the bootloader of the device. This would factory reset the device again and
+        In order to re-enter the production mode, you must, switch back to 4KB mode and then lock the bootloader of the device. This would factory reset the device again and
         restore it to production settings. After the device successfully boots into Android, disable OEM unlocking in Developer options.
         If the device fails to boot into Android or is unstable, re-flash the device with the latest factory images from
         &lt;a href=\"https://developers.google.com/android/images\"&gt;https://developers.google.com/android/images&lt;/a&gt;
diff --git a/res/xml/modes_rule_settings.xml b/res/xml/modes_rule_settings.xml
index 0df9f80..5be206e 100644
--- a/res/xml/modes_rule_settings.xml
+++ b/res/xml/modes_rule_settings.xml
@@ -67,5 +67,9 @@
         <Preference
                 android:key="mode_display_settings"
                 android:title="@string/mode_display_settings_title" />
+
+        <Preference
+                android:key="mode_manual_duration"
+                android:title="@string/zen_category_duration" />
     </PreferenceCategory>
 </PreferenceScreen>
\ No newline at end of file
diff --git a/src/com/android/settings/Settings.java b/src/com/android/settings/Settings.java
index e3bb1a1..24d9525 100644
--- a/src/com/android/settings/Settings.java
+++ b/src/com/android/settings/Settings.java
@@ -303,6 +303,7 @@
     public static class UserSettingsActivity extends SettingsActivity { /* empty */ }
     public static class NotificationAccessSettingsActivity extends SettingsActivity { /* empty */ }
     public static class NotificationAccessDetailsActivity extends SettingsActivity { /* empty */ }
+    public static class ManageAdaptiveNotificationsActivity extends SettingsActivity { /* empty */ }
     public static class VrListenersSettingsActivity extends SettingsActivity { /* empty */ }
     public static class PremiumSmsAccessActivity extends SettingsActivity { /* empty */ }
     public static class PictureInPictureSettingsActivity extends SettingsActivity { /* empty */ }
diff --git a/src/com/android/settings/SettingsApplication.java b/src/com/android/settings/SettingsApplication.java
index 7d5d2c7..5b052f2 100644
--- a/src/com/android/settings/SettingsApplication.java
+++ b/src/com/android/settings/SettingsApplication.java
@@ -31,6 +31,7 @@
 import com.android.settings.biometrics.fingerprint2.BiometricsEnvironment;
 import com.android.settings.core.instrumentation.ElapsedTimeUtils;
 import com.android.settings.development.DeveloperOptionsActivityLifecycle;
+import com.android.settings.flags.Flags;
 import com.android.settings.fuelgauge.BatterySettingsStorage;
 import com.android.settings.homepage.SettingsHomepageActivity;
 import com.android.settings.localepicker.LocaleNotificationDataManager;
@@ -49,7 +50,7 @@
 public class SettingsApplication extends Application {
 
     private WeakReference<SettingsHomepageActivity> mHomeActivity = new WeakReference<>(null);
-    private BiometricsEnvironment mBiometricsEnvironment;
+    @Nullable private BiometricsEnvironment mBiometricsEnvironment;
 
     @Override
     protected void attachBaseContext(Context base) {
@@ -73,7 +74,9 @@
 
         // Set Spa environment.
         setSpaEnvironment();
-        mBiometricsEnvironment = new BiometricsEnvironment(this);
+        if (Flags.fingerprintV2Enrollment()) {
+            mBiometricsEnvironment = new BiometricsEnvironment(this);
+        }
 
         if (ActivityEmbeddingUtils.isSettingsSplitEnabled(this)
                 && FeatureFlagUtils.isEnabled(this,
diff --git a/src/com/android/settings/connecteddevice/audiosharing/AudioSharingCallAudioDialogFragment.java b/src/com/android/settings/connecteddevice/audiosharing/AudioSharingCallAudioDialogFragment.java
index 75b4acf..5e0ec07 100644
--- a/src/com/android/settings/connecteddevice/audiosharing/AudioSharingCallAudioDialogFragment.java
+++ b/src/com/android/settings/connecteddevice/audiosharing/AudioSharingCallAudioDialogFragment.java
@@ -19,6 +19,7 @@
 import android.app.Dialog;
 import android.app.settings.SettingsEnums;
 import android.os.Bundle;
+import android.util.Log;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -66,7 +67,13 @@
             @NonNull List<AudioSharingDeviceItem> deviceItems,
             @NonNull DialogEventListener listener) {
         if (!AudioSharingUtils.isFeatureEnabled()) return;
-        final FragmentManager manager = host.getChildFragmentManager();
+        final FragmentManager manager;
+        try {
+            manager = host.getChildFragmentManager();
+        } catch (IllegalStateException e) {
+            Log.d(TAG, "Fail to show dialog: " + e.getMessage());
+            return;
+        }
         sListener = listener;
         if (manager.findFragmentByTag(TAG) == null) {
             final Bundle bundle = new Bundle();
@@ -79,10 +86,18 @@
     }
 
     @Override
+    @NonNull
     public Dialog onCreateDialog(Bundle savedInstanceState) {
         Bundle arguments = requireArguments();
         List<AudioSharingDeviceItem> deviceItems =
                 arguments.getParcelable(BUNDLE_KEY_DEVICE_ITEMS, List.class);
+        AlertDialog.Builder builder =
+                new AlertDialog.Builder(getActivity())
+                        .setTitle(R.string.audio_sharing_call_audio_title);
+        if (deviceItems == null) {
+            Log.d(TAG, "Create dialog error: null deviceItems");
+            return builder.create();
+        }
         int checkedItem = -1;
         for (AudioSharingDeviceItem item : deviceItems) {
             int fallbackActiveGroupId = AudioSharingUtils.getFallbackActiveGroupId(getContext());
@@ -92,17 +107,14 @@
         }
         String[] choices =
                 deviceItems.stream().map(AudioSharingDeviceItem::getName).toArray(String[]::new);
-        AlertDialog.Builder builder =
-                new AlertDialog.Builder(getActivity())
-                        .setTitle(R.string.audio_sharing_call_audio_title)
-                        .setSingleChoiceItems(
-                                choices,
-                                checkedItem,
-                                (dialog, which) -> {
-                                    if (sListener != null) {
-                                        sListener.onItemClick(deviceItems.get(which));
-                                    }
-                                });
+        builder.setSingleChoiceItems(
+                choices,
+                checkedItem,
+                (dialog, which) -> {
+                    if (sListener != null) {
+                        sListener.onItemClick(deviceItems.get(which));
+                    }
+                });
         return builder.create();
     }
 }
diff --git a/src/com/android/settings/connecteddevice/audiosharing/AudioSharingConfirmDialogFragment.java b/src/com/android/settings/connecteddevice/audiosharing/AudioSharingConfirmDialogFragment.java
index e9013d7..61b1df1 100644
--- a/src/com/android/settings/connecteddevice/audiosharing/AudioSharingConfirmDialogFragment.java
+++ b/src/com/android/settings/connecteddevice/audiosharing/AudioSharingConfirmDialogFragment.java
@@ -21,6 +21,7 @@
 import android.os.Bundle;
 import android.util.Log;
 
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.appcompat.app.AlertDialog;
 import androidx.fragment.app.Fragment;
@@ -44,7 +45,13 @@
      */
     public static void show(Fragment host) {
         if (!AudioSharingUtils.isFeatureEnabled()) return;
-        FragmentManager manager = host.getChildFragmentManager();
+        final FragmentManager manager;
+        try {
+            manager = host.getChildFragmentManager();
+        } catch (IllegalStateException e) {
+            Log.d(TAG, "Fail to show dialog: " + e.getMessage());
+            return;
+        }
         AlertDialog dialog = AudioSharingDialogHelper.getDialogIfShowing(manager, TAG);
         if (dialog != null) {
             Log.d(TAG, "Dialog is showing, return.");
@@ -56,6 +63,7 @@
     }
 
     @Override
+    @NonNull
     public Dialog onCreateDialog(@Nullable Bundle savedInstanceState) {
         AlertDialog dialog =
                 AudioSharingDialogFactory.newBuilder(getActivity())
diff --git a/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogFragment.java b/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogFragment.java
index e787be3..f00cf73 100644
--- a/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogFragment.java
+++ b/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogFragment.java
@@ -77,7 +77,13 @@
             @NonNull DialogEventListener listener,
             @NonNull Pair<Integer, Object>[] eventData) {
         if (!AudioSharingUtils.isFeatureEnabled()) return;
-        final FragmentManager manager = host.getChildFragmentManager();
+        final FragmentManager manager;
+        try {
+            manager = host.getChildFragmentManager();
+        } catch (IllegalStateException e) {
+            Log.d(TAG, "Fail to show dialog: " + e.getMessage());
+            return;
+        }
         sListener = listener;
         sEventData = eventData;
         AlertDialog dialog = AudioSharingDialogHelper.getDialogIfShowing(manager, TAG);
diff --git a/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogHandler.java b/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogHandler.java
index 753daaf..81d7979 100644
--- a/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogHandler.java
+++ b/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogHandler.java
@@ -414,7 +414,13 @@
 
     private void closeOpeningDialogsOtherThan(String tag) {
         if (mHostFragment == null) return;
-        List<Fragment> fragments = mHostFragment.getChildFragmentManager().getFragments();
+        List<Fragment> fragments;
+        try {
+            fragments = mHostFragment.getChildFragmentManager().getFragments();
+        } catch (IllegalStateException e) {
+            Log.d(TAG, "Fail to closeOpeningDialogsOtherThan " + tag + ": " + e.getMessage());
+            return;
+        }
         for (Fragment fragment : fragments) {
             if (fragment instanceof DialogFragment
                     && fragment.getTag() != null
@@ -430,7 +436,13 @@
     public void closeOpeningDialogsForLeaDevice(@NonNull CachedBluetoothDevice cachedDevice) {
         if (mHostFragment == null) return;
         int groupId = AudioSharingUtils.getGroupId(cachedDevice);
-        List<Fragment> fragments = mHostFragment.getChildFragmentManager().getFragments();
+        List<Fragment> fragments;
+        try {
+            fragments = mHostFragment.getChildFragmentManager().getFragments();
+        } catch (IllegalStateException e) {
+            Log.d(TAG, "Fail to closeOpeningDialogsForLeaDevice: " + e.getMessage());
+            return;
+        }
         for (Fragment fragment : fragments) {
             CachedBluetoothDevice device = getCachedBluetoothDeviceFromDialog(fragment);
             if (device != null
@@ -447,7 +459,13 @@
     public void closeOpeningDialogsForNonLeaDevice(@NonNull CachedBluetoothDevice cachedDevice) {
         if (mHostFragment == null) return;
         String address = cachedDevice.getAddress();
-        List<Fragment> fragments = mHostFragment.getChildFragmentManager().getFragments();
+        List<Fragment> fragments;
+        try {
+            fragments = mHostFragment.getChildFragmentManager().getFragments();
+        } catch (IllegalStateException e) {
+            Log.d(TAG, "Fail to closeOpeningDialogsForNonLeaDevice: " + e.getMessage());
+            return;
+        }
         for (Fragment fragment : fragments) {
             CachedBluetoothDevice device = getCachedBluetoothDeviceFromDialog(fragment);
             if (device != null && address != null && address.equals(device.getAddress())) {
diff --git a/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogHelper.java b/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogHelper.java
index 69001aa..010a3ba 100644
--- a/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogHelper.java
+++ b/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogHelper.java
@@ -51,12 +51,13 @@
     public static AlertDialog getDialogIfShowing(
             @NonNull FragmentManager manager, @NonNull String tag) {
         Fragment dialog = manager.findFragmentByTag(tag);
-        return dialog != null
-                        && dialog instanceof DialogFragment
-                        && ((DialogFragment) dialog).getDialog() != null
-                        && ((DialogFragment) dialog).getDialog().isShowing()
-                        && ((DialogFragment) dialog).getDialog() instanceof AlertDialog
+        return dialog instanceof DialogFragment
+                && ((DialogFragment) dialog).getDialog() != null
+                && ((DialogFragment) dialog).getDialog().isShowing()
+                && ((DialogFragment) dialog).getDialog() instanceof AlertDialog
                 ? (AlertDialog) ((DialogFragment) dialog).getDialog()
                 : null;
     }
+
+    private AudioSharingDialogHelper() {}
 }
diff --git a/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDisconnectDialogFragment.java b/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDisconnectDialogFragment.java
index dcd8a3b..66e327b 100644
--- a/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDisconnectDialogFragment.java
+++ b/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDisconnectDialogFragment.java
@@ -84,7 +84,13 @@
             @NonNull DialogEventListener listener,
             @NonNull Pair<Integer, Object>[] eventData) {
         if (!AudioSharingUtils.isFeatureEnabled()) return;
-        FragmentManager manager = host.getChildFragmentManager();
+        final FragmentManager manager;
+        try {
+            manager = host.getChildFragmentManager();
+        } catch (IllegalStateException e) {
+            Log.d(TAG, "Fail to show dialog: " + e.getMessage());
+            return;
+        }
         AlertDialog dialog = AudioSharingDialogHelper.getDialogIfShowing(manager, TAG);
         if (dialog != null) {
             int newGroupId = AudioSharingUtils.getGroupId(newDevice);
diff --git a/src/com/android/settings/connecteddevice/audiosharing/AudioSharingJoinDialogFragment.java b/src/com/android/settings/connecteddevice/audiosharing/AudioSharingJoinDialogFragment.java
index ec669bf..9afa186 100644
--- a/src/com/android/settings/connecteddevice/audiosharing/AudioSharingJoinDialogFragment.java
+++ b/src/com/android/settings/connecteddevice/audiosharing/AudioSharingJoinDialogFragment.java
@@ -81,7 +81,13 @@
             @NonNull DialogEventListener listener,
             @NonNull Pair<Integer, Object>[] eventData) {
         if (!AudioSharingUtils.isFeatureEnabled()) return;
-        final FragmentManager manager = host.getChildFragmentManager();
+        final FragmentManager manager;
+        try {
+            manager = host.getChildFragmentManager();
+        } catch (IllegalStateException e) {
+            Log.d(TAG, "Fail to show dialog: " + e.getMessage());
+            return;
+        }
         sListener = listener;
         sNewDevice = newDevice;
         sEventData = eventData;
diff --git a/src/com/android/settings/connecteddevice/audiosharing/AudioSharingStopDialogFragment.java b/src/com/android/settings/connecteddevice/audiosharing/AudioSharingStopDialogFragment.java
index b8da290..d026fa7 100644
--- a/src/com/android/settings/connecteddevice/audiosharing/AudioSharingStopDialogFragment.java
+++ b/src/com/android/settings/connecteddevice/audiosharing/AudioSharingStopDialogFragment.java
@@ -81,7 +81,13 @@
             @NonNull DialogEventListener listener,
             @NonNull Pair<Integer, Object>[] eventData) {
         if (!AudioSharingUtils.isFeatureEnabled()) return;
-        final FragmentManager manager = host.getChildFragmentManager();
+        final FragmentManager manager;
+        try {
+            manager = host.getChildFragmentManager();
+        } catch (IllegalStateException e) {
+            Log.d(TAG, "Fail to show dialog: " + e.getMessage());
+            return;
+        }
         AlertDialog dialog = AudioSharingDialogHelper.getDialogIfShowing(manager, TAG);
         if (dialog != null) {
             int newGroupId = AudioSharingUtils.getGroupId(newDevice);
diff --git a/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamButtonController.java b/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamButtonController.java
index 2661072..939dd5c 100644
--- a/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamButtonController.java
+++ b/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamButtonController.java
@@ -27,6 +27,7 @@
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 import androidx.lifecycle.DefaultLifecycleObserver;
 import androidx.lifecycle.LifecycleOwner;
 import androidx.preference.PreferenceScreen;
@@ -48,7 +49,9 @@
     private static final String TAG = "AudioStreamButtonController";
     private static final String KEY = "audio_stream_button";
     private static final int SOURCE_ORIGIN_REPOSITORY = SourceOriginForLogging.REPOSITORY.ordinal();
-    private final BluetoothLeBroadcastAssistant.Callback mBroadcastAssistantCallback =
+
+    @VisibleForTesting
+    final BluetoothLeBroadcastAssistant.Callback mBroadcastAssistantCallback =
             new AudioStreamsBroadcastAssistantCallback() {
                 @Override
                 public void onSourceRemoved(BluetoothDevice sink, int sourceId, int reason) {
@@ -97,8 +100,7 @@
                 }
             };
 
-    private final AudioStreamsRepository mAudioStreamsRepository =
-            AudioStreamsRepository.getInstance();
+    private AudioStreamsRepository mAudioStreamsRepository = AudioStreamsRepository.getInstance();
     private final Executor mExecutor;
     private final AudioStreamsHelper mAudioStreamsHelper;
     private final @Nullable LocalBluetoothLeBroadcastAssistant mLeBroadcastAssistant;
@@ -228,4 +230,9 @@
     void init(int broadcastId) {
         mBroadcastId = broadcastId;
     }
+
+    @VisibleForTesting
+    void setAudioStreamsRepositoryForTesting(AudioStreamsRepository repository) {
+        mAudioStreamsRepository = repository;
+    }
 }
diff --git a/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamConfirmDialog.java b/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamConfirmDialog.java
index 6c449a4..148c776 100644
--- a/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamConfirmDialog.java
+++ b/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamConfirmDialog.java
@@ -31,6 +31,7 @@
 import android.util.Log;
 
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 
 import com.android.settings.R;
 import com.android.settings.bluetooth.Utils;
@@ -43,9 +44,12 @@
 
 public class AudioStreamConfirmDialog extends InstrumentedDialogFragment {
     private static final String TAG = "AudioStreamConfirmDialog";
-    private static final int DEFAULT_DEVICE_NAME = R.string.audio_streams_dialog_default_device;
+
+    @VisibleForTesting
+    static final int DEFAULT_DEVICE_NAME = R.string.audio_streams_dialog_default_device;
+
     private Context mContext;
-    @Nullable private Activity mActivity;
+    @VisibleForTesting @Nullable Activity mActivity;
     @Nullable private BluetoothLeBroadcastMetadata mBroadcastMetadata;
     @Nullable private BluetoothDevice mConnectedDevice;
     private int mAudioStreamConfirmDialogId = SettingsEnums.PAGE_UNKNOWN;
diff --git a/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamHeaderController.java b/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamHeaderController.java
index 860e62e..e1a178d 100644
--- a/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamHeaderController.java
+++ b/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamHeaderController.java
@@ -54,7 +54,9 @@
     private final Executor mExecutor;
     private final AudioStreamsHelper mAudioStreamsHelper;
     @Nullable private final LocalBluetoothLeBroadcastAssistant mLeBroadcastAssistant;
-    private final BluetoothLeBroadcastAssistant.Callback mBroadcastAssistantCallback =
+
+    @VisibleForTesting
+    final BluetoothLeBroadcastAssistant.Callback mBroadcastAssistantCallback =
             new AudioStreamsBroadcastAssistantCallback() {
                 @Override
                 public void onSourceRemoved(BluetoothDevice sink, int sourceId, int reason) {
diff --git a/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamMediaService.java b/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamMediaService.java
index f812e06..ad358ed 100644
--- a/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamMediaService.java
+++ b/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamMediaService.java
@@ -16,6 +16,8 @@
 
 package com.android.settings.connecteddevice.audiosharing.audiostreams;
 
+import static java.util.Collections.emptyList;
+
 import android.app.Notification;
 import android.app.NotificationChannel;
 import android.app.NotificationManager;
@@ -50,10 +52,14 @@
 import com.android.settingslib.bluetooth.LocalBluetoothManager;
 import com.android.settingslib.bluetooth.VolumeControlProfile;
 import com.android.settingslib.core.instrumentation.MetricsFeatureProvider;
+import com.android.settingslib.utils.ThreadUtils;
 
-import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 
 public class AudioStreamMediaService extends Service {
     static final String BROADCAST_ID = "audio_stream_media_service_broadcast_id";
@@ -62,118 +68,13 @@
     private static final String TAG = "AudioStreamMediaService";
     private static final int NOTIFICATION_ID = 1;
     private static final int BROADCAST_CONTENT_TEXT = R.string.audio_streams_listening_now;
-    private static final String LEAVE_BROADCAST_ACTION = "leave_broadcast_action";
+    @VisibleForTesting static final String LEAVE_BROADCAST_ACTION = "leave_broadcast_action";
     private static final String LEAVE_BROADCAST_TEXT = "Leave Broadcast";
     private static final String CHANNEL_ID = "bluetooth_notification_channel";
     private static final String DEFAULT_DEVICE_NAME = "";
     private static final int STATIC_PLAYBACK_DURATION = 100;
     private static final int STATIC_PLAYBACK_POSITION = 30;
     private static final int ZERO_PLAYBACK_SPEED = 0;
-    private final AudioStreamsBroadcastAssistantCallback mBroadcastAssistantCallback =
-            new AudioStreamsBroadcastAssistantCallback() {
-                @Override
-                public void onSourceLost(int broadcastId) {
-                    super.onSourceLost(broadcastId);
-                    if (broadcastId == mBroadcastId) {
-                        Log.d(TAG, "onSourceLost() : stopSelf");
-                        if (mNotificationManager != null) {
-                            mNotificationManager.cancel(NOTIFICATION_ID);
-                        }
-                        stopSelf();
-                    }
-                }
-
-                @Override
-                public void onSourceRemoved(BluetoothDevice sink, int sourceId, int reason) {
-                    super.onSourceRemoved(sink, sourceId, reason);
-                    if (mAudioStreamsHelper != null
-                            && mAudioStreamsHelper.getAllConnectedSources().stream()
-                                    .map(BluetoothLeBroadcastReceiveState::getBroadcastId)
-                                    .noneMatch(id -> id == mBroadcastId)) {
-                        Log.d(TAG, "onSourceRemoved() : stopSelf");
-                        if (mNotificationManager != null) {
-                            mNotificationManager.cancel(NOTIFICATION_ID);
-                        }
-                        stopSelf();
-                    }
-                }
-            };
-
-    private final BluetoothCallback mBluetoothCallback =
-            new BluetoothCallback() {
-                @Override
-                public void onBluetoothStateChanged(int bluetoothState) {
-                    if (BluetoothAdapter.STATE_OFF == bluetoothState) {
-                        Log.d(TAG, "onBluetoothStateChanged() : stopSelf");
-                        if (mNotificationManager != null) {
-                            mNotificationManager.cancel(NOTIFICATION_ID);
-                        }
-                        stopSelf();
-                    }
-                }
-
-                @Override
-                public void onProfileConnectionStateChanged(
-                        @NonNull CachedBluetoothDevice cachedDevice,
-                        @ConnectionState int state,
-                        int bluetoothProfile) {
-                    if (state == BluetoothAdapter.STATE_DISCONNECTED
-                            && bluetoothProfile == BluetoothProfile.LE_AUDIO_BROADCAST_ASSISTANT
-                            && mDevices != null) {
-                        mDevices.remove(cachedDevice.getDevice());
-                        cachedDevice
-                                .getMemberDevice()
-                                .forEach(
-                                        m -> {
-                                            // Check nullability to pass NullAway check
-                                            if (mDevices != null) {
-                                                mDevices.remove(m.getDevice());
-                                            }
-                                        });
-                    }
-                    if (mDevices == null || mDevices.isEmpty()) {
-                        Log.d(TAG, "onProfileConnectionStateChanged() : stopSelf");
-                        if (mNotificationManager != null) {
-                            mNotificationManager.cancel(NOTIFICATION_ID);
-                        }
-                        stopSelf();
-                    }
-                }
-            };
-
-    private final BluetoothVolumeControl.Callback mVolumeControlCallback =
-            new BluetoothVolumeControl.Callback() {
-                @Override
-                public void onDeviceVolumeChanged(
-                        @NonNull BluetoothDevice device,
-                        @IntRange(from = -255, to = 255) int volume) {
-                    if (mDevices == null || mDevices.isEmpty()) {
-                        Log.w(TAG, "active device or device has source is null!");
-                        return;
-                    }
-                    if (mDevices.contains(device)) {
-                        Log.d(
-                                TAG,
-                                "onDeviceVolumeChanged() bluetoothDevice : "
-                                        + device
-                                        + " volume: "
-                                        + volume);
-                        if (volume == 0) {
-                            mIsMuted = true;
-                        } else {
-                            mIsMuted = false;
-                            mLatestPositiveVolume = volume;
-                        }
-                        if (mLocalSession != null) {
-                            mLocalSession.setPlaybackState(getPlaybackState());
-                            if (mNotificationManager != null) {
-                                mNotificationManager.notify(NOTIFICATION_ID, buildNotification());
-                            }
-                        }
-                    }
-                }
-            };
-
     private final PlaybackState.Builder mPlayStatePlayingBuilder =
             new PlaybackState.Builder()
                     .setActions(PlaybackState.ACTION_PAUSE | PlaybackState.ACTION_SEEK_TO)
@@ -200,20 +101,24 @@
     private final MetricsFeatureProvider mMetricsFeatureProvider =
             FeatureFactory.getFeatureFactory().getMetricsFeatureProvider();
     private final ExecutorService mExecutor = Executors.newSingleThreadExecutor();
+    private final AtomicBoolean mIsMuted = new AtomicBoolean(false);
+    // Set 25 as default as the volume range from `VolumeControlProfile` is from 0 to 255.
+    // If the initial volume from `onDeviceVolumeChanged` is larger than zero (not muted), we will
+    // override this value. Otherwise, we raise the volume to 25 when the play button is clicked.
+    private final AtomicInteger mLatestPositiveVolume = new AtomicInteger(25);
+    private final AtomicBoolean mHasStopped = new AtomicBoolean(false);
     private int mBroadcastId;
-    @Nullable private ArrayList<BluetoothDevice> mDevices;
+    @Nullable private List<BluetoothDevice> mDevices;
     @Nullable private LocalBluetoothManager mLocalBtManager;
     @Nullable private AudioStreamsHelper mAudioStreamsHelper;
     @Nullable private LocalBluetoothLeBroadcastAssistant mLeBroadcastAssistant;
     @Nullable private VolumeControlProfile mVolumeControl;
     @Nullable private NotificationManager mNotificationManager;
-
-    // Set 25 as default as the volume range from `VolumeControlProfile` is from 0 to 255.
-    // If the initial volume from `onDeviceVolumeChanged` is larger than zero (not muted), we will
-    // override this value. Otherwise, we raise the volume to 25 when the play button is clicked.
-    private int mLatestPositiveVolume = 25;
-    private boolean mIsMuted = false;
-    @VisibleForTesting @Nullable MediaSession mLocalSession;
+    @Nullable private MediaSession mLocalSession;
+    @VisibleForTesting @Nullable AudioStreamsBroadcastAssistantCallback mBroadcastAssistantCallback;
+    @VisibleForTesting @Nullable BluetoothCallback mBluetoothCallback;
+    @VisibleForTesting @Nullable BluetoothVolumeControl.Callback mVolumeControlCallback;
+    @VisibleForTesting @Nullable MediaSession.Callback mMediaSessionCallback;
 
     @Override
     public void onCreate() {
@@ -250,13 +155,16 @@
             mNotificationManager.createNotificationChannel(notificationChannel);
         }
 
+        mBluetoothCallback = new BtCallback();
         mLocalBtManager.getEventManager().registerCallback(mBluetoothCallback);
 
         mVolumeControl = mLocalBtManager.getProfileManager().getVolumeControlProfile();
         if (mVolumeControl != null) {
+            mVolumeControlCallback = new VolumeControlCallback();
             mVolumeControl.registerCallback(mExecutor, mVolumeControlCallback);
         }
 
+        mBroadcastAssistantCallback = new AssistantCallback();
         mLeBroadcastAssistant.registerServiceCallBack(mExecutor, mBroadcastAssistantCallback);
     }
 
@@ -264,25 +172,19 @@
     public void onDestroy() {
         Log.d(TAG, "onDestroy()");
         super.onDestroy();
-
         if (!AudioSharingUtils.isFeatureEnabled()) {
-            Log.d(TAG, "onDestroy() : skip due to feature not enabled");
             return;
         }
         if (mLocalBtManager != null) {
-            Log.d(TAG, "onDestroy() : unregister mBluetoothCallback");
             mLocalBtManager.getEventManager().unregisterCallback(mBluetoothCallback);
         }
-        if (mLeBroadcastAssistant != null) {
-            Log.d(TAG, "onDestroy() : unregister mBroadcastAssistantCallback");
+        if (mLeBroadcastAssistant != null && mBroadcastAssistantCallback != null) {
             mLeBroadcastAssistant.unregisterServiceCallBack(mBroadcastAssistantCallback);
         }
-        if (mVolumeControl != null) {
-            Log.d(TAG, "onDestroy() : unregister mVolumeControlCallback");
+        if (mVolumeControl != null && mVolumeControlCallback != null) {
             mVolumeControl.unregisterCallback(mVolumeControlCallback);
         }
         if (mLocalSession != null) {
-            Log.d(TAG, "onDestroy() : release mLocalSession");
             mLocalSession.release();
             mLocalSession = null;
         }
@@ -291,33 +193,31 @@
     @Override
     public int onStartCommand(@Nullable Intent intent, int flags, int startId) {
         Log.d(TAG, "onStartCommand()");
-
-        mBroadcastId = intent != null ? intent.getIntExtra(BROADCAST_ID, -1) : -1;
+        if (intent == null) {
+            Log.w(TAG, "Intent is null. Service will not start.");
+            mHasStopped.set(true);
+            stopSelf();
+            return START_NOT_STICKY;
+        }
+        mBroadcastId = intent.getIntExtra(BROADCAST_ID, -1);
         if (mBroadcastId == -1) {
             Log.w(TAG, "Invalid broadcast ID. Service will not start.");
-            if (mNotificationManager != null) {
-                mNotificationManager.cancel(NOTIFICATION_ID);
-            }
+            mHasStopped.set(true);
             stopSelf();
             return START_NOT_STICKY;
         }
-
-        if (intent != null) {
-            mDevices = intent.getParcelableArrayListExtra(DEVICES, BluetoothDevice.class);
-        }
-        if (mDevices == null || mDevices.isEmpty()) {
+        var extra = intent.getParcelableArrayListExtra(DEVICES, BluetoothDevice.class);
+        if (extra == null || extra.isEmpty()) {
             Log.w(TAG, "No device. Service will not start.");
-            if (mNotificationManager != null) {
-                mNotificationManager.cancel(NOTIFICATION_ID);
-            }
+            mHasStopped.set(true);
             stopSelf();
             return START_NOT_STICKY;
         }
-        if (intent != null) {
-            createLocalMediaSession(intent.getStringExtra(BROADCAST_TITLE));
-            startForeground(NOTIFICATION_ID, buildNotification());
-        }
-
+        mDevices = Collections.synchronizedList(extra);
+        createLocalMediaSession(intent.getStringExtra(BROADCAST_TITLE));
+        startForeground(NOTIFICATION_ID, buildNotification());
+        // Reset in case the service is previously stopped but not yet destroyed.
+        mHasStopped.set(false);
         return START_NOT_STICKY;
     }
 
@@ -330,78 +230,12 @@
                         .build());
         mLocalSession.setActive(true);
         mLocalSession.setPlaybackState(getPlaybackState());
-        mLocalSession.setCallback(
-                new MediaSession.Callback() {
-                    public void onSeekTo(long pos) {
-                        Log.d(TAG, "onSeekTo: " + pos);
-                        if (mLocalSession != null) {
-                            mLocalSession.setPlaybackState(getPlaybackState());
-                            if (mNotificationManager != null) {
-                                mNotificationManager.notify(NOTIFICATION_ID, buildNotification());
-                            }
-                        }
-                    }
-
-                    @Override
-                    public void onPause() {
-                        if (mDevices == null || mDevices.isEmpty()) {
-                            Log.w(TAG, "active device or device has source is null!");
-                            return;
-                        }
-                        Log.d(
-                                TAG,
-                                "onPause() setting volume for device : "
-                                        + mDevices.get(0)
-                                        + " volume: "
-                                        + 0);
-                        if (mVolumeControl != null) {
-                            mVolumeControl.setDeviceVolume(mDevices.get(0), 0, true);
-                            mMetricsFeatureProvider.action(
-                                    getApplicationContext(),
-                                    SettingsEnums
-                                            .ACTION_AUDIO_STREAM_NOTIFICATION_MUTE_BUTTON_CLICK,
-                                    1);
-                        }
-                    }
-
-                    @Override
-                    public void onPlay() {
-                        if (mDevices == null || mDevices.isEmpty()) {
-                            Log.w(TAG, "active device or device has source is null!");
-                            return;
-                        }
-                        Log.d(
-                                TAG,
-                                "onPlay() setting volume for device : "
-                                        + mDevices.get(0)
-                                        + " volume: "
-                                        + mLatestPositiveVolume);
-                        if (mVolumeControl != null) {
-                            mVolumeControl.setDeviceVolume(
-                                    mDevices.get(0), mLatestPositiveVolume, true);
-                        }
-                        mMetricsFeatureProvider.action(
-                                getApplicationContext(),
-                                SettingsEnums.ACTION_AUDIO_STREAM_NOTIFICATION_MUTE_BUTTON_CLICK,
-                                0);
-                    }
-
-                    @Override
-                    public void onCustomAction(@NonNull String action, Bundle extras) {
-                        Log.d(TAG, "onCustomAction: " + action);
-                        if (action.equals(LEAVE_BROADCAST_ACTION) && mAudioStreamsHelper != null) {
-                            mAudioStreamsHelper.removeSource(mBroadcastId);
-                            mMetricsFeatureProvider.action(
-                                    getApplicationContext(),
-                                    SettingsEnums
-                                            .ACTION_AUDIO_STREAM_NOTIFICATION_LEAVE_BUTTON_CLICK);
-                        }
-                    }
-                });
+        mMediaSessionCallback = new MediaSessionCallback();
+        mLocalSession.setCallback(mMediaSessionCallback);
     }
 
     private PlaybackState getPlaybackState() {
-        return mIsMuted ? mPlayStatePausingBuilder.build() : mPlayStatePlayingBuilder.build();
+        return mIsMuted.get() ? mPlayStatePausingBuilder.build() : mPlayStatePlayingBuilder.build();
     }
 
     private String getDeviceName() {
@@ -442,4 +276,167 @@
     public IBinder onBind(Intent intent) {
         return null;
     }
+
+    private class AssistantCallback extends AudioStreamsBroadcastAssistantCallback {
+        @Override
+        public void onSourceLost(int broadcastId) {
+            super.onSourceLost(broadcastId);
+            handleRemoveSource();
+        }
+
+        @Override
+        public void onSourceRemoved(BluetoothDevice sink, int sourceId, int reason) {
+            super.onSourceRemoved(sink, sourceId, reason);
+            handleRemoveSource();
+        }
+
+        private void handleRemoveSource() {
+            var unused =
+                    ThreadUtils.postOnBackgroundThread(
+                            () -> {
+                                List<BluetoothLeBroadcastReceiveState> connected =
+                                        mAudioStreamsHelper == null
+                                                ? emptyList()
+                                                : mAudioStreamsHelper.getAllConnectedSources();
+                                if (connected.stream()
+                                        .map(BluetoothLeBroadcastReceiveState::getBroadcastId)
+                                        .noneMatch(id -> id == mBroadcastId)) {
+                                    mHasStopped.set(true);
+                                    stopSelf();
+                                }
+                            });
+        }
+    }
+
+    private class VolumeControlCallback implements BluetoothVolumeControl.Callback {
+        @Override
+        public void onDeviceVolumeChanged(
+                @NonNull BluetoothDevice device, @IntRange(from = -255, to = 255) int volume) {
+            if (mDevices == null || mDevices.isEmpty()) {
+                Log.w(TAG, "active device or device has source is null!");
+                return;
+            }
+            Log.d(
+                    TAG,
+                    "onDeviceVolumeChanged() bluetoothDevice : " + device + " volume: " + volume);
+            if (mDevices.contains(device)) {
+                if (volume == 0) {
+                    mIsMuted.set(true);
+                } else {
+                    mIsMuted.set(false);
+                    mLatestPositiveVolume.set(volume);
+                }
+                updateNotification(getPlaybackState());
+            }
+        }
+    }
+
+    private class BtCallback implements BluetoothCallback {
+        @Override
+        public void onBluetoothStateChanged(int bluetoothState) {
+            if (BluetoothAdapter.STATE_OFF == bluetoothState) {
+                Log.d(TAG, "onBluetoothStateChanged() : stopSelf");
+                mHasStopped.set(true);
+                stopSelf();
+            }
+        }
+
+        @Override
+        public void onProfileConnectionStateChanged(
+                @NonNull CachedBluetoothDevice cachedDevice,
+                @ConnectionState int state,
+                int bluetoothProfile) {
+            if (state == BluetoothAdapter.STATE_DISCONNECTED
+                    && bluetoothProfile == BluetoothProfile.LE_AUDIO_BROADCAST_ASSISTANT
+                    && mDevices != null) {
+                mDevices.remove(cachedDevice.getDevice());
+                cachedDevice
+                        .getMemberDevice()
+                        .forEach(
+                                m -> {
+                                    // Check nullability to pass NullAway check
+                                    if (mDevices != null) {
+                                        mDevices.remove(m.getDevice());
+                                    }
+                                });
+            }
+            if (mDevices == null || mDevices.isEmpty()) {
+                Log.d(TAG, "onProfileConnectionStateChanged() : stopSelf");
+                mHasStopped.set(true);
+                stopSelf();
+            }
+        }
+    }
+
+    private class MediaSessionCallback extends MediaSession.Callback {
+        public void onSeekTo(long pos) {
+            Log.d(TAG, "onSeekTo: " + pos);
+            updateNotification(getPlaybackState());
+        }
+
+        @Override
+        public void onPause() {
+            if (mDevices == null || mDevices.isEmpty()) {
+                Log.w(TAG, "active device or device has source is null!");
+                return;
+            }
+            Log.d(
+                    TAG,
+                    "onPause() setting volume for device : " + mDevices.get(0) + " volume: " + 0);
+            setDeviceVolume(mDevices.get(0), /* volume= */ 0);
+        }
+
+        @Override
+        public void onPlay() {
+            if (mDevices == null || mDevices.isEmpty()) {
+                Log.w(TAG, "active device or device has source is null!");
+                return;
+            }
+            Log.d(
+                    TAG,
+                    "onPlay() setting volume for device : "
+                            + mDevices.get(0)
+                            + " volume: "
+                            + mLatestPositiveVolume.get());
+            setDeviceVolume(mDevices.get(0), mLatestPositiveVolume.get());
+        }
+
+        @Override
+        public void onCustomAction(@NonNull String action, Bundle extras) {
+            Log.d(TAG, "onCustomAction: " + action);
+            if (action.equals(LEAVE_BROADCAST_ACTION) && mAudioStreamsHelper != null) {
+                mAudioStreamsHelper.removeSource(mBroadcastId);
+                mMetricsFeatureProvider.action(
+                        getApplicationContext(),
+                        SettingsEnums.ACTION_AUDIO_STREAM_NOTIFICATION_LEAVE_BUTTON_CLICK);
+            }
+        }
+
+        private void setDeviceVolume(BluetoothDevice device, int volume) {
+            int event = SettingsEnums.ACTION_AUDIO_STREAM_NOTIFICATION_MUTE_BUTTON_CLICK;
+            var unused =
+                    ThreadUtils.postOnBackgroundThread(
+                            () -> {
+                                if (mVolumeControl != null) {
+                                    mVolumeControl.setDeviceVolume(device, volume, true);
+                                    mMetricsFeatureProvider.action(
+                                            getApplicationContext(), event, volume == 0 ? 1 : 0);
+                                }
+                            });
+        }
+    }
+
+    private void updateNotification(PlaybackState playbackState) {
+        var unused =
+                ThreadUtils.postOnBackgroundThread(
+                        () -> {
+                            if (mLocalSession != null) {
+                                mLocalSession.setPlaybackState(playbackState);
+                                if (mNotificationManager != null && !mHasStopped.get()) {
+                                    mNotificationManager.notify(
+                                            NOTIFICATION_ID, buildNotification());
+                                }
+                            }
+                        });
+    }
 }
diff --git a/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsDashboardFragment.java b/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsDashboardFragment.java
index ce32cdb..ae5cb6e 100644
--- a/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsDashboardFragment.java
+++ b/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsDashboardFragment.java
@@ -23,7 +23,6 @@
 import android.bluetooth.BluetoothLeBroadcastMetadata;
 import android.content.Context;
 import android.content.Intent;
-import android.os.Bundle;
 import android.util.Log;
 
 import androidx.annotation.Nullable;
@@ -64,11 +63,6 @@
     }
 
     @Override
-    public void onCreate(Bundle savedInstanceState) {
-        super.onCreate(savedInstanceState);
-    }
-
-    @Override
     public void onAttach(Context context) {
         super.onAttach(context);
         use(AudioStreamsScanQrCodeController.class).setFragment(this);
@@ -92,11 +86,6 @@
     }
 
     @Override
-    public void onActivityCreated(Bundle savedInstanceState) {
-        super.onActivityCreated(savedInstanceState);
-    }
-
-    @Override
     public void onActivityResult(int requestCode, int resultCode, @Nullable Intent data) {
         super.onActivityResult(requestCode, resultCode, data);
         if (DEBUG) {
diff --git a/src/com/android/settings/core/gateway/SettingsGateway.java b/src/com/android/settings/core/gateway/SettingsGateway.java
index 11c05f3..734bddc 100644
--- a/src/com/android/settings/core/gateway/SettingsGateway.java
+++ b/src/com/android/settings/core/gateway/SettingsGateway.java
@@ -156,6 +156,7 @@
 import com.android.settings.notification.ConfigureNotificationSettings;
 import com.android.settings.notification.NotificationAccessSettings;
 import com.android.settings.notification.NotificationAssistantPicker;
+import com.android.settings.notification.PoliteNotificationsPreferenceFragment;
 import com.android.settings.notification.SoundSettings;
 import com.android.settings.notification.app.AppBubbleNotificationSettings;
 import com.android.settings.notification.app.AppNotificationSettings;
@@ -314,6 +315,7 @@
             AppInfoDashboardFragment.class.getName(),
             BatterySaverSettings.class.getName(),
             AppNotificationSettings.class.getName(),
+            PoliteNotificationsPreferenceFragment.class.getName(),
             NotificationAssistantPicker.class.getName(),
             ChannelNotificationSettings.class.getName(),
             SatelliteSetting.class.getName(),
diff --git a/src/com/android/settings/datausage/DataUsageList.kt b/src/com/android/settings/datausage/DataUsageList.kt
index af115d9..a293277 100644
--- a/src/com/android/settings/datausage/DataUsageList.kt
+++ b/src/com/android/settings/datausage/DataUsageList.kt
@@ -33,13 +33,10 @@
 import com.android.settings.dashboard.DashboardFragment
 import com.android.settings.datausage.lib.BillingCycleRepository
 import com.android.settings.datausage.lib.NetworkUsageData
-import com.android.settings.network.MobileNetworkRepository
 import com.android.settings.network.SubscriptionUtil
-import com.android.settings.network.telephony.requireSubscriptionManager
-import com.android.settingslib.mobile.dataservice.SubscriptionInfoEntity
+import com.android.settings.network.telephony.SubscriptionRepository
 import com.android.settingslib.spa.framework.util.collectLatestWithLifecycle
 import com.android.settingslib.spaprivileged.framework.common.userManager
-import com.android.settingslib.utils.ThreadUtils
 import kotlin.jvm.optionals.getOrNull
 
 /**
@@ -59,7 +56,6 @@
     private lateinit var billingCycleRepository: BillingCycleRepository
 
     private var usageAmount: Preference? = null
-    private var subscriptionInfoEntity: SubscriptionInfoEntity? = null
     private var dataUsageListAppsController: DataUsageListAppsController? = null
     private var chartDataUsagePreferenceController: ChartDataUsagePreferenceController? = null
     private var dataUsageListHeaderController: DataUsageListHeaderController? = null
@@ -90,7 +86,6 @@
             finish()
             return
         }
-        updateSubscriptionInfoEntity()
         dataUsageListAppsController = use(DataUsageListAppsController::class.java).apply {
             init(template)
         }
@@ -132,6 +127,16 @@
         viewModel.chartDataFlow.collectLatestWithLifecycle(viewLifecycleOwner) { chartData ->
             chartDataUsagePreferenceController?.update(chartData)
         }
+        finishIfSubscriptionDisabled()
+    }
+
+    private fun finishIfSubscriptionDisabled() {
+        if (SubscriptionManager.isUsableSubscriptionId(subId)) {
+            SubscriptionRepository(requireContext()).isSubscriptionEnabledFlow(subId)
+                .collectLatestWithLifecycle(viewLifecycleOwner) { isSubscriptionEnabled ->
+                    if (!isSubscriptionEnabled) finish()
+                }
+        }
     }
 
     override fun getPreferenceScreenResId() = R.xml.data_usage_list
@@ -155,23 +160,12 @@
         }
     }
 
-    private fun updateSubscriptionInfoEntity() {
-        ThreadUtils.postOnBackgroundThread {
-            subscriptionInfoEntity =
-                MobileNetworkRepository.getInstance(context).getSubInfoById(subId.toString())
-        }
-    }
-
     /** Update chart sweeps and cycle list to reflect [NetworkPolicy] for current [template]. */
     private fun updatePolicy(isModifiable: Boolean) {
-        val isBillingCycleModifiable = isModifiable && isActiveSubscription()
-        dataUsageListHeaderController?.setConfigButtonVisible(isBillingCycleModifiable)
-        chartDataUsagePreferenceController?.setBillingCycleModifiable(isBillingCycleModifiable)
+        dataUsageListHeaderController?.setConfigButtonVisible(isModifiable)
+        chartDataUsagePreferenceController?.setBillingCycleModifiable(isModifiable)
     }
 
-    private fun isActiveSubscription(): Boolean =
-            requireContext().requireSubscriptionManager().getActiveSubscriptionInfo(subId) != null
-
     /**
      * Updates the chart and detail data when initial loaded or selected cycle changed.
      */
@@ -187,7 +181,7 @@
     /** Updates applications data usage. */
     private fun updateApps(usageData: NetworkUsageData) {
         dataUsageListAppsController?.update(
-            carrierId = subscriptionInfoEntity?.carrierId,
+            subId = subId,
             startTime = usageData.startTime,
             endTime = usageData.endTime,
         )
diff --git a/src/com/android/settings/datausage/DataUsageListAppsController.kt b/src/com/android/settings/datausage/DataUsageListAppsController.kt
index 93623f4..d8bddde 100644
--- a/src/com/android/settings/datausage/DataUsageListAppsController.kt
+++ b/src/com/android/settings/datausage/DataUsageListAppsController.kt
@@ -20,6 +20,7 @@
 import android.content.Context
 import android.net.NetworkTemplate
 import android.os.Bundle
+import android.telephony.SubscriptionManager
 import androidx.annotation.OpenForTesting
 import androidx.annotation.VisibleForTesting
 import androidx.lifecycle.LifecycleCoroutineScope
@@ -32,6 +33,7 @@
 import com.android.settings.core.SubSettingLauncher
 import com.android.settings.datausage.lib.AppDataUsageRepository
 import com.android.settings.datausage.lib.NetworkUsageData
+import com.android.settings.network.telephony.requireSubscriptionManager
 import com.android.settingslib.AppItem
 import com.android.settingslib.net.UidDetailProvider
 import kotlinx.coroutines.Dispatchers
@@ -74,8 +76,11 @@
         this.cycleData = cycleData
     }
 
-    fun update(carrierId: Int?, startTime: Long, endTime: Long) = lifecycleScope.launch {
+    fun update(subId: Int, startTime: Long, endTime: Long) = lifecycleScope.launch {
         val apps = withContext(Dispatchers.Default) {
+            val carrierId = if (SubscriptionManager.isValidSubscriptionId(subId)) {
+                mContext.requireSubscriptionManager().getActiveSubscriptionInfo(subId)?.carrierId
+            } else null
             repository.getAppPercent(carrierId, startTime, endTime).map { (appItem, percent) ->
                 AppDataUsagePreference(mContext, appItem, percent, uidDetailProvider).apply {
                     setOnPreferenceClickListener {
diff --git a/src/com/android/settings/development/PageAgnosticNotificationService.java b/src/com/android/settings/development/PageAgnosticNotificationService.java
index bce1dd9..40dff0c 100644
--- a/src/com/android/settings/development/PageAgnosticNotificationService.java
+++ b/src/com/android/settings/development/PageAgnosticNotificationService.java
@@ -96,13 +96,13 @@
                                 notifyPendingIntent)
                         .build();
 
+        // TODO:(b/349860833) Change text style to BigTextStyle once the ellipsis issue is fixed.
         Notification.Builder builder =
                 new Notification.Builder(this, NOTIFICATION_CHANNEL_ID)
                         .setContentTitle(title)
                         .setContentText(text)
                         .setOngoing(true)
                         .setSmallIcon(R.drawable.ic_settings_24dp)
-                        .setStyle(new Notification.BigTextStyle().bigText(text))
                         .setContentIntent(notifyPendingIntent)
                         .addAction(action);
 
diff --git a/src/com/android/settings/fuelgauge/BatteryInfo.java b/src/com/android/settings/fuelgauge/BatteryInfo.java
index b54801a..7cf9e44 100644
--- a/src/com/android/settings/fuelgauge/BatteryInfo.java
+++ b/src/com/android/settings/fuelgauge/BatteryInfo.java
@@ -53,7 +53,8 @@
     public int batteryStatus;
     public int pluggedStatus;
     public boolean discharging = true;
-    public boolean isBatteryDefender;
+    public boolean isBatteryDefender = false;
+    public boolean isLongLife = false;
     public boolean isFastCharging;
     public long remainingTimeUs = 0;
     public long averageTimeToDischarge = EstimateKt.AVERAGE_TIME_TO_DISCHARGE_UNKNOWN;
@@ -306,7 +307,7 @@
         info.pluggedStatus = batteryBroadcast.getIntExtra(BatteryManager.EXTRA_PLUGGED, 0);
         info.mCharging = info.pluggedStatus != 0;
         info.averageTimeToDischarge = estimate.getAverageDischargeTime();
-        info.isBatteryDefender =
+        info.isLongLife =
                 batteryBroadcast.getIntExtra(
                                 BatteryManager.EXTRA_CHARGING_STATUS,
                                 BatteryManager.CHARGING_POLICY_DEFAULT)
@@ -319,7 +320,7 @@
         info.isFastCharging =
                 BatteryStatus.getChargingSpeed(context, batteryBroadcast)
                         == BatteryStatus.CHARGING_FAST;
-        if (info.isBatteryDefender) {
+        if (info.isLongLife) {
             info.isBatteryDefender =
                     FeatureFactory.getFeatureFactory()
                             .getPowerUsageFeatureProvider()
diff --git a/src/com/android/settings/fuelgauge/BatteryUtils.java b/src/com/android/settings/fuelgauge/BatteryUtils.java
index 9e08664..7cb5733 100644
--- a/src/com/android/settings/fuelgauge/BatteryUtils.java
+++ b/src/com/android/settings/fuelgauge/BatteryUtils.java
@@ -600,12 +600,12 @@
                             context.getContentResolver(), SETTINGS_GLOBAL_DOCK_DEFENDER_BYPASS, 0)
                     == 1) {
                 return DockDefenderMode.TEMPORARILY_BYPASSED;
-            } else if (batteryInfo.isBatteryDefender
+            } else if (batteryInfo.isLongLife
                     && FeatureFactory.getFeatureFactory()
                             .getPowerUsageFeatureProvider()
                             .isExtraDefend()) {
                 return DockDefenderMode.ACTIVE;
-            } else if (!batteryInfo.isBatteryDefender) {
+            } else if (!batteryInfo.isLongLife) {
                 return DockDefenderMode.FUTURE_BYPASS;
             }
         }
diff --git a/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImpl.java b/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImpl.java
index 65e8435..02dff7e 100644
--- a/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImpl.java
+++ b/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImpl.java
@@ -248,7 +248,7 @@
 
     @Override
     public boolean isBatteryDefend(BatteryInfo info) {
-        return info.isBatteryDefender && !isExtraDefend();
+        return info.isLongLife && !isExtraDefend();
     }
 
     @Override
diff --git a/src/com/android/settings/fuelgauge/batterytip/detectors/BatteryDefenderDetector.java b/src/com/android/settings/fuelgauge/batterytip/detectors/BatteryDefenderDetector.java
index 639b3c7..6ff2665 100644
--- a/src/com/android/settings/fuelgauge/batterytip/detectors/BatteryDefenderDetector.java
+++ b/src/com/android/settings/fuelgauge/batterytip/detectors/BatteryDefenderDetector.java
@@ -21,7 +21,6 @@
 import com.android.settings.fuelgauge.BatteryInfo;
 import com.android.settings.fuelgauge.batterytip.tips.BatteryDefenderTip;
 import com.android.settings.fuelgauge.batterytip.tips.BatteryTip;
-import com.android.settings.overlay.FeatureFactory;
 
 /** Detect whether the battery is overheated */
 public class BatteryDefenderDetector implements BatteryTipDetector {
@@ -35,12 +34,10 @@
 
     @Override
     public BatteryTip detect() {
-        final boolean isBasicBatteryDefend =
-                FeatureFactory.getFeatureFactory()
-                        .getPowerUsageFeatureProvider()
-                        .isBatteryDefend(mBatteryInfo);
         final int state =
-                isBasicBatteryDefend ? BatteryTip.StateType.NEW : BatteryTip.StateType.INVISIBLE;
+                mBatteryInfo.isBatteryDefender
+                        ? BatteryTip.StateType.NEW
+                        : BatteryTip.StateType.INVISIBLE;
         final boolean isPluggedIn = mBatteryInfo.pluggedStatus != 0;
         return new BatteryDefenderTip(state, isPluggedIn);
     }
diff --git a/src/com/android/settings/network/MobileNetworkRepository.java b/src/com/android/settings/network/MobileNetworkRepository.java
index ebb341e..bd892c8 100644
--- a/src/com/android/settings/network/MobileNetworkRepository.java
+++ b/src/com/android/settings/network/MobileNetworkRepository.java
@@ -41,7 +41,6 @@
 import androidx.lifecycle.LifecycleOwner;
 
 import com.android.internal.telephony.flags.Flags;
-import com.android.settings.network.telephony.MobileNetworkUtils;
 import com.android.settings.overlay.FeatureFactory;
 import com.android.settingslib.core.instrumentation.MetricsFeatureProvider;
 import com.android.settingslib.mobile.dataservice.MobileNetworkDatabase;
@@ -435,7 +434,7 @@
                 mMetricsFeatureProvider.action(mContext,
                         SettingsEnums.ACTION_MOBILE_NETWORK_DB_INSERT_SUB_INFO, subId);
                 insertUiccInfo(subId, telephonyManager);
-                insertMobileNetworkInfo(context, subId, telephonyManager);
+                insertMobileNetworkInfo(subId, telephonyManager);
             }
         } else if (DEBUG) {
             Log.d(TAG, "Can not insert subInfo, the entity is null");
@@ -517,20 +516,13 @@
         }
     }
 
-    private void insertMobileNetworkInfo(Context context, int subId,
-            TelephonyManager telephonyManager) {
-        MobileNetworkInfoEntity mobileNetworkInfoEntity = convertToMobileNetworkInfoEntity(context,
-                subId, telephonyManager);
-
+    private void insertMobileNetworkInfo(int subId, TelephonyManager telephonyManager) {
+        MobileNetworkInfoEntity mobileNetworkInfoEntity =
+                convertToMobileNetworkInfoEntity(subId, telephonyManager);
 
         Log.d(TAG, "insertMobileNetworkInfo, mobileNetworkInfoEntity = "
                 + mobileNetworkInfoEntity);
 
-
-        if (mobileNetworkInfoEntity == null) {
-            return;
-        }
-
         if (!sCacheMobileNetworkInfoEntityMap.containsKey(subId)
                 || !sCacheMobileNetworkInfoEntityMap.get(subId).equals(mobileNetworkInfoEntity)) {
             sCacheMobileNetworkInfoEntityMap.put(subId, mobileNetworkInfoEntity);
@@ -540,7 +532,7 @@
         }
     }
 
-    private MobileNetworkInfoEntity convertToMobileNetworkInfoEntity(Context context, int subId,
+    private MobileNetworkInfoEntity convertToMobileNetworkInfoEntity(int subId,
             TelephonyManager telephonyManager) {
         boolean isDataEnabled = false;
         if (telephonyManager != null) {
@@ -549,18 +541,8 @@
             Log.d(TAG, "TelephonyManager is null, subId = " + subId);
         }
 
-        return new MobileNetworkInfoEntity(String.valueOf(subId),
-                MobileNetworkUtils.isContactDiscoveryEnabled(context, subId),
-                MobileNetworkUtils.isContactDiscoveryVisible(context, subId),
-                isDataEnabled,
-                MobileNetworkUtils.isCdmaOptions(context, subId),
-                MobileNetworkUtils.isGsmOptions(context, subId),
-                MobileNetworkUtils.isWorldMode(context, subId),
-                MobileNetworkUtils.shouldDisplayNetworkSelectOptions(context, subId),
-                MobileNetworkUtils.isTdscdmaSupported(context, subId),
-                MobileNetworkUtils.activeNetworkIsCellular(context),
-                SubscriptionUtil.showToggleForPhysicalSim(mSubscriptionManager),
-                /* deprecated isDataRoamingEnabled = */ false
+        return new MobileNetworkInfoEntity(String.valueOf(subId), isDataEnabled,
+                SubscriptionUtil.showToggleForPhysicalSim(mSubscriptionManager)
         );
     }
 
@@ -681,8 +663,7 @@
         public void onUserMobileDataStateChanged(boolean enabled) {
             Log.d(TAG, "onUserMobileDataStateChanged enabled " + enabled + " on SUB " + mSubId);
             sExecutor.execute(() -> {
-                insertMobileNetworkInfo(mContext, mSubId,
-                        getTelephonyManagerBySubId(mContext, mSubId));
+                insertMobileNetworkInfo(mSubId, getTelephonyManagerBySubId(mContext, mSubId));
             });
         }
     }
diff --git a/src/com/android/settings/network/NetworkProviderSettings.java b/src/com/android/settings/network/NetworkProviderSettings.java
index a4193f8..0fcfcb5 100644
--- a/src/com/android/settings/network/NetworkProviderSettings.java
+++ b/src/com/android/settings/network/NetworkProviderSettings.java
@@ -36,6 +36,7 @@
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.provider.Settings;
+import android.telephony.SubscriptionManager;
 import android.telephony.TelephonyManager;
 import android.text.TextUtils;
 import android.util.EventLog;
@@ -370,7 +371,7 @@
         mDataUsagePreference = findPreference(PREF_KEY_DATA_USAGE);
         mDataUsagePreference.setVisible(DataUsageUtils.hasWifiRadio(getContext()));
         mDataUsagePreference.setTemplate(new NetworkTemplate.Builder(NetworkTemplate.MATCH_WIFI)
-                        .build(), 0 /*subId*/);
+                        .build(), SubscriptionManager.INVALID_SUBSCRIPTION_ID);
         mResetInternetPreference = findPreference(PREF_KEY_RESET_INTERNET);
         if (mResetInternetPreference != null) {
             mResetInternetPreference.setVisible(false);
diff --git a/src/com/android/settings/notification/modes/ManualDurationHelper.java b/src/com/android/settings/notification/modes/ManualDurationHelper.java
new file mode 100644
index 0000000..da9f420
--- /dev/null
+++ b/src/com/android/settings/notification/modes/ManualDurationHelper.java
@@ -0,0 +1,123 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.notification.modes;
+
+import android.annotation.Nullable;
+import android.content.Context;
+import android.database.ContentObserver;
+import android.icu.text.MessageFormat;
+import android.net.Uri;
+import android.provider.Settings;
+
+import androidx.annotation.NonNull;
+import androidx.preference.Preference;
+
+import com.android.settings.R;
+
+import java.util.HashMap;
+import java.util.Locale;
+import java.util.Map;
+
+/**
+ * Class to contain shared utilities for reading and observing the Settings ZEN_DURATION value.
+ */
+class ManualDurationHelper {
+    private Context mContext;
+
+    ManualDurationHelper(@NonNull Context context) {
+        mContext = context;
+    }
+
+    int getZenDuration() {
+        return Settings.Secure.getInt(mContext.getContentResolver(), Settings.Secure.ZEN_DURATION,
+                0);
+    }
+
+    /**
+     * Generates a summary of the duration that manual DND will be on when turned on from
+     * quick settings, for example "Until you turn off" or "[number] hours", based on the given
+     * setting value.
+     */
+    public String getSummary() {
+        int zenDuration = getZenDuration();
+        String summary;
+        if (zenDuration < 0) {
+            summary = mContext.getString(R.string.zen_mode_duration_summary_always_prompt);
+        } else if (zenDuration == 0) {
+            summary = mContext.getString(R.string.zen_mode_duration_summary_forever);
+        } else {
+            if (zenDuration >= 60) {
+                MessageFormat msgFormat = new MessageFormat(
+                        mContext.getString(R.string.zen_mode_duration_summary_time_hours),
+                        Locale.getDefault());
+                Map<String, Object> msgArgs = new HashMap<>();
+                msgArgs.put("count", zenDuration / 60);
+                summary = msgFormat.format(msgArgs);
+            } else {
+                MessageFormat msgFormat = new MessageFormat(
+                        mContext.getString(R.string.zen_mode_duration_summary_time_minutes),
+                        Locale.getDefault());
+                Map<String, Object> msgArgs = new HashMap<>();
+                msgArgs.put("count", zenDuration);
+                summary = msgFormat.format(msgArgs);
+            }
+        }
+        return summary;
+    }
+
+    SettingsObserver makeSettingsObserver(@NonNull AbstractZenModePreferenceController controller) {
+        return new SettingsObserver(controller);
+    }
+
+    final class SettingsObserver extends ContentObserver {
+        private static final Uri ZEN_MODE_DURATION_URI = Settings.Secure.getUriFor(
+                Settings.Secure.ZEN_DURATION);
+
+        private final AbstractZenModePreferenceController mPrefController;
+        private Preference mPreference;
+
+        /**
+         * Create a settings observer attached to the provided PreferenceController, whose
+         * updateState method should be called onChange.
+         */
+        SettingsObserver(@NonNull AbstractZenModePreferenceController prefController) {
+            super(mContext.getMainExecutor(), 0);
+            mPrefController = prefController;
+        }
+
+        void setPreference(Preference preference) {
+            mPreference = preference;
+        }
+
+        public void register() {
+            mContext.getContentResolver().registerContentObserver(ZEN_MODE_DURATION_URI, false,
+                    this);
+        }
+
+        public void unregister() {
+            mContext.getContentResolver().unregisterContentObserver(this);
+        }
+
+        @Override
+        public void onChange(boolean selfChange, @Nullable Uri uri) {
+            super.onChange(selfChange, uri);
+            if (ZEN_MODE_DURATION_URI.equals(uri) && mPreference != null) {
+                mPrefController.updateState(mPreference);
+            }
+        }
+    }
+}
diff --git a/src/com/android/settings/notification/modes/ManualDurationPreferenceController.java b/src/com/android/settings/notification/modes/ManualDurationPreferenceController.java
new file mode 100644
index 0000000..073f8ab
--- /dev/null
+++ b/src/com/android/settings/notification/modes/ManualDurationPreferenceController.java
@@ -0,0 +1,86 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.notification.modes;
+
+import android.content.Context;
+
+import androidx.fragment.app.Fragment;
+import androidx.preference.Preference;
+import androidx.preference.PreferenceScreen;
+
+import com.android.settings.notification.zen.SettingsZenDurationDialog;
+import com.android.settingslib.notification.modes.ZenMode;
+import com.android.settingslib.notification.modes.ZenModesBackend;
+
+public class ManualDurationPreferenceController extends AbstractZenModePreferenceController {
+    private static final String TAG = "QsDurationPrefController";
+
+    private final Fragment mParent;
+    private final ManualDurationHelper mDurationHelper;
+    private final ManualDurationHelper.SettingsObserver mSettingsObserver;
+
+    ManualDurationPreferenceController(Context context, String key, Fragment parent,
+            ZenModesBackend backend) {
+        super(context, key, backend);
+        mParent = parent;
+        mDurationHelper = new ManualDurationHelper(context);
+        mSettingsObserver = mDurationHelper.makeSettingsObserver(this);
+    }
+
+    @Override
+    public boolean isAvailable(ZenMode zenMode) {
+        if (!super.isAvailable(zenMode)) {
+            return false;
+        }
+        return zenMode.isManualDnd();
+    }
+
+    // Called by parent fragment onAttach().
+    void registerSettingsObserver() {
+        mSettingsObserver.register();
+    }
+
+    // Called by parent fragment onDetach().
+    void unregisterSettingsObserver() {
+        mSettingsObserver.unregister();
+    }
+
+    @Override
+    public void displayPreference(PreferenceScreen screen) {
+        super.displayPreference(screen);
+        Preference pref = screen.findPreference(getPreferenceKey());
+        if (pref != null) {
+            mSettingsObserver.setPreference(pref);
+        }
+    }
+
+    @Override
+    public void updateState(Preference preference, ZenMode unusedZenMode) {
+        // This controller is a link between a Settings value (ZEN_DURATION) and the manual DND
+        // mode. The status of the zen mode object itself doesn't affect the preference
+        // value, as that comes from settings; that value from settings will determine the
+        // condition that is attached to the mode on manual activation. Thus we ignore the actual
+        // zen mode value provided here.
+        preference.setSummary(mDurationHelper.getSummary());
+        preference.setOnPreferenceClickListener(pref -> {
+            // The new setting value is set by the dialog, so we don't need to do it here.
+            final SettingsZenDurationDialog durationDialog = new SettingsZenDurationDialog();
+            durationDialog.show(mParent.getParentFragmentManager(), TAG);
+            return true;
+        });
+    }
+}
diff --git a/src/com/android/settings/notification/modes/ZenModeButtonPreferenceController.java b/src/com/android/settings/notification/modes/ZenModeButtonPreferenceController.java
index 4a99b33..6b84414 100644
--- a/src/com/android/settings/notification/modes/ZenModeButtonPreferenceController.java
+++ b/src/com/android/settings/notification/modes/ZenModeButtonPreferenceController.java
@@ -18,21 +18,32 @@
 
 import android.annotation.NonNull;
 import android.content.Context;
+import android.provider.Settings;
 import android.widget.Button;
 
+import androidx.fragment.app.Fragment;
 import androidx.preference.Preference;
 
 import com.android.settings.R;
+import com.android.settings.notification.SettingsEnableZenModeDialog;
 import com.android.settingslib.notification.modes.ZenMode;
 import com.android.settingslib.notification.modes.ZenModesBackend;
 import com.android.settingslib.widget.LayoutPreference;
 
+import java.time.Duration;
+
 class ZenModeButtonPreferenceController extends AbstractZenModePreferenceController {
+    private static final String TAG = "ZenModeButtonPrefController";
 
     private Button mZenButton;
+    private Fragment mParent;
+    private ManualDurationHelper mDurationHelper;
 
-    public ZenModeButtonPreferenceController(Context context, String key, ZenModesBackend backend) {
+    ZenModeButtonPreferenceController(Context context, String key, Fragment parent,
+            ZenModesBackend backend) {
         super(context, key, backend);
+        mParent = parent;
+        mDurationHelper = new ManualDurationHelper(context);
     }
 
     @Override
@@ -49,7 +60,23 @@
             if (zenMode.isActive()) {
                 mBackend.deactivateMode(zenMode);
             } else {
-                mBackend.activateMode(zenMode, null);
+                if (zenMode.isManualDnd()) {
+                    // if manual DND, potentially ask for or use desired duration
+                    int zenDuration = mDurationHelper.getZenDuration();
+                    switch (zenDuration) {
+                        case Settings.Secure.ZEN_DURATION_PROMPT:
+                            new SettingsEnableZenModeDialog().show(
+                                    mParent.getParentFragmentManager(), TAG);
+                            break;
+                        case Settings.Secure.ZEN_DURATION_FOREVER:
+                            mBackend.activateMode(zenMode, null);
+                            break;
+                        default:
+                            mBackend.activateMode(zenMode, Duration.ofMinutes(zenDuration));
+                    }
+                } else {
+                    mBackend.activateMode(zenMode, null);
+                }
             }
         });
         if (zenMode.isActive()) {
diff --git a/src/com/android/settings/notification/modes/ZenModeFragment.java b/src/com/android/settings/notification/modes/ZenModeFragment.java
index 63ed839..67815b1 100644
--- a/src/com/android/settings/notification/modes/ZenModeFragment.java
+++ b/src/com/android/settings/notification/modes/ZenModeFragment.java
@@ -35,7 +35,6 @@
 import java.util.List;
 
 public class ZenModeFragment extends ZenModeFragmentBase {
-
     // for mode deletion menu
     private static final int DELETE_MODE = 1;
 
@@ -48,7 +47,8 @@
     protected List<AbstractPreferenceController> createPreferenceControllers(Context context) {
         List<AbstractPreferenceController> prefControllers = new ArrayList<>();
         prefControllers.add(new ZenModeHeaderController(context, "header", this, mBackend));
-        prefControllers.add(new ZenModeButtonPreferenceController(context, "activate", mBackend));
+        prefControllers.add(
+                new ZenModeButtonPreferenceController(context, "activate", this, mBackend));
         prefControllers.add(new ZenModeActionsPreferenceController(context, "actions", mBackend));
         prefControllers.add(new ZenModePeopleLinkPreferenceController(
                 context, "zen_mode_people", mBackend, mHelperBackend));
@@ -64,10 +64,20 @@
                 "zen_automatic_trigger_category", this, mBackend));
         prefControllers.add(new InterruptionFilterPreferenceController(
                 context, "allow_filtering", mBackend));
+        prefControllers.add(new ManualDurationPreferenceController(
+                context, "mode_manual_duration", this, mBackend));
         return prefControllers;
     }
 
     @Override
+    public void onAttach(@NonNull Context context) {
+        super.onAttach(context);
+
+        // allow duration preference controller to listen for settings changes
+        use(ManualDurationPreferenceController.class).registerSettingsObserver();
+    }
+
+    @Override
     public void onStart() {
         super.onStart();
 
@@ -79,6 +89,12 @@
     }
 
     @Override
+    public void onDetach() {
+        use(ManualDurationPreferenceController.class).unregisterSettingsObserver();
+        super.onDetach();
+    }
+
+    @Override
     public int getMetricsCategory() {
         // TODO: b/332937635 - make this the correct metrics category
         return SettingsEnums.NOTIFICATION_ZEN_MODE_AUTOMATION;
diff --git a/src/com/android/settings/search/SearchFeatureProviderImpl.java b/src/com/android/settings/search/SearchFeatureProviderImpl.java
deleted file mode 100644
index 3a62ddf..0000000
--- a/src/com/android/settings/search/SearchFeatureProviderImpl.java
+++ /dev/null
@@ -1,81 +0,0 @@
-/*
- * Copyright (C) 2017 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- *
- */
-
-package com.android.settings.search;
-
-import android.content.Context;
-import android.content.Intent;
-import android.net.Uri;
-import android.provider.Settings;
-import android.text.TextUtils;
-
-import androidx.annotation.NonNull;
-
-import com.android.settingslib.search.SearchIndexableResources;
-import com.android.settingslib.search.SearchIndexableResourcesMobile;
-
-/**
- * FeatureProvider for the refactored search code.
- */
-public class SearchFeatureProviderImpl implements SearchFeatureProvider {
-
-    private SearchIndexableResources mSearchIndexableResources;
-
-    @Override
-    public void verifyLaunchSearchResultPageCaller(@NonNull Context context,
-            @NonNull String callerPackage) {
-        if (TextUtils.isEmpty(callerPackage)) {
-            throw new IllegalArgumentException("ExternalSettingsTrampoline intents "
-                    + "must be called with startActivityForResult");
-        }
-        final boolean isSettingsPackage = TextUtils.equals(callerPackage, context.getPackageName())
-                || TextUtils.equals(getSettingsIntelligencePkgName(context), callerPackage);
-        final boolean isAllowlistedPackage = isSignatureAllowlisted(context, callerPackage);
-        if (isSettingsPackage || isAllowlistedPackage) {
-            return;
-        }
-        throw new SecurityException("Search result intents must be called with from an "
-                + "allowlisted package.");
-    }
-
-    @Override
-    public SearchIndexableResources getSearchIndexableResources() {
-        if (mSearchIndexableResources == null) {
-            mSearchIndexableResources = new SearchIndexableResourcesMobile();
-        }
-        return mSearchIndexableResources;
-    }
-
-    @Override
-    public Intent buildSearchIntent(Context context, int pageId) {
-        return new Intent(Settings.ACTION_APP_SEARCH_SETTINGS)
-                .setPackage(getSettingsIntelligencePkgName(context))
-                .putExtra(Intent.EXTRA_REFERRER, buildReferrer(context, pageId));
-    }
-
-    protected boolean isSignatureAllowlisted(Context context, String callerPackage) {
-        return false;
-    }
-
-    private static Uri buildReferrer(Context context, int pageId) {
-        return new Uri.Builder()
-                .scheme("android-app")
-                .authority(context.getPackageName())
-                .path(String.valueOf(pageId))
-                .build();
-    }
-}
diff --git a/src/com/android/settings/search/SearchFeatureProviderImpl.kt b/src/com/android/settings/search/SearchFeatureProviderImpl.kt
new file mode 100644
index 0000000..2ea9910
--- /dev/null
+++ b/src/com/android/settings/search/SearchFeatureProviderImpl.kt
@@ -0,0 +1,64 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.settings.search
+
+import android.content.Context
+import android.content.Intent
+import android.net.Uri
+import android.provider.Settings
+import com.android.settings.search.SearchIndexableResourcesFactory.createSearchIndexableResources
+import com.android.settingslib.search.SearchIndexableResources
+
+/** FeatureProvider for the refactored search code. */
+open class SearchFeatureProviderImpl : SearchFeatureProvider {
+    private val lazySearchIndexableResources by lazy { createSearchIndexableResources() }
+
+    override fun verifyLaunchSearchResultPageCaller(context: Context, callerPackage: String) {
+        require(callerPackage.isNotEmpty()) {
+            "ExternalSettingsTrampoline intents must be called with startActivityForResult"
+        }
+        val isSettingsPackage = callerPackage == context.packageName
+        if (isSettingsPackage ||
+            callerPackage == getSettingsIntelligencePkgName(context) ||
+            isSignatureAllowlisted(context, callerPackage)) {
+            return
+        }
+        throw SecurityException(
+            "Search result intents must be called with from an allowlisted package.")
+    }
+
+    override fun getSearchIndexableResources(): SearchIndexableResources =
+        lazySearchIndexableResources
+
+    override fun buildSearchIntent(context: Context, pageId: Int): Intent =
+        Intent(Settings.ACTION_APP_SEARCH_SETTINGS)
+            .setPackage(getSettingsIntelligencePkgName(context))
+            .putExtra(Intent.EXTRA_REFERRER, buildReferrer(context, pageId))
+
+    protected open fun isSignatureAllowlisted(context: Context, callerPackage: String): Boolean =
+        false
+
+    companion object {
+        private fun buildReferrer(context: Context, pageId: Int): Uri =
+            Uri.Builder()
+                .scheme("android-app")
+                .authority(context.packageName)
+                .path(pageId.toString())
+                .build()
+    }
+}
diff --git a/src/com/android/settings/search/SearchIndexableResourcesFactory.java b/src/com/android/settings/search/SearchIndexableResourcesFactory.java
new file mode 100644
index 0000000..25f34ed
--- /dev/null
+++ b/src/com/android/settings/search/SearchIndexableResourcesFactory.java
@@ -0,0 +1,34 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.search;
+
+import androidx.annotation.NonNull;
+
+import com.android.settingslib.search.SearchIndexableResources;
+import com.android.settingslib.search.SearchIndexableResourcesMobile;
+
+/**
+ * Creates the {@link SearchIndexableResourcesMobile}.
+ * <p>
+ * Since this class is generated by annotation processor, so it can only be created in Java now.
+ */
+class SearchIndexableResourcesFactory {
+    @NonNull
+    static SearchIndexableResources createSearchIndexableResources() {
+        return new SearchIndexableResourcesMobile();
+    }
+}
diff --git a/src/com/android/settings/spa/SpaBridgeActivity.kt b/src/com/android/settings/spa/SpaBridgeActivity.kt
index 61d8f51..d579fdf 100644
--- a/src/com/android/settings/spa/SpaBridgeActivity.kt
+++ b/src/com/android/settings/spa/SpaBridgeActivity.kt
@@ -17,20 +17,18 @@
 package com.android.settings.spa
 
 import android.app.Activity
-import android.content.Intent
+import android.content.pm.PackageManager.ComponentInfoFlags
+import android.content.pm.PackageManager.GET_META_DATA
 import android.os.Bundle
-import com.android.settings.activityembedding.ActivityEmbeddingUtils
-import com.android.settings.activityembedding.EmbeddedDeepLinkUtils.tryStartMultiPaneDeepLink
-import com.android.settings.spa.SpaDestination.Companion.getDestination
-import com.android.settingslib.spa.framework.util.SESSION_EXTERNAL
-import com.android.settingslib.spa.framework.util.appendSpaParams
+import androidx.annotation.VisibleForTesting
+import com.android.settings.SettingsActivity.META_DATA_KEY_HIGHLIGHT_MENU_KEY
 
 /**
  * Activity used as a bridge to [SpaActivity].
  *
  * Since [SpaActivity] is not exported, [SpaActivity] could not be the target activity of
- * <activity-alias>, otherwise all its pages will be exported.
- * So need this bridge activity to sit in the middle of <activity-alias> and [SpaActivity].
+ * <activity-alias>, otherwise all its pages will be exported. So need this bridge activity to sit
+ * in the middle of <activity-alias> and [SpaActivity].
  */
 class SpaBridgeActivity : Activity() {
     override fun onCreate(savedInstanceState: Bundle?) {
@@ -41,17 +39,28 @@
 
     companion object {
         fun Activity.startSpaActivityFromBridge(destinationFactory: (String) -> String? = { it }) {
-            val (destination, highlightMenuKey) = getDestination(destinationFactory) ?: return
-            val intent = Intent(this, SpaActivity::class.java)
-                .appendSpaParams(
-                    destination = destination,
-                    sessionName = SESSION_EXTERNAL,
-                )
-            if (!ActivityEmbeddingUtils.isEmbeddingActivityEnabled(this) ||
-                !tryStartMultiPaneDeepLink(intent, highlightMenuKey)
-            ) {
-                startActivity(intent)
-            }
+            getDestination(destinationFactory)?.startFromExportedActivity(this)
         }
+
+        @VisibleForTesting
+        fun Activity.getDestination(
+            destinationFactory: (String) -> String? = { it },
+        ): SpaDestination? {
+            val metaData =
+                packageManager
+                    .getActivityInfo(componentName, ComponentInfoFlags.of(GET_META_DATA.toLong()))
+                    .metaData
+            val destination = metaData.getString(META_DATA_KEY_DESTINATION)
+            if (destination.isNullOrBlank()) return null
+            val finalDestination = destinationFactory(destination)
+            if (finalDestination.isNullOrBlank()) return null
+            return SpaDestination(
+                destination = finalDestination,
+                highlightMenuKey = metaData.getString(META_DATA_KEY_HIGHLIGHT_MENU_KEY),
+            )
+        }
+
+        @VisibleForTesting
+        const val META_DATA_KEY_DESTINATION = "com.android.settings.spa.DESTINATION"
     }
 }
diff --git a/src/com/android/settings/spa/SpaDestination.kt b/src/com/android/settings/spa/SpaDestination.kt
index bdec1d8..cb20c37 100644
--- a/src/com/android/settings/spa/SpaDestination.kt
+++ b/src/com/android/settings/spa/SpaDestination.kt
@@ -17,33 +17,26 @@
 package com.android.settings.spa
 
 import android.app.Activity
-import android.content.pm.PackageManager
-import androidx.annotation.VisibleForTesting
-import com.android.settings.SettingsActivity.META_DATA_KEY_HIGHLIGHT_MENU_KEY
+import android.content.Intent
+import com.android.settings.activityembedding.ActivityEmbeddingUtils
+import com.android.settings.activityembedding.EmbeddedDeepLinkUtils.tryStartMultiPaneDeepLink
+import com.android.settingslib.spa.framework.util.SESSION_EXTERNAL
+import com.android.settingslib.spa.framework.util.appendSpaParams
 
 data class SpaDestination(
     val destination: String,
     val highlightMenuKey: String?,
 ) {
-    companion object {
-        fun Activity.getDestination(
-            destinationFactory: (String) -> String? = { it },
-        ): SpaDestination? {
-            val metaData = packageManager.getActivityInfo(
-                componentName,
-                PackageManager.ComponentInfoFlags.of(PackageManager.GET_META_DATA.toLong())
-            ).metaData
-            val destination = metaData.getString(META_DATA_KEY_DESTINATION)
-            if (destination.isNullOrBlank()) return null
-            val finalDestination = destinationFactory(destination)
-            if (finalDestination.isNullOrBlank()) return null
-            return SpaDestination(
-                destination = finalDestination,
-                highlightMenuKey = metaData.getString(META_DATA_KEY_HIGHLIGHT_MENU_KEY),
+    fun startFromExportedActivity(activity: Activity) {
+        val intent = Intent(activity, SpaActivity::class.java)
+            .appendSpaParams(
+                destination = destination,
+                sessionName = SESSION_EXTERNAL,
             )
+        if (!ActivityEmbeddingUtils.isEmbeddingActivityEnabled(activity) ||
+            !activity.tryStartMultiPaneDeepLink(intent, highlightMenuKey)
+        ) {
+            activity.startActivity(intent)
         }
-
-        @VisibleForTesting
-        const val META_DATA_KEY_DESTINATION = "com.android.settings.spa.DESTINATION"
     }
 }
diff --git a/src/com/android/settings/spa/network/NetworkCellularGroupProvider.kt b/src/com/android/settings/spa/network/NetworkCellularGroupProvider.kt
index 05a8f6a..873a2c3 100644
--- a/src/com/android/settings/spa/network/NetworkCellularGroupProvider.kt
+++ b/src/com/android/settings/spa/network/NetworkCellularGroupProvider.kt
@@ -279,14 +279,14 @@
 ) {
     CreatePrimarySimListPreference(
         stringResource(id = R.string.primary_sim_calls_title),
-        primarySimInfo.callsAndSmsList,
+        primarySimInfo.callsList,
         callsSelectedId,
         ImageVector.vectorResource(R.drawable.ic_phone),
         actionSetCalls
     )
     CreatePrimarySimListPreference(
         stringResource(id = R.string.primary_sim_texts_title),
-        primarySimInfo.callsAndSmsList,
+        primarySimInfo.smsList,
         textsSelectedId,
         Icons.AutoMirrored.Outlined.Message,
         actionSetTexts
diff --git a/src/com/android/settings/spa/network/PrimarySimRepository.kt b/src/com/android/settings/spa/network/PrimarySimRepository.kt
index e34a29e..420f870 100644
--- a/src/com/android/settings/spa/network/PrimarySimRepository.kt
+++ b/src/com/android/settings/spa/network/PrimarySimRepository.kt
@@ -22,12 +22,14 @@
 import android.util.Log
 import com.android.settings.R
 import com.android.settings.network.SubscriptionUtil
+import com.android.settings.network.telephony.CarrierConfigRepository
 import com.android.settingslib.spa.widget.preference.ListPreferenceOption
 
 class PrimarySimRepository(private val context: Context) {
 
     data class PrimarySimInfo(
-        val callsAndSmsList: List<ListPreferenceOption>,
+        val callsList: List<ListPreferenceOption>,
+        val smsList: List<ListPreferenceOption>,
         val dataList: List<ListPreferenceOption>,
     )
 
@@ -37,7 +39,8 @@
             return null
         }
 
-        val callsAndSmsList = mutableListOf<ListPreferenceOption>()
+        val callsList = mutableListOf<ListPreferenceOption>()
+        val smsList = mutableListOf<ListPreferenceOption>()
         val dataList = mutableListOf<ListPreferenceOption>()
         for (info in selectableSubscriptionInfoList) {
             val item = ListPreferenceOption(
@@ -45,15 +48,22 @@
                 text = "${info.displayName}",
                 summary = SubscriptionUtil.getBidiFormattedPhoneNumber(context, info) ?: "",
             )
-            callsAndSmsList += item
+            callsList += item
+            smsList += item
             dataList += item
         }
-        callsAndSmsList += ListPreferenceOption(
+
+        val askEveryTime = ListPreferenceOption(
             id = SubscriptionManager.INVALID_SUBSCRIPTION_ID,
             text = context.getString(R.string.sim_calls_ask_first_prefs_title),
         )
+        callsList += askEveryTime
+        if (context.resources
+                .getBoolean(com.android.internal.R.bool.config_sms_ask_every_time_support)) {
+            smsList += askEveryTime
+        }
 
-        return PrimarySimInfo(callsAndSmsList, dataList)
+        return PrimarySimInfo(callsList, smsList, dataList)
     }
 
     private companion object {
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingCallAudioDialogFragmentTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingCallAudioDialogFragmentTest.java
index 7f0c1c9..51ed899 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingCallAudioDialogFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingCallAudioDialogFragmentTest.java
@@ -23,6 +23,7 @@
 import android.app.settings.SettingsEnums;
 import android.bluetooth.BluetoothAdapter;
 import android.bluetooth.BluetoothStatusCodes;
+import android.os.Bundle;
 import android.platform.test.flag.junit.SetFlagsRule;
 
 import androidx.appcompat.app.AlertDialog;
@@ -78,10 +79,6 @@
                 BluetoothStatusCodes.FEATURE_SUPPORTED);
         shadowBluetoothAdapter.setIsLeAudioBroadcastAssistantSupported(
                 BluetoothStatusCodes.FEATURE_SUPPORTED);
-        mFragment = new AudioSharingCallAudioDialogFragment();
-        mParent = new Fragment();
-        FragmentController.setupFragment(
-                mParent, FragmentActivity.class, /* containerViewId= */ 0, /* bundle= */ null);
     }
 
     @After
@@ -91,6 +88,7 @@
 
     @Test
     public void getMetricsCategory_correctValue() {
+        mFragment = new AudioSharingCallAudioDialogFragment();
         assertThat(mFragment.getMetricsCategory())
                 .isEqualTo(SettingsEnums.DIALOG_AUDIO_SHARING_CALL_AUDIO);
     }
@@ -98,21 +96,52 @@
     @Test
     public void onCreateDialog_flagOff_dialogNotExist() {
         mSetFlagsRule.disableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
-        mFragment.show(mParent, new ArrayList<>(), (item) -> {});
+        mParent = new Fragment();
+        FragmentController.setupFragment(
+                mParent, FragmentActivity.class, /* containerViewId= */ 0, /* bundle= */ null);
+        AudioSharingCallAudioDialogFragment.show(mParent, new ArrayList<>(), (item) -> {});
         shadowMainLooper().idle();
         AlertDialog dialog = ShadowAlertDialogCompat.getLatestAlertDialog();
         assertThat(dialog).isNull();
     }
 
     @Test
+    public void onCreateDialog_unattachedFragment_dialogNotExist() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+        mParent = new Fragment();
+        AudioSharingCallAudioDialogFragment.show(mParent, new ArrayList<>(), (item) -> {});
+        shadowMainLooper().idle();
+        AlertDialog dialog = ShadowAlertDialogCompat.getLatestAlertDialog();
+        assertThat(dialog).isNull();
+    }
+
+    @Test
+    public void onCreateDialog_nullDeviceItems_showEmptyDialog() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+        mFragment = new AudioSharingCallAudioDialogFragment();
+        mFragment.setArguments(Bundle.EMPTY);
+        FragmentController.setupFragment(
+                mFragment, FragmentActivity.class, /* containerViewId= */ 0, /* bundle= */ null);
+        AlertDialog dialog = (AlertDialog) mFragment.onCreateDialog(Bundle.EMPTY);
+        dialog.show();
+        shadowMainLooper().idle();
+        assertThat(dialog.isShowing()).isTrue();
+        assertThat(dialog.getListView()).isNull();
+    }
+
+    @Test
     public void onCreateDialog_showCorrectItems() {
         mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+        mParent = new Fragment();
+        FragmentController.setupFragment(
+                mParent, FragmentActivity.class, /* containerViewId= */ 0, /* bundle= */ null);
         ArrayList<AudioSharingDeviceItem> deviceItemList = new ArrayList<>();
         deviceItemList.add(TEST_DEVICE_ITEM1);
         deviceItemList.add(TEST_DEVICE_ITEM2);
-        mFragment.show(mParent, deviceItemList, (item) -> {});
+        AudioSharingCallAudioDialogFragment.show(mParent, deviceItemList, (item) -> {});
         shadowMainLooper().idle();
         AlertDialog dialog = ShadowAlertDialogCompat.getLatestAlertDialog();
+        assertThat(dialog).isNotNull();
         assertThat(dialog.getListView().getCount()).isEqualTo(2);
     }
 }
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingConfirmDialogFragmentTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingConfirmDialogFragmentTest.java
index 32f666f..9aa00fb 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingConfirmDialogFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingConfirmDialogFragmentTest.java
@@ -96,6 +96,15 @@
     }
 
     @Test
+    public void onCreateDialog_unattachedFragment_dialogNotExist() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+        AudioSharingConfirmDialogFragment.show(new Fragment());
+        shadowMainLooper().idle();
+        AlertDialog dialog = ShadowAlertDialogCompat.getLatestAlertDialog();
+        assertThat(dialog).isNull();
+    }
+
+    @Test
     public void onCreateDialog_flagOn_showDialog() {
         mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
         AudioSharingConfirmDialogFragment.show(mParent);
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogFragmentTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogFragmentTest.java
index 39709c1..20c225c 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogFragmentTest.java
@@ -130,7 +130,16 @@
         AudioSharingDialogFragment.show(
                 mParent, new ArrayList<>(), EMPTY_EVENT_LISTENER, TEST_EVENT_DATA_LIST);
         shadowMainLooper().idle();
+        AlertDialog dialog = ShadowAlertDialogCompat.getLatestAlertDialog();
+        assertThat(dialog).isNull();
+    }
 
+    @Test
+    public void onCreateDialog_unattachedFragment_dialogNotExist() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+        AudioSharingDialogFragment.show(
+                new Fragment(), new ArrayList<>(), EMPTY_EVENT_LISTENER, TEST_EVENT_DATA_LIST);
+        shadowMainLooper().idle();
         AlertDialog dialog = ShadowAlertDialogCompat.getLatestAlertDialog();
         assertThat(dialog).isNull();
     }
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogHandlerTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogHandlerTest.java
index 53c214b..e71e876 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogHandlerTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDialogHandlerTest.java
@@ -648,6 +648,15 @@
     }
 
     @Test
+    public void closeOpeningDialogsForLeaDevice_unattachedFragment_doNothing() {
+        mParentFragment = new Fragment();
+        mHandler = new AudioSharingDialogHandler(mContext, mParentFragment);
+        mHandler.closeOpeningDialogsForLeaDevice(mCachedDevice1);
+        shadowOf(Looper.getMainLooper()).idle();
+        verifyNoMoreInteractions(mFeatureFactory.metricsFeatureProvider);
+    }
+
+    @Test
     public void closeOpeningDialogsForLeaDevice_closeDisconnectDialog() {
         // Show disconnect dialog
         setUpBroadcast(true);
@@ -675,6 +684,15 @@
     }
 
     @Test
+    public void closeOpeningDialogsForNonLeaDevice_unattachedFragment_doNothing() {
+        mParentFragment = new Fragment();
+        mHandler = new AudioSharingDialogHandler(mContext, mParentFragment);
+        mHandler.closeOpeningDialogsForNonLeaDevice(mCachedDevice2);
+        shadowOf(Looper.getMainLooper()).idle();
+        verifyNoMoreInteractions(mFeatureFactory.metricsFeatureProvider);
+    }
+
+    @Test
     public void closeOpeningDialogsForNonLeaDevice_closeStopDialog() {
         // Show stop dialog
         setUpBroadcast(true);
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDisconnectDialogFragmentTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDisconnectDialogFragmentTest.java
index 6b984af..86b0d65 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDisconnectDialogFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingDisconnectDialogFragmentTest.java
@@ -154,6 +154,23 @@
     }
 
     @Test
+    public void onCreateDialog_unattachedFragment_dialogNotExist() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+        mDeviceItems = new ArrayList<>();
+        mDeviceItems.add(TEST_DEVICE_ITEM1);
+        mDeviceItems.add(TEST_DEVICE_ITEM2);
+        AudioSharingDisconnectDialogFragment.show(
+                new Fragment(),
+                mDeviceItems,
+                mCachedDevice3,
+                EMPTY_EVENT_LISTENER,
+                TEST_EVENT_DATA_LIST);
+        shadowMainLooper().idle();
+        AlertDialog dialog = ShadowAlertDialogCompat.getLatestAlertDialog();
+        assertThat(dialog).isNull();
+    }
+
+    @Test
     public void onCreateDialog_flagOn_dialogShowBtnForTwoDevices() {
         mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
         mDeviceItems = new ArrayList<>();
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingJoinDialogFragmentTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingJoinDialogFragmentTest.java
index 13ea965..2310d75 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingJoinDialogFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingJoinDialogFragmentTest.java
@@ -164,6 +164,20 @@
     }
 
     @Test
+    public void onCreateDialog_unattachedFragment_dialogNotExist() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+        AudioSharingJoinDialogFragment.show(
+                new Fragment(),
+                new ArrayList<>(),
+                mCachedDevice2,
+                EMPTY_EVENT_LISTENER,
+                TEST_EVENT_DATA_LIST);
+        shadowMainLooper().idle();
+        AlertDialog dialog = ShadowAlertDialogCompat.getLatestAlertDialog();
+        assertThat(dialog).isNull();
+    }
+
+    @Test
     public void onCreateDialog_flagOn_dialogShowTextForSingleDevice() {
         mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
         AudioSharingJoinDialogFragment.show(
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingStopDialogFragmentTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingStopDialogFragmentTest.java
index b6babfb..c0af09f 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingStopDialogFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/AudioSharingStopDialogFragmentTest.java
@@ -149,6 +149,20 @@
     }
 
     @Test
+    public void onCreateDialog_unattachedFragment_dialogNotExist() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+        AudioSharingStopDialogFragment.show(
+                new Fragment(),
+                ImmutableList.of(TEST_DEVICE_ITEM2),
+                mCachedDevice1,
+                EMPTY_EVENT_LISTENER,
+                TEST_EVENT_DATA_LIST);
+        shadowMainLooper().idle();
+        AlertDialog dialog = ShadowAlertDialogCompat.getLatestAlertDialog();
+        assertThat(dialog).isNull();
+    }
+
+    @Test
     public void onCreateDialog_oneDeviceInSharing_showDialogWithCorrectMessage() {
         mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
         AudioSharingStopDialogFragment.show(
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamButtonControllerTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamButtonControllerTest.java
index cbf1432..c6fb361 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamButtonControllerTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamButtonControllerTest.java
@@ -16,22 +16,36 @@
 
 package com.android.settings.connecteddevice.audiosharing.audiostreams;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+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.settings.SettingsEnums;
+import android.bluetooth.BluetoothDevice;
+import android.bluetooth.BluetoothLeBroadcastAssistant;
+import android.bluetooth.BluetoothLeBroadcastMetadata;
 import android.bluetooth.BluetoothLeBroadcastReceiveState;
 import android.content.Context;
 import android.view.View;
 
+import androidx.lifecycle.LifecycleOwner;
 import androidx.preference.PreferenceScreen;
 import androidx.test.core.app.ApplicationProvider;
 
 import com.android.settings.R;
 import com.android.settings.connecteddevice.audiosharing.audiostreams.testshadows.ShadowAudioStreamsHelper;
+import com.android.settings.testutils.FakeFeatureFactory;
 import com.android.settings.testutils.shadow.ShadowThreadUtils;
+import com.android.settingslib.bluetooth.LocalBluetoothLeBroadcastAssistant;
+import com.android.settingslib.core.lifecycle.Lifecycle;
 import com.android.settingslib.widget.ActionButtonsPreference;
 
 import org.junit.After;
@@ -39,14 +53,17 @@
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.annotation.Config;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.Executor;
 
 @RunWith(RobolectricTestRunner.class)
 @Config(
@@ -63,14 +80,23 @@
     @Mock private AudioStreamsHelper mAudioStreamsHelper;
     @Mock private PreferenceScreen mScreen;
     @Mock private BluetoothLeBroadcastReceiveState mBroadcastReceiveState;
+    @Mock private LocalBluetoothLeBroadcastAssistant mAssistant;
+    @Mock private AudioStreamsRepository mRepository;
     @Mock private ActionButtonsPreference mPreference;
+    private Lifecycle mLifecycle;
+    private LifecycleOwner mLifecycleOwner;
+    private FakeFeatureFactory mFeatureFactory;
     private AudioStreamButtonController mController;
 
     @Before
     public void setUp() {
         ShadowAudioStreamsHelper.setUseMock(mAudioStreamsHelper);
+        when(mAudioStreamsHelper.getLeBroadcastAssistant()).thenReturn(mAssistant);
+        mFeatureFactory = FakeFeatureFactory.setupForTest();
         mController = new AudioStreamButtonController(mContext, KEY);
         mController.init(BROADCAST_ID);
+        mLifecycleOwner = () -> mLifecycle;
+        mLifecycle = new Lifecycle(mLifecycleOwner);
         when(mScreen.findPreference(KEY)).thenReturn(mPreference);
         when(mPreference.getContext()).thenReturn(mContext);
         when(mPreference.setButton1Text(anyInt())).thenReturn(mPreference);
@@ -86,6 +112,40 @@
     }
 
     @Test
+    public void onStart_registerCallbacks() {
+        mController.onStart(mLifecycleOwner);
+        verify(mAssistant)
+                .registerServiceCallBack(
+                        any(Executor.class), any(BluetoothLeBroadcastAssistant.Callback.class));
+    }
+
+    @Test
+    public void onStart_profileNull_doNothing() {
+        when(mAudioStreamsHelper.getLeBroadcastAssistant()).thenReturn(null);
+        mController = new AudioStreamButtonController(mContext, KEY);
+        mController.onStart(mLifecycleOwner);
+        verify(mAssistant, never())
+                .registerServiceCallBack(
+                        any(Executor.class), any(BluetoothLeBroadcastAssistant.Callback.class));
+    }
+
+    @Test
+    public void onStop_unregisterCallbacks() {
+        mController.onStop(mLifecycleOwner);
+        verify(mAssistant)
+                .unregisterServiceCallBack(any(BluetoothLeBroadcastAssistant.Callback.class));
+    }
+
+    @Test
+    public void onStop_profileNull_doNothing() {
+        when(mAudioStreamsHelper.getLeBroadcastAssistant()).thenReturn(null);
+        mController = new AudioStreamButtonController(mContext, KEY);
+        mController.onStop(mLifecycleOwner);
+        verify(mAssistant, never())
+                .unregisterServiceCallBack(any(BluetoothLeBroadcastAssistant.Callback.class));
+    }
+
+    @Test
     public void testDisplayPreference_sourceConnected_setDisconnectButton() {
         when(mAudioStreamsHelper.getAllConnectedSources())
                 .thenReturn(List.of(mBroadcastReceiveState));
@@ -96,18 +156,133 @@
         verify(mPreference).setButton1Enabled(true);
         verify(mPreference).setButton1Text(R.string.audio_streams_disconnect);
         verify(mPreference).setButton1Icon(com.android.settings.R.drawable.ic_settings_close);
-        verify(mPreference).setButton1OnClickListener(any(View.OnClickListener.class));
+
+        ArgumentCaptor<View.OnClickListener> listenerCaptor =
+                ArgumentCaptor.forClass(View.OnClickListener.class);
+        verify(mPreference).setButton1OnClickListener(listenerCaptor.capture());
+        var listener = listenerCaptor.getValue();
+
+        assertThat(listener).isNotNull();
+        listener.onClick(mock(View.class));
+        verify(mAudioStreamsHelper).removeSource(BROADCAST_ID);
+        verify(mPreference).setButton1Enabled(false);
+        verify(mFeatureFactory.metricsFeatureProvider)
+                .action(any(), eq(SettingsEnums.ACTION_AUDIO_STREAM_LEAVE_BUTTON_CLICK));
     }
 
     @Test
     public void testDisplayPreference_sourceNotConnected_setConnectButton() {
         when(mAudioStreamsHelper.getAllConnectedSources()).thenReturn(Collections.emptyList());
+        mController.setAudioStreamsRepositoryForTesting(mRepository);
+        var metadataToRejoin = mock(BluetoothLeBroadcastMetadata.class);
+        when(mRepository.getSavedMetadata(any(), anyInt())).thenReturn(metadataToRejoin);
 
         mController.displayPreference(mScreen);
 
         verify(mPreference).setButton1Enabled(true);
         verify(mPreference).setButton1Text(R.string.audio_streams_connect);
         verify(mPreference).setButton1Icon(com.android.settings.R.drawable.ic_add_24dp);
-        verify(mPreference).setButton1OnClickListener(any(View.OnClickListener.class));
+
+        ArgumentCaptor<View.OnClickListener> listenerCaptor =
+                ArgumentCaptor.forClass(View.OnClickListener.class);
+        verify(mPreference).setButton1OnClickListener(listenerCaptor.capture());
+        var listener = listenerCaptor.getValue();
+
+        assertThat(listener).isNotNull();
+        listener.onClick(mock(View.class));
+        verify(mAudioStreamsHelper).addSource(metadataToRejoin);
+        verify(mPreference).setButton1Enabled(false);
+        verify(mFeatureFactory.metricsFeatureProvider)
+                .action(any(), eq(SettingsEnums.ACTION_AUDIO_STREAM_JOIN), anyInt());
+    }
+
+    @Test
+    public void testCallback_onSourceRemoved_updateButton() {
+        when(mAudioStreamsHelper.getAllConnectedSources()).thenReturn(Collections.emptyList());
+
+        mController.displayPreference(mScreen);
+        mController.mBroadcastAssistantCallback.onSourceRemoved(
+                mock(BluetoothDevice.class), /* sourceId= */ 0, /* reason= */ 0);
+
+        // Called twice, once in displayPreference, the other one in callback
+        verify(mPreference, times(2)).setButton1Enabled(true);
+        verify(mPreference, times(2)).setButton1Text(R.string.audio_streams_connect);
+        verify(mPreference, times(2)).setButton1Icon(com.android.settings.R.drawable.ic_add_24dp);
+    }
+
+    @Test
+    public void testCallback_onSourceRemovedFailed_updateButton() {
+        when(mAudioStreamsHelper.getAllConnectedSources())
+                .thenReturn(List.of(mBroadcastReceiveState));
+        when(mBroadcastReceiveState.getBroadcastId()).thenReturn(BROADCAST_ID);
+
+        mController.displayPreference(mScreen);
+        mController.mBroadcastAssistantCallback.onSourceRemoveFailed(
+                mock(BluetoothDevice.class), /* sourceId= */ 0, /* reason= */ 0);
+
+        verify(mFeatureFactory.metricsFeatureProvider)
+                .action(any(), eq(SettingsEnums.ACTION_AUDIO_STREAM_LEAVE_FAILED));
+
+        // Called twice, once in displayPreference, the other one in callback
+        verify(mPreference, times(2)).setButton1Enabled(true);
+        verify(mPreference, times(2)).setButton1Text(R.string.audio_streams_disconnect);
+        verify(mPreference, times(2))
+                .setButton1Icon(com.android.settings.R.drawable.ic_settings_close);
+    }
+
+    @Test
+    public void testCallback_onReceiveStateChanged_updateButton() {
+        when(mAudioStreamsHelper.getAllConnectedSources())
+                .thenReturn(List.of(mBroadcastReceiveState));
+        when(mBroadcastReceiveState.getBroadcastId()).thenReturn(BROADCAST_ID);
+        BluetoothLeBroadcastReceiveState state = mock(BluetoothLeBroadcastReceiveState.class);
+        List<Long> bisSyncState = new ArrayList<>();
+        bisSyncState.add(1L);
+        when(state.getBisSyncState()).thenReturn(bisSyncState);
+
+        mController.displayPreference(mScreen);
+        mController.mBroadcastAssistantCallback.onReceiveStateChanged(
+                mock(BluetoothDevice.class), /* sourceId= */ 0, state);
+
+        verify(mFeatureFactory.metricsFeatureProvider)
+                .action(any(), eq(SettingsEnums.ACTION_AUDIO_STREAM_JOIN_SUCCEED), anyInt());
+
+        // Called twice, once in displayPreference, the other one in callback
+        verify(mPreference, times(2)).setButton1Enabled(true);
+        verify(mPreference, times(2)).setButton1Text(R.string.audio_streams_disconnect);
+        verify(mPreference, times(2))
+                .setButton1Icon(com.android.settings.R.drawable.ic_settings_close);
+    }
+
+    @Test
+    public void testCallback_onSourceAddFailed_updateButton() {
+        when(mAudioStreamsHelper.getAllConnectedSources()).thenReturn(Collections.emptyList());
+
+        mController.displayPreference(mScreen);
+        mController.mBroadcastAssistantCallback.onSourceAddFailed(
+                mock(BluetoothDevice.class),
+                mock(BluetoothLeBroadcastMetadata.class),
+                /* reason= */ 0);
+
+        verify(mFeatureFactory.metricsFeatureProvider)
+                .action(any(), eq(SettingsEnums.ACTION_AUDIO_STREAM_JOIN_FAILED_OTHER), anyInt());
+
+        // Called twice, once in displayPreference, the other one in callback
+        verify(mPreference, times(2)).setButton1Enabled(true);
+        verify(mPreference, times(2)).setButton1Text(R.string.audio_streams_connect);
+        verify(mPreference, times(2)).setButton1Icon(com.android.settings.R.drawable.ic_add_24dp);
+    }
+
+    @Test
+    public void testCallback_onSourceLost_updateButton() {
+        when(mAudioStreamsHelper.getAllConnectedSources()).thenReturn(Collections.emptyList());
+
+        mController.displayPreference(mScreen);
+        mController.mBroadcastAssistantCallback.onSourceLost(/* broadcastId= */ 0);
+
+        // Called twice, once in displayPreference, the other one in callback
+        verify(mPreference, times(2)).setButton1Enabled(true);
+        verify(mPreference, times(2)).setButton1Text(R.string.audio_streams_connect);
+        verify(mPreference, times(2)).setButton1Icon(com.android.settings.R.drawable.ic_add_24dp);
     }
 }
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamConfirmDialogActivityTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamConfirmDialogActivityTest.java
index e967a12..1f05cbb 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamConfirmDialogActivityTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamConfirmDialogActivityTest.java
@@ -18,28 +18,132 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.bluetooth.BluetoothAdapter;
+import android.bluetooth.BluetoothStatusCodes;
+import android.platform.test.flag.junit.SetFlagsRule;
+
+import com.android.settings.testutils.shadow.ShadowBluetoothAdapter;
+import com.android.settings.testutils.shadow.ShadowBluetoothUtils;
+import com.android.settingslib.bluetooth.LocalBluetoothLeBroadcast;
+import com.android.settingslib.bluetooth.LocalBluetoothLeBroadcastAssistant;
+import com.android.settingslib.bluetooth.LocalBluetoothManager;
+import com.android.settingslib.bluetooth.LocalBluetoothProfileManager;
+import com.android.settingslib.bluetooth.VolumeControlProfile;
+import com.android.settingslib.flags.Flags;
+
+import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
 import org.robolectric.Robolectric;
 import org.robolectric.RobolectricTestRunner;
+import org.robolectric.annotation.Config;
+import org.robolectric.shadow.api.Shadow;
 
 @RunWith(RobolectricTestRunner.class)
+@Config(
+        shadows = {
+            ShadowBluetoothAdapter.class,
+            ShadowBluetoothUtils.class,
+        })
 public class AudioStreamConfirmDialogActivityTest {
+    @Rule public final MockitoRule mocks = MockitoJUnit.rule();
+    @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+    @Mock private LocalBluetoothManager mLocalBluetoothManager;
+    @Mock private LocalBluetoothProfileManager mLocalBluetoothProfileManager;
+    @Mock private LocalBluetoothLeBroadcast mBroadcast;
+    @Mock private LocalBluetoothLeBroadcastAssistant mAssistant;
+    @Mock private VolumeControlProfile mVolumeControl;
+    private ShadowBluetoothAdapter mShadowBluetoothAdapter;
     private AudioStreamConfirmDialogActivity mActivity;
 
     @Before
     public void setUp() {
-        mActivity = Robolectric.buildActivity(AudioStreamConfirmDialogActivity.class).get();
+        mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+        mShadowBluetoothAdapter = Shadow.extract(BluetoothAdapter.getDefaultAdapter());
+        mShadowBluetoothAdapter.setEnabled(true);
+        mShadowBluetoothAdapter.setIsLeAudioBroadcastSourceSupported(
+                BluetoothStatusCodes.FEATURE_SUPPORTED);
+        mShadowBluetoothAdapter.setIsLeAudioBroadcastAssistantSupported(
+                BluetoothStatusCodes.FEATURE_SUPPORTED);
+        ShadowBluetoothUtils.sLocalBluetoothManager = mLocalBluetoothManager;
+        when(mLocalBluetoothManager.getProfileManager()).thenReturn(mLocalBluetoothProfileManager);
+        when(mLocalBluetoothProfileManager.getLeAudioBroadcastProfile()).thenReturn(mBroadcast);
+        when(mLocalBluetoothProfileManager.getLeAudioBroadcastAssistantProfile())
+                .thenReturn(mAssistant);
+        when(mLocalBluetoothProfileManager.getVolumeControlProfile()).thenReturn(mVolumeControl);
+        when(mBroadcast.isProfileReady()).thenReturn(true);
+        when(mAssistant.isProfileReady()).thenReturn(true);
+        when(mVolumeControl.isProfileReady()).thenReturn(true);
+    }
+
+    @After
+    public void tearDown() {
+        ShadowBluetoothUtils.reset();
     }
 
     @Test
     public void isValidFragment_returnsTrue() {
+        mActivity = Robolectric.setupActivity(AudioStreamConfirmDialogActivity.class);
         assertThat(mActivity.isValidFragment(AudioStreamConfirmDialog.class.getName())).isTrue();
     }
 
     @Test
     public void isValidFragment_returnsFalse() {
+        mActivity = Robolectric.setupActivity(AudioStreamConfirmDialogActivity.class);
         assertThat(mActivity.isValidFragment("")).isFalse();
     }
+
+    @Test
+    public void isToolbarEnabled_returnsFalse() {
+        mActivity = Robolectric.setupActivity(AudioStreamConfirmDialogActivity.class);
+        assertThat(mActivity.isToolbarEnabled()).isFalse();
+    }
+
+    @Test
+    public void setupActivity_serviceNotReady_registerCallback() {
+        when(mBroadcast.isProfileReady()).thenReturn(false);
+        mActivity = Robolectric.setupActivity(AudioStreamConfirmDialogActivity.class);
+
+        verify(mLocalBluetoothProfileManager).addServiceListener(any());
+    }
+
+    @Test
+    public void setupActivity_serviceNotReady_registerCallback_onServiceCallback() {
+        when(mBroadcast.isProfileReady()).thenReturn(false);
+        mActivity = Robolectric.setupActivity(AudioStreamConfirmDialogActivity.class);
+
+        verify(mLocalBluetoothProfileManager).addServiceListener(any());
+
+        when(mBroadcast.isProfileReady()).thenReturn(true);
+        mActivity.onServiceConnected();
+        verify(mLocalBluetoothProfileManager).removeServiceListener(any());
+
+        mActivity.onServiceDisconnected();
+        // Do nothing.
+    }
+
+    @Test
+    public void setupActivity_serviceReady_doNothing() {
+        mActivity = Robolectric.setupActivity(AudioStreamConfirmDialogActivity.class);
+
+        verify(mLocalBluetoothProfileManager, never()).addServiceListener(any());
+    }
+
+    @Test
+    public void onStop_unregisterCallback() {
+        mActivity = Robolectric.setupActivity(AudioStreamConfirmDialogActivity.class);
+        mActivity.onStop();
+
+        verify(mLocalBluetoothProfileManager).removeServiceListener(any());
+    }
 }
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamConfirmDialogTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamConfirmDialogTest.java
index c7aafe8..601c432 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamConfirmDialogTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamConfirmDialogTest.java
@@ -16,14 +16,21 @@
 
 package com.android.settings.connecteddevice.audiosharing.audiostreams;
 
+import static android.app.settings.SettingsEnums.DIALOG_AUDIO_STREAM_CONFIRM_LISTEN;
+
+import static com.android.settings.connecteddevice.audiosharing.audiostreams.AudioStreamConfirmDialog.DEFAULT_DEVICE_NAME;
 import static com.android.settings.connecteddevice.audiosharing.audiostreams.AudioStreamsDashboardFragment.KEY_BROADCAST_METADATA;
 
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 import static org.robolectric.shadows.ShadowLooper.shadowMainLooper;
 
+import android.app.Dialog;
 import android.app.settings.SettingsEnums;
 import android.bluetooth.BluetoothAdapter;
 import android.bluetooth.BluetoothDevice;
@@ -127,6 +134,8 @@
 
         assertThat(mDialogFragment.getMetricsCategory())
                 .isEqualTo(SettingsEnums.DIALOG_AUDIO_STREAM_CONFIRM_FEATURE_UNSUPPORTED);
+        assertThat(mDialogFragment.mActivity).isNotNull();
+        mDialogFragment.mActivity = spy(mDialogFragment.mActivity);
 
         var dialog = mDialogFragment.getDialog();
         assertThat(dialog).isNotNull();
@@ -152,6 +161,10 @@
         assertThat(rightButton).isNotNull();
         assertThat(rightButton.getVisibility()).isEqualTo(View.VISIBLE);
         assertThat(rightButton.hasOnClickListeners()).isTrue();
+
+        rightButton.callOnClick();
+        assertThat(dialog.isShowing()).isFalse();
+        verify(mDialogFragment.mActivity).finish();
     }
 
     @Test
@@ -165,6 +178,8 @@
 
         assertThat(mDialogFragment.getMetricsCategory())
                 .isEqualTo(SettingsEnums.DIALOG_AUDIO_STREAM_CONFIRM_NO_LE_DEVICE);
+        assertThat(mDialogFragment.mActivity).isNotNull();
+        mDialogFragment.mActivity = spy(mDialogFragment.mActivity);
 
         var dialog = mDialogFragment.getDialog();
         assertThat(dialog).isNotNull();
@@ -184,11 +199,20 @@
         View leftButton = dialog.findViewById(R.id.left_button);
         assertThat(leftButton).isNotNull();
         assertThat(leftButton.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(leftButton.hasOnClickListeners()).isTrue();
+
+        leftButton.callOnClick();
+        assertThat(dialog.isShowing()).isFalse();
+
         Button rightButton = dialog.findViewById(R.id.right_button);
         assertThat(rightButton).isNotNull();
         assertThat(rightButton.getText())
                 .isEqualTo(mContext.getString(R.string.audio_streams_dialog_no_le_device_button));
         assertThat(rightButton.hasOnClickListeners()).isTrue();
+
+        rightButton.callOnClick();
+        assertThat(dialog.isShowing()).isFalse();
+        verify(mDialogFragment.mActivity, times(2)).finish();
     }
 
     @Test
@@ -207,6 +231,8 @@
 
         assertThat(mDialogFragment.getMetricsCategory())
                 .isEqualTo(SettingsEnums.DIALOG_AUDIO_STREAM_CONFIRM_DATA_ERROR);
+        assertThat(mDialogFragment.mActivity).isNotNull();
+        mDialogFragment.mActivity = spy(mDialogFragment.mActivity);
 
         var dialog = mDialogFragment.getDialog();
         assertThat(dialog).isNotNull();
@@ -231,6 +257,10 @@
         assertThat(rightButton).isNotNull();
         assertThat(rightButton.getVisibility()).isEqualTo(View.VISIBLE);
         assertThat(rightButton.hasOnClickListeners()).isTrue();
+
+        rightButton.callOnClick();
+        assertThat(dialog.isShowing()).isFalse();
+        verify(mDialogFragment.mActivity).finish();
     }
 
     @Test
@@ -252,6 +282,8 @@
 
         assertThat(mDialogFragment.getMetricsCategory())
                 .isEqualTo(SettingsEnums.DIALOG_AUDIO_STREAM_CONFIRM_DATA_ERROR);
+        assertThat(mDialogFragment.mActivity).isNotNull();
+        mDialogFragment.mActivity = spy(mDialogFragment.mActivity);
 
         var dialog = mDialogFragment.getDialog();
         assertThat(dialog).isNotNull();
@@ -276,6 +308,10 @@
         assertThat(rightButton).isNotNull();
         assertThat(rightButton.getVisibility()).isEqualTo(View.VISIBLE);
         assertThat(rightButton.hasOnClickListeners()).isTrue();
+
+        rightButton.callOnClick();
+        assertThat(dialog.isShowing()).isFalse();
+        verify(mDialogFragment.mActivity).finish();
     }
 
     @Test
@@ -283,7 +319,7 @@
         List<BluetoothDevice> devices = new ArrayList<>();
         devices.add(mBluetoothDevice);
         when(mAssistant.getDevicesMatchingConnectionStates(any())).thenReturn(devices);
-        when(mBluetoothDevice.getAlias()).thenReturn(DEVICE_NAME);
+        when(mBluetoothDevice.getAlias()).thenReturn("");
 
         Intent intent = new Intent();
         intent.putExtra(KEY_BROADCAST_METADATA, VALID_METADATA);
@@ -296,9 +332,11 @@
         shadowMainLooper().idle();
 
         assertThat(mDialogFragment.getMetricsCategory())
-                .isEqualTo(SettingsEnums.DIALOG_AUDIO_STREAM_CONFIRM_LISTEN);
+                .isEqualTo(DIALOG_AUDIO_STREAM_CONFIRM_LISTEN);
+        assertThat(mDialogFragment.mActivity).isNotNull();
+        mDialogFragment.mActivity = spy(mDialogFragment.mActivity);
 
-        var dialog = mDialogFragment.getDialog();
+        Dialog dialog = mDialogFragment.getDialog();
         assertThat(dialog).isNotNull();
         assertThat(dialog.isShowing()).isTrue();
         TextView title = dialog.findViewById(R.id.dialog_title);
@@ -311,17 +349,27 @@
         assertThat(subtitle1.getVisibility()).isEqualTo(View.VISIBLE);
         TextView subtitle2 = dialog.findViewById(R.id.dialog_subtitle_2);
         assertThat(subtitle2).isNotNull();
+        var defaultName = mContext.getString(DEFAULT_DEVICE_NAME);
         assertThat(subtitle2.getText())
                 .isEqualTo(
                         mContext.getString(
-                                R.string.audio_streams_dialog_control_volume, DEVICE_NAME));
+                                R.string.audio_streams_dialog_control_volume, defaultName));
         View leftButton = dialog.findViewById(R.id.left_button);
         assertThat(leftButton).isNotNull();
         assertThat(leftButton.getVisibility()).isEqualTo(View.VISIBLE);
+        assertThat(leftButton.hasOnClickListeners()).isTrue();
+
+        leftButton.callOnClick();
+        assertThat(dialog.isShowing()).isFalse();
+
         Button rightButton = dialog.findViewById(R.id.right_button);
         assertThat(rightButton).isNotNull();
         assertThat(rightButton.getText())
                 .isEqualTo(mContext.getString(R.string.audio_streams_dialog_listen));
         assertThat(rightButton.hasOnClickListeners()).isTrue();
+
+        rightButton.callOnClick();
+        assertThat(dialog.isShowing()).isFalse();
+        verify(mDialogFragment.mActivity, times(2)).finish();
     }
 }
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamDetailsFragmentTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamDetailsFragmentTest.java
index 724c772..46d481a 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamDetailsFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamDetailsFragmentTest.java
@@ -16,22 +16,48 @@
 
 package com.android.settings.connecteddevice.audiosharing.audiostreams;
 
+import static com.android.settings.connecteddevice.audiosharing.audiostreams.AudioStreamDetailsFragment.BROADCAST_ID_ARG;
+import static com.android.settings.connecteddevice.audiosharing.audiostreams.AudioStreamDetailsFragment.BROADCAST_NAME_ARG;
+
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+
+import android.app.settings.SettingsEnums;
+import android.content.Context;
+import android.os.Bundle;
+
+import androidx.test.core.app.ApplicationProvider;
+
 import com.android.settings.R;
+import com.android.settingslib.core.AbstractPreferenceController;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
 import org.robolectric.RobolectricTestRunner;
 
 @RunWith(RobolectricTestRunner.class)
 public class AudioStreamDetailsFragmentTest {
-    private AudioStreamDetailsFragment mFragment;
+    @Rule public final MockitoRule mocks = MockitoJUnit.rule();
+    private static final String BROADCAST_NAME = "name";
+    private static final int BROADCAST_ID = 1;
+    private final Context mContext = ApplicationProvider.getApplicationContext();
+    @Mock private AudioStreamHeaderController mHeaderController;
+    @Mock private AudioStreamButtonController mButtonController;
+    private TestFragment mFragment;
 
     @Before
     public void setUp() {
-        mFragment = new AudioStreamDetailsFragment();
+        mFragment = spy(new TestFragment());
+        doReturn(mHeaderController).when(mFragment).use(AudioStreamHeaderController.class);
+        doReturn(mButtonController).when(mFragment).use(AudioStreamButtonController.class);
     }
 
     @Test
@@ -44,4 +70,29 @@
     public void getLogTag_returnsCorrectTag() {
         assertThat(mFragment.getLogTag()).isEqualTo(AudioStreamDetailsFragment.TAG);
     }
+
+    @Test
+    public void getMetricsCategory_returnsCorrectEnum() {
+        assertThat(mFragment.getMetricsCategory()).isEqualTo(SettingsEnums.AUDIO_STREAM_DETAIL);
+    }
+
+    @Test
+    public void onAttach_getArguments() {
+        Bundle bundle = new Bundle();
+        bundle.putString(BROADCAST_NAME_ARG, BROADCAST_NAME);
+        bundle.putInt(BROADCAST_ID_ARG, BROADCAST_ID);
+        mFragment.setArguments(bundle);
+
+        mFragment.onAttach(mContext);
+
+        verify(mButtonController).init(BROADCAST_ID);
+        verify(mHeaderController).init(mFragment, BROADCAST_NAME, BROADCAST_ID);
+    }
+
+    public static class TestFragment extends AudioStreamDetailsFragment {
+        @Override
+        protected <T extends AbstractPreferenceController> T use(Class<T> clazz) {
+            return super.use(clazz);
+        }
+    }
 }
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamHeaderControllerTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamHeaderControllerTest.java
index 0cd5d61..327090d 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamHeaderControllerTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamHeaderControllerTest.java
@@ -19,12 +19,20 @@
 import static com.android.settings.connecteddevice.audiosharing.audiostreams.AudioStreamHeaderController.AUDIO_STREAM_HEADER_LISTENING_NOW_SUMMARY;
 import static com.android.settings.connecteddevice.audiosharing.audiostreams.AudioStreamHeaderController.AUDIO_STREAM_HEADER_NOT_LISTENING_SUMMARY;
 
+import static org.mockito.ArgumentMatchers.any;
+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.bluetooth.BluetoothDevice;
+import android.bluetooth.BluetoothLeBroadcastAssistant;
 import android.bluetooth.BluetoothLeBroadcastReceiveState;
 import android.content.Context;
+import android.graphics.drawable.Drawable;
 
+import androidx.lifecycle.LifecycleOwner;
 import androidx.preference.PreferenceScreen;
 import androidx.test.core.app.ApplicationProvider;
 
@@ -32,6 +40,8 @@
 import com.android.settings.connecteddevice.audiosharing.audiostreams.testshadows.ShadowEntityHeaderController;
 import com.android.settings.testutils.shadow.ShadowThreadUtils;
 import com.android.settings.widget.EntityHeaderController;
+import com.android.settingslib.bluetooth.LocalBluetoothLeBroadcastAssistant;
+import com.android.settingslib.core.lifecycle.Lifecycle;
 import com.android.settingslib.widget.LayoutPreference;
 
 import org.junit.After;
@@ -45,8 +55,10 @@
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.annotation.Config;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.Executor;
 
 @RunWith(RobolectricTestRunner.class)
 @Config(
@@ -65,15 +77,21 @@
     @Mock private AudioStreamsHelper mAudioStreamsHelper;
     @Mock private PreferenceScreen mScreen;
     @Mock private BluetoothLeBroadcastReceiveState mBroadcastReceiveState;
+    @Mock private LocalBluetoothLeBroadcastAssistant mAssistant;
     @Mock private AudioStreamDetailsFragment mFragment;
     @Mock private LayoutPreference mPreference;
     @Mock private EntityHeaderController mHeaderController;
+    private Lifecycle mLifecycle;
+    private LifecycleOwner mLifecycleOwner;
     private AudioStreamHeaderController mController;
 
     @Before
     public void setUp() {
         ShadowEntityHeaderController.setUseMock(mHeaderController);
         ShadowAudioStreamsHelper.setUseMock(mAudioStreamsHelper);
+        when(mAudioStreamsHelper.getLeBroadcastAssistant()).thenReturn(mAssistant);
+        mLifecycleOwner = () -> mLifecycle;
+        mLifecycle = new Lifecycle(mLifecycleOwner);
         mController = new AudioStreamHeaderController(mContext, KEY);
         mController.init(mFragment, BROADCAST_NAME, BROADCAST_ID);
         when(mScreen.findPreference(KEY)).thenReturn(mPreference);
@@ -88,6 +106,40 @@
     }
 
     @Test
+    public void onStart_registerCallbacks() {
+        mController.onStart(mLifecycleOwner);
+        verify(mAssistant)
+                .registerServiceCallBack(
+                        any(Executor.class), any(BluetoothLeBroadcastAssistant.Callback.class));
+    }
+
+    @Test
+    public void onStart_profileNull_doNothing() {
+        when(mAudioStreamsHelper.getLeBroadcastAssistant()).thenReturn(null);
+        mController = new AudioStreamHeaderController(mContext, KEY);
+        mController.onStart(mLifecycleOwner);
+        verify(mAssistant, never())
+                .registerServiceCallBack(
+                        any(Executor.class), any(BluetoothLeBroadcastAssistant.Callback.class));
+    }
+
+    @Test
+    public void onStop_unregisterCallbacks() {
+        mController.onStop(mLifecycleOwner);
+        verify(mAssistant)
+                .unregisterServiceCallBack(any(BluetoothLeBroadcastAssistant.Callback.class));
+    }
+
+    @Test
+    public void onStop_profileNull_doNothing() {
+        when(mAudioStreamsHelper.getLeBroadcastAssistant()).thenReturn(null);
+        mController = new AudioStreamHeaderController(mContext, KEY);
+        mController.onStop(mLifecycleOwner);
+        verify(mAssistant, never())
+                .unregisterServiceCallBack(any(BluetoothLeBroadcastAssistant.Callback.class));
+    }
+
+    @Test
     public void testDisplayPreference_sourceConnected_setSummary() {
         when(mAudioStreamsHelper.getAllConnectedSources())
                 .thenReturn(List.of(mBroadcastReceiveState));
@@ -96,9 +148,11 @@
         mController.displayPreference(mScreen);
 
         verify(mHeaderController).setLabel(BROADCAST_NAME);
+        verify(mHeaderController).setIcon(any(Drawable.class));
         verify(mHeaderController)
                 .setSummary(mContext.getString(AUDIO_STREAM_HEADER_LISTENING_NOW_SUMMARY));
         verify(mHeaderController).done(true);
+        verify(mScreen).addPreference(any());
     }
 
     @Test
@@ -108,7 +162,54 @@
         mController.displayPreference(mScreen);
 
         verify(mHeaderController).setLabel(BROADCAST_NAME);
+        verify(mHeaderController).setIcon(any(Drawable.class));
         verify(mHeaderController).setSummary(AUDIO_STREAM_HEADER_NOT_LISTENING_SUMMARY);
         verify(mHeaderController).done(true);
+        verify(mScreen).addPreference(any());
+    }
+
+    @Test
+    public void testCallback_onSourceRemoved_updateButton() {
+        when(mAudioStreamsHelper.getAllConnectedSources()).thenReturn(Collections.emptyList());
+
+        mController.displayPreference(mScreen);
+        mController.mBroadcastAssistantCallback.onSourceRemoved(
+                mock(BluetoothDevice.class), /* sourceId= */ 0, /* reason= */ 0);
+
+        // Called twice, once in displayPreference, the other one in callback
+        verify(mHeaderController, times(2)).setSummary(AUDIO_STREAM_HEADER_NOT_LISTENING_SUMMARY);
+        verify(mHeaderController, times(2)).done(true);
+    }
+
+    @Test
+    public void testCallback_onSourceLost_updateButton() {
+        when(mAudioStreamsHelper.getAllConnectedSources()).thenReturn(Collections.emptyList());
+
+        mController.displayPreference(mScreen);
+        mController.mBroadcastAssistantCallback.onSourceLost(/* broadcastId= */ 1);
+
+        // Called twice, once in displayPreference, the other one in callback
+        verify(mHeaderController, times(2)).setSummary(AUDIO_STREAM_HEADER_NOT_LISTENING_SUMMARY);
+        verify(mHeaderController, times(2)).done(true);
+    }
+
+    @Test
+    public void testCallback_onReceiveStateChanged_updateButton() {
+        when(mAudioStreamsHelper.getAllConnectedSources())
+                .thenReturn(List.of(mBroadcastReceiveState));
+        when(mBroadcastReceiveState.getBroadcastId()).thenReturn(BROADCAST_ID);
+        BluetoothLeBroadcastReceiveState state = mock(BluetoothLeBroadcastReceiveState.class);
+        List<Long> bisSyncState = new ArrayList<>();
+        bisSyncState.add(1L);
+        when(state.getBisSyncState()).thenReturn(bisSyncState);
+
+        mController.displayPreference(mScreen);
+        mController.mBroadcastAssistantCallback.onReceiveStateChanged(
+                mock(BluetoothDevice.class), /* sourceId= */ 0, state);
+
+        // Called twice, once in displayPreference, the other one in callback
+        verify(mHeaderController, times(2))
+                .setSummary(mContext.getString(AUDIO_STREAM_HEADER_LISTENING_NOW_SUMMARY));
+        verify(mHeaderController, times(2)).done(true);
     }
 }
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamMediaServiceTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamMediaServiceTest.java
index b184d88..abdd743 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamMediaServiceTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamMediaServiceTest.java
@@ -18,22 +18,29 @@
 
 import static com.android.settings.connecteddevice.audiosharing.audiostreams.AudioStreamMediaService.BROADCAST_ID;
 import static com.android.settings.connecteddevice.audiosharing.audiostreams.AudioStreamMediaService.DEVICES;
+import static com.android.settings.connecteddevice.audiosharing.audiostreams.AudioStreamMediaService.LEAVE_BROADCAST_ACTION;
 
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.app.Notification;
+import android.app.NotificationChannel;
 import android.app.NotificationManager;
+import android.app.settings.SettingsEnums;
 import android.bluetooth.BluetoothAdapter;
 import android.bluetooth.BluetoothDevice;
+import android.bluetooth.BluetoothProfile;
 import android.bluetooth.BluetoothStatusCodes;
 import android.content.Context;
 import android.content.Intent;
@@ -43,14 +50,20 @@
 import android.media.session.ISession;
 import android.media.session.ISessionController;
 import android.media.session.MediaSessionManager;
+import android.os.Bundle;
+import android.os.IBinder;
 import android.os.RemoteException;
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.util.DisplayMetrics;
 
 import com.android.settings.connecteddevice.audiosharing.audiostreams.testshadows.ShadowAudioStreamsHelper;
+import com.android.settings.testutils.FakeFeatureFactory;
 import com.android.settings.testutils.shadow.ShadowBluetoothAdapter;
 import com.android.settings.testutils.shadow.ShadowBluetoothUtils;
+import com.android.settings.testutils.shadow.ShadowThreadUtils;
 import com.android.settingslib.bluetooth.BluetoothEventManager;
+import com.android.settingslib.bluetooth.CachedBluetoothDevice;
+import com.android.settingslib.bluetooth.CachedBluetoothDeviceManager;
 import com.android.settingslib.bluetooth.LocalBluetoothLeBroadcastAssistant;
 import com.android.settingslib.bluetooth.LocalBluetoothManager;
 import com.android.settingslib.bluetooth.LocalBluetoothProfileManager;
@@ -72,10 +85,12 @@
 import org.robolectric.util.ReflectionHelpers;
 
 import java.util.ArrayList;
+import java.util.Set;
 
 @RunWith(RobolectricTestRunner.class)
 @Config(
         shadows = {
+            ShadowThreadUtils.class,
             ShadowBluetoothAdapter.class,
             ShadowBluetoothUtils.class,
             ShadowAudioStreamsHelper.class,
@@ -83,6 +98,8 @@
 public class AudioStreamMediaServiceTest {
     @Rule public final MockitoRule mMockitoRule = MockitoJUnit.rule();
     @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+    private static final String CHANNEL_ID = "bluetooth_notification_channel";
+    private static final String DEVICE_NAME = "name";
     @Mock private Resources mResources;
     @Mock private LocalBluetoothManager mLocalBtManager;
     @Mock private LocalBluetoothLeBroadcastAssistant mLeBroadcastAssistant;
@@ -91,17 +108,21 @@
     @Mock private MediaSessionManager mMediaSessionManager;
     @Mock private BluetoothEventManager mBluetoothEventManager;
     @Mock private LocalBluetoothProfileManager mLocalBluetoothProfileManager;
+    @Mock private CachedBluetoothDeviceManager mCachedDeviceManager;
     @Mock private VolumeControlProfile mVolumeControlProfile;
+    @Mock private CachedBluetoothDevice mCachedBluetoothDevice;
     @Mock private BluetoothDevice mDevice;
     @Mock private ISession mISession;
     @Mock private ISessionController mISessionController;
     @Mock private PackageManager mPackageManager;
     @Mock private DisplayMetrics mDisplayMetrics;
     @Mock private Context mContext;
+    private FakeFeatureFactory mFeatureFactory;
     private AudioStreamMediaService mAudioStreamMediaService;
 
     @Before
     public void setUp() {
+        mFeatureFactory = FakeFeatureFactory.setupForTest();
         ShadowAudioStreamsHelper.setUseMock(mAudioStreamsHelper);
         when(mAudioStreamsHelper.getLeBroadcastAssistant()).thenReturn(mLeBroadcastAssistant);
         ShadowBluetoothAdapter shadowBluetoothAdapter =
@@ -114,6 +135,9 @@
         ShadowBluetoothUtils.sLocalBluetoothManager = mLocalBtManager;
         when(mLocalBtManager.getEventManager()).thenReturn(mBluetoothEventManager);
         when(mLocalBtManager.getProfileManager()).thenReturn(mLocalBluetoothProfileManager);
+        when(mLocalBtManager.getCachedDeviceManager()).thenReturn(mCachedDeviceManager);
+        when(mCachedDeviceManager.findDevice(any())).thenReturn(mCachedBluetoothDevice);
+        when(mCachedBluetoothDevice.getName()).thenReturn(DEVICE_NAME);
         when(mLocalBluetoothProfileManager.getVolumeControlProfile())
                 .thenReturn(mVolumeControlProfile);
 
@@ -169,6 +193,25 @@
     }
 
     @Test
+    public void onCreate_flagOn_createNewChannel() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+        when(mNotificationManager.getNotificationChannel(anyString())).thenReturn(null);
+
+        mAudioStreamMediaService.onCreate();
+
+        ArgumentCaptor<NotificationChannel> notificationChannelCapture =
+                ArgumentCaptor.forClass(NotificationChannel.class);
+        verify(mNotificationManager)
+                .createNotificationChannel(notificationChannelCapture.capture());
+        NotificationChannel newChannel = notificationChannelCapture.getValue();
+        assertThat(newChannel).isNotNull();
+        assertThat(newChannel.getId()).isEqualTo(CHANNEL_ID);
+        assertThat(newChannel.getName())
+                .isEqualTo(mContext.getString(com.android.settings.R.string.bluetooth));
+        assertThat(newChannel.getImportance()).isEqualTo(NotificationManager.IMPORTANCE_HIGH);
+    }
+
+    @Test
     public void onDestroy_flagOff_doNothing() {
         mSetFlagsRule.disableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
 
@@ -183,8 +226,15 @@
     @Test
     public void onDestroy_flagOn_cleanup() {
         mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+        var devices = new ArrayList<BluetoothDevice>();
+        devices.add(mDevice);
+
+        Intent intent = new Intent();
+        intent.putExtra(BROADCAST_ID, 1);
+        intent.putParcelableArrayListExtra(DEVICES, devices);
 
         mAudioStreamMediaService.onCreate();
+        mAudioStreamMediaService.onStartCommand(intent, /* flags= */ 0, /* startId= */ 0);
         mAudioStreamMediaService.onDestroy();
 
         verify(mBluetoothEventManager).unregisterCallback(any());
@@ -196,7 +246,6 @@
     public void onStartCommand_noBroadcastId_stopSelf() {
         mAudioStreamMediaService.onStartCommand(new Intent(), /* flags= */ 0, /* startId= */ 0);
 
-        assertThat(mAudioStreamMediaService.mLocalSession).isNull();
         verify(mAudioStreamMediaService).stopSelf();
     }
 
@@ -207,7 +256,6 @@
 
         mAudioStreamMediaService.onStartCommand(intent, /* flags= */ 0, /* startId= */ 0);
 
-        assertThat(mAudioStreamMediaService.mLocalSession).isNull();
         verify(mAudioStreamMediaService).stopSelf();
     }
 
@@ -222,12 +270,179 @@
 
         mAudioStreamMediaService.onStartCommand(intent, /* flags= */ 0, /* startId= */ 0);
 
-        assertThat(mAudioStreamMediaService.mLocalSession).isNotNull();
-        verify(mAudioStreamMediaService, never()).stopSelf();
+        ArgumentCaptor<Notification> notificationCapture =
+                ArgumentCaptor.forClass(Notification.class);
+        verify(mAudioStreamMediaService).startForeground(anyInt(), notificationCapture.capture());
+        var notification = notificationCapture.getValue();
+        assertThat(notification.getSmallIcon()).isNotNull();
+        assertThat(notification.isStyle(Notification.MediaStyle.class)).isTrue();
 
-        ArgumentCaptor<Notification> notification = ArgumentCaptor.forClass(Notification.class);
-        verify(mAudioStreamMediaService).startForeground(anyInt(), notification.capture());
-        assertThat(notification.getValue().getSmallIcon()).isNotNull();
-        assertThat(notification.getValue().isStyle(Notification.MediaStyle.class)).isTrue();
+        verify(mAudioStreamMediaService, never()).stopSelf();
+    }
+
+    @Test
+    public void assistantCallback_onSourceLost_stopSelf() {
+        mAudioStreamMediaService.onCreate();
+
+        assertThat(mAudioStreamMediaService.mBroadcastAssistantCallback).isNotNull();
+        mAudioStreamMediaService.mBroadcastAssistantCallback.onSourceLost(/* broadcastId= */ 0);
+
+        verify(mAudioStreamMediaService).stopSelf();
+    }
+
+    @Test
+    public void assistantCallback_onSourceRemoved_stopSelf() {
+        mAudioStreamMediaService.onCreate();
+
+        assertThat(mAudioStreamMediaService.mBroadcastAssistantCallback).isNotNull();
+        mAudioStreamMediaService.mBroadcastAssistantCallback.onSourceRemoved(
+                mDevice, /* sourceId= */ 0, /* reason= */ 0);
+
+        verify(mAudioStreamMediaService).stopSelf();
+    }
+
+    @Test
+    public void bluetoothCallback_onBluetoothOff_stopSelf() {
+        mAudioStreamMediaService.onCreate();
+
+        assertThat(mAudioStreamMediaService.mBluetoothCallback).isNotNull();
+        mAudioStreamMediaService.mBluetoothCallback.onBluetoothStateChanged(
+                BluetoothAdapter.STATE_OFF);
+
+        verify(mAudioStreamMediaService).stopSelf();
+    }
+
+    @Test
+    public void bluetoothCallback_onDeviceDisconnect_stopSelf() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+        mAudioStreamMediaService.onCreate();
+        assertThat(mAudioStreamMediaService.mBluetoothCallback).isNotNull();
+        mAudioStreamMediaService.onStartCommand(setupIntent(), /* flags= */ 0, /* startId= */ 0);
+
+        mAudioStreamMediaService.mBluetoothCallback.onProfileConnectionStateChanged(
+                mCachedBluetoothDevice,
+                BluetoothAdapter.STATE_DISCONNECTED,
+                BluetoothProfile.LE_AUDIO_BROADCAST_ASSISTANT);
+
+        verify(mAudioStreamMediaService).stopSelf();
+    }
+
+    @Test
+    public void bluetoothCallback_onMemberDeviceDisconnect_stopSelf() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+        when(mCachedBluetoothDevice.getDevice()).thenReturn(mock(BluetoothDevice.class));
+        CachedBluetoothDevice member = mock(CachedBluetoothDevice.class);
+        when(mCachedBluetoothDevice.getMemberDevice()).thenReturn(Set.of(member));
+        when(member.getDevice()).thenReturn(mDevice);
+        var devices = new ArrayList<BluetoothDevice>();
+        devices.add(mDevice);
+
+        Intent intent = new Intent();
+        intent.putExtra(BROADCAST_ID, 1);
+        intent.putParcelableArrayListExtra(DEVICES, devices);
+
+        mAudioStreamMediaService.onCreate();
+        assertThat(mAudioStreamMediaService.mBluetoothCallback).isNotNull();
+        mAudioStreamMediaService.onStartCommand(intent, /* flags= */ 0, /* startId= */ 0);
+        mAudioStreamMediaService.mBluetoothCallback.onProfileConnectionStateChanged(
+                mCachedBluetoothDevice,
+                BluetoothAdapter.STATE_DISCONNECTED,
+                BluetoothProfile.LE_AUDIO_BROADCAST_ASSISTANT);
+
+        verify(mAudioStreamMediaService).stopSelf();
+    }
+
+    @Test
+    public void mediaSessionCallback_onSeekTo_updateNotification() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+
+        mAudioStreamMediaService.onCreate();
+        mAudioStreamMediaService.onStartCommand(setupIntent(), /* flags= */ 0, /* startId= */ 0);
+        assertThat(mAudioStreamMediaService.mMediaSessionCallback).isNotNull();
+        mAudioStreamMediaService.mMediaSessionCallback.onSeekTo(100);
+
+        verify(mNotificationManager).notify(anyInt(), any());
+    }
+
+    @Test
+    public void mediaSessionCallback_onPause_setVolume() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+
+        mAudioStreamMediaService.onCreate();
+        mAudioStreamMediaService.onStartCommand(setupIntent(), /* flags= */ 0, /* startId= */ 0);
+        assertThat(mAudioStreamMediaService.mMediaSessionCallback).isNotNull();
+        mAudioStreamMediaService.mMediaSessionCallback.onPause();
+
+        verify(mVolumeControlProfile).setDeviceVolume(any(), anyInt(), anyBoolean());
+        verify(mFeatureFactory.metricsFeatureProvider)
+                .action(
+                        any(),
+                        eq(SettingsEnums.ACTION_AUDIO_STREAM_NOTIFICATION_MUTE_BUTTON_CLICK),
+                        eq(1));
+    }
+
+    @Test
+    public void mediaSessionCallback_onPlay_setVolume() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+
+        mAudioStreamMediaService.onCreate();
+        mAudioStreamMediaService.onStartCommand(setupIntent(), /* flags= */ 0, /* startId= */ 0);
+        assertThat(mAudioStreamMediaService.mMediaSessionCallback).isNotNull();
+        mAudioStreamMediaService.mMediaSessionCallback.onPlay();
+
+        verify(mVolumeControlProfile).setDeviceVolume(any(), anyInt(), anyBoolean());
+        verify(mFeatureFactory.metricsFeatureProvider)
+                .action(
+                        any(),
+                        eq(SettingsEnums.ACTION_AUDIO_STREAM_NOTIFICATION_MUTE_BUTTON_CLICK),
+                        eq(0));
+    }
+
+    @Test
+    public void mediaSessionCallback_onCustomAction_leaveBroadcast() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+
+        mAudioStreamMediaService.onCreate();
+        mAudioStreamMediaService.onStartCommand(setupIntent(), /* flags= */ 0, /* startId= */ 0);
+        assertThat(mAudioStreamMediaService.mMediaSessionCallback).isNotNull();
+        mAudioStreamMediaService.mMediaSessionCallback.onCustomAction(
+                LEAVE_BROADCAST_ACTION, Bundle.EMPTY);
+
+        verify(mAudioStreamsHelper).removeSource(anyInt());
+        verify(mFeatureFactory.metricsFeatureProvider)
+                .action(
+                        any(),
+                        eq(SettingsEnums.ACTION_AUDIO_STREAM_NOTIFICATION_LEAVE_BUTTON_CLICK));
+    }
+
+    @Test
+    public void volumeControlCallback_onDeviceVolumeChanged_updateNotification() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_LE_AUDIO_SHARING);
+
+        mAudioStreamMediaService.onCreate();
+        assertThat(mAudioStreamMediaService.mVolumeControlCallback).isNotNull();
+        mAudioStreamMediaService.onStartCommand(setupIntent(), /* flags= */ 0, /* startId= */ 0);
+        mAudioStreamMediaService.mVolumeControlCallback.onDeviceVolumeChanged(
+                mDevice, /* volume= */ 0);
+
+        verify(mNotificationManager).notify(anyInt(), any());
+    }
+
+    @Test
+    public void onBind_returnNull() {
+        IBinder binder = mAudioStreamMediaService.onBind(new Intent());
+
+        assertThat(binder).isNull();
+    }
+
+    private Intent setupIntent() {
+        when(mCachedBluetoothDevice.getDevice()).thenReturn(mDevice);
+        var devices = new ArrayList<BluetoothDevice>();
+        devices.add(mDevice);
+
+        Intent intent = new Intent();
+        intent.putExtra(BROADCAST_ID, 1);
+        intent.putParcelableArrayListExtra(DEVICES, devices);
+        return intent;
     }
 }
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsDashboardFragmentTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsDashboardFragmentTest.java
index 9058ab6..dd3d8b7 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsDashboardFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsDashboardFragmentTest.java
@@ -22,18 +22,25 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 
 import android.app.Activity;
+import android.app.settings.SettingsEnums;
+import android.bluetooth.BluetoothLeBroadcastMetadata;
 import android.content.Context;
 import android.content.Intent;
+import android.os.Bundle;
 
 import androidx.test.core.app.ApplicationProvider;
 
 import com.android.settings.R;
+import com.android.settings.testutils.FakeFeatureFactory;
 import com.android.settingslib.core.AbstractPreferenceController;
 
 import org.junit.Before;
@@ -53,12 +60,14 @@
                     + "MD:BgNwVGVzdA==;AS:1;PI:A0;NS:1;BS:3;NB:2;SM:BQNUZXN0BARlbmc=;;";
 
     private Context mContext;
+    private FakeFeatureFactory mFeatureFactory;
     private AudioStreamsProgressCategoryController mController;
     private TestFragment mTestFragment;
 
     @Before
     public void setUp() {
         mContext = ApplicationProvider.getApplicationContext();
+        mFeatureFactory = FakeFeatureFactory.setupForTest();
         mTestFragment = spy(new TestFragment());
         doReturn(mContext).when(mTestFragment).getContext();
         mController = spy(new AudioStreamsProgressCategoryController(mContext, "key"));
@@ -114,6 +123,28 @@
         mTestFragment.onActivityResult(
                 REQUEST_SCAN_BT_BROADCAST_QR_CODE, Activity.RESULT_OK, intent);
         verify(mController).setSourceFromQrCode(any(), any());
+        verify(mFeatureFactory.metricsFeatureProvider)
+                .action(
+                        any(),
+                        eq(SettingsEnums.ACTION_AUDIO_STREAM_QR_CODE_SCAN_SUCCEED),
+                        anyInt());
+    }
+
+    @Test
+    public void onAttach_hasArgument() {
+        BluetoothLeBroadcastMetadata data = mock(BluetoothLeBroadcastMetadata.class);
+        Bundle bundle = new Bundle();
+        bundle.putParcelable(KEY_BROADCAST_METADATA, data);
+        mTestFragment.setArguments(bundle);
+
+        mTestFragment.onAttach(mContext);
+
+        verify(mController).setSourceFromQrCode(eq(data), any());
+        verify(mFeatureFactory.metricsFeatureProvider)
+                .action(
+                        any(),
+                        eq(SettingsEnums.ACTION_AUDIO_STREAM_QR_CODE_SCAN_SUCCEED),
+                        anyInt());
     }
 
     public static class TestFragment extends AudioStreamsDashboardFragment {
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsDialogFragmentTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsDialogFragmentTest.java
index e83dade..efdd389 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsDialogFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsDialogFragmentTest.java
@@ -73,7 +73,7 @@
     }
 
     @Test
-    public void testShowDialog() {
+    public void testShowDialog_dismissAll() {
         FragmentController.setupFragment(mFragment);
         AudioStreamsDialogFragment.show(mFragment, mDialogBuilder, SettingsEnums.PAGE_UNKNOWN);
         ShadowLooper.idleMainLooper();
@@ -81,5 +81,8 @@
         var dialog = ShadowAlertDialog.getLatestAlertDialog();
         assertThat(dialog).isNotNull();
         assertThat(dialog.isShowing()).isTrue();
+
+        AudioStreamsDialogFragment.dismissAll(mFragment);
+        assertThat(dialog.isShowing()).isFalse();
     }
 }
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsHelperTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsHelperTest.java
new file mode 100644
index 0000000..66ef5fb
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsHelperTest.java
@@ -0,0 +1,239 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.connecteddevice.audiosharing.audiostreams;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.bluetooth.BluetoothDevice;
+import android.bluetooth.BluetoothLeBroadcastMetadata;
+import android.bluetooth.BluetoothLeBroadcastReceiveState;
+import android.content.Context;
+
+import androidx.test.core.app.ApplicationProvider;
+
+import com.android.settings.testutils.shadow.ShadowThreadUtils;
+import com.android.settingslib.bluetooth.CachedBluetoothDevice;
+import com.android.settingslib.bluetooth.CachedBluetoothDeviceManager;
+import com.android.settingslib.bluetooth.LocalBluetoothLeBroadcastAssistant;
+import com.android.settingslib.bluetooth.LocalBluetoothManager;
+import com.android.settingslib.bluetooth.LocalBluetoothProfileManager;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.annotation.Config;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+@RunWith(RobolectricTestRunner.class)
+@Config(
+        shadows = {
+            ShadowThreadUtils.class,
+        })
+public class AudioStreamsHelperTest {
+    @Rule public final MockitoRule mMockitoRule = MockitoJUnit.rule();
+    private static final int GROUP_ID = 1;
+    private static final int BROADCAST_ID_1 = 1;
+    private static final int BROADCAST_ID_2 = 2;
+    private static final String BROADCAST_NAME = "name";
+    private final Context mContext = spy(ApplicationProvider.getApplicationContext());
+    @Mock private LocalBluetoothManager mLocalBluetoothManager;
+    @Mock private LocalBluetoothProfileManager mLocalBluetoothProfileManager;
+    @Mock private LocalBluetoothLeBroadcastAssistant mAssistant;
+    @Mock private CachedBluetoothDeviceManager mDeviceManager;
+    @Mock private BluetoothLeBroadcastMetadata mMetadata;
+    @Mock private CachedBluetoothDevice mCachedDevice;
+    @Mock private BluetoothDevice mDevice;
+    private AudioStreamsHelper mHelper;
+
+    @Before
+    public void setUp() {
+        when(mLocalBluetoothManager.getProfileManager()).thenReturn(mLocalBluetoothProfileManager);
+        when(mLocalBluetoothManager.getCachedDeviceManager()).thenReturn(mDeviceManager);
+        when(mLocalBluetoothProfileManager.getLeAudioBroadcastAssistantProfile())
+                .thenReturn(mAssistant);
+        mHelper = spy(new AudioStreamsHelper(mLocalBluetoothManager));
+    }
+
+    @Test
+    public void addSource_noDevice_doNothing() {
+        when(mAssistant.getDevicesMatchingConnectionStates(any()))
+                .thenReturn(Collections.emptyList());
+        mHelper.addSource(mMetadata);
+
+        verify(mAssistant, never()).addSource(any(), any(), anyBoolean());
+    }
+
+    @Test
+    public void addSource_hasDevice() {
+        List<BluetoothDevice> devices = new ArrayList<>();
+        devices.add(mDevice);
+        when(mAssistant.getDevicesMatchingConnectionStates(any())).thenReturn(devices);
+        when(mDeviceManager.findDevice(any())).thenReturn(mCachedDevice);
+        when(mCachedDevice.getDevice()).thenReturn(mDevice);
+        when(mCachedDevice.getGroupId()).thenReturn(GROUP_ID);
+
+        mHelper.addSource(mMetadata);
+
+        verify(mAssistant).addSource(eq(mDevice), eq(mMetadata), anyBoolean());
+    }
+
+    @Test
+    public void removeSource_noDevice_doNothing() {
+        when(mAssistant.getDevicesMatchingConnectionStates(any()))
+                .thenReturn(Collections.emptyList());
+        mHelper.removeSource(BROADCAST_ID_1);
+
+        verify(mAssistant, never()).removeSource(any(), anyInt());
+    }
+
+    @Test
+    public void removeSource_noConnectedSource_doNothing() {
+        List<BluetoothDevice> devices = new ArrayList<>();
+        devices.add(mDevice);
+        when(mAssistant.getDevicesMatchingConnectionStates(any())).thenReturn(devices);
+        BluetoothLeBroadcastReceiveState source = mock(BluetoothLeBroadcastReceiveState.class);
+        when(source.getBroadcastId()).thenReturn(BROADCAST_ID_2);
+        when(mDeviceManager.findDevice(any())).thenReturn(mCachedDevice);
+        when(mCachedDevice.getDevice()).thenReturn(mDevice);
+        when(mCachedDevice.getGroupId()).thenReturn(GROUP_ID);
+        when(mAssistant.getAllSources(any())).thenReturn(ImmutableList.of(source));
+
+        mHelper.removeSource(BROADCAST_ID_1);
+
+        verify(mAssistant, never()).removeSource(any(), anyInt());
+    }
+
+    @Test
+    public void removeSource_hasConnectedSource() {
+        List<BluetoothDevice> devices = new ArrayList<>();
+        devices.add(mDevice);
+        when(mAssistant.getDevicesMatchingConnectionStates(any())).thenReturn(devices);
+        BluetoothLeBroadcastReceiveState source = mock(BluetoothLeBroadcastReceiveState.class);
+        when(source.getBroadcastId()).thenReturn(BROADCAST_ID_2);
+        when(mDeviceManager.findDevice(any())).thenReturn(mCachedDevice);
+        when(mCachedDevice.getDevice()).thenReturn(mDevice);
+        when(mCachedDevice.getGroupId()).thenReturn(GROUP_ID);
+        when(mAssistant.getAllSources(any())).thenReturn(ImmutableList.of(source));
+        List<Long> bisSyncState = new ArrayList<>();
+        bisSyncState.add(1L);
+        when(source.getBisSyncState()).thenReturn(bisSyncState);
+
+        mHelper.removeSource(BROADCAST_ID_2);
+
+        verify(mAssistant).removeSource(eq(mDevice), anyInt());
+    }
+
+    @Test
+    public void removeSource_memberHasConnectedSource() {
+        List<BluetoothDevice> devices = new ArrayList<>();
+        var memberDevice = mock(BluetoothDevice.class);
+        devices.add(mDevice);
+        devices.add(memberDevice);
+        when(mAssistant.getDevicesMatchingConnectionStates(any())).thenReturn(devices);
+        BluetoothLeBroadcastReceiveState source = mock(BluetoothLeBroadcastReceiveState.class);
+        when(source.getBroadcastId()).thenReturn(BROADCAST_ID_2);
+        when(mDeviceManager.findDevice(any())).thenReturn(mCachedDevice);
+        when(mCachedDevice.getDevice()).thenReturn(mDevice);
+        var memberCachedDevice = mock(CachedBluetoothDevice.class);
+        when(memberCachedDevice.getDevice()).thenReturn(memberDevice);
+        when(mCachedDevice.getMemberDevice()).thenReturn(ImmutableSet.of(memberCachedDevice));
+        when(mCachedDevice.getGroupId()).thenReturn(GROUP_ID);
+        when(mAssistant.getAllSources(mDevice)).thenReturn(ImmutableList.of());
+        when(mAssistant.getAllSources(memberDevice)).thenReturn(ImmutableList.of(source));
+        List<Long> bisSyncState = new ArrayList<>();
+        bisSyncState.add(1L);
+        when(source.getBisSyncState()).thenReturn(bisSyncState);
+
+        mHelper.removeSource(BROADCAST_ID_2);
+
+        verify(mAssistant).removeSource(eq(memberDevice), anyInt());
+    }
+
+    @Test
+    public void getAllConnectedSources_noAssistant() {
+        when(mLocalBluetoothProfileManager.getLeAudioBroadcastAssistantProfile()).thenReturn(null);
+        mHelper = new AudioStreamsHelper(mLocalBluetoothManager);
+
+        assertThat(mHelper.getAllConnectedSources()).isEmpty();
+    }
+
+    @Test
+    public void getAllConnectedSources_returnSource() {
+        List<BluetoothDevice> devices = new ArrayList<>();
+        devices.add(mDevice);
+        when(mAssistant.getDevicesMatchingConnectionStates(any())).thenReturn(devices);
+        BluetoothLeBroadcastReceiveState source = mock(BluetoothLeBroadcastReceiveState.class);
+        when(mDeviceManager.findDevice(any())).thenReturn(mCachedDevice);
+        when(mCachedDevice.getDevice()).thenReturn(mDevice);
+        when(mCachedDevice.getGroupId()).thenReturn(GROUP_ID);
+        when(mAssistant.getAllSources(any())).thenReturn(ImmutableList.of(source));
+        List<Long> bisSyncState = new ArrayList<>();
+        bisSyncState.add(1L);
+        when(source.getBisSyncState()).thenReturn(bisSyncState);
+
+        var list = mHelper.getAllConnectedSources();
+        assertThat(list).isNotEmpty();
+        assertThat(list.get(0)).isEqualTo(source);
+    }
+
+    @Test
+    public void startMediaService_noDevice_doNothing() {
+        mHelper.startMediaService(mContext, BROADCAST_ID_1, BROADCAST_NAME);
+
+        verify(mContext, never()).startService(any());
+    }
+
+    @Test
+    public void startMediaService_hasDevice() {
+        List<BluetoothDevice> devices = new ArrayList<>();
+        devices.add(mDevice);
+        when(mAssistant.getDevicesMatchingConnectionStates(any())).thenReturn(devices);
+        BluetoothLeBroadcastReceiveState source = mock(BluetoothLeBroadcastReceiveState.class);
+        when(mDeviceManager.findDevice(any())).thenReturn(mCachedDevice);
+        when(mCachedDevice.getDevice()).thenReturn(mDevice);
+        when(mCachedDevice.getGroupId()).thenReturn(GROUP_ID);
+        when(mAssistant.getAllSources(any())).thenReturn(ImmutableList.of(source));
+        List<Long> bisSyncState = new ArrayList<>();
+        bisSyncState.add(1L);
+        when(source.getBisSyncState()).thenReturn(bisSyncState);
+
+        mHelper.startMediaService(mContext, BROADCAST_ID_1, BROADCAST_NAME);
+
+        verify(mContext).startService(any());
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsProgressCategoryControllerTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsProgressCategoryControllerTest.java
index d43ec81..fd1b649 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsProgressCategoryControllerTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsProgressCategoryControllerTest.java
@@ -168,6 +168,25 @@
     }
 
     @Test
+    public void testShowToast_noError() {
+        mController.showToast(BROADCAST_NAME_1);
+    }
+
+    @Test
+    public void testOnStop_unregister() {
+        mController.onStop(mLifecycleOwner);
+
+        verify(mBluetoothEventManager).unregisterCallback(any());
+    }
+
+    @Test
+    public void testGetFragment_returnFragment() {
+        mController.setFragment(mFragment);
+
+        assertThat(mController.getFragment()).isEqualTo(mFragment);
+    }
+
+    @Test
     public void testOnStart_initNoDevice_showDialog() {
         when(mLeBroadcastAssistant.isSearchInProgress()).thenReturn(true);
 
@@ -245,6 +264,25 @@
     }
 
     @Test
+    public void testOnStart_initHasDevice_scanningInProgress() {
+        // Setup a device
+        ShadowAudioStreamsHelper.setCachedBluetoothDeviceInSharingOrLeConnected(mDevice);
+        when(mLeBroadcastAssistant.isSearchInProgress()).thenReturn(true);
+
+        mController.onStart(mLifecycleOwner);
+        shadowOf(Looper.getMainLooper()).idle();
+
+        verify(mLeBroadcastAssistant).registerServiceCallBack(any(), any());
+        verify(mLeBroadcastAssistant).stopSearchingForSources();
+        verify(mLeBroadcastAssistant).startSearchingForSources(any());
+
+        var dialog = ShadowAlertDialog.getLatestAlertDialog();
+        assertThat(dialog).isNull();
+
+        verify(mController, never()).moveToState(any(), any());
+    }
+
+    @Test
     public void testOnStart_handleSourceFromQrCode() {
         // Setup a device
         ShadowAudioStreamsHelper.setCachedBluetoothDeviceInSharingOrLeConnected(mDevice);
@@ -383,6 +421,49 @@
     }
 
     @Test
+    public void testHandleSourceAddRequest_updateMetadataAndState() {
+        // Setup a device
+        ShadowAudioStreamsHelper.setCachedBluetoothDeviceInSharingOrLeConnected(mDevice);
+
+        var metadata =
+                BluetoothLeBroadcastMetadataExt.INSTANCE.convertToBroadcastMetadata(VALID_METADATA);
+        assertThat(metadata).isNotNull();
+        var metadataWithNoCode =
+                new BluetoothLeBroadcastMetadata.Builder(metadata)
+                        .setBroadcastId(NEWLY_FOUND_BROADCAST_ID)
+                        .setBroadcastName(BROADCAST_NAME_1)
+                        .build();
+        // A new source is found
+        mController.handleSourceFound(metadataWithNoCode);
+
+        ArgumentCaptor<AudioStreamPreference> preferenceCaptor =
+                ArgumentCaptor.forClass(AudioStreamPreference.class);
+        ArgumentCaptor<AudioStreamsProgressCategoryController.AudioStreamState> stateCaptor =
+                ArgumentCaptor.forClass(
+                        AudioStreamsProgressCategoryController.AudioStreamState.class);
+
+        // moving state to SYNCED
+        verify(mController).moveToState(preferenceCaptor.capture(), stateCaptor.capture());
+        var preference = preferenceCaptor.getValue();
+        var state = stateCaptor.getValue();
+
+        assertThat(preference).isNotNull();
+        assertThat(preference.getAudioStreamBroadcastId()).isEqualTo(NEWLY_FOUND_BROADCAST_ID);
+        assertThat(state).isEqualTo(SYNCED);
+
+        var updatedMetadata =
+                new BluetoothLeBroadcastMetadata.Builder(metadataWithNoCode)
+                        .setBroadcastCode(BROADCAST_CODE)
+                        .build();
+        mController.handleSourceAddRequest(preference, updatedMetadata);
+        // state updated to ADD_SOURCE_WAIT_FOR_RESPONSE
+        assertThat(preference.getAudioStreamBroadcastId()).isEqualTo(NEWLY_FOUND_BROADCAST_ID);
+        assertThat(preference.getAudioStreamMetadata().getBroadcastCode())
+                .isEqualTo(BROADCAST_CODE);
+        assertThat(preference.getAudioStreamState()).isEqualTo(ADD_SOURCE_WAIT_FOR_RESPONSE);
+    }
+
+    @Test
     public void testHandleSourceFound_sameIdWithSourceFromQrCode_updateMetadataAndState() {
         // Setup a device
         ShadowAudioStreamsHelper.setCachedBluetoothDeviceInSharingOrLeConnected(mDevice);
@@ -519,6 +600,42 @@
     }
 
     @Test
+    public void testHandleSourceLost_sourceConnected_doNothing() {
+        // Setup a device
+        ShadowAudioStreamsHelper.setCachedBluetoothDeviceInSharingOrLeConnected(mDevice);
+
+        // Setup mPreference so it's not null
+        mController.displayPreference(mScreen);
+
+        // A new source found
+        when(mMetadata.getBroadcastId()).thenReturn(NEWLY_FOUND_BROADCAST_ID);
+        mController.handleSourceFound(mMetadata);
+        shadowOf(Looper.getMainLooper()).idle();
+
+        // A new source found is lost, but the source is still connected
+        BluetoothLeBroadcastReceiveState connected = createConnectedMock(NEWLY_FOUND_BROADCAST_ID);
+        when(mAudioStreamsHelper.getAllConnectedSources()).thenReturn(ImmutableList.of(connected));
+        mController.handleSourceLost(NEWLY_FOUND_BROADCAST_ID);
+        shadowOf(Looper.getMainLooper()).idle();
+
+        ArgumentCaptor<AudioStreamPreference> preferenceToAdd =
+                ArgumentCaptor.forClass(AudioStreamPreference.class);
+        ArgumentCaptor<AudioStreamsProgressCategoryController.AudioStreamState> state =
+                ArgumentCaptor.forClass(
+                        AudioStreamsProgressCategoryController.AudioStreamState.class);
+
+        // Verify a new preference is created with state SYNCED.
+        verify(mController).moveToState(preferenceToAdd.capture(), state.capture());
+        assertThat(preferenceToAdd.getValue()).isNotNull();
+        assertThat(preferenceToAdd.getValue().getAudioStreamBroadcastId())
+                .isEqualTo(NEWLY_FOUND_BROADCAST_ID);
+        assertThat(state.getValue()).isEqualTo(SYNCED);
+
+        // No preference is removed.
+        verify(mPreference, never()).removePreference(any());
+    }
+
+    @Test
     public void testHandleSourceRemoved_removed() {
         // Setup a device
         ShadowAudioStreamsHelper.setCachedBluetoothDeviceInSharingOrLeConnected(mDevice);
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsProgressCategoryPreferenceTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsProgressCategoryPreferenceTest.java
index 337d64d..76bd5ec 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsProgressCategoryPreferenceTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsProgressCategoryPreferenceTest.java
@@ -53,6 +53,8 @@
 
     @Test
     public void addAudioStreamPreference_singlePreference() {
+        mPreference = spy(new AudioStreamsProgressCategoryPreference(mContext, null));
+        when(mPreference.getPreferenceManager()).thenReturn(mPreferenceManager);
         AudioStreamPreference first = new AudioStreamPreference(mContext, null);
         mPreference.addAudioStreamPreference(first, (p1, p2) -> 0);
 
@@ -62,6 +64,8 @@
 
     @Test
     public void addAudioStreamPreference_multiPreference_sorted() {
+        mPreference = spy(new AudioStreamsProgressCategoryPreference(mContext, null, 0));
+        when(mPreference.getPreferenceManager()).thenReturn(mPreferenceManager);
         Comparator<AudioStreamPreference> c =
                 Comparator.comparingInt(AudioStreamPreference::getOrder);
         AudioStreamPreference first = new AudioStreamPreference(mContext, null);
@@ -78,6 +82,8 @@
 
     @Test
     public void removeAudioStreamPreferences_shouldBeEmpty() {
+        mPreference = spy(new AudioStreamsProgressCategoryPreference(mContext, null, 0, 0));
+        when(mPreference.getPreferenceManager()).thenReturn(mPreferenceManager);
         Comparator<AudioStreamPreference> c =
                 Comparator.comparingInt(AudioStreamPreference::getOrder);
         AudioStreamPreference first = new AudioStreamPreference(mContext, null);
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsQrCodeFragmentTest.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsQrCodeFragmentTest.java
index 7d85b7a..06e4837 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsQrCodeFragmentTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/AudioStreamsQrCodeFragmentTest.java
@@ -98,6 +98,22 @@
     }
 
     @Test
+    public void onCreateView_noProfile_noQrCode() {
+        when(mBtProfileManager.getLeAudioBroadcastProfile()).thenReturn(null);
+        FragmentController.setupFragment(
+                mFragment, FragmentActivity.class, /* containerViewId= */ 0, /* bundle= */ null);
+        View view = mFragment.getView();
+
+        assertThat(view).isNotNull();
+        ImageView qrCodeView = view.findViewById(R.id.qrcode_view);
+        TextView passwordView = view.requireViewById(R.id.password);
+        assertThat(qrCodeView).isNotNull();
+        assertThat(qrCodeView.getDrawable()).isNull();
+        assertThat(passwordView).isNotNull();
+        assertThat(passwordView.getText().toString()).isEqualTo("");
+    }
+
+    @Test
     public void onCreateView_noMetadata_noQrCode() {
         List<BluetoothLeBroadcastMetadata> list = new ArrayList<>();
         when(mBroadcast.getAllBroadcastMetadata()).thenReturn(list);
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/testshadows/ShadowAudioStreamsHelper.java b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/testshadows/ShadowAudioStreamsHelper.java
index 13c19ca..051eda7 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/testshadows/ShadowAudioStreamsHelper.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/audiosharing/audiostreams/testshadows/ShadowAudioStreamsHelper.java
@@ -16,6 +16,7 @@
 
 package com.android.settings.connecteddevice.audiosharing.audiostreams.testshadows;
 
+import android.bluetooth.BluetoothLeBroadcastMetadata;
 import android.bluetooth.BluetoothLeBroadcastReceiveState;
 
 import androidx.annotation.Nullable;
@@ -69,4 +70,16 @@
     public LocalBluetoothLeBroadcastAssistant getLeBroadcastAssistant() {
         return sMockHelper.getLeBroadcastAssistant();
     }
+
+    /** Removes sources from LE broadcasts associated for all active sinks based on broadcast Id. */
+    @Implementation
+    public void removeSource(int broadcastId) {
+        sMockHelper.removeSource(broadcastId);
+    }
+
+    /** Adds the specified LE broadcast source to all active sinks. */
+    @Implementation
+    public void addSource(BluetoothLeBroadcastMetadata source) {
+        sMockHelper.addSource(source);
+    }
 }
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/BatteryInfoTest.java b/tests/robotests/src/com/android/settings/fuelgauge/BatteryInfoTest.java
index 7bafc6d..b7e6590 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/BatteryInfoTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/BatteryInfoTest.java
@@ -789,6 +789,40 @@
                 expectedChargeLabel);
     }
 
+    @Test
+    public void getBatteryInfo_longlife_shouldSetLonglife() {
+        var batteryIntent = createIntentForLongLifeTest(/* hasLongLife= */ true);
+
+        var batteryInfo =
+                BatteryInfo.getBatteryInfo(
+                        mContext,
+                        batteryIntent,
+                        mBatteryUsageStats,
+                        /* estimate= */ MOCK_ESTIMATE,
+                        /* elapsedRealtimeUs= */ 0L,
+                        /* shortString= */ false,
+                        /* currentTimeMs= */ 0L);
+
+        assertThat(batteryInfo.isLongLife).isTrue();
+    }
+
+    @Test
+    public void getBatteryInfo_noLonglife_shouldNotLonglife() {
+        var batteryIntent = createIntentForLongLifeTest(/* hasLongLife= */ false);
+
+        var batteryInfo =
+                BatteryInfo.getBatteryInfo(
+                        mContext,
+                        batteryIntent,
+                        mBatteryUsageStats,
+                        /* estimate= */ MOCK_ESTIMATE,
+                        /* elapsedRealtimeUs= */ 0L,
+                        /* shortString= */ false,
+                        /* currentTimeMs= */ 0L);
+
+        assertThat(batteryInfo.isLongLife).isFalse();
+    }
+
     private enum ChargingSpeed {
         FAST,
         REGULAR,
@@ -801,6 +835,15 @@
         DOCKED
     }
 
+    private Intent createIntentForLongLifeTest(Boolean hasLongLife) {
+        return new Intent(Intent.ACTION_BATTERY_CHANGED)
+                .putExtra(
+                        BatteryManager.EXTRA_CHARGING_STATUS,
+                        hasLongLife
+                                ? BatteryManager.CHARGING_POLICY_ADAPTIVE_LONGLIFE
+                                : BatteryManager.CHARGING_POLICY_DEFAULT);
+    }
+
     private Intent createIntentForGetBatteryInfoTest(
             ChargingType chargingType, ChargingSpeed chargingSpeed, int batteryLevel) {
         return createBatteryIntent(
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImplTest.java b/tests/robotests/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImplTest.java
index 6b32ff5..cecf8f0 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImplTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImplTest.java
@@ -165,7 +165,7 @@
 
     @Test
     public void isBatteryDefend_defenderModeAndExtraDefendAreFalse_returnFalse() {
-        mBatteryInfo.isBatteryDefender = false;
+        mBatteryInfo.isLongLife = false;
         doReturn(false).when(mPowerFeatureProvider).isExtraDefend();
 
         assertThat(mPowerFeatureProvider.isBatteryDefend(mBatteryInfo)).isFalse();
@@ -173,7 +173,7 @@
 
     @Test
     public void isBatteryDefend_defenderModeIsFalse_returnFalse() {
-        mBatteryInfo.isBatteryDefender = false;
+        mBatteryInfo.isLongLife = false;
         doReturn(true).when(mPowerFeatureProvider).isExtraDefend();
 
         assertThat(mPowerFeatureProvider.isBatteryDefend(mBatteryInfo)).isFalse();
@@ -181,7 +181,7 @@
 
     @Test
     public void isBatteryDefend_defenderModeAndExtraDefendAreTrue_returnFalse() {
-        mBatteryInfo.isBatteryDefender = true;
+        mBatteryInfo.isLongLife = true;
         doReturn(true).when(mPowerFeatureProvider).isExtraDefend();
 
         assertThat(mPowerFeatureProvider.isBatteryDefend(mBatteryInfo)).isFalse();
@@ -189,7 +189,7 @@
 
     @Test
     public void isBatteryDefend_extraDefendIsFalse_returnTrue() {
-        mBatteryInfo.isBatteryDefender = true;
+        mBatteryInfo.isLongLife = true;
         doReturn(false).when(mPowerFeatureProvider).isExtraDefend();
 
         assertThat(mPowerFeatureProvider.isBatteryDefend(mBatteryInfo)).isTrue();
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/batterytip/detectors/BatteryDefenderDetectorTest.java b/tests/robotests/src/com/android/settings/fuelgauge/batterytip/detectors/BatteryDefenderDetectorTest.java
index ab1ceb5..7643c41 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/batterytip/detectors/BatteryDefenderDetectorTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/batterytip/detectors/BatteryDefenderDetectorTest.java
@@ -18,15 +18,12 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.mockito.Mockito.when;
-
 import android.content.Context;
 
 import androidx.test.core.app.ApplicationProvider;
 
 import com.android.settings.fuelgauge.BatteryInfo;
 import com.android.settings.fuelgauge.batterytip.tips.BatteryTip;
-import com.android.settings.testutils.FakeFeatureFactory;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -41,28 +38,23 @@
     @Mock private BatteryInfo mBatteryInfo;
     private BatteryDefenderDetector mBatteryDefenderDetector;
 
-    private FakeFeatureFactory mFakeFeatureFactory;
-
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         final Context context = ApplicationProvider.getApplicationContext();
         mBatteryDefenderDetector = new BatteryDefenderDetector(mBatteryInfo, context);
-        mFakeFeatureFactory = FakeFeatureFactory.setupForTest();
     }
 
     @Test
     public void detect_notBatteryDefend_tipInvisible() {
-        when(mFakeFeatureFactory.powerUsageFeatureProvider.isBatteryDefend(mBatteryInfo))
-                .thenReturn(false);
+        mBatteryInfo.isBatteryDefender = false;
 
         assertThat(mBatteryDefenderDetector.detect().isVisible()).isFalse();
     }
 
     @Test
     public void detect_isBatteryDefend_tipNew() {
-        when(mFakeFeatureFactory.powerUsageFeatureProvider.isBatteryDefend(mBatteryInfo))
-                .thenReturn(true);
+        mBatteryInfo.isBatteryDefender = true;
 
         assertThat(mBatteryDefenderDetector.detect().getState())
                 .isEqualTo(BatteryTip.StateType.NEW);
diff --git a/tests/robotests/src/com/android/settings/network/NetworkProviderSettingsTest.java b/tests/robotests/src/com/android/settings/network/NetworkProviderSettingsTest.java
index fb29e05..0161178 100644
--- a/tests/robotests/src/com/android/settings/network/NetworkProviderSettingsTest.java
+++ b/tests/robotests/src/com/android/settings/network/NetworkProviderSettingsTest.java
@@ -51,6 +51,7 @@
 import android.os.PowerManager;
 import android.os.UserManager;
 import android.provider.Settings;
+import android.telephony.SubscriptionManager;
 import android.view.ContextMenu;
 import android.view.Menu;
 import android.view.MenuItem;
@@ -319,7 +320,8 @@
         mNetworkProviderSettings.onCreate(Bundle.EMPTY);
 
         verify(mDataUsagePreference).setVisible(true);
-        verify(mDataUsagePreference).setTemplate(any(), eq(0) /*subId*/);
+        verify(mDataUsagePreference)
+                .setTemplate(any(), eq(SubscriptionManager.INVALID_SUBSCRIPTION_ID));
     }
 
     @Test
diff --git a/tests/robotests/src/com/android/settings/notification/modes/ManualDurationHelperTest.java b/tests/robotests/src/com/android/settings/notification/modes/ManualDurationHelperTest.java
new file mode 100644
index 0000000..18ee2cf
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/notification/modes/ManualDurationHelperTest.java
@@ -0,0 +1,77 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.notification.modes;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.content.ContentResolver;
+import android.content.Context;
+import android.provider.Settings;
+
+import com.android.settings.R;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+
+@RunWith(RobolectricTestRunner.class)
+public class ManualDurationHelperTest {
+    private Context mContext;
+    private ContentResolver mContentResolver;
+
+    private ManualDurationHelper mHelper;
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+        mContext = RuntimeEnvironment.application;
+        mContentResolver = RuntimeEnvironment.application.getContentResolver();
+
+        mHelper = new ManualDurationHelper(mContext);
+    }
+
+    @Test
+    public void getDurationSummary_durationForever() {
+        Settings.Secure.putInt(mContentResolver, Settings.Secure.ZEN_DURATION,
+                Settings.Secure.ZEN_DURATION_FOREVER);
+        assertThat(mHelper.getSummary()).isEqualTo(
+                mContext.getString(R.string.zen_mode_duration_summary_forever));
+    }
+
+    @Test
+    public void getDurationSummary_durationPrompt() {
+        Settings.Secure.putInt(mContentResolver, Settings.Secure.ZEN_DURATION,
+                Settings.Secure.ZEN_DURATION_PROMPT);
+        assertThat(mHelper.getSummary()).isEqualTo(
+                mContext.getString(R.string.zen_mode_duration_summary_always_prompt));
+    }
+
+    @Test
+    public void getDurationSummary_durationCustom() {
+        // minutes
+        Settings.Secure.putInt(mContentResolver, Settings.Secure.ZEN_DURATION, 45);
+        assertThat(mHelper.getSummary()).isEqualTo("45 minutes");
+
+        // hours
+        Settings.Secure.putInt(mContentResolver, Settings.Secure.ZEN_DURATION, 300);
+        assertThat(mHelper.getSummary()).isEqualTo("5 hours");
+    }
+
+}
diff --git a/tests/robotests/src/com/android/settings/notification/modes/ManualDurationPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/notification/modes/ManualDurationPreferenceControllerTest.java
new file mode 100644
index 0000000..0a600c0
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/notification/modes/ManualDurationPreferenceControllerTest.java
@@ -0,0 +1,91 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.notification.modes;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.app.AutomaticZenRule;
+import android.app.Flags;
+import android.content.ContentResolver;
+import android.content.Context;
+import android.net.Uri;
+import android.platform.test.annotations.EnableFlags;
+import android.platform.test.flag.junit.SetFlagsRule;
+import android.provider.Settings;
+
+import androidx.fragment.app.Fragment;
+import androidx.preference.Preference;
+
+import com.android.settingslib.notification.modes.ZenMode;
+import com.android.settingslib.notification.modes.ZenModesBackend;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+
+@RunWith(RobolectricTestRunner.class)
+@EnableFlags(Flags.FLAG_MODES_UI)
+public class ManualDurationPreferenceControllerTest {
+    @Rule
+    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+
+    private Context mContext;
+    private ContentResolver mContentResolver;
+
+    @Mock
+    private ZenModesBackend mBackend;
+
+    @Mock
+    private Fragment mParent;
+
+    private ManualDurationPreferenceController mPrefController;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+
+        mContext = RuntimeEnvironment.application;
+        mContentResolver = RuntimeEnvironment.application.getContentResolver();
+        mPrefController = new ManualDurationPreferenceController(mContext, "key", mParent,
+                mBackend);
+    }
+
+    @Test
+    public void testIsAvailable_onlyForManualDnd() {
+        assertThat(mPrefController.isAvailable(TestModeBuilder.EXAMPLE)).isFalse();
+
+        ZenMode manualDnd = ZenMode.manualDndMode(
+                new AutomaticZenRule.Builder("id", Uri.EMPTY).build(), false);
+        assertThat(mPrefController.isAvailable(manualDnd)).isTrue();
+    }
+
+    @Test
+    public void testUpdateState_durationSummary() {
+        Settings.Secure.putInt(mContentResolver, Settings.Secure.ZEN_DURATION,
+                45 /* minutes */);
+
+        Preference pref = new Preference(mContext);
+        mPrefController.updateState(pref, TestModeBuilder.EXAMPLE);
+
+        assertThat(pref.getSummary().toString()).contains("45");
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/notification/modes/ZenModeButtonPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/notification/modes/ZenModeButtonPreferenceControllerTest.java
index 625f231..5869c6b 100644
--- a/tests/robotests/src/com/android/settings/notification/modes/ZenModeButtonPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/notification/modes/ZenModeButtonPreferenceControllerTest.java
@@ -23,12 +23,18 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.app.AutomaticZenRule;
 import android.app.Flags;
+import android.content.ContentResolver;
 import android.content.Context;
+import android.net.Uri;
 import android.platform.test.annotations.EnableFlags;
 import android.platform.test.flag.junit.SetFlagsRule;
+import android.provider.Settings;
 import android.widget.Button;
 
+import androidx.fragment.app.Fragment;
+
 import com.android.settingslib.notification.modes.ZenMode;
 import com.android.settingslib.notification.modes.ZenModesBackend;
 import com.android.settingslib.widget.LayoutPreference;
@@ -42,6 +48,8 @@
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
 
+import java.time.Duration;
+
 @EnableFlags(Flags.FLAG_MODES_UI)
 @RunWith(RobolectricTestRunner.class)
 public final class ZenModeButtonPreferenceControllerTest {
@@ -51,19 +59,24 @@
     @Rule
     public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
 
-
     private Context mContext;
+    private ContentResolver mContentResolver;
+
     @Mock
     private ZenModesBackend mBackend;
 
+    @Mock
+    private Fragment mParent;
+
     @Before
     public void setup() {
         MockitoAnnotations.initMocks(this);
 
         mContext = RuntimeEnvironment.application;
+        mContentResolver = RuntimeEnvironment.application.getContentResolver();
 
         mController = new ZenModeButtonPreferenceController(
-                mContext, "something", mBackend);
+                mContext, "something", mParent, mBackend);
     }
 
     @Test
@@ -162,4 +175,34 @@
         button.callOnClick();
         verify(mBackend).activateMode(zenMode, null);
     }
+
+    @Test
+    public void updateStateThenClick_withDuration() {
+        Settings.Secure.putInt(mContentResolver, Settings.Secure.ZEN_DURATION,
+                45 /* minutes */);
+        Button button = new Button(mContext);
+        LayoutPreference pref = mock(LayoutPreference.class);
+        when(pref.findViewById(anyInt())).thenReturn(button);
+        ZenMode zenMode = ZenMode.manualDndMode(
+                new AutomaticZenRule.Builder("manual", Uri.EMPTY).build(), false);
+
+        mController.updateZenMode(pref, zenMode);
+        button.callOnClick();
+        verify(mBackend).activateMode(zenMode, Duration.ofMinutes(45));
+    }
+
+    @Test
+    public void updateStateThenClick_durationForever() {
+        Settings.Secure.putInt(mContentResolver, Settings.Secure.ZEN_DURATION,
+                Settings.Secure.ZEN_DURATION_FOREVER);
+        Button button = new Button(mContext);
+        LayoutPreference pref = mock(LayoutPreference.class);
+        when(pref.findViewById(anyInt())).thenReturn(button);
+        ZenMode zenMode = ZenMode.manualDndMode(
+                new AutomaticZenRule.Builder("manual", Uri.EMPTY).build(), false);
+
+        mController.updateZenMode(pref, zenMode);
+        button.callOnClick();
+        verify(mBackend).activateMode(zenMode, null);
+    }
 }
\ No newline at end of file
diff --git a/tests/robotests/src/com/android/settings/search/SearchFeatureProviderImplTest.java b/tests/robotests/src/com/android/settings/search/SearchFeatureProviderImplTest.java
index 8a7419b..599649b 100644
--- a/tests/robotests/src/com/android/settings/search/SearchFeatureProviderImplTest.java
+++ b/tests/robotests/src/com/android/settings/search/SearchFeatureProviderImplTest.java
@@ -124,8 +124,8 @@
     }
 
     @Test(expected = IllegalArgumentException.class)
-    public void verifyLaunchSearchResultPageCaller_nullCaller_shouldCrash() {
-        mProvider.verifyLaunchSearchResultPageCaller(mActivity, null /* caller */);
+    public void verifyLaunchSearchResultPageCaller_emptyCaller_shouldCrash() {
+        mProvider.verifyLaunchSearchResultPageCaller(mActivity, "");
     }
 
     @Test(expected = SecurityException.class)
diff --git a/tests/robotests/src/com/android/settings/security/trustagent/TrustAgentsPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/security/trustagent/TrustAgentsPreferenceControllerTest.java
index 8339798..766855c 100644
--- a/tests/robotests/src/com/android/settings/security/trustagent/TrustAgentsPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/security/trustagent/TrustAgentsPreferenceControllerTest.java
@@ -21,14 +21,10 @@
 import android.app.admin.DevicePolicyManager;
 import android.content.ComponentName;
 import android.content.Context;
-import android.content.Intent;
-import android.content.pm.ApplicationInfo;
-import android.content.pm.PackageManager;
-import android.content.pm.ResolveInfo;
-import android.content.pm.ServiceInfo;
+import android.content.IntentFilter;
+import android.content.pm.PackageInfo;
 import android.service.trust.TrustAgentService;
 
-import android.text.TextUtils;
 import androidx.preference.PreferenceManager;
 import androidx.preference.PreferenceScreen;
 
@@ -38,34 +34,30 @@
 import com.android.settings.testutils.shadow.ShadowRestrictedLockUtilsInternal;
 import com.android.settingslib.RestrictedSwitchPreference;
 
-import org.junit.After;
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
 import org.robolectric.Shadows;
 import org.robolectric.annotation.Config;
-import org.robolectric.annotation.Implementation;
-import org.robolectric.annotation.Implements;
 import org.robolectric.shadows.ShadowApplicationPackageManager;
 
-import java.util.ArrayList;
-import java.util.List;
-
 @RunWith(RobolectricTestRunner.class)
 @Config(shadows = {
         ShadowLockPatternUtils.class,
         ShadowRestrictedLockUtilsInternal.class,
-        ShadowDevicePolicyManager.class,
-        ShadowApplicationPackageManager.class,
-        TrustAgentsPreferenceControllerTest.ShadowTrustAgentManager.class
+        ShadowDevicePolicyManager.class, ShadowApplicationPackageManager.class
 })
 public class TrustAgentsPreferenceControllerTest {
-
-    private static final Intent TEST_INTENT =
-            new Intent(TrustAgentService.SERVICE_INTERFACE);
+    private static final ComponentName TRUST_AGENT_A = new ComponentName(
+            "test.data.packageA", "clzAAA");
+    private static final ComponentName TRUST_AGENT_B = new ComponentName(
+            "test.data.packageB", "clzBBB");
+    private static final ComponentName TRUST_AGENT_C = new ComponentName(
+            "test.data.packageC", "clzCCC");
+    private static final ComponentName TRUST_AGENT_D = new ComponentName(
+            "test.data.packageD", "clzDDD");
 
     private Context mContext;
     private ShadowApplicationPackageManager mPackageManager;
@@ -84,11 +76,6 @@
         mPreferenceScreen.setKey("pref_key");
     }
 
-    @After
-    public void tearDown() {
-        ShadowTrustAgentManager.clearPermissionGrantedList();
-    }
-
     @Test
     public void getAvailabilityStatus_byDefault_shouldBeShown() {
         assertThat(mController.getAvailabilityStatus())
@@ -97,8 +84,7 @@
 
     @Test
     public void onStart_noTrustAgent_shouldNotAddPreference() {
-        final List<ResolveInfo> availableAgents = createFakeAvailableAgents();
-        mPackageManager.setResolveInfosForIntent(TEST_INTENT, availableAgents);
+        installFakeAvailableAgents(/* grantPermission= */ false);
 
         mController.displayPreference(mPreferenceScreen);
         mController.onStart();
@@ -106,57 +92,34 @@
         assertThat(mPreferenceScreen.getPreferenceCount()).isEqualTo(0);
     }
 
-    @Ignore("b/313612480")
     @Test
-    public void
-    onStart_hasAUninstalledTrustAgent_shouldRemoveOnePreferenceAndLeaveTwoPreferences() {
-        final List<ResolveInfo> availableAgents = createFakeAvailableAgents();
-        final ResolveInfo uninstalledTrustAgent = availableAgents.get(0);
-
-        for (ResolveInfo rInfo : availableAgents) {
-            ShadowTrustAgentManager.grantPermissionToResolveInfo(rInfo);
-        }
-        mPackageManager.setResolveInfosForIntent(TEST_INTENT, availableAgents);
+    public void onStart_uninstalledTrustAgent_shouldRemoveOnePreferenceAndLeaveTwoPreferences() {
+        installFakeAvailableAgents(/* grantPermission= */ true);
         mController.displayPreference(mPreferenceScreen);
         mController.onStart();
-        availableAgents.remove(uninstalledTrustAgent);
+        uninstallAgent(TRUST_AGENT_A);
 
-        mPackageManager.setResolveInfosForIntent(TEST_INTENT, availableAgents);
         mController.onStart();
 
         assertThat(mPreferenceScreen.getPreferenceCount()).isEqualTo(2);
     }
 
-    @Ignore("b/313612480")
     @Test
     public void onStart_hasANewTrustAgent_shouldAddOnePreferenceAndHaveFourPreferences() {
-        final List<ResolveInfo> availableAgents = createFakeAvailableAgents();
-        final ComponentName newComponentName = new ComponentName("test.data.packageD", "clzDDD");
-        final ResolveInfo newTrustAgent = createFakeResolveInfo(newComponentName);
-        for (ResolveInfo rInfo : availableAgents) {
-            ShadowTrustAgentManager.grantPermissionToResolveInfo(rInfo);
-        }
-        mPackageManager.setResolveInfosForIntent(TEST_INTENT, availableAgents);
+        installFakeAvailableAgents(/* grantPermission= */ true);
         mController.displayPreference(mPreferenceScreen);
         mController.onStart();
-        availableAgents.add(newTrustAgent);
-        ShadowTrustAgentManager.grantPermissionToResolveInfo(newTrustAgent);
+        installFakeAvailableAgent(TRUST_AGENT_D, /* grantPermission= */ true);
 
-        mPackageManager.setResolveInfosForIntent(TEST_INTENT, availableAgents);
         mController.onStart();
 
         assertThat(mPreferenceScreen.getPreferenceCount()).isEqualTo(4);
     }
 
-    @Ignore("b/313612480")
     @Test
     public void onStart_hasUnrestrictedTrustAgent_shouldAddThreeChangeablePreferences() {
         ShadowRestrictedLockUtilsInternal.setKeyguardDisabledFeatures(0);
-        final List<ResolveInfo> availableAgents = createFakeAvailableAgents();
-        for (ResolveInfo rInfo : availableAgents) {
-            ShadowTrustAgentManager.grantPermissionToResolveInfo(rInfo);
-        }
-        mPackageManager.setResolveInfosForIntent(TEST_INTENT, availableAgents);
+        installFakeAvailableAgents(/* grantPermission= */ true);
 
         mController.displayPreference(mPreferenceScreen);
         mController.onStart();
@@ -169,14 +132,9 @@
         }
     }
 
-    @Ignore("b/313612480")
     @Test
-    public void onStart_hasRestrictedTructAgent_shouldAddThreeUnchangeablePreferences() {
-        final List<ResolveInfo> availableAgents = createFakeAvailableAgents();
-        for (ResolveInfo rInfo : availableAgents) {
-            ShadowTrustAgentManager.grantPermissionToResolveInfo(rInfo);
-        }
-        mPackageManager.setResolveInfosForIntent(TEST_INTENT, availableAgents);
+    public void onStart_hasRestrictedTrustAgent_shouldAddThreeUnchangeablePreferences() {
+        installFakeAvailableAgents(/* grantPermission= */ true);
         ShadowRestrictedLockUtilsInternal.setKeyguardDisabledFeatures(
                 DevicePolicyManager.KEYGUARD_DISABLE_TRUST_AGENTS);
 
@@ -191,51 +149,30 @@
         }
     }
 
-    private List<ResolveInfo> createFakeAvailableAgents() {
-        final List<ComponentName> componentNames = new ArrayList<>();
-        componentNames.add(new ComponentName("test.data.packageA", "clzAAA"));
-        componentNames.add(new ComponentName("test.data.packageB", "clzBBB"));
-        componentNames.add(new ComponentName("test.data.packageC", "clzCCC"));
-        final List<ResolveInfo> result = new ArrayList<>();
-        for (ComponentName cn : componentNames) {
-            final ResolveInfo ri = createFakeResolveInfo(cn);
-            result.add(ri);
-        }
-        return result;
+    private void installFakeAvailableAgents(boolean grantPermission) {
+        installFakeAvailableAgent(TRUST_AGENT_A, grantPermission);
+        installFakeAvailableAgent(TRUST_AGENT_B, grantPermission);
+        installFakeAvailableAgent(TRUST_AGENT_C, grantPermission);
     }
 
-    private ResolveInfo createFakeResolveInfo(ComponentName cn) {
-        final ResolveInfo ri = new ResolveInfo();
-        ri.serviceInfo = new ServiceInfo();
-        ri.serviceInfo.packageName = cn.getPackageName();
-        ri.serviceInfo.name = cn.getClassName();
-        ri.serviceInfo.applicationInfo = new ApplicationInfo();
-        ri.serviceInfo.applicationInfo.packageName = cn.getPackageName();
-        ri.serviceInfo.applicationInfo.name = cn.getClassName();
-        return ri;
+    private void installFakeAvailableAgent(ComponentName name,
+            boolean grantPermission) {
+        mPackageManager.addServiceIfNotPresent(name);
+        mPackageManager.addIntentFilterForService(name,
+                new IntentFilter(TrustAgentService.SERVICE_INTERFACE));
+        if (!grantPermission) {
+            return;
+        }
+        PackageInfo pkgInfo = mPackageManager.getInternalMutablePackageInfo(
+                name.getPackageName());
+        pkgInfo.requestedPermissions =
+                new String[]{android.Manifest.permission.PROVIDE_TRUST_AGENT};
+        pkgInfo.requestedPermissionsFlags =
+                new int[]{PackageInfo.REQUESTED_PERMISSION_GRANTED};
     }
 
-    @Implements(TrustAgentManager.class)
-    public static class ShadowTrustAgentManager {
-        private final static List<ResolveInfo> sPermissionGrantedList = new ArrayList<>();
-
-        @Implementation
-        protected boolean shouldProvideTrust(ResolveInfo resolveInfo, PackageManager pm) {
-            for (ResolveInfo info : sPermissionGrantedList) {
-                if (info.serviceInfo.equals(resolveInfo.serviceInfo)) {
-                    return true;
-                }
-            }
-
-            return false;
-        }
-
-        private static void grantPermissionToResolveInfo(ResolveInfo rInfo) {
-            sPermissionGrantedList.add(rInfo);
-        }
-
-        private static void clearPermissionGrantedList() {
-            sPermissionGrantedList.clear();
-        }
+    private void uninstallAgent(ComponentName name) {
+        mPackageManager.removeService(name);
+        mPackageManager.removePackage(name.getPackageName());
     }
 }
diff --git a/tests/spa_unit/src/com/android/settings/spa/SpaBridgeActivityTest.kt b/tests/spa_unit/src/com/android/settings/spa/SpaBridgeActivityTest.kt
new file mode 100644
index 0000000..e29bd96
--- /dev/null
+++ b/tests/spa_unit/src/com/android/settings/spa/SpaBridgeActivityTest.kt
@@ -0,0 +1,95 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.spa
+
+import android.app.Activity
+import android.content.ComponentName
+import android.content.pm.ActivityInfo
+import android.content.pm.PackageManager
+import android.os.Bundle
+import androidx.core.os.bundleOf
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.settings.SettingsActivity.META_DATA_KEY_HIGHLIGHT_MENU_KEY
+import com.android.settings.spa.SpaBridgeActivity.Companion.META_DATA_KEY_DESTINATION
+import com.android.settings.spa.SpaBridgeActivity.Companion.getDestination
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doAnswer
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
+
+@RunWith(AndroidJUnit4::class)
+class SpaBridgeActivityTest {
+    private var activityMetadata: Bundle = bundleOf()
+
+    private val mockPackageManager =
+        mock<PackageManager> {
+            on {
+                getActivityInfo(eq(COMPONENT_NAME), any<PackageManager.ComponentInfoFlags>())
+            } doAnswer { ActivityInfo().apply { metaData = activityMetadata } }
+        }
+
+    private val activity =
+        mock<Activity> {
+            on { componentName } doReturn COMPONENT_NAME
+            on { packageManager } doReturn mockPackageManager
+        }
+
+    @Test
+    fun getDestination_noDestination_returnNull() {
+        activityMetadata = bundleOf()
+
+        val destination = activity.getDestination()
+
+        assertThat(destination).isNull()
+    }
+
+    @Test
+    fun getDestination_withoutHighlightMenuKey() {
+        activityMetadata = bundleOf(META_DATA_KEY_DESTINATION to DESTINATION)
+
+        val (destination, highlightMenuKey) = activity.getDestination()!!
+
+        assertThat(destination).isEqualTo(DESTINATION)
+        assertThat(highlightMenuKey).isNull()
+    }
+
+    @Test
+    fun getDestination_withHighlightMenuKey() {
+        activityMetadata =
+            bundleOf(
+                META_DATA_KEY_DESTINATION to DESTINATION,
+                META_DATA_KEY_HIGHLIGHT_MENU_KEY to HIGHLIGHT_MENU_KEY,
+            )
+
+        val (destination, highlightMenuKey) = activity.getDestination()!!
+
+        assertThat(destination).isEqualTo(DESTINATION)
+        assertThat(highlightMenuKey).isEqualTo(HIGHLIGHT_MENU_KEY)
+    }
+
+    private companion object {
+        const val PACKAGE_NAME = "package.name"
+        const val ACTIVITY_NAME = "ActivityName"
+        val COMPONENT_NAME = ComponentName(PACKAGE_NAME, ACTIVITY_NAME)
+        const val DESTINATION = "Destination"
+        const val HIGHLIGHT_MENU_KEY = "apps"
+    }
+}
diff --git a/tests/spa_unit/src/com/android/settings/spa/SpaDestinationTest.kt b/tests/spa_unit/src/com/android/settings/spa/SpaDestinationTest.kt
index 0b9eb22..ee658c1 100644
--- a/tests/spa_unit/src/com/android/settings/spa/SpaDestinationTest.kt
+++ b/tests/spa_unit/src/com/android/settings/spa/SpaDestinationTest.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 The Android Open Source Project
+ * Copyright (C) 2024 The Android Open Source Project
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,81 +17,32 @@
 package com.android.settings.spa
 
 import android.app.Activity
-import android.content.ComponentName
-import android.content.pm.ActivityInfo
-import android.content.pm.PackageManager
-import android.os.Bundle
-import androidx.core.os.bundleOf
 import androidx.test.ext.junit.runners.AndroidJUnit4
-import com.android.settings.SettingsActivity.META_DATA_KEY_HIGHLIGHT_MENU_KEY
-import com.android.settings.spa.SpaDestination.Companion.META_DATA_KEY_DESTINATION
-import com.android.settings.spa.SpaDestination.Companion.getDestination
-import com.google.common.truth.Truth.assertThat
+import com.android.settingslib.spa.framework.util.KEY_DESTINATION
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.kotlin.any
-import org.mockito.kotlin.doAnswer
-import org.mockito.kotlin.doReturn
-import org.mockito.kotlin.eq
+import org.mockito.kotlin.argThat
 import org.mockito.kotlin.mock
+import org.mockito.kotlin.verify
 
 @RunWith(AndroidJUnit4::class)
 class SpaDestinationTest {
-    private var activityMetadata: Bundle = bundleOf()
 
-    private val mockPackageManager = mock<PackageManager> {
-        on {
-            getActivityInfo(
-                eq(COMPONENT_NAME),
-                any<PackageManager.ComponentInfoFlags>()
-            )
-        } doAnswer {
-            ActivityInfo().apply { metaData = activityMetadata }
-        }
-    }
-
-    private val activity = mock<Activity> {
-        on { componentName } doReturn COMPONENT_NAME
-        on { packageManager } doReturn mockPackageManager
-    }
+    private val activity = mock<Activity>()
 
     @Test
-    fun getDestination_noDestination_returnNull() {
-        activityMetadata = bundleOf()
+    fun startFromExportedActivity() {
+        val spaDestination = SpaDestination(destination = DESTINATION, highlightMenuKey = null)
 
-        val destination = activity.getDestination()
+        spaDestination.startFromExportedActivity(activity)
 
-        assertThat(destination).isNull()
-    }
-
-    @Test
-    fun getDestination_withoutHighlightMenuKey() {
-        activityMetadata = bundleOf(META_DATA_KEY_DESTINATION to DESTINATION)
-
-        val (destination, highlightMenuKey) = activity.getDestination()!!
-
-        assertThat(destination).isEqualTo(DESTINATION)
-        assertThat(highlightMenuKey).isNull()
-    }
-
-    @Test
-    fun getDestination_withHighlightMenuKey() {
-        activityMetadata = bundleOf(
-            META_DATA_KEY_DESTINATION to DESTINATION,
-            META_DATA_KEY_HIGHLIGHT_MENU_KEY to HIGHLIGHT_MENU_KEY,
-        )
-
-        val (destination, highlightMenuKey) = activity.getDestination()!!
-
-        assertThat(destination).isEqualTo(DESTINATION)
-        assertThat(highlightMenuKey).isEqualTo(HIGHLIGHT_MENU_KEY)
+        verify(activity).startActivity(argThat {
+            component!!.className == SpaActivity::class.qualifiedName
+            getStringExtra(KEY_DESTINATION) == DESTINATION
+        })
     }
 
     private companion object {
-        const val PACKAGE_NAME = "package.name"
-        const val ACTIVITY_NAME = "ActivityName"
-        val COMPONENT_NAME = ComponentName(PACKAGE_NAME, ACTIVITY_NAME)
         const val DESTINATION = "Destination"
-        const val HIGHLIGHT_MENU_KEY = "apps"
     }
 }
diff --git a/tests/spa_unit/src/com/android/settings/spa/network/PrimarySimRepositoryTest.kt b/tests/spa_unit/src/com/android/settings/spa/network/PrimarySimRepositoryTest.kt
new file mode 100644
index 0000000..459afa9
--- /dev/null
+++ b/tests/spa_unit/src/com/android/settings/spa/network/PrimarySimRepositoryTest.kt
@@ -0,0 +1,206 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.spa.network
+
+import android.content.Context
+import android.content.res.Resources
+import android.telephony.SubscriptionInfo
+import android.telephony.SubscriptionManager
+import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.settings.R
+import com.android.settingslib.spa.widget.preference.ListPreferenceOption
+import com.google.common.truth.Truth.assertThat
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doAnswer
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.spy
+import org.mockito.kotlin.stub
+
+@RunWith(AndroidJUnit4::class)
+class PrimarySimRepositoryTest {
+    @get:Rule
+    val composeTestRule = createComposeRule()
+
+    private val mockSubscriptionManager = mock<SubscriptionManager> {
+        on { addOnSubscriptionsChangedListener(any(), any()) } doAnswer {
+            val listener = it.arguments[1] as SubscriptionManager.OnSubscriptionsChangedListener
+            listener.onSubscriptionsChanged()
+        }
+        on { getPhoneNumber(SUB_ID_1) } doReturn NUMBER_1
+        on { getPhoneNumber(SUB_ID_2) } doReturn NUMBER_2
+    }
+
+    private val context: Context = spy(ApplicationProvider.getApplicationContext()) {
+        on { getSystemService(SubscriptionManager::class.java) } doReturn mockSubscriptionManager
+    }
+    private val spyResources: Resources = spy(context.resources)
+
+    @Test
+    fun getPrimarySimInfo_oneSim_returnNull() {
+        val simList = listOf(
+            SUB_INFO_1,
+        )
+
+        val primarySimInfo = PrimarySimRepository(context).getPrimarySimInfo(simList)
+
+        assertThat(primarySimInfo).isNull()
+    }
+
+    @Test
+    fun getPrimarySimInfo_verifyCallsList() {
+        val simList = listOf(
+            SUB_INFO_1,
+            SUB_INFO_2
+        )
+        val expectedList = listOf(
+            ListPreferenceOption(
+                id = SUB_INFO_1.subscriptionId,
+                text = "${SUB_INFO_1.displayName}",
+                summary = NUMBER_1
+            ),
+            ListPreferenceOption(
+                id = SUB_INFO_2.subscriptionId,
+                text = "${SUB_INFO_2.displayName}",
+                summary = NUMBER_2
+            ),
+            ListPreferenceOption(
+                id = SubscriptionManager.INVALID_SUBSCRIPTION_ID,
+                text = context.getString(R.string.sim_calls_ask_first_prefs_title),
+            ),
+        )
+
+        val primarySimInfo = PrimarySimRepository(context).getPrimarySimInfo(simList)
+
+        assertThat(primarySimInfo).isNotNull()
+        assertThat(primarySimInfo?.callsList).isEqualTo(expectedList)
+    }
+
+    @Test
+    fun getPrimarySimInfo_verifySmsList() {
+        val simList = listOf(
+            SUB_INFO_1,
+            SUB_INFO_2
+        )
+        val expectedList = listOf(
+            ListPreferenceOption(
+                id = SUB_INFO_1.subscriptionId,
+                text = "${SUB_INFO_1.displayName}",
+                summary = NUMBER_1
+            ),
+            ListPreferenceOption(
+                id = SUB_INFO_2.subscriptionId,
+                text = "${SUB_INFO_2.displayName}",
+                summary = NUMBER_2
+            ),
+            ListPreferenceOption(
+                id = SubscriptionManager.INVALID_SUBSCRIPTION_ID,
+                text = context.getString(R.string.sim_calls_ask_first_prefs_title),
+            ),
+        )
+
+        val primarySimInfo = PrimarySimRepository(context).getPrimarySimInfo(simList)
+
+        assertThat(primarySimInfo).isNotNull()
+        assertThat(primarySimInfo?.smsList).isEqualTo(expectedList)
+    }
+
+    @Test
+    fun getPrimarySimInfo_noAskEveryTime_verifySmsList() {
+        val simList = listOf(
+            SUB_INFO_1,
+            SUB_INFO_2
+        )
+        context.stub {
+            on { resources } doReturn spyResources
+        }
+        spyResources.stub {
+            on {
+                getBoolean(com.android.internal.R.bool.config_sms_ask_every_time_support)
+            } doReturn false
+        }
+        val expectedList = listOf(
+            ListPreferenceOption(
+                id = SUB_INFO_1.subscriptionId,
+                text = "${SUB_INFO_1.displayName}",
+                summary = NUMBER_1
+            ),
+            ListPreferenceOption(
+                id = SUB_INFO_2.subscriptionId,
+                text = "${SUB_INFO_2.displayName}",
+                summary = NUMBER_2
+            ),
+        )
+
+        val primarySimInfo = PrimarySimRepository(context).getPrimarySimInfo(simList)
+
+        assertThat(primarySimInfo).isNotNull()
+        assertThat(primarySimInfo?.smsList).isEqualTo(expectedList)
+    }
+
+    @Test
+    fun getPrimarySimInfo_verifyDataList() {
+        val simList = listOf(
+            SUB_INFO_1,
+            SUB_INFO_2
+        )
+        val expectedList = listOf(
+            ListPreferenceOption(
+                id = SUB_INFO_1.subscriptionId,
+                text = "${SUB_INFO_1.displayName}",
+                summary = NUMBER_1
+            ),
+            ListPreferenceOption(
+                id = SUB_INFO_2.subscriptionId,
+                text = "${SUB_INFO_2.displayName}",
+                summary = NUMBER_2
+            ),
+        )
+
+        val primarySimInfo = PrimarySimRepository(context).getPrimarySimInfo(simList)
+
+        assertThat(primarySimInfo).isNotNull()
+        assertThat(primarySimInfo?.dataList).isEqualTo(expectedList)
+    }
+
+    private companion object {
+        const val SUB_ID_1 = 1
+        const val SUB_ID_2 = 2
+        const val DISPLAY_NAME_1 = "Sub 1"
+        const val DISPLAY_NAME_2 = "Sub 2"
+        const val NUMBER_1 = "000000001"
+        const val NUMBER_2 = "000000002"
+        const val MCC = "310"
+
+        val SUB_INFO_1: SubscriptionInfo = SubscriptionInfo.Builder().apply {
+            setId(SUB_ID_1)
+            setDisplayName(DISPLAY_NAME_1)
+            setMcc(MCC)
+        }.build()
+
+        val SUB_INFO_2: SubscriptionInfo = SubscriptionInfo.Builder().apply {
+            setId(SUB_ID_2)
+            setDisplayName(DISPLAY_NAME_2)
+            setMcc(MCC)
+        }.build()
+    }
+}
diff --git a/tests/unit/Android.bp b/tests/unit/Android.bp
index 55df480..be43f8e 100644
--- a/tests/unit/Android.bp
+++ b/tests/unit/Android.bp
@@ -32,7 +32,6 @@
         "truth",
         "kotlinx_coroutines_test",
         "Settings-testutils2",
-        "MediaDrmSettingsFlagsLib",
         "servicestests-utils",
         // Don't add SettingsLib libraries here - you can use them directly as they are in the
         // instrumented Settings app.
diff --git a/tests/unit/src/com/android/settings/development/mediadrm/ForceSwSecureCryptoFallbackPreferenceControllerTest.java b/tests/unit/src/com/android/settings/development/mediadrm/ForceSwSecureCryptoFallbackPreferenceControllerTest.java
index 2a0cd05..5eb76f1 100644
--- a/tests/unit/src/com/android/settings/development/mediadrm/ForceSwSecureCryptoFallbackPreferenceControllerTest.java
+++ b/tests/unit/src/com/android/settings/development/mediadrm/ForceSwSecureCryptoFallbackPreferenceControllerTest.java
@@ -29,6 +29,8 @@
 import android.content.Context;
 
 import com.android.settings.media_drm.Flags;
+import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
 import android.platform.test.flag.junit.SetFlagsRule;
 
 import androidx.preference.SwitchPreference;
@@ -65,10 +67,10 @@
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_FORCE_L3_ENABLED)
     public void updateState_flagEnabled_checkPreference() {
-        mSetFlagsRule.enableFlags(Flags.FLAG_FORCE_L3_ENABLED);
         mController.updateState(mPreference);
-        assumeTrue(mPreference.isEnabled());
+        assertThat(mPreference.isEnabled()).isTrue();
         assertThat(mPreference.isChecked()).isFalse();
         assertThat(WidevineProperties.forcel3_enabled().orElse(false)).isFalse();
 
@@ -85,33 +87,22 @@
         assertThat(WidevineProperties.forcel3_enabled().orElse(false)).isFalse();
         assertThat(mPreference.isEnabled()).isTrue();
         assertThat(mPreference.isChecked()).isFalse();
-
-        // Test flag rollback
-        mController.setChecked(true);
-        mController.updateState(mPreference);
-        assertThat(mPreference.isChecked()).isTrue();
-        assertThat(WidevineProperties.forcel3_enabled().orElse(false)).isTrue();
-        mSetFlagsRule.disableFlags(Flags.FLAG_FORCE_L3_ENABLED);
-        mController.updateState(mPreference);
-        assertThat(mPreference.isEnabled()).isFalse();
-        assertThat(mPreference.isChecked()).isFalse();
-        assertThat(WidevineProperties.forcel3_enabled().orElse(false)).isFalse();
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_FORCE_L3_ENABLED)
     public void updateState_flagDisabled_checkPreference() {
-        mSetFlagsRule.disableFlags(Flags.FLAG_FORCE_L3_ENABLED);
         mController.updateState(mPreference);
         assertThat(mPreference.isEnabled()).isFalse();
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_FORCE_L3_ENABLED)
     public void updateState_checkWidevine() throws Exception {
         try (MediaDrm drm = new MediaDrm(WIDEVINE_UUID)) {
             assumeTrue(drm.getPropertyString("securityLevel").equals("L1"));
-            mSetFlagsRule.enableFlags(Flags.FLAG_FORCE_L3_ENABLED);
             mController.updateState(mPreference);
-            assumeTrue(mPreference.isEnabled());
+            assertThat(mPreference.isEnabled()).isTrue();
         } catch (UnsupportedSchemeException ex) {
             assumeNoException(ex);
         }
@@ -139,11 +130,11 @@
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_FORCE_L3_ENABLED)
     public void updateState_checkWhenWidevineReady() throws Exception {
         try (MediaDrm drm = new MediaDrm(WIDEVINE_UUID)) {
             if (drm.getPropertyString("securityLevel").equals("L1")) {
                 String version = drm.getPropertyString(MediaDrm.PROPERTY_VERSION);
-                mSetFlagsRule.enableFlags(Flags.FLAG_FORCE_L3_ENABLED);
                 mController.updateState(mPreference);
                 if (Integer.parseInt(version.split("\\.", 2)[0]) >= 19) {
                     assertThat(mPreference.isEnabled()).isTrue();
diff --git a/tests/unit/src/com/android/settings/network/telephony/MobileDataPreferenceControllerTest.java b/tests/unit/src/com/android/settings/network/telephony/MobileDataPreferenceControllerTest.java
index 152091a..f31e274 100644
--- a/tests/unit/src/com/android/settings/network/telephony/MobileDataPreferenceControllerTest.java
+++ b/tests/unit/src/com/android/settings/network/telephony/MobileDataPreferenceControllerTest.java
@@ -131,8 +131,7 @@
 
     private MobileNetworkInfoEntity setupMobileNetworkInfoEntity(String subId,
             boolean isDatEnabled) {
-        return new MobileNetworkInfoEntity(subId, false, false, isDatEnabled, false, false, false,
-                false, false, false, false, false);
+        return new MobileNetworkInfoEntity(subId, isDatEnabled, false);
     }
 
     @Test