Merge "Fix another race condition for provider ui launching."
diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index 2e2ee23..eb2d2ca 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -3097,7 +3097,7 @@
   public class VirtualSensor {
     method @NonNull public String getName();
     method public int getType();
-    method @RequiresPermission(android.Manifest.permission.CREATE_VIRTUAL_DEVICE) public void sendSensorEvent(@NonNull android.companion.virtual.sensor.VirtualSensorEvent);
+    method @RequiresPermission(android.Manifest.permission.CREATE_VIRTUAL_DEVICE) public void sendEvent(@NonNull android.companion.virtual.sensor.VirtualSensorEvent);
   }
 
   public static interface VirtualSensor.SensorStateChangeCallback {
diff --git a/core/java/android/app/ResourcesManager.java b/core/java/android/app/ResourcesManager.java
index d275c83..32217610 100644
--- a/core/java/android/app/ResourcesManager.java
+++ b/core/java/android/app/ResourcesManager.java
@@ -305,9 +305,13 @@
             for (int i = mCachedApkAssets.size() - 1; i >= 0; i--) {
                 final ApkKey key = mCachedApkAssets.keyAt(i);
                 if (key.path.equals(path)) {
-                    WeakReference<ApkAssets> apkAssetsRef = mCachedApkAssets.removeAt(i);
-                    if (apkAssetsRef != null && apkAssetsRef.get() != null) {
-                        apkAssetsRef.get().close();
+                    final WeakReference<ApkAssets> apkAssetsRef = mCachedApkAssets.removeAt(i);
+                    if (apkAssetsRef == null) {
+                        continue;
+                    }
+                    final ApkAssets apkAssets = apkAssetsRef.get();
+                    if (apkAssets != null) {
+                        apkAssets.close();
                     }
                 }
             }
@@ -446,16 +450,14 @@
         ApkAssets apkAssets;
 
         // Optimistically check if this ApkAssets exists somewhere else.
+        final WeakReference<ApkAssets> apkAssetsRef;
         synchronized (mLock) {
-            final WeakReference<ApkAssets> apkAssetsRef = mCachedApkAssets.get(key);
-            if (apkAssetsRef != null) {
-                apkAssets = apkAssetsRef.get();
-                if (apkAssets != null && apkAssets.isUpToDate()) {
-                    return apkAssets;
-                } else {
-                    // Clean up the reference.
-                    mCachedApkAssets.remove(key);
-                }
+            apkAssetsRef = mCachedApkAssets.get(key);
+        }
+        if (apkAssetsRef != null) {
+            apkAssets = apkAssetsRef.get();
+            if (apkAssets != null && apkAssets.isUpToDate()) {
+                return apkAssets;
             }
         }
 
diff --git a/core/java/android/app/SystemServiceRegistry.java b/core/java/android/app/SystemServiceRegistry.java
index 762ac23..3730a6f 100644
--- a/core/java/android/app/SystemServiceRegistry.java
+++ b/core/java/android/app/SystemServiceRegistry.java
@@ -51,6 +51,8 @@
 import android.app.usage.UsageStatsManager;
 import android.app.wallpapereffectsgeneration.IWallpaperEffectsGenerationManager;
 import android.app.wallpapereffectsgeneration.WallpaperEffectsGenerationManager;
+import android.app.wearable.IWearableSensingManager;
+import android.app.wearable.WearableSensingManager;
 import android.apphibernation.AppHibernationManager;
 import android.appwidget.AppWidgetManager;
 import android.bluetooth.BluetoothFrameworkInitializer;
@@ -1544,6 +1546,17 @@
                                 IAmbientContextManager.Stub.asInterface(iBinder);
                         return new AmbientContextManager(ctx.getOuterContext(), manager);
                     }});
