diff --git a/api/Android.bp b/api/Android.bp
index cd1997c..6a04f0d 100644
--- a/api/Android.bp
+++ b/api/Android.bp
@@ -368,8 +368,6 @@
         "--hide CallbackInterface",
         // Disable HiddenSuperclass, as Metalava handles this fine (it should be hidden by default)
         "--hide HiddenSuperclass",
-        "--hide-package android.audio.policy.configuration.V7_0",
-        "--hide-package com.android.server",
         "--manifest $(location :frameworks-base-core-AndroidManifest.xml)",
     ],
     filter_packages: packages_to_document,
diff --git a/core/api/current.txt b/core/api/current.txt
index f2c59dac..889b627 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -6844,6 +6844,9 @@
     method public android.app.Notification.MessagingStyle.Message setData(String, android.net.Uri);
   }
 
+  @FlaggedApi("android.app.api_rich_ongoing") public abstract static class Notification.RichOngoingStyle extends android.app.Notification.Style {
+  }
+
   public abstract static class Notification.Style {
     ctor @Deprecated public Notification.Style();
     method public android.app.Notification build();
diff --git a/core/java/android/app/Notification.java b/core/java/android/app/Notification.java
index db979a5..ef09dc4 100644
--- a/core/java/android/app/Notification.java
+++ b/core/java/android/app/Notification.java
@@ -10979,6 +10979,18 @@
     }
 
     /**
+     * An object that can apply a rich ongoing notification style to a {@link Notification.Builder}
+     * object.
+     */
+    @FlaggedApi(Flags.FLAG_API_RICH_ONGOING)
+    public abstract static class RichOngoingStyle extends Notification.Style {
+        /**
+         * @hide
+         */
+        public RichOngoingStyle() {}
+    }
+
+    /**
      * Notification style for custom views that are decorated by the system
      *
      * <p>Instead of providing a notification that is completely custom, a developer can set this
diff --git a/core/java/android/audio/policy/configuration/V7_0/package-info.java b/core/java/android/audio/policy/configuration/V7_0/package-info.java
new file mode 100644
index 0000000..8f7425f
--- /dev/null
+++ b/core/java/android/audio/policy/configuration/V7_0/package-info.java
@@ -0,0 +1,23 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * Hide the android.audio.policy.configuration.V7_0 API as that is managed
+ * separately.
+ *
+ * @hide
+ */
+package android.audio.policy.configuration.V7_0;
diff --git a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
index 6abba8b..50fb8d5 100644
--- a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
@@ -104,8 +104,8 @@
     private final Runnable mCacheUpdater;
 
     private final int[] mDefaultLogLevelCounts = new int[LogLevel.values().length];
-    private final Map<IProtoLogGroup, int[]> mLogLevelCounts = new ArrayMap<>();
-    private final Map<IProtoLogGroup, Integer> mCollectStackTraceGroupCounts = new ArrayMap<>();
+    private final Map<String, int[]> mLogLevelCounts = new ArrayMap<>();
+    private final Map<String, Integer> mCollectStackTraceGroupCounts = new ArrayMap<>();
 
     private final Lock mBackgroundServiceLock = new ReentrantLock();
     private ExecutorService mBackgroundLoggingService = Executors.newSingleThreadExecutor();
@@ -226,7 +226,7 @@
 
     @Override
     public boolean isEnabled(IProtoLogGroup group, LogLevel level) {
-        final int[] groupLevelCount = mLogLevelCounts.get(group);
+        final int[] groupLevelCount = mLogLevelCounts.get(group.name());
         return (groupLevelCount == null && mDefaultLogLevelCounts[level.ordinal()] > 0)
                 || (groupLevelCount != null && groupLevelCount[level.ordinal()] > 0)
                 || group.isLogToLogcat();
@@ -279,7 +279,7 @@
         if (isProtoEnabled()) {
             long tsNanos = SystemClock.elapsedRealtimeNanos();
             final String stacktrace;
-            if (mCollectStackTraceGroupCounts.getOrDefault(group, 0) > 0) {
+            if (mCollectStackTraceGroupCounts.getOrDefault(group.name(), 0) > 0) {
                 stacktrace = collectStackTrace();
             } else {
                 stacktrace = null;
@@ -739,15 +739,8 @@
         final Set<String> overriddenGroupTags = config.getGroupTagsWithOverriddenConfigs();
 
         for (String overriddenGroupTag : overriddenGroupTags) {
-            IProtoLogGroup group = mLogGroups.get(overriddenGroupTag);
-
-            if (group == null) {
-                throw new IllegalArgumentException("Trying to set config for \""
-                        + overriddenGroupTag + "\" that isn't registered");
-            }
-
-            mLogLevelCounts.putIfAbsent(group, new int[LogLevel.values().length]);
-            final int[] logLevelsCountsForGroup = mLogLevelCounts.get(group);
+            mLogLevelCounts.putIfAbsent(overriddenGroupTag, new int[LogLevel.values().length]);
+            final int[] logLevelsCountsForGroup = mLogLevelCounts.get(overriddenGroupTag);
 
             final LogLevel logFromLevel = config.getConfigFor(overriddenGroupTag).logFrom;
             for (int i = logFromLevel.ordinal(); i < LogLevel.values().length; i++) {
@@ -755,13 +748,13 @@
             }
 
             if (config.getConfigFor(overriddenGroupTag).collectStackTrace) {
-                mCollectStackTraceGroupCounts.put(group,
-                        mCollectStackTraceGroupCounts.getOrDefault(group, 0) + 1);
+                mCollectStackTraceGroupCounts.put(overriddenGroupTag,
+                        mCollectStackTraceGroupCounts.getOrDefault(overriddenGroupTag, 0) + 1);
             }
 
             if (config.getConfigFor(overriddenGroupTag).collectStackTrace) {
-                mCollectStackTraceGroupCounts.put(group,
-                        mCollectStackTraceGroupCounts.getOrDefault(group, 0) + 1);
+                mCollectStackTraceGroupCounts.put(overriddenGroupTag,
+                        mCollectStackTraceGroupCounts.getOrDefault(overriddenGroupTag, 0) + 1);
             }
         }
 
@@ -781,24 +774,22 @@
         final Set<String> overriddenGroupTags = config.getGroupTagsWithOverriddenConfigs();
 
         for (String overriddenGroupTag : overriddenGroupTags) {
-            IProtoLogGroup group = mLogGroups.get(overriddenGroupTag);
-
-            final int[] logLevelsCountsForGroup = mLogLevelCounts.get(group);
+            final int[] logLevelsCountsForGroup = mLogLevelCounts.get(overriddenGroupTag);
 
             final LogLevel logFromLevel = config.getConfigFor(overriddenGroupTag).logFrom;
-            for (int i = defaultLogFrom.ordinal(); i < LogLevel.values().length; i++) {
+            for (int i = logFromLevel.ordinal(); i < LogLevel.values().length; i++) {
                 logLevelsCountsForGroup[i]--;
             }
             if (Arrays.stream(logLevelsCountsForGroup).allMatch(it -> it == 0)) {
-                mLogLevelCounts.remove(group);
+                mLogLevelCounts.remove(overriddenGroupTag);
             }
 
             if (config.getConfigFor(overriddenGroupTag).collectStackTrace) {
-                mCollectStackTraceGroupCounts.put(group,
-                        mCollectStackTraceGroupCounts.get(group) - 1);
+                mCollectStackTraceGroupCounts.put(overriddenGroupTag,
+                        mCollectStackTraceGroupCounts.get(overriddenGroupTag) - 1);
 
-                if (mCollectStackTraceGroupCounts.get(group) == 0) {
-                    mCollectStackTraceGroupCounts.remove(group);
+                if (mCollectStackTraceGroupCounts.get(overriddenGroupTag) == 0) {
+                    mCollectStackTraceGroupCounts.remove(overriddenGroupTag);
                 }
             }
         }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index 18157d6..dc3e2d0 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -159,18 +159,6 @@
             }
         }
 
-    private val transitionAreaHeight
-        get() =
-            context.resources.getDimensionPixelSize(
-                com.android.wm.shell.R.dimen.desktop_mode_fullscreen_from_desktop_height
-            )
-
-    private val transitionAreaWidth
-        get() =
-            context.resources.getDimensionPixelSize(
-                com.android.wm.shell.R.dimen.desktop_mode_transition_area_width
-            )
-
     /** Task id of the task currently being dragged from fullscreen/split. */
     val draggingTaskId
         get() = dragToDesktopTransitionHandler.draggingTaskId
@@ -776,12 +764,15 @@
             newTaskIdInFront ?: "null"
         )
 
-        if (Flags.enableDesktopWindowingWallpaperActivity()) {
-            // Add translucent wallpaper activity to show the wallpaper underneath
-            addWallpaperActivity(wct)
-        } else {
-            // Move home to front
-            moveHomeTask(wct, toTop = true)
+        // Currently, we only handle the desktop on the default display really.
+        if (displayId == DEFAULT_DISPLAY) {
+            if (Flags.enableDesktopWindowingWallpaperActivity()) {
+                // Add translucent wallpaper activity to show the wallpaper underneath
+                addWallpaperActivity(wct)
+            } else {
+                // Move home to front
+                moveHomeTask(wct, toTop = true)
+            }
         }
 
         val nonMinimizedTasksOrderedFrontToBack =
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
index 6cabbf9..4d1b6ba 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
@@ -311,6 +311,23 @@
   }
 
   @Test
+  @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+  fun showDesktopApps_onSecondaryDisplay_desktopWallpaperEnabled_shouldNotShowWallpaper() {
+    val task1 = setUpFreeformTask(SECOND_DISPLAY)
+    val task2 = setUpFreeformTask(SECOND_DISPLAY)
+    markTaskHidden(task1)
+    markTaskHidden(task2)
+
+    controller.showDesktopApps(SECOND_DISPLAY, RemoteTransition(TestRemoteTransition()))
+
+    val wct = getLatestWct(type = TRANSIT_TO_FRONT, handlerClass = OneShotRemoteHandler::class.java)
+    assertThat(wct.hierarchyOps).hasSize(2)
+    // Expect order to be from bottom: task1, task2 (no wallpaper intent)
+    wct.assertReorderAt(index = 0, task1)
+    wct.assertReorderAt(index = 1, task2)
+  }
+
+  @Test
   @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
   fun showDesktopApps_appsAlreadyVisible_bringsToFront_desktopWallpaperDisabled() {
     val homeTask = setUpHomeTask()
@@ -330,6 +347,22 @@
   }
 
   @Test
+  @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+  fun showDesktopApps_onSecondaryDisplay_desktopWallpaperDisabled_shouldNotMoveLauncher() {
+    val task1 = setUpFreeformTask(SECOND_DISPLAY)
+    val task2 = setUpFreeformTask(SECOND_DISPLAY)
+    markTaskHidden(task1)
+    markTaskHidden(task2)
+
+    controller.showDesktopApps(SECOND_DISPLAY, RemoteTransition(TestRemoteTransition()))
+
+    val wct = getLatestWct(type = TRANSIT_TO_FRONT, handlerClass = OneShotRemoteHandler::class.java)
+    assertThat(wct.hierarchyOps).hasSize(2)
+    wct.assertReorderAt(index = 0, task1)
+    wct.assertReorderAt(index = 1, task2)
+  }
+
+  @Test
   @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
   fun showDesktopApps_appsAlreadyVisible_bringsToFront_desktopWallpaperEnabled() {
     val task1 = setUpFreeformTask()
diff --git a/nfc/java/android/nfc/NfcAdapter.java b/nfc/java/android/nfc/NfcAdapter.java
index 9ce1c82..395f81d 100644
--- a/nfc/java/android/nfc/NfcAdapter.java
+++ b/nfc/java/android/nfc/NfcAdapter.java
@@ -963,22 +963,9 @@
             throw new UnsupportedOperationException("You need a context on NfcAdapter to use the "
                     + " NFC extras APIs");
         }
-        try {
-            return sService.getNfcDtaInterface(mContext.getPackageName());
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return null;
-            }
-            try {
-                return sService.getNfcDtaInterface(mContext.getPackageName());
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return null;
-        }
+        return callServiceReturn(() ->  sService.getNfcDtaInterface(mContext.getPackageName()),
+                null);
+
     }
 
     /**
@@ -1095,22 +1082,8 @@
     @SystemApi
     @FlaggedApi(Flags.FLAG_ENABLE_NFC_MAINLINE)
     public @AdapterState int getAdapterState() {
-        try {
-            return sService.getState();
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return NfcAdapter.STATE_OFF;
-            }
-            try {
-                return sService.getState();
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return NfcAdapter.STATE_OFF;
-        }
+        return callServiceReturn(() ->  sService.getState(), NfcAdapter.STATE_OFF);
+
     }
 
     /**
@@ -1134,22 +1107,8 @@
     @FlaggedApi(Flags.FLAG_NFC_STATE_CHANGE)
     @RequiresPermission(android.Manifest.permission.WRITE_SECURE_SETTINGS)
     public boolean enable() {
-        try {
-            return sService.enable(mContext.getPackageName());
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return false;
-            }
-            try {
-                return sService.enable(mContext.getPackageName());
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return false;
-        }
+        return callServiceReturn(() ->  sService.enable(mContext.getPackageName()), false);
+
     }
 
     /**
@@ -1175,22 +1134,9 @@
     @FlaggedApi(Flags.FLAG_NFC_STATE_CHANGE)
     @RequiresPermission(android.Manifest.permission.WRITE_SECURE_SETTINGS)
     public boolean disable() {
-        try {
-            return sService.disable(true, mContext.getPackageName());
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return false;
-            }
-            try {
-                return sService.disable(true, mContext.getPackageName());
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return false;
-        }
+        return callServiceReturn(() ->  sService.disable(true, mContext.getPackageName()),
+                false);
+
     }
 
     /**
@@ -1200,22 +1146,9 @@
     @SystemApi
     @RequiresPermission(android.Manifest.permission.WRITE_SECURE_SETTINGS)
     public boolean disable(boolean persist) {
-        try {
-            return sService.disable(persist, mContext.getPackageName());
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return false;
-            }
-            try {
-                return sService.disable(persist, mContext.getPackageName());
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return false;
-        }
+        return callServiceReturn(() ->  sService.disable(persist, mContext.getPackageName()),
+                false);
+
     }
 
     /**
@@ -1241,12 +1174,7 @@
      */
     @FlaggedApi(Flags.FLAG_NFC_OBSERVE_MODE)
     public boolean isObserveModeSupported() {
-        try {
-            return sService.isObserveModeSupported();
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            return false;
-        }
+        return callServiceReturn(() ->  sService.isObserveModeSupported(), false);
     }
 
     /**
@@ -1257,12 +1185,7 @@
 
     @FlaggedApi(Flags.FLAG_NFC_OBSERVE_MODE)
     public boolean isObserveModeEnabled() {
-        try {
-            return sService.isObserveModeEnabled();
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            return false;
-        }
+        return callServiceReturn(() ->  sService.isObserveModeEnabled(), false);
     }
 
     /**
@@ -1286,12 +1209,8 @@
             throw new UnsupportedOperationException("You need a context on NfcAdapter to use the "
                     + " observe mode APIs");
         }
-        try {
-            return sService.setObserveMode(enabled, mContext.getPackageName());
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            return false;
-        }
+        return callServiceReturn(() ->  sService.setObserveMode(enabled, mContext.getPackageName()),
+                false);
     }
 
     /**
@@ -2057,22 +1976,8 @@
         if (!sHasNfcFeature && !sHasCeFeature) {
             throw new UnsupportedOperationException();
         }
-        try {
-            return sService.setNfcSecure(enable);
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return false;
-            }
-            try {
-                return sService.setNfcSecure(enable);
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return false;
-        }
+        return callServiceReturn(() ->  sService.setNfcSecure(enable), false);
+
     }
 
     /**
@@ -2088,22 +1993,8 @@
         if (!sHasNfcFeature && !sHasCeFeature) {
             throw new UnsupportedOperationException();
         }
-        try {
-            return sService.deviceSupportsNfcSecure();
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return false;
-            }
-            try {
-                return sService.deviceSupportsNfcSecure();
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return false;
-        }
+        return callServiceReturn(() ->  sService.deviceSupportsNfcSecure(), false);
+
     }
 
     /**
@@ -2121,22 +2012,8 @@
         if (!sHasNfcFeature && !sHasCeFeature) {
             throw new UnsupportedOperationException();
         }
-        try {
-            return sService.getNfcAntennaInfo();
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return null;
-            }
-            try {
-                return sService.getNfcAntennaInfo();
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return null;
-        }
+        return callServiceReturn(() ->  sService.getNfcAntennaInfo(), null);
+
     }
 
     /**
@@ -2154,22 +2031,8 @@
         if (!sHasNfcFeature && !sHasCeFeature) {
             throw new UnsupportedOperationException();
         }
-        try {
-            return sService.isNfcSecureEnabled();
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return false;
-            }
-            try {
-                return sService.isNfcSecureEnabled();
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return false;
-        }
+        return callServiceReturn(() ->  sService.isNfcSecureEnabled(), false);
+
     }
 
     /**
@@ -2185,22 +2048,8 @@
         if (!sHasNfcFeature) {
             throw new UnsupportedOperationException();
         }
-        try {
-            return sService.enableReaderOption(enable);
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return false;
-            }
-            try {
-                return sService.enableReaderOption(enable);
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return false;
-        }
+        return callServiceReturn(() ->  sService.enableReaderOption(enable), false);
+
     }
 
     /**
@@ -2214,22 +2063,8 @@
         if (!sHasNfcFeature) {
             throw new UnsupportedOperationException();
         }
-        try {
-            return sService.isReaderOptionSupported();
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return false;
-            }
-            try {
-                return sService.isReaderOptionSupported();
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return false;
-        }
+        return callServiceReturn(() ->  sService.isReaderOptionSupported(), false);
+
     }
 
     /**
@@ -2245,22 +2080,8 @@
         if (!sHasNfcFeature) {
             throw new UnsupportedOperationException();
         }
-        try {
-            return sService.isReaderOptionEnabled();
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return false;
-            }
-            try {
-                return sService.isReaderOptionEnabled();
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return false;
-        }
+        return callServiceReturn(() ->  sService.isReaderOptionEnabled(), false);
+
     }
 
     /**
@@ -2388,11 +2209,9 @@
         synchronized (mLock) {
             mTagRemovedListener = iListener;
         }
-        try {
-            return sService.ignore(tag.getServiceHandle(), debounceMs, iListener);
-        } catch (RemoteException e) {
-            return false;
-        }
+        final ITagRemovedCallback.Stub passedListener = iListener;
+        return callServiceReturn(() ->
+                sService.ignore(tag.getServiceHandle(), debounceMs, passedListener), false);
     }
 
     /**
@@ -2509,22 +2328,9 @@
             throw new UnsupportedOperationException("You need a context on NfcAdapter to use the "
                     + " NFC extras APIs");
         }
-        try {
-            return sService.getNfcAdapterExtrasInterface(mContext.getPackageName());
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return null;
-            }
-            try {
-                return sService.getNfcAdapterExtrasInterface(mContext.getPackageName());
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return null;
-        }
+        return callServiceReturn(() ->
+                sService.getNfcAdapterExtrasInterface(mContext.getPackageName()), null);
+
     }
 
     void enforceResumed(Activity activity) {
@@ -2569,22 +2375,8 @@
         if (!sHasNfcFeature && !sHasCeFeature) {
             throw new UnsupportedOperationException();
         }
-        try {
-            return sService.setControllerAlwaysOn(value);
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return false;
-            }
-            try {
-                return sService.setControllerAlwaysOn(value);
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return false;
-        }
+        return callServiceReturn(() ->  sService.setControllerAlwaysOn(value), false);
+
     }
 
     /**
@@ -2600,22 +2392,8 @@
     @SystemApi
     @RequiresPermission(android.Manifest.permission.NFC_SET_CONTROLLER_ALWAYS_ON)
     public boolean isControllerAlwaysOn() {
-        try {
-            return sService.isControllerAlwaysOn();
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return false;
-            }
-            try {
-                return sService.isControllerAlwaysOn();
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return false;
-        }
+        return callServiceReturn(() ->  sService.isControllerAlwaysOn(), false);
+
     }
 
     /**
@@ -2634,22 +2412,8 @@
         if (!sHasNfcFeature && !sHasCeFeature) {
             throw new UnsupportedOperationException();
         }
-        try {
-            return sService.isControllerAlwaysOnSupported();
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return false;
-            }
-            try {
-                return sService.isControllerAlwaysOnSupported();
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return false;
-        }
+        return callServiceReturn(() ->  sService.isControllerAlwaysOnSupported(), false);
+
     }
 
     /**
@@ -2719,21 +2483,9 @@
             Log.e(TAG, "TagIntentAppPreference is not supported");
             throw new UnsupportedOperationException();
         }
-        try {
-            return sService.setTagIntentAppPreferenceForUser(userId, pkg, allow);
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            try {
-                return sService.setTagIntentAppPreferenceForUser(userId, pkg, allow);
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return TAG_INTENT_APP_PREF_RESULT_UNAVAILABLE;
-        }
+        return callServiceReturn(() ->
+                sService.setTagIntentAppPreferenceForUser(userId, pkg, allow),
+                        TAG_INTENT_APP_PREF_RESULT_UNAVAILABLE);
     }
 
 
@@ -2808,22 +2560,8 @@
         if (!sHasNfcFeature) {
             throw new UnsupportedOperationException();
         }
-        try {
-            return sService.isTagIntentAppPreferenceSupported();
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return false;
-            }
-            try {
-                return sService.isTagIntentAppPreferenceSupported();
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return false;
-        }
+        return callServiceReturn(() ->  sService.isTagIntentAppPreferenceSupported(), false);
+
     }
 
    /**
@@ -2836,26 +2574,10 @@
     @TestApi
     @FlaggedApi(Flags.FLAG_NFC_READ_POLLING_LOOP)
     public void notifyPollingLoop(@NonNull PollingFrame pollingFrame) {
-        try {
-            if (sService == null) {
-                attemptDeadServiceRecovery(null);
-            }
-            sService.notifyPollingLoop(pollingFrame);
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return;
-            }
-            try {
-                sService.notifyPollingLoop(pollingFrame);
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-        }
+        callService(() ->  sService.notifyPollingLoop(pollingFrame));
     }
 
+
    /**
      * Notifies the system of new HCE data for tests.
      *
@@ -2863,11 +2585,19 @@
      */
     @FlaggedApi(Flags.FLAG_NFC_READ_POLLING_LOOP)
     public void notifyTestHceData(int technology, byte[] data) {
+        callService(() ->  sService.notifyTestHceData(technology, data));
+    }
+
+    interface ServiceCall {
+        void call() throws RemoteException;
+    }
+
+    void callService(ServiceCall call) {
         try {
             if (sService == null) {
                 attemptDeadServiceRecovery(null);
             }
-            sService.notifyTestHceData(technology, data);
+            call.call();
         } catch (RemoteException e) {
             attemptDeadServiceRecovery(e);
             // Try one more time
@@ -2876,12 +2606,36 @@
                 return;
             }
             try {
-                sService.notifyTestHceData(technology, data);
-            } catch (RemoteException e2) {
+                call.call();
+            } catch (RemoteException ee) {
                 Log.e(TAG, "Failed to recover NFC Service.");
             }
         }
     }
+    interface ServiceCallReturn<T> {
+        T call() throws RemoteException;
+    }
+    <T> T callServiceReturn(ServiceCallReturn<T> call, T defaultReturn) {
+        try {
+            if (sService == null) {
+                attemptDeadServiceRecovery(null);
+            }
+            return call.call();
+        } catch (RemoteException e) {
+            attemptDeadServiceRecovery(e);
+            // Try one more time
+            if (sService == null) {
+                Log.e(TAG, "Failed to recover NFC Service.");
+                return defaultReturn;
+            }
+            try {
+                return call.call();
+            } catch (RemoteException ee) {
+                Log.e(TAG, "Failed to recover NFC Service.");
+            }
+        }
+        return defaultReturn;
+    }
 
    /**
      * Notifies the system of a an HCE session being deactivated.
@@ -2891,24 +2645,7 @@
     @TestApi
     @FlaggedApi(Flags.FLAG_NFC_READ_POLLING_LOOP)
     public void notifyHceDeactivated() {
-        try {
-            if (sService == null) {
-                attemptDeadServiceRecovery(null);
-            }
-            sService.notifyHceDeactivated();
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return;
-            }
-            try {
-                sService.notifyHceDeactivated();
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-        }
+        callService(() ->  sService.notifyHceDeactivated());
     }
 
     /**
@@ -2924,22 +2661,7 @@
         if (!sHasNfcWlcFeature) {
             throw new UnsupportedOperationException();
         }
-        try {
-            return sService.setWlcEnabled(enable);
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return false;
-            }
-            try {
-                return sService.setWlcEnabled(enable);
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return false;
-        }
+        return callServiceReturn(() ->  sService.setWlcEnabled(enable), false);
     }
 
     /**
@@ -2954,22 +2676,8 @@
         if (!sHasNfcWlcFeature) {
             throw new UnsupportedOperationException();
         }
-        try {
-            return sService.isWlcEnabled();
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return false;
-            }
-            try {
-                return sService.isWlcEnabled();
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return false;
-        }
+        return callServiceReturn(() ->  sService.isWlcEnabled(), false);
+
     }
 
     /**
@@ -3048,22 +2756,8 @@
         if (!sHasNfcWlcFeature) {
             throw new UnsupportedOperationException();
         }
-        try {
-            return sService.getWlcListenerDeviceInfo();
-        } catch (RemoteException e) {
-            attemptDeadServiceRecovery(e);
-            // Try one more time
-            if (sService == null) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-                return null;
-            }
-            try {
-                return sService.getWlcListenerDeviceInfo();
-            } catch (RemoteException ee) {
-                Log.e(TAG, "Failed to recover NFC Service.");
-            }
-            return null;
-        }
+        return callServiceReturn(() ->  sService.getWlcListenerDeviceInfo(), null);
+
     }
 
     /**
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/media/controls/ui/composable/MediaCarousel.kt b/packages/SystemUI/compose/features/src/com/android/systemui/media/controls/ui/composable/MediaCarousel.kt
index 581f3a5..d629eec 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/media/controls/ui/composable/MediaCarousel.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/media/controls/ui/composable/MediaCarousel.kt
@@ -16,8 +16,10 @@
 
 package com.android.systemui.media.controls.ui.composable
 
+import android.view.View
 import android.view.ViewGroup
 import android.widget.FrameLayout
+import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.height
 import androidx.compose.runtime.Composable
@@ -26,7 +28,6 @@
 import androidx.compose.ui.platform.LocalDensity
 import androidx.compose.ui.res.dimensionResource
 import androidx.compose.ui.viewinterop.AndroidView
-import androidx.core.view.contains
 import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.SceneScope
 import com.android.systemui.media.controls.ui.controller.MediaCarouselController
@@ -36,7 +37,8 @@
 
 private object MediaCarousel {
     object Elements {
-        internal val Content = ElementKey("MediaCarouselContent")
+        internal val Content =
+            ElementKey(debugName = "MediaCarouselContent", scenePicker = MediaScenePicker)
     }
 }
 
@@ -61,40 +63,43 @@
     mediaHost.measurementInput = MeasurementInput(layoutWidth, layoutHeight)
     carouselController.setSceneContainerSize(layoutWidth, layoutHeight)
 
-    AndroidView(
-        modifier =
-            modifier
-                .element(MediaCarousel.Elements.Content)
-                .height(mediaHeight)
-                .fillMaxWidth()
-                .layout { measurable, constraints ->
-                    val placeable = measurable.measure(constraints)
+    MovableElement(
+        key = MediaCarousel.Elements.Content,
+        modifier = modifier.height(mediaHeight).fillMaxWidth()
+    ) {
+        content {
+            AndroidView(
+                modifier =
+                    Modifier.fillMaxSize().layout { measurable, constraints ->
+                        val placeable = measurable.measure(constraints)
 
-                    // Notify controller to size the carousel for the current space
-                    mediaHost.measurementInput = MeasurementInput(placeable.width, placeable.height)
-                    carouselController.setSceneContainerSize(placeable.width, placeable.height)
+                        // Notify controller to size the carousel for the current space
+                        mediaHost.measurementInput =
+                            MeasurementInput(placeable.width, placeable.height)
+                        carouselController.setSceneContainerSize(placeable.width, placeable.height)
 
-                    layout(placeable.width, placeable.height) { placeable.placeRelative(0, 0) }
+                        layout(placeable.width, placeable.height) { placeable.placeRelative(0, 0) }
+                    },
+                factory = { context ->
+                    FrameLayout(context).apply {
+                        layoutParams =
+                            FrameLayout.LayoutParams(
+                                FrameLayout.LayoutParams.MATCH_PARENT,
+                                FrameLayout.LayoutParams.MATCH_PARENT,
+                            )
+                    }
                 },
-        factory = { context ->
-            FrameLayout(context).apply {
-                val mediaFrame = carouselController.mediaFrame
-                (mediaFrame.parent as? ViewGroup)?.removeView(mediaFrame)
-                addView(mediaFrame)
-                layoutParams =
-                    FrameLayout.LayoutParams(
-                        FrameLayout.LayoutParams.MATCH_PARENT,
-                        FrameLayout.LayoutParams.MATCH_PARENT,
-                    )
-            }
-        },
-        update = {
-            if (it.contains(carouselController.mediaFrame)) {
-                return@AndroidView
-            }
-            val mediaFrame = carouselController.mediaFrame
-            (mediaFrame.parent as? ViewGroup)?.removeView(mediaFrame)
-            it.addView(mediaFrame)
-        },
-    )
+                update = { it.setView(carouselController.mediaFrame) },
+                onRelease = { it.removeAllViews() }
+            )
+        }
+    }
+}
+
+private fun ViewGroup.setView(view: View) {
+    if (view.parent == this) {
+        return
+    }
+    (view.parent as? ViewGroup)?.removeView(view)
+    addView(view)
 }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/media/controls/ui/composable/MediaScenePicker.kt b/packages/SystemUI/compose/features/src/com/android/systemui/media/controls/ui/composable/MediaScenePicker.kt
new file mode 100644
index 0000000..0398133
--- /dev/null
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/media/controls/ui/composable/MediaScenePicker.kt
@@ -0,0 +1,72 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.media.controls.ui.composable
+
+import com.android.compose.animation.scene.ElementKey
+import com.android.compose.animation.scene.ElementScenePicker
+import com.android.compose.animation.scene.SceneKey
+import com.android.compose.animation.scene.TransitionState
+import com.android.systemui.scene.shared.model.Scenes
+
+/** [ElementScenePicker] implementation for the media carousel object. */
+object MediaScenePicker : ElementScenePicker {
+
+    private val shadeLockscreenFraction = 0.65f
+    private val scenes =
+        setOf(
+            Scenes.Lockscreen,
+            Scenes.Shade,
+            Scenes.QuickSettings,
+            Scenes.QuickSettingsShade,
+            Scenes.Communal
+        )
+
+    override fun sceneDuringTransition(
+        element: ElementKey,
+        transition: TransitionState.Transition,
+        fromSceneZIndex: Float,
+        toSceneZIndex: Float
+    ): SceneKey? {
+        return when {
+            // TODO: 352052894 - update with the actual scene picking
+            transition.isTransitioning(from = Scenes.Lockscreen, to = Scenes.Shade) -> {
+                if (transition.progress < shadeLockscreenFraction) {
+                    Scenes.Lockscreen
+                } else {
+                    Scenes.Shade
+                }
+            }
+
+            // TODO: 345467290 - update with the actual scene picking
+            transition.isTransitioning(from = Scenes.Shade, to = Scenes.Lockscreen) -> {
+                if (transition.progress < 1f - shadeLockscreenFraction) {
+                    Scenes.Shade
+                } else {
+                    Scenes.Lockscreen
+                }
+            }
+
+            // TODO: 345467290 - update with the actual scene picking
+            transition.isTransitioningBetween(Scenes.QuickSettings, Scenes.Shade) -> {
+                Scenes.QuickSettings
+            }
+
+            // TODO: 340216785 - update with the actual scene picking
+            else -> pickSingleSceneIn(scenes, transition, element)
+        }
+    }
+}
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MultiPointerDraggable.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MultiPointerDraggable.kt
index cdcfc84..615d393 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MultiPointerDraggable.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MultiPointerDraggable.kt
@@ -241,43 +241,50 @@
                 }
             }
 