+        registerService(Context.WEARABLE_SENSING_SERVICE, WearableSensingManager.class,
+                new CachedServiceFetcher<WearableSensingManager>() {
+                    @Override
+                    public WearableSensingManager createService(ContextImpl ctx)
+                            throws ServiceNotFoundException {
+                        IBinder iBinder = ServiceManager.getServiceOrThrow(
+                                Context.WEARABLE_SENSING_SERVICE);
+                        IWearableSensingManager manager =
+                                IWearableSensingManager.Stub.asInterface(iBinder);
+                        return new WearableSensingManager(ctx.getOuterContext(), manager);
+                    }});
 
         sInitializing = true;
         try {
diff --git a/core/java/android/companion/virtual/sensor/VirtualSensor.java b/core/java/android/companion/virtual/sensor/VirtualSensor.java
index a184481..58a5387 100644
--- a/core/java/android/companion/virtual/sensor/VirtualSensor.java
+++ b/core/java/android/companion/virtual/sensor/VirtualSensor.java
@@ -20,6 +20,7 @@
 import android.annotation.RequiresPermission;
 import android.annotation.SystemApi;
 import android.companion.virtual.IVirtualDevice;
+import android.hardware.Sensor;
 import android.os.IBinder;
 import android.os.RemoteException;
 
@@ -68,8 +69,10 @@
     }
 
     /**
-     * Returns the
-     * <a href="https://source.android.com/devices/sensors/sensor-types">type</a> of the sensor.
+     * Returns the type of the sensor.
+     *
+     * @see Sensor#getType()
+     * @see <a href="https://source.android.com/devices/sensors/sensor-types">Sensor types</a>
      */
     public int getType() {
         return mType;
@@ -87,7 +90,7 @@
      * Send a sensor event to the system.
      */
     @RequiresPermission(android.Manifest.permission.CREATE_VIRTUAL_DEVICE)
-    public void sendSensorEvent(@NonNull VirtualSensorEvent event) {
+    public void sendEvent(@NonNull VirtualSensorEvent event) {
         try {
             mVirtualDevice.sendSensorEvent(mToken, event);
         } catch (RemoteException e) {
diff --git a/core/java/android/companion/virtual/sensor/VirtualSensorConfig.java b/core/java/android/companion/virtual/sensor/VirtualSensorConfig.java
index 7982fa5..eb2f9dd 100644
--- a/core/java/android/companion/virtual/sensor/VirtualSensorConfig.java
+++ b/core/java/android/companion/virtual/sensor/VirtualSensorConfig.java
@@ -23,6 +23,7 @@
 import android.annotation.Nullable;
 import android.annotation.SuppressLint;
 import android.annotation.SystemApi;
+import android.hardware.Sensor;
 import android.os.Parcel;
 import android.os.Parcelable;
 
@@ -77,8 +78,10 @@
     }
 
     /**
-     * Returns the
-     * <a href="https://source.android.com/devices/sensors/sensor-types">type</a> of the sensor.
+     * Returns the type of the sensor.
+     *
+     * @see Sensor#getType()
+     * @see <a href="https://source.android.com/devices/sensors/sensor-types">Sensor types</a>
      */
     public int getType() {
         return mType;
@@ -150,8 +153,7 @@
         /**
          * Creates a new builder.
          *
-         * @param type The
-         * <a href="https://source.android.com/devices/sensors/sensor-types">type</a> of the sensor.
+         * @param type The type of the sensor, matching {@link Sensor#getType}.
          * @param name The name of the sensor. Must be unique among all sensors with the same type
          * that belong to the same virtual device.
          */
diff --git a/core/java/android/companion/virtual/sensor/VirtualSensorEvent.java b/core/java/android/companion/virtual/sensor/VirtualSensorEvent.java
index 8f8860e..01b4975 100644
--- a/core/java/android/companion/virtual/sensor/VirtualSensorEvent.java
+++ b/core/java/android/companion/virtual/sensor/VirtualSensorEvent.java
@@ -18,6 +18,7 @@
 
 import android.annotation.NonNull;
 import android.annotation.SystemApi;
+import android.hardware.Sensor;
 import android.os.Parcel;
 import android.os.Parcelable;
 import android.os.SystemClock;
@@ -60,9 +61,11 @@
     }
 
     /**
-     * Returns the values of this sensor event. The length and contents depend on the
-     * <a href="https://source.android.com/devices/sensors/sensor-types">sensor type</a>.
+     * Returns the values of this sensor event. The length and contents depend on the sensor type.
+     *
      * @see android.hardware.SensorEvent#values
+     * @see Sensor#getType()
+     * @see <a href="https://source.android.com/devices/sensors/sensor-types">Sensor types</a>
      */
     @NonNull
     public float[] getValues() {
@@ -90,7 +93,9 @@
 
         /**
          * Creates a new builder.
-         * @param values the values of the sensor event. @see android.hardware.SensorEvent#values
+         *
+         * @param values the values of the sensor event.
+         * @see android.hardware.SensorEvent#values
          */
         public Builder(@NonNull float[] values) {
             mValues = values;
diff --git a/core/java/android/hardware/input/VirtualDpad.java b/core/java/android/hardware/input/VirtualDpad.java
index 4d61553..8133472 100644
--- a/core/java/android/hardware/input/VirtualDpad.java
+++ b/core/java/android/hardware/input/VirtualDpad.java
@@ -32,8 +32,8 @@
 /**
  * A virtual dpad representing a key input mechanism on a remote device.
  *
- * This registers an InputDevice that is interpreted like a physically-connected device and
- * dispatches received events to it.
+ * <p>This registers an InputDevice that is interpreted like a physically-connected device and
+ * dispatches received events to it.</p>
  *
  * @hide
  */
@@ -44,6 +44,7 @@
             Collections.unmodifiableSet(
                     new HashSet<>(
                             Arrays.asList(
+                                    KeyEvent.KEYCODE_BACK,
                                     KeyEvent.KEYCODE_DPAD_UP,
                                     KeyEvent.KEYCODE_DPAD_DOWN,
                                     KeyEvent.KEYCODE_DPAD_LEFT,
@@ -58,8 +59,15 @@
     /**
      * Sends a key event to the system.
      *
-     * Supported key codes are KEYCODE_DPAD_UP, KEYCODE_DPAD_DOWN, KEYCODE_DPAD_LEFT,
-     * KEYCODE_DPAD_RIGHT and KEYCODE_DPAD_CENTER,
+     * <p>Supported key codes are:
+     * <ul>
+     *     <li>{@link KeyEvent.KEYCODE_DPAD_UP}</li>
+     *     <li>{@link KeyEvent.KEYCODE_DPAD_DOWN}</li>
+     *     <li>{@link KeyEvent.KEYCODE_DPAD_LEFT}</li>
+     *     <li>{@link KeyEvent.KEYCODE_DPAD_RIGHT}</li>
+     *     <li>{@link KeyEvent.KEYCODE_DPAD_CENTER}</li>
+     *     <li>{@link KeyEvent.KEYCODE_BACK}</li>
+     * </ul>
      *
      * @param event the event to send
      */
diff --git a/core/java/android/service/controls/ControlsProviderService.java b/core/java/android/service/controls/ControlsProviderService.java
index d2a4ae2..9396a88 100644
--- a/core/java/android/service/controls/ControlsProviderService.java
+++ b/core/java/android/service/controls/ControlsProviderService.java
@@ -69,6 +69,18 @@
             "android.service.controls.META_DATA_PANEL_ACTIVITY";
 
     /**
+     * Boolean extra containing the value of
+     * {@link android.provider.Settings.Secure#LOCKSCREEN_ALLOW_TRIVIAL_CONTROLS}.
+     *
+     * This is passed with the intent when the panel specified by {@link #META_DATA_PANEL_ACTIVITY}
+     * is launched.
+     *
+     * @hide
+     */
+    public static final String EXTRA_LOCKSCREEN_ALLOW_TRIVIAL_CONTROLS =
+            "android.service.controls.EXTRA_LOCKSCREEN_ALLOW_TRIVIAL_CONTROLS";
+
+    /**
      * @hide
      */
     @SdkConstant(SdkConstantType.BROADCAST_INTENT_ACTION)
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
index 400039b..9329d02 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
@@ -529,10 +529,24 @@
             @SplitPosition int splitPosition, float splitRatio, RemoteAnimationAdapter adapter,
             InstanceId instanceId) {
         Intent fillInIntent = null;
-        if (launchSameComponentAdjacently(pendingIntent, splitPosition, taskId)
-                && supportMultiInstancesSplit(pendingIntent.getIntent().getComponent())) {
-            fillInIntent = new Intent();
-            fillInIntent.addFlags(FLAG_ACTIVITY_MULTIPLE_TASK);
+        if (launchSameComponentAdjacently(pendingIntent, splitPosition, taskId)) {
+            if (supportMultiInstancesSplit(pendingIntent.getIntent().getComponent())) {
+                fillInIntent = new Intent();
+                fillInIntent.addFlags(FLAG_ACTIVITY_MULTIPLE_TASK);
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN, "Adding MULTIPLE_TASK");
+            } else {
+                try {
+                    adapter.getRunner().onAnimationCancelled(false /* isKeyguardOccluded */);
+                    ActivityTaskManager.getService().startActivityFromRecents(taskId, options2);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error starting remote animation", e);
+                }
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN,
+                        "Cancel entering split as not supporting multi-instances");
+                Toast.makeText(mContext, R.string.dock_multi_instances_not_supported_text,
+                        Toast.LENGTH_SHORT).show();
+                return;
+            }
         }
         mStageCoordinator.startIntentAndTaskWithLegacyTransition(pendingIntent, fillInIntent,
                 options1, taskId, options2, splitPosition, splitRatio, adapter, instanceId);
@@ -542,10 +556,17 @@
             int taskId, @Nullable Bundle options2, @SplitPosition int splitPosition,
             float splitRatio, @Nullable RemoteTransition remoteTransition, InstanceId instanceId) {
         Intent fillInIntent = null;
-        if (launchSameComponentAdjacently(pendingIntent, splitPosition, taskId)
-                && supportMultiInstancesSplit(pendingIntent.getIntent().getComponent())) {
-            fillInIntent = new Intent();
-            fillInIntent.addFlags(FLAG_ACTIVITY_MULTIPLE_TASK);
+        if (launchSameComponentAdjacently(pendingIntent, splitPosition, taskId)) {
+            if (supportMultiInstancesSplit(pendingIntent.getIntent().getComponent())) {
+                fillInIntent = new Intent();
+                fillInIntent.addFlags(FLAG_ACTIVITY_MULTIPLE_TASK);
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN, "Adding MULTIPLE_TASK");
+            } else {
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN,
+                        "Cancel entering split as not supporting multi-instances");
+                Toast.makeText(mContext, R.string.dock_multi_instances_not_supported_text,
+                        Toast.LENGTH_SHORT).show();
+            }
         }
         mStageCoordinator.startIntentAndTask(pendingIntent, fillInIntent, options1, taskId,
                 options2, splitPosition, splitRatio, remoteTransition, instanceId);
@@ -557,12 +578,26 @@
             float splitRatio, RemoteAnimationAdapter adapter, InstanceId instanceId) {
         Intent fillInIntent1 = null;
         Intent fillInIntent2 = null;
-        if (launchSameComponentAdjacently(pendingIntent1, pendingIntent2)
-                && supportMultiInstancesSplit(pendingIntent1.getIntent().getComponent())) {
-            fillInIntent1 = new Intent();
-            fillInIntent1.addFlags(FLAG_ACTIVITY_MULTIPLE_TASK);
-            fillInIntent2 = new Intent();
-            fillInIntent2.addFlags(FLAG_ACTIVITY_MULTIPLE_TASK);
+        if (launchSameComponentAdjacently(pendingIntent1, pendingIntent2)) {
+            if (supportMultiInstancesSplit(pendingIntent1.getIntent().getComponent())) {
+                fillInIntent1 = new Intent();
+                fillInIntent1.addFlags(FLAG_ACTIVITY_MULTIPLE_TASK);
+                fillInIntent2 = new Intent();
+                fillInIntent2.addFlags(FLAG_ACTIVITY_MULTIPLE_TASK);
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN, "Adding MULTIPLE_TASK");
+            } else {
+                try {
+                    adapter.getRunner().onAnimationCancelled(false /* isKeyguardOccluded */);
+                    pendingIntent1.send();
+                } catch (RemoteException | PendingIntent.CanceledException e) {
+                    Slog.e(TAG, "Error starting remote animation", e);
+                }
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN,
+                        "Cancel entering split as not supporting multi-instances");
+                Toast.makeText(mContext, R.string.dock_multi_instances_not_supported_text,
+                        Toast.LENGTH_SHORT).show();
+                return;
+            }
         }
         mStageCoordinator.startIntentsWithLegacyTransition(pendingIntent1, fillInIntent1, options1,
                 pendingIntent2, fillInIntent2, options2, splitPosition, splitRatio, adapter,
@@ -601,6 +636,8 @@
                 mStageCoordinator.switchSplitPosition("startIntent");
                 return;
             } else {
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN,
+                        "Cancel entering split as not supporting multi-instances");
                 Toast.makeText(mContext, R.string.dock_multi_instances_not_supported_text,
                         Toast.LENGTH_SHORT).show();
                 return;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
index 38fb2df..da8dc87 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
@@ -456,8 +456,6 @@
     void startIntentLegacy(PendingIntent intent, Intent fillInIntent, @SplitPosition int position,
             @Nullable Bundle options) {
         final boolean isEnteringSplit = !isSplitActive();
-        final WindowContainerTransaction evictWct = new WindowContainerTransaction();
-        prepareEvictChildTasks(position, evictWct);
 
         LegacyTransitions.ILegacyTransition transition = new LegacyTransitions.ILegacyTransition() {
             @Override
@@ -465,22 +463,21 @@
                     RemoteAnimationTarget[] wallpapers, RemoteAnimationTarget[] nonApps,
                     IRemoteAnimationFinishedCallback finishedCallback,
                     SurfaceControl.Transaction t) {
-                if (isEnteringSplit) {
-                    boolean openingToSide = false;
-                    if (apps != null) {
-                        for (int i = 0; i < apps.length; ++i) {
-                            if (apps[i].mode == MODE_OPENING
-                                    && mSideStage.containsTask(apps[i].taskId)) {
-                                openingToSide = true;
-                                break;
-                            }
+                boolean openingToSide = false;
+                if (apps != null) {
+                    for (int i = 0; i < apps.length; ++i) {
+                        if (apps[i].mode == MODE_OPENING
+                                && mSideStage.containsTask(apps[i].taskId)) {
+                            openingToSide = true;
+                            break;
                         }
                     }
-                    if (!openingToSide) {
-                        mMainExecutor.execute(() -> exitSplitScreen(
-                                mSideStage.getChildCount() == 0 ? mMainStage : mSideStage,
-                                EXIT_REASON_UNKNOWN));
-                    }
+                }
+
+                if (isEnteringSplit && !openingToSide) {
+                    mMainExecutor.execute(() -> exitSplitScreen(
+                            mSideStage.getChildCount() == 0 ? mMainStage : mSideStage,
+                            EXIT_REASON_UNKNOWN));
                 }
 
                 if (apps != null) {
@@ -500,7 +497,12 @@
                     }
                 }
 
-                mSyncQueue.queue(evictWct);
+
+                if (!isEnteringSplit && openingToSide) {
+                    final WindowContainerTransaction evictWct = new WindowContainerTransaction();
+                    prepareEvictNonOpeningChildTasks(position, apps, evictWct);
+                    mSyncQueue.queue(evictWct);
+                }
             }
         };
 
diff --git a/media/java/android/media/RingtoneManager.java b/media/java/android/media/RingtoneManager.java
index f15f443..1e270b1 100644
--- a/media/java/android/media/RingtoneManager.java
+++ b/media/java/android/media/RingtoneManager.java
@@ -826,6 +826,19 @@
         if(!isInternalRingtoneUri(ringtoneUri)) {
             ringtoneUri = ContentProvider.maybeAddUserId(ringtoneUri, context.getUserId());
         }
+
+        final String mimeType = resolver.getType(ringtoneUri);
+        if (mimeType == null) {
+            Log.e(TAG, "setActualDefaultRingtoneUri for URI:" + ringtoneUri
+                    + " ignored: failure to find mimeType (no access from this context?)");
+            return;
+        }
+        if (!(mimeType.startsWith("audio/") || mimeType.equals("application/ogg"))) {
+            Log.e(TAG, "setActualDefaultRingtoneUri for URI:" + ringtoneUri
+                    + " ignored: associated mimeType:" + mimeType + " is not an audio type");
+            return;
+        }
+
         Settings.System.putStringForUser(resolver, setting,
                 ringtoneUri != null ? ringtoneUri.toString() : null, context.getUserId());
 
diff --git a/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/search/SpaSearchContract.kt b/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/search/SpaSearchContract.kt
new file mode 100644
index 0000000..2301f04
--- /dev/null
+++ b/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/search/SpaSearchContract.kt
@@ -0,0 +1,97 @@
+/*
+ * Copyright (C) 2022 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.settingslib.spa.search
+
+/**
+ * Intent action used to identify SpaSearchProvider instances. This is used in the {@code
+ * <intent-filter>} of a {@code <provider>}.
+ */
+const val PROVIDER_INTERFACE = "android.content.action.SPA_SEARCH_PROVIDER"
+
+/** ContentProvider path for search static data */
+const val SEARCH_STATIC_DATA = "search_static_data"
+
+/** ContentProvider path for search dynamic data */
+const val SEARCH_DYNAMIC_DATA = "search_dynamic_data"
+
+/** ContentProvider path for search immutable status */
+const val SEARCH_IMMUTABLE_STATUS = "search_immutable_status"
+
+/** ContentProvider path for search mutable status */
+const val SEARCH_MUTABLE_STATUS = "search_mutable_status"
+
+/** Enum to define all column names in provider. */
+enum class ColumnEnum(val id: String) {
+    ENTRY_ID("entryId"),
+    SEARCH_TITLE("searchTitle"),
+    SEARCH_KEYWORD("searchKw"),
+    SEARCH_PATH("searchPath"),
+    INTENT_TARGET_PACKAGE("intentTargetPackage"),
+    INTENT_TARGET_CLASS("intentTargetClass"),
+    INTENT_EXTRAS("intentExtras"),
+    SLICE_URI("sliceUri"),
+    LEGACY_KEY("legacyKey"),
+    ENTRY_DISABLED("entryDisabled"),
+}
+
+/** Enum to define all queries supported in the provider. */
+@SuppressWarnings("Immutable")
+enum class QueryEnum(
+    val queryPath: String,
+    val columnNames: List<ColumnEnum>
+) {
+    SEARCH_STATIC_DATA_QUERY(
+        SEARCH_STATIC_DATA,
+        listOf(
+            ColumnEnum.ENTRY_ID,
+            ColumnEnum.SEARCH_TITLE,
+            ColumnEnum.SEARCH_KEYWORD,
+            ColumnEnum.SEARCH_PATH,
+            ColumnEnum.INTENT_TARGET_PACKAGE,
+            ColumnEnum.INTENT_TARGET_CLASS,
+            ColumnEnum.INTENT_EXTRAS,
+            ColumnEnum.SLICE_URI,
+            ColumnEnum.LEGACY_KEY
+        )
+    ),
+    SEARCH_DYNAMIC_DATA_QUERY(
+        SEARCH_DYNAMIC_DATA,
+        listOf(
+            ColumnEnum.ENTRY_ID,
+            ColumnEnum.SEARCH_TITLE,
+            ColumnEnum.SEARCH_KEYWORD,
+            ColumnEnum.SEARCH_PATH,
+            ColumnEnum.INTENT_TARGET_PACKAGE,
+            ColumnEnum.INTENT_TARGET_CLASS,
+            ColumnEnum.SLICE_URI,
+            ColumnEnum.LEGACY_KEY
+        )
+    ),
+    SEARCH_IMMUTABLE_STATUS_DATA_QUERY(
+        SEARCH_IMMUTABLE_STATUS,
+        listOf(
+            ColumnEnum.ENTRY_ID,
+            ColumnEnum.ENTRY_DISABLED,
+        )
+    ),
+    SEARCH_MUTABLE_STATUS_DATA_QUERY(
+        SEARCH_MUTABLE_STATUS,
+        listOf(
+            ColumnEnum.ENTRY_ID,
+            ColumnEnum.ENTRY_DISABLED,
+        )
+    ),
+}
diff --git a/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/search/SpaSearchProvider.kt b/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/search/SpaSearchProvider.kt
index 02aed1c..21bc75a 100644
--- a/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/search/SpaSearchProvider.kt
+++ b/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/search/SpaSearchProvider.kt
@@ -19,22 +19,22 @@
 import android.content.ContentProvider
 import android.content.ContentValues
 import android.content.Context
-import android.content.Intent
 import android.content.UriMatcher
 import android.content.pm.ProviderInfo
 import android.database.Cursor
 import android.database.MatrixCursor
 import android.net.Uri
+import android.os.Parcel
+import android.os.Parcelable
 import android.util.Log
 import androidx.annotation.VisibleForTesting
-import com.android.settingslib.spa.framework.common.ColumnEnum
-import com.android.settingslib.spa.framework.common.QueryEnum
+import com.android.settingslib.spa.framework.common.EntryStatusData
 import com.android.settingslib.spa.framework.common.SettingsEntry
 import com.android.settingslib.spa.framework.common.SpaEnvironmentFactory
-import com.android.settingslib.spa.framework.common.addUri
-import com.android.settingslib.spa.framework.common.getColumns
 import com.android.settingslib.spa.framework.util.SESSION_SEARCH
 import com.android.settingslib.spa.framework.util.createIntent
+import com.android.settingslib.spa.slice.fromEntry
+
 
 private const val TAG = "SpaSearchProvider"
 
@@ -42,18 +42,25 @@
  * The content provider to return entry related data, which can be used for search and hierarchy.
  * One can query the provider result by:
  *   $ adb shell content query --uri content://<AuthorityPath>/<QueryPath>
- * For gallery, AuthorityPath = com.android.spa.gallery.provider
- * For Settings, AuthorityPath = com.android.settings.spa.provider
+ * For gallery, AuthorityPath = com.android.spa.gallery.search.provider
+ * For Settings, AuthorityPath = com.android.settings.spa.search.provider"
  * Some examples:
- *   $ adb shell content query --uri content://<AuthorityPath>/search_static
- *   $ adb shell content query --uri content://<AuthorityPath>/search_dynamic
- *   $ adb shell content query --uri content://<AuthorityPath>/search_mutable_status
+ *   $ adb shell content query --uri content://<AuthorityPath>/search_static_data
+ *   $ adb shell content query --uri content://<AuthorityPath>/search_dynamic_data
  *   $ adb shell content query --uri content://<AuthorityPath>/search_immutable_status
+ *   $ adb shell content query --uri content://<AuthorityPath>/search_mutable_status
  */
 class SpaSearchProvider : ContentProvider() {
     private val spaEnvironment get() = SpaEnvironmentFactory.instance
     private val uriMatcher = UriMatcher(UriMatcher.NO_MATCH)
 
+    private val queryMatchCode = mapOf(
+        SEARCH_STATIC_DATA to 301,
+        SEARCH_DYNAMIC_DATA to 302,
+        SEARCH_MUTABLE_STATUS to 303,
+        SEARCH_IMMUTABLE_STATUS to 304
+    )
+
     override fun delete(uri: Uri, selection: String?, selectionArgs: Array<String>?): Int {
         TODO("Implement this to handle requests to delete one or more rows")
     }
@@ -85,10 +92,9 @@
 
     override fun attachInfo(context: Context?, info: ProviderInfo?) {
         if (info != null) {
-            QueryEnum.SEARCH_STATIC_DATA_QUERY.addUri(uriMatcher, info.authority)
-            QueryEnum.SEARCH_DYNAMIC_DATA_QUERY.addUri(uriMatcher, info.authority)
-            QueryEnum.SEARCH_MUTABLE_STATUS_DATA_QUERY.addUri(uriMatcher, info.authority)
-            QueryEnum.SEARCH_IMMUTABLE_STATUS_DATA_QUERY.addUri(uriMatcher, info.authority)
+            for (entry in queryMatchCode) {
+                uriMatcher.addURI(info.authority, entry.key, entry.value)
+            }
         }
         super.attachInfo(context, info)
     }
@@ -102,11 +108,11 @@
     ): Cursor? {
         return try {
             when (uriMatcher.match(uri)) {
-                QueryEnum.SEARCH_STATIC_DATA_QUERY.queryMatchCode -> querySearchStaticData()
-                QueryEnum.SEARCH_DYNAMIC_DATA_QUERY.queryMatchCode -> querySearchDynamicData()
-                QueryEnum.SEARCH_MUTABLE_STATUS_DATA_QUERY.queryMatchCode ->
+                queryMatchCode[SEARCH_STATIC_DATA] -> querySearchStaticData()
+                queryMatchCode[SEARCH_DYNAMIC_DATA] -> querySearchDynamicData()
+                queryMatchCode[SEARCH_MUTABLE_STATUS] ->
                     querySearchMutableStatusData()
-                QueryEnum.SEARCH_IMMUTABLE_STATUS_DATA_QUERY.queryMatchCode ->
+                queryMatchCode[SEARCH_IMMUTABLE_STATUS] ->
                     querySearchImmutableStatusData()
                 else -> throw UnsupportedOperationException("Unknown Uri $uri")
             }
@@ -167,23 +173,45 @@
 
         // Fetch search data. We can add runtime arguments later if necessary
         val searchData = entry.getSearchData() ?: return
-        val intent = entry.createIntent(SESSION_SEARCH) ?: Intent()
-        cursor.newRow()
-            .add(ColumnEnum.ENTRY_ID.id, entry.id)
-            .add(ColumnEnum.ENTRY_INTENT_URI.id, intent.toUri(Intent.URI_INTENT_SCHEME))
+        val intent = entry.createIntent(SESSION_SEARCH)
+        val row = cursor.newRow().add(ColumnEnum.ENTRY_ID.id, entry.id)
             .add(ColumnEnum.SEARCH_TITLE.id, searchData.title)
             .add(ColumnEnum.SEARCH_KEYWORD.id, searchData.keyword)
             .add(
                 ColumnEnum.SEARCH_PATH.id,
                 entryRepository.getEntryPathWithTitle(entry.id, searchData.title)
             )
+        intent?.let {
+            row.add(ColumnEnum.INTENT_TARGET_PACKAGE.id, spaEnvironment.appContext.packageName)
+                .add(ColumnEnum.INTENT_TARGET_CLASS.id, spaEnvironment.browseActivityClass?.name)
+                .add(ColumnEnum.INTENT_EXTRAS.id, marshall(intent.extras))
+        }
+        if (entry.hasSliceSupport)
+            row.add(
+                ColumnEnum.SLICE_URI.id, Uri.Builder()
+                    .fromEntry(entry, spaEnvironment.sliceProviderAuthorities)
+            )
+        // TODO: support legacy key
     }
 
     private fun fetchStatusData(entry: SettingsEntry, cursor: MatrixCursor) {
         // Fetch status data. We can add runtime arguments later if necessary
-        val statusData = entry.getStatusData() ?: return
+        val statusData = entry.getStatusData() ?: EntryStatusData()
         cursor.newRow()
             .add(ColumnEnum.ENTRY_ID.id, entry.id)
-            .add(ColumnEnum.SEARCH_STATUS_DISABLED.id, statusData.isDisabled)
+            .add(ColumnEnum.ENTRY_DISABLED.id, statusData.isDisabled)
+    }
+
+    private fun QueryEnum.getColumns(): Array<String> {
+        return columnNames.map { it.id }.toTypedArray()
+    }
+
+    private fun marshall(parcelable: Parcelable?): ByteArray? {
+        if (parcelable == null) return null
+        val parcel = Parcel.obtain()
+        parcelable.writeToParcel(parcel, 0)
+        val bytes = parcel.marshall()
+        parcel.recycle()
+        return bytes
     }
 }
diff --git a/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/widget/scaffold/CustomizedAppBar.kt b/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/widget/scaffold/CustomizedAppBar.kt
index 7db1ca1..ae88ed7 100644
--- a/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/widget/scaffold/CustomizedAppBar.kt
+++ b/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/widget/scaffold/CustomizedAppBar.kt
@@ -132,7 +132,10 @@
                     .asPaddingValues()
                     .horizontalValues()
             )
-            .padding(horizontal = SettingsDimension.itemPaddingAround),
+            .padding(
+                start = SettingsDimension.itemPaddingAround,
+                end = SettingsDimension.itemPaddingEnd,
+            ),
         overflow = TextOverflow.Ellipsis,
         maxLines = maxLines,
     )
diff --git a/packages/SettingsLib/Spa/tests/src/com/android/settingslib/spa/search/SpaSearchProviderTest.kt b/packages/SettingsLib/Spa/tests/src/com/android/settingslib/spa/search/SpaSearchProviderTest.kt
index cdb0f3a..831aded 100644
--- a/packages/SettingsLib/Spa/tests/src/com/android/settingslib/spa/search/SpaSearchProviderTest.kt
+++ b/packages/SettingsLib/Spa/tests/src/com/android/settingslib/spa/search/SpaSearchProviderTest.kt
@@ -18,15 +18,14 @@
 
 import android.content.Context
 import android.database.Cursor
+import android.os.Bundle
+import android.os.Parcel
 import androidx.test.core.app.ApplicationProvider
 import androidx.test.ext.junit.runners.AndroidJUnit4
-import com.android.settingslib.spa.framework.common.ColumnEnum
-import com.android.settingslib.spa.framework.common.QueryEnum
 import com.android.settingslib.spa.framework.common.SettingsEntryBuilder
 import com.android.settingslib.spa.framework.common.SettingsPage
 import com.android.settingslib.spa.framework.common.SpaEnvironmentFactory
 import com.android.settingslib.spa.framework.common.createSettingsPage
-import com.android.settingslib.spa.framework.common.getIndex
 import com.android.settingslib.spa.tests.testutils.SpaEnvironmentForTest
 import com.android.settingslib.spa.tests.testutils.SppForSearch
 import com.google.common.truth.Truth
@@ -39,51 +38,145 @@
     private val spaEnvironment =
         SpaEnvironmentForTest(context, listOf(SppForSearch.createSettingsPage()))
     private val searchProvider = SpaSearchProvider()
+    private val pageOwner = spaEnvironment.createPage("SppForSearch")
 
     @Test
     fun testQuerySearchStatusData() {
         SpaEnvironmentFactory.reset(spaEnvironment)
-        val pageOwner = spaEnvironment.createPage("SppForSearch")
 
         val immutableStatus = searchProvider.querySearchImmutableStatusData()
-        Truth.assertThat(immutableStatus.count).isEqualTo(1)
+        Truth.assertThat(immutableStatus.count).isEqualTo(2)
         immutableStatus.moveToFirst()
         immutableStatus.checkValue(
             QueryEnum.SEARCH_IMMUTABLE_STATUS_DATA_QUERY,
             ColumnEnum.ENTRY_ID,
+            pageOwner.getEntryId("SearchStaticWithNoStatus")
+        )
+        immutableStatus.checkValue(
+            QueryEnum.SEARCH_IMMUTABLE_STATUS_DATA_QUERY,
+            ColumnEnum.ENTRY_DISABLED,
+            false.toString()
+        )
+
+        immutableStatus.moveToNext()
+        immutableStatus.checkValue(
+            QueryEnum.SEARCH_IMMUTABLE_STATUS_DATA_QUERY,
+            ColumnEnum.ENTRY_ID,
             pageOwner.getEntryId("SearchDynamicWithImmutableStatus")
         )
+        immutableStatus.checkValue(
+            QueryEnum.SEARCH_IMMUTABLE_STATUS_DATA_QUERY, ColumnEnum.ENTRY_DISABLED, true.toString()
+        )
 
         val mutableStatus = searchProvider.querySearchMutableStatusData()
         Truth.assertThat(mutableStatus.count).isEqualTo(2)
         mutableStatus.moveToFirst()
         mutableStatus.checkValue(
-            QueryEnum.SEARCH_IMMUTABLE_STATUS_DATA_QUERY,
+            QueryEnum.SEARCH_MUTABLE_STATUS_DATA_QUERY,
             ColumnEnum.ENTRY_ID,
             pageOwner.getEntryId("SearchStaticWithMutableStatus")
         )
+        mutableStatus.checkValue(
+            QueryEnum.SEARCH_MUTABLE_STATUS_DATA_QUERY, ColumnEnum.ENTRY_DISABLED, false.toString()
+        )
 
         mutableStatus.moveToNext()
         mutableStatus.checkValue(
-            QueryEnum.SEARCH_IMMUTABLE_STATUS_DATA_QUERY,
+            QueryEnum.SEARCH_MUTABLE_STATUS_DATA_QUERY,
             ColumnEnum.ENTRY_ID,
             pageOwner.getEntryId("SearchDynamicWithMutableStatus")
         )
+        mutableStatus.checkValue(
+            QueryEnum.SEARCH_MUTABLE_STATUS_DATA_QUERY, ColumnEnum.ENTRY_DISABLED, true.toString()
+        )
     }
 
     @Test
     fun testQuerySearchIndexData() {
         SpaEnvironmentFactory.reset(spaEnvironment)
+
         val staticData = searchProvider.querySearchStaticData()
         Truth.assertThat(staticData.count).isEqualTo(2)
+        staticData.moveToFirst()
+        staticData.checkValue(
+            QueryEnum.SEARCH_STATIC_DATA_QUERY,
+            ColumnEnum.ENTRY_ID,
+            pageOwner.getEntryId("SearchStaticWithNoStatus")
+        )
+        staticData.checkValue(
+            QueryEnum.SEARCH_STATIC_DATA_QUERY, ColumnEnum.SEARCH_TITLE, "SearchStaticWithNoStatus"
+        )
+        staticData.checkValue(
+            QueryEnum.SEARCH_STATIC_DATA_QUERY, ColumnEnum.SEARCH_KEYWORD, listOf("").toString()
+        )
+        staticData.checkValue(
+            QueryEnum.SEARCH_STATIC_DATA_QUERY,
+            ColumnEnum.SEARCH_PATH,
+            listOf("SearchStaticWithNoStatus", "SppForSearch").toString()
+        )
+        staticData.checkValue(
+            QueryEnum.SEARCH_STATIC_DATA_QUERY,
+            ColumnEnum.INTENT_TARGET_PACKAGE,
+            spaEnvironment.appContext.packageName
+        )
+        staticData.checkValue(
+            QueryEnum.SEARCH_STATIC_DATA_QUERY,
+            ColumnEnum.INTENT_TARGET_CLASS,
+            "com.android.settingslib.spa.tests.testutils.BlankActivity"
+        )
+
+        // Check extras in intent
+        val bundle =
+            staticData.getExtras(QueryEnum.SEARCH_STATIC_DATA_QUERY, ColumnEnum.INTENT_EXTRAS)
+        Truth.assertThat(bundle).isNotNull()
+        Truth.assertThat(bundle!!.size()).isEqualTo(3)
+        Truth.assertThat(bundle.getString("spaActivityDestination")).isEqualTo("SppForSearch")
+        Truth.assertThat(bundle.getString("highlightEntry"))
+            .isEqualTo(pageOwner.getEntryId("SearchStaticWithNoStatus"))
+        Truth.assertThat(bundle.getString("sessionSource")).isEqualTo("search")
+
+        staticData.moveToNext()
+        staticData.checkValue(
+            QueryEnum.SEARCH_STATIC_DATA_QUERY,
+            ColumnEnum.ENTRY_ID,
+            pageOwner.getEntryId("SearchStaticWithMutableStatus")
+        )
 
         val dynamicData = searchProvider.querySearchDynamicData()
         Truth.assertThat(dynamicData.count).isEqualTo(2)
+        dynamicData.moveToFirst()
+        dynamicData.checkValue(
+            QueryEnum.SEARCH_DYNAMIC_DATA_QUERY,
+            ColumnEnum.ENTRY_ID,
+            pageOwner.getEntryId("SearchDynamicWithMutableStatus")
+        )
+
+        dynamicData.moveToNext()
+        dynamicData.checkValue(
+            QueryEnum.SEARCH_DYNAMIC_DATA_QUERY,
+            ColumnEnum.ENTRY_ID,
+            pageOwner.getEntryId("SearchDynamicWithImmutableStatus")
+        )
+        dynamicData.checkValue(
+            QueryEnum.SEARCH_DYNAMIC_DATA_QUERY,
+            ColumnEnum.SEARCH_KEYWORD,
+            listOf("kw1", "kw2").toString()
+        )
     }
 }
 
 private fun Cursor.checkValue(query: QueryEnum, column: ColumnEnum, value: String) {
-    Truth.assertThat(getString(query.getIndex(column))).isEqualTo(value)
+    Truth.assertThat(getString(query.columnNames.indexOf(column))).isEqualTo(value)
+}
+
+private fun Cursor.getExtras(query: QueryEnum, column: ColumnEnum): Bundle? {
+    val extrasByte = getBlob(query.columnNames.indexOf(column)) ?: return null
+    val parcel = Parcel.obtain()
+    parcel.unmarshall(extrasByte, 0, extrasByte.size)
+    parcel.setDataPosition(0)
+    val bundle = Bundle()
+    bundle.readFromParcel(parcel)
+    return bundle
 }
 
 private fun SettingsPage.getEntryId(name: String): String {
diff --git a/packages/SystemUI/animation/Android.bp b/packages/SystemUI/animation/Android.bp
index 17ad55f..8acc2f8 100644
--- a/packages/SystemUI/animation/Android.bp
+++ b/packages/SystemUI/animation/Android.bp
@@ -44,23 +44,3 @@
     manifest: "AndroidManifest.xml",
     kotlincflags: ["-Xjvm-default=all"],
 }
-
-android_test {
-    name: "SystemUIAnimationLibTests",
-
-    static_libs: [
-        "SystemUIAnimationLib",
-        "androidx.test.ext.junit",
-        "androidx.test.rules",
-        "testables",
-    ],
-    libs: [
-        "android.test.base",
-    ],
-    srcs: [
-        "**/*.java",
-        "**/*.kt",
-    ],
-    kotlincflags: ["-Xjvm-default=all"],
-    test_suites: ["general-tests"],
-}
diff --git a/packages/SystemUI/animation/TEST_MAPPING b/packages/SystemUI/animation/TEST_MAPPING
deleted file mode 100644
index 3dc8510..0000000
--- a/packages/SystemUI/animation/TEST_MAPPING
+++ /dev/null
@@ -1,7 +0,0 @@
-{
-  "presubmit": [
-    {
-      "name": "SystemUIAnimationLibTests"
-    }
-  ]
-}
diff --git a/packages/SystemUI/src/com/android/systemui/ScreenDecorations.java b/packages/SystemUI/src/com/android/systemui/ScreenDecorations.java
index 02a6d7b..e6f559b 100644
--- a/packages/SystemUI/src/com/android/systemui/ScreenDecorations.java
+++ b/packages/SystemUI/src/com/android/systemui/ScreenDecorations.java
@@ -210,8 +210,10 @@
                 (FaceScanningOverlay) getOverlayView(mFaceScanningViewId);
         if (faceScanningOverlay != null) {
             faceScanningOverlay.setHideOverlayRunnable(() -> {
+                Trace.beginSection("ScreenDecorations#hideOverlayRunnable");
                 updateOverlayWindowVisibilityIfViewExists(
                         faceScanningOverlay.findViewById(mFaceScanningViewId));
+                Trace.endSection();
             });
             faceScanningOverlay.enableShowProtection(false);
         }
@@ -273,16 +275,18 @@
             if (mOverlays == null || !shouldOptimizeVisibility()) {
                 return;
             }
-
+            Trace.beginSection("ScreenDecorations#updateOverlayWindowVisibilityIfViewExists");
             for (final OverlayWindow overlay : mOverlays) {
                 if (overlay == null) {
                     continue;
                 }
                 if (overlay.getView(view.getId()) != null) {
                     overlay.getRootView().setVisibility(getWindowVisibility(overlay, true));
+                    Trace.endSection();
                     return;
                 }
             }
+            Trace.endSection();
         });
     }
 
@@ -370,6 +374,7 @@
     }
 
     private void startOnScreenDecorationsThread() {
+        Trace.beginSection("ScreenDecorations#startOnScreenDecorationsThread");
         mWindowManager = mContext.getSystemService(WindowManager.class);
         mDisplayManager = mContext.getSystemService(DisplayManager.class);
         mContext.getDisplay().getDisplayInfo(mDisplayInfo);
@@ -472,6 +477,7 @@
 
         mDisplayManager.registerDisplayListener(mDisplayListener, mHandler);
         updateConfiguration();
+        Trace.endSection();
     }
 
     @VisibleForTesting
@@ -521,6 +527,12 @@
     }
 
     private void setupDecorations() {
+        Trace.beginSection("ScreenDecorations#setupDecorations");
+        setupDecorationsInner();
+        Trace.endSection();
+    }
+
+    private void setupDecorationsInner() {
         if (hasRoundedCorners() || shouldDrawCutout() || isPrivacyDotEnabled()
                 || mFaceScanningFactory.getHasProviders()) {
 
@@ -573,7 +585,11 @@
                 return;
             }
 
-            mMainExecutor.execute(() -> mTunerService.addTunable(this, SIZE));
+            mMainExecutor.execute(() -> {
+                Trace.beginSection("ScreenDecorations#addTunable");
+                mTunerService.addTunable(this, SIZE);
+                Trace.endSection();
+            });
 
             // Watch color inversion and invert the overlay as needed.
             if (mColorInversionSetting == null) {
@@ -593,7 +609,11 @@
             mUserTracker.addCallback(mUserChangedCallback, mExecutor);
             mIsRegistered = true;
         } else {
-            mMainExecutor.execute(() -> mTunerService.removeTunable(this));
+            mMainExecutor.execute(() -> {
+                Trace.beginSection("ScreenDecorations#removeTunable");
+                mTunerService.removeTunable(this);
+                Trace.endSection();
+            });
 
             if (mColorInversionSetting != null) {
                 mColorInversionSetting.setListening(false);
@@ -939,6 +959,7 @@
         }
 
         mExecutor.execute(() -> {
+            Trace.beginSection("ScreenDecorations#onConfigurationChanged");
             int oldRotation = mRotation;
             mPendingConfigChange = false;
             updateConfiguration();
@@ -951,6 +972,7 @@
                 // the updated rotation).
                 updateLayoutParams();
             }
+            Trace.endSection();
         });
     }
 
@@ -1119,6 +1141,7 @@
             if (mOverlays == null || !SIZE.equals(key)) {
                 return;
             }
+            Trace.beginSection("ScreenDecorations#onTuningChanged");
             try {
                 final int sizeFactor = Integer.parseInt(newValue);
                 mRoundedCornerResDelegate.setTuningSizeFactor(sizeFactor);
@@ -1132,6 +1155,7 @@
                     R.id.rounded_corner_bottom_right
             });
             updateHwLayerRoundedCornerExistAndSize();
+            Trace.endSection();
         });
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/controls/ui/ControlsUiControllerImpl.kt b/packages/SystemUI/src/com/android/systemui/controls/ui/ControlsUiControllerImpl.kt
index 4c8e1ac..a07c716 100644
--- a/packages/SystemUI/src/com/android/systemui/controls/ui/ControlsUiControllerImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/controls/ui/ControlsUiControllerImpl.kt
@@ -28,6 +28,7 @@
 import android.graphics.drawable.Drawable
 import android.graphics.drawable.LayerDrawable
 import android.service.controls.Control
+import android.service.controls.ControlsProviderService
 import android.util.Log
 import android.view.ContextThemeWrapper
 import android.view.LayoutInflater
@@ -48,6 +49,7 @@
 import com.android.systemui.R
 import com.android.systemui.controls.ControlsMetricsLogger
 import com.android.systemui.controls.ControlsServiceInfo