-            awaitPointerEventScope {
-                while (isActive) {
-                    try {
-                        detectDragGestures(
-                            orientation = orientation,
-                            startDragImmediately = startDragImmediately,
-                            onDragStart = { startedPosition, overSlop, pointersDown ->
-                                velocityTracker.resetTracking()
-                                onDragStarted(startedPosition, overSlop, pointersDown)
-                            },
-                            onDrag = { controller, change, amount ->
-                                velocityTracker.addPointerInputChange(change)
-                                controller.onDrag(amount)
-                            },
-                            onDragEnd = { controller ->
-                                val viewConfiguration = currentValueOf(LocalViewConfiguration)
-                                val maxVelocity =
-                                    viewConfiguration.maximumFlingVelocity.let { Velocity(it, it) }
-                                val velocity = velocityTracker.calculateVelocity(maxVelocity)
-                                controller.onStop(
-                                    velocity =
-                                        when (orientation) {
-                                            Orientation.Horizontal -> velocity.x
-                                            Orientation.Vertical -> velocity.y
-                                        },
-                                    canChangeScene = true,
-                                )
-                            },
-                            onDragCancel = { controller ->
-                                controller.onStop(velocity = 0f, canChangeScene = true)
-                            },
-                            swipeDetector = swipeDetector
-                        )
-                    } catch (exception: CancellationException) {
-                        // If the coroutine scope is active, we can just restart the drag cycle.
-                        if (!isActive) {
-                            throw exception
+            // The order is important here: we want to make sure that the previous PointerEventScope
+            // is initialized first. This ensures that the following PointerEventScope doesn't
+            // receive more events than the first one.
+            launch {
+                awaitPointerEventScope {
+                    while (isActive) {
+                        try {
+                            detectDragGestures(
+                                orientation = orientation,
+                                startDragImmediately = startDragImmediately,
+                                onDragStart = { startedPosition, overSlop, pointersDown ->
+                                    velocityTracker.resetTracking()
+                                    onDragStarted(startedPosition, overSlop, pointersDown)
+                                },
+                                onDrag = { controller, change, amount ->
+                                    velocityTracker.addPointerInputChange(change)
+                                    controller.onDrag(amount)
+                                },
+                                onDragEnd = { controller ->
+                                    val viewConfiguration = currentValueOf(LocalViewConfiguration)
+                                    val maxVelocity =
+                                        viewConfiguration.maximumFlingVelocity.let {
+                                            Velocity(it, it)
+                                        }
+                                    val velocity = velocityTracker.calculateVelocity(maxVelocity)
+                                    controller.onStop(
+                                        velocity =
+                                            when (orientation) {
+                                                Orientation.Horizontal -> velocity.x
+                                                Orientation.Vertical -> velocity.y
+                                            },
+                                        canChangeScene = true,
+                                    )
+                                },
+                                onDragCancel = { controller ->
+                                    controller.onStop(velocity = 0f, canChangeScene = true)
+                                },
+                                swipeDetector = swipeDetector
+                            )
+                        } catch (exception: CancellationException) {
+                            // If the coroutine scope is active, we can just restart the drag cycle.
+                            if (!isActive) {
+                                throw exception
+                            }
                         }
                     }
                 }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PunchHole.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PunchHole.kt
deleted file mode 100644
index b346a70..0000000
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PunchHole.kt
+++ /dev/null
@@ -1,178 +0,0 @@
-/*
- * Copyright (C) 2023 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.compose.animation.scene
-
-import androidx.compose.runtime.Stable
-import androidx.compose.ui.Modifier
-import androidx.compose.ui.geometry.Offset
-import androidx.compose.ui.geometry.Size
-import androidx.compose.ui.graphics.BlendMode
-import androidx.compose.ui.graphics.Color
-import androidx.compose.ui.graphics.CompositingStrategy
-import androidx.compose.ui.graphics.Outline
-import androidx.compose.ui.graphics.RectangleShape
-import androidx.compose.ui.graphics.Shape
-import androidx.compose.ui.graphics.drawOutline
-import androidx.compose.ui.graphics.drawscope.ContentDrawScope
-import androidx.compose.ui.graphics.drawscope.DrawScope
-import androidx.compose.ui.graphics.drawscope.translate
-import androidx.compose.ui.layout.LayoutCoordinates
-import androidx.compose.ui.layout.Measurable
-import androidx.compose.ui.layout.MeasureResult
-import androidx.compose.ui.layout.MeasureScope
-import androidx.compose.ui.node.DelegatingNode
-import androidx.compose.ui.node.DrawModifierNode
-import androidx.compose.ui.node.GlobalPositionAwareModifierNode
-import androidx.compose.ui.node.LayoutModifierNode
-import androidx.compose.ui.node.ModifierNodeElement
-import androidx.compose.ui.unit.Constraints
-import androidx.compose.ui.unit.LayoutDirection
-import androidx.compose.ui.unit.toSize
-
-/**
- * Punch a hole in this node with the given [size], [offset] and [shape].
- *
- * Punching a hole in an element will "remove" any pixel drawn by that element in the hole area.
- * This can be used to make content drawn below an opaque element visible. For example, if we have
- * [this lockscreen scene](http://shortn/_VYySFnJDhN) drawn below
- * [this shade scene](http://shortn/_fpxGUk0Rg7) and punch a hole in the latter using the big clock
- * time bounds and a RoundedCornerShape(10dp), [this](http://shortn/_qt80IvORFj) would be the
- * result.
- */
-@Stable
-fun Modifier.punchHole(
-    size: () -> Size,
-    offset: () -> Offset,
-    shape: Shape = RectangleShape,
-): Modifier = this.then(PunchHoleElement(size, offset, shape))
-
-/**
- * Punch a hole in this node using the bounds of [coords] and the given [shape].
- *
- * You can use [androidx.compose.ui.layout.onGloballyPositioned] to get the last coordinates of a
- * node.
- */
-@Stable
-fun Modifier.punchHole(
-    coords: () -> LayoutCoordinates?,
-    shape: Shape = RectangleShape,
-): Modifier = this.then(PunchHoleWithBoundsElement(coords, shape))
-
-private data class PunchHoleElement(
-    private val size: () -> Size,
-    private val offset: () -> Offset,
-    private val shape: Shape,
-) : ModifierNodeElement<PunchHoleNode>() {
-    override fun create(): PunchHoleNode = PunchHoleNode(size, offset, { shape })
-
-    override fun update(node: PunchHoleNode) {
-        node.size = size
-        node.offset = offset
-        node.shape = { shape }
-    }
-}
-
-private class PunchHoleNode(
-    var size: () -> Size,
-    var offset: () -> Offset,
-    var shape: () -> Shape,
-) : Modifier.Node(), DrawModifierNode, LayoutModifierNode {
-    private var lastSize: Size = Size.Unspecified
-    private var lastLayoutDirection: LayoutDirection = LayoutDirection.Ltr
-    private var lastOutline: Outline? = null
-
-    override fun MeasureScope.measure(
-        measurable: Measurable,
-        constraints: Constraints
-    ): MeasureResult {
-        return measurable.measure(constraints).run {
-            layout(width, height) {
-                placeWithLayer(0, 0) { compositingStrategy = CompositingStrategy.Offscreen }
-            }
-        }
-    }
-
-    override fun ContentDrawScope.draw() {
-        drawContent()
-
-        val holeSize = size()
-        if (holeSize != Size.Zero) {
-            val offset = offset()
-            translate(offset.x, offset.y) { drawHole(holeSize) }
-        }
-    }
-
-    private fun DrawScope.drawHole(size: Size) {
-        if (shape == RectangleShape) {
-            drawRect(Color.Black, size = size, blendMode = BlendMode.DstOut)
-            return
-        }
-
-        val outline =
-            if (size == lastSize && layoutDirection == lastLayoutDirection) {
-                lastOutline!!
-            } else {
-                val newOutline = shape().createOutline(size, layoutDirection, this)
-                lastSize = size
-                lastLayoutDirection = layoutDirection
-                lastOutline = newOutline
-                newOutline
-            }
-
-        drawOutline(
-            outline,
-            Color.Black,
-            blendMode = BlendMode.DstOut,
-        )
-    }
-}
-
-private data class PunchHoleWithBoundsElement(
-    private val coords: () -> LayoutCoordinates?,
-    private val shape: Shape,
-) : ModifierNodeElement<PunchHoleWithBoundsNode>() {
-    override fun create(): PunchHoleWithBoundsNode = PunchHoleWithBoundsNode(coords, shape)
-
-    override fun update(node: PunchHoleWithBoundsNode) {
-        node.holeCoords = coords
-        node.shape = shape
-    }
-}
-
-private class PunchHoleWithBoundsNode(
-    var holeCoords: () -> LayoutCoordinates?,
-    var shape: Shape,
-) : DelegatingNode(), DrawModifierNode, GlobalPositionAwareModifierNode {
-    private val delegate = delegate(PunchHoleNode(::holeSize, ::holeOffset, ::shape))
-    private var lastCoords: LayoutCoordinates? = null
-
-    override fun onGloballyPositioned(coordinates: LayoutCoordinates) {
-        this.lastCoords = coordinates
-    }
-
-    override fun ContentDrawScope.draw() = with(delegate) { draw() }
-
-    private fun holeSize(): Size {
-        return holeCoords()?.size?.toSize() ?: Size.Zero
-    }
-
-    private fun holeOffset(): Offset {
-        val holeCoords = holeCoords() ?: return Offset.Zero
-        val lastCoords = lastCoords ?: error("draw() was called before onGloballyPositioned()")
-        return lastCoords.localPositionOf(holeCoords, relativeToSource = Offset.Zero)
-    }
-}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.kt
index 5a39de8..444f63a 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.kt
@@ -34,6 +34,9 @@
 import androidx.lifecycle.LifecycleRegistry
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.app.viewcapture.ViewCapture
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager
+import com.android.app.viewcapture.ViewCaptureFactory
 import com.android.compose.animation.scene.ObservableTransitionState
 import com.android.internal.logging.UiEventLogger
 import com.android.keyguard.KeyguardUpdateMonitor
@@ -79,6 +82,7 @@
 import org.mockito.Mockito
 import org.mockito.Mockito.clearInvocations
 import org.mockito.Mockito.isNull
+import org.mockito.Mockito.spy
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
 
@@ -114,11 +118,11 @@
 
     @Mock
     lateinit var mDreamComplicationComponentFactory:
-        com.android.systemui.dreams.complication.dagger.ComplicationComponent.Factory
+            com.android.systemui.dreams.complication.dagger.ComplicationComponent.Factory
 
     @Mock
     lateinit var mDreamComplicationComponent:
-        com.android.systemui.dreams.complication.dagger.ComplicationComponent
+            com.android.systemui.dreams.complication.dagger.ComplicationComponent
 
     @Mock lateinit var mHideComplicationTouchHandler: HideComplicationTouchHandler
 
@@ -154,8 +158,12 @@
 
     @Mock lateinit var mDreamOverlayCallbackController: DreamOverlayCallbackController
 
+    @Mock lateinit var mLazyViewCapture: Lazy<ViewCapture>
+
+    private lateinit var mViewCaptureAwareWindowManager: ViewCaptureAwareWindowManager
     private lateinit var bouncerRepository: FakeKeyguardBouncerRepository
     private lateinit var communalRepository: FakeCommunalSceneRepository
+    private var viewCaptureSpy = spy(ViewCaptureFactory.getInstance(context))
 
     @Captor var mViewCaptor: ArgumentCaptor<View>? = null
     private lateinit var mService: DreamOverlayService
@@ -192,13 +200,16 @@
         whenever(mDreamOverlayContainerViewController.containerView)
             .thenReturn(mDreamOverlayContainerView)
         whenever(mScrimManager.getCurrentController()).thenReturn(mScrimController)
+        whenever(mLazyViewCapture.value).thenReturn(viewCaptureSpy)
         mWindowParams = WindowManager.LayoutParams()
+        mViewCaptureAwareWindowManager = ViewCaptureAwareWindowManager(mWindowManager,
+                mLazyViewCapture, isViewCaptureEnabled = false)
         mService =
             DreamOverlayService(
                 mContext,
                 mLifecycleOwner,
                 mMainExecutor,
-                mWindowManager,
+                mViewCaptureAwareWindowManager,
                 mComplicationComponentFactory,
                 mDreamComplicationComponentFactory,
                 mDreamOverlayComponentFactory,
@@ -246,7 +257,7 @@
         mMainExecutor.runAllReady()
         verify(mUiEventLogger).log(DreamOverlayService.DreamOverlayEvent.DREAM_OVERLAY_ENTER_START)
         verify(mUiEventLogger)
-            .log(DreamOverlayService.DreamOverlayEvent.DREAM_OVERLAY_COMPLETE_START)
+                .log(DreamOverlayService.DreamOverlayEvent.DREAM_OVERLAY_COMPLETE_START)
     }
 
     @Test
diff --git a/packages/SystemUI/res/xml/large_screen_shade_header.xml b/packages/SystemUI/res/xml/large_screen_shade_header.xml
index fe61c46..eb0aae9 100644
--- a/packages/SystemUI/res/xml/large_screen_shade_header.xml
+++ b/packages/SystemUI/res/xml/large_screen_shade_header.xml
@@ -33,6 +33,7 @@
             android:layout_width="wrap_content"
             android:layout_height="0dp"
             android:layout_marginStart="8dp"
+            app:layout_constraintBaseline_toBaselineOf="@id/clock"
             app:layout_constraintBottom_toBottomOf="parent"
             app:layout_constraintStart_toEndOf="@id/clock"
             app:layout_constraintTop_toTopOf="parent" />
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/FrameworkServicesModule.java b/packages/SystemUI/src/com/android/systemui/dagger/FrameworkServicesModule.java
index b0fc60e..a6deca7 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/FrameworkServicesModule.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/FrameworkServicesModule.java
@@ -16,6 +16,9 @@
 
 package com.android.systemui.dagger;
 
+import static com.android.systemui.Flags.enableViewCaptureTracing;
+import static com.android.systemui.util.ConvenienceExtensionsKt.toKotlinLazy;
+
 import android.annotation.Nullable;
 import android.annotation.SuppressLint;
 import android.app.ActivityManager;
@@ -111,6 +114,9 @@
 import androidx.asynclayoutinflater.view.AsyncLayoutInflater;
 import androidx.core.app.NotificationManagerCompat;
 
+import com.android.app.viewcapture.ViewCapture;
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager;
+import com.android.app.viewcapture.ViewCaptureFactory;
 import com.android.internal.app.IBatteryStats;
 import com.android.internal.appwidget.IAppWidgetService;
 import com.android.internal.jank.InteractionJankMonitor;
@@ -125,6 +131,7 @@
 import com.android.systemui.user.utils.UserScopedService;
 import com.android.systemui.user.utils.UserScopedServiceImpl;
 
+import dagger.Lazy;
 import dagger.Module;
 import dagger.Provides;
 
@@ -680,6 +687,15 @@
 
     @Provides
     @Singleton
+    static ViewCaptureAwareWindowManager provideViewCaptureAwareWindowManager(
+            WindowManager windowManager, Lazy<ViewCapture> daggerLazyViewCapture) {
+        return new ViewCaptureAwareWindowManager(windowManager,
+                /* lazyViewCapture= */ toKotlinLazy(daggerLazyViewCapture),
+                /* isViewCaptureEnabled= */ enableViewCaptureTracing());
+    }
+
+    @Provides
+    @Singleton
     static PermissionManager providePermissionManager(Context context) {
         PermissionManager pm = context.getSystemService(PermissionManager.class);
         if (pm != null) {
@@ -764,4 +780,10 @@
         return IDeviceIdleController.Stub.asInterface(
                 ServiceManager.getService(Context.DEVICE_IDLE_CONTROLLER));
     }
+
+    @Provides
+    @Singleton
+    static ViewCapture provideViewCapture(Context context) {
+        return ViewCaptureFactory.getInstance(context);
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayService.java b/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayService.java
index c6c57479..83fa001 100644
--- a/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayService.java
+++ b/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayService.java
@@ -45,6 +45,7 @@
 import androidx.lifecycle.ServiceLifecycleDispatcher;
 import androidx.lifecycle.ViewModelStore;
 
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager;
 import com.android.dream.lowlight.dagger.LowLightDreamModule;
 import com.android.internal.logging.UiEvent;
 import com.android.internal.logging.UiEventLogger;
@@ -97,7 +98,7 @@
     @Nullable
     private final ComponentName mHomeControlPanelDreamComponent;
     private final UiEventLogger mUiEventLogger;
-    private final WindowManager mWindowManager;
+    private final ViewCaptureAwareWindowManager mWindowManager;
     private final String mWindowTitle;
 
     // A reference to the {@link Window} used to hold the dream overlay.
@@ -244,7 +245,7 @@
             Context context,
             DreamOverlayLifecycleOwner lifecycleOwner,
             @Main DelayableExecutor executor,
-            WindowManager windowManager,
+            ViewCaptureAwareWindowManager viewCaptureAwareWindowManager,
             ComplicationComponent.Factory complicationComponentFactory,
             com.android.systemui.dreams.complication.dagger.ComplicationComponent.Factory
                     dreamComplicationComponentFactory,
@@ -267,7 +268,7 @@
         super(executor);
         mContext = context;
         mExecutor = executor;
-        mWindowManager = windowManager;
+        mWindowManager = viewCaptureAwareWindowManager;
         mKeyguardUpdateMonitor = keyguardUpdateMonitor;
         mScrimManager = scrimManager;
         mLowLightDreamComponent = lowLightDreamComponent;
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
index c4b70d8..9f33113 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
@@ -81,6 +81,7 @@
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.keyguard.domain.interactor.KeyguardEnabledInteractor;
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor;
+import com.android.systemui.keyguard.domain.interactor.KeyguardWakeDirectlyToGoneInteractor;
 import com.android.systemui.keyguard.ui.binder.KeyguardSurfaceBehindParamsApplier;
 import com.android.systemui.keyguard.ui.binder.KeyguardSurfaceBehindViewBinder;
 import com.android.systemui.keyguard.ui.binder.WindowManagerLockscreenVisibilityViewBinder;
@@ -317,7 +318,7 @@
 
     private final WindowManagerOcclusionManager mWmOcclusionManager;
     private final KeyguardEnabledInteractor mKeyguardEnabledInteractor;
-
+    private final KeyguardWakeDirectlyToGoneInteractor mKeyguardWakeDirectlyToGoneInteractor;
     private final Lazy<FoldGracePeriodProvider> mFoldGracePeriodProvider = new Lazy<>() {
         @Override
         public FoldGracePeriodProvider get() {
@@ -344,7 +345,8 @@
             @Main Executor mainExecutor,
             KeyguardInteractor keyguardInteractor,
             KeyguardEnabledInteractor keyguardEnabledInteractor,
-            Lazy<KeyguardStateCallbackStartable> keyguardStateCallbackStartableLazy) {
+            Lazy<KeyguardStateCallbackStartable> keyguardStateCallbackStartableLazy,
+            KeyguardWakeDirectlyToGoneInteractor keyguardWakeDirectlyToGoneInteractor) {
         super();
         mKeyguardViewMediator = keyguardViewMediator;
         mKeyguardLifecyclesDispatcher = keyguardLifecyclesDispatcher;
@@ -372,6 +374,7 @@
 
         mWmOcclusionManager = windowManagerOcclusionManager;
         mKeyguardEnabledInteractor = keyguardEnabledInteractor;
+        mKeyguardWakeDirectlyToGoneInteractor = keyguardWakeDirectlyToGoneInteractor;
     }
 
     @Override
@@ -486,6 +489,7 @@
         public void onDreamingStarted() {
             trace("onDreamingStarted");
             checkPermission();
+            mKeyguardWakeDirectlyToGoneInteractor.onDreamingStarted();
             mKeyguardInteractor.setDreaming(true);
             mKeyguardViewMediator.onDreamingStarted();
         }
@@ -494,6 +498,7 @@
         public void onDreamingStopped() {
             trace("onDreamingStopped");
             checkPermission();
+            mKeyguardWakeDirectlyToGoneInteractor.onDreamingStopped();
             mKeyguardInteractor.setDreaming(false);
             mKeyguardViewMediator.onDreamingStopped();
         }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt
index f837d8e..ae751db 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt
@@ -127,6 +127,30 @@
      */
     val isKeyguardEnabled: StateFlow<Boolean>
 
+    /**
+     * Whether we can transition directly back to GONE from AOD/DOZING without any authentication
+     * events (such as a fingerprint wake and unlock), even though authentication would normally be
+     * required. This means that if you tap the screen or press the power button, you'll return
+     * directly to the unlocked app content without seeing the lockscreen, even if a secure
+     * authentication method (PIN/password/biometrics) is set.
+     *
+     * This is true in these cases:
+     * - The screen timed out, but the "lock after screen timeout" duration (default 5 seconds) has
+     *   not yet elapsed.
+     * - The power button was pressed, but "power button instantly locks" is not enabled, and the
+     *   "lock after screen timeout" duration has not elapsed.
+     *
+     * Note that this value specifically tells us if we can *ignore* authentication that would
+     * otherwise be required to transition from AOD/DOZING -> GONE. AOD/DOZING -> GONE is also
+     * possible if keyguard is disabled, either from an app request or because security is set to
+     * "none", but in that case, auth is not required so this boolean is not relevant.
+     *
+     * See [KeyguardWakeToGoneInteractor].
+     */
+    val canIgnoreAuthAndReturnToGone: StateFlow<Boolean>
+
+    fun setCanIgnoreAuthAndReturnToGone(canWake: Boolean)
+
     /** Is the always-on display available to be used? */
     val isAodAvailable: StateFlow<Boolean>
 
@@ -386,6 +410,13 @@
         MutableStateFlow(!lockPatternUtils.isLockScreenDisabled(userTracker.userId))
     override val isKeyguardEnabled: StateFlow<Boolean> = _isKeyguardEnabled.asStateFlow()
 
+    private val _canIgnoreAuthAndReturnToGone = MutableStateFlow(false)
+    override val canIgnoreAuthAndReturnToGone = _canIgnoreAuthAndReturnToGone.asStateFlow()
+
+    override fun setCanIgnoreAuthAndReturnToGone(canWakeToGone: Boolean) {
+        _canIgnoreAuthAndReturnToGone.value = canWakeToGone
+    }
+
     private val _isDozing = MutableStateFlow(statusBarStateController.isDozing)
     override val isDozing: StateFlow<Boolean> = _isDozing.asStateFlow()
 
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractor.kt
index 1426cba..893835a 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractor.kt
@@ -53,6 +53,7 @@
     powerInteractor: PowerInteractor,
     keyguardOcclusionInteractor: KeyguardOcclusionInteractor,
     val deviceEntryRepository: DeviceEntryRepository,
+    private val wakeToGoneInteractor: KeyguardWakeDirectlyToGoneInteractor,
 ) :
     TransitionInteractor(
         fromState = KeyguardState.AOD,
@@ -98,6 +99,7 @@
                     keyguardInteractor.primaryBouncerShowing,
                     keyguardInteractor.isKeyguardOccluded,
                     canDismissLockscreen,
+                    wakeToGoneInteractor.canWakeDirectlyToGone,
                 )
                 .collect {
                     (
@@ -107,6 +109,7 @@
                         primaryBouncerShowing,
                         isKeyguardOccludedLegacy,
                         canDismissLockscreen,
+                        canWakeDirectlyToGone,
                     ) ->
                     if (!maybeHandleInsecurePowerGesture()) {
                         val shouldTransitionToLockscreen =
@@ -131,8 +134,7 @@
 
                         val shouldTransitionToGone =
                             (!KeyguardWmStateRefactor.isEnabled && canDismissLockscreen) ||
-                                (KeyguardWmStateRefactor.isEnabled &&
-                                    !deviceEntryRepository.isLockscreenEnabled())
+                                (KeyguardWmStateRefactor.isEnabled && canWakeDirectlyToGone)
 
                         if (shouldTransitionToGone) {
                             // TODO(b/336576536): Check if adaptation for scene framework is needed
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromDozingTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromDozingTransitionInteractor.kt
index 76e88a2..aee65a8 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromDozingTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromDozingTransitionInteractor.kt
@@ -55,6 +55,7 @@
     private val communalInteractor: CommunalInteractor,
     keyguardOcclusionInteractor: KeyguardOcclusionInteractor,
     val deviceEntryRepository: DeviceEntryRepository,
+    private val wakeToGoneInteractor: KeyguardWakeDirectlyToGoneInteractor,
 ) :
     TransitionInteractor(
         fromState = KeyguardState.DOZING,
@@ -181,7 +182,7 @@
                 .sample(
                     communalInteractor.isIdleOnCommunal,
                     keyguardInteractor.biometricUnlockState,
-                    canTransitionToGoneOnWake,
+                    wakeToGoneInteractor.canWakeDirectlyToGone,
                     keyguardInteractor.primaryBouncerShowing,
                 )
                 .collect {
@@ -189,27 +190,14 @@
                         _,
                         isIdleOnCommunal,
                         biometricUnlockState,
-                        canDismissLockscreen,
+                        canWakeDirectlyToGone,
                         primaryBouncerShowing) ->
                     if (
                         !maybeStartTransitionToOccludedOrInsecureCamera() &&
                             // Handled by dismissFromDozing().
                             !isWakeAndUnlock(biometricUnlockState.mode)
                     ) {
-                        if (!KeyguardWmStateRefactor.isEnabled && canDismissLockscreen) {
-                            if (SceneContainerFlag.isEnabled) {
-                                // TODO(b/336576536): Check if adaptation for scene framework is
-                                // needed
-                            } else {
-                                startTransitionTo(
-                                    KeyguardState.GONE,
-                                    ownerReason = "waking from dozing"
-                                )
-                            }
-                        } else if (
-                            KeyguardWmStateRefactor.isEnabled &&
-                                !deviceEntryRepository.isLockscreenEnabled()
-                        ) {
+                        if (canWakeDirectlyToGone) {
                             if (SceneContainerFlag.isEnabled) {
                                 // TODO(b/336576536): Check if adaptation for scene framework is
                                 // needed
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractor.kt
index 0e76487..cfb161c 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractor.kt
@@ -23,6 +23,7 @@
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.deviceentry.domain.interactor.DeviceEntryInteractor
 import com.android.systemui.keyguard.KeyguardWmStateRefactor
 import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
 import com.android.systemui.keyguard.shared.model.BiometricUnlockMode
@@ -37,11 +38,14 @@
 import kotlin.time.Duration.Companion.seconds
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.FlowPreview
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.debounce
+import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.launch
 
+@OptIn(ExperimentalCoroutinesApi::class)
 @SysUISingleton
 class FromDreamingTransitionInteractor
 @Inject
@@ -56,6 +60,7 @@
     private val glanceableHubTransitions: GlanceableHubTransitions,
     powerInteractor: PowerInteractor,
     keyguardOcclusionInteractor: KeyguardOcclusionInteractor,
+    private val deviceEntryInteractor: DeviceEntryInteractor,
 ) :
     TransitionInteractor(
         fromState = KeyguardState.DREAMING,
@@ -72,7 +77,7 @@
         listenForDreamingToOccluded()
         listenForDreamingToGoneWhenDismissable()
         listenForDreamingToGoneFromBiometricUnlock()
-        listenForDreamingToLockscreen()
+        listenForDreamingToLockscreenOrGone()
         listenForDreamingToAodOrDozing()
         listenForTransitionToCamera(scope, keyguardInteractor)
         listenForDreamingToGlanceableHub()
@@ -132,17 +137,7 @@
 
     @OptIn(FlowPreview::class)
     private fun listenForDreamingToOccluded() {
-        if (KeyguardWmStateRefactor.isEnabled) {
-            scope.launch {
-                combine(
-                        keyguardInteractor.isDreaming,
-                        keyguardOcclusionInteractor.isShowWhenLockedActivityOnTop,
-                        ::Pair
-                    )
-                    .filterRelevantKeyguardStateAnd { (isDreaming, _) -> !isDreaming }
-                    .collect { maybeStartTransitionToOccludedOrInsecureCamera() }
-            }
-        } else {
+        if (!KeyguardWmStateRefactor.isEnabled) {
             scope.launch {
                 combine(
                         keyguardInteractor.isKeyguardOccluded,
@@ -168,21 +163,41 @@
         }
     }
 
-    private fun listenForDreamingToLockscreen() {
+    private fun listenForDreamingToLockscreenOrGone() {
         if (!KeyguardWmStateRefactor.isEnabled) {
             return
         }
 
         scope.launch {
-            keyguardOcclusionInteractor.isShowWhenLockedActivityOnTop
-                .filterRelevantKeyguardStateAnd { onTop -> !onTop }
-                .collect { startTransitionTo(KeyguardState.LOCKSCREEN) }
+            keyguardInteractor.isDreaming
+                .filter { !it }
+                .sample(deviceEntryInteractor.isUnlocked, ::Pair)
+                .collect { (_, dismissable) ->
+                    // TODO(b/349837588): Add check for -> OCCLUDED.
+                    if (dismissable) {
+                        startTransitionTo(
+                            KeyguardState.GONE,
+                            ownerReason = "No longer dreaming; dismissable"
+                        )
+                    } else {
+                        startTransitionTo(
+                            KeyguardState.LOCKSCREEN,
+                            ownerReason = "No longer dreaming"
+                        )
+                    }
+                }
         }
     }
 
     private fun listenForDreamingToGoneWhenDismissable() {
-        // TODO(b/336576536): Check if adaptation for scene framework is needed
-        if (SceneContainerFlag.isEnabled) return
+        if (SceneContainerFlag.isEnabled) {
+            return // TODO(b/336576536): Check if adaptation for scene framework is needed
+        }
+
+        if (KeyguardWmStateRefactor.isEnabled) {
+            return
+        }
+
         scope.launch {
             keyguardInteractor.isAbleToDream
                 .sampleCombine(
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardWakeDirectlyToGoneInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardWakeDirectlyToGoneInteractor.kt
new file mode 100644
index 0000000..f0bf402
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardWakeDirectlyToGoneInteractor.kt
@@ -0,0 +1,372 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.keyguard.domain.interactor
+
+import android.annotation.SuppressLint
+import android.app.AlarmManager
+import android.app.PendingIntent
+import android.content.BroadcastReceiver
+import android.content.Context
+import android.content.Intent
+import android.content.IntentFilter
+import android.provider.Settings
+import android.provider.Settings.Secure
+import com.android.internal.widget.LockPatternUtils
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.keyguard.KeyguardViewMediator
+import com.android.systemui.keyguard.KeyguardWmStateRefactor
+import com.android.systemui.keyguard.data.repository.KeyguardRepository
+import com.android.systemui.keyguard.shared.model.BiometricUnlockMode
+import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.keyguard.shared.model.KeyguardState.Companion.deviceIsAsleepInState
+import com.android.systemui.keyguard.shared.model.KeyguardState.Companion.deviceIsAwakeInState
+import com.android.systemui.power.domain.interactor.PowerInteractor
+import com.android.systemui.power.shared.model.WakeSleepReason
+import com.android.systemui.user.domain.interactor.SelectedUserInteractor
+import com.android.systemui.util.kotlin.sample
+import com.android.systemui.util.settings.SecureSettings
+import com.android.systemui.util.settings.SystemSettings
+import com.android.systemui.util.time.SystemClock
+import javax.inject.Inject
+import kotlin.math.max
+import kotlin.math.min
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.distinctUntilChangedBy
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.launch
+
+/**
+ * Logic related to the ability to wake directly to GONE from asleep (AOD/DOZING), without going
+ * through LOCKSCREEN or a BOUNCER state.
+ *
+ * This is possible in the following scenarios:
+ * - The lockscreen is disabled, either from an app request (SUW does this), or by the security
+ *   "None" setting.
+ * - A biometric authentication event occurred while we were asleep (fingerprint auth, etc). This
+ *   specifically is referred to throughout the codebase as "wake and unlock".
+ * - The screen timed out, but the "lock after screen timeout" duration has not elapsed.
+ * - The power button was pressed, but "power button instantly locks" is disabled and the "lock
+ *   after screen timeout" duration has not elapsed.
+ *
+ * In these cases, no (further) authentication is required, and we can transition directly from
+ * AOD/DOZING -> GONE.
+ */
+@SysUISingleton
+class KeyguardWakeDirectlyToGoneInteractor
+@Inject
+constructor(
+    @Application private val scope: CoroutineScope,
+    private val context: Context,
+    private val repository: KeyguardRepository,
+    private val systemClock: SystemClock,
+    private val alarmManager: AlarmManager,
+    private val transitionInteractor: KeyguardTransitionInteractor,
+    private val powerInteractor: PowerInteractor,
+    private val secureSettings: SecureSettings,
+    private val lockPatternUtils: LockPatternUtils,
+    private val systemSettings: SystemSettings,
+    private val selectedUserInteractor: SelectedUserInteractor,
+) {
+
+    /**
+     * Whether the lockscreen was disabled as of the last wake/sleep event, according to
+     * LockPatternUtils.
+     *
+     * This will always be true if [repository.isKeyguardServiceEnabled]=false, but it can also be
+     * true when the keyguard service is enabled if the lockscreen has been disabled via adb using
+     * the `adb shell locksettings set-disabled true` command, which is often done in tests.
+     *
+     * Unlike keyguardServiceEnabled, changes to this value should *not* immediately show or hide
+     * the keyguard. If the lockscreen is disabled in this way, it will just not show on the next
+     * sleep/wake.
+     */
+    private val isLockscreenDisabled: Flow<Boolean> =
+        powerInteractor.isAwake.map { isLockscreenDisabled() }
+
+    /**
+     * Whether we can wake from AOD/DOZING directly to GONE, bypassing LOCKSCREEN/BOUNCER states.
+     *
+     * This is possible in the following cases:
+     * - Keyguard is disabled, either from an app request or from security being set to "None".
+     * - We're wake and unlocking (fingerprint auth occurred while asleep).
+     * - We're allowed to ignore auth and return to GONE, due to timeouts not elapsing.
+     */
+    val canWakeDirectlyToGone =
+        combine(
+                repository.isKeyguardEnabled,
+                isLockscreenDisabled,
+                repository.biometricUnlockState,
+                repository.canIgnoreAuthAndReturnToGone,
+            ) {
+                keyguardEnabled,
+                isLockscreenDisabled,
+                biometricUnlockState,
+                canIgnoreAuthAndReturnToGone ->
+                (!keyguardEnabled || isLockscreenDisabled) ||
+                    BiometricUnlockMode.isWakeAndUnlock(biometricUnlockState.mode) ||
+                    canIgnoreAuthAndReturnToGone
+            }
+            .distinctUntilChanged()
+
+    /**
+     * Counter that is incremented every time we wake up or stop dreaming. Upon sleeping/dreaming,
+     * we put the current value of this counter into the intent extras of the timeout alarm intent.
+     * If this value has changed by the time we receive the intent, it is discarded since it's out
+     * of date.
+     */
+    var timeoutCounter = 0
+
+    var isAwake = false
+
+    private val broadcastReceiver: BroadcastReceiver =
+        object : BroadcastReceiver() {
+            override fun onReceive(context: Context, intent: Intent) {
+                if (DELAYED_KEYGUARD_ACTION == intent.action) {
+                    val sequence = intent.getIntExtra(SEQ_EXTRA_KEY, 0)
+                    synchronized(this) {
+                        if (timeoutCounter == sequence) {
+                            // If the sequence # matches, we have not woken up or stopped dreaming
+                            // since
+                            // the alarm was set. That means this is still relevant - the lock
+                            // timeout
+                            // has elapsed, so let the repository know that we can no longer return
+                            // to
+                            // GONE without authenticating.
+                            repository.setCanIgnoreAuthAndReturnToGone(false)
+                        }
+                    }
+                }
+            }
+        }
+
+    init {
+        setOrCancelAlarmFromWakefulness()
+        listenForWakeToClearCanIgnoreAuth()
+        registerBroadcastReceiver()
+    }
+
+    fun onDreamingStarted() {
+        // If we start dreaming while awake, lock after the normal timeout.
+        if (isAwake) {
+            setResetCanIgnoreAuthAlarm()
+        }
+    }
+
+    fun onDreamingStopped() {
+        // Cancel the timeout if we stop dreaming while awake.
+        if (isAwake) {
+            cancelCanIgnoreAuthAlarm()
+        }
+    }
+
+    private fun setOrCancelAlarmFromWakefulness() {
+        scope.launch {
+            powerInteractor.detailedWakefulness
+                .distinctUntilChangedBy { it.isAwake() }
+                .sample(transitionInteractor.currentKeyguardState, ::Pair)
+                .collect { (wakefulness, currentState) ->
+                    // Save isAwake for use in onDreamingStarted/onDreamingStopped.
+                    this@KeyguardWakeDirectlyToGoneInteractor.isAwake = wakefulness.isAwake()
+
+                    // If we're sleeping from GONE, check the timeout and lock instantly settings.
+                    // These are not relevant if we're coming from non-GONE states.
+                    if (!isAwake && currentState == KeyguardState.GONE) {
+                        val lockTimeoutDuration = getCanIgnoreAuthAndReturnToGoneDuration()
+
+                        // If the screen timed out and went to sleep, and the lock timeout is > 0ms,
+                        // then we can return to GONE until that duration elapses. If the power
+                        // button was pressed but "instantly locks" is disabled, then we can also
+                        // return to GONE until the timeout duration elapses.
+                        if (
+                            (wakefulness.lastSleepReason == WakeSleepReason.TIMEOUT &&
+                                lockTimeoutDuration > 0) ||
+                                (wakefulness.lastSleepReason == WakeSleepReason.POWER_BUTTON &&
+                                    !willLockImmediately())
+                        ) {
+
+                            // Let the repository know that we can return to GONE until we notify
+                            // it otherwise.
+                            repository.setCanIgnoreAuthAndReturnToGone(true)
+                            setResetCanIgnoreAuthAlarm()
+                        }
+                    } else if (isAwake) {
+                        // If we're waking up, ignore the alarm if it goes off since it's no longer
+                        // relevant. Once a wake KeyguardTransition is started, we'll also clear the
+                        // canIgnoreAuthAndReturnToGone value in listenForWakeToClearCanIgnoreAuth.
+                        cancelCanIgnoreAuthAlarm()
+                    }
+                }
+        }
+    }
+
+    /** Clears the canIgnoreAuthAndReturnToGone value upon waking. */
+    private fun listenForWakeToClearCanIgnoreAuth() {
+        scope.launch {
+            transitionInteractor
+                .isInTransitionWhere(
+                    fromStatePredicate = { deviceIsAsleepInState(it) },
+                    toStatePredicate = { deviceIsAwakeInState(it) },
+                )
+                .collect {
+                    // This value is reset when the timeout alarm fires, but if the device is woken
+                    // back up before then, it needs to be reset here. The alarm is cancelled
+                    // immediately upon waking up, but since this value is used by keyguard
+                    // transition internals to decide whether we can transition to GONE, wait until
+                    // that decision is made before resetting it.
+                    repository.setCanIgnoreAuthAndReturnToGone(false)
+                }
+        }
+    }
+
+    /**
+     * Registers the broadcast receiver to receive the alarm intent.
+     *
+     * TODO(b/351817381): Investigate using BroadcastDispatcher vs. ignoring this lint warning.
+     */
+    @SuppressLint("WrongConstant", "RegisterReceiverViaContext")
+    private fun registerBroadcastReceiver() {
+        val delayedActionFilter = IntentFilter()
+        delayedActionFilter.addAction(KeyguardViewMediator.DELAYED_KEYGUARD_ACTION)
+        // TODO(b/346803756): Listen for DELAYED_LOCK_PROFILE_ACTION.
+        delayedActionFilter.priority = IntentFilter.SYSTEM_HIGH_PRIORITY
+        context.registerReceiver(
+            broadcastReceiver,
+            delayedActionFilter,
+            SYSTEMUI_PERMISSION,
+            null /* scheduler */,
+            Context.RECEIVER_EXPORTED_UNAUDITED
+        )
+    }
+
+    /** Set an alarm for */
+    private fun setResetCanIgnoreAuthAlarm() {
+        if (!KeyguardWmStateRefactor.isEnabled) {
+            return
+        }
+
+        val intent =
+            Intent(DELAYED_KEYGUARD_ACTION).apply {
+                setPackage(context.packageName)
+                putExtra(SEQ_EXTRA_KEY, timeoutCounter)
+                addFlags(Intent.FLAG_RECEIVER_FOREGROUND)
+            }
+
+        val sender =
+            PendingIntent.getBroadcast(
+                context,
+                0,
+                intent,
+                PendingIntent.FLAG_CANCEL_CURRENT or PendingIntent.FLAG_IMMUTABLE
+            )
+
+        val time = systemClock.elapsedRealtime() + getCanIgnoreAuthAndReturnToGoneDuration()
+        alarmManager.setExactAndAllowWhileIdle(AlarmManager.ELAPSED_REALTIME_WAKEUP, time, sender)
+
+        // TODO(b/346803756): Migrate support for child profiles.
+    }
+
+    /**
+     * Cancel the timeout by incrementing the counter so that we ignore the intent when it's
+     * received.
+     */
+    private fun cancelCanIgnoreAuthAlarm() {
+        timeoutCounter++
+    }
+
+    /**
+     * Whether pressing the power button locks the device immediately; vs. waiting for a specified
+     * timeout first.
+     */
+    private fun willLockImmediately(
+        userId: Int = selectedUserInteractor.getSelectedUserId()
+    ): Boolean {
+        return lockPatternUtils.getPowerButtonInstantlyLocks(userId) ||
+            !lockPatternUtils.isSecure(userId)
+    }
+
+    /**
+     * Returns whether the lockscreen is disabled, either because the keyguard service is disabled
+     * or because an adb command has disabled the lockscreen.
+     */
+    private fun isLockscreenDisabled(
+        userId: Int = selectedUserInteractor.getSelectedUserId()
+    ): Boolean {
+        return lockPatternUtils.isLockScreenDisabled(userId)
+    }
+
+    /**
+     * Returns the duration within which we can return to GONE without auth after a screen timeout
+     * (or power button press, if lock instantly is disabled).
+     *
+     * This takes into account the user's settings as well as device policy maximums.
+     */
+    private fun getCanIgnoreAuthAndReturnToGoneDuration(
+        userId: Int = selectedUserInteractor.getSelectedUserId()
+    ): Long {
+        // The timeout duration from settings (Security > Device Unlock > Gear icon > "Lock after
+        // screen timeout".
+        val durationSetting: Long =
+            secureSettings
+                .getIntForUser(
+                    Secure.LOCK_SCREEN_LOCK_AFTER_TIMEOUT,
+                    KEYGUARD_CAN_IGNORE_AUTH_DURATION,
+                    userId
+                )
+                .toLong()
+
+        // Device policy maximum timeout.
+        val durationDevicePolicyMax =
+            lockPatternUtils.devicePolicyManager.getMaximumTimeToLock(null, userId)
+
+        return if (durationDevicePolicyMax <= 0) {
+            durationSetting
+        } else {
+            var displayTimeout =
+                systemSettings
+                    .getIntForUser(
+                        Settings.System.SCREEN_OFF_TIMEOUT,
+                        KEYGUARD_DISPLAY_TIMEOUT_DELAY_DEFAULT,
+                        userId
+                    )
+                    .toLong()
+
+            // Ignore negative values. I don't know why this would be negative, but this check has
+            // been around since 2016 and I see no upside to removing it.
+            displayTimeout = max(displayTimeout, 0)
+
+            // Respect the shorter of: the device policy (maximum duration between last user action
+            // and fully locking) or the "Lock after screen timeout" setting.
+            max(min(durationDevicePolicyMax - displayTimeout, durationSetting), 0)
+        }
+    }
+
+    companion object {
+        private const val DELAYED_KEYGUARD_ACTION =
+            "com.android.internal.policy.impl.PhoneWindowManager.DELAYED_KEYGUARD"
+        private const val DELAYED_LOCK_PROFILE_ACTION =
+            "com.android.internal.policy.impl.PhoneWindowManager.DELAYED_LOCK"
+        private const val SYSTEMUI_PERMISSION = "com.android.systemui.permission.SELF"
+        private const val SEQ_EXTRA_KEY = "count"
+
+        private const val KEYGUARD_CAN_IGNORE_AUTH_DURATION = 5000
+        private const val KEYGUARD_DISPLAY_TIMEOUT_DELAY_DEFAULT = 30000
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractor.kt
index 523370c..e1b333d 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractor.kt
@@ -21,14 +21,17 @@
 import com.android.compose.animation.scene.ObservableTransitionState
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.deviceentry.domain.interactor.DeviceEntryInteractor
+import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
 import com.android.systemui.keyguard.shared.model.BiometricUnlockMode
 import com.android.systemui.keyguard.shared.model.Edge
 import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.keyguard.shared.model.KeyguardState.Companion.deviceIsAsleepInState
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.scene.domain.interactor.SceneInteractor
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.statusbar.notification.domain.interactor.NotificationLaunchAnimationInteractor
+import com.android.systemui.util.kotlin.Utils.Companion.toTriple
 import com.android.systemui.util.kotlin.sample
 import com.android.systemui.utils.coroutines.flow.flatMapLatestConflated
 import dagger.Lazy
@@ -41,11 +44,13 @@
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
 
+@OptIn(ExperimentalCoroutinesApi::class)
 @SysUISingleton
 class WindowManagerLockscreenVisibilityInteractor
 @Inject
 constructor(
     keyguardInteractor: KeyguardInteractor,
+    transitionRepository: KeyguardTransitionRepository,
     transitionInteractor: KeyguardTransitionInteractor,
     surfaceBehindInteractor: KeyguardSurfaceBehindInteractor,
     fromLockscreenInteractor: FromLockscreenTransitionInteractor,
@@ -54,9 +59,15 @@
     notificationLaunchAnimationInteractor: NotificationLaunchAnimationInteractor,
     sceneInteractor: Lazy<SceneInteractor>,
     deviceEntryInteractor: Lazy<DeviceEntryInteractor>,
+    wakeToGoneInteractor: KeyguardWakeDirectlyToGoneInteractor,
 ) {
     private val defaultSurfaceBehindVisibility =
-        transitionInteractor.finishedKeyguardState.map(::isSurfaceVisible)
+        combine(
+            transitionInteractor.finishedKeyguardState,
+            wakeToGoneInteractor.canWakeDirectlyToGone,
+        ) { finishedState, canWakeDirectlyToGone ->
+            isSurfaceVisible(finishedState) || canWakeDirectlyToGone
+        }
 
     /**
      * Surface visibility provided by the From*TransitionInteractor responsible for the currently
@@ -204,9 +215,13 @@
         if (SceneContainerFlag.isEnabled) {
             isDeviceNotEntered
         } else {
-            transitionInteractor.currentKeyguardState
-                .sample(transitionInteractor.startedStepWithPrecedingStep, ::Pair)
-                .map { (currentState, startedWithPrev) ->
+            combine(
+                    transitionInteractor.currentKeyguardState,
+                    wakeToGoneInteractor.canWakeDirectlyToGone,
+                    ::Pair
+                )
+                .sample(transitionInteractor.startedStepWithPrecedingStep, ::toTriple)
+                .map { (currentState, canWakeDirectlyToGone, startedWithPrev) ->
                     val startedFromStep = startedWithPrev.previousValue
                     val startedStep = startedWithPrev.newValue
                     val returningToGoneAfterCancellation =
@@ -214,16 +229,33 @@
                             startedFromStep.transitionState == TransitionState.CANCELED &&
                             startedFromStep.from == KeyguardState.GONE
 
-                    if (!returningToGoneAfterCancellation) {
-                        // By default, apply the lockscreen visibility of the current state.
-                        deviceEntryInteractor.get().isLockscreenEnabled() &&
-                            KeyguardState.lockscreenVisibleInState(currentState)
-                    } else {
-                        // If we're transitioning to GONE after a prior canceled transition from
-                        // GONE, then this is the camera launch transition from an asleep state back
-                        // to GONE. We don't want to show the lockscreen since we're aborting the
-                        // lock and going back to GONE.
+                    val transitionInfo = transitionRepository.currentTransitionInfoInternal.value
+                    val wakingDirectlyToGone =
+                        deviceIsAsleepInState(transitionInfo.from) &&
+                            transitionInfo.to == KeyguardState.GONE
+
+                    if (returningToGoneAfterCancellation || wakingDirectlyToGone) {
+                        // GONE -> AOD/DOZING (cancel) -> GONE is the camera launch transition,
+                        // which means we never want to show the lockscreen throughout the
+                        // transition. Same for waking directly to gone, due to the lockscreen being
+                        // disabled or because the device was woken back up before the lock timeout
+                        // duration elapsed.
                         KeyguardState.lockscreenVisibleInState(KeyguardState.GONE)
+                    } else if (canWakeDirectlyToGone) {
+                        // Never show the lockscreen if we can wake directly to GONE. This means
+                        // that the lock timeout has not yet elapsed, or the keyguard is disabled.
+                        // In either case, we don't show the activity lock screen until one of those
+                        // conditions changes.
+                        false
+                    } else if (
+                        currentState == KeyguardState.DREAMING &&
+                            deviceEntryInteractor.get().isUnlocked.value
+                    ) {
+                        // Dreams dismiss keyguard and return to GONE if they can.
+                        false
+                    } else {
+                        // Otherwise, use the visibility of the current state.
+                        KeyguardState.lockscreenVisibleInState(currentState)
                     }
                 }
                 .distinctUntilChanged()
diff --git a/packages/SystemUI/src/com/android/systemui/power/shared/model/WakeSleepReason.kt b/packages/SystemUI/src/com/android/systemui/power/shared/model/WakeSleepReason.kt
index 75055668..776a8f4 100644
--- a/packages/SystemUI/src/com/android/systemui/power/shared/model/WakeSleepReason.kt
+++ b/packages/SystemUI/src/com/android/systemui/power/shared/model/WakeSleepReason.kt
@@ -54,7 +54,10 @@
     OTHER(isTouch = false, PowerManager.WAKE_REASON_UNKNOWN),
 
     /** Device goes to sleep due to folding of a foldable device. */
-    FOLD(isTouch = false, PowerManager.GO_TO_SLEEP_REASON_DEVICE_FOLD);
+    FOLD(isTouch = false, PowerManager.GO_TO_SLEEP_REASON_DEVICE_FOLD),
+
+    /** Device goes to sleep because it timed out. */
+    TIMEOUT(isTouch = false, PowerManager.GO_TO_SLEEP_REASON_TIMEOUT);
 
     companion object {
         fun fromPowerManagerWakeReason(reason: Int): WakeSleepReason {
@@ -75,6 +78,7 @@
         fun fromPowerManagerSleepReason(reason: Int): WakeSleepReason {
             return when (reason) {
                 PowerManager.GO_TO_SLEEP_REASON_POWER_BUTTON -> POWER_BUTTON
+                PowerManager.GO_TO_SLEEP_REASON_TIMEOUT -> TIMEOUT
                 PowerManager.GO_TO_SLEEP_REASON_DEVICE_FOLD -> FOLD
                 else -> OTHER
             }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/Tile.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/Tile.kt
index 2f87b01..3fdd7f7 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/Tile.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/Tile.kt
@@ -22,7 +22,6 @@
 import android.service.quicksettings.Tile.STATE_ACTIVE
 import android.service.quicksettings.Tile.STATE_INACTIVE
 import android.text.TextUtils
-import android.util.Log
 import androidx.appcompat.content.res.AppCompatResources
 import androidx.compose.animation.graphics.ExperimentalAnimationGraphicsApi
 import androidx.compose.animation.graphics.res.animatedVectorResource
@@ -465,7 +464,6 @@
     animateToEnd: Boolean = false,
     modifier: Modifier = Modifier,
 ) {
-    Log.d("Fabian", "Recomposing tile icon")
     val iconModifier = modifier.size(dimensionResource(id = R.dimen.qs_icon_size))
     val context = LocalContext.current
     val loadedDrawable =
diff --git a/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java b/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
index 04de2c2..c1caeed 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
+++ b/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
@@ -1774,8 +1774,9 @@
         // the small clock here
         // With migrateClocksToBlueprint, weather clock will have behaviors similar to other clocks
         if (!MigrateClocksToBlueprint.isEnabled()) {
+            boolean bypassEnabled = mKeyguardBypassController.getBypassEnabled();
             if (mKeyguardStatusViewController.isLargeClockBlockingNotificationShelf()
-                    && hasVisibleNotifications() && isOnAod()) {
+                    && hasVisibleNotifications() && (isOnAod() || bypassEnabled)) {
                 return SMALL;
             }
         }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModel.kt
index 4183cdd..f5e17df 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModel.kt
@@ -18,7 +18,6 @@
 
 import android.content.Context
 import androidx.annotation.DrawableRes
-import com.android.systemui.animation.DialogTransitionAnimator
 import com.android.systemui.common.shared.model.ContentDescription
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.dagger.SysUISingleton
@@ -58,7 +57,6 @@
     private val mediaProjectionChipInteractor: MediaProjectionChipInteractor,
     private val mediaRouterChipInteractor: MediaRouterChipInteractor,
     private val systemClock: SystemClock,
-    private val dialogTransitionAnimator: DialogTransitionAnimator,
     private val endMediaProjectionDialogHelper: EndMediaProjectionDialogHelper,
 ) : OngoingActivityChipViewModel {
     /**
@@ -175,7 +173,6 @@
             startTimeMs = systemClock.elapsedRealtime(),
             createDialogLaunchOnClickListener(
                 createCastScreenToOtherDeviceDialogDelegate(state),
-                dialogTransitionAnimator,
             ),
         )
     }
@@ -191,7 +188,6 @@
             colors = ColorsModel.Red,
             createDialogLaunchOnClickListener(
                 createGenericCastToOtherDeviceDialogDelegate(deviceName),
-                dialogTransitionAnimator,
             ),
         )
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModel.kt
index df25d57..e201652 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModel.kt
@@ -19,7 +19,6 @@
 import android.app.ActivityManager
 import android.content.Context
 import androidx.annotation.DrawableRes
-import com.android.systemui.animation.DialogTransitionAnimator
 import com.android.systemui.common.shared.model.ContentDescription
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.dagger.SysUISingleton
@@ -52,7 +51,6 @@
     private val interactor: ScreenRecordChipInteractor,
     private val systemClock: SystemClock,
     private val endMediaProjectionDialogHelper: EndMediaProjectionDialogHelper,
-    private val dialogTransitionAnimator: DialogTransitionAnimator,
 ) : OngoingActivityChipViewModel {
     override val chip: StateFlow<OngoingActivityChipModel> =
         interactor.screenRecordState
@@ -78,7 +76,6 @@
                             startTimeMs = systemClock.elapsedRealtime(),
                             createDialogLaunchOnClickListener(
                                 createDelegate(state.recordedTask),
-                                dialogTransitionAnimator,
                             ),
                         )
                     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModel.kt
index c097720..45260e18 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModel.kt
@@ -18,7 +18,6 @@
 
 import android.content.Context
 import androidx.annotation.DrawableRes
-import com.android.systemui.animation.DialogTransitionAnimator
 import com.android.systemui.common.shared.model.ContentDescription
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.dagger.SysUISingleton
@@ -52,7 +51,6 @@
     private val context: Context,
     private val mediaProjectionChipInteractor: MediaProjectionChipInteractor,
     private val systemClock: SystemClock,
-    private val dialogTransitionAnimator: DialogTransitionAnimator,
     private val endMediaProjectionDialogHelper: EndMediaProjectionDialogHelper,
 ) : OngoingActivityChipViewModel {
     override val chip: StateFlow<OngoingActivityChipModel> =
@@ -89,10 +87,7 @@
             colors = ColorsModel.Red,
             // TODO(b/332662551): Maybe use a MediaProjection API to fetch this time.
             startTimeMs = systemClock.elapsedRealtime(),
-            createDialogLaunchOnClickListener(
-                createShareToAppDialogDelegate(state),
-                dialogTransitionAnimator
-            ),
+            createDialogLaunchOnClickListener(createShareToAppDialogDelegate(state)),
         )
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipViewModel.kt
index 0dbf5d6..65f94ac 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipViewModel.kt
@@ -17,10 +17,7 @@
 package com.android.systemui.statusbar.chips.ui.viewmodel
 
 import android.view.View
-import com.android.systemui.animation.DialogTransitionAnimator
-import com.android.systemui.res.R
 import com.android.systemui.statusbar.chips.ui.model.OngoingActivityChipModel
-import com.android.systemui.statusbar.chips.ui.view.ChipBackgroundContainer
 import com.android.systemui.statusbar.phone.SystemUIDialog
 import kotlinx.coroutines.flow.StateFlow
 
@@ -36,19 +33,10 @@
         /** Creates a chip click listener that launches a dialog created by [dialogDelegate]. */
         fun createDialogLaunchOnClickListener(
             dialogDelegate: SystemUIDialog.Delegate,
-            dialogTransitionAnimator: DialogTransitionAnimator,
         ): View.OnClickListener {
             return View.OnClickListener { view ->
                 val dialog = dialogDelegate.createDialog()
-                val launchableView =
-                    view.requireViewById<ChipBackgroundContainer>(
-                        R.id.ongoing_activity_chip_background
-                    )
-                // TODO(b/343699052): This makes a beautiful animate-in, but the
-                //  animate-out looks odd because the dialog animates back into the chip
-                //  but then the chip disappears. If we aren't able to address
-                //  b/343699052 in time for launch, we should just use `dialog.show`.
-                dialogTransitionAnimator.showFromView(dialog, launchableView)
+                dialog.show()
             }
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/RichOngoingNotificationContentExtractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/RichOngoingNotificationContentExtractor.kt
index b5ea861..b8af369 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/RichOngoingNotificationContentExtractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/RichOngoingNotificationContentExtractor.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.statusbar.notification.row
 
 import android.app.Notification
+import android.app.Notification.RichOngoingStyle
 import android.app.PendingIntent
 import android.content.Context
 import android.util.Log
@@ -68,12 +69,14 @@
         builder: Notification.Builder,
         systemUIContext: Context,
         packageContext: Context
-    ): RichOngoingContentModel? =
+    ): RichOngoingContentModel? {
+        if (builder.style !is RichOngoingStyle) return null
+
         try {
             val sbn = entry.sbn
             val notification = sbn.notification
             val icon = IconModel(notification.smallIcon)
-            if (sbn.packageName == "com.google.android.deskclock") {
+            return if (sbn.packageName == "com.google.android.deskclock") {
                 when (notification.channelId) {
                     "Timers v2" -> {
                         parseTimerNotification(notification, icon)
@@ -90,8 +93,9 @@
             } else null
         } catch (e: Exception) {
             Log.e("RONs", "Error parsing RON", e)
-            null
+            return null
         }
+    }
 
     /**
      * FOR PROTOTYPING ONLY: create a RON TimerContentModel using the time information available
diff --git a/packages/SystemUI/src/com/android/systemui/util/ConvenienceExtensions.kt b/packages/SystemUI/src/com/android/systemui/util/ConvenienceExtensions.kt
index ec7aabb..f2132248 100644
--- a/packages/SystemUI/src/com/android/systemui/util/ConvenienceExtensions.kt
+++ b/packages/SystemUI/src/com/android/systemui/util/ConvenienceExtensions.kt
@@ -20,6 +20,7 @@
 import android.util.IndentingPrintWriter
 import android.view.View
 import android.view.ViewGroup
+import dagger.Lazy
 import java.io.PrintWriter
 
 /** [Sequence] that yields all of the direct children of this [ViewGroup] */
@@ -56,3 +57,8 @@
         getBoundsOnScreen(bounds)
         return bounds
     }
+
+/** Extension method to convert [dagger.Lazy] to [kotlin.Lazy] for object of any class [T]. */
+fun <T> Lazy<T>.toKotlinLazy(): kotlin.Lazy<T> {
+    return lazy { this.get() }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractorTest.kt
index 42ab25f..032794c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractorTest.kt
@@ -56,6 +56,7 @@
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
+import org.junit.Ignore
 import org.junit.runner.RunWith
 import org.mockito.Mockito.reset
 import org.mockito.Mockito.spy
@@ -97,6 +98,7 @@
 
     @Test
     @EnableFlags(Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR)
+    @Ignore("Until b/349837588 is fixed")
     fun testTransitionToOccluded_ifDreamEnds_occludingActivityOnTop() =
         testScope.runTest {
             kosmos.fakeKeyguardRepository.setDreaming(true)
@@ -156,6 +158,7 @@
             reset(transitionRepository)
 
             kosmos.keyguardOcclusionRepository.setShowWhenLockedActivityInfo(onTop = false)
+            kosmos.fakeKeyguardRepository.setDreaming(false)
             runCurrent()
 
             assertThat(transitionRepository)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/InWindowLauncherUnlockAnimationInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/InWindowLauncherUnlockAnimationInteractorTest.kt
index 459e41d..ea5a41f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/InWindowLauncherUnlockAnimationInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/InWindowLauncherUnlockAnimationInteractorTest.kt
@@ -18,25 +18,22 @@
 
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
-import com.android.systemui.SysUITestModule
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.TestMocksModule
-import com.android.systemui.biometrics.domain.BiometricsDomainLayerModule
 import com.android.systemui.coroutines.collectValues
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.keyguard.data.repository.FakeKeyguardSurfaceBehindRepository
 import com.android.systemui.keyguard.data.repository.FakeKeyguardTransitionRepository
-import com.android.systemui.keyguard.data.repository.InWindowLauncherUnlockAnimationRepository
+import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
+import com.android.systemui.keyguard.data.repository.inWindowLauncherUnlockAnimationRepository
+import com.android.systemui.keyguard.data.repository.keyguardSurfaceBehindRepository
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.TransitionStep
 import com.android.systemui.keyguard.util.mockTopActivityClassName
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.kosmos.testScope
 import com.android.systemui.shared.system.ActivityManagerWrapper
-import com.android.systemui.user.domain.UserDomainLayerModule
-import dagger.BindsInstance
-import dagger.Component
+import com.android.systemui.shared.system.activityManagerWrapper
+import com.android.systemui.testKosmos
 import junit.framework.Assert.assertEquals
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
@@ -49,10 +46,16 @@
 @RunWith(AndroidJUnit4::class)
 @kotlinx.coroutines.ExperimentalCoroutinesApi
 class InWindowLauncherUnlockAnimationInteractorTest : SysuiTestCase() {
-    private lateinit var underTest: InWindowLauncherUnlockAnimationInteractor
-
-    private lateinit var testComponent: TestComponent
-    private lateinit var testScope: TestScope
+    private val kosmos = testKosmos()
+    private val underTest =
+        InWindowLauncherUnlockAnimationInteractor(
+            kosmos.inWindowLauncherUnlockAnimationRepository,
+            kosmos.applicationCoroutineScope,
+            kosmos.keyguardTransitionInteractor,
+            { kosmos.keyguardSurfaceBehindRepository },
+            kosmos.activityManagerWrapper,
+        )
+    private val testScope = kosmos.testScope
     private lateinit var transitionRepository: FakeKeyguardTransitionRepository
     @Mock private lateinit var activityManagerWrapper: ActivityManagerWrapper
 
@@ -62,19 +65,9 @@
     fun setUp() {
         MockitoAnnotations.initMocks(this)
 
-        testComponent =
-            DaggerInWindowLauncherUnlockAnimationInteractorTest_TestComponent.factory()
-                .create(
-                    test = this,
-                    mocks =
-                        TestMocksModule(
-                            activityManagerWrapper = activityManagerWrapper,
-                        ),
-                )
-        underTest = testComponent.underTest
-        testScope = testComponent.testScope
-        transitionRepository = testComponent.transitionRepository
+        transitionRepository = kosmos.fakeKeyguardTransitionRepository
 
+        activityManagerWrapper = kosmos.activityManagerWrapper
         activityManagerWrapper.mockTopActivityClassName(launcherClassName)
     }
 
@@ -92,7 +85,7 @@
             )
 
             // Put launcher on top
-            testComponent.inWindowLauncherUnlockAnimationRepository.setLauncherActivityClass(
+            kosmos.inWindowLauncherUnlockAnimationRepository.setLauncherActivityClass(
                 launcherClassName
             )
             activityManagerWrapper.mockTopActivityClassName(launcherClassName)
@@ -175,7 +168,7 @@
             )
 
             // Put not launcher on top
-            testComponent.inWindowLauncherUnlockAnimationRepository.setLauncherActivityClass(
+            kosmos.inWindowLauncherUnlockAnimationRepository.setLauncherActivityClass(
                 launcherClassName
             )
             activityManagerWrapper.mockTopActivityClassName("not_launcher")
@@ -252,7 +245,7 @@
             )
 
             // Put launcher on top
-            testComponent.inWindowLauncherUnlockAnimationRepository.setLauncherActivityClass(
+            kosmos.inWindowLauncherUnlockAnimationRepository.setLauncherActivityClass(
                 launcherClassName
             )
             activityManagerWrapper.mockTopActivityClassName(launcherClassName)
@@ -296,7 +289,7 @@
             )
 
             // Put Launcher on top and begin transitioning to GONE.
-            testComponent.inWindowLauncherUnlockAnimationRepository.setLauncherActivityClass(
+            kosmos.inWindowLauncherUnlockAnimationRepository.setLauncherActivityClass(
                 launcherClassName
             )
             activityManagerWrapper.mockTopActivityClassName(launcherClassName)
@@ -316,7 +309,7 @@
                 values
             )
 
-            testComponent.surfaceBehindRepository.setSurfaceRemoteAnimationTargetAvailable(true)
+            kosmos.keyguardSurfaceBehindRepository.setSurfaceRemoteAnimationTargetAvailable(true)
             runCurrent()
 
             assertEquals(
@@ -360,7 +353,7 @@
             )
 
             // Put Launcher on top and begin transitioning to GONE.
-            testComponent.inWindowLauncherUnlockAnimationRepository.setLauncherActivityClass(
+            kosmos.inWindowLauncherUnlockAnimationRepository.setLauncherActivityClass(
                 launcherClassName
             )
             activityManagerWrapper.mockTopActivityClassName(launcherClassName)
@@ -402,7 +395,7 @@
             )
 
             // Put Launcher on top and begin transitioning to GONE.
-            testComponent.inWindowLauncherUnlockAnimationRepository.setLauncherActivityClass(
+            kosmos.inWindowLauncherUnlockAnimationRepository.setLauncherActivityClass(
                 launcherClassName
             )
             activityManagerWrapper.mockTopActivityClassName(launcherClassName)
@@ -427,7 +420,7 @@
                     to = KeyguardState.AOD,
                 )
             )
-            testComponent.surfaceBehindRepository.setSurfaceRemoteAnimationTargetAvailable(true)
+            kosmos.keyguardSurfaceBehindRepository.setSurfaceRemoteAnimationTargetAvailable(true)
             runCurrent()
 
             assertEquals(
@@ -437,29 +430,4 @@
                 values
             )
         }
-
-    @SysUISingleton
-    @Component(
-        modules =
-            [
-                SysUITestModule::class,
-                BiometricsDomainLayerModule::class,
-                UserDomainLayerModule::class,
-            ]
-    )
-    interface TestComponent {
-        val underTest: InWindowLauncherUnlockAnimationInteractor
-        val testScope: TestScope
-        val transitionRepository: FakeKeyguardTransitionRepository
-        val surfaceBehindRepository: FakeKeyguardSurfaceBehindRepository
-        val inWindowLauncherUnlockAnimationRepository: InWindowLauncherUnlockAnimationRepository
-
-        @Component.Factory
-        interface Factory {
-            fun create(
-                @BindsInstance test: SysuiTestCase,
-                mocks: TestMocksModule,
-            ): TestComponent
-        }
-    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardWakeDirectlyToGoneInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardWakeDirectlyToGoneInteractorTest.kt
new file mode 100644
index 0000000..7e249e8
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardWakeDirectlyToGoneInteractorTest.kt
@@ -0,0 +1,378 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.keyguard.domain.interactor
+
+import android.app.AlarmManager
+import android.app.admin.alarmManager
+import android.app.admin.devicePolicyManager
+import android.content.BroadcastReceiver
+import android.content.Intent
+import android.content.mockedContext
+import android.os.PowerManager
+import android.os.UserHandle
+import android.platform.test.annotations.EnableFlags
+import android.provider.Settings
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.internal.widget.lockPatternUtils
+import com.android.systemui.Flags
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectValues
+import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
+import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
+import com.android.systemui.keyguard.shared.model.BiometricUnlockMode
+import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAsleepForTest
+import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAwakeForTest
+import com.android.systemui.power.domain.interactor.powerInteractor
+import com.android.systemui.testKosmos
+import com.android.systemui.util.settings.fakeSettings
+import junit.framework.Assert.assertEquals
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentMatchers.anyInt
+import org.mockito.ArgumentMatchers.anyLong
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doAnswer
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class KeyguardWakeDirectlyToGoneInteractorTest : SysuiTestCase() {
+
+    private var lastRegisteredBroadcastReceiver: BroadcastReceiver? = null
+    private val kosmos =
+        testKosmos().apply {
+            whenever(mockedContext.user).thenReturn(mock<UserHandle>())
+            doAnswer { invocation ->
+                    lastRegisteredBroadcastReceiver = invocation.arguments[0] as BroadcastReceiver
+                }
+                .whenever(mockedContext)
+                .registerReceiver(any(), any(), any(), any(), any())
+        }
+
+    private val testScope = kosmos.testScope
+    private val underTest = kosmos.keyguardWakeDirectlyToGoneInteractor
+    private val lockPatternUtils = kosmos.lockPatternUtils
+    private val repository = kosmos.fakeKeyguardRepository
+    private val transitionRepository = kosmos.fakeKeyguardTransitionRepository
+
+    @Test
+    @EnableFlags(Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR)
+    fun testCanWakeDirectlyToGone_keyguardServiceEnabledThenDisabled() =
+        testScope.runTest {
+            val canWake by collectValues(underTest.canWakeDirectlyToGone)
+
+            assertEquals(
+                listOf(
+                    false, // Defaults to false.
+                ),
+                canWake
+            )
+
+            repository.setKeyguardEnabled(false)
+            runCurrent()
+
+            assertEquals(
+                listOf(
+                    false, // Default to false.
+                    true, // True now that keyguard service is disabled
+                ),
+                canWake
+            )
+
+            repository.setKeyguardEnabled(true)
+            runCurrent()
+
+            assertEquals(
+                listOf(
+                    false,
+                    true,
+                    false,
+                ),
+                canWake
+            )
+        }
+
+    @Test
+    @EnableFlags(Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR)
+    fun testCanWakeDirectlyToGone_lockscreenDisabledThenEnabled() =
+        testScope.runTest {
+            val canWake by collectValues(underTest.canWakeDirectlyToGone)
+
+            assertEquals(
+                listOf(
+                    false, // Defaults to false.
+                ),
+                canWake
+            )
+
+            whenever(lockPatternUtils.isLockScreenDisabled(anyInt())).thenReturn(true)
+            runCurrent()
+
+            assertEquals(
+                listOf(
+                    // Still false - isLockScreenDisabled only causes canWakeDirectlyToGone to
+                    // update on the next wake/sleep event.
+                    false,
+                ),
+                canWake
+            )
+
+            kosmos.powerInteractor.setAsleepForTest()
+            runCurrent()
+
+            assertEquals(
+                listOf(
+                    false,
+                    // True since we slept after setting isLockScreenDisabled=true
+                    true,
+                ),
+                canWake
+            )
+
+            kosmos.powerInteractor.setAwakeForTest()
+            runCurrent()
+
+            kosmos.powerInteractor.setAsleepForTest()
+            runCurrent()
+
+            assertEquals(
+                listOf(
+                    false,
+                    true,
+                ),
+                canWake
+            )
+
+            whenever(lockPatternUtils.isLockScreenDisabled(anyInt())).thenReturn(false)
+            kosmos.powerInteractor.setAwakeForTest()
+            runCurrent()
+
+            assertEquals(
+                listOf(
+                    false,
+                    true,
+                    false,
+                ),
+                canWake
+            )
+        }
+
+    @Test
+    @EnableFlags(Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR)
+    fun testCanWakeDirectlyToGone_wakeAndUnlock() =
+        testScope.runTest {
+            val canWake by collectValues(underTest.canWakeDirectlyToGone)
+
+            assertEquals(
+                listOf(
+                    false, // Defaults to false.
+                ),
+                canWake
+            )
+
+            repository.setBiometricUnlockState(BiometricUnlockMode.WAKE_AND_UNLOCK)
+            runCurrent()
+
+            assertEquals(listOf(false, true), canWake)
+
+            repository.setBiometricUnlockState(BiometricUnlockMode.NONE)
+            runCurrent()
+
+            assertEquals(listOf(false, true, false), canWake)
+        }
+
+    @Test
+    @EnableFlags(Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR)
+    fun testCanWakeDirectlyToGone_andSetsAlarm_ifPowerButtonDoesNotLockImmediately() =
+        testScope.runTest {
+            val canWake by collectValues(underTest.canWakeDirectlyToGone)
+
+            assertEquals(
+                listOf(
+                    false, // Defaults to false.
+                ),
+                canWake
+            )
+
+            repository.setCanIgnoreAuthAndReturnToGone(true)
+            runCurrent()
+
+            assertEquals(listOf(false, true), canWake)
+
+            repository.setCanIgnoreAuthAndReturnToGone(false)
+            runCurrent()
+
+            assertEquals(listOf(false, true, false), canWake)
+        }
+
+    @Test
+    @EnableFlags(Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR)
+    fun testSetsCanIgnoreAuth_andSetsAlarm_whenTimingOut() =
+        testScope.runTest {
+            val canWake by collectValues(underTest.canWakeDirectlyToGone)
+
+            assertEquals(
+                listOf(
+                    false, // Defaults to false.
+                ),
+                canWake
+            )
+
+            whenever(kosmos.devicePolicyManager.getMaximumTimeToLock(eq(null), anyInt()))
+                .thenReturn(-1)
+            kosmos.fakeSettings.putInt(Settings.Secure.LOCK_SCREEN_LOCK_AFTER_TIMEOUT, 500)
+
+            transitionRepository.sendTransitionSteps(
+                from = KeyguardState.LOCKSCREEN,
+                to = KeyguardState.GONE,
+                testScope,
+            )
+
+            kosmos.powerInteractor.setAsleepForTest(
+                sleepReason = PowerManager.GO_TO_SLEEP_REASON_TIMEOUT
+            )
+            runCurrent()
+
+            assertEquals(
+                listOf(
+                    false,
+                    true,
+                ),
+                canWake
+            )
+
+            verify(kosmos.alarmManager)
+                .setExactAndAllowWhileIdle(
+                    eq(AlarmManager.ELAPSED_REALTIME_WAKEUP),
+                    anyLong(),
+                    any(),
+                )
+        }
+
+    @Test
+    @EnableFlags(Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR)
+    fun testCancelsFirstAlarm_onWake_withSecondAlarmSet() =
+        testScope.runTest {
+            val canWake by collectValues(underTest.canWakeDirectlyToGone)
+
+            assertEquals(
+                listOf(
+                    false, // Defaults to false.
+                ),
+                canWake
+            )
+
+            whenever(kosmos.devicePolicyManager.getMaximumTimeToLock(eq(null), anyInt()))
+                .thenReturn(-1)
+            kosmos.fakeSettings.putInt(Settings.Secure.LOCK_SCREEN_LOCK_AFTER_TIMEOUT, 500)
+
+            transitionRepository.sendTransitionSteps(
+                from = KeyguardState.LOCKSCREEN,
+                to = KeyguardState.GONE,
+                testScope,
+            )
+
+            kosmos.powerInteractor.setAsleepForTest(
+                sleepReason = PowerManager.GO_TO_SLEEP_REASON_TIMEOUT
+            )
+            transitionRepository.sendTransitionSteps(
+                from = KeyguardState.LOCKSCREEN,
+                to = KeyguardState.AOD,
+                testScope = testScope,
+            )
+            runCurrent()
+
+            assertEquals(
+                listOf(
+                    false,
+                    // Timed out, so we can ignore auth/return to GONE.
+                    true,
+                ),
+                canWake
+            )
+
+            verify(kosmos.alarmManager)
+                .setExactAndAllowWhileIdle(
+                    eq(AlarmManager.ELAPSED_REALTIME_WAKEUP),
+                    anyLong(),
+                    any(),
+                )
+
+            kosmos.powerInteractor.setAwakeForTest()
+            transitionRepository.sendTransitionSteps(
+                from = KeyguardState.AOD,
+                to = KeyguardState.GONE,
+                testScope = testScope,
+            )
+            runCurrent()
+
+            assertEquals(
+                listOf(
+                    false,
+                    true,
+                    // Should be canceled by the wakeup, but there would still be an
+                    // alarm in flight that should be canceled.
+                    false,
+                ),
+                canWake
+            )
+
+            kosmos.powerInteractor.setAsleepForTest(
+                sleepReason = PowerManager.GO_TO_SLEEP_REASON_TIMEOUT
+            )
+            runCurrent()
+
+            assertEquals(
+                listOf(
+                    false,
+                    true,
+                    false,
+                    // Back to sleep.
+                    true,
+                ),
+                canWake
+            )
+
+            // Simulate the first sleep's alarm coming in.
+            lastRegisteredBroadcastReceiver?.onReceive(
+                kosmos.mockedContext,
+                Intent("com.android.internal.policy.impl.PhoneWindowManager.DELAYED_KEYGUARD")
+            )
+            runCurrent()
+
+            // It should not have any effect.
+            assertEquals(
+                listOf(
+                    false,
+                    true,
+                    false,
+                    true,
+                ),
+                canWake
+            )
+        }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/binder/InWindowLauncherUnlockAnimationManagerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/binder/InWindowLauncherUnlockAnimationManagerTest.kt
index c7f4416..0cfc20d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/binder/InWindowLauncherUnlockAnimationManagerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/binder/InWindowLauncherUnlockAnimationManagerTest.kt
@@ -18,17 +18,15 @@
 
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
-import com.android.systemui.SysUITestModule
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.biometrics.domain.BiometricsDomainLayerModule
-import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.keyguard.domain.interactor.inWindowLauncherUnlockAnimationInteractor
 import com.android.systemui.keyguard.ui.view.InWindowLauncherUnlockAnimationManager
+import com.android.systemui.keyguard.ui.viewmodel.InWindowLauncherAnimationViewModel
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.kosmos.testScope
 import com.android.systemui.shared.system.smartspace.ILauncherUnlockAnimationController
-import com.android.systemui.user.domain.UserDomainLayerModule
+import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.any
-import dagger.BindsInstance
-import dagger.Component
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
@@ -45,10 +43,9 @@
 @RunWith(AndroidJUnit4::class)
 @kotlinx.coroutines.ExperimentalCoroutinesApi
 class InWindowLauncherUnlockAnimationManagerTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
     private lateinit var underTest: InWindowLauncherUnlockAnimationManager
-
-    private lateinit var testComponent: TestComponent
-    private lateinit var testScope: TestScope
+    private val testScope = kosmos.testScope
 
     @Mock private lateinit var launcherUnlockAnimationController: ILauncherUnlockAnimationController
 
@@ -56,14 +53,14 @@
     fun setUp() {
         MockitoAnnotations.initMocks(this)
 
-        testComponent =
-            DaggerInWindowLauncherUnlockAnimationManagerTest_TestComponent.factory()
-                .create(
-                    test = this,
-                )
-        underTest = testComponent.underTest
-        testScope = testComponent.testScope
-
+        underTest =
+            InWindowLauncherUnlockAnimationManager(
+                kosmos.inWindowLauncherUnlockAnimationInteractor,
+                InWindowLauncherAnimationViewModel(
+                    kosmos.inWindowLauncherUnlockAnimationInteractor
+                ),
+                kosmos.applicationCoroutineScope
+            )
         underTest.setLauncherUnlockController("launcherClass", launcherUnlockAnimationController)
     }
 
@@ -114,25 +111,4 @@
 
             verifyNoMoreInteractions(launcherUnlockAnimationController)
         }
-
-    @SysUISingleton
-    @Component(
-        modules =
-            [
-                SysUITestModule::class,
-                BiometricsDomainLayerModule::class,
-                UserDomainLayerModule::class,
-            ]
-    )
-    interface TestComponent {
-        val underTest: InWindowLauncherUnlockAnimationManager
-        val testScope: TestScope
-
-        @Component.Factory
-        interface Factory {
-            fun create(
-                @BindsInstance test: SysuiTestCase,
-            ): TestComponent
-        }
-    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/LockscreenShadeTransitionControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/LockscreenShadeTransitionControllerTest.kt
index 69e8f47..9e6a498 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/LockscreenShadeTransitionControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/LockscreenShadeTransitionControllerTest.kt
@@ -6,26 +6,23 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.ExpandHelper
-import com.android.systemui.SysUITestModule
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.TestMocksModule
-import com.android.systemui.biometrics.domain.BiometricsDomainLayerModule
 import com.android.systemui.classifier.FalsingCollectorFake
 import com.android.systemui.classifier.FalsingManagerFake
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.flags.FakeFeatureFlagsClassicModule
 import com.android.systemui.flags.Flags
+import com.android.systemui.flags.fakeFeatureFlagsClassic
 import com.android.systemui.keyguard.domain.interactor.NaturalScrollingSettingObserver
+import com.android.systemui.kosmos.testScope
 import com.android.systemui.media.controls.ui.controller.MediaHierarchyManager
 import com.android.systemui.plugins.qs.QS
-import com.android.systemui.power.domain.interactor.PowerInteractor
 import com.android.systemui.qs.ui.adapter.FakeQSSceneAdapter
 import com.android.systemui.res.R
+import com.android.systemui.shade.data.repository.shadeRepository
 import com.android.systemui.shade.domain.interactor.ShadeLockscreenInteractor
-import com.android.systemui.shade.data.repository.FakeShadeRepository
-import com.android.systemui.shade.domain.interactor.ShadeInteractor
+import com.android.systemui.shade.domain.interactor.shadeInteractor
 import com.android.systemui.statusbar.disableflags.data.model.DisableFlagsModel
-import com.android.systemui.statusbar.disableflags.data.repository.FakeDisableFlagsRepository
+import com.android.systemui.statusbar.disableflags.data.repository.disableFlagsRepository
+import com.android.systemui.statusbar.disableflags.data.repository.fakeDisableFlagsRepository
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
 import com.android.systemui.statusbar.notification.row.NotificationTestHelper
 import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayout
@@ -33,17 +30,16 @@
 import com.android.systemui.statusbar.phone.CentralSurfaces
 import com.android.systemui.statusbar.phone.KeyguardBypassController
 import com.android.systemui.statusbar.phone.ScrimController
-import com.android.systemui.statusbar.policy.FakeConfigurationController
 import com.android.systemui.statusbar.policy.ResourcesSplitShadeStateController
-import com.android.systemui.user.domain.UserDomainLayerModule
+import com.android.systemui.statusbar.policy.configurationController
+import com.android.systemui.statusbar.policy.fakeConfigurationController
+import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.argumentCaptor
 import com.android.systemui.util.mockito.mock
-import dagger.BindsInstance
-import dagger.Component
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
 import org.junit.After
 import org.junit.Assert.assertFalse
 import org.junit.Assert.assertNotNull
@@ -65,8 +61,8 @@
 import org.mockito.Mockito.never
 import org.mockito.Mockito.verify
 import org.mockito.Mockito.verifyZeroInteractions
-import org.mockito.junit.MockitoJUnit
 import org.mockito.Mockito.`when` as whenever
+import org.mockito.junit.MockitoJUnit
 
 private fun <T> anyObject(): T {
     return Mockito.anyObject<T>()
@@ -77,15 +73,14 @@
 @RunWith(AndroidJUnit4::class)
 @OptIn(ExperimentalCoroutinesApi::class)
 class LockscreenShadeTransitionControllerTest : SysuiTestCase() {
-
+    private val kosmos =
+        testKosmos().apply {
+            fakeFeatureFlagsClassic.apply { set(Flags.FULL_SCREEN_USER_SWITCHER, false) }
+        }
     private lateinit var transitionController: LockscreenShadeTransitionController
-    private lateinit var testComponent: TestComponent
-    private val configurationController
-        get() = testComponent.configurationController
-    private val disableFlagsRepository
-        get() = testComponent.disableFlagsRepository
-    private val testScope
-        get() = testComponent.testScope
+    private val configurationController = kosmos.fakeConfigurationController
+    private val disableFlagsRepository = kosmos.fakeDisableFlagsRepository
+    private val testScope = kosmos.testScope
 
     private val qsSceneAdapter = FakeQSSceneAdapter({ mock() })
 
@@ -134,26 +129,6 @@
         whenever(keyguardBypassController.bypassEnabled).thenReturn(false)
         whenever(naturalScrollingSettingObserver.isNaturalScrollingEnabled).thenReturn(true)
 
-        testComponent =
-            DaggerLockscreenShadeTransitionControllerTest_TestComponent.factory()
-                .create(
-                    test = this,
-                    featureFlags =
-                        FakeFeatureFlagsClassicModule {
-                            set(Flags.FULL_SCREEN_USER_SWITCHER, false)
-                        },
-                    mocks =
-                        TestMocksModule(
-                            notificationShadeDepthController = depthController,
-                            keyguardBypassController = keyguardBypassController,
-                            mediaHierarchyManager = mediaHierarchyManager,
-                            notificationLockscreenUserManager = lockScreenUserManager,
-                            notificationStackScrollLayoutController = nsslController,
-                            scrimController = scrimController,
-                            statusBarStateController = statusbarStateController,
-                        )
-                )
-
         transitionController =
             LockscreenShadeTransitionController(
                 statusBarStateController = statusbarStateController,
@@ -191,10 +166,10 @@
                 falsingManager = FalsingManagerFake(),
                 dumpManager = mock(),
                 qsTransitionControllerFactory = { qsTransitionController },
-                shadeRepository = testComponent.shadeRepository,
-                shadeInteractor = testComponent.shadeInteractor,
+                shadeRepository = kosmos.shadeRepository,
+                shadeInteractor = kosmos.shadeInteractor,
                 splitShadeStateController = ResourcesSplitShadeStateController(),
-                shadeLockscreenInteractorLazy = {shadeLockscreenInteractor},
+                shadeLockscreenInteractorLazy = { shadeLockscreenInteractor },
                 naturalScrollingSettingObserver = naturalScrollingSettingObserver,
                 lazyQSSceneAdapter = { qsSceneAdapter }
             )
@@ -214,387 +189,424 @@
     }
 
     @Test
-    fun testCantDragDownWhenQSExpanded() {
-        assertTrue("Can't drag down on keyguard", transitionController.canDragDown())
-        whenever(qS.isFullyCollapsed).thenReturn(false)
-        assertFalse("Can drag down when QS is expanded", transitionController.canDragDown())
-    }
+    fun testCantDragDownWhenQSExpanded() =
+        testScope.runTest {
+            assertTrue("Can't drag down on keyguard", transitionController.canDragDown())
+            whenever(qS.isFullyCollapsed).thenReturn(false)
+            assertFalse("Can drag down when QS is expanded", transitionController.canDragDown())
+        }
 
     @Test
-    fun testCanDragDownInLockedDownShade() {
-        whenever(statusbarStateController.state).thenReturn(StatusBarState.SHADE_LOCKED)
-        assertFalse("Can drag down in shade locked", transitionController.canDragDown())
-        whenever(nsslController.isInLockedDownShade).thenReturn(true)
-        assertTrue("Can't drag down in locked down shade", transitionController.canDragDown())
-    }
+    fun testCanDragDownInLockedDownShade() =
+        testScope.runTest {
+            whenever(statusbarStateController.state).thenReturn(StatusBarState.SHADE_LOCKED)
+            assertFalse("Can drag down in shade locked", transitionController.canDragDown())
+            whenever(nsslController.isInLockedDownShade).thenReturn(true)
+            assertTrue("Can't drag down in locked down shade", transitionController.canDragDown())
+        }
 
     @Test
-    fun testGoingToLockedShade() {
-        transitionController.goToLockedShade(null)
-        verify(statusbarStateController).setState(StatusBarState.SHADE_LOCKED)
-    }
+    fun testGoingToLockedShade() =
+        testScope.runTest {
+            transitionController.goToLockedShade(null)
+            verify(statusbarStateController).setState(StatusBarState.SHADE_LOCKED)
+        }
 
     @Test
-    fun testWakingToShadeLockedWhenDozing() {
-        whenever(statusbarStateController.isDozing).thenReturn(true)
-        transitionController.goToLockedShade(null)
-        verify(statusbarStateController).setState(StatusBarState.SHADE_LOCKED)
-        assertTrue("Not waking to shade locked", transitionController.isWakingToShadeLocked)
-    }
+    fun testWakingToShadeLockedWhenDozing() =
+        testScope.runTest {
+            whenever(statusbarStateController.isDozing).thenReturn(true)
+            transitionController.goToLockedShade(null)
+            verify(statusbarStateController).setState(StatusBarState.SHADE_LOCKED)
+            assertTrue("Not waking to shade locked", transitionController.isWakingToShadeLocked)
+        }
 
     @Test
-    fun testNotWakingToShadeLockedWhenNotDozing() {
-        whenever(statusbarStateController.isDozing).thenReturn(false)
-        transitionController.goToLockedShade(null)
-        verify(statusbarStateController).setState(StatusBarState.SHADE_LOCKED)
-        assertFalse(
-            "Waking to shade locked when not dozing",
-            transitionController.isWakingToShadeLocked
-        )
-    }
-
-    @Test
-    fun testGoToLockedShadeOnlyOnKeyguard() {
-        whenever(statusbarStateController.state).thenReturn(StatusBarState.SHADE_LOCKED)
-        transitionController.goToLockedShade(null)
-        whenever(statusbarStateController.state).thenReturn(StatusBarState.SHADE)
-        transitionController.goToLockedShade(null)
-        verify(statusbarStateController, never()).setState(anyInt())
-    }
-
-    @Test
-    fun testDontGoWhenShadeDisabled() {
-        disableFlagsRepository.disableFlags.value =
-            DisableFlagsModel(
-                disable2 = DISABLE2_NOTIFICATION_SHADE,
+    fun testNotWakingToShadeLockedWhenNotDozing() =
+        testScope.runTest {
+            whenever(statusbarStateController.isDozing).thenReturn(false)
+            transitionController.goToLockedShade(null)
+            verify(statusbarStateController).setState(StatusBarState.SHADE_LOCKED)
+            assertFalse(
+                "Waking to shade locked when not dozing",
+                transitionController.isWakingToShadeLocked
             )
-        testScope.runCurrent()
-        transitionController.goToLockedShade(null)
-        verify(statusbarStateController, never()).setState(anyInt())
-    }
+        }
 
     @Test
-    fun testUserExpandsViewOnGoingToFullShade() {
-        assertFalse("Row shouldn't be user expanded yet", row.isUserExpanded)
-        transitionController.goToLockedShade(row)
-        assertTrue("Row wasn't user expanded on drag down", row.isUserExpanded)
-    }
+    fun testGoToLockedShadeOnlyOnKeyguard() =
+        testScope.runTest {
+            whenever(statusbarStateController.state).thenReturn(StatusBarState.SHADE_LOCKED)
+            transitionController.goToLockedShade(null)
+            whenever(statusbarStateController.state).thenReturn(StatusBarState.SHADE)
+            transitionController.goToLockedShade(null)
+            verify(statusbarStateController, never()).setState(anyInt())
+        }
 
     @Test
-    fun testTriggeringBouncerNoNotificationsOnLockscreen() {
-        whenever(lockScreenUserManager.shouldShowLockscreenNotifications()).thenReturn(false)
-        transitionController.goToLockedShade(null)
-        verify(statusbarStateController, never()).setState(anyInt())
-        verify(statusbarStateController).setLeaveOpenOnKeyguardHide(true)
-        verify(centralSurfaces).showBouncerWithDimissAndCancelIfKeyguard(anyObject(), anyObject())
-    }
+    fun testDontGoWhenShadeDisabled() =
+        testScope.runTest {
+            disableFlagsRepository.disableFlags.value =
+                DisableFlagsModel(
+                    disable2 = DISABLE2_NOTIFICATION_SHADE,
+                )
+            testScope.runCurrent()
+            transitionController.goToLockedShade(null)
+            verify(statusbarStateController, never()).setState(anyInt())
+        }
 
     @Test
-    fun testGoToLockedShadeCreatesQSAnimation() {
-        transitionController.goToLockedShade(null)
-        verify(statusbarStateController).setState(StatusBarState.SHADE_LOCKED)
-        verify(shadeLockscreenInteractor).transitionToExpandedShade(anyLong())
-        assertNotNull(transitionController.dragDownAnimator)
-    }
+    fun testUserExpandsViewOnGoingToFullShade() =
+        testScope.runTest {
+            assertFalse("Row shouldn't be user expanded yet", row.isUserExpanded)
+            transitionController.goToLockedShade(row)
+            assertTrue("Row wasn't user expanded on drag down", row.isUserExpanded)
+        }
 
     @Test
-    fun testGoToLockedShadeDoesntCreateQSAnimation() {
-        transitionController.goToLockedShade(null, needsQSAnimation = false)
-        verify(statusbarStateController).setState(StatusBarState.SHADE_LOCKED)
-        verify(shadeLockscreenInteractor).transitionToExpandedShade(anyLong())
-        assertNull(transitionController.dragDownAnimator)
-    }
+    fun testTriggeringBouncerNoNotificationsOnLockscreen() =
+        testScope.runTest {
+            whenever(lockScreenUserManager.shouldShowLockscreenNotifications()).thenReturn(false)
+            transitionController.goToLockedShade(null)
+            verify(statusbarStateController, never()).setState(anyInt())
+            verify(statusbarStateController).setLeaveOpenOnKeyguardHide(true)
+            verify(centralSurfaces)
+                .showBouncerWithDimissAndCancelIfKeyguard(anyObject(), anyObject())
+        }
 
     @Test
-    fun testGoToLockedShadeAlwaysCreatesQSAnimationInSplitShade() {
-        enableSplitShade()
-        transitionController.goToLockedShade(null, needsQSAnimation = true)
-        verify(shadeLockscreenInteractor).transitionToExpandedShade(anyLong())
-        assertNotNull(transitionController.dragDownAnimator)
-    }
+    fun testGoToLockedShadeCreatesQSAnimation() =
+        testScope.runTest {
+            transitionController.goToLockedShade(null)
+            verify(statusbarStateController).setState(StatusBarState.SHADE_LOCKED)
+            verify(shadeLockscreenInteractor).transitionToExpandedShade(anyLong())
+            assertNotNull(transitionController.dragDownAnimator)
+        }
 
     @Test
-    fun testGoToLockedShadeCancelDoesntLeaveShadeOpenOnKeyguardHide() {
-        whenever(lockScreenUserManager.shouldShowLockscreenNotifications()).thenReturn(false)
-        whenever(lockScreenUserManager.isLockscreenPublicMode(any())).thenReturn(true)
-        transitionController.goToLockedShade(null)
-        val captor = argumentCaptor<Runnable>()
-        verify(centralSurfaces).showBouncerWithDimissAndCancelIfKeyguard(isNull(), captor.capture())
-        captor.value.run()
-        verify(statusbarStateController).setLeaveOpenOnKeyguardHide(false)
-    }
+    fun testGoToLockedShadeDoesntCreateQSAnimation() =
+        testScope.runTest {
+            transitionController.goToLockedShade(null, needsQSAnimation = false)
+            verify(statusbarStateController).setState(StatusBarState.SHADE_LOCKED)
+            verify(shadeLockscreenInteractor).transitionToExpandedShade(anyLong())
+            assertNull(transitionController.dragDownAnimator)
+        }
 
     @Test
-    fun testDragDownAmountDoesntCallOutInLockedDownShade() {
-        whenever(nsslController.isInLockedDownShade).thenReturn(true)
-        transitionController.dragDownAmount = 10f
-        verify(nsslController, never()).setTransitionToFullShadeAmount(anyFloat())
-        verify(mediaHierarchyManager, never()).setTransitionToFullShadeAmount(anyFloat())
-        verify(scrimController, never()).setTransitionToFullShadeProgress(anyFloat(), anyFloat())
-        verify(transitionControllerCallback, never())
-            .setTransitionToFullShadeAmount(anyFloat(), anyBoolean(), anyLong())
-        verify(qsTransitionController, never()).dragDownAmount = anyFloat()
-    }
+    fun testGoToLockedShadeAlwaysCreatesQSAnimationInSplitShade() =
+        testScope.runTest {
+            enableSplitShade()
+            transitionController.goToLockedShade(null, needsQSAnimation = true)
+            verify(shadeLockscreenInteractor).transitionToExpandedShade(anyLong())
+            assertNotNull(transitionController.dragDownAnimator)
+        }
 
     @Test
-    fun testDragDownAmountCallsOut() {
-        transitionController.dragDownAmount = 10f
-        verify(nsslController).setTransitionToFullShadeAmount(anyFloat())
-        verify(mediaHierarchyManager).setTransitionToFullShadeAmount(anyFloat())
-        verify(scrimController).setTransitionToFullShadeProgress(anyFloat(), anyFloat())
-        verify(transitionControllerCallback)
-            .setTransitionToFullShadeAmount(anyFloat(), anyBoolean(), anyLong())
-        verify(qsTransitionController).dragDownAmount = 10f
-        verify(depthController).transitionToFullShadeProgress = anyFloat()
-    }
+    fun testGoToLockedShadeCancelDoesntLeaveShadeOpenOnKeyguardHide() =
+        testScope.runTest {
+            whenever(lockScreenUserManager.shouldShowLockscreenNotifications()).thenReturn(false)
+            whenever(lockScreenUserManager.isLockscreenPublicMode(any())).thenReturn(true)
+            transitionController.goToLockedShade(null)
+            val captor = argumentCaptor<Runnable>()
+            verify(centralSurfaces)
+                .showBouncerWithDimissAndCancelIfKeyguard(isNull(), captor.capture())
+            captor.value.run()
+            verify(statusbarStateController).setLeaveOpenOnKeyguardHide(false)
+        }
 
     @Test
-    fun testDragDownAmount_depthDistanceIsZero_setsProgressToZero() {
-        context
-            .getOrCreateTestableResources()
-            .addOverride(R.dimen.lockscreen_shade_depth_controller_transition_distance, 0)
-        configurationController.notifyConfigurationChanged()
-
-        transitionController.dragDownAmount = 10f
-
-        verify(depthController).transitionToFullShadeProgress = 0f
-    }
+    fun testDragDownAmountDoesntCallOutInLockedDownShade() =
+        testScope.runTest {
+            whenever(nsslController.isInLockedDownShade).thenReturn(true)
+            transitionController.dragDownAmount = 10f
+            verify(nsslController, never()).setTransitionToFullShadeAmount(anyFloat())
+            verify(mediaHierarchyManager, never()).setTransitionToFullShadeAmount(anyFloat())
+            verify(scrimController, never())
+                .setTransitionToFullShadeProgress(anyFloat(), anyFloat())
+            verify(transitionControllerCallback, never())
+                .setTransitionToFullShadeAmount(anyFloat(), anyBoolean(), anyLong())
+            verify(qsTransitionController, never()).dragDownAmount = anyFloat()
+        }
 
     @Test
-    fun testDragDownAmount_depthDistanceNonZero_setsProgressBasedOnDistance() {
-        context
-            .getOrCreateTestableResources()
-            .addOverride(R.dimen.lockscreen_shade_depth_controller_transition_distance, 100)
-        configurationController.notifyConfigurationChanged()
-
-        transitionController.dragDownAmount = 10f
-
-        verify(depthController).transitionToFullShadeProgress = 0.1f
-    }
+    fun testDragDownAmountCallsOut() =
+        testScope.runTest {
+            transitionController.dragDownAmount = 10f
+            verify(nsslController).setTransitionToFullShadeAmount(anyFloat())
+            verify(mediaHierarchyManager).setTransitionToFullShadeAmount(anyFloat())
+            verify(scrimController).setTransitionToFullShadeProgress(anyFloat(), anyFloat())
+            verify(transitionControllerCallback)
+                .setTransitionToFullShadeAmount(anyFloat(), anyBoolean(), anyLong())
+            verify(qsTransitionController).dragDownAmount = 10f
+            verify(depthController).transitionToFullShadeProgress = anyFloat()
+        }
 
     @Test
-    fun setDragAmount_setsKeyguardTransitionProgress() {
-        transitionController.dragDownAmount = 10f
+    fun testDragDownAmount_depthDistanceIsZero_setsProgressToZero() =
+        testScope.runTest {
+            context
+                .getOrCreateTestableResources()
+                .addOverride(R.dimen.lockscreen_shade_depth_controller_transition_distance, 0)
+            configurationController.notifyConfigurationChanged()
 
-        verify(shadeLockscreenInteractor).setKeyguardTransitionProgress(anyFloat(), anyInt())
-    }
+            transitionController.dragDownAmount = 10f
+
+            verify(depthController).transitionToFullShadeProgress = 0f
+        }
 
     @Test
-    fun setDragAmount_setsKeyguardAlphaBasedOnDistance() {
-        val alphaDistance =
-            context.resources.getDimensionPixelSize(
-                R.dimen.lockscreen_shade_npvc_keyguard_content_alpha_transition_distance
-            )
-        transitionController.dragDownAmount = 10f
+    fun testDragDownAmount_depthDistanceNonZero_setsProgressBasedOnDistance() =
+        testScope.runTest {
+            context
+                .getOrCreateTestableResources()
+                .addOverride(R.dimen.lockscreen_shade_depth_controller_transition_distance, 100)
+            configurationController.notifyConfigurationChanged()
 
-        val expectedAlpha = 1 - 10f / alphaDistance
-        verify(shadeLockscreenInteractor).setKeyguardTransitionProgress(eq(expectedAlpha), anyInt())
-    }
+            transitionController.dragDownAmount = 10f
+
+            verify(depthController).transitionToFullShadeProgress = 0.1f
+        }
 
     @Test
-    fun setDragAmount_notInSplitShade_setsKeyguardTranslationToZero() {
-        val mediaTranslationY = 123
-        disableSplitShade()
-        whenever(mediaHierarchyManager.isCurrentlyInGuidedTransformation()).thenReturn(true)
-        whenever(mediaHierarchyManager.getGuidedTransformationTranslationY())
-            .thenReturn(mediaTranslationY)
+    fun setDragAmount_setsKeyguardTransitionProgress() =
+        testScope.runTest {
+            transitionController.dragDownAmount = 10f
 
-        transitionController.dragDownAmount = 10f
-
-        verify(shadeLockscreenInteractor).setKeyguardTransitionProgress(anyFloat(), eq(0))
-    }
+            verify(shadeLockscreenInteractor).setKeyguardTransitionProgress(anyFloat(), anyInt())
+        }
 
     @Test
-    fun setDragAmount_inSplitShade_setsKeyguardTranslationBasedOnMediaTranslation() {
-        val mediaTranslationY = 123
-        enableSplitShade()
-        whenever(mediaHierarchyManager.isCurrentlyInGuidedTransformation()).thenReturn(true)
-        whenever(mediaHierarchyManager.getGuidedTransformationTranslationY())
-            .thenReturn(mediaTranslationY)
+    fun setDragAmount_setsKeyguardAlphaBasedOnDistance() =
+        testScope.runTest {
+            val alphaDistance =
+                context.resources.getDimensionPixelSize(
+                    R.dimen.lockscreen_shade_npvc_keyguard_content_alpha_transition_distance
+                )
+            transitionController.dragDownAmount = 10f
 
-        transitionController.dragDownAmount = 10f
+            val expectedAlpha = 1 - 10f / alphaDistance
+            verify(shadeLockscreenInteractor)
+                .setKeyguardTransitionProgress(eq(expectedAlpha), anyInt())
+        }
 
-        verify(shadeLockscreenInteractor)
+    @Test
+    fun setDragAmount_notInSplitShade_setsKeyguardTranslationToZero() =
+        testScope.runTest {
+            val mediaTranslationY = 123
+            disableSplitShade()
+            whenever(mediaHierarchyManager.isCurrentlyInGuidedTransformation()).thenReturn(true)
+            whenever(mediaHierarchyManager.getGuidedTransformationTranslationY())
+                .thenReturn(mediaTranslationY)
+
+            transitionController.dragDownAmount = 10f
+
+            verify(shadeLockscreenInteractor).setKeyguardTransitionProgress(anyFloat(), eq(0))
+        }
+
+    @Test
+    fun setDragAmount_inSplitShade_setsKeyguardTranslationBasedOnMediaTranslation() =
+        testScope.runTest {
+            val mediaTranslationY = 123
+            enableSplitShade()
+            whenever(mediaHierarchyManager.isCurrentlyInGuidedTransformation()).thenReturn(true)
+            whenever(mediaHierarchyManager.getGuidedTransformationTranslationY())
+                .thenReturn(mediaTranslationY)
+
+            transitionController.dragDownAmount = 10f
+
+            verify(shadeLockscreenInteractor)
                 .setKeyguardTransitionProgress(anyFloat(), eq(mediaTranslationY))
-    }
+        }
 
     @Test
-    fun setDragAmount_inSplitShade_mediaNotShowing_setsKeyguardTranslationBasedOnDistance() {
-        enableSplitShade()
-        whenever(mediaHierarchyManager.isCurrentlyInGuidedTransformation()).thenReturn(false)
-        whenever(mediaHierarchyManager.getGuidedTransformationTranslationY()).thenReturn(123)
+    fun setDragAmount_inSplitShade_mediaNotShowing_setsKeyguardTranslationBasedOnDistance() =
+        testScope.runTest {
+            enableSplitShade()
+            whenever(mediaHierarchyManager.isCurrentlyInGuidedTransformation()).thenReturn(false)
+            whenever(mediaHierarchyManager.getGuidedTransformationTranslationY()).thenReturn(123)
 
-        transitionController.dragDownAmount = 10f
+            transitionController.dragDownAmount = 10f
 
-        val distance =
-            context.resources.getDimensionPixelSize(
-                R.dimen.lockscreen_shade_keyguard_transition_distance
+            val distance =
+                context.resources.getDimensionPixelSize(
+                    R.dimen.lockscreen_shade_keyguard_transition_distance
+                )
+            val offset =
+                context.resources.getDimensionPixelSize(
+                    R.dimen.lockscreen_shade_keyguard_transition_vertical_offset
+                )
+            val expectedTranslation = 10f / distance * offset
+            verify(shadeLockscreenInteractor)
+                .setKeyguardTransitionProgress(anyFloat(), eq(expectedTranslation.toInt()))
+        }
+
+    @Test
+    fun setDragDownAmount_setsValueOnMediaHierarchyManager() =
+        testScope.runTest {
+            transitionController.dragDownAmount = 10f
+
+            verify(mediaHierarchyManager).setTransitionToFullShadeAmount(10f)
+        }
+
+    @Test
+    fun setDragAmount_setsScrimProgressBasedOnScrimDistance() =
+        testScope.runTest {
+            val distance = 10
+            context.orCreateTestableResources.addOverride(
+                R.dimen.lockscreen_shade_scrim_transition_distance,
+                distance
             )
-        val offset =
-            context.resources.getDimensionPixelSize(
-                R.dimen.lockscreen_shade_keyguard_transition_vertical_offset
+            configurationController.notifyConfigurationChanged()
+
+            transitionController.dragDownAmount = 5f
+
+            verify(scrimController)
+                .transitionToFullShadeProgress(
+                    progress = eq(0.5f),
+                    lockScreenNotificationsProgress = anyFloat()
+                )
+        }
+
+    @Test
+    fun setDragAmount_setsNotificationsScrimProgressBasedOnNotificationsScrimDistanceAndDelay() =
+        testScope.runTest {
+            val distance = 100
+            val delay = 10
+            context.orCreateTestableResources.addOverride(
+                R.dimen.lockscreen_shade_notifications_scrim_transition_distance,
+                distance
             )
-        val expectedTranslation = 10f / distance * offset
-        verify(shadeLockscreenInteractor)
-            .setKeyguardTransitionProgress(anyFloat(), eq(expectedTranslation.toInt()))
-    }
-
-    @Test
-    fun setDragDownAmount_setsValueOnMediaHierarchyManager() {
-        transitionController.dragDownAmount = 10f
-
-        verify(mediaHierarchyManager).setTransitionToFullShadeAmount(10f)
-    }
-
-    @Test
-    fun setDragAmount_setsScrimProgressBasedOnScrimDistance() {
-        val distance = 10
-        context.orCreateTestableResources.addOverride(
-            R.dimen.lockscreen_shade_scrim_transition_distance,
-            distance
-        )
-        configurationController.notifyConfigurationChanged()
-
-        transitionController.dragDownAmount = 5f
-
-        verify(scrimController)
-            .transitionToFullShadeProgress(
-                progress = eq(0.5f),
-                lockScreenNotificationsProgress = anyFloat()
+            context.orCreateTestableResources.addOverride(
+                R.dimen.lockscreen_shade_notifications_scrim_transition_delay,
+                delay
             )
-    }
+            configurationController.notifyConfigurationChanged()
+
+            transitionController.dragDownAmount = 20f
+
+            verify(scrimController)
+                .transitionToFullShadeProgress(
+                    progress = anyFloat(),
+                    lockScreenNotificationsProgress = eq(0.1f)
+                )
+        }
 
     @Test
-    fun setDragAmount_setsNotificationsScrimProgressBasedOnNotificationsScrimDistanceAndDelay() {
-        val distance = 100
-        val delay = 10
-        context.orCreateTestableResources.addOverride(
-            R.dimen.lockscreen_shade_notifications_scrim_transition_distance,
-            distance
-        )
-        context.orCreateTestableResources.addOverride(
-            R.dimen.lockscreen_shade_notifications_scrim_transition_delay,
-            delay
-        )
-        configurationController.notifyConfigurationChanged()
-
-        transitionController.dragDownAmount = 20f
-
-        verify(scrimController)
-            .transitionToFullShadeProgress(
-                progress = anyFloat(),
-                lockScreenNotificationsProgress = eq(0.1f)
+    fun setDragAmount_dragAmountLessThanNotifDelayDistance_setsNotificationsScrimProgressToZero() =
+        testScope.runTest {
+            val distance = 100
+            val delay = 50
+            context.orCreateTestableResources.addOverride(
+                R.dimen.lockscreen_shade_notifications_scrim_transition_distance,
+                distance
             )
-    }
-
-    @Test
-    fun setDragAmount_dragAmountLessThanNotifDelayDistance_setsNotificationsScrimProgressToZero() {
-        val distance = 100
-        val delay = 50
-        context.orCreateTestableResources.addOverride(
-            R.dimen.lockscreen_shade_notifications_scrim_transition_distance,
-            distance
-        )
-        context.orCreateTestableResources.addOverride(
-            R.dimen.lockscreen_shade_notifications_scrim_transition_delay,
-            delay
-        )
-        configurationController.notifyConfigurationChanged()
-
-        transitionController.dragDownAmount = 20f
-
-        verify(scrimController)
-            .transitionToFullShadeProgress(
-                progress = anyFloat(),
-                lockScreenNotificationsProgress = eq(0f)
+            context.orCreateTestableResources.addOverride(
+                R.dimen.lockscreen_shade_notifications_scrim_transition_delay,
+                delay
             )
-    }
+            configurationController.notifyConfigurationChanged()
+
+            transitionController.dragDownAmount = 20f
+
+            verify(scrimController)
+                .transitionToFullShadeProgress(
+                    progress = anyFloat(),
+                    lockScreenNotificationsProgress = eq(0f)
+                )
+        }
 
     @Test
-    fun setDragAmount_dragAmountMoreThanTotalDistance_setsNotificationsScrimProgressToOne() {
-        val distance = 100
-        val delay = 50
-        context.orCreateTestableResources.addOverride(
-            R.dimen.lockscreen_shade_notifications_scrim_transition_distance,
-            distance
-        )
-        context.orCreateTestableResources.addOverride(
-            R.dimen.lockscreen_shade_notifications_scrim_transition_delay,
-            delay
-        )
-        configurationController.notifyConfigurationChanged()
-
-        transitionController.dragDownAmount = 999999f
-
-        verify(scrimController)
-            .transitionToFullShadeProgress(
-                progress = anyFloat(),
-                lockScreenNotificationsProgress = eq(1f)
+    fun setDragAmount_dragAmountMoreThanTotalDistance_setsNotificationsScrimProgressToOne() =
+        testScope.runTest {
+            val distance = 100
+            val delay = 50
+            context.orCreateTestableResources.addOverride(
+                R.dimen.lockscreen_shade_notifications_scrim_transition_distance,
+                distance
             )
-    }
-
-    @Test
-    fun setDragDownAmount_inSplitShade_setsValueOnMediaHierarchyManager() {
-        enableSplitShade()
-
-        transitionController.dragDownAmount = 10f
-
-        verify(mediaHierarchyManager).setTransitionToFullShadeAmount(10f)
-    }
-
-    @Test
-    fun setDragAmount_notInSplitShade_forwardsToSingleShadeOverScroller() {
-        disableSplitShade()
-
-        transitionController.dragDownAmount = 10f
-
-        verify(singleShadeOverScroller).expansionDragDownAmount = 10f
-        verifyZeroInteractions(splitShadeOverScroller)
-    }
-
-    @Test
-    fun setDragAmount_inSplitShade_forwardsToSplitShadeOverScroller() {
-        enableSplitShade()
-
-        transitionController.dragDownAmount = 10f
-
-        verify(splitShadeOverScroller).expansionDragDownAmount = 10f
-        verifyZeroInteractions(singleShadeOverScroller)
-    }
-
-    @Test
-    fun setDragDownAmount_inSplitShade_setsKeyguardStatusBarAlphaBasedOnDistance() {
-        val alphaDistance =
-            context.resources.getDimensionPixelSize(
-                R.dimen.lockscreen_shade_npvc_keyguard_content_alpha_transition_distance
+            context.orCreateTestableResources.addOverride(
+                R.dimen.lockscreen_shade_notifications_scrim_transition_delay,
+                delay
             )
-        val dragDownAmount = 10f
-        enableSplitShade()
+            configurationController.notifyConfigurationChanged()
 
-        transitionController.dragDownAmount = dragDownAmount
+            transitionController.dragDownAmount = 999999f
 
-        val expectedAlpha = 1 - dragDownAmount / alphaDistance
-        verify(shadeLockscreenInteractor).setKeyguardStatusBarAlpha(expectedAlpha)
-    }
+            verify(scrimController)
+                .transitionToFullShadeProgress(
+                    progress = anyFloat(),
+                    lockScreenNotificationsProgress = eq(1f)
+                )
+        }
 
     @Test
-    fun setDragDownAmount_notInSplitShade_setsKeyguardStatusBarAlphaToMinusOne() {
-        disableSplitShade()
+    fun setDragDownAmount_inSplitShade_setsValueOnMediaHierarchyManager() =
+        testScope.runTest {
+            enableSplitShade()
 
-        transitionController.dragDownAmount = 10f
+            transitionController.dragDownAmount = 10f
 
-        verify(shadeLockscreenInteractor).setKeyguardStatusBarAlpha(-1f)
-    }
+            verify(mediaHierarchyManager).setTransitionToFullShadeAmount(10f)
+        }
 
     @Test
-    fun nullQs_canDragDownFromAdapter() {
-        transitionController.qS = null
+    fun setDragAmount_notInSplitShade_forwardsToSingleShadeOverScroller() =
+        testScope.runTest {
+            disableSplitShade()
 
-        qsSceneAdapter.isQsFullyCollapsed = true
-        assertTrue("Can't drag down on keyguard", transitionController.canDragDown())
-        qsSceneAdapter.isQsFullyCollapsed = false
-        assertFalse("Can drag down when QS is expanded", transitionController.canDragDown())
-    }
+            transitionController.dragDownAmount = 10f
+
+            verify(singleShadeOverScroller).expansionDragDownAmount = 10f
+            verifyZeroInteractions(splitShadeOverScroller)
+        }
+
+    @Test
+    fun setDragAmount_inSplitShade_forwardsToSplitShadeOverScroller() =
+        testScope.runTest {
+            enableSplitShade()
+
+            transitionController.dragDownAmount = 10f
+
+            verify(splitShadeOverScroller).expansionDragDownAmount = 10f
+            verifyZeroInteractions(singleShadeOverScroller)
+        }
+
+    @Test
+    fun setDragDownAmount_inSplitShade_setsKeyguardStatusBarAlphaBasedOnDistance() =
+        testScope.runTest {
+            val alphaDistance =
+                context.resources.getDimensionPixelSize(
+                    R.dimen.lockscreen_shade_npvc_keyguard_content_alpha_transition_distance
+                )
+            val dragDownAmount = 10f
+            enableSplitShade()
+
+            transitionController.dragDownAmount = dragDownAmount
+
+            val expectedAlpha = 1 - dragDownAmount / alphaDistance
+            verify(shadeLockscreenInteractor).setKeyguardStatusBarAlpha(expectedAlpha)
+        }
+
+    @Test
+    fun setDragDownAmount_notInSplitShade_setsKeyguardStatusBarAlphaToMinusOne() =
+        testScope.runTest {
+            disableSplitShade()
+
+            transitionController.dragDownAmount = 10f
+
+            verify(shadeLockscreenInteractor).setKeyguardStatusBarAlpha(-1f)
+        }
+
+    @Test
+    fun nullQs_canDragDownFromAdapter() =
+        testScope.runTest {
+            transitionController.qS = null
+
+            qsSceneAdapter.isQsFullyCollapsed = true
+            assertTrue("Can't drag down on keyguard", transitionController.canDragDown())
+            qsSceneAdapter.isQsFullyCollapsed = false
+            assertFalse("Can drag down when QS is expanded", transitionController.canDragDown())
+        }
 
     private fun enableSplitShade() {
         setSplitShadeEnabled(true)
@@ -619,32 +631,4 @@
     ) {
         setTransitionToFullShadeProgress(progress, lockScreenNotificationsProgress)
     }
-
-    @SysUISingleton
-    @Component(
-        modules =
-            [
-                SysUITestModule::class,
-                UserDomainLayerModule::class,
-                BiometricsDomainLayerModule::class,
-            ]
-    )
-    interface TestComponent {
-
-        val configurationController: FakeConfigurationController
-        val disableFlagsRepository: FakeDisableFlagsRepository
-        val powerInteractor: PowerInteractor
-        val shadeInteractor: ShadeInteractor
-        val shadeRepository: FakeShadeRepository
-        val testScope: TestScope
-
-        @Component.Factory
-        interface Factory {
-            fun create(
-                    @BindsInstance test: SysuiTestCase,
-                    featureFlags: FakeFeatureFlagsClassicModule,
-                    mocks: TestMocksModule,
-            ): TestComponent
-        }
-    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelTest.kt
index fe29140..18a62cf 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelTest.kt
@@ -19,7 +19,6 @@
 import android.view.View
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.animation.mockDialogTransitionAnimator
 import com.android.systemui.common.shared.model.ContentDescription
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.coroutines.collectLastValue
@@ -38,7 +37,6 @@
 import com.android.systemui.statusbar.chips.mediaprojection.domain.interactor.MediaProjectionChipInteractorTest.Companion.setUpPackageManagerForMediaProjection
 import com.android.systemui.statusbar.chips.ui.model.ColorsModel
 import com.android.systemui.statusbar.chips.ui.model.OngoingActivityChipModel
-import com.android.systemui.statusbar.chips.ui.view.ChipBackgroundContainer
 import com.android.systemui.statusbar.phone.SystemUIDialog
 import com.android.systemui.statusbar.phone.mockSystemUIDialogFactory
 import com.android.systemui.statusbar.policy.CastDevice
@@ -47,9 +45,7 @@
 import kotlin.test.Test
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
-import org.mockito.ArgumentMatchers
 import org.mockito.kotlin.any
-import org.mockito.kotlin.eq
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.verify
 import org.mockito.kotlin.whenever
@@ -65,17 +61,6 @@
     private val mockScreenCastDialog = mock<SystemUIDialog>()
     private val mockGenericCastDialog = mock<SystemUIDialog>()
 
-    private val chipBackgroundView = mock<ChipBackgroundContainer>()
-    private val chipView =
-        mock<View>().apply {
-            whenever(
-                    this.requireViewById<ChipBackgroundContainer>(
-                        R.id.ongoing_activity_chip_background
-                    )
-                )
-                .thenReturn(chipBackgroundView)
-        }
-
     private val underTest = kosmos.castToOtherDeviceChipViewModel
 
     @Before
@@ -306,14 +291,8 @@
             val clickListener = ((latest as OngoingActivityChipModel.Shown).onClickListener)
             assertThat(clickListener).isNotNull()
 
-            clickListener!!.onClick(chipView)
-            verify(kosmos.mockDialogTransitionAnimator)
-                .showFromView(
-                    eq(mockScreenCastDialog),
-                    eq(chipBackgroundView),
-                    eq(null),
-                    ArgumentMatchers.anyBoolean(),
-                )
+            clickListener!!.onClick(mock<View>())
+            verify(mockScreenCastDialog).show()
         }
 
     @Test
@@ -330,14 +309,8 @@
             val clickListener = ((latest as OngoingActivityChipModel.Shown).onClickListener)
             assertThat(clickListener).isNotNull()
 
-            clickListener!!.onClick(chipView)
-            verify(kosmos.mockDialogTransitionAnimator)
-                .showFromView(
-                    eq(mockScreenCastDialog),
-                    eq(chipBackgroundView),
-                    eq(null),
-                    ArgumentMatchers.anyBoolean(),
-                )
+            clickListener!!.onClick(mock<View>())
+            verify(mockScreenCastDialog).show()
         }
 
     @Test
@@ -359,13 +332,7 @@
             val clickListener = ((latest as OngoingActivityChipModel.Shown).onClickListener)
             assertThat(clickListener).isNotNull()
 
-            clickListener!!.onClick(chipView)
-            verify(kosmos.mockDialogTransitionAnimator)
-                .showFromView(
-                    eq(mockGenericCastDialog),
-                    eq(chipBackgroundView),
-                    eq(null),
-                    ArgumentMatchers.anyBoolean(),
-                )
+            clickListener!!.onClick(mock<View>())
+            verify(mockGenericCastDialog).show()
         }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModelTest.kt
index 0a06cc7..90f94c8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModelTest.kt
@@ -19,7 +19,6 @@
 import android.view.View
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.animation.mockDialogTransitionAnimator
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.kosmos.Kosmos
@@ -34,7 +33,6 @@
 import com.android.systemui.statusbar.chips.screenrecord.ui.view.EndScreenRecordingDialogDelegate
 import com.android.systemui.statusbar.chips.ui.model.ColorsModel
 import com.android.systemui.statusbar.chips.ui.model.OngoingActivityChipModel
-import com.android.systemui.statusbar.chips.ui.view.ChipBackgroundContainer
 import com.android.systemui.statusbar.phone.SystemUIDialog
 import com.android.systemui.statusbar.phone.mockSystemUIDialogFactory
 import com.android.systemui.util.time.fakeSystemClock
@@ -42,9 +40,7 @@
 import kotlin.test.Test
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
-import org.mockito.ArgumentMatchers
 import org.mockito.kotlin.any
-import org.mockito.kotlin.eq
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.verify
 import org.mockito.kotlin.whenever
@@ -58,17 +54,6 @@
     private val systemClock = kosmos.fakeSystemClock
     private val mockSystemUIDialog = mock<SystemUIDialog>()
 
-    private val chipBackgroundView = mock<ChipBackgroundContainer>()
-    private val chipView =
-        mock<View>().apply {
-            whenever(
-                    this.requireViewById<ChipBackgroundContainer>(
-                        R.id.ongoing_activity_chip_background
-                    )
-                )
-                .thenReturn(chipBackgroundView)
-        }
-
     private val underTest = kosmos.screenRecordChipViewModel
 
     @Before
@@ -197,15 +182,9 @@
             val clickListener = ((latest as OngoingActivityChipModel.Shown).onClickListener)
             assertThat(clickListener).isNotNull()
 
-            clickListener!!.onClick(chipView)
+            clickListener!!.onClick(mock<View>())
             // EndScreenRecordingDialogDelegate will test that the dialog has the right message
-            verify(kosmos.mockDialogTransitionAnimator)
-                .showFromView(
-                    eq(mockSystemUIDialog),
-                    eq(chipBackgroundView),
-                    eq(null),
-                    ArgumentMatchers.anyBoolean(),
-                )
+            verify(mockSystemUIDialog).show()
         }
 
     @Test
@@ -219,15 +198,9 @@
             val clickListener = ((latest as OngoingActivityChipModel.Shown).onClickListener)
             assertThat(clickListener).isNotNull()
 
-            clickListener!!.onClick(chipView)
+            clickListener!!.onClick(mock<View>())
             // EndScreenRecordingDialogDelegate will test that the dialog has the right message
-            verify(kosmos.mockDialogTransitionAnimator)
-                .showFromView(
-                    eq(mockSystemUIDialog),
-                    eq(chipBackgroundView),
-                    eq(null),
-                    ArgumentMatchers.anyBoolean(),
-                )
+            verify(mockSystemUIDialog).show()
         }
 
     @Test
@@ -244,14 +217,8 @@
             val clickListener = ((latest as OngoingActivityChipModel.Shown).onClickListener)
             assertThat(clickListener).isNotNull()
 
-            clickListener!!.onClick(chipView)
+            clickListener!!.onClick(mock<View>())
             // EndScreenRecordingDialogDelegate will test that the dialog has the right message
-            verify(kosmos.mockDialogTransitionAnimator)
-                .showFromView(
-                    eq(mockSystemUIDialog),
-                    eq(chipBackgroundView),
-                    eq(null),
-                    ArgumentMatchers.anyBoolean(),
-                )
+            verify(mockSystemUIDialog).show()
         }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModelTest.kt
index 3028d00..29fd792 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModelTest.kt
@@ -19,7 +19,6 @@
 import android.view.View
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.animation.mockDialogTransitionAnimator
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.kosmos.Kosmos
@@ -35,7 +34,6 @@
 import com.android.systemui.statusbar.chips.sharetoapp.ui.view.EndShareToAppDialogDelegate
 import com.android.systemui.statusbar.chips.ui.model.ColorsModel
 import com.android.systemui.statusbar.chips.ui.model.OngoingActivityChipModel
-import com.android.systemui.statusbar.chips.ui.view.ChipBackgroundContainer
 import com.android.systemui.statusbar.phone.SystemUIDialog
 import com.android.systemui.statusbar.phone.mockSystemUIDialogFactory
 import com.android.systemui.util.time.fakeSystemClock
@@ -43,9 +41,7 @@
 import kotlin.test.Test
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
-import org.mockito.ArgumentMatchers
 import org.mockito.kotlin.any
-import org.mockito.kotlin.eq
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.verify
 import org.mockito.kotlin.whenever
@@ -59,17 +55,6 @@
 
     private val mockShareDialog = mock<SystemUIDialog>()
 
-    private val chipBackgroundView = mock<ChipBackgroundContainer>()
-    private val chipView =
-        mock<View>().apply {
-            whenever(
-                    this.requireViewById<ChipBackgroundContainer>(
-                        R.id.ongoing_activity_chip_background
-                    )
-                )
-                .thenReturn(chipBackgroundView)
-        }
-
     private val underTest = kosmos.shareToAppChipViewModel
 
     @Before
@@ -193,14 +178,8 @@
             val clickListener = ((latest as OngoingActivityChipModel.Shown).onClickListener)
             assertThat(clickListener).isNotNull()
 
-            clickListener!!.onClick(chipView)
-            verify(kosmos.mockDialogTransitionAnimator)
-                .showFromView(
-                    eq(mockShareDialog),
-                    eq(chipBackgroundView),
-                    eq(null),
-                    ArgumentMatchers.anyBoolean(),
-                )
+            clickListener!!.onClick(mock<View>())
+            verify(mockShareDialog).show()
         }
 
     @Test
@@ -216,13 +195,7 @@
             val clickListener = ((latest as OngoingActivityChipModel.Shown).onClickListener)
             assertThat(clickListener).isNotNull()
 
-            clickListener!!.onClick(chipView)
-            verify(kosmos.mockDialogTransitionAnimator)
-                .showFromView(
-                    eq(mockShareDialog),
-                    eq(chipBackgroundView),
-                    eq(null),
-                    ArgumentMatchers.anyBoolean(),
-                )
+            clickListener!!.onClick(mock<View>())
+            verify(mockShareDialog).show()
         }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipViewModelTest.kt
index c9c7359..2e0c773 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipViewModelTest.kt
@@ -19,50 +19,25 @@
 import android.view.View
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.animation.DialogTransitionAnimator
-import com.android.systemui.res.R
-import com.android.systemui.statusbar.chips.ui.view.ChipBackgroundContainer
 import com.android.systemui.statusbar.chips.ui.viewmodel.OngoingActivityChipViewModel.Companion.createDialogLaunchOnClickListener
 import com.android.systemui.statusbar.phone.SystemUIDialog
 import kotlin.test.Test
-import org.mockito.ArgumentMatchers.anyBoolean
-import org.mockito.kotlin.eq
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.verify
-import org.mockito.kotlin.whenever
 
 @SmallTest
 class OngoingActivityChipViewModelTest : SysuiTestCase() {
     private val mockSystemUIDialog = mock<SystemUIDialog>()
     private val dialogDelegate = SystemUIDialog.Delegate { mockSystemUIDialog }
-    private val dialogTransitionAnimator = mock<DialogTransitionAnimator>()
-
-    private val chipBackgroundView = mock<ChipBackgroundContainer>()
-    private val chipView =
-        mock<View>().apply {
-            whenever(
-                    this.requireViewById<ChipBackgroundContainer>(
-                        R.id.ongoing_activity_chip_background
-                    )
-                )
-                .thenReturn(chipBackgroundView)
-        }
 
     @Test
     fun createDialogLaunchOnClickListener_showsDialogOnClick() {
-        val clickListener =
-            createDialogLaunchOnClickListener(dialogDelegate, dialogTransitionAnimator)
+        val clickListener = createDialogLaunchOnClickListener(dialogDelegate)
 
         // Dialogs must be created on the main thread
         context.mainExecutor.execute {
-            clickListener.onClick(chipView)
-            verify(dialogTransitionAnimator)
-                .showFromView(
-                    eq(mockSystemUIDialog),
-                    eq(chipBackgroundView),
-                    eq(null),
-                    anyBoolean(),
-                )
+            clickListener.onClick(mock<View>())
+            verify(mockSystemUIDialog).show()
         }
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractorTest.kt
index 26f5370..f07303e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractorTest.kt
@@ -17,21 +17,19 @@
 
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
-import com.android.systemui.SysUITestComponent
-import com.android.systemui.SysUITestModule
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.TestMocksModule
-import com.android.systemui.biometrics.domain.BiometricsDomainLayerModule
-import com.android.systemui.collectLastValue
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.deviceentry.data.repository.FakeDeviceEntryRepository
-import com.android.systemui.runTest
-import com.android.systemui.statusbar.data.repository.NotificationListenerSettingsRepository
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.deviceentry.data.repository.fakeDeviceEntryRepository
+import com.android.systemui.deviceentry.domain.interactor.deviceEntryInteractor
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.statusbar.data.repository.notificationListenerSettingsRepository
 import com.android.systemui.statusbar.notification.data.model.activeNotificationModel
-import com.android.systemui.statusbar.notification.data.repository.ActiveNotificationListRepository
 import com.android.systemui.statusbar.notification.data.repository.ActiveNotificationsStore
-import com.android.systemui.statusbar.notification.domain.interactor.HeadsUpNotificationIconInteractor
-import com.android.systemui.statusbar.notification.domain.interactor.NotificationsKeyguardInteractor
+import com.android.systemui.statusbar.notification.data.repository.activeNotificationListRepository
+import com.android.systemui.statusbar.notification.data.repository.notificationsKeyguardViewStateRepository
+import com.android.systemui.statusbar.notification.domain.interactor.activeNotificationsInteractor
+import com.android.systemui.statusbar.notification.domain.interactor.headsUpNotificationIconInteractor
 import com.android.systemui.statusbar.notification.shared.byIsAmbient
 import com.android.systemui.statusbar.notification.shared.byIsLastMessageFromReply
 import com.android.systemui.statusbar.notification.shared.byIsPulsing
@@ -39,15 +37,15 @@
 import com.android.systemui.statusbar.notification.shared.byIsSilent
 import com.android.systemui.statusbar.notification.shared.byIsSuppressedFromStatusBar
 import com.android.systemui.statusbar.notification.shared.byKey
-import com.android.systemui.user.domain.UserDomainLayerModule
+import com.android.systemui.statusbar.notification.stack.domain.interactor.notificationsKeyguardInteractor
+import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.eq
-import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
-import com.android.wm.shell.bubbles.Bubbles
+import com.android.wm.shell.bubbles.bubbles
+import com.android.wm.shell.bubbles.bubblesOptional
 import com.google.common.truth.Truth.assertThat
-import dagger.BindsInstance
-import dagger.Component
-import java.util.Optional
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -55,29 +53,22 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class NotificationIconsInteractorTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+    private val activeNotificationListRepository = kosmos.activeNotificationListRepository
+    private val notificationsKeyguardInteractor = kosmos.notificationsKeyguardInteractor
 
-    private val bubbles: Bubbles = mock()
-
-    @Component(modules = [SysUITestModule::class])
-    @SysUISingleton
-    interface TestComponent : SysUITestComponent<NotificationIconsInteractor> {
-
-        val activeNotificationListRepository: ActiveNotificationListRepository
-        val notificationsKeyguardInteractor: NotificationsKeyguardInteractor
-
-        @Component.Factory
-        interface Factory {
-            fun create(@BindsInstance test: SysuiTestCase, mocks: TestMocksModule): TestComponent
-        }
-    }
-
-    val testComponent: TestComponent =
-        DaggerNotificationIconsInteractorTest_TestComponent.factory()
-            .create(test = this, mocks = TestMocksModule(bubbles = Optional.of(bubbles)))
+    private val underTest =
+        NotificationIconsInteractor(
+            kosmos.activeNotificationsInteractor,
+            kosmos.bubblesOptional,
+            kosmos.headsUpNotificationIconInteractor,
+            kosmos.notificationsKeyguardViewStateRepository
+        )
 
     @Before
     fun setup() {
-        testComponent.apply {
+        testScope.apply {
             activeNotificationListRepository.activeNotifications.value =
                 ActiveNotificationsStore.Builder()
                     .apply { testIcons.forEach(::addIndividualNotif) }
@@ -87,22 +78,22 @@
 
     @Test
     fun filteredEntrySet() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.filteredNotifSet())
             assertThat(filteredSet).containsExactlyElementsIn(testIcons)
         }
 
     @Test
     fun filteredEntrySet_noExpandedBubbles() =
-        testComponent.runTest {
-            whenever(bubbles.isBubbleExpanded(eq("notif1"))).thenReturn(true)
+        testScope.runTest {
+            whenever(kosmos.bubbles.isBubbleExpanded(eq("notif1"))).thenReturn(true)
             val filteredSet by collectLastValue(underTest.filteredNotifSet())
             assertThat(filteredSet).comparingElementsUsing(byKey).doesNotContain("notif1")
         }
 
     @Test
     fun filteredEntrySet_noAmbient() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.filteredNotifSet(showAmbient = false))
             assertThat(filteredSet).comparingElementsUsing(byIsAmbient).doesNotContain(true)
             assertThat(filteredSet)
@@ -112,21 +103,21 @@
 
     @Test
     fun filteredEntrySet_noLowPriority() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.filteredNotifSet(showLowPriority = false))
             assertThat(filteredSet).comparingElementsUsing(byIsSilent).doesNotContain(true)
         }
 
     @Test
     fun filteredEntrySet_noDismissed() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.filteredNotifSet(showDismissed = false))
             assertThat(filteredSet).comparingElementsUsing(byIsRowDismissed).doesNotContain(true)
         }
 
     @Test
     fun filteredEntrySet_noRepliedMessages() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by
                 collectLastValue(underTest.filteredNotifSet(showRepliedMessages = false))
             assertThat(filteredSet)
@@ -136,7 +127,7 @@
 
     @Test
     fun filteredEntrySet_noPulsing_notifsNotFullyHidden() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.filteredNotifSet(showPulsing = false))
             notificationsKeyguardInteractor.setNotificationsFullyHidden(false)
             assertThat(filteredSet).comparingElementsUsing(byIsPulsing).doesNotContain(true)
@@ -144,65 +135,46 @@
 
     @Test
     fun filteredEntrySet_noPulsing_notifsFullyHidden() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.filteredNotifSet(showPulsing = false))
             notificationsKeyguardInteractor.setNotificationsFullyHidden(true)
             assertThat(filteredSet).comparingElementsUsing(byIsPulsing).contains(true)
         }
 }
 
+@OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class AlwaysOnDisplayNotificationIconsInteractorTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
 
-    private val bubbles: Bubbles = mock()
-
-    @Component(
-        modules =
-            [
-                SysUITestModule::class,
-                BiometricsDomainLayerModule::class,
-                UserDomainLayerModule::class,
-            ]
-    )
-    @SysUISingleton
-    interface TestComponent : SysUITestComponent<AlwaysOnDisplayNotificationIconsInteractor> {
-
-        val activeNotificationListRepository: ActiveNotificationListRepository
-        val deviceEntryRepository: FakeDeviceEntryRepository
-        val notificationsKeyguardInteractor: NotificationsKeyguardInteractor
-
-        @Component.Factory
-        interface Factory {
-            fun create(@BindsInstance test: SysuiTestCase, mocks: TestMocksModule): TestComponent
-        }
-    }
-
-    private val testComponent: TestComponent =
-        DaggerAlwaysOnDisplayNotificationIconsInteractorTest_TestComponent.factory()
-            .create(test = this, mocks = TestMocksModule(bubbles = Optional.of(bubbles)))
+    private val underTest =
+        AlwaysOnDisplayNotificationIconsInteractor(
+            kosmos.testDispatcher,
+            kosmos.deviceEntryInteractor,
+            kosmos.notificationIconsInteractor,
+        )
 
     @Before
     fun setup() {
-        testComponent.apply {
-            activeNotificationListRepository.activeNotifications.value =
-                ActiveNotificationsStore.Builder()
-                    .apply { testIcons.forEach(::addIndividualNotif) }
-                    .build()
-        }
+        kosmos.activeNotificationListRepository.activeNotifications.value =
+            ActiveNotificationsStore.Builder()
+                .apply { testIcons.forEach(::addIndividualNotif) }
+                .build()
     }
 
     @Test
     fun filteredEntrySet_noExpandedBubbles() =
-        testComponent.runTest {
-            whenever(bubbles.isBubbleExpanded(eq("notif1"))).thenReturn(true)
+        testScope.runTest {
+            whenever(kosmos.bubbles.isBubbleExpanded(eq("notif1"))).thenReturn(true)
             val filteredSet by collectLastValue(underTest.aodNotifs)
             assertThat(filteredSet).comparingElementsUsing(byKey).doesNotContain("notif1")
         }
 
     @Test
     fun filteredEntrySet_noAmbient() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.aodNotifs)
             assertThat(filteredSet).comparingElementsUsing(byIsAmbient).doesNotContain(true)
             assertThat(filteredSet)
@@ -212,14 +184,14 @@
 
     @Test
     fun filteredEntrySet_noDismissed() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.aodNotifs)
             assertThat(filteredSet).comparingElementsUsing(byIsRowDismissed).doesNotContain(true)
         }
 
     @Test
     fun filteredEntrySet_noRepliedMessages() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.aodNotifs)
             assertThat(filteredSet)
                 .comparingElementsUsing(byIsLastMessageFromReply)
@@ -228,37 +200,37 @@
 
     @Test
     fun filteredEntrySet_showPulsing_notifsNotFullyHidden_bypassDisabled() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.aodNotifs)
-            deviceEntryRepository.setBypassEnabled(false)
-            notificationsKeyguardInteractor.setNotificationsFullyHidden(false)
+            kosmos.fakeDeviceEntryRepository.setBypassEnabled(false)
+            kosmos.notificationsKeyguardInteractor.setNotificationsFullyHidden(false)
             assertThat(filteredSet).comparingElementsUsing(byIsPulsing).contains(true)
         }
 
     @Test
     fun filteredEntrySet_showPulsing_notifsFullyHidden_bypassDisabled() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.aodNotifs)
-            deviceEntryRepository.setBypassEnabled(false)
-            notificationsKeyguardInteractor.setNotificationsFullyHidden(true)
+            kosmos.fakeDeviceEntryRepository.setBypassEnabled(false)
+            kosmos.notificationsKeyguardInteractor.setNotificationsFullyHidden(true)
             assertThat(filteredSet).comparingElementsUsing(byIsPulsing).contains(true)
         }
 
     @Test
     fun filteredEntrySet_noPulsing_notifsNotFullyHidden_bypassEnabled() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.aodNotifs)
-            deviceEntryRepository.setBypassEnabled(true)
-            notificationsKeyguardInteractor.setNotificationsFullyHidden(false)
+            kosmos.fakeDeviceEntryRepository.setBypassEnabled(true)
+            kosmos.notificationsKeyguardInteractor.setNotificationsFullyHidden(false)
             assertThat(filteredSet).comparingElementsUsing(byIsPulsing).doesNotContain(true)
         }
 
     @Test
     fun filteredEntrySet_showPulsing_notifsFullyHidden_bypassEnabled() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.aodNotifs)
-            deviceEntryRepository.setBypassEnabled(true)
-            notificationsKeyguardInteractor.setNotificationsFullyHidden(true)
+            kosmos.fakeDeviceEntryRepository.setBypassEnabled(true)
+            kosmos.notificationsKeyguardInteractor.setNotificationsFullyHidden(true)
             assertThat(filteredSet).comparingElementsUsing(byIsPulsing).contains(true)
         }
 }
@@ -266,32 +238,19 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class StatusBarNotificationIconsInteractorTest : SysuiTestCase() {
-
-    private val bubbles: Bubbles = mock()
-
-    @Component(modules = [SysUITestModule::class])
-    @SysUISingleton
-    interface TestComponent : SysUITestComponent<StatusBarNotificationIconsInteractor> {
-
-        val activeNotificationListRepository: ActiveNotificationListRepository
-        val headsUpIconsInteractor: HeadsUpNotificationIconInteractor
-        val notificationsKeyguardInteractor: NotificationsKeyguardInteractor
-        val notificationListenerSettingsRepository: NotificationListenerSettingsRepository
-
-        @Component.Factory
-        interface Factory {
-            fun create(@BindsInstance test: SysuiTestCase, mocks: TestMocksModule): TestComponent
-        }
-    }
-
-    val testComponent: TestComponent =
-        DaggerStatusBarNotificationIconsInteractorTest_TestComponent.factory()
-            .create(test = this, mocks = TestMocksModule(bubbles = Optional.of(bubbles)))
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+    private val underTest =
+        StatusBarNotificationIconsInteractor(
+            kosmos.testDispatcher,
+            kosmos.notificationIconsInteractor,
+            kosmos.notificationListenerSettingsRepository,
+        )
 
     @Before
     fun setup() {
-        testComponent.apply {
-            activeNotificationListRepository.activeNotifications.value =
+        testScope.apply {
+            kosmos.activeNotificationListRepository.activeNotifications.value =
                 ActiveNotificationsStore.Builder()
                     .apply { testIcons.forEach(::addIndividualNotif) }
                     .build()
@@ -300,15 +259,15 @@
 
     @Test
     fun filteredEntrySet_noExpandedBubbles() =
-        testComponent.runTest {
-            whenever(bubbles.isBubbleExpanded(eq("notif1"))).thenReturn(true)
+        testScope.runTest {
+            whenever(kosmos.bubbles.isBubbleExpanded(eq("notif1"))).thenReturn(true)
             val filteredSet by collectLastValue(underTest.statusBarNotifs)
             assertThat(filteredSet).comparingElementsUsing(byKey).doesNotContain("notif1")
         }
 
     @Test
     fun filteredEntrySet_noAmbient() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.statusBarNotifs)
             assertThat(filteredSet).comparingElementsUsing(byIsAmbient).doesNotContain(true)
             assertThat(filteredSet)
@@ -318,30 +277,30 @@
 
     @Test
     fun filteredEntrySet_noLowPriority_whenDontShowSilentIcons() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.statusBarNotifs)
-            notificationListenerSettingsRepository.showSilentStatusIcons.value = false
+            kosmos.notificationListenerSettingsRepository.showSilentStatusIcons.value = false
             assertThat(filteredSet).comparingElementsUsing(byIsSilent).doesNotContain(true)
         }
 
     @Test
     fun filteredEntrySet_showLowPriority_whenShowSilentIcons() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.statusBarNotifs)
-            notificationListenerSettingsRepository.showSilentStatusIcons.value = true
+            kosmos.notificationListenerSettingsRepository.showSilentStatusIcons.value = true
             assertThat(filteredSet).comparingElementsUsing(byIsSilent).contains(true)
         }
 
     @Test
     fun filteredEntrySet_noDismissed() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.statusBarNotifs)
             assertThat(filteredSet).comparingElementsUsing(byIsRowDismissed).doesNotContain(true)
         }
 
     @Test
     fun filteredEntrySet_noRepliedMessages() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.statusBarNotifs)
             assertThat(filteredSet)
                 .comparingElementsUsing(byIsLastMessageFromReply)
@@ -350,9 +309,9 @@
 
     @Test
     fun filteredEntrySet_includesIsolatedIcon() =
-        testComponent.runTest {
+        testScope.runTest {
             val filteredSet by collectLastValue(underTest.statusBarNotifs)
-            headsUpIconsInteractor.setIsolatedIconNotificationKey("notif5")
+            kosmos.headsUpNotificationIconInteractor.setIsolatedIconNotificationKey("notif5")
             assertThat(filteredSet).comparingElementsUsing(byKey).contains("notif5")
         }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerAlwaysOnDisplayViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerAlwaysOnDisplayViewModelTest.kt
index 894e02e..1f4e80e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerAlwaysOnDisplayViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerAlwaysOnDisplayViewModelTest.kt
@@ -16,111 +16,81 @@
 
 package com.android.systemui.statusbar.notification.icon.ui.viewmodel
 
+import android.content.res.mainResources
 import android.platform.test.annotations.DisableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR
 import com.android.systemui.Flags.FLAG_NEW_AOD_TRANSITION
-import com.android.systemui.SysUITestComponent
-import com.android.systemui.SysUITestModule
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.TestMocksModule
-import com.android.systemui.biometrics.domain.BiometricsDomainLayerModule
-import com.android.systemui.collectLastValue
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.flags.FakeFeatureFlagsClassicModule
+import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.flags.Flags
-import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository
-import com.android.systemui.keyguard.data.repository.FakeKeyguardTransitionRepository
+import com.android.systemui.flags.fakeFeatureFlagsClassic
+import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
+import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
+import com.android.systemui.keyguard.domain.interactor.keyguardInteractor
+import com.android.systemui.keyguard.domain.interactor.keyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.DozeStateModel
 import com.android.systemui.keyguard.shared.model.DozeTransitionModel
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.TransitionStep
-import com.android.systemui.power.data.repository.FakePowerRepository
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.power.data.repository.fakePowerRepository
 import com.android.systemui.power.shared.model.WakeSleepReason
 import com.android.systemui.power.shared.model.WakefulnessState
-import com.android.systemui.runCurrent
-import com.android.systemui.runTest
-import com.android.systemui.statusbar.phone.DozeParameters
-import com.android.systemui.statusbar.phone.ScreenOffAnimationController
-import com.android.systemui.statusbar.policy.data.repository.FakeDeviceProvisioningRepository
-import com.android.systemui.user.domain.UserDomainLayerModule
-import com.android.systemui.util.mockito.mock
+import com.android.systemui.shade.domain.interactor.shadeInteractor
+import com.android.systemui.statusbar.notification.icon.domain.interactor.alwaysOnDisplayNotificationIconsInteractor
+import com.android.systemui.statusbar.phone.dozeParameters
+import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
-import dagger.BindsInstance
-import dagger.Component
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 
+@OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class NotificationIconContainerAlwaysOnDisplayViewModelTest : SysuiTestCase() {
-
-    @SysUISingleton
-    @Component(
-        modules =
-            [
-                SysUITestModule::class,
-                BiometricsDomainLayerModule::class,
-                UserDomainLayerModule::class,
-            ]
-    )
-    interface TestComponent :
-        SysUITestComponent<NotificationIconContainerAlwaysOnDisplayViewModel> {
-
-        val deviceProvisioningRepository: FakeDeviceProvisioningRepository
-        val keyguardRepository: FakeKeyguardRepository
-        val keyguardTransitionRepository: FakeKeyguardTransitionRepository
-        val powerRepository: FakePowerRepository
-
-        @Component.Factory
-        interface Factory {
-            fun create(
-                @BindsInstance test: SysuiTestCase,
-                mocks: TestMocksModule,
-                featureFlags: FakeFeatureFlagsClassicModule,
-            ): TestComponent
+    private val kosmos =
+        testKosmos().apply {
+            fakeFeatureFlagsClassic.apply { set(Flags.FULL_SCREEN_USER_SWITCHER, value = false) }
         }
-    }
 
-    private val dozeParams: DozeParameters = mock()
-    private val screenOffAnimController: ScreenOffAnimationController = mock()
-
-    private val testComponent: TestComponent =
-        DaggerNotificationIconContainerAlwaysOnDisplayViewModelTest_TestComponent.factory()
-            .create(
-                test = this,
-                featureFlags =
-                    FakeFeatureFlagsClassicModule {
-                        set(Flags.FULL_SCREEN_USER_SWITCHER, value = false)
-                    },
-                mocks =
-                    TestMocksModule(
-                        dozeParameters = dozeParams,
-                        screenOffAnimationController = screenOffAnimController,
-                    ),
-            )
+    val underTest =
+        NotificationIconContainerAlwaysOnDisplayViewModel(
+            kosmos.testDispatcher,
+            kosmos.alwaysOnDisplayNotificationIconsInteractor,
+            kosmos.keyguardInteractor,
+            kosmos.keyguardTransitionInteractor,
+            kosmos.mainResources,
+            kosmos.shadeInteractor,
+        )
+    val testScope = kosmos.testScope
+    val keyguardRepository = kosmos.fakeKeyguardRepository
+    val keyguardTransitionRepository = kosmos.fakeKeyguardTransitionRepository
+    val powerRepository = kosmos.fakePowerRepository
 
     @Before
     fun setup() {
-        testComponent.apply {
-            keyguardRepository.setKeyguardShowing(true)
-            keyguardRepository.setKeyguardOccluded(false)
-            powerRepository.updateWakefulness(
-                rawState = WakefulnessState.AWAKE,
-                lastWakeReason = WakeSleepReason.OTHER,
-                lastSleepReason = WakeSleepReason.OTHER,
-            )
-        }
+        keyguardRepository.setKeyguardShowing(true)
+        keyguardRepository.setKeyguardOccluded(false)
+        kosmos.fakePowerRepository.updateWakefulness(
+            rawState = WakefulnessState.AWAKE,
+            lastWakeReason = WakeSleepReason.OTHER,
+            lastSleepReason = WakeSleepReason.OTHER,
+        )
         mSetFlagsRule.enableFlags(FLAG_NEW_AOD_TRANSITION)
     }
 
     @Test
     fun animationsEnabled_isFalse_whenDeviceAsleepAndNotPulsing() =
-        testComponent.runTest {
+        testScope.runTest {
             powerRepository.updateWakefulness(
                 rawState = WakefulnessState.ASLEEP,
                 lastWakeReason = WakeSleepReason.POWER_BUTTON,
@@ -143,7 +113,7 @@
 
     @Test
     fun animationsEnabled_isTrue_whenDeviceAsleepAndPulsing() =
-        testComponent.runTest {
+        testScope.runTest {
             powerRepository.updateWakefulness(
                 rawState = WakefulnessState.ASLEEP,
                 lastWakeReason = WakeSleepReason.POWER_BUTTON,
@@ -166,7 +136,7 @@
 
     @Test
     fun animationsEnabled_isFalse_whenStartingToSleepAndNotControlScreenOff() =
-        testComponent.runTest {
+        testScope.runTest {
             powerRepository.updateWakefulness(
                 rawState = WakefulnessState.STARTING_TO_SLEEP,
                 lastWakeReason = WakeSleepReason.POWER_BUTTON,
@@ -179,7 +149,7 @@
                     transitionState = TransitionState.STARTED,
                 )
             )
-            whenever(dozeParams.shouldControlScreenOff()).thenReturn(false)
+            whenever(kosmos.dozeParameters.shouldControlScreenOff()).thenReturn(false)
             val animationsEnabled by collectLastValue(underTest.areContainerChangesAnimated)
             runCurrent()
             assertThat(animationsEnabled).isFalse()
@@ -187,7 +157,7 @@
 
     @Test
     fun animationsEnabled_isTrue_whenStartingToSleepAndControlScreenOff() =
-        testComponent.runTest {
+        testScope.runTest {
             val animationsEnabled by collectLastValue(underTest.areContainerChangesAnimated)
             assertThat(animationsEnabled).isTrue()
 
@@ -203,13 +173,13 @@
                     transitionState = TransitionState.STARTED,
                 )
             )
-            whenever(dozeParams.shouldControlScreenOff()).thenReturn(true)
+            whenever(kosmos.dozeParameters.shouldControlScreenOff()).thenReturn(true)
             assertThat(animationsEnabled).isTrue()
         }
 
     @Test
     fun animationsEnabled_isTrue_whenNotAsleep() =
-        testComponent.runTest {
+        testScope.runTest {
             powerRepository.updateWakefulness(
                 rawState = WakefulnessState.AWAKE,
                 lastWakeReason = WakeSleepReason.POWER_BUTTON,
@@ -228,7 +198,7 @@
     @Test
     @DisableFlags(FLAG_KEYGUARD_WM_STATE_REFACTOR)
     fun animationsEnabled_isTrue_whenKeyguardIsShowing() =
-        testComponent.runTest {
+        testScope.runTest {
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
                     transitionState = TransitionState.STARTED,
@@ -257,7 +227,7 @@
 
     @Test
     fun tintAlpha_isZero_whenNotOnAodOrDozing() =
-        testComponent.runTest {
+        testScope.runTest {
             val tintAlpha by collectLastValue(underTest.tintAlpha)
             runCurrent()
             keyguardTransitionRepository.sendTransitionSteps(
@@ -271,7 +241,7 @@
 
     @Test
     fun tintAlpha_isOne_whenOnAod() =
-        testComponent.runTest {
+        testScope.runTest {
             val tintAlpha by collectLastValue(underTest.tintAlpha)
             runCurrent()
             keyguardTransitionRepository.sendTransitionSteps(
@@ -285,7 +255,7 @@
 
     @Test
     fun tintAlpha_isOne_whenDozing() =
-        testComponent.runTest {
+        testScope.runTest {
             val tintAlpha by collectLastValue(underTest.tintAlpha)
             runCurrent()
             keyguardTransitionRepository.sendTransitionSteps(
@@ -298,7 +268,7 @@
 
     @Test
     fun tintAlpha_isOne_whenTransitionFromAodToDoze() =
-        testComponent.runTest {
+        testScope.runTest {
             keyguardTransitionRepository.sendTransitionSteps(
                 from = KeyguardState.GONE,
                 to = KeyguardState.AOD,
@@ -332,7 +302,7 @@
 
     @Test
     fun tintAlpha_isFraction_midTransitionToAod() =
-        testComponent.runTest {
+        testScope.runTest {
             val tintAlpha by collectLastValue(underTest.tintAlpha)
             runCurrent()
 
@@ -361,7 +331,7 @@
 
     @Test
     fun iconAnimationsEnabled_whenOnLockScreen() =
-        testComponent.runTest {
+        testScope.runTest {
             val iconAnimationsEnabled by collectLastValue(underTest.areIconAnimationsEnabled)
             runCurrent()
 
@@ -376,7 +346,7 @@
 
     @Test
     fun iconAnimationsDisabled_whenOnAod() =
-        testComponent.runTest {
+        testScope.runTest {
             val iconAnimationsEnabled by collectLastValue(underTest.areIconAnimationsEnabled)
             runCurrent()
 
@@ -391,7 +361,7 @@
 
     @Test
     fun iconAnimationsDisabled_whenDozing() =
-        testComponent.runTest {
+        testScope.runTest {
             val iconAnimationsEnabled by collectLastValue(underTest.areIconAnimationsEnabled)
             runCurrent()
 
diff --git a/packages/SystemUI/tests/utils/src/android/app/admin/AlarmManagerKosmos.kt b/packages/SystemUI/tests/utils/src/android/app/admin/AlarmManagerKosmos.kt
new file mode 100644
index 0000000..a7b5873
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/android/app/admin/AlarmManagerKosmos.kt
@@ -0,0 +1,23 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.app.admin
+
+import android.app.AlarmManager
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.util.mockito.mock
+
+var Kosmos.alarmManager by Kosmos.Fixture { mock<AlarmManager>() }
diff --git a/packages/SystemUI/tests/utils/src/com/android/app/admin/DevicePolicyManagerKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/app/admin/DevicePolicyManagerKosmos.kt
new file mode 100644
index 0000000..f51e122
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/app/admin/DevicePolicyManagerKosmos.kt
@@ -0,0 +1,22 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.app.admin
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.util.mockito.mock
+
+val Kosmos.devicePolicyManager by Kosmos.Fixture { mock<android.app.admin.DevicePolicyManager>() }
diff --git a/packages/SystemUI/tests/utils/src/com/android/internal/widget/LockPatternUtilsKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/internal/widget/LockPatternUtilsKosmos.kt
index d9ea5e9..b511270 100644
--- a/packages/SystemUI/tests/utils/src/com/android/internal/widget/LockPatternUtilsKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/internal/widget/LockPatternUtilsKosmos.kt
@@ -16,7 +16,14 @@
 
 package com.android.internal.widget
 
+import android.app.admin.devicePolicyManager
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.whenever
 
-var Kosmos.lockPatternUtils by Kosmos.Fixture { mock<LockPatternUtils>() }
+var Kosmos.lockPatternUtils by
+    Kosmos.Fixture {
+        mock<LockPatternUtils>().apply {
+            whenever(this.devicePolicyManager).thenReturn(this@Fixture.devicePolicyManager)
+        }
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt
index 5bae6ec..87143ef 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt
@@ -135,6 +135,9 @@
 
     private var isShowKeyguardWhenReenabled: Boolean = false
 
+    private val _canIgnoreAuthAndReturnToGone = MutableStateFlow(false)
+    override val canIgnoreAuthAndReturnToGone = _canIgnoreAuthAndReturnToGone.asStateFlow()
+
     override fun setQuickSettingsVisible(isVisible: Boolean) {
         _isQuickSettingsVisible.value = isVisible
     }
@@ -278,6 +281,10 @@
     override fun isShowKeyguardWhenReenabled(): Boolean {
         return isShowKeyguardWhenReenabled
     }
+
+    override fun setCanIgnoreAuthAndReturnToGone(canWake: Boolean) {
+        _canIgnoreAuthAndReturnToGone.value = canWake
+    }
 }
 
 @Module
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractorKosmos.kt
index ae138c8..ef789d1 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractorKosmos.kt
@@ -37,5 +37,6 @@
             powerInteractor = powerInteractor,
             keyguardOcclusionInteractor = keyguardOcclusionInteractor,
             deviceEntryRepository = deviceEntryRepository,
+            wakeToGoneInteractor = keyguardWakeDirectlyToGoneInteractor,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/FromDozingTransitionInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/FromDozingTransitionInteractorKosmos.kt
index e7e007f..446652c 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/FromDozingTransitionInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/FromDozingTransitionInteractorKosmos.kt
@@ -39,5 +39,6 @@
             powerInteractor = powerInteractor,
             keyguardOcclusionInteractor = keyguardOcclusionInteractor,
             deviceEntryRepository = deviceEntryRepository,
+            wakeToGoneInteractor = keyguardWakeDirectlyToGoneInteractor,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractorKosmos.kt
index a9be06d..6c3de44 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractorKosmos.kt
@@ -16,13 +16,16 @@
 
 package com.android.systemui.keyguard.domain.interactor
 
+import com.android.systemui.deviceentry.domain.interactor.deviceEntryInteractor
 import com.android.systemui.keyguard.data.repository.keyguardTransitionRepository
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.applicationCoroutineScope
 import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.power.domain.interactor.powerInteractor
 import com.android.systemui.statusbar.domain.interactor.keyguardOcclusionInteractor
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 
+@OptIn(ExperimentalCoroutinesApi::class)
 var Kosmos.fromDreamingTransitionInteractor by
     Kosmos.Fixture {
         FromDreamingTransitionInteractor(
@@ -36,5 +39,6 @@
             glanceableHubTransitions = glanceableHubTransitions,
             powerInteractor = powerInteractor,
             keyguardOcclusionInteractor = keyguardOcclusionInteractor,
+            deviceEntryInteractor = deviceEntryInteractor,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardWakeDirectlyToGoneInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardWakeDirectlyToGoneInteractorKosmos.kt
new file mode 100644
index 0000000..63e168d
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardWakeDirectlyToGoneInteractorKosmos.kt
@@ -0,0 +1,45 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.keyguard.domain.interactor
+
+import android.app.admin.alarmManager
+import android.content.mockedContext
+import com.android.internal.widget.lockPatternUtils
+import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.power.domain.interactor.powerInteractor
+import com.android.systemui.user.domain.interactor.selectedUserInteractor
+import com.android.systemui.util.settings.fakeSettings
+import com.android.systemui.util.time.systemClock
+
+val Kosmos.keyguardWakeDirectlyToGoneInteractor by
+    Kosmos.Fixture {
+        KeyguardWakeDirectlyToGoneInteractor(
+            applicationCoroutineScope,
+            mockedContext,
+            fakeKeyguardRepository,
+            systemClock,
+            alarmManager,
+            keyguardTransitionInteractor,
+            powerInteractor,
+            fakeSettings,
+            lockPatternUtils,
+            fakeSettings,
+            selectedUserInteractor,
+        )
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractorKosmos.kt
index bd9c0be..8bb2fce 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractorKosmos.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.keyguard.domain.interactor
 
 import com.android.systemui.deviceentry.domain.interactor.deviceEntryInteractor
+import com.android.systemui.keyguard.data.repository.keyguardTransitionRepository
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.scene.domain.interactor.sceneInteractor
 import com.android.systemui.statusbar.notification.domain.interactor.notificationLaunchAnimationInteractor
@@ -25,6 +26,7 @@
     Kosmos.Fixture {
         WindowManagerLockscreenVisibilityInteractor(
             keyguardInteractor = keyguardInteractor,
+            transitionRepository = keyguardTransitionRepository,
             transitionInteractor = keyguardTransitionInteractor,
             surfaceBehindInteractor = keyguardSurfaceBehindInteractor,
             fromLockscreenInteractor = fromLockscreenTransitionInteractor,
@@ -33,5 +35,6 @@
             notificationLaunchAnimationInteractor = notificationLaunchAnimationInteractor,
             sceneInteractor = { sceneInteractor },
             deviceEntryInteractor = { deviceEntryInteractor },
+            wakeToGoneInteractor = keyguardWakeDirectlyToGoneInteractor,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelKosmos.kt
index 144fe26..2335f21 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelKosmos.kt
@@ -17,7 +17,6 @@
 package com.android.systemui.statusbar.chips.casttootherdevice.ui.viewmodel
 
 import android.content.applicationContext
-import com.android.systemui.animation.mockDialogTransitionAnimator
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.applicationCoroutineScope
 import com.android.systemui.statusbar.chips.casttootherdevice.domain.interactor.mediaRouterChipInteractor
@@ -34,6 +33,5 @@
             mediaRouterChipInteractor = mediaRouterChipInteractor,
             systemClock = fakeSystemClock,
             endMediaProjectionDialogHelper = endMediaProjectionDialogHelper,
-            dialogTransitionAnimator = mockDialogTransitionAnimator,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModelKosmos.kt
index 1d06947..2773f82 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModelKosmos.kt
@@ -17,7 +17,6 @@
 package com.android.systemui.statusbar.chips.screenrecord.ui.viewmodel
 
 import android.content.applicationContext
-import com.android.systemui.animation.mockDialogTransitionAnimator
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.applicationCoroutineScope
 import com.android.systemui.statusbar.chips.mediaprojection.ui.view.endMediaProjectionDialogHelper
@@ -31,7 +30,6 @@
             context = applicationContext,
             interactor = screenRecordChipInteractor,
             endMediaProjectionDialogHelper = endMediaProjectionDialogHelper,
-            dialogTransitionAnimator = mockDialogTransitionAnimator,
             systemClock = fakeSystemClock,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModelKosmos.kt
index 2e475a3..1b3108c 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModelKosmos.kt
@@ -17,7 +17,6 @@
 package com.android.systemui.statusbar.chips.sharetoapp.ui.viewmodel
 
 import android.content.applicationContext
-import com.android.systemui.animation.mockDialogTransitionAnimator
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.applicationCoroutineScope
 import com.android.systemui.statusbar.chips.mediaprojection.domain.interactor.mediaProjectionChipInteractor
@@ -32,6 +31,5 @@
             mediaProjectionChipInteractor = mediaProjectionChipInteractor,
             systemClock = fakeSystemClock,
             endMediaProjectionDialogHelper = endMediaProjectionDialogHelper,
-            dialogTransitionAnimator = mockDialogTransitionAnimator,
         )
     }
diff --git a/services/core/java/com/android/server/Watchdog.java b/services/core/java/com/android/server/Watchdog.java
index 21947ba..95dbaae 100644
--- a/services/core/java/com/android/server/Watchdog.java
+++ b/services/core/java/com/android/server/Watchdog.java
@@ -1016,6 +1016,7 @@
             // Trigger the kernel to dump all blocked threads, and backtraces on all CPUs to the
             // kernel log
             doSysRq('w');
+            doSysRq('m');
             doSysRq('l');
         }
 
diff --git a/services/core/java/com/android/server/media/projection/MediaProjectionManagerService.java b/services/core/java/com/android/server/media/projection/MediaProjectionManagerService.java
index e3d5c54..803b125 100644
--- a/services/core/java/com/android/server/media/projection/MediaProjectionManagerService.java
+++ b/services/core/java/com/android/server/media/projection/MediaProjectionManagerService.java
@@ -17,6 +17,7 @@
 package com.android.server.media.projection;
 
 import static android.Manifest.permission.MANAGE_MEDIA_PROJECTION;
+import static android.Manifest.permission.RECORD_SENSITIVE_CONTENT;
 import static android.app.ActivityManagerInternal.MEDIA_PROJECTION_TOKEN_EVENT_CREATED;
 import static android.app.ActivityManagerInternal.MEDIA_PROJECTION_TOKEN_EVENT_DESTROYED;
 import static android.content.Intent.FLAG_ACTIVITY_EXCLUDE_FROM_RECENTS;
@@ -34,10 +35,12 @@
 import android.annotation.EnforcePermission;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.annotation.RequiresPermission;
 import android.app.ActivityManagerInternal;
 import android.app.ActivityOptions.LaunchCookie;
 import android.app.AppOpsManager;
 import android.app.IProcessObserver;
+import android.app.KeyguardManager;
 import android.app.compat.CompatChanges;
 import android.compat.annotation.ChangeId;
 import android.compat.annotation.EnabledSince;
@@ -78,6 +81,7 @@
 import com.android.internal.util.ArrayUtils;
 import com.android.internal.util.DumpUtils;
 import com.android.server.LocalServices;
+import com.android.server.SystemConfig;
 import com.android.server.SystemService;
 import com.android.server.Watchdog;
 import com.android.server.wm.WindowManagerInternal;
@@ -132,6 +136,7 @@
     private final ActivityManagerInternal mActivityManagerInternal;
     private final PackageManager mPackageManager;
     private final WindowManagerInternal mWmInternal;
+    private final KeyguardManager mKeyguardManager;
 
     private final MediaRouter mMediaRouter;
     private final MediaRouterCallback mMediaRouterCallback;
@@ -147,7 +152,9 @@
         this(context, new Injector());
     }
 
-    @VisibleForTesting MediaProjectionManagerService(Context context, Injector injector) {
+    @RequiresPermission(Manifest.permission.SUBSCRIBE_TO_KEYGUARD_LOCKED_STATE)
+    @VisibleForTesting
+    MediaProjectionManagerService(Context context, Injector injector) {
         super(context);
         mContext = context;
         mInjector = injector;
@@ -163,9 +170,47 @@
         mMediaRouter = (MediaRouter) mContext.getSystemService(Context.MEDIA_ROUTER_SERVICE);
         mMediaRouterCallback = new MediaRouterCallback();
         mMediaProjectionMetricsLogger = injector.mediaProjectionMetricsLogger(context);
+        mKeyguardManager = (KeyguardManager) mContext.getSystemService(Context.KEYGUARD_SERVICE);
+        mKeyguardManager.addKeyguardLockedStateListener(
+                mContext.getMainExecutor(), this::onKeyguardLockedStateChanged);
         Watchdog.getInstance().addMonitor(this);
     }
 
+    /**
+     * In order to record the keyguard, the MediaProjection package must be either:
+     *   - a holder of RECORD_SENSITIVE_CONTENT permission, or
+     *   - be one of the bugreport whitelisted packages
+     */
+    private boolean canCaptureKeyguard() {
+        if (!android.companion.virtualdevice.flags.Flags.mediaProjectionKeyguardRestrictions()) {
+            return true;
+        }
+        synchronized (mLock) {
+            if (mProjectionGrant == null || mProjectionGrant.packageName == null) {
+                return false;
+            }
+            if (mPackageManager.checkPermission(RECORD_SENSITIVE_CONTENT,
+                    mProjectionGrant.packageName)
+                    == PackageManager.PERMISSION_GRANTED) {
+                return true;
+            }
+            return SystemConfig.getInstance().getBugreportWhitelistedPackages()
+                    .contains(mProjectionGrant.packageName);
+        }
+    }
+
+    @VisibleForTesting
+    void onKeyguardLockedStateChanged(boolean isKeyguardLocked) {
+        if (!isKeyguardLocked) return;
+        synchronized (mLock) {
+            if (mProjectionGrant != null && !canCaptureKeyguard()) {
+                Slog.d(TAG, "Content Recording: Stopped MediaProjection"
+                        + " due to keyguard lock");
+                mProjectionGrant.stop();
+            }
+        }
+    }
+
     /** Functional interface for providing time. */
     @VisibleForTesting
     interface Clock {
@@ -1252,6 +1297,11 @@
         @Override
         public void notifyVirtualDisplayCreated(int displayId) {
             notifyVirtualDisplayCreated_enforcePermission();
+            if (mKeyguardManager.isKeyguardLocked() && !canCaptureKeyguard()) {
+                Slog.w(TAG, "Content Recording: Keyguard locked, aborting MediaProjection");
+                stop();
+                return;
+            }
             synchronized (mLock) {
                 mVirtualDisplayId = displayId;
 
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index 5e644d3..834c17e 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -232,6 +232,7 @@
 import static com.android.server.wm.ActivityTaskManagerService.RELAUNCH_REASON_NONE;
 import static com.android.server.wm.ActivityTaskManagerService.RELAUNCH_REASON_WINDOWING_MODE_RESIZE;
 import static com.android.server.wm.ActivityTaskManagerService.getInputDispatchingTimeoutMillisLocked;
+import static com.android.server.wm.DesktopModeLaunchParamsModifier.canEnterDesktopMode;
 import static com.android.server.wm.IdentifierProto.HASH_CODE;
 import static com.android.server.wm.IdentifierProto.TITLE;
 import static com.android.server.wm.IdentifierProto.USER_ID;
@@ -9283,18 +9284,24 @@
     }
 
     void updateSizeCompatScale(Rect resolvedAppBounds, Rect containerAppBounds) {
-        // Only allow to scale down.
         mSizeCompatScale = mAppCompatController.getTransparentPolicy()
                 .findOpaqueNotFinishingActivityBelow()
                 .map(activityRecord -> activityRecord.mSizeCompatScale)
-                .orElseGet(() -> {
-                    final int contentW = resolvedAppBounds.width();
-                    final int contentH = resolvedAppBounds.height();
-                    final int viewportW = containerAppBounds.width();
-                    final int viewportH = containerAppBounds.height();
-                    return (contentW <= viewportW && contentH <= viewportH) ? 1f : Math.min(
-                            (float) viewportW / contentW, (float) viewportH / contentH);
-                });
+                .orElseGet(() -> calculateSizeCompatScale(resolvedAppBounds, containerAppBounds));
+    }
+
+    private float calculateSizeCompatScale(Rect resolvedAppBounds, Rect containerAppBounds) {
+        final int contentW = resolvedAppBounds.width();
+        final int contentH = resolvedAppBounds.height();
+        final int viewportW = containerAppBounds.width();
+        final int viewportH = containerAppBounds.height();
+        // Allow an application to be up-scaled if its window is smaller than its
+        // original container or if it's a freeform window in desktop mode.
+        boolean shouldAllowUpscaling = !(contentW <= viewportW && contentH <= viewportH)
+                || (canEnterDesktopMode(mAtmService.mContext)
+                    && getWindowingMode() == WINDOWING_MODE_FREEFORM);
+        return shouldAllowUpscaling ? Math.min(
+                (float) viewportW / contentW, (float) viewportH / contentH) : 1f;
     }
 
     private boolean isInSizeCompatModeForBounds(final Rect appBounds, final Rect containerBounds) {
diff --git a/services/tests/servicestests/src/com/android/server/media/projection/MediaProjectionManagerServiceTest.java b/services/tests/servicestests/src/com/android/server/media/projection/MediaProjectionManagerServiceTest.java
index e64397d..316b5fa 100644
--- a/services/tests/servicestests/src/com/android/server/media/projection/MediaProjectionManagerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/media/projection/MediaProjectionManagerServiceTest.java
@@ -17,6 +17,7 @@
 package com.android.server.media.projection;
 
 
+import static android.Manifest.permission.RECORD_SENSITIVE_CONTENT;
 import static android.app.WindowConfiguration.WINDOWING_MODE_MULTI_WINDOW;
 import static android.content.pm.ApplicationInfo.PRIVATE_FLAG_PRIVILEGED;
 import static android.media.projection.MediaProjectionManager.TYPE_MIRRORING;
@@ -50,6 +51,7 @@
 
 import android.app.ActivityManagerInternal;
 import android.app.ActivityOptions.LaunchCookie;
+import android.app.KeyguardManager;
 import android.content.Context;
 import android.content.ContextWrapper;
 import android.content.pm.ApplicationInfo;
@@ -66,7 +68,9 @@
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.os.test.TestLooper;
+import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.Presubmit;
+import android.platform.test.flag.junit.SetFlagsRule;
 import android.view.ContentRecordingSession;
 import android.view.ContentRecordingSession.RecordContent;
 
@@ -81,6 +85,7 @@
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -151,6 +156,9 @@
     private ContentRecordingSession mWaitingDisplaySession =
             createDisplaySession(DEFAULT_DISPLAY);
 
+    @Rule
+    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+
     @Mock
     private ActivityManagerInternal mAmInternal;
     @Mock
@@ -158,6 +166,8 @@
     @Mock
     private PackageManager mPackageManager;
     @Mock
+    private KeyguardManager mKeyguardManager;
+    @Mock
     private IMediaProjectionWatcherCallback mWatcherCallback;
     @Mock
     private MediaProjectionMetricsLogger mMediaProjectionMetricsLogger;
@@ -177,6 +187,7 @@
         mContext = spy(new ContextWrapper(
                 InstrumentationRegistry.getInstrumentation().getTargetContext()));
         doReturn(mPackageManager).when(mContext).getPackageManager();
+        doReturn(mKeyguardManager).when(mContext).getSystemService(eq(Context.KEYGUARD_SERVICE));
 
         mClock = new OffsettableClock.Stopped();
         mWaitingDisplaySession.setWaitingForConsent(true);
@@ -246,6 +257,39 @@
         assertThat(stoppedCallback2).isFalse();
     }
 
+    @EnableFlags(android.companion.virtualdevice.flags
+            .Flags.FLAG_MEDIA_PROJECTION_KEYGUARD_RESTRICTIONS)
+    @Test
+    public void testCreateProjection_keyguardLocked() throws Exception {
+        MediaProjectionManagerService.MediaProjection projection = startProjectionPreconditions();
+
+        doReturn(true).when(mKeyguardManager).isKeyguardLocked();
+        doReturn(PackageManager.PERMISSION_DENIED).when(mPackageManager)
+                .checkPermission(RECORD_SENSITIVE_CONTENT, projection.packageName);
+        projection.start(mIMediaProjectionCallback);
+        projection.notifyVirtualDisplayCreated(10);
+
+        assertThat(mService.getActiveProjectionInfo()).isNull();
+        assertThat(mIMediaProjectionCallback.mLatch.await(5, TimeUnit.SECONDS)).isTrue();
+    }
+
+    @EnableFlags(android.companion.virtualdevice.flags
+            .Flags.FLAG_MEDIA_PROJECTION_KEYGUARD_RESTRICTIONS)
+    @Test
+    public void testCreateProjection_keyguardLocked_packageAllowlisted()
+            throws NameNotFoundException {
+        MediaProjectionManagerService.MediaProjection projection = startProjectionPreconditions();
+
+        doReturn(true).when(mKeyguardManager).isKeyguardLocked();
+        doReturn(PackageManager.PERMISSION_GRANTED).when(mPackageManager)
+                .checkPermission(RECORD_SENSITIVE_CONTENT, projection.packageName);
+        projection.start(mIMediaProjectionCallback);
+        projection.notifyVirtualDisplayCreated(10);
+
+        // The projection was started because it was allowed to capture the keyguard.
+        assertThat(mService.getActiveProjectionInfo()).isNotNull();
+    }
+
     @Test
     public void testCreateProjection_attemptReuse_noPriorProjectionGrant()
             throws NameNotFoundException {
@@ -317,6 +361,48 @@
         assertThat(secondProjection).isNotEqualTo(projection);
     }
 
+    @EnableFlags(android.companion.virtualdevice.flags
+            .Flags.FLAG_MEDIA_PROJECTION_KEYGUARD_RESTRICTIONS)
+    @Test
+    public void testKeyguardLocked_stopsActiveProjection() throws Exception {
+        MediaProjectionManagerService service =
+                new MediaProjectionManagerService(mContext, mMediaProjectionMetricsLoggerInjector);
+        MediaProjectionManagerService.MediaProjection projection =
+                startProjectionPreconditions(service);
+        projection.start(mIMediaProjectionCallback);
+
+        assertThat(service.getActiveProjectionInfo()).isNotNull();
+
+        doReturn(PackageManager.PERMISSION_DENIED).when(mPackageManager)
+                .checkPermission(RECORD_SENSITIVE_CONTENT, projection.packageName);
+        service.onKeyguardLockedStateChanged(true);
+
+        verify(mMediaProjectionMetricsLogger).logStopped(UID, TARGET_UID_UNKNOWN);
+        assertThat(service.getActiveProjectionInfo()).isNull();
+        assertThat(mIMediaProjectionCallback.mLatch.await(5, TimeUnit.SECONDS)).isTrue();
+    }
+
+    @EnableFlags(android.companion.virtualdevice.flags
+            .Flags.FLAG_MEDIA_PROJECTION_KEYGUARD_RESTRICTIONS)
+    @Test
+    public void testKeyguardLocked_packageAllowlisted_doesNotStopActiveProjection()
+            throws NameNotFoundException {
+        MediaProjectionManagerService service =
+                new MediaProjectionManagerService(mContext, mMediaProjectionMetricsLoggerInjector);
+        MediaProjectionManagerService.MediaProjection projection =
+                startProjectionPreconditions(service);
+        projection.start(mIMediaProjectionCallback);
+
+        assertThat(service.getActiveProjectionInfo()).isNotNull();
+
+        doReturn(PackageManager.PERMISSION_GRANTED).when(mPackageManager).checkPermission(
+                RECORD_SENSITIVE_CONTENT, projection.packageName);
+        service.onKeyguardLockedStateChanged(true);
+
+        verifyZeroInteractions(mMediaProjectionMetricsLogger);
+        assertThat(service.getActiveProjectionInfo()).isNotNull();
+    }
+
     @Test
     public void stop_noActiveProjections_doesNotLog() throws Exception {
         MediaProjectionManagerService service =
diff --git a/services/tests/wmtests/src/com/android/server/wm/SizeCompatTests.java b/services/tests/wmtests/src/com/android/server/wm/SizeCompatTests.java
index 5fe8524..ae88b1b 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SizeCompatTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SizeCompatTests.java
@@ -104,6 +104,7 @@
 import android.os.Binder;
 import android.os.RemoteException;
 import android.os.UserHandle;
+import android.platform.test.annotations.DisableFlags;
 import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.Presubmit;
 import android.provider.DeviceConfig;
@@ -495,7 +496,7 @@
         // Activity is sandboxed; it is in size compat mode since it is not resizable and has a
         // max aspect ratio.
         assertActivityMaxBoundsSandboxed();
-        assertScaled();
+        assertDownScaled();
     }
 
     @Test
@@ -515,7 +516,7 @@
         // The bounds should be [100, 0 - 1100, 2500].
         assertEquals(origBounds.width(), currentBounds.width());
         assertEquals(origBounds.height(), currentBounds.height());
-        assertScaled();
+        assertDownScaled();
 
         // The scale is 2000/2500=0.8. The horizontal centered offset is (1000-(1000*0.8))/2=100.
         final float scale = (float) display.mBaseDisplayHeight / currentBounds.height();
@@ -555,7 +556,7 @@
         assertEquals(origBounds.width(), currentBounds.width());
         assertEquals(origBounds.height(), currentBounds.height());
         assertEquals(offsetX, mActivity.getBounds().left);
-        assertScaled();
+        assertDownScaled();
         // Activity is sandboxed due to size compat mode.
         assertActivityMaxBoundsSandboxed();
 
@@ -693,7 +694,7 @@
         // The configuration bounds [820, 0 - 1820, 2500] should keep the same.
         assertEquals(originalBounds.width(), currentBounds.width());
         assertEquals(originalBounds.height(), currentBounds.height());
-        assertScaled();
+        assertDownScaled();
         // Activity max bounds are sandboxed due to size compat mode on the new display.
         assertActivityMaxBoundsSandboxed();
 
@@ -752,7 +753,7 @@
         assertEquals(origAppBounds.width(), appBounds.width());
         assertEquals(origAppBounds.height(), appBounds.height());
         // The activity is 1000x1400 and the display is 2500x1000.
-        assertScaled();
+        assertDownScaled();
         final float scale = mActivity.getCompatScale();
         // The position in configuration should be in app coordinates.
         final Rect screenBounds = mActivity.getBounds();
@@ -849,7 +850,7 @@
         // Size compatibility mode is able to handle orientation change so the process shouldn't be
         // restarted and the override configuration won't be cleared.
         verify(mActivity, never()).restartProcessIfVisible();
-        assertScaled();
+        assertDownScaled();
         // Activity max bounds are sandboxed due to size compat mode, even if is not visible.
         assertActivityMaxBoundsSandboxed();
 
@@ -1624,6 +1625,85 @@
                 activity.getBounds().width(), 0.5);
     }
 
+
+    /**
+     * Test that a freeform unresizeable activity can be down-scaled to fill its smaller parent
+     * bounds.
+     */
+    @Test
+    public void testCompatScaling_freeformUnresizeableApp_largerThanParent_downScaled() {
+        final int dw = 600;
+        final int dh = 800;
+        final DisplayContent display = new TestDisplayContent.Builder(mAtm, dw, dh)
+                .setWindowingMode(WINDOWING_MODE_FREEFORM)
+                .build();
+        setUpApp(display);
+        prepareUnresizable(mActivity, /* maxAspect */ 0f, SCREEN_ORIENTATION_PORTRAIT);
+        mActivity.setWindowingMode(WINDOWING_MODE_FREEFORM);
+        assertFalse(mActivity.inSizeCompatMode());
+
+        // Resize app to make original app bounds larger than parent bounds.
+        mTask.getWindowConfiguration().setAppBounds(
+                new Rect(0, 0, dw - 300, dh - 400));
+        mActivity.onConfigurationChanged(mTask.getConfiguration());
+        // App should enter size compat mode and be down-scaled to fill new parent bounds.
+        assertDownScaled();
+    }
+
+    /**
+     * Test that when desktop mode is enabled, a freeform unresizeable activity can be up-scaled to
+     * fill its larger parent bounds.
+     */
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE)
+    public void testCompatScaling_freeformUnresizeableApp_smallerThanParent_upScaled() {
+        doReturn(true).when(() ->
+                DesktopModeLaunchParamsModifier.canEnterDesktopMode(any()));
+        final int dw = 600;
+        final int dh = 800;
+        final DisplayContent display = new TestDisplayContent.Builder(mAtm, dw, dh)
+                .setWindowingMode(WINDOWING_MODE_FREEFORM)
+                .build();
+        setUpApp(display);
+        prepareUnresizable(mActivity, /* maxAspect */ 0f, SCREEN_ORIENTATION_PORTRAIT);
+        mActivity.setWindowingMode(WINDOWING_MODE_FREEFORM);
+        assertFalse(mActivity.inSizeCompatMode());
+
+        // Resize app to make original app bounds smaller than parent bounds.
+        mTask.getWindowConfiguration().setAppBounds(
+                new Rect(0, 0, dw + 300, dh + 400));
+        mActivity.onConfigurationChanged(mTask.getConfiguration());
+        // App should enter size compat mode and be up-scaled to fill parent bounds.
+        assertUpScaled();
+    }
+
+    /**
+     * Test that when desktop mode is disabled, a freeform unresizeable activity cannot be up-scaled
+     * despite its larger parent bounds.
+     */
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE)
+    public void testSizeCompatScaling_freeformUnresizeableApp_smallerThanParent_notScaled() {
+        final int dw = 600;
+        final int dh = 800;
+        final DisplayContent display = new TestDisplayContent.Builder(mAtm, dw, dh)
+                .setWindowingMode(WINDOWING_MODE_FREEFORM)
+                .build();
+        setUpApp(display);
+        prepareUnresizable(mActivity, /* maxAspect */ 0f, SCREEN_ORIENTATION_PORTRAIT);
+        mActivity.setWindowingMode(WINDOWING_MODE_FREEFORM);
+        assertFalse(mActivity.inSizeCompatMode());
+        final Rect originalAppBounds = mActivity.getBounds();
+
+        // Resize app to make original app bounds smaller than parent bounds.
+        mTask.getWindowConfiguration().setAppBounds(
+                new Rect(0, 0, dw + 300, dh + 400));
+        mActivity.onConfigurationChanged(mTask.getConfiguration());
+        // App should enter size compat mode but remain its original size.
+        assertTrue(mActivity.inSizeCompatMode());
+        assertEquals(originalAppBounds, mActivity.getBounds());
+    }
+
     @Test
     public void testGetLetterboxInnerBounds_noScalingApplied() {
         // Set up a display in portrait and ignoring orientation request.
@@ -1659,7 +1739,7 @@
         // App should be in size compat.
         assertFalse(mActivity.mAppCompatController.getAppCompatAspectRatioPolicy()
                 .isLetterboxedForFixedOrientationAndAspectRatio());
-        assertScaled();
+        assertDownScaled();
         assertThat(mActivity.inSizeCompatMode()).isTrue();
         assertActivityMaxBoundsSandboxed();
 
@@ -2000,7 +2080,7 @@
 
         // After we rotate, the activity should go in the size-compat mode and report the same
         // configuration values.
-        assertScaled();
+        assertDownScaled();
         assertEquals(originalScreenWidthDp, mActivity.getConfiguration().smallestScreenWidthDp);
         assertEquals(originalScreenWidthDp, mActivity.getConfiguration().screenWidthDp);
         assertEquals(originalScreenHeighthDp, mActivity.getConfiguration().screenHeightDp);
@@ -2775,7 +2855,7 @@
         // App should be in size compat.
         assertFalse(mActivity.mAppCompatController.getAppCompatAspectRatioPolicy()
                 .isLetterboxedForFixedOrientationAndAspectRatio());
-        assertScaled();
+        assertDownScaled();
         assertEquals(activityBounds.width(), newActivityBounds.width());
         assertEquals(activityBounds.height(), newActivityBounds.height());
         assertActivityMaxBoundsSandboxed();
@@ -2807,7 +2887,7 @@
         // App should be in size compat.
         assertFalse(mActivity.mAppCompatController.getAppCompatAspectRatioPolicy()
                 .isLetterboxedForFixedOrientationAndAspectRatio());
-        assertScaled();
+        assertDownScaled();
         assertThat(mActivity.inSizeCompatMode()).isTrue();
         assertActivityMaxBoundsSandboxed();
 
@@ -2955,7 +3035,7 @@
         // App should be in size compat.
         assertFalse(mActivity.mAppCompatController.getAppCompatAspectRatioPolicy()
                 .isLetterboxedForFixedOrientationAndAspectRatio());
-        assertScaled();
+        assertDownScaled();
         assertThat(mActivity.inSizeCompatMode()).isTrue();
         // Activity max bounds are sandboxed due to size compat mode.
         assertActivityMaxBoundsSandboxed();
@@ -2967,7 +3047,7 @@
         verify(mActivity, never()).clearSizeCompatMode();
         assertFalse(mActivity.mAppCompatController.getAppCompatAspectRatioPolicy()
                 .isLetterboxedForFixedOrientationAndAspectRatio());
-        assertScaled();
+        assertDownScaled();
         assertEquals(activityBounds, mActivity.getBounds());
         // Activity max bounds are sandboxed due to size compat.
         assertActivityMaxBoundsSandboxed();
@@ -2995,7 +3075,7 @@
         // App should be in size compat.
         assertFalse(mActivity.mAppCompatController.getAppCompatAspectRatioPolicy()
                 .isLetterboxedForFixedOrientationAndAspectRatio());
-        assertScaled();
+        assertDownScaled();
         assertActivityMaxBoundsSandboxed();
 
         // Rotate display to landscape.
@@ -3032,7 +3112,7 @@
         // App should be in size compat.
         assertFalse(mActivity.mAppCompatController.getAppCompatAspectRatioPolicy()
                 .isLetterboxedForFixedOrientationAndAspectRatio());
-        assertScaled();
+        assertDownScaled();
         assertActivityMaxBoundsSandboxed();
 
         // Rotate display to portrait.
@@ -3224,7 +3304,7 @@
         assertEquals(WINDOWING_MODE_MULTI_WINDOW, mActivity.getWindowingMode());
 
         // Non-resizable activity in size compat mode
-        assertScaled();
+        assertDownScaled();
         final Rect newBounds = new Rect(mActivity.getWindowConfiguration().getBounds());
         assertEquals(originalBounds.width(), newBounds.width());
         assertEquals(originalBounds.height(), newBounds.height());
@@ -3288,7 +3368,7 @@
         assertEquals(WINDOWING_MODE_MULTI_WINDOW, mActivity.getWindowingMode());
 
         // Non-resizable activity in size compat mode
-        assertScaled();
+        assertDownScaled();
         final Rect newBounds = new Rect(mActivity.getWindowConfiguration().getBounds());
         assertEquals(originalBounds.width(), newBounds.width());
         assertEquals(originalBounds.height(), newBounds.height());
@@ -3329,7 +3409,7 @@
         organizer.mPrimary.setBounds(0, 0, 1000, 800);
 
         // Non-resizable activity should be in size compat mode.
-        assertScaled();
+        assertDownScaled();
         assertEquals(mActivity.getBounds(), new Rect(60, 0, 940, 800));
 
         recomputeNaturalConfigurationOfUnresizableActivity();
@@ -3906,7 +3986,7 @@
         // Force activity to scaled down for size compat mode.
         resizeDisplay(mTask.mDisplayContent, 700, 1400);
         assertTrue(mActivity.inSizeCompatMode());
-        assertScaled();
+        assertDownScaled();
         assertEquals(sizeCompatScaled, mActivity.getBounds());
     }
 
@@ -4406,7 +4486,7 @@
         resizeDisplay(mTask.mDisplayContent, 1400, 700);
 
         assertTrue(mActivity.inSizeCompatMode());
-        assertScaled();
+        assertDownScaled();
         assertEquals(sizeCompatScaled, mActivity.getBounds());
     }
 
@@ -4672,7 +4752,7 @@
         // Target min aspect ratio must be larger than parent aspect ratio to be applied.
         final float targetMinAspectRatio = 3.0f;
 
-        // Create fixed portait activity with min aspect ratio greater than parent aspect ratio.
+        // Create fixed portrait activity with min aspect ratio greater than parent aspect ratio.
         final ActivityRecord fixedOrientationActivity = new ActivityBuilder(mAtm)
                 .setTask(task).setScreenOrientation(SCREEN_ORIENTATION_PORTRAIT)
                 .setMinAspectRatio(targetMinAspectRatio).build();
@@ -4686,7 +4766,7 @@
         final Rect minAspectRatioAppBounds = new Rect(minAspectRatioActivity.getConfiguration()
                 .windowConfiguration.getAppBounds());
 
-        // Create unresizeable fixed portait activity with min aspect ratio greater than parent
+        // Create unresizeable fixed portrait activity with min aspect ratio greater than parent
         // aspect ratio.
         final ActivityRecord sizeCompatActivity = new ActivityBuilder(mAtm)
                 .setTask(task).setResizeMode(RESIZE_MODE_UNRESIZEABLE)
@@ -4719,7 +4799,7 @@
         // Activity should enter size compat with old density after display density change.
         display.setForcedDensity(newDensity, UserHandle.USER_CURRENT);
 
-        assertScaled();
+        assertDownScaled();
         assertEquals(origDensity, mActivity.getConfiguration().densityDpi);
 
         // Activity should exit size compat with new density.
@@ -4958,14 +5038,25 @@
         }
     }
 
-    private void assertScaled() {
-        assertScaled(mActivity);
+    private void assertUpScaled() {
+        assertScaled(mActivity, /* upScalingExpected */ true);
     }
 
-    /** Asserts that the size of activity is larger than its parent so it is scaling. */
-    private void assertScaled(ActivityRecord activity) {
+    private void assertDownScaled() {
+        assertScaled(mActivity, /* upScalingExpected */ false);
+    }
+
+    /**
+     * Asserts that the size of an activity differs from its parent and so it is scaling (either up
+     * or down).
+     */
+    private void assertScaled(ActivityRecord activity, boolean upScalingExpected) {
         assertTrue(activity.inSizeCompatMode());
-        assertNotEquals(1f, activity.getCompatScale(), 0.0001f /* delta */);
+        if (upScalingExpected) {
+            assertTrue(activity.getCompatScale() > 1f);
+        } else {
+            assertTrue(activity.getCompatScale() < 1f);
+        }
     }
 
     private void assertFitted() {