+import com.android.systemui.controls.ControlsSettingsRepository
 import com.android.systemui.controls.CustomIconCache
 import com.android.systemui.controls.controller.ControlsController
 import com.android.systemui.controls.controller.StructureInfo
@@ -96,6 +98,7 @@
         private val userFileManager: UserFileManager,
         private val userTracker: UserTracker,
         private val taskViewFactory: Optional<TaskViewFactory>,
+        private val controlsSettingsRepository: ControlsSettingsRepository,
         dumpManager: DumpManager
 ) : ControlsUiController, Dumpable {
 
@@ -354,7 +357,6 @@
                 } else {
                     items[0]
                 }
-
         maybeUpdateSelectedItem(selectionItem)
 
         createControlsSpaceFrame()
@@ -374,11 +376,20 @@
     }
 
     private fun createPanelView(componentName: ComponentName) {
-        val pendingIntent = PendingIntent.getActivity(
+        val setting = controlsSettingsRepository
+                .allowActionOnTrivialControlsInLockscreen.value
+        val pendingIntent = PendingIntent.getActivityAsUser(
                 context,
                 0,
-                Intent().setComponent(componentName),
-                PendingIntent.FLAG_IMMUTABLE
+                Intent()
+                        .setComponent(componentName)
+                        .putExtra(
+                                ControlsProviderService.EXTRA_LOCKSCREEN_ALLOW_TRIVIAL_CONTROLS,
+                                setting
+                        ),
+                PendingIntent.FLAG_IMMUTABLE or PendingIntent.FLAG_UPDATE_CURRENT,
+                null,
+                userTracker.userHandle
         )
 
         parent.requireViewById<View>(R.id.controls_scroll_view).visibility = View.GONE
@@ -698,6 +709,8 @@
             println("hidden: $hidden")
             println("selectedItem: $selectedItem")
             println("lastSelections: $lastSelections")
+            println("setting: ${controlsSettingsRepository
+                    .allowActionOnTrivialControlsInLockscreen.value}")
         }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSTileHost.java b/packages/SystemUI/src/com/android/systemui/qs/QSTileHost.java
index 6240c10..cad296b 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSTileHost.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSTileHost.java
@@ -608,7 +608,7 @@
 
         if (TextUtils.isEmpty(tileList)) {
             tileList = res.getString(R.string.quick_settings_tiles);
-            if (DEBUG) Log.d(TAG, "Loaded tile specs from config: " + tileList);
+            if (DEBUG) Log.d(TAG, "Loaded tile specs from default config: " + tileList);
         } else {
             if (DEBUG) Log.d(TAG, "Loaded tile specs from setting: " + tileList);
         }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/WorkModeTile.java b/packages/SystemUI/src/com/android/systemui/qs/tiles/WorkModeTile.java
index 7130294..a6c7781 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/WorkModeTile.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/WorkModeTile.java
@@ -27,6 +27,7 @@
 import android.view.View;
 import android.widget.Switch;
 
+import androidx.annotation.MainThread;
 import androidx.annotation.Nullable;
 
 import com.android.internal.logging.MetricsLogger;
@@ -91,11 +92,13 @@
     }
 
     @Override
+    @MainThread
     public void onManagedProfileChanged() {
         refreshState(mProfileController.isWorkModeEnabled());
     }
 
     @Override
+    @MainThread
     public void onManagedProfileRemoved() {
         mHost.removeTile(getTileSpec());
         mHost.unmarkTileAsAutoAdded(getTileSpec());
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileController.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileController.java
index 4969a1c..6811bf6 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileController.java
@@ -14,6 +14,8 @@
 
 package com.android.systemui.statusbar.phone;
 
+import androidx.annotation.MainThread;
+
 import com.android.systemui.statusbar.phone.ManagedProfileController.Callback;
 import com.android.systemui.statusbar.policy.CallbackController;
 
@@ -25,8 +27,20 @@
 
     boolean isWorkModeEnabled();
 
-    public interface Callback {
+    /**
+     * Callback to get updates about work profile status.
+     */
+    interface Callback {
+        /**
+         * Called when managed profile change is detected. This always runs on the main thread.
+         */
+        @MainThread
         void onManagedProfileChanged();
+
+        /**
+         * Called when managed profile removal is detected. This always runs on the main thread.
+         */
+        @MainThread
         void onManagedProfileRemoved();
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java
index 4beb87d..abdf827 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java
@@ -33,31 +33,28 @@
 
 import javax.inject.Inject;
 
-/**
- */
 @SysUISingleton
 public class ManagedProfileControllerImpl implements ManagedProfileController {
 
     private final List<Callback> mCallbacks = new ArrayList<>();
-
+    private final UserTrackerCallback mUserTrackerCallback = new UserTrackerCallback();
     private final Context mContext;
     private final Executor mMainExecutor;
     private final UserManager mUserManager;
     private final UserTracker mUserTracker;
     private final LinkedList<UserInfo> mProfiles;
+
     private boolean mListening;
     private int mCurrentUser;
 
-    /**
-     */
     @Inject
     public ManagedProfileControllerImpl(Context context, @Main Executor mainExecutor,
-            UserTracker userTracker) {
+            UserTracker userTracker, UserManager userManager) {
         mContext = context;
         mMainExecutor = mainExecutor;
-        mUserManager = UserManager.get(mContext);
+        mUserManager = userManager;
         mUserTracker = userTracker;
-        mProfiles = new LinkedList<UserInfo>();
+        mProfiles = new LinkedList<>();
     }
 
     @Override
@@ -100,16 +97,22 @@
                 }
             }
             if (mProfiles.size() == 0 && hadProfile && (user == mCurrentUser)) {
-                for (Callback callback : mCallbacks) {
-                    callback.onManagedProfileRemoved();
-                }
+                mMainExecutor.execute(this::notifyManagedProfileRemoved);
             }
             mCurrentUser = user;
         }
     }
 
+    private void notifyManagedProfileRemoved() {
+        for (Callback callback : mCallbacks) {
+            callback.onManagedProfileRemoved();
+        }
+    }
+
     public boolean hasActiveProfile() {
-        if (!mListening) reloadManagedProfiles();
+        if (!mListening || mUserTracker.getUserId() != mCurrentUser) {
+            reloadManagedProfiles();
+        }
         synchronized (mProfiles) {
             return mProfiles.size() > 0;
         }
@@ -134,28 +137,28 @@
         mListening = listening;
         if (listening) {
             reloadManagedProfiles();
-            mUserTracker.addCallback(mUserChangedCallback, mMainExecutor);
+            mUserTracker.addCallback(mUserTrackerCallback, mMainExecutor);
         } else {
-            mUserTracker.removeCallback(mUserChangedCallback);
+            mUserTracker.removeCallback(mUserTrackerCallback);
         }
     }
 
-    private final UserTracker.Callback mUserChangedCallback =
-            new UserTracker.Callback() {
-                @Override
-                public void onUserChanged(int newUser, @NonNull Context userContext) {
-                    reloadManagedProfiles();
-                    for (Callback callback : mCallbacks) {
-                        callback.onManagedProfileChanged();
-                    }
-                }
+    private final class UserTrackerCallback implements UserTracker.Callback {
 
-                @Override
-                public void onProfilesChanged(@NonNull List<UserInfo> profiles) {
-                    reloadManagedProfiles();
-                    for (Callback callback : mCallbacks) {
-                        callback.onManagedProfileChanged();
-                    }
-                }
-            };
+        @Override
+        public void onUserChanged(int newUser, @NonNull Context userContext) {
+            reloadManagedProfiles();
+            for (Callback callback : mCallbacks) {
+                callback.onManagedProfileChanged();
+            }
+        }
+
+        @Override
+        public void onProfilesChanged(@NonNull List<UserInfo> profiles) {
+            reloadManagedProfiles();
+            for (Callback callback : mCallbacks) {
+                callback.onManagedProfileChanged();
+            }
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/dagger/StatusBarPipelineModule.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/dagger/StatusBarPipelineModule.kt
index e7afc50..c350c78 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/dagger/StatusBarPipelineModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/dagger/StatusBarPipelineModule.kt
@@ -25,7 +25,7 @@
 import com.android.systemui.statusbar.pipeline.airplane.ui.viewmodel.AirplaneModeViewModel
 import com.android.systemui.statusbar.pipeline.airplane.ui.viewmodel.AirplaneModeViewModelImpl
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionsRepository
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionsRepositoryImpl
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileRepositorySwitcher
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.UserSetupRepository
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.UserSetupRepositoryImpl
 import com.android.systemui.statusbar.pipeline.mobile.domain.interactor.MobileIconsInteractor
@@ -63,7 +63,7 @@
 
     @Binds
     abstract fun mobileConnectionsRepository(
-        impl: MobileConnectionsRepositoryImpl
+        impl: MobileRepositorySwitcher
     ): MobileConnectionsRepository
 
     @Binds abstract fun userSetupRepository(impl: UserSetupRepositoryImpl): UserSetupRepository
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionRepository.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionRepository.kt
index 581842b..f094563 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionRepository.kt
@@ -16,44 +16,13 @@
 
 package com.android.systemui.statusbar.pipeline.mobile.data.repository
 
-import android.content.Context
-import android.database.ContentObserver
-import android.provider.Settings.Global
-import android.telephony.CellSignalStrength
-import android.telephony.CellSignalStrengthCdma
-import android.telephony.ServiceState
-import android.telephony.SignalStrength
 import android.telephony.SubscriptionInfo
 import android.telephony.SubscriptionManager
 import android.telephony.TelephonyCallback
-import android.telephony.TelephonyDisplayInfo
-import android.telephony.TelephonyDisplayInfo.OVERRIDE_NETWORK_TYPE_NONE
 import android.telephony.TelephonyManager
-import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
-import com.android.systemui.dagger.qualifiers.Application
-import com.android.systemui.dagger.qualifiers.Background
-import com.android.systemui.statusbar.pipeline.mobile.data.model.DefaultNetworkType
 import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileSubscriptionModel
-import com.android.systemui.statusbar.pipeline.mobile.data.model.OverrideNetworkType
-import com.android.systemui.statusbar.pipeline.mobile.data.model.toDataConnectionType
-import com.android.systemui.statusbar.pipeline.shared.ConnectivityPipelineLogger
-import com.android.systemui.statusbar.pipeline.shared.ConnectivityPipelineLogger.Companion.logOutputChange
-import com.android.systemui.util.settings.GlobalSettings
-import java.lang.IllegalStateException
-import javax.inject.Inject
-import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.asExecutor
-import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.MutableSharedFlow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.mapLatest
-import kotlinx.coroutines.flow.merge
-import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.flow.stateIn
 
 /**
  * Every mobile line of service can be identified via a [SubscriptionInfo] object. We set up a
@@ -80,183 +49,3 @@
      */
     val isDefaultDataSubscription: StateFlow<Boolean>
 }
-
-@Suppress("EXPERIMENTAL_IS_NOT_ENABLED")
-@OptIn(ExperimentalCoroutinesApi::class)
-class MobileConnectionRepositoryImpl(
-    private val context: Context,
-    private val subId: Int,
-    private val telephonyManager: TelephonyManager,
-    private val globalSettings: GlobalSettings,
-    defaultDataSubId: StateFlow<Int>,
-    globalMobileDataSettingChangedEvent: Flow<Unit>,
-    bgDispatcher: CoroutineDispatcher,
-    logger: ConnectivityPipelineLogger,
-    scope: CoroutineScope,
-) : MobileConnectionRepository {
-    init {
-        if (telephonyManager.subscriptionId != subId) {
-            throw IllegalStateException(
-                "TelephonyManager should be created with subId($subId). " +
-                    "Found ${telephonyManager.subscriptionId} instead."
-            )
-        }
-    }
-
-    private val telephonyCallbackEvent = MutableSharedFlow<Unit>(extraBufferCapacity = 1)
-
-    override val subscriptionModelFlow: StateFlow<MobileSubscriptionModel> = run {
-        var state = MobileSubscriptionModel()
-        conflatedCallbackFlow {
-                // TODO (b/240569788): log all of these into the connectivity logger
-                val callback =
-                    object :
-                        TelephonyCallback(),
-                        TelephonyCallback.ServiceStateListener,
-                        TelephonyCallback.SignalStrengthsListener,
-                        TelephonyCallback.DataConnectionStateListener,
-                        TelephonyCallback.DataActivityListener,
-                        TelephonyCallback.CarrierNetworkListener,
-                        TelephonyCallback.DisplayInfoListener {
-                        override fun onServiceStateChanged(serviceState: ServiceState) {
-                            state = state.copy(isEmergencyOnly = serviceState.isEmergencyOnly)
-                            trySend(state)
-                        }
-
-                        override fun onSignalStrengthsChanged(signalStrength: SignalStrength) {
-                            val cdmaLevel =
-                                signalStrength
-                                    .getCellSignalStrengths(CellSignalStrengthCdma::class.java)
-                                    .let { strengths ->
-                                        if (!strengths.isEmpty()) {
-                                            strengths[0].level
-                                        } else {
-                                            CellSignalStrength.SIGNAL_STRENGTH_NONE_OR_UNKNOWN
-                                        }
-                                    }
-
-                            val primaryLevel = signalStrength.level
-
-                            state =
-                                state.copy(
-                                    cdmaLevel = cdmaLevel,
-                                    primaryLevel = primaryLevel,
-                                    isGsm = signalStrength.isGsm,
-                                )
-                            trySend(state)
-                        }
-
-                        override fun onDataConnectionStateChanged(
-                            dataState: Int,
-                            networkType: Int
-                        ) {
-                            state =
-                                state.copy(dataConnectionState = dataState.toDataConnectionType())
-                            trySend(state)
-                        }
-
-                        override fun onDataActivity(direction: Int) {
-                            state = state.copy(dataActivityDirection = direction)
-                            trySend(state)
-                        }
-
-                        override fun onCarrierNetworkChange(active: Boolean) {
-                            state = state.copy(carrierNetworkChangeActive = active)
-                            trySend(state)
-                        }
-
-                        override fun onDisplayInfoChanged(
-                            telephonyDisplayInfo: TelephonyDisplayInfo
-                        ) {
-                            val networkType =
-                                if (
-                                    telephonyDisplayInfo.overrideNetworkType ==
-                                        OVERRIDE_NETWORK_TYPE_NONE
-                                ) {
-                                    DefaultNetworkType(telephonyDisplayInfo.networkType)
-                                } else {
-                                    OverrideNetworkType(telephonyDisplayInfo.overrideNetworkType)
-                                }
-                            state = state.copy(resolvedNetworkType = networkType)
-                            trySend(state)
-                        }
-                    }
-                telephonyManager.registerTelephonyCallback(bgDispatcher.asExecutor(), callback)
-                awaitClose { telephonyManager.unregisterTelephonyCallback(callback) }
-            }
-            .onEach { telephonyCallbackEvent.tryEmit(Unit) }
-            .logOutputChange(logger, "MobileSubscriptionModel")
-            .stateIn(scope, SharingStarted.WhileSubscribed(), state)
-    }
-
-    /** Produces whenever the mobile data setting changes for this subId */
-    private val localMobileDataSettingChangedEvent: Flow<Unit> = conflatedCallbackFlow {
-        val observer =
-            object : ContentObserver(null) {
-                override fun onChange(selfChange: Boolean) {
-                    trySend(Unit)
-                }
-            }
-
-        globalSettings.registerContentObserver(
-            globalSettings.getUriFor("${Global.MOBILE_DATA}$subId"),
-            /* notifyForDescendants */ true,
-            observer
-        )
-
-        awaitClose { context.contentResolver.unregisterContentObserver(observer) }
-    }
-
-    /**
-     * There are a few cases where we will need to poll [TelephonyManager] so we can update some
-     * internal state where callbacks aren't provided. Any of those events should be merged into
-     * this flow, which can be used to trigger the polling.
-     */
-    private val telephonyPollingEvent: Flow<Unit> =
-        merge(
-            telephonyCallbackEvent,
-            localMobileDataSettingChangedEvent,
-            globalMobileDataSettingChangedEvent,
-        )
-
-    override val dataEnabled: StateFlow<Boolean> =
-        telephonyPollingEvent
-            .mapLatest { dataConnectionAllowed() }
-            .stateIn(scope, SharingStarted.WhileSubscribed(), dataConnectionAllowed())
-
-    private fun dataConnectionAllowed(): Boolean = telephonyManager.isDataConnectionAllowed
-
-    override val isDefaultDataSubscription: StateFlow<Boolean> =
-        defaultDataSubId
-            .mapLatest { it == subId }
-            .stateIn(scope, SharingStarted.WhileSubscribed(), defaultDataSubId.value == subId)
-
-    class Factory
-    @Inject
-    constructor(
-        private val context: Context,
-        private val telephonyManager: TelephonyManager,
-        private val logger: ConnectivityPipelineLogger,
-        private val globalSettings: GlobalSettings,
-        @Background private val bgDispatcher: CoroutineDispatcher,
-        @Application private val scope: CoroutineScope,
-    ) {
-        fun build(
-            subId: Int,
-            defaultDataSubId: StateFlow<Int>,
-            globalMobileDataSettingChangedEvent: Flow<Unit>,
-        ): MobileConnectionRepository {
-            return MobileConnectionRepositoryImpl(
-                context,
-                subId,
-                telephonyManager.createForSubscriptionId(subId),
-                globalSettings,
-                defaultDataSubId,
-                globalMobileDataSettingChangedEvent,
-                bgDispatcher,
-                logger,
-                scope,
-            )
-        }
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepository.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepository.kt
index c3c1f14..14200f0 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepository.kt
@@ -16,53 +16,14 @@
 
 package com.android.systemui.statusbar.pipeline.mobile.data.repository
 
-import android.annotation.SuppressLint
-import android.content.Context
-import android.content.IntentFilter
-import android.database.ContentObserver
-import android.net.ConnectivityManager
-import android.net.ConnectivityManager.NetworkCallback
-import android.net.Network
-import android.net.NetworkCapabilities
-import android.net.NetworkCapabilities.NET_CAPABILITY_VALIDATED
-import android.net.NetworkCapabilities.TRANSPORT_CELLULAR
 import android.provider.Settings
-import android.provider.Settings.Global.MOBILE_DATA
-import android.telephony.CarrierConfigManager
 import android.telephony.SubscriptionInfo
 import android.telephony.SubscriptionManager
-import android.telephony.SubscriptionManager.INVALID_SUBSCRIPTION_ID
-import android.telephony.TelephonyCallback
-import android.telephony.TelephonyCallback.ActiveDataSubscriptionIdListener
-import android.telephony.TelephonyManager
-import androidx.annotation.VisibleForTesting
-import com.android.internal.telephony.PhoneConstants
 import com.android.settingslib.mobile.MobileMappings
 import com.android.settingslib.mobile.MobileMappings.Config
-import com.android.systemui.broadcast.BroadcastDispatcher
-import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.dagger.qualifiers.Application
-import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileConnectivityModel
-import com.android.systemui.statusbar.pipeline.shared.ConnectivityPipelineLogger
-import com.android.systemui.util.settings.GlobalSettings
-import javax.inject.Inject
-import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.asExecutor
-import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.MutableSharedFlow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.distinctUntilChanged
-import kotlinx.coroutines.flow.mapLatest
-import kotlinx.coroutines.flow.merge
-import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.flow.stateIn
-import kotlinx.coroutines.withContext
 
 /**
  * Repo for monitoring the complete active subscription info list, to be consumed and filtered based
@@ -90,202 +51,3 @@
     /** Observe changes to the [Settings.Global.MOBILE_DATA] setting */
     val globalMobileDataSettingChangedEvent: Flow<Unit>
 }
-
-@Suppress("EXPERIMENTAL_IS_NOT_ENABLED")
-@OptIn(ExperimentalCoroutinesApi::class)
-@SysUISingleton
-class MobileConnectionsRepositoryImpl
-@Inject
-constructor(
-    private val connectivityManager: ConnectivityManager,
-    private val subscriptionManager: SubscriptionManager,
-    private val telephonyManager: TelephonyManager,
-    private val logger: ConnectivityPipelineLogger,
-    broadcastDispatcher: BroadcastDispatcher,
-    private val globalSettings: GlobalSettings,
-    private val context: Context,
-    @Background private val bgDispatcher: CoroutineDispatcher,
-    @Application private val scope: CoroutineScope,
-    private val mobileConnectionRepositoryFactory: MobileConnectionRepositoryImpl.Factory
-) : MobileConnectionsRepository {
-    private val subIdRepositoryCache: MutableMap<Int, MobileConnectionRepository> = mutableMapOf()
-
-    /**
-     * State flow that emits the set of mobile data subscriptions, each represented by its own
-     * [SubscriptionInfo]. We probably only need the [SubscriptionInfo.getSubscriptionId] of each
-     * info object, but for now we keep track of the infos themselves.
-     */
-    override val subscriptionsFlow: StateFlow<List<SubscriptionInfo>> =
-        conflatedCallbackFlow {
-                val callback =
-                    object : SubscriptionManager.OnSubscriptionsChangedListener() {
-                        override fun onSubscriptionsChanged() {
-                            trySend(Unit)
-                        }
-                    }
-
-                subscriptionManager.addOnSubscriptionsChangedListener(
-                    bgDispatcher.asExecutor(),
-                    callback,
-                )
-
-                awaitClose { subscriptionManager.removeOnSubscriptionsChangedListener(callback) }
-            }
-            .mapLatest { fetchSubscriptionsList() }
-            .onEach { infos -> dropUnusedReposFromCache(infos) }
-            .stateIn(scope, started = SharingStarted.WhileSubscribed(), listOf())
-
-    /** StateFlow that keeps track of the current active mobile data subscription */
-    override val activeMobileDataSubscriptionId: StateFlow<Int> =
-        conflatedCallbackFlow {
-                val callback =
-                    object : TelephonyCallback(), ActiveDataSubscriptionIdListener {
-                        override fun onActiveDataSubscriptionIdChanged(subId: Int) {
-                            trySend(subId)
-                        }
-                    }
-
-                telephonyManager.registerTelephonyCallback(bgDispatcher.asExecutor(), callback)
-                awaitClose { telephonyManager.unregisterTelephonyCallback(callback) }
-            }
-            .stateIn(scope, started = SharingStarted.WhileSubscribed(), INVALID_SUBSCRIPTION_ID)
-
-    private val defaultDataSubIdChangeEvent: MutableSharedFlow<Unit> =
-        MutableSharedFlow(extraBufferCapacity = 1)
-
-    override val defaultDataSubId: StateFlow<Int> =
-        broadcastDispatcher
-            .broadcastFlow(
-                IntentFilter(TelephonyManager.ACTION_DEFAULT_DATA_SUBSCRIPTION_CHANGED)
-            ) { intent, _ ->
-                intent.getIntExtra(PhoneConstants.SUBSCRIPTION_KEY, INVALID_SUBSCRIPTION_ID)
-            }
-            .distinctUntilChanged()
-            .onEach { defaultDataSubIdChangeEvent.tryEmit(Unit) }
-            .stateIn(
-                scope,
-                SharingStarted.WhileSubscribed(),
-                SubscriptionManager.getDefaultDataSubscriptionId()
-            )
-
-    private val carrierConfigChangedEvent =
-        broadcastDispatcher.broadcastFlow(
-            IntentFilter(CarrierConfigManager.ACTION_CARRIER_CONFIG_CHANGED)
-        )
-
-    /**
-     * [Config] is an object that tracks relevant configuration flags for a given subscription ID.
-     * In the case of [MobileMappings], it's hard-coded to check the default data subscription's
-     * config, so this will apply to every icon that we care about.
-     *
-     * Relevant bits in the config are things like
-     * [CarrierConfigManager.KEY_SHOW_4G_FOR_LTE_DATA_ICON_BOOL]
-     *
-     * This flow will produce whenever the default data subscription or the carrier config changes.
-     */
-    override val defaultDataSubRatConfig: StateFlow<Config> =
-        merge(defaultDataSubIdChangeEvent, carrierConfigChangedEvent)
-            .mapLatest { Config.readConfig(context) }
-            .stateIn(
-                scope,
-                SharingStarted.WhileSubscribed(),
-                initialValue = Config.readConfig(context)
-            )
-
-    override fun getRepoForSubId(subId: Int): MobileConnectionRepository {
-        if (!isValidSubId(subId)) {
-            throw IllegalArgumentException(
-                "subscriptionId $subId is not in the list of valid subscriptions"
-            )
-        }
-
-        return subIdRepositoryCache[subId]
-            ?: createRepositoryForSubId(subId).also { subIdRepositoryCache[subId] = it }
-    }
-
-    /**
-     * In single-SIM devices, the [MOBILE_DATA] setting is phone-wide. For multi-SIM, the individual
-     * connection repositories also observe the URI for [MOBILE_DATA] + subId.
-     */
-    override val globalMobileDataSettingChangedEvent: Flow<Unit> = conflatedCallbackFlow {
-        val observer =
-            object : ContentObserver(null) {
-                override fun onChange(selfChange: Boolean) {
-                    trySend(Unit)
-                }
-            }
-
-        globalSettings.registerContentObserver(
-            globalSettings.getUriFor(MOBILE_DATA),
-            true,
-            observer
-        )
-
-        awaitClose { context.contentResolver.unregisterContentObserver(observer) }
-    }
-
-    @SuppressLint("MissingPermission")
-    override val defaultMobileNetworkConnectivity: StateFlow<MobileConnectivityModel> =
-        conflatedCallbackFlow {
-                val callback =
-                    object : NetworkCallback(FLAG_INCLUDE_LOCATION_INFO) {
-                        override fun onLost(network: Network) {
-                            // Send a disconnected model when lost. Maybe should create a sealed
-                            // type or null here?
-                            trySend(MobileConnectivityModel())
-                        }
-
-                        override fun onCapabilitiesChanged(
-                            network: Network,
-                            caps: NetworkCapabilities
-                        ) {
-                            trySend(
-                                MobileConnectivityModel(
-                                    isConnected = caps.hasTransport(TRANSPORT_CELLULAR),
-                                    isValidated = caps.hasCapability(NET_CAPABILITY_VALIDATED),
-                                )
-                            )
-                        }
-                    }
-
-                connectivityManager.registerDefaultNetworkCallback(callback)
-
-                awaitClose { connectivityManager.unregisterNetworkCallback(callback) }
-            }
-            .stateIn(scope, SharingStarted.WhileSubscribed(), MobileConnectivityModel())
-
-    private fun isValidSubId(subId: Int): Boolean {
-        subscriptionsFlow.value.forEach {
-            if (it.subscriptionId == subId) {
-                return true
-            }
-        }
-
-        return false
-    }
-
-    @VisibleForTesting fun getSubIdRepoCache() = subIdRepositoryCache
-
-    private fun createRepositoryForSubId(subId: Int): MobileConnectionRepository {
-        return mobileConnectionRepositoryFactory.build(
-            subId,
-            defaultDataSubId,
-            globalMobileDataSettingChangedEvent,
-        )
-    }
-
-    private fun dropUnusedReposFromCache(newInfos: List<SubscriptionInfo>) {
-        // Remove any connection repository from the cache that isn't in the new set of IDs. They
-        // will get garbage collected once their subscribers go away
-        val currentValidSubscriptionIds = newInfos.map { it.subscriptionId }
-
-        subIdRepositoryCache.keys.forEach {
-            if (!currentValidSubscriptionIds.contains(it)) {
-                subIdRepositoryCache.remove(it)
-            }
-        }
-    }
-
-    private suspend fun fetchSubscriptionsList(): List<SubscriptionInfo> =
-        withContext(bgDispatcher) { subscriptionManager.completeActiveSubscriptionInfoList }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcher.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcher.kt
new file mode 100644
index 0000000..e214005
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcher.kt
@@ -0,0 +1,162 @@
+/*
+ * Copyright (C) 2022 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.systemui.statusbar.pipeline.mobile.data.repository
+
+import android.os.Bundle
+import android.telephony.SubscriptionInfo
+import androidx.annotation.VisibleForTesting
+import com.android.settingslib.mobile.MobileMappings
+import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.demomode.DemoMode
+import com.android.systemui.demomode.DemoModeController
+import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileConnectivityModel
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.DemoMobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.prod.MobileConnectionsRepositoryImpl
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.mapLatest
+import kotlinx.coroutines.flow.stateIn
+
+/**
+ * A provider for the [MobileConnectionsRepository] interface that can choose between the Demo and
+ * Prod concrete implementations at runtime. It works by defining a base flow, [activeRepo], which
+ * switches based on the latest information from [DemoModeController], and switches every flow in
+ * the interface to point to the currently-active provider. This allows us to put the demo mode
+ * interface in its own repository, completely separate from the real version, while still using all
+ * of the prod implementations for the rest of the pipeline (interactors and onward). Looks
+ * something like this:
+ *
+ * ```
+ * RealRepository
+ *                 │
+ *                 ├──►RepositorySwitcher──►RealInteractor──►RealViewModel
+ *                 │
+ * DemoRepository
+ * ```
+ *
+ * NOTE: because the UI layer for mobile icons relies on a nested-repository structure, it is likely
+ * that we will have to drain the subscription list whenever demo mode changes. Otherwise if a real
+ * subscription list [1] is replaced with a demo subscription list [1], the view models will not see
+ * a change (due to `distinctUntilChanged`) and will not refresh their data providers to the demo
+ * implementation.
+ */
+@Suppress("EXPERIMENTAL_IS_NOT_ENABLED")
+@OptIn(ExperimentalCoroutinesApi::class)
+class MobileRepositorySwitcher
+@Inject
+constructor(
+    @Application scope: CoroutineScope,
+    val realRepository: MobileConnectionsRepositoryImpl,
+    val demoMobileConnectionsRepository: DemoMobileConnectionsRepository,
+    demoModeController: DemoModeController,
+) : MobileConnectionsRepository {
+
+    val isDemoMode: StateFlow<Boolean> =
+        conflatedCallbackFlow {
+                val callback =
+                    object : DemoMode {
+                        override fun dispatchDemoCommand(command: String?, args: Bundle?) {
+                            // Nothing, we just care about on/off
+                        }
+
+                        override fun onDemoModeStarted() {
+                            demoMobileConnectionsRepository.startProcessingCommands()
+                            trySend(true)
+                        }
+
+                        override fun onDemoModeFinished() {
+                            demoMobileConnectionsRepository.stopProcessingCommands()
+                            trySend(false)
+                        }
+                    }
+
+                demoModeController.addCallback(callback)
+                awaitClose { demoModeController.removeCallback(callback) }
+            }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), demoModeController.isInDemoMode)
+
+    // Convenient definition flow for the currently active repo (based on demo mode or not)
+    @VisibleForTesting
+    internal val activeRepo: StateFlow<MobileConnectionsRepository> =
+        isDemoMode
+            .mapLatest { demoMode ->
+                if (demoMode) {
+                    demoMobileConnectionsRepository
+                } else {
+                    realRepository
+                }
+            }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), realRepository)
+
+    override val subscriptionsFlow: StateFlow<List<SubscriptionInfo>> =
+        activeRepo
+            .flatMapLatest { it.subscriptionsFlow }
+            .stateIn(
+                scope,
+                SharingStarted.WhileSubscribed(),
+                realRepository.subscriptionsFlow.value
+            )
+
+    override val activeMobileDataSubscriptionId: StateFlow<Int> =
+        activeRepo
+            .flatMapLatest { it.activeMobileDataSubscriptionId }
+            .stateIn(
+                scope,
+                SharingStarted.WhileSubscribed(),
+                realRepository.activeMobileDataSubscriptionId.value
+            )
+
+    override val defaultDataSubRatConfig: StateFlow<MobileMappings.Config> =
+        activeRepo
+            .flatMapLatest { it.defaultDataSubRatConfig }
+            .stateIn(
+                scope,
+                SharingStarted.WhileSubscribed(),
+                realRepository.defaultDataSubRatConfig.value
+            )
+
+    override val defaultDataSubId: StateFlow<Int> =
+        activeRepo
+            .flatMapLatest { it.defaultDataSubId }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), realRepository.defaultDataSubId.value)
+
+    override val defaultMobileNetworkConnectivity: StateFlow<MobileConnectivityModel> =
+        activeRepo
+            .flatMapLatest { it.defaultMobileNetworkConnectivity }
+            .stateIn(
+                scope,
+                SharingStarted.WhileSubscribed(),
+                realRepository.defaultMobileNetworkConnectivity.value
+            )
+
+    override val globalMobileDataSettingChangedEvent: Flow<Unit> =
+        activeRepo.flatMapLatest { it.globalMobileDataSettingChangedEvent }
+
+    override fun getRepoForSubId(subId: Int): MobileConnectionRepository {
+        if (isDemoMode.value) {
+            return demoMobileConnectionsRepository.getRepoForSubId(subId)
+        }
+        return realRepository.getRepoForSubId(subId)
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepository.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepository.kt
new file mode 100644
index 0000000..5f2feb2
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepository.kt
@@ -0,0 +1,277 @@
+/*
+ * Copyright (C) 2022 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.systemui.statusbar.pipeline.mobile.data.repository.demo
+
+import android.content.Context
+import android.telephony.Annotation
+import android.telephony.SubscriptionInfo
+import android.telephony.SubscriptionManager.INVALID_SUBSCRIPTION_ID
+import android.telephony.TelephonyDisplayInfo.OVERRIDE_NETWORK_TYPE_NR_ADVANCED
+import android.telephony.TelephonyManager.NETWORK_TYPE_GSM
+import android.telephony.TelephonyManager.NETWORK_TYPE_LTE
+import android.telephony.TelephonyManager.NETWORK_TYPE_NR
+import android.telephony.TelephonyManager.NETWORK_TYPE_UMTS
+import android.telephony.TelephonyManager.NETWORK_TYPE_UNKNOWN
+import android.util.Log
+import com.android.settingslib.SignalIcon
+import com.android.settingslib.mobile.MobileMappings
+import com.android.settingslib.mobile.TelephonyIcons
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.statusbar.pipeline.mobile.data.model.DataConnectionState
+import com.android.systemui.statusbar.pipeline.mobile.data.model.DefaultNetworkType
+import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileConnectivityModel
+import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileSubscriptionModel
+import com.android.systemui.statusbar.pipeline.mobile.data.model.OverrideNetworkType
+import com.android.systemui.statusbar.pipeline.mobile.data.model.ResolvedNetworkType
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel.Mobile
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel.MobileDisabled
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.filterNotNull
+import kotlinx.coroutines.flow.mapLatest
+import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.launch
+
+/** This repository vends out data based on demo mode commands */
+@OptIn(ExperimentalCoroutinesApi::class)
+class DemoMobileConnectionsRepository
+@Inject
+constructor(
+    private val dataSource: DemoModeMobileConnectionDataSource,
+    @Application private val scope: CoroutineScope,
+    context: Context,
+) : MobileConnectionsRepository {
+
+    private var demoCommandJob: Job? = null
+
+    private val connectionRepoCache = mutableMapOf<Int, DemoMobileConnectionRepository>()
+    private val subscriptionInfoCache = mutableMapOf<Int, SubscriptionInfo>()
+    val demoModeFinishedEvent = MutableSharedFlow<Unit>(extraBufferCapacity = 1)
+
+    private val _subscriptions = MutableStateFlow<List<SubscriptionInfo>>(listOf())
+    override val subscriptionsFlow =
+        _subscriptions
+            .onEach { infos -> dropUnusedReposFromCache(infos) }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), _subscriptions.value)
+
+    private fun dropUnusedReposFromCache(newInfos: List<SubscriptionInfo>) {
+        // Remove any connection repository from the cache that isn't in the new set of IDs. They
+        // will get garbage collected once their subscribers go away
+        val currentValidSubscriptionIds = newInfos.map { it.subscriptionId }
+
+        connectionRepoCache.keys.forEach {
+            if (!currentValidSubscriptionIds.contains(it)) {
+                connectionRepoCache.remove(it)
+            }
+        }
+    }
+
+    private fun maybeCreateSubscription(subId: Int) {
+        if (!subscriptionInfoCache.containsKey(subId)) {
+            createSubscriptionForSubId(subId, subId).also { subscriptionInfoCache[subId] = it }
+
+            _subscriptions.value = subscriptionInfoCache.values.toList()
+        }
+    }
+
+    /** Mimics the old NetworkControllerImpl for now */
+    private fun createSubscriptionForSubId(subId: Int, slotIndex: Int): SubscriptionInfo {
+        return SubscriptionInfo(
+            subId,
+            "",
+            slotIndex,
+            "",
+            "",
+            0,
+            0,
+            "",
+            0,
+            null,
+            null,
+            null,
+            "",
+            false,
+            null,
+            null,
+        )
+    }
+
+    // TODO(b/261029387): add a command for this value
+    override val activeMobileDataSubscriptionId =
+        subscriptionsFlow
+            .mapLatest { infos ->
+                // For now, active is just the first in the list
+                infos.firstOrNull()?.subscriptionId ?: INVALID_SUBSCRIPTION_ID
+            }
+            .stateIn(
+                scope,
+                SharingStarted.WhileSubscribed(),
+                subscriptionsFlow.value.firstOrNull()?.subscriptionId ?: INVALID_SUBSCRIPTION_ID
+            )
+
+    /** Demo mode doesn't currently support modifications to the mobile mappings */
+    override val defaultDataSubRatConfig =
+        MutableStateFlow(MobileMappings.Config.readConfig(context))
+
+    // TODO(b/261029387): add a command for this value
+    override val defaultDataSubId =
+        activeMobileDataSubscriptionId.stateIn(
+            scope,
+            SharingStarted.WhileSubscribed(),
+            INVALID_SUBSCRIPTION_ID
+        )
+
+    // TODO(b/261029387): not yet supported
+    override val defaultMobileNetworkConnectivity = MutableStateFlow(MobileConnectivityModel())
+
+    override fun getRepoForSubId(subId: Int): DemoMobileConnectionRepository {
+        return connectionRepoCache[subId]
+            ?: DemoMobileConnectionRepository(subId).also { connectionRepoCache[subId] = it }
+    }
+
+    override val globalMobileDataSettingChangedEvent = MutableStateFlow(Unit)
+
+    fun startProcessingCommands() {
+        demoCommandJob =
+            scope.launch {
+                dataSource.mobileEvents.filterNotNull().collect { event -> processEvent(event) }
+            }
+    }
+
+    fun stopProcessingCommands() {
+        demoCommandJob?.cancel()
+        _subscriptions.value = listOf()
+        connectionRepoCache.clear()
+        subscriptionInfoCache.clear()
+    }
+
+    private fun processEvent(event: FakeNetworkEventModel) {
+        when (event) {
+            is Mobile -> {
+                processEnabledMobileState(event)
+            }
+            is MobileDisabled -> {
+                processDisabledMobileState(event)
+            }
+        }
+    }
+
+    private fun processEnabledMobileState(state: Mobile) {
+        // get or create the connection repo, and set its values
+        val subId = state.subId ?: DEFAULT_SUB_ID
+        maybeCreateSubscription(subId)
+
+        val connection = getRepoForSubId(subId)
+        // This is always true here, because we split out disabled states at the data-source level
+        connection.dataEnabled.value = true
+        connection.isDefaultDataSubscription.value = state.dataType != null
+
+        connection.subscriptionModelFlow.value = state.toMobileSubscriptionModel()
+    }
+
+    private fun processDisabledMobileState(state: MobileDisabled) {
+        if (_subscriptions.value.isEmpty()) {
+            // Nothing to do here
+            return
+        }
+
+        val subId =
+            state.subId
+                ?: run {
+                    // For sake of usability, we can allow for no subId arg if there is only one
+                    // subscription
+                    if (_subscriptions.value.size > 1) {
+                        Log.d(
+                            TAG,
+                            "processDisabledMobileState: Unable to infer subscription to " +
+                                "disable. Specify subId using '-e slot <subId>'" +
+                                "Known subIds: [${subIdsString()}]"
+                        )
+                        return
+                    }
+
+                    // Use the only existing subscription as our arg, since there is only one
+                    _subscriptions.value[0].subscriptionId
+                }
+
+        removeSubscription(subId)
+    }
+
+    private fun removeSubscription(subId: Int) {
+        val currentSubscriptions = _subscriptions.value
+        subscriptionInfoCache.remove(subId)
+        _subscriptions.value = currentSubscriptions.filter { it.subscriptionId != subId }
+    }
+
+    private fun subIdsString(): String =
+        _subscriptions.value.joinToString(",") { it.subscriptionId.toString() }
+
+    companion object {
+        private const val TAG = "DemoMobileConnectionsRepo"
+
+        private const val DEFAULT_SUB_ID = 1
+    }
+}
+
+private fun Mobile.toMobileSubscriptionModel(): MobileSubscriptionModel {
+    return MobileSubscriptionModel(
+        isEmergencyOnly = false, // TODO(b/261029387): not yet supported
+        isGsm = false, // TODO(b/261029387): not yet supported
+        cdmaLevel = level ?: 0,
+        primaryLevel = level ?: 0,
+        dataConnectionState = DataConnectionState.Connected, // TODO(b/261029387): not yet supported
+        dataActivityDirection = activity,
+        carrierNetworkChangeActive = carrierNetworkChange,
+        // TODO(b/261185097): once mobile mappings can be mocked at this layer, we can build our
+        //  own demo map
+        resolvedNetworkType = dataType.toResolvedNetworkType()
+    )
+}
+
+@Annotation.NetworkType
+private fun SignalIcon.MobileIconGroup?.toNetworkType(): Int =
+    when (this) {
+        TelephonyIcons.THREE_G -> NETWORK_TYPE_GSM
+        TelephonyIcons.LTE -> NETWORK_TYPE_LTE
+        TelephonyIcons.FOUR_G -> NETWORK_TYPE_UMTS
+        TelephonyIcons.NR_5G -> NETWORK_TYPE_NR
+        TelephonyIcons.NR_5G_PLUS -> OVERRIDE_NETWORK_TYPE_NR_ADVANCED
+        else -> NETWORK_TYPE_UNKNOWN
+    }
+
+private fun SignalIcon.MobileIconGroup?.toResolvedNetworkType(): ResolvedNetworkType =
+    when (this) {
+        TelephonyIcons.NR_5G_PLUS -> OverrideNetworkType(toNetworkType())
+        else -> DefaultNetworkType(toNetworkType())
+    }
+
+class DemoMobileConnectionRepository(val subId: Int) : MobileConnectionRepository {
+    override val subscriptionModelFlow = MutableStateFlow(MobileSubscriptionModel())
+
+    override val dataEnabled = MutableStateFlow(true)
+
+    override val isDefaultDataSubscription = MutableStateFlow(true)
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoModeMobileConnectionDataSource.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoModeMobileConnectionDataSource.kt
new file mode 100644
index 0000000..da55787
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoModeMobileConnectionDataSource.kt
@@ -0,0 +1,141 @@
+/*
+ * Copyright (C) 2022 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.systemui.statusbar.pipeline.mobile.data.repository.demo
+
+import android.os.Bundle
+import android.telephony.Annotation.DataActivityType
+import android.telephony.TelephonyManager.DATA_ACTIVITY_IN
+import android.telephony.TelephonyManager.DATA_ACTIVITY_INOUT
+import android.telephony.TelephonyManager.DATA_ACTIVITY_NONE
+import android.telephony.TelephonyManager.DATA_ACTIVITY_OUT
+import com.android.settingslib.SignalIcon.MobileIconGroup
+import com.android.settingslib.mobile.TelephonyIcons
+import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.demomode.DemoMode
+import com.android.systemui.demomode.DemoMode.COMMAND_NETWORK
+import com.android.systemui.demomode.DemoModeController
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel.Mobile
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel.MobileDisabled
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.shareIn
+
+/**
+ * Data source that can map from demo mode commands to inputs into the
+ * [DemoMobileConnectionsRepository]'s flows
+ */
+@SysUISingleton
+class DemoModeMobileConnectionDataSource
+@Inject
+constructor(
+    demoModeController: DemoModeController,
+    @Application scope: CoroutineScope,
+) {
+    private val demoCommandStream: Flow<Bundle> = conflatedCallbackFlow {
+        val callback =
+            object : DemoMode {
+                override fun demoCommands(): List<String> = listOf(COMMAND_NETWORK)
+
+                override fun dispatchDemoCommand(command: String, args: Bundle) {
+                    trySend(args)
+                }
+
+                override fun onDemoModeFinished() {
+                    // Handled elsewhere
+                }
+
+                override fun onDemoModeStarted() {
+                    // Handled elsewhere
+                }
+            }
+
+        demoModeController.addCallback(callback)
+        awaitClose { demoModeController.removeCallback(callback) }
+    }
+
+    // If the args contains "mobile", then all of the args are relevant. It's just the way demo mode
+    // commands work and it's a little silly
+    private val _mobileCommands = demoCommandStream.map { args -> args.toMobileEvent() }
+    val mobileEvents = _mobileCommands.shareIn(scope, SharingStarted.WhileSubscribed())
+
+    private fun Bundle.toMobileEvent(): FakeNetworkEventModel? {
+        val mobile = getString("mobile") ?: return null
+        return if (mobile == "show") {
+            activeMobileEvent()
+        } else {
+            MobileDisabled(subId = getString("slot")?.toInt())
+        }
+    }
+
+    /** Parse a valid mobile command string into a network event */
+    private fun Bundle.activeMobileEvent(): Mobile {
+        // There are many key/value pairs supported by mobile demo mode. Bear with me here
+        val level = getString("level")?.toInt()
+        val dataType = getString("datatype")?.toDataType()
+        val slot = getString("slot")?.toInt()
+        val carrierId = getString("carrierid")?.toInt()
+        val inflateStrength = getString("inflate")?.toBoolean()
+        val activity = getString("activity")?.toActivity()
+        val carrierNetworkChange = getString("carriernetworkchange") == "show"
+
+        return Mobile(
+            level = level,
+            dataType = dataType,
+            subId = slot,
+            carrierId = carrierId,
+            inflateStrength = inflateStrength,
+            activity = activity,
+            carrierNetworkChange = carrierNetworkChange,
+        )
+    }
+}
+
+private fun String.toDataType(): MobileIconGroup =
+    when (this) {
+        "1x" -> TelephonyIcons.ONE_X
+        "3g" -> TelephonyIcons.THREE_G
+        "4g" -> TelephonyIcons.FOUR_G
+        "4g+" -> TelephonyIcons.FOUR_G_PLUS
+        "5g" -> TelephonyIcons.NR_5G
+        "5ge" -> TelephonyIcons.LTE_CA_5G_E
+        "5g+" -> TelephonyIcons.NR_5G_PLUS
+        "e" -> TelephonyIcons.E
+        "g" -> TelephonyIcons.G
+        "h" -> TelephonyIcons.H
+        "h+" -> TelephonyIcons.H_PLUS
+        "lte" -> TelephonyIcons.LTE
+        "lte+" -> TelephonyIcons.LTE_PLUS
+        "dis" -> TelephonyIcons.DATA_DISABLED
+        "not" -> TelephonyIcons.NOT_DEFAULT_DATA
+        else -> TelephonyIcons.UNKNOWN
+    }
+
+@DataActivityType
+private fun String.toActivity(): Int =
+    when (this) {
+        "inout" -> DATA_ACTIVITY_INOUT
+        "in" -> DATA_ACTIVITY_IN
+        "out" -> DATA_ACTIVITY_OUT
+        else -> DATA_ACTIVITY_NONE
+    }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/model/FakeNetworkEventModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/model/FakeNetworkEventModel.kt
new file mode 100644
index 0000000..3f3acaf
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/model/FakeNetworkEventModel.kt
@@ -0,0 +1,43 @@
+/*
+ * Copyright (C) 2022 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.systemui.statusbar.pipeline.mobile.data.repository.demo.model
+
+import android.telephony.Annotation.DataActivityType
+import com.android.settingslib.SignalIcon
+
+/**
+ * Model for the demo commands, ported from [NetworkControllerImpl]
+ *
+ * Nullable fields represent optional command line arguments
+ */
+sealed interface FakeNetworkEventModel {
+    data class Mobile(
+        val level: Int?,
+        val dataType: SignalIcon.MobileIconGroup?,
+        // Null means the default (chosen by the repository)
+        val subId: Int?,
+        val carrierId: Int?,
+        val inflateStrength: Boolean?,
+        @DataActivityType val activity: Int?,
+        val carrierNetworkChange: Boolean,
+    ) : FakeNetworkEventModel
+
+    data class MobileDisabled(
+        // Null means the default (chosen by the repository)
+        val subId: Int?
+    ) : FakeNetworkEventModel
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionRepositoryImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionRepositoryImpl.kt
new file mode 100644
index 0000000..4c1cf4a
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionRepositoryImpl.kt
@@ -0,0 +1,235 @@
+/*
+ * Copyright (C) 2022 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.systemui.statusbar.pipeline.mobile.data.repository.prod
+
+import android.content.Context
+import android.database.ContentObserver
+import android.provider.Settings.Global
+import android.telephony.CellSignalStrength
+import android.telephony.CellSignalStrengthCdma
+import android.telephony.ServiceState
+import android.telephony.SignalStrength
+import android.telephony.TelephonyCallback
+import android.telephony.TelephonyDisplayInfo
+import android.telephony.TelephonyDisplayInfo.OVERRIDE_NETWORK_TYPE_NONE
+import android.telephony.TelephonyManager
+import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.statusbar.pipeline.mobile.data.model.DefaultNetworkType
+import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileSubscriptionModel
+import com.android.systemui.statusbar.pipeline.mobile.data.model.OverrideNetworkType
+import com.android.systemui.statusbar.pipeline.mobile.data.model.toDataConnectionType
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionRepository
+import com.android.systemui.statusbar.pipeline.shared.ConnectivityPipelineLogger
+import com.android.systemui.statusbar.pipeline.shared.ConnectivityPipelineLogger.Companion.logOutputChange
+import com.android.systemui.util.settings.GlobalSettings
+import java.lang.IllegalStateException
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.asExecutor
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.mapLatest
+import kotlinx.coroutines.flow.merge
+import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.flow.stateIn
+
+@Suppress("EXPERIMENTAL_IS_NOT_ENABLED")
+@OptIn(ExperimentalCoroutinesApi::class)
+class MobileConnectionRepositoryImpl(
+    private val context: Context,
+    private val subId: Int,
+    private val telephonyManager: TelephonyManager,
+    private val globalSettings: GlobalSettings,
+    defaultDataSubId: StateFlow<Int>,
+    globalMobileDataSettingChangedEvent: Flow<Unit>,
+    bgDispatcher: CoroutineDispatcher,
+    logger: ConnectivityPipelineLogger,
+    scope: CoroutineScope,
+) : MobileConnectionRepository {
+    init {
+        if (telephonyManager.subscriptionId != subId) {
+            throw IllegalStateException(
+                "TelephonyManager should be created with subId($subId). " +
+                    "Found ${telephonyManager.subscriptionId} instead."
+            )
+        }
+    }
+
+    private val telephonyCallbackEvent = MutableSharedFlow<Unit>(extraBufferCapacity = 1)
+
+    override val subscriptionModelFlow: StateFlow<MobileSubscriptionModel> = run {
+        var state = MobileSubscriptionModel()
+        conflatedCallbackFlow {
+                // TODO (b/240569788): log all of these into the connectivity logger
+                val callback =
+                    object :
+                        TelephonyCallback(),
+                        TelephonyCallback.ServiceStateListener,
+                        TelephonyCallback.SignalStrengthsListener,
+                        TelephonyCallback.DataConnectionStateListener,
+                        TelephonyCallback.DataActivityListener,
+                        TelephonyCallback.CarrierNetworkListener,
+                        TelephonyCallback.DisplayInfoListener {
+                        override fun onServiceStateChanged(serviceState: ServiceState) {
+                            state = state.copy(isEmergencyOnly = serviceState.isEmergencyOnly)
+                            trySend(state)
+                        }
+
+                        override fun onSignalStrengthsChanged(signalStrength: SignalStrength) {
+                            val cdmaLevel =
+                                signalStrength
+                                    .getCellSignalStrengths(CellSignalStrengthCdma::class.java)
+                                    .let { strengths ->
+                                        if (!strengths.isEmpty()) {
+                                            strengths[0].level
+                                        } else {
+                                            CellSignalStrength.SIGNAL_STRENGTH_NONE_OR_UNKNOWN
+                                        }
+                                    }
+
+                            val primaryLevel = signalStrength.level
+
+                            state =
+                                state.copy(
+                                    cdmaLevel = cdmaLevel,
+                                    primaryLevel = primaryLevel,
+                                    isGsm = signalStrength.isGsm,
+                                )
+                            trySend(state)
+                        }
+
+                        override fun onDataConnectionStateChanged(
+                            dataState: Int,
+                            networkType: Int
+                        ) {
+                            state =
+                                state.copy(dataConnectionState = dataState.toDataConnectionType())
+                            trySend(state)
+                        }
+
+                        override fun onDataActivity(direction: Int) {
+                            state = state.copy(dataActivityDirection = direction)
+                            trySend(state)
+                        }
+
+                        override fun onCarrierNetworkChange(active: Boolean) {
+                            state = state.copy(carrierNetworkChangeActive = active)
+                            trySend(state)
+                        }
+
+                        override fun onDisplayInfoChanged(
+                            telephonyDisplayInfo: TelephonyDisplayInfo
+                        ) {
+                            val networkType =
+                                if (
+                                    telephonyDisplayInfo.overrideNetworkType ==
+                                        OVERRIDE_NETWORK_TYPE_NONE
+                                ) {
+                                    DefaultNetworkType(telephonyDisplayInfo.networkType)
+                                } else {
+                                    OverrideNetworkType(telephonyDisplayInfo.overrideNetworkType)
+                                }
+                            state = state.copy(resolvedNetworkType = networkType)
+                            trySend(state)
+                        }
+                    }
+                telephonyManager.registerTelephonyCallback(bgDispatcher.asExecutor(), callback)
+                awaitClose { telephonyManager.unregisterTelephonyCallback(callback) }
+            }
+            .onEach { telephonyCallbackEvent.tryEmit(Unit) }
+            .logOutputChange(logger, "MobileSubscriptionModel")
+            .stateIn(scope, SharingStarted.WhileSubscribed(), state)
+    }
+
+    /** Produces whenever the mobile data setting changes for this subId */
+    private val localMobileDataSettingChangedEvent: Flow<Unit> = conflatedCallbackFlow {
+        val observer =
+            object : ContentObserver(null) {
+                override fun onChange(selfChange: Boolean) {
+                    trySend(Unit)
+                }
+            }
+
+        globalSettings.registerContentObserver(
+            globalSettings.getUriFor("${Global.MOBILE_DATA}$subId"),
+            /* notifyForDescendants */ true,
+            observer
+        )
+
+        awaitClose { context.contentResolver.unregisterContentObserver(observer) }
+    }
+
+    /**
+     * There are a few cases where we will need to poll [TelephonyManager] so we can update some
+     * internal state where callbacks aren't provided. Any of those events should be merged into
+     * this flow, which can be used to trigger the polling.
+     */
+    private val telephonyPollingEvent: Flow<Unit> =
+        merge(
+            telephonyCallbackEvent,
+            localMobileDataSettingChangedEvent,
+            globalMobileDataSettingChangedEvent,
+        )
+
+    override val dataEnabled: StateFlow<Boolean> =
+        telephonyPollingEvent
+            .mapLatest { dataConnectionAllowed() }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), dataConnectionAllowed())
+
+    private fun dataConnectionAllowed(): Boolean = telephonyManager.isDataConnectionAllowed
+
+    override val isDefaultDataSubscription: StateFlow<Boolean> =
+        defaultDataSubId
+            .mapLatest { it == subId }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), defaultDataSubId.value == subId)
+
+    class Factory
+    @Inject
+    constructor(
+        private val context: Context,
+        private val telephonyManager: TelephonyManager,
+        private val logger: ConnectivityPipelineLogger,
+        private val globalSettings: GlobalSettings,
+        @Background private val bgDispatcher: CoroutineDispatcher,
+        @Application private val scope: CoroutineScope,
+    ) {
+        fun build(
+            subId: Int,
+            defaultDataSubId: StateFlow<Int>,
+            globalMobileDataSettingChangedEvent: Flow<Unit>,
+        ): MobileConnectionRepository {
+            return MobileConnectionRepositoryImpl(
+                context,
+                subId,
+                telephonyManager.createForSubscriptionId(subId),
+                globalSettings,
+                defaultDataSubId,
+                globalMobileDataSettingChangedEvent,
+                bgDispatcher,
+                logger,
+                scope,
+            )
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryImpl.kt
new file mode 100644
index 0000000..08d6010
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryImpl.kt
@@ -0,0 +1,265 @@
+/*
+ * Copyright (C) 2022 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.systemui.statusbar.pipeline.mobile.data.repository.prod
+
+import android.annotation.SuppressLint
+import android.content.Context
+import android.content.IntentFilter
+import android.database.ContentObserver
+import android.net.ConnectivityManager
+import android.net.ConnectivityManager.NetworkCallback
+import android.net.Network
+import android.net.NetworkCapabilities
+import android.net.NetworkCapabilities.NET_CAPABILITY_VALIDATED
+import android.net.NetworkCapabilities.TRANSPORT_CELLULAR
+import android.provider.Settings.Global.MOBILE_DATA
+import android.telephony.CarrierConfigManager
+import android.telephony.SubscriptionInfo
+import android.telephony.SubscriptionManager
+import android.telephony.SubscriptionManager.INVALID_SUBSCRIPTION_ID
+import android.telephony.TelephonyCallback
+import android.telephony.TelephonyCallback.ActiveDataSubscriptionIdListener
+import android.telephony.TelephonyManager
+import androidx.annotation.VisibleForTesting
+import com.android.internal.telephony.PhoneConstants
+import com.android.settingslib.mobile.MobileMappings
+import com.android.settingslib.mobile.MobileMappings.Config
+import com.android.systemui.broadcast.BroadcastDispatcher
+import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileConnectivityModel
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.shared.ConnectivityPipelineLogger
+import com.android.systemui.util.settings.GlobalSettings
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.asExecutor
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.mapLatest
+import kotlinx.coroutines.flow.merge
+import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.withContext
+
+@Suppress("EXPERIMENTAL_IS_NOT_ENABLED")
+@OptIn(ExperimentalCoroutinesApi::class)
+@SysUISingleton
+class MobileConnectionsRepositoryImpl
+@Inject
+constructor(
+    private val connectivityManager: ConnectivityManager,
+    private val subscriptionManager: SubscriptionManager,
+    private val telephonyManager: TelephonyManager,
+    private val logger: ConnectivityPipelineLogger,
+    broadcastDispatcher: BroadcastDispatcher,
+    private val globalSettings: GlobalSettings,
+    private val context: Context,
+    @Background private val bgDispatcher: CoroutineDispatcher,
+    @Application private val scope: CoroutineScope,
+    private val mobileConnectionRepositoryFactory: MobileConnectionRepositoryImpl.Factory
+) : MobileConnectionsRepository {
+    private val subIdRepositoryCache: MutableMap<Int, MobileConnectionRepository> = mutableMapOf()
+
+    /**
+     * State flow that emits the set of mobile data subscriptions, each represented by its own
+     * [SubscriptionInfo]. We probably only need the [SubscriptionInfo.getSubscriptionId] of each
+     * info object, but for now we keep track of the infos themselves.
+     */
+    override val subscriptionsFlow: StateFlow<List<SubscriptionInfo>> =
+        conflatedCallbackFlow {
+                val callback =
+                    object : SubscriptionManager.OnSubscriptionsChangedListener() {
+                        override fun onSubscriptionsChanged() {
+                            trySend(Unit)
+                        }
+                    }
+
+                subscriptionManager.addOnSubscriptionsChangedListener(
+                    bgDispatcher.asExecutor(),
+                    callback,
+                )
+
+                awaitClose { subscriptionManager.removeOnSubscriptionsChangedListener(callback) }
+            }
+            .mapLatest { fetchSubscriptionsList() }
+            .onEach { infos -> dropUnusedReposFromCache(infos) }
+            .stateIn(scope, started = SharingStarted.WhileSubscribed(), listOf())
+
+    /** StateFlow that keeps track of the current active mobile data subscription */
+    override val activeMobileDataSubscriptionId: StateFlow<Int> =
+        conflatedCallbackFlow {
+                val callback =
+                    object : TelephonyCallback(), ActiveDataSubscriptionIdListener {
+                        override fun onActiveDataSubscriptionIdChanged(subId: Int) {
+                            trySend(subId)
+                        }
+                    }
+
+                telephonyManager.registerTelephonyCallback(bgDispatcher.asExecutor(), callback)
+                awaitClose { telephonyManager.unregisterTelephonyCallback(callback) }
+            }
+            .stateIn(scope, started = SharingStarted.WhileSubscribed(), INVALID_SUBSCRIPTION_ID)
+
+    private val defaultDataSubIdChangeEvent: MutableSharedFlow<Unit> =
+        MutableSharedFlow(extraBufferCapacity = 1)
+
+    override val defaultDataSubId: StateFlow<Int> =
+        broadcastDispatcher
+            .broadcastFlow(
+                IntentFilter(TelephonyManager.ACTION_DEFAULT_DATA_SUBSCRIPTION_CHANGED)
+            ) { intent, _ ->
+                intent.getIntExtra(PhoneConstants.SUBSCRIPTION_KEY, INVALID_SUBSCRIPTION_ID)
+            }
+            .distinctUntilChanged()
+            .onEach { defaultDataSubIdChangeEvent.tryEmit(Unit) }
+            .stateIn(
+                scope,
+                SharingStarted.WhileSubscribed(),
+                SubscriptionManager.getDefaultDataSubscriptionId()
+            )
+
+    private val carrierConfigChangedEvent =
+        broadcastDispatcher.broadcastFlow(
+            IntentFilter(CarrierConfigManager.ACTION_CARRIER_CONFIG_CHANGED)
+        )
+
+    /**
+     * [Config] is an object that tracks relevant configuration flags for a given subscription ID.
+     * In the case of [MobileMappings], it's hard-coded to check the default data subscription's
+     * config, so this will apply to every icon that we care about.
+     *
+     * Relevant bits in the config are things like
+     * [CarrierConfigManager.KEY_SHOW_4G_FOR_LTE_DATA_ICON_BOOL]
+     *
+     * This flow will produce whenever the default data subscription or the carrier config changes.
+     */
+    override val defaultDataSubRatConfig: StateFlow<Config> =
+        merge(defaultDataSubIdChangeEvent, carrierConfigChangedEvent)
+            .mapLatest { Config.readConfig(context) }
+            .stateIn(
+                scope,
+                SharingStarted.WhileSubscribed(),
+                initialValue = Config.readConfig(context)
+            )
+
+    override fun getRepoForSubId(subId: Int): MobileConnectionRepository {
+        if (!isValidSubId(subId)) {
+            throw IllegalArgumentException(
+                "subscriptionId $subId is not in the list of valid subscriptions"
+            )
+        }
+
+        return subIdRepositoryCache[subId]
+            ?: createRepositoryForSubId(subId).also { subIdRepositoryCache[subId] = it }
+    }
+
+    /**
+     * In single-SIM devices, the [MOBILE_DATA] setting is phone-wide. For multi-SIM, the individual
+     * connection repositories also observe the URI for [MOBILE_DATA] + subId.
+     */
+    override val globalMobileDataSettingChangedEvent: Flow<Unit> = conflatedCallbackFlow {
+        val observer =
+            object : ContentObserver(null) {
+                override fun onChange(selfChange: Boolean) {
+                    trySend(Unit)
+                }
+            }
+
+        globalSettings.registerContentObserver(
+            globalSettings.getUriFor(MOBILE_DATA),
+            true,
+            observer
+        )
+
+        awaitClose { context.contentResolver.unregisterContentObserver(observer) }
+    }
+
+    @SuppressLint("MissingPermission")
+    override val defaultMobileNetworkConnectivity: StateFlow<MobileConnectivityModel> =
+        conflatedCallbackFlow {
+                val callback =
+                    object : NetworkCallback(FLAG_INCLUDE_LOCATION_INFO) {
+                        override fun onLost(network: Network) {
+                            // Send a disconnected model when lost. Maybe should create a sealed
+                            // type or null here?
+                            trySend(MobileConnectivityModel())
+                        }
+
+                        override fun onCapabilitiesChanged(
+                            network: Network,
+                            caps: NetworkCapabilities
+                        ) {
+                            trySend(
+                                MobileConnectivityModel(
+                                    isConnected = caps.hasTransport(TRANSPORT_CELLULAR),
+                                    isValidated = caps.hasCapability(NET_CAPABILITY_VALIDATED),
+                                )
+                            )
+                        }
+                    }
+
+                connectivityManager.registerDefaultNetworkCallback(callback)
+
+                awaitClose { connectivityManager.unregisterNetworkCallback(callback) }
+            }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), MobileConnectivityModel())
+
+    private fun isValidSubId(subId: Int): Boolean {
+        subscriptionsFlow.value.forEach {
+            if (it.subscriptionId == subId) {
+                return true
+            }
+        }
+
+        return false
+    }
+
+    @VisibleForTesting fun getSubIdRepoCache() = subIdRepositoryCache
+
+    private fun createRepositoryForSubId(subId: Int): MobileConnectionRepository {
+        return mobileConnectionRepositoryFactory.build(
+            subId,
+            defaultDataSubId,
+            globalMobileDataSettingChangedEvent,
+        )
+    }
+
+    private fun dropUnusedReposFromCache(newInfos: List<SubscriptionInfo>) {
+        // Remove any connection repository from the cache that isn't in the new set of IDs. They
+        // will get garbage collected once their subscribers go away
+        val currentValidSubscriptionIds = newInfos.map { it.subscriptionId }
+
+        subIdRepositoryCache.keys.forEach {
+            if (!currentValidSubscriptionIds.contains(it)) {
+                subIdRepositoryCache.remove(it)
+            }
+        }
+    }
+
+    private suspend fun fetchSubscriptionsList(): List<SubscriptionInfo> =
+        withContext(bgDispatcher) { subscriptionManager.completeActiveSubscriptionInfoList }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt
index 162c915..b2ec27c 100644
--- a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt
+++ b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt
@@ -40,6 +40,8 @@
 import com.android.systemui.statusbar.LightRevealEffect
 import com.android.systemui.statusbar.LightRevealScrim
 import com.android.systemui.statusbar.LinearLightRevealEffect
+import com.android.systemui.unfold.UnfoldLightRevealOverlayAnimation.AddOverlayReason.FOLD
+import com.android.systemui.unfold.UnfoldLightRevealOverlayAnimation.AddOverlayReason.UNFOLD
 import com.android.systemui.unfold.UnfoldTransitionProgressProvider.TransitionProgressListener
 import com.android.systemui.unfold.updates.RotationChangeProvider
 import com.android.systemui.unfold.util.ScaleAwareTransitionProgressProvider.Companion.areAnimationsEnabled
@@ -125,7 +127,7 @@
         try {
             // Add the view only if we are unfolding and this is the first screen on
             if (!isFolded && !isUnfoldHandled && contentResolver.areAnimationsEnabled()) {
-                executeInBackground { addView(onOverlayReady) }
+                executeInBackground { addOverlay(onOverlayReady, reason = UNFOLD) }
                 isUnfoldHandled = true
             } else {
                 // No unfold transition, immediately report that overlay is ready
@@ -137,7 +139,7 @@
         }
     }
 
-    private fun addView(onOverlayReady: Runnable? = null) {
+    private fun addOverlay(onOverlayReady: Runnable? = null, reason: AddOverlayReason) {
         if (!::wwm.isInitialized) {
             // Surface overlay is not created yet on the first SysUI launch
             onOverlayReady?.run()
@@ -152,7 +154,10 @@
             LightRevealScrim(context, null).apply {
                 revealEffect = createLightRevealEffect()
                 isScrimOpaqueChangedListener = Consumer {}
-                revealAmount = 0f
+                revealAmount = when (reason) {
+                    FOLD -> TRANSPARENT
+                    UNFOLD -> BLACK
+                }
             }
 
         val params = getLayoutParams()
@@ -228,7 +233,7 @@
     }
 
     private fun getUnfoldedDisplayInfo(): DisplayInfo =
-        displayManager.displays
+        displayManager.getDisplays(DisplayManager.DISPLAY_CATEGORY_ALL_INCLUDING_DISABLED)
             .asSequence()
             .map { DisplayInfo().apply { it.getDisplayInfo(this) } }
             .filter { it.type == Display.TYPE_INTERNAL }
@@ -247,7 +252,7 @@
         override fun onTransitionStarted() {
             // Add view for folding case (when unfolding the view is added earlier)
             if (scrimView == null) {
-                executeInBackground { addView() }
+                executeInBackground { addOverlay(reason = FOLD) }
             }
             // Disable input dispatching during transition.
             InputManager.getInstance().cancelCurrentTouch()
@@ -294,11 +299,17 @@
             }
         )
 
+    private enum class AddOverlayReason { FOLD, UNFOLD }
+
     private companion object {
-        private const val ROTATION_ANIMATION_OVERLAY_Z_INDEX = Integer.MAX_VALUE
+        const val ROTATION_ANIMATION_OVERLAY_Z_INDEX = Integer.MAX_VALUE
 
         // Put the unfold overlay below the rotation animation screenshot to hide the moment
         // when it is rotated but the rotation of the other windows hasn't happen yet
-        private const val UNFOLD_OVERLAY_LAYER_Z_INDEX = ROTATION_ANIMATION_OVERLAY_Z_INDEX - 1
+        const val UNFOLD_OVERLAY_LAYER_Z_INDEX = ROTATION_ANIMATION_OVERLAY_Z_INDEX - 1
+
+        // constants for revealAmount.
+        const val TRANSPARENT = 1f
+        const val BLACK = 0f
     }
 }
diff --git a/packages/SystemUI/animation/tests/com/android/systemui/animation/InterpolatorsAndroidXTest.kt b/packages/SystemUI/tests/src/com/android/systemui/animation/InterpolatorsAndroidXTest.kt
similarity index 94%
rename from packages/SystemUI/animation/tests/com/android/systemui/animation/InterpolatorsAndroidXTest.kt
rename to packages/SystemUI/tests/src/com/android/systemui/animation/InterpolatorsAndroidXTest.kt
index 389eed0..2c680be 100644
--- a/packages/SystemUI/animation/tests/com/android/systemui/animation/InterpolatorsAndroidXTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/animation/InterpolatorsAndroidXTest.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.animation
 
 import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
 import java.lang.reflect.Modifier
 import junit.framework.Assert.assertEquals
 import org.junit.Test
@@ -25,7 +26,7 @@
 
 @SmallTest
 @RunWith(JUnit4::class)
-class InterpolatorsAndroidXTest {
+class InterpolatorsAndroidXTest : SysuiTestCase() {
 
     @Test
     fun testInterpolatorsAndInterpolatorsAndroidXPublicMethodsAreEqual() {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/controls/ui/ControlsUiControllerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/controls/ui/ControlsUiControllerImplTest.kt
index e679b13..d965e33 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/controls/ui/ControlsUiControllerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/controls/ui/ControlsUiControllerImplTest.kt
@@ -16,15 +16,26 @@
 
 package com.android.systemui.controls.ui
 
+import android.app.PendingIntent
 import android.content.ComponentName
 import android.content.Context
+import android.content.pm.ApplicationInfo
+import android.content.pm.ServiceInfo
+import android.os.UserHandle
+import android.service.controls.ControlsProviderService
 import android.testing.AndroidTestingRunner
 import android.testing.TestableLooper
+import android.util.AttributeSet
+import android.view.LayoutInflater
+import android.view.View
 import android.widget.FrameLayout
 import androidx.test.filters.SmallTest
+import com.android.systemui.R
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.controls.ControlsMetricsLogger
+import com.android.systemui.controls.ControlsServiceInfo
 import com.android.systemui.controls.CustomIconCache
+import com.android.systemui.controls.FakeControlsSettingsRepository
 import com.android.systemui.controls.controller.ControlsController
 import com.android.systemui.controls.controller.StructureInfo
 import com.android.systemui.controls.management.ControlsListingController
@@ -38,19 +49,26 @@
 import com.android.systemui.util.FakeSharedPreferences
 import com.android.systemui.util.concurrency.FakeExecutor
 import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.argumentCaptor
+import com.android.systemui.util.mockito.capture
+import com.android.systemui.util.mockito.eq
+import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.time.FakeSystemClock
+import com.android.wm.shell.TaskView
 import com.android.wm.shell.TaskViewFactory
 import com.google.common.truth.Truth.assertThat
 import dagger.Lazy
 import java.util.Optional
+import java.util.function.Consumer
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mock
 import org.mockito.Mockito.anyInt
 import org.mockito.Mockito.anyString
-import org.mockito.Mockito.mock
+import org.mockito.Mockito.clearInvocations
 import org.mockito.Mockito.never
+import org.mockito.Mockito.spy
 import org.mockito.Mockito.verify
 import org.mockito.Mockito.`when`
 import org.mockito.MockitoAnnotations
@@ -70,9 +88,9 @@
     @Mock lateinit var userFileManager: UserFileManager
     @Mock lateinit var userTracker: UserTracker
     @Mock lateinit var taskViewFactory: TaskViewFactory
-    @Mock lateinit var activityContext: Context
     @Mock lateinit var dumpManager: DumpManager
     val sharedPreferences = FakeSharedPreferences()
+    lateinit var controlsSettingsRepository: FakeControlsSettingsRepository
 
     var uiExecutor = FakeExecutor(FakeSystemClock())
     var bgExecutor = FakeExecutor(FakeSystemClock())
@@ -83,6 +101,17 @@
     fun setup() {
         MockitoAnnotations.initMocks(this)
 
+        controlsSettingsRepository = FakeControlsSettingsRepository()
+
+        // This way, it won't be cloned every time `LayoutInflater.fromContext` is called, but we
+        // need to clone it once so we don't modify the original one.
+        mContext.addMockSystemService(
+            Context.LAYOUT_INFLATER_SERVICE,
+            mContext.baseContext
+                .getSystemService(LayoutInflater::class.java)!!
+                .cloneInContext(mContext)
+        )
+
         parent = FrameLayout(mContext)
 
         underTest =
@@ -100,6 +129,7 @@
                 userFileManager,
                 userTracker,
                 Optional.of(taskViewFactory),
+                controlsSettingsRepository,
                 dumpManager
             )
         `when`(
@@ -113,11 +143,12 @@
         `when`(userFileManager.getSharedPreferences(anyString(), anyInt(), anyInt()))
             .thenReturn(sharedPreferences)
         `when`(userTracker.userId).thenReturn(0)
+        `when`(userTracker.userHandle).thenReturn(UserHandle.of(0))
     }
 
     @Test
     fun testGetPreferredStructure() {
-        val structureInfo = mock(StructureInfo::class.java)
+        val structureInfo = mock<StructureInfo>()
         underTest.getPreferredSelectedItem(listOf(structureInfo))
         verify(userFileManager)
             .getSharedPreferences(
@@ -189,14 +220,195 @@
     @Test
     fun testPanelDoesNotRefreshControls() {
         val panel = SelectedItem.PanelItem("App name", ComponentName("pkg", "cls"))
+        setUpPanel(panel)
+
+        underTest.show(parent, {}, context)
+        verify(controlsController, never()).refreshStatus(any(), any())
+    }
+
+    @Test
+    fun testPanelCallsTaskViewFactoryCreate() {
+        mockLayoutInflater()
+        val panel = SelectedItem.PanelItem("App name", ComponentName("pkg", "cls"))
+        val serviceInfo = setUpPanel(panel)
+
+        underTest.show(parent, {}, context)
+
+        val captor = argumentCaptor<ControlsListingController.ControlsListingCallback>()
+
+        verify(controlsListingController).addCallback(capture(captor))
+
+        captor.value.onServicesUpdated(listOf(serviceInfo))
+        FakeExecutor.exhaustExecutors(uiExecutor, bgExecutor)
+
+        verify(taskViewFactory).create(eq(context), eq(uiExecutor), any())
+    }
+
+    @Test
+    fun testPanelControllerStartActivityWithCorrectArguments() {
+        mockLayoutInflater()
+        controlsSettingsRepository.setAllowActionOnTrivialControlsInLockscreen(true)
+
+        val panel = SelectedItem.PanelItem("App name", ComponentName("pkg", "cls"))
+        val serviceInfo = setUpPanel(panel)
+
+        underTest.show(parent, {}, context)
+
+        val captor = argumentCaptor<ControlsListingController.ControlsListingCallback>()
+
+        verify(controlsListingController).addCallback(capture(captor))
+
+        captor.value.onServicesUpdated(listOf(serviceInfo))
+        FakeExecutor.exhaustExecutors(uiExecutor, bgExecutor)
+
+        val pendingIntent = verifyPanelCreatedAndStartTaskView()
+
+        with(pendingIntent) {
+            assertThat(isActivity).isTrue()
+            assertThat(intent.component).isEqualTo(serviceInfo.panelActivity)
+            assertThat(
+                    intent.getBooleanExtra(
+                        ControlsProviderService.EXTRA_LOCKSCREEN_ALLOW_TRIVIAL_CONTROLS,
+                        false
+                    )
+                )
+                .isTrue()
+        }
+    }
+
+    @Test
+    fun testPendingIntentExtrasAreModified() {
+        mockLayoutInflater()
+        controlsSettingsRepository.setAllowActionOnTrivialControlsInLockscreen(true)
+
+        val panel = SelectedItem.PanelItem("App name", ComponentName("pkg", "cls"))
+        val serviceInfo = setUpPanel(panel)
+
+        underTest.show(parent, {}, context)
+
+        val captor = argumentCaptor<ControlsListingController.ControlsListingCallback>()
+
+        verify(controlsListingController).addCallback(capture(captor))
+
+        captor.value.onServicesUpdated(listOf(serviceInfo))
+        FakeExecutor.exhaustExecutors(uiExecutor, bgExecutor)
+
+        val pendingIntent = verifyPanelCreatedAndStartTaskView()
+        assertThat(
+                pendingIntent.intent.getBooleanExtra(
+                    ControlsProviderService.EXTRA_LOCKSCREEN_ALLOW_TRIVIAL_CONTROLS,
+                    false
+                )
+            )
+            .isTrue()
+
+        underTest.hide()
+
+        clearInvocations(controlsListingController, taskViewFactory)
+        controlsSettingsRepository.setAllowActionOnTrivialControlsInLockscreen(false)
+        underTest.show(parent, {}, context)
+
+        verify(controlsListingController).addCallback(capture(captor))
+        captor.value.onServicesUpdated(listOf(serviceInfo))
+        FakeExecutor.exhaustExecutors(uiExecutor, bgExecutor)
+
+        val newPendingIntent = verifyPanelCreatedAndStartTaskView()
+        assertThat(
+                newPendingIntent.intent.getBooleanExtra(
+                    ControlsProviderService.EXTRA_LOCKSCREEN_ALLOW_TRIVIAL_CONTROLS,
+                    false
+                )
+            )
+            .isFalse()
+    }
+
+    private fun setUpPanel(panel: SelectedItem.PanelItem): ControlsServiceInfo {
+        val activity = ComponentName("pkg", "activity")
         sharedPreferences
             .edit()
             .putString("controls_component", panel.componentName.flattenToString())
             .putString("controls_structure", panel.appName.toString())
             .putBoolean("controls_is_panel", true)
             .commit()
+        return ControlsServiceInfo(panel.componentName, panel.appName, activity)
+    }
 
-        underTest.show(parent, {}, activityContext)
-        verify(controlsController, never()).refreshStatus(any(), any())
+    private fun verifyPanelCreatedAndStartTaskView(): PendingIntent {
+        val taskViewConsumerCaptor = argumentCaptor<Consumer<TaskView>>()
+        verify(taskViewFactory).create(eq(context), eq(uiExecutor), capture(taskViewConsumerCaptor))
+
+        val taskView: TaskView = mock {
+            `when`(this.post(any())).thenAnswer {
+                uiExecutor.execute(it.arguments[0] as Runnable)
+                true
+            }
+        }
+        // calls PanelTaskViewController#launchTaskView
+        taskViewConsumerCaptor.value.accept(taskView)
+        val listenerCaptor = argumentCaptor<TaskView.Listener>()
+        verify(taskView).setListener(any(), capture(listenerCaptor))
+        listenerCaptor.value.onInitialized()
+        FakeExecutor.exhaustExecutors(uiExecutor, bgExecutor)
+
+        val pendingIntentCaptor = argumentCaptor<PendingIntent>()
+        verify(taskView).startActivity(capture(pendingIntentCaptor), any(), any(), any())
+        return pendingIntentCaptor.value
+    }
+
+    private fun ControlsServiceInfo(
+        componentName: ComponentName,
+        label: CharSequence,
+        panelComponentName: ComponentName? = null
+    ): ControlsServiceInfo {
+        val serviceInfo =
+            ServiceInfo().apply {
+                applicationInfo = ApplicationInfo()
+                packageName = componentName.packageName
+                name = componentName.className
+            }
+        return spy(ControlsServiceInfo(mContext, serviceInfo)).apply {
+            `when`(loadLabel()).thenReturn(label)
+            `when`(loadIcon()).thenReturn(mock())
+            `when`(panelActivity).thenReturn(panelComponentName)
+        }
+    }
+
+    private fun mockLayoutInflater() {
+        LayoutInflater.from(context)
+            .setPrivateFactory(
+                object : LayoutInflater.Factory2 {
+                    override fun onCreateView(
+                        view: View?,
+                        name: String,
+                        context: Context,
+                        attrs: AttributeSet
+                    ): View? {
+                        return onCreateView(name, context, attrs)
+                    }
+
+                    override fun onCreateView(
+                        name: String,
+                        context: Context,
+                        attrs: AttributeSet
+                    ): View? {
+                        if (FrameLayout::class.java.simpleName.equals(name)) {
+                            val mock: FrameLayout = mock {
+                                `when`(this.context).thenReturn(context)
+                                `when`(this.id).thenReturn(R.id.controls_panel)
+                                `when`(this.requireViewById<View>(any())).thenCallRealMethod()
+                                `when`(this.findViewById<View>(R.id.controls_panel))
+                                    .thenReturn(this)
+                                `when`(this.post(any())).thenAnswer {
+                                    uiExecutor.execute(it.arguments[0] as Runnable)
+                                    true
+                                }
+                            }
+                            return mock
+                        } else {
+                            return null
+                        }
+                    }
+                }
+            )
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt
new file mode 100644
index 0000000..7eba3b46
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt
@@ -0,0 +1,88 @@
+/*
+ * Copyright (C) 2022 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.systemui.statusbar.phone
+
+import android.content.pm.UserInfo
+import android.os.UserManager
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.settings.UserTracker
+import com.android.systemui.util.concurrency.FakeExecutor
+import com.android.systemui.util.time.FakeSystemClock
+import junit.framework.Assert
+import org.junit.Before
+import org.junit.Test
+import org.mockito.Mock
+import org.mockito.Mockito.`when`
+import org.mockito.MockitoAnnotations
+
+@SmallTest
+class ManagedProfileControllerImplTest : SysuiTestCase() {
+
+    private val mainExecutor: FakeExecutor = FakeExecutor(FakeSystemClock())
+
+    private lateinit var controller: ManagedProfileControllerImpl
+
+    @Mock private lateinit var userTracker: UserTracker
+    @Mock private lateinit var userManager: UserManager
+
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+
+        controller = ManagedProfileControllerImpl(context, mainExecutor, userTracker, userManager)
+    }
+
+    @Test
+    fun hasWorkingProfile_isWorkModeEnabled_returnsTrue() {
+        `when`(userTracker.userId).thenReturn(1)
+        setupWorkingProfile(1)
+
+        Assert.assertEquals(true, controller.hasActiveProfile())
+    }
+
+    @Test
+    fun noWorkingProfile_isWorkModeEnabled_returnsFalse() {
+        `when`(userTracker.userId).thenReturn(1)
+
+        Assert.assertEquals(false, controller.hasActiveProfile())
+    }
+
+    @Test
+    fun listeningUserChanges_isWorkModeEnabled_returnsTrue() {
+        `when`(userTracker.userId).thenReturn(1)
+        controller.addCallback(TestCallback)
+        `when`(userTracker.userId).thenReturn(2)
+        setupWorkingProfile(2)
+
+        Assert.assertEquals(true, controller.hasActiveProfile())
+    }
+
+    private fun setupWorkingProfile(userId: Int) {
+        `when`(userManager.getEnabledProfiles(userId))
+            .thenReturn(
+                listOf(UserInfo(userId, "test_user", "", 0, UserManager.USER_TYPE_PROFILE_MANAGED))
+            )
+    }
+
+    private object TestCallback : ManagedProfileController.Callback {
+
+        override fun onManagedProfileChanged() = Unit
+
+        override fun onManagedProfileRemoved() = Unit
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcherTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcherTest.kt
new file mode 100644
index 0000000..96a280a
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcherTest.kt
@@ -0,0 +1,219 @@
+/*
+ * Copyright (C) 2022 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.systemui.statusbar.pipeline.mobile.data.repository
+
+import android.net.ConnectivityManager
+import android.telephony.SubscriptionInfo
+import android.telephony.SubscriptionManager
+import android.telephony.TelephonyManager
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.demomode.DemoMode
+import com.android.systemui.demomode.DemoModeController
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.DemoMobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.DemoModeMobileConnectionDataSource
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.validMobileEvent
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.prod.MobileConnectionsRepositoryImpl
+import com.android.systemui.statusbar.pipeline.shared.ConnectivityPipelineLogger
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.kotlinArgumentCaptor
+import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.whenever
+import com.android.systemui.util.settings.FakeSettings
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.runBlocking
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+import org.mockito.Mock
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+
+/**
+ * The switcher acts as a dispatcher to either the `prod` or `demo` versions of the repository
+ * interface it's switching on. These tests just need to verify that the entire interface properly
+ * switches over when the value of `demoMode` changes
+ */
+@Suppress("EXPERIMENTAL_IS_NOT_ENABLED")
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class MobileRepositorySwitcherTest : SysuiTestCase() {
+    private lateinit var underTest: MobileRepositorySwitcher
+    private lateinit var realRepo: MobileConnectionsRepositoryImpl
+    private lateinit var demoRepo: DemoMobileConnectionsRepository
+    private lateinit var mockDataSource: DemoModeMobileConnectionDataSource
+
+    @Mock private lateinit var connectivityManager: ConnectivityManager
+    @Mock private lateinit var subscriptionManager: SubscriptionManager
+    @Mock private lateinit var telephonyManager: TelephonyManager
+    @Mock private lateinit var logger: ConnectivityPipelineLogger
+    @Mock private lateinit var demoModeController: DemoModeController
+
+    private val globalSettings = FakeSettings()
+    private val fakeNetworkEventsFlow = MutableStateFlow<FakeNetworkEventModel?>(null)
+
+    private val scope = CoroutineScope(IMMEDIATE)
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+
+        // Never start in demo mode
+        whenever(demoModeController.isInDemoMode).thenReturn(false)
+
+        mockDataSource =
+            mock<DemoModeMobileConnectionDataSource>().also {
+                whenever(it.mobileEvents).thenReturn(fakeNetworkEventsFlow)
+            }
+
+        realRepo =
+            MobileConnectionsRepositoryImpl(
+                connectivityManager,
+                subscriptionManager,
+                telephonyManager,
+                logger,
+                fakeBroadcastDispatcher,
+                globalSettings,
+                context,
+                IMMEDIATE,
+                scope,
+                mock(),
+            )
+
+        demoRepo =
+            DemoMobileConnectionsRepository(
+                dataSource = mockDataSource,
+                scope = scope,
+                context = context,
+            )
+
+        underTest =
+            MobileRepositorySwitcher(
+                scope = scope,
+                realRepository = realRepo,
+                demoMobileConnectionsRepository = demoRepo,
+                demoModeController = demoModeController,
+            )
+    }
+
+    @After
+    fun tearDown() {
+        scope.cancel()
+    }
+
+    @Test
+    fun `active repo matches demo mode setting`() =
+        runBlocking(IMMEDIATE) {
+            whenever(demoModeController.isInDemoMode).thenReturn(false)
+
+            var latest: MobileConnectionsRepository? = null
+            val job = underTest.activeRepo.onEach { latest = it }.launchIn(this)
+
+            assertThat(latest).isEqualTo(realRepo)
+
+            startDemoMode()
+
+            assertThat(latest).isEqualTo(demoRepo)
+
+            finishDemoMode()
+
+            assertThat(latest).isEqualTo(realRepo)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `subscription list updates when demo mode changes`() =
+        runBlocking(IMMEDIATE) {
+            whenever(demoModeController.isInDemoMode).thenReturn(false)
+
+            whenever(subscriptionManager.completeActiveSubscriptionInfoList)
+                .thenReturn(listOf(SUB_1, SUB_2))
+
+            var latest: List<SubscriptionInfo>? = null
+            val job = underTest.subscriptionsFlow.onEach { latest = it }.launchIn(this)
+
+            // The real subscriptions has 2 subs
+            whenever(subscriptionManager.completeActiveSubscriptionInfoList)
+                .thenReturn(listOf(SUB_1, SUB_2))
+            getSubscriptionCallback().onSubscriptionsChanged()
+
+            assertThat(latest).isEqualTo(listOf(SUB_1, SUB_2))
+
+            // Demo mode turns on, and we should see only the demo subscriptions
+            startDemoMode()
+            fakeNetworkEventsFlow.value = validMobileEvent(subId = 3)
+
+            // Demo mobile connections repository makes arbitrarily-formed subscription info
+            // objects, so just validate the data we care about
+            assertThat(latest).hasSize(1)
+            assertThat(latest!![0].subscriptionId).isEqualTo(3)
+
+            finishDemoMode()
+
+            assertThat(latest).isEqualTo(listOf(SUB_1, SUB_2))
+
+            job.cancel()
+        }
+
+    private fun startDemoMode() {
+        whenever(demoModeController.isInDemoMode).thenReturn(true)
+        getDemoModeCallback().onDemoModeStarted()
+    }
+
+    private fun finishDemoMode() {
+        whenever(demoModeController.isInDemoMode).thenReturn(false)
+        getDemoModeCallback().onDemoModeFinished()
+    }
+
+    private fun getSubscriptionCallback(): SubscriptionManager.OnSubscriptionsChangedListener {
+        val callbackCaptor =
+            kotlinArgumentCaptor<SubscriptionManager.OnSubscriptionsChangedListener>()
+        verify(subscriptionManager)
+            .addOnSubscriptionsChangedListener(any(), callbackCaptor.capture())
+        return callbackCaptor.value
+    }
+
+    private fun getDemoModeCallback(): DemoMode {
+        val captor = kotlinArgumentCaptor<DemoMode>()
+        verify(demoModeController).addCallback(captor.capture())
+        return captor.value
+    }
+
+    companion object {
+        private val IMMEDIATE = Dispatchers.Main.immediate
+
+        private const val SUB_1_ID = 1
+        private val SUB_1 =
+            mock<SubscriptionInfo>().also { whenever(it.subscriptionId).thenReturn(SUB_1_ID) }
+
+        private const val SUB_2_ID = 2
+        private val SUB_2 =
+            mock<SubscriptionInfo>().also { whenever(it.subscriptionId).thenReturn(SUB_2_ID) }
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionParameterizedTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionParameterizedTest.kt
new file mode 100644
index 0000000..bf5ecd8
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionParameterizedTest.kt
@@ -0,0 +1,249 @@
+/*
+ * Copyright (C) 2022 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.systemui.statusbar.pipeline.mobile.data.repository.demo
+
+import android.telephony.Annotation
+import android.telephony.TelephonyManager
+import androidx.test.filters.SmallTest
+import com.android.settingslib.SignalIcon
+import com.android.settingslib.mobile.TelephonyIcons
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.statusbar.pipeline.mobile.data.model.DataConnectionState
+import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileSubscriptionModel
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel
+import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.whenever
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runTest
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.Parameterized
+import org.junit.runners.Parameterized.Parameters
+
+/**
+ * Parameterized test for all of the common values of [FakeNetworkEventModel]. This test simply
+ * verifies that passing the given model to [DemoMobileConnectionsRepository] results in the correct
+ * flows emitting from the given connection.
+ */
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(Parameterized::class)
+internal class DemoMobileConnectionParameterizedTest(private val testCase: TestCase) :
+    SysuiTestCase() {
+    private val testDispatcher = UnconfinedTestDispatcher()
+    private val testScope = TestScope(testDispatcher)
+
+    private val fakeNetworkEventFlow = MutableStateFlow<FakeNetworkEventModel?>(null)
+
+    private lateinit var connectionsRepo: DemoMobileConnectionsRepository
+    private lateinit var underTest: DemoMobileConnectionRepository
+    private lateinit var mockDataSource: DemoModeMobileConnectionDataSource
+
+    @Before
+    fun setUp() {
+        // The data source only provides one API, so we can mock it with a flow here for convenience
+        mockDataSource =
+            mock<DemoModeMobileConnectionDataSource>().also {
+                whenever(it.mobileEvents).thenReturn(fakeNetworkEventFlow)
+            }
+
+        connectionsRepo =
+            DemoMobileConnectionsRepository(
+                dataSource = mockDataSource,
+                scope = testScope.backgroundScope,
+                context = context,
+            )
+
+        connectionsRepo.startProcessingCommands()
+    }
+
+    @After
+    fun tearDown() {
+        testScope.cancel()
+    }
+
+    @Test
+    fun demoNetworkData() =
+        testScope.runTest {
+            val networkModel =
+                FakeNetworkEventModel.Mobile(
+                    level = testCase.level,
+                    dataType = testCase.dataType,
+                    subId = testCase.subId,
+                    carrierId = testCase.carrierId,
+                    inflateStrength = testCase.inflateStrength,
+                    activity = testCase.activity,
+                    carrierNetworkChange = testCase.carrierNetworkChange,
+                )
+
+            fakeNetworkEventFlow.value = networkModel
+            underTest = connectionsRepo.getRepoForSubId(subId)
+
+            assertConnection(underTest, networkModel)
+        }
+
+    private fun assertConnection(
+        conn: DemoMobileConnectionRepository,
+        model: FakeNetworkEventModel
+    ) {
+        when (model) {
+            is FakeNetworkEventModel.Mobile -> {
+                val subscriptionModel: MobileSubscriptionModel = conn.subscriptionModelFlow.value
+                assertThat(conn.subId).isEqualTo(model.subId)
+                assertThat(subscriptionModel.cdmaLevel).isEqualTo(model.level)
+                assertThat(subscriptionModel.primaryLevel).isEqualTo(model.level)
+                assertThat(subscriptionModel.dataActivityDirection).isEqualTo(model.activity)
+                assertThat(subscriptionModel.carrierNetworkChangeActive)
+                    .isEqualTo(model.carrierNetworkChange)
+
+                // TODO(b/261029387): check these once we start handling them
+                assertThat(subscriptionModel.isEmergencyOnly).isFalse()
+                assertThat(subscriptionModel.isGsm).isFalse()
+                assertThat(subscriptionModel.dataConnectionState)
+                    .isEqualTo(DataConnectionState.Connected)
+            }
+            // MobileDisabled isn't combinatorial in nature, and is tested in
+            // DemoMobileConnectionsRepositoryTest.kt
+            else -> {}
+        }
+    }
+
+    /** Matches [FakeNetworkEventModel] */
+    internal data class TestCase(
+        val level: Int,
+        val dataType: SignalIcon.MobileIconGroup,
+        val subId: Int,
+        val carrierId: Int,
+        val inflateStrength: Boolean,
+        @Annotation.DataActivityType val activity: Int,
+        val carrierNetworkChange: Boolean,
+    ) {
+        override fun toString(): String {
+            return "INPUT(level=$level, " +
+                "dataType=${dataType.name}, " +
+                "subId=$subId, " +
+                "carrierId=$carrierId, " +
+                "inflateStrength=$inflateStrength, " +
+                "activity=$activity, " +
+                "carrierNetworkChange=$carrierNetworkChange)"
+        }
+
+        // Convenience for iterating test data and creating new cases
+        fun modifiedBy(
+            level: Int? = null,
+            dataType: SignalIcon.MobileIconGroup? = null,
+            subId: Int? = null,
+            carrierId: Int? = null,
+            inflateStrength: Boolean? = null,
+            @Annotation.DataActivityType activity: Int? = null,
+            carrierNetworkChange: Boolean? = null,
+        ): TestCase =
+            TestCase(
+                level = level ?: this.level,
+                dataType = dataType ?: this.dataType,
+                subId = subId ?: this.subId,
+                carrierId = carrierId ?: this.carrierId,
+                inflateStrength = inflateStrength ?: this.inflateStrength,
+                activity = activity ?: this.activity,
+                carrierNetworkChange = carrierNetworkChange ?: this.carrierNetworkChange
+            )
+    }
+
+    companion object {
+        private val subId = 1
+
+        private val booleanList = listOf(true, false)
+        private val levels = listOf(0, 1, 2, 3)
+        private val dataTypes =
+            listOf(
+                TelephonyIcons.THREE_G,
+                TelephonyIcons.LTE,
+                TelephonyIcons.FOUR_G,
+                TelephonyIcons.NR_5G,
+                TelephonyIcons.NR_5G_PLUS,
+            )
+        private val carrierIds = listOf(1, 10, 100)
+        private val inflateStrength = booleanList
+        private val activity =
+            listOf(
+                TelephonyManager.DATA_ACTIVITY_NONE,
+                TelephonyManager.DATA_ACTIVITY_IN,
+                TelephonyManager.DATA_ACTIVITY_OUT,
+                TelephonyManager.DATA_ACTIVITY_INOUT
+            )
+        private val carrierNetworkChange = booleanList
+
+        @Parameters(name = "{0}") @JvmStatic fun data() = testData()
+
+        /**
+         * Generate some test data. For the sake of convenience, we'll parameterize only non-null
+         * network event data. So given the lists of test data:
+         * ```
+         *    list1 = [1, 2, 3]
+         *    list2 = [false, true]
+         *    list3 = [a, b, c]
+         * ```
+         * We'll generate test cases for:
+         *
+         * Test (1, false, a) Test (2, false, a) Test (3, false, a) Test (1, true, a) Test (1,
+         * false, b) Test (1, false, c)
+         *
+         * NOTE: this is not a combinatorial product of all of the possible sets of parameters.
+         * Since this test is built to exercise demo mode, the general approach is to define a
+         * fully-formed "base case", and from there to make sure to use every valid parameter once,
+         * by defining the rest of the test cases against the base case. Specific use-cases can be
+         * added to the non-parameterized test, or manually below the generated test cases.
+         */
+        private fun testData(): List<TestCase> {
+            val testSet = mutableSetOf<TestCase>()
+
+            val baseCase =
+                TestCase(
+                    levels.first(),
+                    dataTypes.first(),
+                    subId,
+                    carrierIds.first(),
+                    inflateStrength.first(),
+                    activity.first(),
+                    carrierNetworkChange.first()
+                )
+
+            val tail =
+                sequenceOf(
+                        levels.map { baseCase.modifiedBy(level = it) },
+                        dataTypes.map { baseCase.modifiedBy(dataType = it) },
+                        carrierIds.map { baseCase.modifiedBy(carrierId = it) },
+                        inflateStrength.map { baseCase.modifiedBy(inflateStrength = it) },
+                        activity.map { baseCase.modifiedBy(activity = it) },
+                        carrierNetworkChange.map { baseCase.modifiedBy(carrierNetworkChange = it) },
+                    )
+                    .flatten()
+
+            testSet.add(baseCase)
+            tail.toCollection(testSet)
+
+            return testSet.toList()
+        }
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepositoryTest.kt
new file mode 100644
index 0000000..a8f6993
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepositoryTest.kt
@@ -0,0 +1,306 @@
+/*
+ * Copyright (C) 2022 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.systemui.statusbar.pipeline.mobile.data.repository.demo
+
+import android.telephony.SubscriptionInfo
+import android.telephony.TelephonyManager.DATA_ACTIVITY_INOUT
+import android.telephony.TelephonyManager.UNKNOWN_CARRIER_ID
+import androidx.test.filters.SmallTest
+import com.android.settingslib.SignalIcon
+import com.android.settingslib.mobile.TelephonyIcons.THREE_G
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.statusbar.pipeline.mobile.data.model.DataConnectionState
+import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileSubscriptionModel
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel.MobileDisabled
+import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.whenever
+import com.google.common.truth.Truth.assertThat
+import junit.framework.Assert
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+class DemoMobileConnectionsRepositoryTest : SysuiTestCase() {
+    private val testDispatcher = UnconfinedTestDispatcher()
+    private val testScope = TestScope(testDispatcher)
+
+    private val fakeNetworkEventFlow = MutableStateFlow<FakeNetworkEventModel?>(null)
+
+    private lateinit var underTest: DemoMobileConnectionsRepository
+    private lateinit var mockDataSource: DemoModeMobileConnectionDataSource
+
+    @Before
+    fun setUp() {
+        // The data source only provides one API, so we can mock it with a flow here for convenience
+        mockDataSource =
+            mock<DemoModeMobileConnectionDataSource>().also {
+                whenever(it.mobileEvents).thenReturn(fakeNetworkEventFlow)
+            }
+
+        underTest =
+            DemoMobileConnectionsRepository(
+                dataSource = mockDataSource,
+                scope = testScope.backgroundScope,
+                context = context,
+            )
+
+        underTest.startProcessingCommands()
+    }
+
+    @Test
+    fun `network event - create new subscription`() =
+        testScope.runTest {
+            var latest: List<SubscriptionInfo>? = null
+            val job = underTest.subscriptionsFlow.onEach { latest = it }.launchIn(this)
+
+            assertThat(latest).isEmpty()
+
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 1)
+
+            assertThat(latest).hasSize(1)
+            assertThat(latest!![0].subscriptionId).isEqualTo(1)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `network event - reuses subscription when same Id`() =
+        testScope.runTest {
+            var latest: List<SubscriptionInfo>? = null
+            val job = underTest.subscriptionsFlow.onEach { latest = it }.launchIn(this)
+
+            assertThat(latest).isEmpty()
+
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 1, level = 1)
+
+            assertThat(latest).hasSize(1)
+            assertThat(latest!![0].subscriptionId).isEqualTo(1)
+
+            // Second network event comes in with the same subId, does not create a new subscription
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 1, level = 2)
+
+            assertThat(latest).hasSize(1)
+            assertThat(latest!![0].subscriptionId).isEqualTo(1)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `multiple subscriptions`() =
+        testScope.runTest {
+            var latest: List<SubscriptionInfo>? = null
+            val job = underTest.subscriptionsFlow.onEach { latest = it }.launchIn(this)
+
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 1)
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 2)
+
+            assertThat(latest).hasSize(2)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `mobile disabled event - disables connection - subId specified - single conn`() =
+        testScope.runTest {
+            var latest: List<SubscriptionInfo>? = null
+            val job = underTest.subscriptionsFlow.onEach { latest = it }.launchIn(this)
+
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 1, level = 1)
+
+            fakeNetworkEventFlow.value = MobileDisabled(subId = 1)
+
+            assertThat(latest).hasSize(0)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `mobile disabled event - disables connection - subId not specified - single conn`() =
+        testScope.runTest {
+            var latest: List<SubscriptionInfo>? = null
+            val job = underTest.subscriptionsFlow.onEach { latest = it }.launchIn(this)
+
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 1, level = 1)
+
+            fakeNetworkEventFlow.value = MobileDisabled(subId = null)
+
+            assertThat(latest).hasSize(0)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `mobile disabled event - disables connection - subId specified - multiple conn`() =
+        testScope.runTest {
+            var latest: List<SubscriptionInfo>? = null
+            val job = underTest.subscriptionsFlow.onEach { latest = it }.launchIn(this)
+
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 1, level = 1)
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 2, level = 1)
+
+            fakeNetworkEventFlow.value = MobileDisabled(subId = 2)
+
+            assertThat(latest).hasSize(1)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `mobile disabled event - subId not specified - multiple conn - ignores command`() =
+        testScope.runTest {
+            var latest: List<SubscriptionInfo>? = null
+            val job = underTest.subscriptionsFlow.onEach { latest = it }.launchIn(this)
+
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 1, level = 1)
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 2, level = 1)
+
+            fakeNetworkEventFlow.value = MobileDisabled(subId = null)
+
+            assertThat(latest).hasSize(2)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `demo connection - single subscription`() =
+        testScope.runTest {
+            var currentEvent: FakeNetworkEventModel = validMobileEvent(subId = 1)
+            var connections: List<DemoMobileConnectionRepository>? = null
+            val job =
+                underTest.subscriptionsFlow
+                    .onEach { infos ->
+                        connections =
+                            infos.map { info -> underTest.getRepoForSubId(info.subscriptionId) }
+                    }
+                    .launchIn(this)
+
+            fakeNetworkEventFlow.value = currentEvent
+
+            assertThat(connections).hasSize(1)
+            val connection1 = connections!![0]
+
+            assertConnection(connection1, currentEvent)
+
+            // Exercise the whole api
+
+            currentEvent = validMobileEvent(subId = 1, level = 2)
+            fakeNetworkEventFlow.value = currentEvent
+            assertConnection(connection1, currentEvent)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `demo connection - two connections - update second - no affect on first`() =
+        testScope.runTest {
+            var currentEvent1 = validMobileEvent(subId = 1)
+            var connection1: DemoMobileConnectionRepository? = null
+            var currentEvent2 = validMobileEvent(subId = 2)
+            var connection2: DemoMobileConnectionRepository? = null
+            var connections: List<DemoMobileConnectionRepository>? = null
+            val job =
+                underTest.subscriptionsFlow
+                    .onEach { infos ->
+                        connections =
+                            infos.map { info -> underTest.getRepoForSubId(info.subscriptionId) }
+                    }
+                    .launchIn(this)
+
+            fakeNetworkEventFlow.value = currentEvent1
+            fakeNetworkEventFlow.value = currentEvent2
+            assertThat(connections).hasSize(2)
+            connections!!.forEach {
+                if (it.subId == 1) {
+                    connection1 = it
+                } else if (it.subId == 2) {
+                    connection2 = it
+                } else {
+                    Assert.fail("Unexpected subscription")
+                }
+            }
+
+            assertConnection(connection1!!, currentEvent1)
+            assertConnection(connection2!!, currentEvent2)
+
+            // WHEN the event changes for connection 2, it updates, and connection 1 stays the same
+            currentEvent2 = validMobileEvent(subId = 2, activity = DATA_ACTIVITY_INOUT)
+            fakeNetworkEventFlow.value = currentEvent2
+            assertConnection(connection1!!, currentEvent1)
+            assertConnection(connection2!!, currentEvent2)
+
+            // and vice versa
+            currentEvent1 = validMobileEvent(subId = 1, inflateStrength = true)
+            fakeNetworkEventFlow.value = currentEvent1
+            assertConnection(connection1!!, currentEvent1)
+            assertConnection(connection2!!, currentEvent2)
+
+            job.cancel()
+        }
+
+    private fun assertConnection(
+        conn: DemoMobileConnectionRepository,
+        model: FakeNetworkEventModel
+    ) {
+        when (model) {
+            is FakeNetworkEventModel.Mobile -> {
+                val subscriptionModel: MobileSubscriptionModel = conn.subscriptionModelFlow.value
+                assertThat(conn.subId).isEqualTo(model.subId)
+                assertThat(subscriptionModel.cdmaLevel).isEqualTo(model.level)
+                assertThat(subscriptionModel.primaryLevel).isEqualTo(model.level)
+                assertThat(subscriptionModel.dataActivityDirection).isEqualTo(model.activity)
+                assertThat(subscriptionModel.carrierNetworkChangeActive)
+                    .isEqualTo(model.carrierNetworkChange)
+
+                // TODO(b/261029387) check these once we start handling them
+                assertThat(subscriptionModel.isEmergencyOnly).isFalse()
+                assertThat(subscriptionModel.isGsm).isFalse()
+                assertThat(subscriptionModel.dataConnectionState)
+                    .isEqualTo(DataConnectionState.Connected)
+            }
+            else -> {}
+        }
+    }
+}
+
+/** Convenience to create a valid fake network event with minimal params */
+fun validMobileEvent(
+    level: Int? = 1,
+    dataType: SignalIcon.MobileIconGroup? = THREE_G,
+    subId: Int? = 1,
+    carrierId: Int? = UNKNOWN_CARRIER_ID,
+    inflateStrength: Boolean? = false,
+    activity: Int? = null,
+    carrierNetworkChange: Boolean = false,
+): FakeNetworkEventModel =
+    FakeNetworkEventModel.Mobile(
+        level = level,
+        dataType = dataType,
+        subId = subId,
+        carrierId = carrierId,
+        inflateStrength = inflateStrength,
+        activity = activity,
+        carrierNetworkChange = carrierNetworkChange,
+    )
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionRepositoryTest.kt
similarity index 99%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionRepositoryTest.kt
rename to packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionRepositoryTest.kt
index aa7ab0d..c8df5ac 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionRepositoryTest.kt
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.statusbar.pipeline.mobile.data.repository
+package com.android.systemui.statusbar.pipeline.mobile.data.repository.prod
 
 import android.os.UserHandle
 import android.provider.Settings
@@ -40,6 +40,7 @@
 import com.android.systemui.statusbar.pipeline.mobile.data.model.DefaultNetworkType
 import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileSubscriptionModel
 import com.android.systemui.statusbar.pipeline.mobile.data.model.OverrideNetworkType
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.shared.ConnectivityPipelineLogger
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.argumentCaptor
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryTest.kt
similarity index 99%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepositoryTest.kt
rename to packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryTest.kt
index a953a3d..359ea18 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryTest.kt
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.statusbar.pipeline.mobile.data.repository
+package com.android.systemui.statusbar.pipeline.mobile.data.repository.prod
 
 import android.content.Intent
 import android.net.ConnectivityManager
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/TestableAlertDialogTest.kt b/packages/SystemUI/tests/src/com/android/systemui/util/TestableAlertDialogTest.kt
new file mode 100644
index 0000000..01dd60a
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/util/TestableAlertDialogTest.kt
@@ -0,0 +1,333 @@
+/*
+ * Copyright (C) 2022 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.systemui.util
+
+import android.content.DialogInterface
+import android.content.DialogInterface.BUTTON_NEGATIVE
+import android.content.DialogInterface.BUTTON_NEUTRAL
+import android.content.DialogInterface.BUTTON_POSITIVE
+import android.testing.AndroidTestingRunner
+import android.testing.TestableLooper
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.mock
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mockito.anyInt
+import org.mockito.Mockito.inOrder
+import org.mockito.Mockito.never
+import org.mockito.Mockito.verify
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+@TestableLooper.RunWithLooper
+class TestableAlertDialogTest : SysuiTestCase() {
+
+    @Test
+    fun dialogNotShowingWhenCreated() {
+        val dialog = TestableAlertDialog(context)
+
+        assertThat(dialog.isShowing).isFalse()
+    }
+
+    @Test
+    fun dialogShownDoesntCrash() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.show()
+    }
+
+    @Test
+    fun dialogShowing() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.show()
+
+        assertThat(dialog.isShowing).isTrue()
+    }
+
+    @Test
+    fun showListenerCalled() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnShowListener = mock()
+        dialog.setOnShowListener(listener)
+
+        dialog.show()
+
+        verify(listener).onShow(dialog)
+    }
+
+    @Test
+    fun showListenerRemoved() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnShowListener = mock()
+        dialog.setOnShowListener(listener)
+        dialog.setOnShowListener(null)
+
+        dialog.show()
+
+        verify(listener, never()).onShow(any())
+    }
+
+    @Test
+    fun dialogHiddenNotShowing() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.show()
+        dialog.hide()
+
+        assertThat(dialog.isShowing).isFalse()
+    }
+
+    @Test
+    fun dialogDismissNotShowing() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.show()
+        dialog.dismiss()
+
+        assertThat(dialog.isShowing).isFalse()
+    }
+
+    @Test
+    fun dismissListenerCalled_ifShowing() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnDismissListener = mock()
+        dialog.setOnDismissListener(listener)
+
+        dialog.show()
+        dialog.dismiss()
+
+        verify(listener).onDismiss(dialog)
+    }
+
+    @Test
+    fun dismissListenerNotCalled_ifNotShowing() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnDismissListener = mock()
+        dialog.setOnDismissListener(listener)
+
+        dialog.dismiss()
+
+        verify(listener, never()).onDismiss(any())
+    }
+
+    @Test
+    fun dismissListenerRemoved() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnDismissListener = mock()
+        dialog.setOnDismissListener(listener)
+        dialog.setOnDismissListener(null)
+
+        dialog.show()
+        dialog.dismiss()
+
+        verify(listener, never()).onDismiss(any())
+    }
+
+    @Test
+    fun cancelListenerCalled_showing() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnCancelListener = mock()
+        dialog.setOnCancelListener(listener)
+
+        dialog.show()
+        dialog.cancel()
+
+        verify(listener).onCancel(dialog)
+    }
+
+    @Test
+    fun cancelListenerCalled_notShowing() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnCancelListener = mock()
+        dialog.setOnCancelListener(listener)
+
+        dialog.cancel()
+
+        verify(listener).onCancel(dialog)
+    }
+
+    @Test
+    fun dismissCalledOnCancel_showing() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnDismissListener = mock()
+        dialog.setOnDismissListener(listener)
+
+        dialog.show()
+        dialog.cancel()
+
+        verify(listener).onDismiss(dialog)
+    }
+
+    @Test
+    fun dialogCancelNotShowing() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.show()
+        dialog.cancel()
+
+        assertThat(dialog.isShowing).isFalse()
+    }
+
+    @Test
+    fun cancelListenerRemoved() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnCancelListener = mock()
+        dialog.setOnCancelListener(listener)
+        dialog.setOnCancelListener(null)
+
+        dialog.show()
+        dialog.cancel()
+
+        verify(listener, never()).onCancel(any())
+    }
+
+    @Test
+    fun positiveButtonClick() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnClickListener = mock()
+        dialog.setButton(BUTTON_POSITIVE, "", listener)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_POSITIVE)
+
+        verify(listener).onClick(dialog, BUTTON_POSITIVE)
+    }
+
+    @Test
+    fun positiveButtonListener_noCalledWhenClickOtherButtons() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnClickListener = mock()
+        dialog.setButton(BUTTON_POSITIVE, "", listener)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_NEUTRAL)
+        dialog.clickButton(BUTTON_NEGATIVE)
+
+        verify(listener, never()).onClick(any(), anyInt())
+    }
+
+    @Test
+    fun negativeButtonClick() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnClickListener = mock()
+        dialog.setButton(BUTTON_NEGATIVE, "", listener)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_NEGATIVE)
+
+        verify(listener).onClick(dialog, DialogInterface.BUTTON_NEGATIVE)
+    }
+
+    @Test
+    fun negativeButtonListener_noCalledWhenClickOtherButtons() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnClickListener = mock()
+        dialog.setButton(BUTTON_NEGATIVE, "", listener)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_NEUTRAL)
+        dialog.clickButton(BUTTON_POSITIVE)
+
+        verify(listener, never()).onClick(any(), anyInt())
+    }
+
+    @Test
+    fun neutralButtonClick() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnClickListener = mock()
+        dialog.setButton(BUTTON_NEUTRAL, "", listener)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_NEUTRAL)
+
+        verify(listener).onClick(dialog, BUTTON_NEUTRAL)
+    }
+
+    @Test
+    fun neutralButtonListener_noCalledWhenClickOtherButtons() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnClickListener = mock()
+        dialog.setButton(BUTTON_NEUTRAL, "", listener)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_POSITIVE)
+        dialog.clickButton(BUTTON_NEGATIVE)
+
+        verify(listener, never()).onClick(any(), anyInt())
+    }
+
+    @Test
+    fun sameClickListenerCalledCorrectly() {
+        val dialog = TestableAlertDialog(context)
+        val listener: DialogInterface.OnClickListener = mock()
+        dialog.setButton(BUTTON_POSITIVE, "", listener)
+        dialog.setButton(BUTTON_NEUTRAL, "", listener)
+        dialog.setButton(BUTTON_NEGATIVE, "", listener)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_POSITIVE)
+        dialog.clickButton(BUTTON_NEGATIVE)
+        dialog.clickButton(BUTTON_NEUTRAL)
+
+        val inOrder = inOrder(listener)
+        inOrder.verify(listener).onClick(dialog, BUTTON_POSITIVE)
+        inOrder.verify(listener).onClick(dialog, BUTTON_NEGATIVE)
+        inOrder.verify(listener).onClick(dialog, BUTTON_NEUTRAL)
+    }
+
+    @Test(expected = IllegalArgumentException::class)
+    fun clickBadButton() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.clickButton(10000)
+    }
+
+    @Test
+    fun clickButtonDismisses_positive() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_POSITIVE)
+
+        assertThat(dialog.isShowing).isFalse()
+    }
+
+    @Test
+    fun clickButtonDismisses_negative() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_NEGATIVE)
+
+        assertThat(dialog.isShowing).isFalse()
+    }
+
+    @Test
+    fun clickButtonDismisses_neutral() {
+        val dialog = TestableAlertDialog(context)
+
+        dialog.show()
+        dialog.clickButton(BUTTON_NEUTRAL)
+
+        assertThat(dialog.isShowing).isFalse()
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/util/TestableAlertDialog.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/util/TestableAlertDialog.kt
new file mode 100644
index 0000000..4d79554
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/util/TestableAlertDialog.kt
@@ -0,0 +1,141 @@
+/*
+ * Copyright (C) 2022 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.systemui.util
+
+import android.app.AlertDialog
+import android.content.Context
+import android.content.DialogInterface
+import java.lang.IllegalArgumentException
+
+/**
+ * [AlertDialog] that is easier to test. Due to [AlertDialog] being a class and not an interface,
+ * there are some things that cannot be avoided, like the creation of a [Handler] on the main thread
+ * (and therefore needing a prepared [Looper] in the test).
+ *
+ * It bypasses calls to show, clicks on buttons, cancel and dismiss so it all can happen bounded in
+ * the test. It tries to be as close in behavior as a real [AlertDialog].
+ *
+ * It will only call [onCreate] as part of its lifecycle, but not any of the other lifecycle methods
+ * in [Dialog].
+ *
+ * In order to test clicking on buttons, use [clickButton] instead of calling [View.callOnClick] on
+ * the view returned by [getButton] to bypass the internal [Handler].
+ */
+class TestableAlertDialog(context: Context) : AlertDialog(context) {
+
+    private var _onDismissListener: DialogInterface.OnDismissListener? = null
+    private var _onCancelListener: DialogInterface.OnCancelListener? = null
+    private var _positiveButtonClickListener: DialogInterface.OnClickListener? = null
+    private var _negativeButtonClickListener: DialogInterface.OnClickListener? = null
+    private var _neutralButtonClickListener: DialogInterface.OnClickListener? = null
+    private var _onShowListener: DialogInterface.OnShowListener? = null
+    private var _dismissOverride: Runnable? = null
+
+    private var showing = false
+    private var visible = false
+    private var created = false
+
+    override fun show() {
+        if (!created) {
+            created = true
+            onCreate(null)
+        }
+        if (isShowing) return
+        showing = true
+        visible = true
+        _onShowListener?.onShow(this)
+    }
+
+    override fun hide() {
+        visible = false
+    }
+
+    override fun isShowing(): Boolean {
+        return visible && showing
+    }
+
+    override fun dismiss() {
+        if (!showing) {
+            return
+        }
+        if (_dismissOverride != null) {
+            _dismissOverride?.run()
+            return
+        }
+        _onDismissListener?.onDismiss(this)
+        showing = false
+    }
+
+    override fun cancel() {
+        _onCancelListener?.onCancel(this)
+        dismiss()
+    }
+
+    override fun setOnDismissListener(listener: DialogInterface.OnDismissListener?) {
+        _onDismissListener = listener
+    }
+
+    override fun setOnCancelListener(listener: DialogInterface.OnCancelListener?) {
+        _onCancelListener = listener
+    }
+
+    override fun setOnShowListener(listener: DialogInterface.OnShowListener?) {
+        _onShowListener = listener
+    }
+
+    override fun takeCancelAndDismissListeners(
+        msg: String?,
+        cancel: DialogInterface.OnCancelListener?,
+        dismiss: DialogInterface.OnDismissListener?
+    ): Boolean {
+        _onCancelListener = cancel
+        _onDismissListener = dismiss
+        return true
+    }
+
+    override fun setButton(
+        whichButton: Int,
+        text: CharSequence?,
+        listener: DialogInterface.OnClickListener?
+    ) {
+        super.setButton(whichButton, text, listener)
+        when (whichButton) {
+            DialogInterface.BUTTON_POSITIVE -> _positiveButtonClickListener = listener
+            DialogInterface.BUTTON_NEGATIVE -> _negativeButtonClickListener = listener
+            DialogInterface.BUTTON_NEUTRAL -> _neutralButtonClickListener = listener
+            else -> Unit
+        }
+    }
+
+    /**
+     * Click one of the buttons in the [AlertDialog] and call the corresponding listener.
+     *
+     * Button ids are from [DialogInterface].
+     */
+    fun clickButton(whichButton: Int) {
+        val listener =
+            when (whichButton) {
+                DialogInterface.BUTTON_POSITIVE -> _positiveButtonClickListener
+                DialogInterface.BUTTON_NEGATIVE -> _negativeButtonClickListener
+                DialogInterface.BUTTON_NEUTRAL -> _neutralButtonClickListener
+                else -> throw IllegalArgumentException("Wrong button $whichButton")
+            }
+        listener?.onClick(this, whichButton)
+        dismiss()
+    }
+}
diff --git a/services/core/Android.bp b/services/core/Android.bp
index a61a61b..b6f0237 100644
--- a/services/core/Android.bp
+++ b/services/core/Android.bp
@@ -144,6 +144,7 @@
 
     static_libs: [
         "android.hardware.authsecret-V1.0-java",
+        "android.hardware.authsecret-V1-java",
         "android.hardware.boot-V1.0-java",
         "android.hardware.boot-V1.1-java",
         "android.hardware.boot-V1.2-java",
diff --git a/services/core/java/com/android/server/locksettings/LockSettingsService.java b/services/core/java/com/android/server/locksettings/LockSettingsService.java
index 0ae3a02..99d77d6 100644
--- a/services/core/java/com/android/server/locksettings/LockSettingsService.java
+++ b/services/core/java/com/android/server/locksettings/LockSettingsService.java
@@ -69,7 +69,7 @@
 import android.content.pm.UserInfo;
 import android.database.ContentObserver;
 import android.database.sqlite.SQLiteDatabase;
-import android.hardware.authsecret.V1_0.IAuthSecret;
+import android.hardware.authsecret.IAuthSecret;
 import android.hardware.biometrics.BiometricManager;
 import android.hardware.face.Face;
 import android.hardware.face.FaceManager;
@@ -265,7 +265,8 @@
     protected boolean mHasSecureLockScreen;
 
     protected IGateKeeperService mGateKeeperService;
-    protected IAuthSecret mAuthSecretService;
+    protected IAuthSecret mAuthSecretServiceAidl;
+    protected android.hardware.authsecret.V1_0.IAuthSecret mAuthSecretServiceHidl;
 
     private static final String GSI_RUNNING_PROP = "ro.gsid.image_running";
 
@@ -833,12 +834,19 @@
     }
 
     private void getAuthSecretHal() {
-        try {
-            mAuthSecretService = IAuthSecret.getService(/* retry */ true);
-        } catch (NoSuchElementException e) {
-            Slog.i(TAG, "Device doesn't implement AuthSecret HAL");
-        } catch (RemoteException e) {
-            Slog.w(TAG, "Failed to get AuthSecret HAL", e);
+        mAuthSecretServiceAidl = IAuthSecret.Stub.asInterface(ServiceManager.
+                                 waitForDeclaredService(IAuthSecret.DESCRIPTOR + "/default"));
+        if (mAuthSecretServiceAidl == null) {
+            Slog.i(TAG, "Device doesn't implement AuthSecret HAL(aidl), try to get hidl version");
+
+            try {
+                mAuthSecretServiceHidl =
+                    android.hardware.authsecret.V1_0.IAuthSecret.getService(/* retry */ true);
+            } catch (NoSuchElementException e) {
+                Slog.i(TAG, "Device doesn't implement AuthSecret HAL(hidl)");
+            } catch (RemoteException e) {
+                Slog.w(TAG, "Failed to get AuthSecret HAL(hidl)", e);
+            }
         }
     }
 
@@ -2601,17 +2609,25 @@
         // If the given user is the primary user, pass the auth secret to the HAL.  Only the system
         // user can be primary.  Check for the system user ID before calling getUserInfo(), as other
         // users may still be under construction.
-        if (mAuthSecretService != null && userId == UserHandle.USER_SYSTEM &&
+        if (userId == UserHandle.USER_SYSTEM &&
                 mUserManager.getUserInfo(userId).isPrimary()) {
-            try {
-                final byte[] rawSecret = sp.deriveVendorAuthSecret();
-                final ArrayList<Byte> secret = new ArrayList<>(rawSecret.length);
-                for (int i = 0; i < rawSecret.length; ++i) {
-                    secret.add(rawSecret[i]);
+            final byte[] rawSecret = sp.deriveVendorAuthSecret();
+            if (mAuthSecretServiceAidl != null) {
+                try {
+                    mAuthSecretServiceAidl.setPrimaryUserCredential(rawSecret);
+                } catch (RemoteException e) {
+                    Slog.w(TAG, "Failed to pass primary user secret to AuthSecret HAL(aidl)", e);
                 }
-                mAuthSecretService.primaryUserCredential(secret);
-            } catch (RemoteException e) {
-                Slog.w(TAG, "Failed to pass primary user secret to AuthSecret HAL", e);
+            } else if (mAuthSecretServiceHidl != null) {
+                try {
+                    final ArrayList<Byte> secret = new ArrayList<>(rawSecret.length);
+                    for (int i = 0; i < rawSecret.length; ++i) {
+                        secret.add(rawSecret[i]);
+                    }
+                    mAuthSecretServiceHidl.primaryUserCredential(secret);
+                } catch (RemoteException e) {
+                    Slog.w(TAG, "Failed to pass primary user secret to AuthSecret HAL(hidl)", e);
+                }
             }
         }
     }
diff --git a/services/core/java/com/android/server/wearable/WearableSensingManagerService.java b/services/core/java/com/android/server/wearable/WearableSensingManagerService.java
index e155a06..707d0044 100644
--- a/services/core/java/com/android/server/wearable/WearableSensingManagerService.java
+++ b/services/core/java/com/android/server/wearable/WearableSensingManagerService.java
@@ -161,6 +161,32 @@
         return null;
     }
 
+    // Used in testing.
+    void provideDataStream(@UserIdInt int userId, ParcelFileDescriptor parcelFileDescriptor,
+            RemoteCallback callback) {
+        synchronized (mLock) {
+            final WearableSensingManagerPerUserService mService = getServiceForUserLocked(userId);
+            if (mService != null) {
+                mService.onProvideDataStream(parcelFileDescriptor, callback);
+            } else {
+                Slog.w(TAG, "Service not available.");
+            }
+        }
+    }
+
+    // Used in testing.
+    void provideData(@UserIdInt int userId, PersistableBundle data, SharedMemory sharedMemory,
+            RemoteCallback callback) {
+        synchronized (mLock) {
+            final WearableSensingManagerPerUserService mService = getServiceForUserLocked(userId);
+            if (mService != null) {
+                mService.onProvidedData(data, sharedMemory, callback);
+            } else {
+                Slog.w(TAG, "Service not available.");
+            }
+        }
+    }
+
     private final class WearableSensingManagerInternal extends IWearableSensingManager.Stub {
         final WearableSensingManagerPerUserService mService = getServiceForUserLocked(
                 UserHandle.getCallingUserId());
@@ -205,6 +231,8 @@
         @Override
         public void onShellCommand(FileDescriptor in, FileDescriptor out, FileDescriptor err,
                 String[] args, ShellCallback callback, ResultReceiver resultReceiver) {
+            new WearableSensingShellCommand(WearableSensingManagerService.this).exec(
+                    this, in, out, err, args, callback, resultReceiver);
         }
     }
 }
diff --git a/services/core/java/com/android/server/wearable/WearableSensingShellCommand.java b/services/core/java/com/android/server/wearable/WearableSensingShellCommand.java
new file mode 100644
index 0000000..842bccb
--- /dev/null
+++ b/services/core/java/com/android/server/wearable/WearableSensingShellCommand.java
@@ -0,0 +1,212 @@
+/*
+ * Copyright (C) 2022 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.server.wearable;
+
+import android.annotation.NonNull;
+import android.app.wearable.WearableSensingManager;
+import android.content.ComponentName;
+import android.os.Binder;
+import android.os.ParcelFileDescriptor;
+import android.os.PersistableBundle;
+import android.os.RemoteCallback;
+import android.os.ShellCommand;
+import android.util.Slog;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.io.PrintWriter;
+
+final class WearableSensingShellCommand extends ShellCommand {
+    private static final String TAG = WearableSensingShellCommand.class.getSimpleName();
+
+    static final TestableCallbackInternal sTestableCallbackInternal =
+            new TestableCallbackInternal();
+
+    @NonNull
+    private final WearableSensingManagerService mService;
+
+    private static ParcelFileDescriptor[] sPipe;
+
+    WearableSensingShellCommand(@NonNull WearableSensingManagerService service) {
+        mService = service;
+    }
+
+    /** Callbacks for WearableSensingService results used internally for testing. */
+    static class TestableCallbackInternal {
+        private int mLastStatus;
+
+        public int getLastStatus() {
+            return mLastStatus;
+        }
+
+        @NonNull
+        private RemoteCallback createRemoteStatusCallback() {
+            return new RemoteCallback(result -> {
+                int status = result.getInt(WearableSensingManager.STATUS_RESPONSE_BUNDLE_KEY);
+                final long token = Binder.clearCallingIdentity();
+                try {
+                    mLastStatus = status;
+                } finally {
+                    Binder.restoreCallingIdentity(token);
+                }
+            });
+        }
+    }
+
+    @Override
+    public int onCommand(String cmd) {
+        if (cmd == null) {
+            return handleDefaultCommands(cmd);
+        }
+
+        switch (cmd) {
+            case "create-data-stream":
+                return createDataStream();
+            case "destroy-data-stream":
+                return destroyDataStream();
+            case "provide-data-stream":
+                return provideDataStream();
+            case "write-to-data-stream":
+                return writeToDataStream();
+            case "provide-data":
+                return provideData();
+            case "get-last-status-code":
+                return getLastStatusCode();
+            case "get-bound-package":
+                return getBoundPackageName();
+            case "set-temporary-service":
+                return setTemporaryService();
+            default:
+                return handleDefaultCommands(cmd);
+        }
+    }
+
+    @Override
+    public void onHelp() {
+        PrintWriter pw = getOutPrintWriter();
+        pw.println("WearableSensingCommands commands: ");
+        pw.println("  help");
+        pw.println("    Print this help text.");
+        pw.println();
+        pw.println("  create-data-stream: Creates a data stream to be provided.");
+        pw.println("  destroy-data-stream: Destroys a data stream if one was previously created.");
+        pw.println("  provide-data-stream USER_ID: "
+                + "Provides data stream to WearableSensingService.");
+        pw.println("  write-to-data-stream STRING: writes string to data stream.");
+        pw.println("  provide-data USER_ID KEY INTEGER: provide integer as data with key.");
+        pw.println("  get-last-status-code: Prints the latest request status code.");
+        pw.println("  get-bound-package USER_ID:"
+                + "     Print the bound package that implements the service.");
+        pw.println("  set-temporary-service USER_ID [PACKAGE_NAME] [COMPONENT_NAME DURATION]");
+        pw.println("    Temporarily (for DURATION ms) changes the service implementation.");
+        pw.println("    To reset, call with just the USER_ID argument.");
+    }
+
+    private int createDataStream() {
+        Slog.d(TAG, "createDataStream");
+        try {
+            sPipe = ParcelFileDescriptor.createPipe();
+        } catch (IOException e) {
+            Slog.d(TAG, "Failed to createDataStream.", e);
+        }
+        return 0;
+    }
+
+    private int destroyDataStream() {
+        Slog.d(TAG, "destroyDataStream");
+        try {
+            if (sPipe != null) {
+                sPipe[0].close();
+                sPipe[1].close();
+            }
+        } catch (IOException e) {
+            Slog.d(TAG, "Failed to destroyDataStream.", e);
+        }
+        return 0;
+    }
+
+    private int provideDataStream() {
+        Slog.d(TAG, "provideDataStream");
+        if (sPipe != null) {
+            final int userId = Integer.parseInt(getNextArgRequired());
+            mService.provideDataStream(userId, sPipe[0],
+                    sTestableCallbackInternal.createRemoteStatusCallback());
+        }
+        return 0;
+    }
+
+    private int writeToDataStream() {
+        Slog.d(TAG, "writeToDataStream");
+        if (sPipe != null) {
+            final String value = getNextArgRequired();
+            try {
+                ParcelFileDescriptor writePipe = sPipe[1].dup();
+                OutputStream os = new ParcelFileDescriptor.AutoCloseOutputStream(writePipe);
+                os.write(value.getBytes());
+            } catch (IOException e) {
+                Slog.d(TAG, "Failed to writeToDataStream.", e);
+            }
+        }
+        return 0;
+    }
+
+    private int provideData() {
+        Slog.d(TAG, "provideData");
+        final int userId = Integer.parseInt(getNextArgRequired());
+        final String key = getNextArgRequired();
+        final int value = Integer.parseInt(getNextArgRequired());
+        PersistableBundle data = new PersistableBundle();
+        data.putInt(key, value);
+
+        mService.provideData(userId, data, null,
+                sTestableCallbackInternal.createRemoteStatusCallback());
+        return 0;
+    }
+
+    private int getLastStatusCode() {
+        Slog.d(TAG, "getLastStatusCode");
+        final PrintWriter resultPrinter = getOutPrintWriter();
+        int lastStatus = sTestableCallbackInternal.getLastStatus();
+        resultPrinter.println(lastStatus);
+        return 0;
+    }
+
+    private int setTemporaryService() {
+        final PrintWriter out = getOutPrintWriter();
+        final int userId = Integer.parseInt(getNextArgRequired());
+        final String serviceName = getNextArg();
+        if (serviceName == null) {
+            mService.resetTemporaryService(userId);
+            out.println("WearableSensingManagerService temporary reset. ");
+            return 0;
+        }
+
+        final int duration = Integer.parseInt(getNextArgRequired());
+        mService.setTemporaryService(userId, serviceName, duration);
+        out.println("WearableSensingService temporarily set to " + serviceName
+                + " for " + duration + "ms");
+        return 0;
+    }
+
+    private int getBoundPackageName() {
+        final PrintWriter resultPrinter = getOutPrintWriter();
+        final int userId = Integer.parseInt(getNextArgRequired());
+        final ComponentName componentName = mService.getComponentName(userId);
+        resultPrinter.println(componentName == null ? "" : componentName.getPackageName());
+        return 0;
+    }
+}
diff --git a/services/java/com/android/server/SystemServer.java b/services/java/com/android/server/SystemServer.java
index c346b2f..e96eff21 100644
--- a/services/java/com/android/server/SystemServer.java
+++ b/services/java/com/android/server/SystemServer.java
@@ -211,6 +211,7 @@
 import com.android.server.utils.TimingsTraceAndSlog;
 import com.android.server.vibrator.VibratorManagerService;
 import com.android.server.vr.VrManagerService;
+import com.android.server.wearable.WearableSensingManagerService;
 import com.android.server.webkit.WebViewUpdateService;
 import com.android.server.wm.ActivityTaskManagerService;
 import com.android.server.wm.WindowManagerGlobalLock;
@@ -1830,6 +1831,7 @@
             startSystemCaptionsManagerService(context, t);
             startTextToSpeechManagerService(context, t);
             startAmbientContextService(t);
+            startWearableSensingService(t);
 
             // System Speech Recognition Service
             t.traceBegin("StartSpeechRecognitionManagerService");
@@ -3170,6 +3172,12 @@
         t.traceEnd();
     }
 
+    private void startWearableSensingService(@NonNull TimingsTraceAndSlog t) {
+        t.traceBegin("startWearableSensingService");
+        mSystemServiceManager.startService(WearableSensingManagerService.class);
+        t.traceEnd();
+    }
+
     private static void startSystemUi(Context context, WindowManagerService windowManager) {
         PackageManagerInternal pm = LocalServices.getService(PackageManagerInternal.class);
         Intent intent = new Intent();
diff --git a/services/tests/servicestests/src/com/android/server/locksettings/BaseLockSettingsServiceTests.java b/services/tests/servicestests/src/com/android/server/locksettings/BaseLockSettingsServiceTests.java
index 55ab4a0..15eaf5d 100644
--- a/services/tests/servicestests/src/com/android/server/locksettings/BaseLockSettingsServiceTests.java
+++ b/services/tests/servicestests/src/com/android/server/locksettings/BaseLockSettingsServiceTests.java
@@ -36,7 +36,7 @@
 import android.content.Context;
 import android.content.pm.PackageManager;
 import android.content.pm.UserInfo;
-import android.hardware.authsecret.V1_0.IAuthSecret;
+import android.hardware.authsecret.IAuthSecret;
 import android.hardware.face.FaceManager;
 import android.hardware.fingerprint.FingerprintManager;
 import android.os.FileUtils;
diff --git a/services/tests/servicestests/src/com/android/server/locksettings/LockSettingsServiceTestable.java b/services/tests/servicestests/src/com/android/server/locksettings/LockSettingsServiceTestable.java
index eccfa06..d3b647d 100644
--- a/services/tests/servicestests/src/com/android/server/locksettings/LockSettingsServiceTestable.java
+++ b/services/tests/servicestests/src/com/android/server/locksettings/LockSettingsServiceTestable.java
@@ -22,7 +22,7 @@
 import android.app.admin.DeviceStateCache;
 import android.content.Context;
 import android.content.pm.UserInfo;
-import android.hardware.authsecret.V1_0.IAuthSecret;
+import android.hardware.authsecret.IAuthSecret;
 import android.os.Handler;
 import android.os.Parcel;
 import android.os.Process;
@@ -154,7 +154,7 @@
                 storageManager, spManager, gsiService, recoverableKeyStoreManager,
                 userManagerInternal, deviceStateCache));
         mGateKeeperService = gatekeeper;
-        mAuthSecretService = authSecretService;
+        mAuthSecretServiceAidl = authSecretService;
     }
 
     @Override
@@ -199,4 +199,4 @@
         UserInfo userInfo = mUserManager.getUserInfo(userId);
         return userInfo.isCloneProfile() || userInfo.isManagedProfile();
     }
-}
\ No newline at end of file
+}
diff --git a/services/tests/servicestests/src/com/android/server/locksettings/SyntheticPasswordTests.java b/services/tests/servicestests/src/com/android/server/locksettings/SyntheticPasswordTests.java
index b9cafa4..a441824 100644
--- a/services/tests/servicestests/src/com/android/server/locksettings/SyntheticPasswordTests.java
+++ b/services/tests/servicestests/src/com/android/server/locksettings/SyntheticPasswordTests.java
@@ -59,6 +59,7 @@
 
 import java.io.File;
 import java.util.ArrayList;
+import java.util.Arrays;
 
 
 /**
@@ -229,9 +230,11 @@
                 badPassword, PRIMARY_USER_ID, 0 /* flags */).getResponseCode());
 
         // Check the same secret was passed each time
-        ArgumentCaptor<ArrayList<Byte>> secret = ArgumentCaptor.forClass(ArrayList.class);
-        verify(mAuthSecretService, atLeastOnce()).primaryUserCredential(secret.capture());
-        assertEquals(1, secret.getAllValues().stream().distinct().count());
+        ArgumentCaptor<byte[]> secret = ArgumentCaptor.forClass(byte[].class);
+        verify(mAuthSecretService, atLeastOnce()).setPrimaryUserCredential(secret.capture());
+        for (byte[] val : secret.getAllValues()) {
+          assertArrayEquals(val, secret.getAllValues().get(0));
+        }
     }
 
     @Test
@@ -242,7 +245,7 @@
         reset(mAuthSecretService);
         assertEquals(VerifyCredentialResponse.RESPONSE_OK, mService.verifyCredential(
                 password, PRIMARY_USER_ID, 0 /* flags */).getResponseCode());
-        verify(mAuthSecretService).primaryUserCredential(any(ArrayList.class));
+        verify(mAuthSecretService).setPrimaryUserCredential(any(byte[].class));
     }
 
     @Test
@@ -252,7 +255,7 @@
         initializeCredential(password, SECONDARY_USER_ID);
         assertEquals(VerifyCredentialResponse.RESPONSE_OK, mService.verifyCredential(
                 password, SECONDARY_USER_ID, 0 /* flags */).getResponseCode());
-        verify(mAuthSecretService, never()).primaryUserCredential(any(ArrayList.class));
+        verify(mAuthSecretService, never()).setPrimaryUserCredential(any(byte[].class));
     }
 
     @Test
@@ -263,7 +266,7 @@
 
         reset(mAuthSecretService);
         mLocalService.unlockUserKeyIfUnsecured(PRIMARY_USER_ID);
-        verify(mAuthSecretService).primaryUserCredential(any(ArrayList.class));
+        verify(mAuthSecretService).setPrimaryUserCredential(any(byte[].class));
     }
 
     @Test
@@ -591,7 +594,7 @@
         initializeCredential(password, PRIMARY_USER_ID);
         assertEquals(VerifyCredentialResponse.RESPONSE_OK, mService.verifyCredential(
                 password, PRIMARY_USER_ID, 0 /* flags */).getResponseCode());
-        verify(mAuthSecretService, never()).primaryUserCredential(any(ArrayList.class));
+        verify(mAuthSecretService, never()).setPrimaryUserCredential(any(byte[].class));
     }
 
     @Test