Merge "Fix idle maintenance debug command to actually trigger it" into main
diff --git a/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerService.java b/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerService.java
index 903248f..900c902 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerService.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerService.java
@@ -159,7 +159,6 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 import java.util.function.Predicate;
 
@@ -302,8 +301,6 @@
     private final ConnectivityController mConnectivityController;
     /** Need directly for sending uid state changes */
     private final DeviceIdleJobsController mDeviceIdleJobsController;
-    /** Need directly for sending exempted bucket changes */
-    private final FlexibilityController mFlexibilityController;
     /** Needed to get next estimated launch time. */
     private final PrefetchController mPrefetchController;
     /** Needed to get remaining quota time. */
@@ -516,10 +513,6 @@
                     if (name == null) {
                         continue;
                     }
-                    if (DEBUG) {
-                        Slog.d(TAG, "DeviceConfig " + name
-                                + " changed to " + properties.getString(name, null));
-                    }
                     switch (name) {
                         case Constants.KEY_ENABLE_API_QUOTAS:
                         case Constants.KEY_ENABLE_EXECUTION_SAFEGUARDS_UDC:
@@ -2539,17 +2532,17 @@
         mControllers = new ArrayList<StateController>();
         mPrefetchController = new PrefetchController(this);
         mControllers.add(mPrefetchController);
-        mFlexibilityController =
+        final FlexibilityController flexibilityController =
                 new FlexibilityController(this, mPrefetchController);
-        mControllers.add(mFlexibilityController);
+        mControllers.add(flexibilityController);
         mConnectivityController =
-                new ConnectivityController(this, mFlexibilityController);
+                new ConnectivityController(this, flexibilityController);
         mControllers.add(mConnectivityController);
         mControllers.add(new TimeController(this));
-        final IdleController idleController = new IdleController(this, mFlexibilityController);
+        final IdleController idleController = new IdleController(this, flexibilityController);
         mControllers.add(idleController);
         final BatteryController batteryController =
-                new BatteryController(this, mFlexibilityController);
+                new BatteryController(this, flexibilityController);
         mControllers.add(batteryController);
         mStorageController = new StorageController(this);
         mControllers.add(mStorageController);
@@ -3198,13 +3191,6 @@
     }
 
     @Override
-    public void onExemptedBucketChanged(@NonNull ArraySet<JobStatus> changedJobs) {
-        if (changedJobs.size() > 0) {
-            mFlexibilityController.onExemptedBucketChanged(changedJobs);
-        }
-    }
-
-    @Override
     public void onRestrictionStateChanged(@NonNull JobRestriction restriction,
             boolean stopOvertimeJobs) {
         mHandler.obtainMessage(MSG_CHECK_JOB).sendToTarget();
@@ -3511,10 +3497,7 @@
                 }
 
                 final boolean shouldForceBatchJob;
-                if (job.overrideState > JobStatus.OVERRIDE_NONE) {
-                    // The job should run for some test. Don't force batch it.
-                    shouldForceBatchJob = false;
-                } else if (job.shouldTreatAsExpeditedJob() || job.shouldTreatAsUserInitiatedJob()) {
+                if (job.shouldTreatAsExpeditedJob() || job.shouldTreatAsUserInitiatedJob()) {
                     // Never batch expedited or user-initiated jobs, even for RESTRICTED apps.
                     shouldForceBatchJob = false;
                 } else if (job.getEffectiveStandbyBucket() == RESTRICTED_INDEX) {
@@ -4967,8 +4950,6 @@
         Slog.d(TAG, "executeRunCommand(): " + pkgName + "/" + namespace + "/" + userId
                 + " " + jobId + " s=" + satisfied + " f=" + force);
 
-        final CountDownLatch delayLatch = new CountDownLatch(1);
-        final JobStatus js;
         try {
             final int uid = AppGlobals.getPackageManager().getPackageUid(pkgName, 0,
                     userId != UserHandle.USER_ALL ? userId : UserHandle.USER_SYSTEM);
@@ -4977,7 +4958,7 @@
             }
 
             synchronized (mLock) {
-                js = mJobs.getJobByUidAndJobId(uid, namespace, jobId);
+                final JobStatus js = mJobs.getJobByUidAndJobId(uid, namespace, jobId);
                 if (js == null) {
                     return JobSchedulerShellCommand.CMD_ERR_NO_JOB;
                 }
@@ -4988,71 +4969,23 @@
                 // Re-evaluate constraints after the override is set in case one of the overridden
                 // constraints was preventing another constraint from thinking it needed to update.
                 for (int c = mControllers.size() - 1; c >= 0; --c) {
-                    mControllers.get(c).evaluateStateLocked(js);
+                    mControllers.get(c).reevaluateStateLocked(uid);
                 }
 
                 if (!js.isConstraintsSatisfied()) {
-                    if (js.hasConnectivityConstraint()
-                            && !js.isConstraintSatisfied(JobStatus.CONSTRAINT_CONNECTIVITY)
-                            && js.wouldBeReadyWithConstraint(JobStatus.CONSTRAINT_CONNECTIVITY)) {
-                        // Because of how asynchronous the connectivity signals are, JobScheduler
-                        // may not get the connectivity satisfaction signal immediately. In this
-                        // case, wait a few seconds to see if it comes in before saying the
-                        // connectivity constraint isn't satisfied.
-                        mHandler.postDelayed(
-                                checkConstraintRunnableForTesting(
-                                        mHandler, js, delayLatch, 5, 1000),
-                                1000);
-                    } else {
-                        // There's no asynchronous signal to wait for. We can immediately say the
-                        // job's constraints aren't satisfied and return.
-                        js.overrideState = JobStatus.OVERRIDE_NONE;
-                        return JobSchedulerShellCommand.CMD_ERR_CONSTRAINTS;
-                    }
-                } else {
-                    delayLatch.countDown();
+                    js.overrideState = JobStatus.OVERRIDE_NONE;
+                    return JobSchedulerShellCommand.CMD_ERR_CONSTRAINTS;
                 }
+
+                queueReadyJobsForExecutionLocked();
+                maybeRunPendingJobsLocked();
             }
         } catch (RemoteException e) {
             // can't happen
-            return 0;
-        }
-
-        // Choose to block the return until we're sure about the state of the connectivity job
-        // so that tests can expect a reliable state after calling the run command.
-        try {
-            delayLatch.await(7L, TimeUnit.SECONDS);
-        } catch (InterruptedException e) {
-            Slog.e(TAG, "Couldn't wait for asynchronous constraint change", e);
-        }
-
-        synchronized (mLock) {
-            if (!js.isConstraintsSatisfied()) {
-                js.overrideState = JobStatus.OVERRIDE_NONE;
-                return JobSchedulerShellCommand.CMD_ERR_CONSTRAINTS;
-            }
-
-            queueReadyJobsForExecutionLocked();
-            maybeRunPendingJobsLocked();
         }
         return 0;
     }
 
-    private static Runnable checkConstraintRunnableForTesting(@NonNull final Handler handler,
-            @NonNull final JobStatus js, @NonNull final CountDownLatch latch,
-            final int remainingAttempts, final long delayMs) {
-        return () -> {
-            if (remainingAttempts <= 0 || js.isConstraintsSatisfied()) {
-                latch.countDown();
-                return;
-            }
-            handler.postDelayed(
-                    checkConstraintRunnableForTesting(
-                            handler, js, latch, remainingAttempts - 1, delayMs),
-                    delayMs);
-        };
-    }
-
     // Shell command infrastructure: immediately timeout currently executing jobs
     int executeStopCommand(PrintWriter pw, String pkgName, int userId,
             @Nullable String namespace, boolean hasJobId, int jobId,
diff --git a/apex/jobscheduler/service/java/com/android/server/job/StateChangedListener.java b/apex/jobscheduler/service/java/com/android/server/job/StateChangedListener.java
index 411a24d..50064bd 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/StateChangedListener.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/StateChangedListener.java
@@ -62,12 +62,6 @@
 
     /**
      * Called when these jobs are added or removed from the
-     * {@link android.app.usage.UsageStatsManager#STANDBY_BUCKET_EXEMPTED} bucket.
-     */
-    void onExemptedBucketChanged(@NonNull ArraySet<JobStatus> jobs);
-
-    /**
-     * Called when these jobs are added or removed from the
      * {@link android.app.usage.UsageStatsManager#STANDBY_BUCKET_RESTRICTED} bucket.
      */
     void onRestrictedBucketChanged(@NonNull List<JobStatus> jobs);
diff --git a/apex/jobscheduler/service/java/com/android/server/job/controllers/FlexibilityController.java b/apex/jobscheduler/service/java/com/android/server/job/controllers/FlexibilityController.java
index e4cb569..0e67b9a 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/controllers/FlexibilityController.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/controllers/FlexibilityController.java
@@ -20,7 +20,6 @@
 import static android.text.format.DateUtils.HOUR_IN_MILLIS;
 import static android.text.format.DateUtils.MINUTE_IN_MILLIS;
 
-import static com.android.server.job.JobSchedulerService.EXEMPTED_INDEX;
 import static com.android.server.job.JobSchedulerService.sElapsedRealtimeClock;
 import static com.android.server.job.controllers.JobStatus.CONSTRAINT_BATTERY_NOT_LOW;
 import static com.android.server.job.controllers.JobStatus.CONSTRAINT_CHARGING;
@@ -181,12 +180,8 @@
                 }
             };
 
-    private static final int MSG_CHECK_ALL_JOBS = 0;
-    /** Check the jobs in {@link #mJobsToCheck} */
-    private static final int MSG_CHECK_JOBS = 1;
-
-    @GuardedBy("mLock")
-    private final ArraySet<JobStatus> mJobsToCheck = new ArraySet<>();
+    private static final int MSG_UPDATE_JOBS = 0;
+    private static final int MSG_UPDATE_JOB = 1;
 
     public FlexibilityController(
             JobSchedulerService service, PrefetchController prefetchController) {
@@ -271,14 +266,7 @@
     @GuardedBy("mLock")
     boolean isFlexibilitySatisfiedLocked(JobStatus js) {
         return !mFlexibilityEnabled
-                // Exclude all jobs of the TOP app
                 || mService.getUidBias(js.getSourceUid()) == JobInfo.BIAS_TOP_APP
-                // Only exclude DEFAULT+ priority jobs for BFGS+ apps
-                || (mService.getUidBias(js.getSourceUid()) >= JobInfo.BIAS_BOUND_FOREGROUND_SERVICE
-                        && js.getEffectivePriority() >= JobInfo.PRIORITY_DEFAULT)
-                // Only exclude DEFAULT+ priority jobs for EXEMPTED apps
-                || (js.getStandbyBucket() == EXEMPTED_INDEX
-                        && js.getEffectivePriority() >= JobInfo.PRIORITY_DEFAULT)
                 || hasEnoughSatisfiedConstraintsLocked(js)
                 || mService.isCurrentlyRunningLocked(js);
     }
@@ -383,19 +371,11 @@
 
                 // Push the job update to the handler to avoid blocking other controllers and
                 // potentially batch back-to-back controller state updates together.
-                mHandler.obtainMessage(MSG_CHECK_ALL_JOBS).sendToTarget();
+                mHandler.obtainMessage(MSG_UPDATE_JOBS).sendToTarget();
             }
         }
     }
 
-    /** Called with a set of apps who have been added to or removed from the exempted bucket. */
-    public void onExemptedBucketChanged(@NonNull ArraySet<JobStatus> changedJobs) {
-        synchronized (mLock) {
-            mJobsToCheck.addAll(changedJobs);
-            mHandler.sendEmptyMessage(MSG_CHECK_JOBS);
-        }
-    }
-
     /** Checks if the given constraint is satisfied in the flexibility controller. */
     @VisibleForTesting
     boolean isConstraintSatisfied(int constraint) {
@@ -505,9 +485,7 @@
     @Override
     @GuardedBy("mLock")
     public void onUidBiasChangedLocked(int uid, int prevBias, int newBias) {
-        if (prevBias < JobInfo.BIAS_BOUND_FOREGROUND_SERVICE
-                && newBias < JobInfo.BIAS_BOUND_FOREGROUND_SERVICE) {
-            // All changes are below BFGS. There's no significant change to care about.
+        if (prevBias != JobInfo.BIAS_TOP_APP && newBias != JobInfo.BIAS_TOP_APP) {
             return;
         }
         final long nowElapsed = sElapsedRealtimeClock.millis();
@@ -732,8 +710,7 @@
                     }
                     mFlexibilityTracker.setNumDroppedFlexibleConstraints(js,
                             js.getNumAppliedFlexibleConstraints());
-                    mJobsToCheck.add(js);
-                    mHandler.sendEmptyMessage(MSG_CHECK_JOBS);
+                    mHandler.obtainMessage(MSG_UPDATE_JOB, js).sendToTarget();
                     return;
                 }
                 if (nextTimeElapsed == NO_LIFECYCLE_END) {
@@ -784,11 +761,10 @@
         @Override
         public void handleMessage(Message msg) {
             switch (msg.what) {
-                case MSG_CHECK_ALL_JOBS:
-                    removeMessages(MSG_CHECK_ALL_JOBS);
+                case MSG_UPDATE_JOBS:
+                    removeMessages(MSG_UPDATE_JOBS);
 
                     synchronized (mLock) {
-                        mJobsToCheck.clear();
                         final long nowElapsed = sElapsedRealtimeClock.millis();
                         final ArraySet<JobStatus> changedJobs = new ArraySet<>();
 
@@ -814,25 +790,19 @@
                     }
                     break;
 
-                case MSG_CHECK_JOBS:
+                case MSG_UPDATE_JOB:
                     synchronized (mLock) {
-                        final long nowElapsed = sElapsedRealtimeClock.millis();
-                        ArraySet<JobStatus> changedJobs = new ArraySet<>();
-
-                        for (int i = mJobsToCheck.size() - 1; i >= 0; --i) {
-                            final JobStatus js = mJobsToCheck.valueAt(i);
-                            if (DEBUG) {
-                                Slog.d(TAG, "Checking on " + js.toShortString());
-                            }
-                            if (js.setFlexibilityConstraintSatisfied(
-                                    nowElapsed, isFlexibilitySatisfiedLocked(js))) {
-                                changedJobs.add(js);
-                            }
+                        final JobStatus js = (JobStatus) msg.obj;
+                        if (DEBUG) {
+                            Slog.d("blah", "Checking on " + js.toShortString());
                         }
-
-                        mJobsToCheck.clear();
-                        if (changedJobs.size() > 0) {
-                            mStateChangedListener.onControllerStateChanged(changedJobs);
+                        final long nowElapsed = sElapsedRealtimeClock.millis();
+                        if (js.setFlexibilityConstraintSatisfied(
+                                nowElapsed, isFlexibilitySatisfiedLocked(js))) {
+                            // TODO(141645789): add method that will take a single job
+                            ArraySet<JobStatus> changedJob = new ArraySet<>();
+                            changedJob.add(js);
+                            mStateChangedListener.onControllerStateChanged(changedJob);
                         }
                     }
                     break;
@@ -1015,10 +985,7 @@
             pw.println(":");
             pw.increaseIndent();
 
-            pw.print(KEY_APPLIED_CONSTRAINTS, APPLIED_CONSTRAINTS);
-            pw.print("(");
-            JobStatus.dumpConstraints(pw, APPLIED_CONSTRAINTS);
-            pw.println(")");
+            pw.print(KEY_APPLIED_CONSTRAINTS, APPLIED_CONSTRAINTS).println();
             pw.print(KEY_DEADLINE_PROXIMITY_LIMIT, DEADLINE_PROXIMITY_LIMIT_MS).println();
             pw.print(KEY_FALLBACK_FLEXIBILITY_DEADLINE, FALLBACK_FLEXIBILITY_DEADLINE_MS).println();
             pw.print(KEY_MIN_TIME_BETWEEN_FLEXIBILITY_ALARMS_MS,
diff --git a/apex/jobscheduler/service/java/com/android/server/job/controllers/JobStatus.java b/apex/jobscheduler/service/java/com/android/server/job/controllers/JobStatus.java
index a095a16..0d5d11e 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/controllers/JobStatus.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/controllers/JobStatus.java
@@ -106,8 +106,11 @@
     public static final long NO_LATEST_RUNTIME = Long.MAX_VALUE;
     public static final long NO_EARLIEST_RUNTIME = 0L;
 
+    @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
     public static final int CONSTRAINT_CHARGING = JobInfo.CONSTRAINT_FLAG_CHARGING; // 1 < 0
+    @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
     public static final int CONSTRAINT_IDLE = JobInfo.CONSTRAINT_FLAG_DEVICE_IDLE;  // 1 << 2
+    @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
     public static final int CONSTRAINT_BATTERY_NOT_LOW =
             JobInfo.CONSTRAINT_FLAG_BATTERY_NOT_LOW; // 1 << 1
     @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
@@ -2191,7 +2194,7 @@
      * @return Whether or not this job would be ready to run if it had the specified constraint
      * granted, based on its requirements.
      */
-    public boolean wouldBeReadyWithConstraint(int constraint) {
+    boolean wouldBeReadyWithConstraint(int constraint) {
         return readinessStatusWithConstraint(constraint, true);
     }
 
diff --git a/apex/jobscheduler/service/java/com/android/server/job/controllers/QuotaController.java b/apex/jobscheduler/service/java/com/android/server/job/controllers/QuotaController.java
index 04da781..8ddbf69 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/controllers/QuotaController.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/controllers/QuotaController.java
@@ -2511,7 +2511,6 @@
                     + " to bucketIndex " + bucketIndex);
         }
         List<JobStatus> restrictedChanges = new ArrayList<>();
-        ArraySet<JobStatus> exemptedChanges = new ArraySet<>();
         synchronized (mLock) {
             ShrinkableDebits debits = mEJStats.get(userId, packageName);
             if (debits != null) {
@@ -2531,10 +2530,6 @@
                         && bucketIndex != js.getStandbyBucket()) {
                     restrictedChanges.add(js);
                 }
-                if ((bucketIndex == EXEMPTED_INDEX || js.getStandbyBucket() == EXEMPTED_INDEX)
-                        && bucketIndex != js.getStandbyBucket()) {
-                    exemptedChanges.add(js);
-                }
                 js.setStandbyBucket(bucketIndex);
             }
             Timer timer = mPkgTimers.get(userId, packageName);
@@ -2549,9 +2544,6 @@
                     maybeUpdateConstraintForPkgLocked(
                             sElapsedRealtimeClock.millis(), userId, packageName));
         }
-        if (exemptedChanges.size() > 0) {
-            mStateChangedListener.onExemptedBucketChanged(exemptedChanges);
-        }
         if (restrictedChanges.size() > 0) {
             mStateChangedListener.onRestrictedBucketChanged(restrictedChanges);
         }
diff --git a/apex/jobscheduler/service/java/com/android/server/usage/AppIdleHistory.java b/apex/jobscheduler/service/java/com/android/server/usage/AppIdleHistory.java
index 913a76a..4d4e340 100644
--- a/apex/jobscheduler/service/java/com/android/server/usage/AppIdleHistory.java
+++ b/apex/jobscheduler/service/java/com/android/server/usage/AppIdleHistory.java
@@ -591,6 +591,16 @@
         if (idle) {
             newBucket = IDLE_BUCKET_CUTOFF;
             reason = REASON_MAIN_FORCED_BY_USER;
+            final AppUsageHistory appHistory = getAppUsageHistory(packageName, userId,
+                    elapsedRealtime);
+            // Wipe all expiry times that could raise the bucket on reevaluation.
+            if (appHistory.bucketExpiryTimesMs != null) {
+                for (int i = appHistory.bucketExpiryTimesMs.size() - 1; i >= 0; --i) {
+                    if (appHistory.bucketExpiryTimesMs.keyAt(i) < newBucket) {
+                        appHistory.bucketExpiryTimesMs.removeAt(i);
+                    }
+                }
+            }
         } else {
             newBucket = STANDBY_BUCKET_ACTIVE;
             // This is to pretend that the app was just used, don't freeze the state anymore.
diff --git a/apex/jobscheduler/service/java/com/android/server/usage/AppStandbyController.java b/apex/jobscheduler/service/java/com/android/server/usage/AppStandbyController.java
index bd5c006..e589c21 100644
--- a/apex/jobscheduler/service/java/com/android/server/usage/AppStandbyController.java
+++ b/apex/jobscheduler/service/java/com/android/server/usage/AppStandbyController.java
@@ -59,6 +59,7 @@
 import static com.android.server.usage.AppIdleHistory.STANDBY_BUCKET_UNKNOWN;
 
 import android.annotation.CurrentTimeMillisLong;
+import android.annotation.DurationMillisLong;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.UserIdInt;
@@ -2146,6 +2147,15 @@
         }
     }
 
+    /**
+     * Flush the handler.
+     * Returns true if successfully flushed within the timeout, otherwise return false.
+     */
+    @VisibleForTesting
+    boolean flushHandler(@DurationMillisLong long timeoutMillis) {
+        return mHandler.runWithScissors(() -> {}, timeoutMillis);
+    }
+
     @Override
     public void flushToDisk() {
         synchronized (mAppIdleLock) {
diff --git a/core/java/android/app/OWNERS b/core/java/android/app/OWNERS
index 772b0b4..47d19ed 100644
--- a/core/java/android/app/OWNERS
+++ b/core/java/android/app/OWNERS
@@ -88,6 +88,10 @@
 # Pinner
 per-file pinner-client.aconfig = file:/core/java/android/app/pinner/OWNERS
 
+# BackgroundInstallControlManager
+per-file BackgroundInstallControlManager.java = file:/services/core/java/com/android/server/pm/OWNERS
+per-file background_install_control_manager.aconfig = file:/services/core/java/com/android/server/pm/OWNERS
+
 # ResourcesManager
 per-file ResourcesManager.java = file:RESOURCES_OWNERS
 
diff --git a/core/java/android/companion/CompanionDeviceManager.java b/core/java/android/companion/CompanionDeviceManager.java
index b11840e..879656a 100644
--- a/core/java/android/companion/CompanionDeviceManager.java
+++ b/core/java/android/companion/CompanionDeviceManager.java
@@ -220,6 +220,12 @@
      */
     public static final int MESSAGE_REQUEST_PING = 0x63807378; // ?PIN
     /**
+     * Test message type without a response.
+     *
+     * @hide
+     */
+    public static final int MESSAGE_ONEWAY_PING = 0x43807378; // +PIN
+    /**
      * Message header assigned to the remote authentication handshakes.
      *
      * @hide
@@ -237,6 +243,18 @@
      * @hide
      */
     public static final int MESSAGE_REQUEST_PERMISSION_RESTORE = 0x63826983; // ?RES
+    /**
+     * Message header assigned to the one-way message sent from the wearable device.
+     *
+     * @hide
+     */
+    public static final int MESSAGE_ONEWAY_FROM_WEARABLE = 0x43708287; // +FRW
+    /**
+     * Message header assigned to the one-way message sent to the wearable device.
+     *
+     * @hide
+     */
+    public static final int MESSAGE_ONEWAY_TO_WEARABLE = 0x43847987; // +TOW
 
     /**
      * The length limit of Association tag.
diff --git a/core/java/android/view/View.java b/core/java/android/view/View.java
index cbafd1c..e4b709e 100644
--- a/core/java/android/view/View.java
+++ b/core/java/android/view/View.java
@@ -33056,7 +33056,10 @@
     }
 
     private float getSizePercentage() {
-        if (mResources == null || getVisibility() != VISIBLE) {
+        float alpha = mTransformationInfo != null ? mTransformationInfo.mAlpha : 1;
+        int visibility = mViewFlags & VISIBILITY_MASK;
+
+        if (mResources == null || alpha == 0 || visibility != VISIBLE) {
             return 0;
         }
 
@@ -33081,25 +33084,26 @@
 
     private void votePreferredFrameRate() {
         // use toolkitSetFrameRate flag to gate the change
-        ViewRootImpl viewRootImpl = getViewRootImpl();
-        float sizePercentage = getSizePercentage();
-        int frameRateCateogry = calculateFrameRateCategory(sizePercentage);
-        if (sToolkitSetFrameRateReadOnlyFlagValue && viewRootImpl != null
-                && sizePercentage > 0) {
-            if (mPreferredFrameRate < 0) {
-                if (mPreferredFrameRate == REQUESTED_FRAME_RATE_CATEGORY_NO_PREFERENCE) {
-                    frameRateCateogry = FRAME_RATE_CATEGORY_NO_PREFERENCE;
-                } else if (mPreferredFrameRate == REQUESTED_FRAME_RATE_CATEGORY_LOW) {
-                    frameRateCateogry = FRAME_RATE_CATEGORY_LOW;
-                } else if (mPreferredFrameRate == REQUESTED_FRAME_RATE_CATEGORY_NORMAL) {
-                    frameRateCateogry = FRAME_RATE_CATEGORY_NORMAL;
-                } else if (mPreferredFrameRate == REQUESTED_FRAME_RATE_CATEGORY_HIGH) {
-                    frameRateCateogry = FRAME_RATE_CATEGORY_HIGH;
+        if (sToolkitSetFrameRateReadOnlyFlagValue) {
+            ViewRootImpl viewRootImpl = getViewRootImpl();
+            float sizePercentage = getSizePercentage();
+            int frameRateCateogry = calculateFrameRateCategory(sizePercentage);
+            if (viewRootImpl != null && sizePercentage > 0) {
+                if (mPreferredFrameRate < 0) {
+                    if (mPreferredFrameRate == REQUESTED_FRAME_RATE_CATEGORY_NO_PREFERENCE) {
+                        frameRateCateogry = FRAME_RATE_CATEGORY_NO_PREFERENCE;
+                    } else if (mPreferredFrameRate == REQUESTED_FRAME_RATE_CATEGORY_LOW) {
+                        frameRateCateogry = FRAME_RATE_CATEGORY_LOW;
+                    } else if (mPreferredFrameRate == REQUESTED_FRAME_RATE_CATEGORY_NORMAL) {
+                        frameRateCateogry = FRAME_RATE_CATEGORY_NORMAL;
+                    } else if (mPreferredFrameRate == REQUESTED_FRAME_RATE_CATEGORY_HIGH) {
+                        frameRateCateogry = FRAME_RATE_CATEGORY_HIGH;
+                    }
+                } else {
+                    viewRootImpl.votePreferredFrameRate(mPreferredFrameRate);
                 }
-            } else {
-                viewRootImpl.votePreferredFrameRate(mPreferredFrameRate);
+                viewRootImpl.votePreferredFrameRateCategory(frameRateCateogry);
             }
-            viewRootImpl.votePreferredFrameRateCategory(frameRateCateogry);
         }
     }
 
diff --git a/core/tests/companiontests/src/android/companion/SystemDataTransportTest.java b/core/tests/companiontests/src/android/companion/SystemDataTransportTest.java
index 2b4123a..73d7fe9 100644
--- a/core/tests/companiontests/src/android/companion/SystemDataTransportTest.java
+++ b/core/tests/companiontests/src/android/companion/SystemDataTransportTest.java
@@ -16,6 +16,9 @@
 
 package android.companion;
 
+import static android.companion.CompanionDeviceManager.MESSAGE_ONEWAY_PING;
+import static android.companion.CompanionDeviceManager.MESSAGE_REQUEST_PING;
+
 import android.content.Context;
 import android.os.SystemClock;
 import android.test.InstrumentationTestCase;
@@ -36,16 +39,22 @@
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.Random;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 
+/**
+ * Tests that CDM can intake incoming messages in the system data transport and output results.
+ *
+ * Build/Install/Run: atest CompanionTests:SystemDataTransportTest
+ */
 public class SystemDataTransportTest extends InstrumentationTestCase {
     private static final String TAG = "SystemDataTransportTest";
 
     private static final int MESSAGE_INVALID = 0xF00DCAFE;
-
-    private static final int MESSAGE_REQUEST_INVALID = 0x63636363; // ????
-    private static final int MESSAGE_REQUEST_PING = 0x63807378; // ?PIN
-
+    private static final int MESSAGE_ONEWAY_INVALID = 0x43434343; // ++++
     private static final int MESSAGE_RESPONSE_INVALID = 0x33333333; // !!!!
+    private static final int MESSAGE_REQUEST_INVALID = 0x63636363; // ????
+
     private static final int MESSAGE_RESPONSE_SUCCESS = 0x33838567; // !SUC
     private static final int MESSAGE_RESPONSE_FAILURE = 0x33706573; // !FAI
 
@@ -122,8 +131,6 @@
         new Random().nextBytes(blob);
 
         final byte[] input = generatePacket(MESSAGE_REQUEST_PING, /* sequence */ 1, blob);
-        final byte[] expected = generatePacket(MESSAGE_RESPONSE_SUCCESS, /* sequence */ 1, blob);
-        assertTransportBehavior(input, expected);
     }
 
     public void testMultiplePingPing() {
@@ -176,6 +183,43 @@
         testPingHandRolled();
     }
 
+    public void testInvalidOnewayMessages() throws InterruptedException {
+        // Add a callback
+        final CountDownLatch received = new CountDownLatch(1);
+        mCdm.addOnMessageReceivedListener(Runnable::run, MESSAGE_ONEWAY_INVALID,
+                (id, data) -> received.countDown());
+
+        final byte[] input = generatePacket(MESSAGE_ONEWAY_INVALID, /* sequence */ 1);
+        final ByteArrayInputStream in = new ByteArrayInputStream(input);
+        final ByteArrayOutputStream out = new ByteArrayOutputStream();
+        mCdm.attachSystemDataTransport(mAssociationId, in, out);
+
+        // Assert that a one-way message was ignored (does not trigger a callback)
+        assertFalse(received.await(5, TimeUnit.SECONDS));
+
+        // There should not be a response to one-way messages
+        assertEquals(0, out.toByteArray().length);
+    }
+
+
+    public void testOnewayMessages() throws InterruptedException {
+        // Add a callback
+        final CountDownLatch received = new CountDownLatch(1);
+        mCdm.addOnMessageReceivedListener(Runnable::run, MESSAGE_ONEWAY_PING,
+                (id, data) -> received.countDown());
+
+        final byte[] input = generatePacket(MESSAGE_ONEWAY_PING, /* sequence */ 1);
+        final ByteArrayInputStream in = new ByteArrayInputStream(input);
+        final ByteArrayOutputStream out = new ByteArrayOutputStream();
+        mCdm.attachSystemDataTransport(mAssociationId, in, out);
+
+        // Assert that a one-way message was received
+        assertTrue(received.await(1, TimeUnit.SECONDS));
+
+        // There should not be a response to one-way messages
+        assertEquals(0, out.toByteArray().length);
+    }
+
     public static byte[] concat(byte[]... blobs) {
         int length = 0;
         for (byte[] blob : blobs) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleStackView.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleStackView.java
index 02a505e..470a825 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleStackView.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleStackView.java
@@ -2323,7 +2323,8 @@
         updateOverflowVisibility();
         updatePointerPosition(false /* forIme */);
         mExpandedAnimationController.expandFromStack(() -> {
-            if (mIsExpanded && mExpandedBubble.getExpandedView() != null) {
+            if (mIsExpanded && mExpandedBubble != null
+                    && mExpandedBubble.getExpandedView() != null) {
                 maybeShowManageEdu();
             }
             updateOverflowDotVisibility(true /* expanding */);
@@ -2384,7 +2385,7 @@
         }
         mExpandedViewContainer.setAnimationMatrix(mExpandedViewContainerMatrix);
 
-        if (mExpandedBubble.getExpandedView() != null) {
+        if (mExpandedBubble != null && mExpandedBubble.getExpandedView() != null) {
             mExpandedBubble.getExpandedView().setContentAlpha(0f);
             mExpandedBubble.getExpandedView().setBackgroundAlpha(0f);
 
diff --git a/location/java/android/location/flags/gnss.aconfig b/location/java/android/location/flags/gnss.aconfig
index b6055e8..f4b1056 100644
--- a/location/java/android/location/flags/gnss.aconfig
+++ b/location/java/android/location/flags/gnss.aconfig
@@ -20,3 +20,10 @@
     description: "Flag for GnssMeasurementRequest WorkSource API"
     bug: "295235160"
 }
+
+flag {
+    name: "release_supl_connection_on_timeout"
+    namespace: "location"
+    description: "Flag for releasing SUPL connection on timeout"
+    bug: "315024652"
+}
diff --git a/packages/SettingsProvider/Android.bp b/packages/SettingsProvider/Android.bp
index adebdcd..d5814e3 100644
--- a/packages/SettingsProvider/Android.bp
+++ b/packages/SettingsProvider/Android.bp
@@ -59,10 +59,10 @@
         // Note we statically link SettingsProviderLib to do some unit tests.  It's not accessible otherwise
         // because this test is not an instrumentation test. (because the target runs in the system process.)
         "SettingsProviderLib",
-
         "androidx.test.rules",
         "flag-junit",
         "junit",
+        "libaconfig_java_proto_lite",
         "mockito-target-minus-junit4",
         "platform-test-annotations",
         "truth",
diff --git a/packages/SettingsProvider/src/com/android/providers/settings/SettingsState.java b/packages/SettingsProvider/src/com/android/providers/settings/SettingsState.java
index 8f459c6..73c2e22 100644
--- a/packages/SettingsProvider/src/com/android/providers/settings/SettingsState.java
+++ b/packages/SettingsProvider/src/com/android/providers/settings/SettingsState.java
@@ -18,6 +18,9 @@
 
 import static android.os.Process.FIRST_APPLICATION_UID;
 
+import android.aconfig.Aconfig.flag_state;
+import android.aconfig.Aconfig.parsed_flag;
+import android.aconfig.Aconfig.parsed_flags;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.content.Context;
@@ -147,6 +150,17 @@
      */
     private static final String CONFIG_STAGED_PREFIX = "staged/";
 
+    private static final List<String> sAconfigTextProtoFilesOnDevice = List.of(
+            "/system/etc/aconfig_flags.pb",
+            "/system_ext/etc/aconfig_flags.pb",
+            "/product/etc/aconfig_flags.pb",
+            "/vendor/etc/aconfig_flags.pb");
+
+    /**
+     * This tag is applied to all aconfig default value-loaded flags.
+     */
+    private static final String BOOT_LOADED_DEFAULT_TAG = "BOOT_LOADED_DEFAULT";
+
     // This was used in version 120 and before.
     private static final String NULL_VALUE_OLD_STYLE = "null";
 
@@ -315,6 +329,59 @@
 
         synchronized (mLock) {
             readStateSyncLocked();
+
+            if (Flags.loadAconfigDefaults()) {
+                // Only load aconfig defaults if this is the first boot, the XML
+                // file doesn't exist yet, or this device is on its first boot after
+                // an OTA.
+                boolean shouldLoadAconfigValues = isConfigSettingsKey(mKey)
+                        && (!file.exists()
+                                || mContext.getPackageManager().isDeviceUpgrading());
+                if (shouldLoadAconfigValues) {
+                    loadAconfigDefaultValuesLocked();
+                }
+            }
+        }
+    }
+
+    @GuardedBy("mLock")
+    private void loadAconfigDefaultValuesLocked() {
+        for (String fileName : sAconfigTextProtoFilesOnDevice) {
+            try (FileInputStream inputStream = new FileInputStream(fileName)) {
+                byte[] contents = inputStream.readAllBytes();
+                loadAconfigDefaultValues(contents);
+            } catch (IOException e) {
+                Slog.e(LOG_TAG, "failed to read protobuf", e);
+            }
+        }
+    }
+
+    @VisibleForTesting
+    @GuardedBy("mLock")
+    public void loadAconfigDefaultValues(byte[] fileContents) {
+        try {
+            parsed_flags parsedFlags = parsed_flags.parseFrom(fileContents);
+
+            if (parsedFlags == null) {
+                Slog.e(LOG_TAG, "failed to parse aconfig protobuf");
+                return;
+            }
+
+            for (parsed_flag flag : parsedFlags.getParsedFlagList()) {
+                String flagName = flag.getNamespace() + "/"
+                        + flag.getPackage() + "." + flag.getName();
+                String value = flag.getState() == flag_state.ENABLED ? "true" : "false";
+
+                Setting existingSetting = getSettingLocked(flagName);
+                boolean isDefaultLoaded = existingSetting.getTag() != null
+                        && existingSetting.getTag().equals(BOOT_LOADED_DEFAULT_TAG);
+                if (existingSetting.getValue() == null || isDefaultLoaded) {
+                    insertSettingLocked(flagName, value, BOOT_LOADED_DEFAULT_TAG,
+                            false, flag.getPackage());
+                }
+            }
+        } catch (IOException e) {
+            Slog.e(LOG_TAG, "failed to parse protobuf", e);
         }
     }
 
diff --git a/packages/SettingsProvider/src/com/android/providers/settings/device_config_service.aconfig b/packages/SettingsProvider/src/com/android/providers/settings/device_config_service.aconfig
index 27ce0d4..ecac5ee 100644
--- a/packages/SettingsProvider/src/com/android/providers/settings/device_config_service.aconfig
+++ b/packages/SettingsProvider/src/com/android/providers/settings/device_config_service.aconfig
@@ -6,3 +6,11 @@
     description: "When enabled, allows setting and displaying local overrides via adb."
     bug: "298392357"
 }
+
+flag {
+    name: "load_aconfig_defaults"
+    namespace: "core_experiments_team_internal"
+    description: "When enabled, loads aconfig default values into DeviceConfig on boot."
+    bug: "311155098"
+    is_fixed_read_only: true
+}
diff --git a/packages/SettingsProvider/test/src/com/android/providers/settings/SettingsStateTest.java b/packages/SettingsProvider/test/src/com/android/providers/settings/SettingsStateTest.java
index 02a7bc1..24625ea 100644
--- a/packages/SettingsProvider/test/src/com/android/providers/settings/SettingsStateTest.java
+++ b/packages/SettingsProvider/test/src/com/android/providers/settings/SettingsStateTest.java
@@ -15,6 +15,9 @@
  */
 package com.android.providers.settings;
 
+import android.aconfig.Aconfig;
+import android.aconfig.Aconfig.parsed_flag;
+import android.aconfig.Aconfig.parsed_flags;
 import android.os.Looper;
 import android.test.AndroidTestCase;
 import android.util.Xml;
@@ -84,6 +87,86 @@
         super.tearDown();
     }
 
+    public void testLoadValidAconfigProto() {
+        int configKey = SettingsState.makeKey(SettingsState.SETTINGS_TYPE_CONFIG, 0);
+        Object lock = new Object();
+        SettingsState settingsState = new SettingsState(
+                getContext(), lock, mSettingsFile, configKey,
+                SettingsState.MAX_BYTES_PER_APP_PACKAGE_UNLIMITED, Looper.getMainLooper());
+
+        parsed_flags flags = parsed_flags
+                .newBuilder()
+                .addParsedFlag(parsed_flag
+                    .newBuilder()
+                        .setPackage("com.android.flags")
+                        .setName("flag1")
+                        .setNamespace("test_namespace")
+                        .setDescription("test flag")
+                        .addBug("12345678")
+                        .setState(Aconfig.flag_state.DISABLED)
+                        .setPermission(Aconfig.flag_permission.READ_WRITE))
+                .addParsedFlag(parsed_flag
+                    .newBuilder()
+                        .setPackage("com.android.flags")
+                        .setName("flag2")
+                        .setNamespace("test_namespace")
+                        .setDescription("another test flag")
+                        .addBug("12345678")
+                        .setState(Aconfig.flag_state.ENABLED)
+                        .setPermission(Aconfig.flag_permission.READ_WRITE))
+                .build();
+
+        synchronized (lock) {
+            settingsState.loadAconfigDefaultValues(flags.toByteArray());
+            settingsState.persistSettingsLocked();
+        }
+        settingsState.waitForHandler();
+
+        synchronized (lock) {
+            assertEquals("false",
+                    settingsState.getSettingLocked(
+                        "test_namespace/com.android.flags.flag1").getValue());
+            assertEquals("true",
+                    settingsState.getSettingLocked(
+                        "test_namespace/com.android.flags.flag2").getValue());
+        }
+    }
+
+    public void testSkipLoadingAconfigFlagWithMissingFields() {
+        int configKey = SettingsState.makeKey(SettingsState.SETTINGS_TYPE_CONFIG, 0);
+        Object lock = new Object();
+        SettingsState settingsState = new SettingsState(
+                getContext(), lock, mSettingsFile, configKey,
+                SettingsState.MAX_BYTES_PER_APP_PACKAGE_UNLIMITED, Looper.getMainLooper());
+
+        parsed_flags flags = parsed_flags
+                .newBuilder()
+                .addParsedFlag(parsed_flag
+                    .newBuilder()
+                        .setDescription("test flag")
+                        .addBug("12345678")
+                        .setState(Aconfig.flag_state.DISABLED)
+                        .setPermission(Aconfig.flag_permission.READ_WRITE))
+                .build();
+
+        synchronized (lock) {
+            settingsState.loadAconfigDefaultValues(flags.toByteArray());
+            settingsState.persistSettingsLocked();
+        }
+        settingsState.waitForHandler();
+
+        synchronized (lock) {
+            assertEquals(null,
+                    settingsState.getSettingLocked(
+                        "test_namespace/com.android.flags.flag1").getValue());
+        }
+    }
+
+    public void testInvalidAconfigProtoDoesNotCrash() {
+        SettingsState settingsState = getSettingStateObject();
+        settingsState.loadAconfigDefaultValues("invalid protobuf".getBytes());
+    }
+
     public void testIsBinary() {
         assertFalse(SettingsState.isBinary(" abc 日本語"));
 
diff --git a/packages/SystemUI/AndroidManifest.xml b/packages/SystemUI/AndroidManifest.xml
index 1a35f04..a03fa9b 100644
--- a/packages/SystemUI/AndroidManifest.xml
+++ b/packages/SystemUI/AndroidManifest.xml
@@ -256,6 +256,9 @@
     <!-- launcher apps -->
     <uses-permission android:name="android.permission.ACCESS_SHORTCUTS" />
 
+    <!-- Permission to start Launcher's widget picker activity. -->
+    <uses-permission android:name="android.permission.START_WIDGET_PICKER_ACTIVITY" />
+
     <uses-permission android:name="android.permission.MODIFY_THEME_OVERLAY" />
 
     <!-- accessibility -->
@@ -974,15 +977,6 @@
             android:excludeFromRecents="true"
             android:visibleToInstantApps="true"/>
 
-        <activity android:name="com.android.systemui.communal.widgets.WidgetPickerActivity"
-            android:theme="@style/Theme.EditWidgetsActivity"
-            android:excludeFromRecents="true"
-            android:autoRemoveFromRecents="true"
-            android:showOnLockScreen="true"
-            android:launchMode="singleTop"
-            android:exported="false">
-        </activity>
-
         <activity android:name="com.android.systemui.communal.widgets.EditWidgetsActivity"
             android:theme="@style/Theme.EditWidgetsActivity"
             android:excludeFromRecents="true"
diff --git a/packages/SystemUI/aconfig/systemui.aconfig b/packages/SystemUI/aconfig/systemui.aconfig
index 8b27960..b4c3850 100644
--- a/packages/SystemUI/aconfig/systemui.aconfig
+++ b/packages/SystemUI/aconfig/systemui.aconfig
@@ -261,3 +261,10 @@
         "prefer using alpha to distinguish network activity."
     bug: "310715220"
 }
+
+flag {
+    name: "haptic_volume_slider"
+    namespace: "systemui"
+    description: "Adds haptic feedback to the volume slider."
+    bug: "316953430"
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/authentication/data/repository/AuthenticationRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/authentication/data/repository/AuthenticationRepositoryTest.kt
index 64ddbc7..c961be9 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/authentication/data/repository/AuthenticationRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/authentication/data/repository/AuthenticationRepositoryTest.kt
@@ -35,8 +35,10 @@
 import com.android.systemui.statusbar.pipeline.mobile.util.FakeMobileMappingsProxy
 import com.android.systemui.user.data.repository.FakeUserRepository
 import com.android.systemui.util.mockito.whenever
+import com.android.systemui.util.time.FakeSystemClock
 import com.google.common.truth.Truth.assertThat
 import java.util.function.Function
+import kotlin.time.Duration.Companion.seconds
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.runCurrent
@@ -58,6 +60,7 @@
 
     private val testUtils = SceneTestUtils(this)
     private val testScope = testUtils.testScope
+    private val clock = FakeSystemClock()
     private val userRepository = FakeUserRepository()
     private lateinit var mobileConnectionsRepository: FakeMobileConnectionsRepository
 
@@ -78,8 +81,10 @@
         underTest =
             AuthenticationRepositoryImpl(
                 applicationScope = testScope.backgroundScope,
-                getSecurityMode = getSecurityMode,
                 backgroundDispatcher = testUtils.testDispatcher,
+                flags = testUtils.sceneContainerFlags,
+                clock = clock,
+                getSecurityMode = getSecurityMode,
                 userRepository = userRepository,
                 lockPatternUtils = lockPatternUtils,
                 broadcastDispatcher = fakeBroadcastDispatcher,
@@ -141,22 +146,6 @@
         }
 
     @Test
-    fun reportAuthenticationAttempt_emitsAuthenticationChallengeResult() =
-        testScope.runTest {
-            val authenticationChallengeResults by
-                collectValues(underTest.authenticationChallengeResult)
-
-            runCurrent()
-            underTest.reportAuthenticationAttempt(true)
-            runCurrent()
-            underTest.reportAuthenticationAttempt(false)
-            runCurrent()
-            underTest.reportAuthenticationAttempt(true)
-
-            assertThat(authenticationChallengeResults).isEqualTo(listOf(true, false, true))
-        }
-
-    @Test
     fun isPinEnhancedPrivacyEnabled() =
         testScope.runTest {
             whenever(lockPatternUtils.isPinEnhancedPrivacyEnabled(USER_INFOS[0].id))
@@ -172,6 +161,45 @@
             assertThat(values.last()).isTrue()
         }
 
+    @Test
+    fun lockoutEndTimestamp() =
+        testScope.runTest {
+            val lockoutEndMs = clock.elapsedRealtime() + 30.seconds.inWholeMilliseconds
+            whenever(lockPatternUtils.getLockoutAttemptDeadline(USER_INFOS[0].id))
+                .thenReturn(lockoutEndMs)
+            whenever(lockPatternUtils.getLockoutAttemptDeadline(USER_INFOS[1].id)).thenReturn(0)
+
+            // Switch to a user who is not locked-out.
+            userRepository.setSelectedUserInfo(USER_INFOS[1])
+            assertThat(underTest.lockoutEndTimestamp).isNull()
+
+            // Switch back to the locked-out user, verify the timestamp is up-to-date.
+            userRepository.setSelectedUserInfo(USER_INFOS[0])
+            assertThat(underTest.lockoutEndTimestamp).isEqualTo(lockoutEndMs)
+
+            // After the lockout expires, null is returned.
+            clock.setElapsedRealtime(lockoutEndMs)
+            assertThat(underTest.lockoutEndTimestamp).isNull()
+        }
+
+    @Test
+    fun hasLockoutOccurred() =
+        testScope.runTest {
+            val hasLockoutOccurred by collectLastValue(underTest.hasLockoutOccurred)
+            assertThat(hasLockoutOccurred).isFalse()
+
+            underTest.reportLockoutStarted(1000)
+            assertThat(hasLockoutOccurred).isTrue()
+
+            clock.setElapsedRealtime(clock.elapsedRealtime() + 60.seconds.inWholeMilliseconds)
+
+            underTest.reportAuthenticationAttempt(isSuccessful = false)
+            assertThat(hasLockoutOccurred).isTrue()
+
+            underTest.reportAuthenticationAttempt(isSuccessful = true)
+            assertThat(hasLockoutOccurred).isFalse()
+        }
+
     private fun setSecurityModeAndDispatchBroadcast(
         securityMode: KeyguardSecurityModel.SecurityMode,
     ) {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/authentication/domain/interactor/AuthenticationInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/authentication/domain/interactor/AuthenticationInteractorTest.kt
index 08cd7ed..c113b37 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/authentication/domain/interactor/AuthenticationInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/authentication/domain/interactor/AuthenticationInteractorTest.kt
@@ -21,14 +21,18 @@
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.authentication.data.repository.FakeAuthenticationRepository
-import com.android.systemui.authentication.shared.model.AuthenticationLockoutModel
-import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
+import com.android.systemui.authentication.shared.model.AuthenticationMethodModel.None
+import com.android.systemui.authentication.shared.model.AuthenticationMethodModel.Password
+import com.android.systemui.authentication.shared.model.AuthenticationMethodModel.Pattern
+import com.android.systemui.authentication.shared.model.AuthenticationMethodModel.Pin
 import com.android.systemui.authentication.shared.model.AuthenticationPatternCoordinate
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.scene.SceneTestUtils
 import com.google.common.truth.Truth.assertThat
+import kotlin.time.Duration.Companion.seconds
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.advanceTimeBy
+import kotlinx.coroutines.test.currentTime
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
@@ -43,21 +47,23 @@
     private val testScope = utils.testScope
     private val underTest = utils.authenticationInteractor()
 
+    private val onAuthenticationResult by
+        testScope.collectLastValue(underTest.onAuthenticationResult)
+    private val failedAuthenticationAttempts by
+        testScope.collectLastValue(underTest.failedAuthenticationAttempts)
+
     @Test
     fun authenticationMethod() =
         testScope.runTest {
             val authMethod by collectLastValue(underTest.authenticationMethod)
             runCurrent()
-            assertThat(authMethod).isEqualTo(AuthenticationMethodModel.Pin)
-            assertThat(underTest.getAuthenticationMethod()).isEqualTo(AuthenticationMethodModel.Pin)
+            assertThat(authMethod).isEqualTo(Pin)
+            assertThat(underTest.getAuthenticationMethod()).isEqualTo(Pin)
 
-            utils.authenticationRepository.setAuthenticationMethod(
-                AuthenticationMethodModel.Password
-            )
+            utils.authenticationRepository.setAuthenticationMethod(Password)
 
-            assertThat(authMethod).isEqualTo(AuthenticationMethodModel.Password)
-            assertThat(underTest.getAuthenticationMethod())
-                .isEqualTo(AuthenticationMethodModel.Password)
+            assertThat(authMethod).isEqualTo(Password)
+            assertThat(underTest.getAuthenticationMethod()).isEqualTo(Password)
         }
 
     @Test
@@ -66,51 +72,45 @@
             val authMethod by collectLastValue(underTest.authenticationMethod)
             runCurrent()
 
-            utils.authenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.None)
+            utils.authenticationRepository.setAuthenticationMethod(None)
 
-            assertThat(authMethod).isEqualTo(AuthenticationMethodModel.None)
-            assertThat(underTest.getAuthenticationMethod())
-                .isEqualTo(AuthenticationMethodModel.None)
+            assertThat(authMethod).isEqualTo(None)
+            assertThat(underTest.getAuthenticationMethod()).isEqualTo(None)
         }
 
     @Test
     fun authenticate_withCorrectPin_succeeds() =
         testScope.runTest {
-            val lockout by collectLastValue(underTest.lockout)
-            utils.authenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.Pin)
+            utils.authenticationRepository.setAuthenticationMethod(Pin)
 
-            assertThat(underTest.authenticate(FakeAuthenticationRepository.DEFAULT_PIN))
-                .isEqualTo(AuthenticationResult.SUCCEEDED)
-            assertThat(lockout).isNull()
-            assertThat(utils.authenticationRepository.lockoutStartedReportCount).isEqualTo(0)
+            assertSucceeded(underTest.authenticate(FakeAuthenticationRepository.DEFAULT_PIN))
         }
 
     @Test
     fun authenticate_withIncorrectPin_fails() =
         testScope.runTest {
-            utils.authenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.Pin)
+            utils.authenticationRepository.setAuthenticationMethod(Pin)
 
-            assertThat(underTest.authenticate(listOf(9, 8, 7, 6, 5, 4)))
-                .isEqualTo(AuthenticationResult.FAILED)
+            assertFailed(underTest.authenticate(listOf(9, 8, 7, 6, 5, 4)))
         }
 
     @Test(expected = IllegalArgumentException::class)
     fun authenticate_withEmptyPin_throwsException() =
         testScope.runTest {
-            utils.authenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.Pin)
+            utils.authenticationRepository.setAuthenticationMethod(Pin)
             underTest.authenticate(listOf())
         }
 
     @Test
     fun authenticate_withCorrectMaxLengthPin_succeeds() =
         testScope.runTest {
-            val pin = List(16) { 9 }
+            val correctMaxLengthPin = List(16) { 9 }
             utils.authenticationRepository.apply {
-                setAuthenticationMethod(AuthenticationMethodModel.Pin)
-                overrideCredential(pin)
+                setAuthenticationMethod(Pin)
+                overrideCredential(correctMaxLengthPin)
             }
 
-            assertThat(underTest.authenticate(pin)).isEqualTo(AuthenticationResult.SUCCEEDED)
+            assertSucceeded(underTest.authenticate(correctMaxLengthPin))
         }
 
     @Test
@@ -122,88 +122,64 @@
             // If the policy changes, there is work to do in SysUI.
             assertThat(DevicePolicyManager.MAX_PASSWORD_LENGTH).isLessThan(17)
 
-            utils.authenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.Pin)
-            assertThat(underTest.authenticate(List(17) { 9 }))
-                .isEqualTo(AuthenticationResult.FAILED)
+            utils.authenticationRepository.setAuthenticationMethod(Pin)
+
+            assertFailed(underTest.authenticate(List(17) { 9 }))
         }
 
     @Test
     fun authenticate_withCorrectPassword_succeeds() =
         testScope.runTest {
-            val lockout by collectLastValue(underTest.lockout)
-            utils.authenticationRepository.setAuthenticationMethod(
-                AuthenticationMethodModel.Password
-            )
+            utils.authenticationRepository.setAuthenticationMethod(Password)
 
-            assertThat(underTest.authenticate("password".toList()))
-                .isEqualTo(AuthenticationResult.SUCCEEDED)
-            assertThat(lockout).isNull()
-            assertThat(utils.authenticationRepository.lockoutStartedReportCount).isEqualTo(0)
+            assertSucceeded(underTest.authenticate("password".toList()))
         }
 
     @Test
     fun authenticate_withIncorrectPassword_fails() =
         testScope.runTest {
-            utils.authenticationRepository.setAuthenticationMethod(
-                AuthenticationMethodModel.Password
-            )
+            utils.authenticationRepository.setAuthenticationMethod(Password)
 
-            assertThat(underTest.authenticate("alohomora".toList()))
-                .isEqualTo(AuthenticationResult.FAILED)
+            assertFailed(underTest.authenticate("alohomora".toList()))
         }
 
     @Test
     fun authenticate_withCorrectPattern_succeeds() =
         testScope.runTest {
-            utils.authenticationRepository.setAuthenticationMethod(
-                AuthenticationMethodModel.Pattern
-            )
+            utils.authenticationRepository.setAuthenticationMethod(Pattern)
 
-            assertThat(underTest.authenticate(FakeAuthenticationRepository.PATTERN))
-                .isEqualTo(AuthenticationResult.SUCCEEDED)
+            assertSucceeded(underTest.authenticate(FakeAuthenticationRepository.PATTERN))
         }
 
     @Test
     fun authenticate_withIncorrectPattern_fails() =
         testScope.runTest {
-            utils.authenticationRepository.setAuthenticationMethod(
-                AuthenticationMethodModel.Pattern
-            )
-
-            assertThat(
-                    underTest.authenticate(
-                        listOf(
-                            AuthenticationPatternCoordinate(x = 2, y = 0),
-                            AuthenticationPatternCoordinate(x = 2, y = 1),
-                            AuthenticationPatternCoordinate(x = 2, y = 2),
-                            AuthenticationPatternCoordinate(x = 1, y = 2),
-                        )
-                    )
+            utils.authenticationRepository.setAuthenticationMethod(Pattern)
+            val wrongPattern =
+                listOf(
+                    AuthenticationPatternCoordinate(x = 2, y = 0),
+                    AuthenticationPatternCoordinate(x = 2, y = 1),
+                    AuthenticationPatternCoordinate(x = 2, y = 2),
+                    AuthenticationPatternCoordinate(x = 1, y = 2),
                 )
-                .isEqualTo(AuthenticationResult.FAILED)
+
+            assertFailed(underTest.authenticate(wrongPattern))
         }
 
     @Test
     fun tryAutoConfirm_withAutoConfirmPinAndShorterPin_returnsNull() =
         testScope.runTest {
             val isAutoConfirmEnabled by collectLastValue(underTest.isAutoConfirmEnabled)
-            val lockout by collectLastValue(underTest.lockout)
             utils.authenticationRepository.apply {
-                setAuthenticationMethod(AuthenticationMethodModel.Pin)
+                setAuthenticationMethod(Pin)
                 setAutoConfirmFeatureEnabled(true)
             }
             assertThat(isAutoConfirmEnabled).isTrue()
+            val shorterPin =
+                FakeAuthenticationRepository.DEFAULT_PIN.toMutableList().apply { removeLast() }
 
-            assertThat(
-                    underTest.authenticate(
-                        FakeAuthenticationRepository.DEFAULT_PIN.toMutableList().apply {
-                            removeLast()
-                        },
-                        tryAutoConfirm = true
-                    )
-                )
-                .isEqualTo(AuthenticationResult.SKIPPED)
-            assertThat(lockout).isNull()
+            assertSkipped(underTest.authenticate(shorterPin, tryAutoConfirm = true))
+            assertThat(underTest.lockoutEndTimestamp).isNull()
             assertThat(utils.authenticationRepository.lockoutStartedReportCount).isEqualTo(0)
         }
 
@@ -212,18 +188,17 @@
         testScope.runTest {
             val isAutoConfirmEnabled by collectLastValue(underTest.isAutoConfirmEnabled)
             utils.authenticationRepository.apply {
-                setAuthenticationMethod(AuthenticationMethodModel.Pin)
+                setAuthenticationMethod(Pin)
                 setAutoConfirmFeatureEnabled(true)
             }
             assertThat(isAutoConfirmEnabled).isTrue()
 
-            assertThat(
-                    underTest.authenticate(
-                        FakeAuthenticationRepository.DEFAULT_PIN.map { it + 1 },
-                        tryAutoConfirm = true
-                    )
-                )
-                .isEqualTo(AuthenticationResult.FAILED)
+            val wrongPin = FakeAuthenticationRepository.DEFAULT_PIN.map { it + 1 }
+
+            assertFailed(
+                underTest.authenticate(wrongPin, tryAutoConfirm = true),
+                assertNoResultEvents = true,
+            )
         }
 
     @Test
@@ -231,18 +206,17 @@
         testScope.runTest {
             val isAutoConfirmEnabled by collectLastValue(underTest.isAutoConfirmEnabled)
             utils.authenticationRepository.apply {
-                setAuthenticationMethod(AuthenticationMethodModel.Pin)
+                setAuthenticationMethod(Pin)
                 setAutoConfirmFeatureEnabled(true)
             }
             assertThat(isAutoConfirmEnabled).isTrue()
 
-            assertThat(
-                    underTest.authenticate(
-                        FakeAuthenticationRepository.DEFAULT_PIN + listOf(7),
-                        tryAutoConfirm = true
-                    )
-                )
-                .isEqualTo(AuthenticationResult.FAILED)
+            val longerPin = FakeAuthenticationRepository.DEFAULT_PIN + listOf(7)
+
+            assertFailed(
+                underTest.authenticate(longerPin, tryAutoConfirm = true),
+                assertNoResultEvents = true,
+            )
         }
 
     @Test
@@ -250,69 +224,54 @@
         testScope.runTest {
             val isAutoConfirmEnabled by collectLastValue(underTest.isAutoConfirmEnabled)
             utils.authenticationRepository.apply {
-                setAuthenticationMethod(AuthenticationMethodModel.Pin)
+                setAuthenticationMethod(Pin)
                 setAutoConfirmFeatureEnabled(true)
             }
             assertThat(isAutoConfirmEnabled).isTrue()
 
-            assertThat(
-                    underTest.authenticate(
-                        FakeAuthenticationRepository.DEFAULT_PIN,
-                        tryAutoConfirm = true
-                    )
-                )
-                .isEqualTo(AuthenticationResult.SUCCEEDED)
+            val correctPin = FakeAuthenticationRepository.DEFAULT_PIN
+
+            assertSucceeded(underTest.authenticate(correctPin, tryAutoConfirm = true))
         }
 
     @Test
     fun tryAutoConfirm_withAutoConfirmCorrectPinButDuringLockout_returnsNull() =
         testScope.runTest {
             val isAutoConfirmEnabled by collectLastValue(underTest.isAutoConfirmEnabled)
-            val isUnlocked by collectLastValue(utils.deviceEntryRepository.isUnlocked)
             val hintedPinLength by collectLastValue(underTest.hintedPinLength)
             utils.authenticationRepository.apply {
-                setAuthenticationMethod(AuthenticationMethodModel.Pin)
+                setAuthenticationMethod(Pin)
                 setAutoConfirmFeatureEnabled(true)
-                setLockoutDuration(42)
+                reportLockoutStarted(42)
             }
 
-            val authResult =
-                underTest.authenticate(
-                    FakeAuthenticationRepository.DEFAULT_PIN,
-                    tryAutoConfirm = true
-                )
+            val correctPin = FakeAuthenticationRepository.DEFAULT_PIN
 
-            assertThat(authResult).isEqualTo(AuthenticationResult.SKIPPED)
+            assertSkipped(underTest.authenticate(correctPin, tryAutoConfirm = true))
             assertThat(isAutoConfirmEnabled).isFalse()
-            assertThat(isUnlocked).isFalse()
             assertThat(hintedPinLength).isNull()
+            assertThat(underTest.lockoutEndTimestamp).isNotNull()
         }
 
     @Test
     fun tryAutoConfirm_withoutAutoConfirmButCorrectPin_returnsNull() =
         testScope.runTest {
             utils.authenticationRepository.apply {
-                setAuthenticationMethod(AuthenticationMethodModel.Pin)
+                setAuthenticationMethod(Pin)
                 setAutoConfirmFeatureEnabled(false)
             }
-            assertThat(
-                    underTest.authenticate(
-                        FakeAuthenticationRepository.DEFAULT_PIN,
-                        tryAutoConfirm = true
-                    )
-                )
-                .isEqualTo(AuthenticationResult.SKIPPED)
+
+            val correctPin = FakeAuthenticationRepository.DEFAULT_PIN
+
+            assertSkipped(underTest.authenticate(correctPin, tryAutoConfirm = true))
         }
 
     @Test
     fun tryAutoConfirm_withoutCorrectPassword_returnsNull() =
         testScope.runTest {
-            utils.authenticationRepository.setAuthenticationMethod(
-                AuthenticationMethodModel.Password
-            )
+            utils.authenticationRepository.setAuthenticationMethod(Password)
 
-            assertThat(underTest.authenticate("password".toList(), tryAutoConfirm = true))
-                .isEqualTo(AuthenticationResult.SKIPPED)
+            assertSkipped(underTest.authenticate("password".toList(), tryAutoConfirm = true))
         }
 
     @Test
@@ -337,7 +296,6 @@
     fun isAutoConfirmEnabled_featureEnabledButDisabledByLockout() =
         testScope.runTest {
             val isAutoConfirmEnabled by collectLastValue(underTest.isAutoConfirmEnabled)
-            val lockout by collectLastValue(underTest.lockout)
             utils.authenticationRepository.setAutoConfirmFeatureEnabled(true)
 
             // The feature is enabled.
@@ -345,92 +303,104 @@
 
             // Make many wrong attempts to trigger lockout.
             repeat(FakeAuthenticationRepository.MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT) {
-                underTest.authenticate(listOf(5, 6, 7)) // Wrong PIN
+                assertFailed(underTest.authenticate(listOf(5, 6, 7))) // Wrong PIN
             }
-            assertThat(lockout).isNotNull()
+            assertThat(underTest.lockoutEndTimestamp).isNotNull()
             assertThat(utils.authenticationRepository.lockoutStartedReportCount).isEqualTo(1)
 
             // Lockout disabled auto-confirm.
             assertThat(isAutoConfirmEnabled).isFalse()
 
             // Move the clock forward one more second, to completely finish the lockout period:
-            advanceTimeBy(FakeAuthenticationRepository.LOCKOUT_DURATION_MS + 1000L)
-            assertThat(lockout).isNull()
+            advanceTimeBy(
+                FakeAuthenticationRepository.LOCKOUT_DURATION_SECONDS.seconds.plus(1.seconds)
+            )
+            assertThat(underTest.lockoutEndTimestamp).isNull()
 
             // Auto-confirm is still disabled, because lockout occurred at least once in this
             // session.
             assertThat(isAutoConfirmEnabled).isFalse()
 
             // Correct PIN and unlocks successfully, resetting the 'session'.
-            assertThat(underTest.authenticate(FakeAuthenticationRepository.DEFAULT_PIN))
-                .isEqualTo(AuthenticationResult.SUCCEEDED)
+            assertSucceeded(underTest.authenticate(FakeAuthenticationRepository.DEFAULT_PIN))
 
             // Auto-confirm is re-enabled.
             assertThat(isAutoConfirmEnabled).isTrue()
-
-            assertThat(utils.authenticationRepository.lockoutStartedReportCount).isEqualTo(1)
         }
 
     @Test
-    fun lockout() =
+    fun failedAuthenticationAttempts() =
         testScope.runTest {
-            val lockout by collectLastValue(underTest.lockout)
-            utils.authenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.Pin)
-            underTest.authenticate(FakeAuthenticationRepository.DEFAULT_PIN)
-            assertThat(lockout).isNull()
+            val failedAuthenticationAttempts by
+                collectLastValue(underTest.failedAuthenticationAttempts)
+
+            utils.authenticationRepository.setAuthenticationMethod(Pin)
+            val correctPin = FakeAuthenticationRepository.DEFAULT_PIN
+
+            assertSucceeded(underTest.authenticate(correctPin))
+            assertThat(failedAuthenticationAttempts).isEqualTo(0)
+
+            // Make many wrong attempts, leading to lockout:
+            repeat(FakeAuthenticationRepository.MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT) { index ->
+                underTest.authenticate(listOf(5, 6, 7)) // Wrong PIN
+                assertThat(failedAuthenticationAttempts).isEqualTo(index + 1)
+            }
+
+            // Correct PIN, but locked out, so doesn't attempt it:
+            assertSkipped(underTest.authenticate(correctPin), assertNoResultEvents = false)
+            assertThat(failedAuthenticationAttempts)
+                .isEqualTo(FakeAuthenticationRepository.MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT)
+
+            // Move the clock forward to finish the lockout period:
+            advanceTimeBy(FakeAuthenticationRepository.LOCKOUT_DURATION_SECONDS.seconds)
+            assertThat(failedAuthenticationAttempts)
+                .isEqualTo(FakeAuthenticationRepository.MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT)
+
+            // Correct PIN and no longer locked out so unlocks successfully:
+            assertSucceeded(underTest.authenticate(correctPin))
+            assertThat(failedAuthenticationAttempts).isEqualTo(0)
+        }
+
+    @Test
+    fun lockoutEndTimestamp() =
+        testScope.runTest {
+            utils.authenticationRepository.setAuthenticationMethod(Pin)
+            val correctPin = FakeAuthenticationRepository.DEFAULT_PIN
+
+            underTest.authenticate(correctPin)
+            assertThat(underTest.lockoutEndTimestamp).isNull()
 
             // Make many wrong attempts, but just shy of what's needed to get locked out:
             repeat(FakeAuthenticationRepository.MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT - 1) {
                 underTest.authenticate(listOf(5, 6, 7)) // Wrong PIN
-                assertThat(lockout).isNull()
+                assertThat(underTest.lockoutEndTimestamp).isNull()
             }
 
             // Make one more wrong attempt, leading to lockout:
             underTest.authenticate(listOf(5, 6, 7)) // Wrong PIN
-            assertThat(lockout)
-                .isEqualTo(
-                    AuthenticationLockoutModel(
-                        failedAttemptCount =
-                            FakeAuthenticationRepository.MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT,
-                        remainingSeconds = FakeAuthenticationRepository.LOCKOUT_DURATION_SECONDS,
-                    )
-                )
+
+            val expectedLockoutEndTimestamp =
+                testScope.currentTime + FakeAuthenticationRepository.LOCKOUT_DURATION_MS
+            assertThat(underTest.lockoutEndTimestamp).isEqualTo(expectedLockoutEndTimestamp)
             assertThat(utils.authenticationRepository.lockoutStartedReportCount).isEqualTo(1)
 
             // Correct PIN, but locked out, so doesn't attempt it:
-            assertThat(underTest.authenticate(FakeAuthenticationRepository.DEFAULT_PIN))
-                .isEqualTo(AuthenticationResult.SKIPPED)
-            assertThat(lockout)
-                .isEqualTo(
-                    AuthenticationLockoutModel(
-                        failedAttemptCount =
-                            FakeAuthenticationRepository.MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT,
-                        remainingSeconds = FakeAuthenticationRepository.LOCKOUT_DURATION_SECONDS,
-                    )
-                )
+            assertSkipped(underTest.authenticate(correctPin), assertNoResultEvents = false)
+            assertThat(underTest.lockoutEndTimestamp).isEqualTo(expectedLockoutEndTimestamp)
 
             // Move the clock forward to ALMOST skip the lockout, leaving one second to go:
-            val lockoutTimeoutSec = FakeAuthenticationRepository.LOCKOUT_DURATION_SECONDS
-            repeat(FakeAuthenticationRepository.LOCKOUT_DURATION_SECONDS - 1) { time ->
-                advanceTimeBy(1000)
-                assertThat(lockout)
-                    .isEqualTo(
-                        AuthenticationLockoutModel(
-                            failedAttemptCount =
-                                FakeAuthenticationRepository.MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT,
-                            remainingSeconds = lockoutTimeoutSec - (time + 1),
-                        )
-                    )
+            repeat(FakeAuthenticationRepository.LOCKOUT_DURATION_SECONDS - 1) {
+                advanceTimeBy(1.seconds)
+                assertThat(underTest.lockoutEndTimestamp).isEqualTo(expectedLockoutEndTimestamp)
             }
 
             // Move the clock forward one more second, to completely finish the lockout period:
-            advanceTimeBy(1000)
-            assertThat(lockout).isNull()
+            advanceTimeBy(1.seconds)
+            assertThat(underTest.lockoutEndTimestamp).isNull()
 
             // Correct PIN and no longer locked out so unlocks successfully:
-            assertThat(underTest.authenticate(FakeAuthenticationRepository.DEFAULT_PIN))
-                .isEqualTo(AuthenticationResult.SUCCEEDED)
-            assertThat(lockout).isNull()
+            assertSucceeded(underTest.authenticate(correctPin))
+            assertThat(underTest.lockoutEndTimestamp).isNull()
         }
 
     @Test
@@ -438,7 +408,7 @@
         testScope.runTest {
             val hintedPinLength by collectLastValue(underTest.hintedPinLength)
             utils.authenticationRepository.apply {
-                setAuthenticationMethod(AuthenticationMethodModel.Pin)
+                setAuthenticationMethod(Pin)
                 setAutoConfirmFeatureEnabled(false)
             }
 
@@ -450,7 +420,7 @@
         testScope.runTest {
             val hintedPinLength by collectLastValue(underTest.hintedPinLength)
             utils.authenticationRepository.apply {
-                setAuthenticationMethod(AuthenticationMethodModel.Pin)
+                setAuthenticationMethod(Pin)
                 overrideCredential(
                     buildList {
                         repeat(utils.authenticationRepository.hintedPinLength - 1) { add(it + 1) }
@@ -467,7 +437,7 @@
         testScope.runTest {
             val hintedPinLength by collectLastValue(underTest.hintedPinLength)
             utils.authenticationRepository.apply {
-                setAuthenticationMethod(AuthenticationMethodModel.Pin)
+                setAuthenticationMethod(Pin)
                 setAutoConfirmFeatureEnabled(true)
                 overrideCredential(
                     buildList {
@@ -484,7 +454,7 @@
         testScope.runTest {
             val hintedPinLength by collectLastValue(underTest.hintedPinLength)
             utils.authenticationRepository.apply {
-                setAuthenticationMethod(AuthenticationMethodModel.Pin)
+                setAuthenticationMethod(Pin)
                 overrideCredential(
                     buildList {
                         repeat(utils.authenticationRepository.hintedPinLength + 1) { add(it + 1) }
@@ -499,18 +469,45 @@
     @Test
     fun authenticate_withTooShortPassword() =
         testScope.runTest {
-            utils.authenticationRepository.setAuthenticationMethod(
-                AuthenticationMethodModel.Password
-            )
-            assertThat(
-                    underTest.authenticate(
-                        buildList {
-                            repeat(utils.authenticationRepository.minPasswordLength - 1) { time ->
-                                add("$time")
-                            }
-                        }
-                    )
-                )
-                .isEqualTo(AuthenticationResult.SKIPPED)
+            utils.authenticationRepository.setAuthenticationMethod(Password)
+
+            val tooShortPassword = buildList {
+                repeat(utils.authenticationRepository.minPasswordLength - 1) { time ->
+                    add("$time")
+                }
+            }
+            assertSkipped(underTest.authenticate(tooShortPassword))
         }
+
+    private fun assertSucceeded(authenticationResult: AuthenticationResult) {
+        assertThat(authenticationResult).isEqualTo(AuthenticationResult.SUCCEEDED)
+        assertThat(onAuthenticationResult).isTrue()
+        assertThat(underTest.lockoutEndTimestamp).isNull()
+        assertThat(utils.authenticationRepository.lockoutStartedReportCount).isEqualTo(0)
+        assertThat(failedAuthenticationAttempts).isEqualTo(0)
+    }
+
+    private fun assertFailed(
+        authenticationResult: AuthenticationResult,
+        assertNoResultEvents: Boolean = false,
+    ) {
+        assertThat(authenticationResult).isEqualTo(AuthenticationResult.FAILED)
+        if (assertNoResultEvents) {
+            assertThat(onAuthenticationResult).isNull()
+        } else {
+            assertThat(onAuthenticationResult).isFalse()
+        }
+    }
+
+    private fun assertSkipped(
+        authenticationResult: AuthenticationResult,
+        assertNoResultEvents: Boolean = true,
+    ) {
+        assertThat(authenticationResult).isEqualTo(AuthenticationResult.SKIPPED)
+        if (assertNoResultEvents) {
+            assertThat(onAuthenticationResult).isNull()
+        } else {
+            assertThat(onAuthenticationResult).isNotNull()
+        }
+    }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractorTest.kt
index 9b1df7c..99c1874 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractorTest.kt
@@ -21,15 +21,15 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.authentication.data.repository.FakeAuthenticationRepository
 import com.android.systemui.authentication.domain.interactor.AuthenticationResult
-import com.android.systemui.authentication.shared.model.AuthenticationLockoutModel
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
 import com.android.systemui.authentication.shared.model.AuthenticationPatternCoordinate
 import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.coroutines.collectValues
 import com.android.systemui.keyguard.domain.interactor.KeyguardFaceAuthInteractor
 import com.android.systemui.res.R
 import com.android.systemui.scene.SceneTestUtils
 import com.google.common.truth.Truth.assertThat
-import kotlin.time.Duration.Companion.milliseconds
+import kotlin.time.Duration.Companion.seconds
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.advanceTimeBy
 import kotlinx.coroutines.test.runCurrent
@@ -79,7 +79,7 @@
             utils.authenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.Pin)
             runCurrent()
             underTest.clearMessage()
-            assertThat(message).isEmpty()
+            assertThat(message).isNull()
 
             underTest.resetMessage()
             assertThat(message).isEqualTo(MESSAGE_ENTER_YOUR_PIN)
@@ -149,7 +149,7 @@
             // Incomplete input.
             assertThat(underTest.authenticate(listOf(1, 2), tryAutoConfirm = true))
                 .isEqualTo(AuthenticationResult.SKIPPED)
-            assertThat(message).isEmpty()
+            assertThat(message).isNull()
 
             // Correct input.
             assertThat(
@@ -159,7 +159,7 @@
                     )
                 )
                 .isEqualTo(AuthenticationResult.SKIPPED)
-            assertThat(message).isEmpty()
+            assertThat(message).isNull()
         }
 
     @Test
@@ -246,57 +246,40 @@
         }
 
     @Test
-    fun lockout() =
+    fun lockoutStarted() =
         testScope.runTest {
-            val lockout by collectLastValue(underTest.lockout)
+            val lockoutStartedEvents by collectValues(underTest.onLockoutStarted)
             val message by collectLastValue(underTest.message)
+
             utils.authenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.Pin)
-            assertThat(lockout).isNull()
+            assertThat(lockoutStartedEvents).isEmpty()
+
+            // Try the wrong PIN repeatedly, until lockout is triggered:
             repeat(FakeAuthenticationRepository.MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT) { times ->
                 // Wrong PIN.
                 assertThat(underTest.authenticate(listOf(6, 7, 8, 9)))
                     .isEqualTo(AuthenticationResult.FAILED)
                 if (times < FakeAuthenticationRepository.MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT - 1) {
-                    assertThat(message).isEqualTo(MESSAGE_WRONG_PIN)
+                    assertThat(lockoutStartedEvents).isEmpty()
+                    assertThat(message).isNotEmpty()
                 }
             }
-            assertThat(lockout)
-                .isEqualTo(
-                    AuthenticationLockoutModel(
-                        failedAttemptCount =
-                            FakeAuthenticationRepository.MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT,
-                        remainingSeconds = FakeAuthenticationRepository.LOCKOUT_DURATION_SECONDS,
-                    )
-                )
-            assertTryAgainMessage(
-                message,
-                FakeAuthenticationRepository.LOCKOUT_DURATION_MS.milliseconds.inWholeSeconds.toInt()
-            )
+            assertThat(authenticationInteractor.lockoutEndTimestamp).isNotNull()
+            assertThat(lockoutStartedEvents.size).isEqualTo(1)
+            assertThat(message).isNull()
 
-            // Correct PIN, but locked out, so doesn't change away from the bouncer scene:
-            assertThat(underTest.authenticate(FakeAuthenticationRepository.DEFAULT_PIN))
-                .isEqualTo(AuthenticationResult.SKIPPED)
-            assertTryAgainMessage(
-                message,
-                FakeAuthenticationRepository.LOCKOUT_DURATION_MS.milliseconds.inWholeSeconds.toInt()
-            )
+            // Advance the time to finish the lockout:
+            advanceTimeBy(FakeAuthenticationRepository.LOCKOUT_DURATION_SECONDS.seconds)
+            assertThat(authenticationInteractor.lockoutEndTimestamp).isNull()
+            assertThat(message).isNull()
+            assertThat(lockoutStartedEvents.size).isEqualTo(1)
 
-            lockout?.remainingSeconds?.let { seconds ->
-                repeat(seconds) { time ->
-                    advanceTimeBy(1000)
-                    val remainingTimeSec = seconds - time - 1
-                    if (remainingTimeSec > 0) {
-                        assertTryAgainMessage(message, remainingTimeSec)
-                    }
-                }
+            // Trigger lockout again:
+            repeat(FakeAuthenticationRepository.MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT) {
+                // Wrong PIN.
+                underTest.authenticate(listOf(6, 7, 8, 9))
             }
-            assertThat(message).isEqualTo("")
-            assertThat(lockout).isNull()
-
-            // Correct PIN and no longer locked out so changes to the Gone scene:
-            assertThat(underTest.authenticate(FakeAuthenticationRepository.DEFAULT_PIN))
-                .isEqualTo(AuthenticationResult.SUCCEEDED)
-            assertThat(lockout).isNull()
+            assertThat(lockoutStartedEvents.size).isEqualTo(2)
         }
 
     @Test
@@ -326,13 +309,6 @@
             verify(keyguardFaceAuthInteractor).onPrimaryBouncerUserInput()
         }
 
-    private fun assertTryAgainMessage(
-        message: String?,
-        time: Int,
-    ) {
-        assertThat(message).isEqualTo("Try again in $time seconds.")
-    }
-
     companion object {
         private const val MESSAGE_ENTER_YOUR_PIN = "Enter your PIN"
         private const val MESSAGE_ENTER_YOUR_PASSWORD = "Enter your password"
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelTest.kt
index 16a9359..4be9b0a 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModelTest.kt
@@ -33,6 +33,7 @@
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.test.advanceTimeBy
+import kotlinx.coroutines.test.currentTime
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
@@ -135,19 +136,47 @@
     fun message() =
         testScope.runTest {
             val message by collectLastValue(underTest.message)
-            val lockout by collectLastValue(bouncerInteractor.lockout)
             utils.authenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.Pin)
             assertThat(message?.isUpdateAnimated).isTrue()
 
             repeat(FakeAuthenticationRepository.MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT) {
-                // Wrong PIN.
-                bouncerInteractor.authenticate(listOf(3, 4, 5, 6))
+                bouncerInteractor.authenticate(WRONG_PIN)
             }
             assertThat(message?.isUpdateAnimated).isFalse()
 
-            lockout?.remainingSeconds?.let { remainingSeconds ->
-                advanceTimeBy(remainingSeconds.seconds.inWholeMilliseconds)
+            val lockoutEndMs = authenticationInteractor.lockoutEndTimestamp ?: 0
+            advanceTimeBy(lockoutEndMs - testScope.currentTime)
+            assertThat(message?.isUpdateAnimated).isTrue()
+        }
+
+    @Test
+    fun lockoutMessage() =
+        testScope.runTest {
+            val authMethodViewModel by collectLastValue(underTest.authMethodViewModel)
+            val message by collectLastValue(underTest.message)
+            utils.authenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.Pin)
+            assertThat(utils.authenticationRepository.lockoutEndTimestamp).isNull()
+            assertThat(authMethodViewModel?.lockoutMessageId).isNotNull()
+
+            repeat(FakeAuthenticationRepository.MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT) { times ->
+                bouncerInteractor.authenticate(WRONG_PIN)
+                if (times < FakeAuthenticationRepository.MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT - 1) {
+                    assertThat(message?.text).isEqualTo(bouncerInteractor.message.value)
+                    assertThat(message?.isUpdateAnimated).isTrue()
+                }
             }
+            val lockoutSeconds = FakeAuthenticationRepository.LOCKOUT_DURATION_SECONDS
+            assertTryAgainMessage(message?.text, lockoutSeconds)
+            assertThat(message?.isUpdateAnimated).isFalse()
+
+            repeat(FakeAuthenticationRepository.LOCKOUT_DURATION_SECONDS) { time ->
+                advanceTimeBy(1.seconds)
+                val remainingSeconds = lockoutSeconds - time - 1
+                if (remainingSeconds > 0) {
+                    assertTryAgainMessage(message?.text, remainingSeconds)
+                }
+            }
+            assertThat(message?.text).isEmpty()
             assertThat(message?.isUpdateAnimated).isTrue()
         }
 
@@ -160,32 +189,30 @@
                         authViewModel?.isInputEnabled ?: emptyFlow()
                     }
                 )
-            val lockout by collectLastValue(bouncerInteractor.lockout)
             utils.authenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.Pin)
             assertThat(isInputEnabled).isTrue()
 
             repeat(FakeAuthenticationRepository.MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT) {
-                // Wrong PIN.
-                bouncerInteractor.authenticate(listOf(3, 4, 5, 6))
+                bouncerInteractor.authenticate(WRONG_PIN)
             }
             assertThat(isInputEnabled).isFalse()
 
-            lockout?.remainingSeconds?.let { remainingSeconds ->
-                advanceTimeBy(remainingSeconds.seconds.inWholeMilliseconds)
-            }
+            val lockoutEndMs = authenticationInteractor.lockoutEndTimestamp ?: 0
+            advanceTimeBy(lockoutEndMs - testScope.currentTime)
             assertThat(isInputEnabled).isTrue()
         }
 
     @Test
     fun dialogMessage() =
         testScope.runTest {
+            val authMethodViewModel by collectLastValue(underTest.authMethodViewModel)
             val dialogMessage by collectLastValue(underTest.dialogMessage)
             utils.authenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.Pin)
+            assertThat(authMethodViewModel?.lockoutMessageId).isNotNull()
 
             repeat(FakeAuthenticationRepository.MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT) {
-                // Wrong PIN.
                 assertThat(dialogMessage).isNull()
-                bouncerInteractor.authenticate(listOf(3, 4, 5, 6))
+                bouncerInteractor.authenticate(WRONG_PIN)
             }
             assertThat(dialogMessage).isNotEmpty()
 
@@ -241,4 +268,15 @@
             AuthenticationMethodModel.Sim,
         )
     }
+
+    private fun assertTryAgainMessage(
+        message: String?,
+        time: Int,
+    ) {
+        assertThat(message).isEqualTo("Try again in $time seconds.")
+    }
+
+    companion object {
+        private val WRONG_PIN = FakeAuthenticationRepository.DEFAULT_PIN.map { it + 1 }
+    }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PasswordBouncerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PasswordBouncerViewModelTest.kt
index 6d6baa5..64e6e57 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PasswordBouncerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PasswordBouncerViewModelTest.kt
@@ -19,7 +19,6 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.authentication.shared.model.AuthenticationLockoutModel
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.coroutines.collectValues
@@ -28,6 +27,7 @@
 import com.android.systemui.scene.shared.model.SceneKey
 import com.android.systemui.scene.shared.model.SceneModel
 import com.google.common.truth.Truth.assertThat
+import kotlin.time.Duration.Companion.seconds
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.asStateFlow
@@ -45,11 +45,7 @@
 
     private val utils = SceneTestUtils(this)
     private val testScope = utils.testScope
-    private val authenticationRepository = utils.authenticationRepository
-    private val authenticationInteractor =
-        utils.authenticationInteractor(
-            repository = authenticationRepository,
-        )
+    private val authenticationInteractor = utils.authenticationInteractor()
     private val sceneInteractor = utils.sceneInteractor()
     private val bouncerInteractor =
         utils.bouncerInteractor(
@@ -61,12 +57,13 @@
             authenticationInteractor = authenticationInteractor,
             actionButtonInteractor = utils.bouncerActionButtonInteractor(),
         )
+    private val isInputEnabled = MutableStateFlow(true)
 
     private val underTest =
         PasswordBouncerViewModel(
             viewModelScope = testScope.backgroundScope,
             interactor = bouncerInteractor,
-            isInputEnabled = MutableStateFlow(true).asStateFlow(),
+            isInputEnabled.asStateFlow(),
         )
 
     @Before
@@ -123,8 +120,7 @@
     @Test
     fun onAuthenticateKeyPressed_whenCorrect() =
         testScope.runTest {
-            val authResult by
-                collectLastValue(authenticationInteractor.authenticationChallengeResult)
+            val authResult by collectLastValue(authenticationInteractor.onAuthenticationResult)
             lockDeviceAndOpenPasswordBouncer()
 
             underTest.onPasswordInputChanged("password")
@@ -169,8 +165,7 @@
     @Test
     fun onAuthenticateKeyPressed_correctAfterWrong() =
         testScope.runTest {
-            val authResult by
-                collectLastValue(authenticationInteractor.authenticationChallengeResult)
+            val authResult by collectLastValue(authenticationInteractor.onAuthenticationResult)
             val message by collectLastValue(bouncerViewModel.message)
             val password by collectLastValue(underTest.password)
             lockDeviceAndOpenPasswordBouncer()
@@ -333,19 +328,15 @@
     ) {
         if (isLockedOut) {
             repeat(failedAttemptCount) {
-                authenticationRepository.reportAuthenticationAttempt(false)
+                utils.authenticationRepository.reportAuthenticationAttempt(false)
             }
-            val remainingTimeSeconds = 30
-            authenticationRepository.setLockoutDuration(remainingTimeSeconds * 1000)
-            authenticationRepository.lockout.value =
-                AuthenticationLockoutModel(
-                    failedAttemptCount = failedAttemptCount,
-                    remainingSeconds = remainingTimeSeconds,
-                )
+            utils.authenticationRepository.reportLockoutStarted(
+                30.seconds.inWholeMilliseconds.toInt()
+            )
         } else {
-            authenticationRepository.reportAuthenticationAttempt(true)
-            authenticationRepository.lockout.value = null
+            utils.authenticationRepository.reportAuthenticationAttempt(true)
         }
+        isInputEnabled.value = !isLockedOut
 
         runCurrent()
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PatternBouncerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PatternBouncerViewModelTest.kt
index 8971423..ed76099 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PatternBouncerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PatternBouncerViewModelTest.kt
@@ -111,8 +111,7 @@
     @Test
     fun onDragEnd_whenCorrect() =
         testScope.runTest {
-            val authResult by
-                collectLastValue(authenticationInteractor.authenticationChallengeResult)
+            val authResult by collectLastValue(authenticationInteractor.onAuthenticationResult)
             val selectedDots by collectLastValue(underTest.selectedDots)
             val currentDot by collectLastValue(underTest.currentDot)
             lockDeviceAndOpenPatternBouncer()
@@ -334,8 +333,7 @@
     @Test
     fun onDragEnd_correctAfterWrong() =
         testScope.runTest {
-            val authResult by
-                collectLastValue(authenticationInteractor.authenticationChallengeResult)
+            val authResult by collectLastValue(authenticationInteractor.onAuthenticationResult)
             val message by collectLastValue(bouncerViewModel.message)
             val selectedDots by collectLastValue(underTest.selectedDots)
             val currentDot by collectLastValue(underTest.currentDot)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PinBouncerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PinBouncerViewModelTest.kt
index c30e405..db98d76 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PinBouncerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/PinBouncerViewModelTest.kt
@@ -201,8 +201,7 @@
     @Test
     fun onAuthenticateButtonClicked_whenCorrect() =
         testScope.runTest {
-            val authResult by
-                collectLastValue(authenticationInteractor.authenticationChallengeResult)
+            val authResult by collectLastValue(authenticationInteractor.onAuthenticationResult)
             lockDeviceAndOpenPinBouncer()
 
             FakeAuthenticationRepository.DEFAULT_PIN.forEach(underTest::onPinButtonClicked)
@@ -236,8 +235,7 @@
     @Test
     fun onAuthenticateButtonClicked_correctAfterWrong() =
         testScope.runTest {
-            val authResult by
-                collectLastValue(authenticationInteractor.authenticationChallengeResult)
+            val authResult by collectLastValue(authenticationInteractor.onAuthenticationResult)
             val message by collectLastValue(bouncerViewModel.message)
             val pin by collectLastValue(underTest.pinInput.map { it.getPin() })
             lockDeviceAndOpenPinBouncer()
@@ -265,8 +263,7 @@
     fun onAutoConfirm_whenCorrect() =
         testScope.runTest {
             utils.authenticationRepository.setAutoConfirmFeatureEnabled(true)
-            val authResult by
-                collectLastValue(authenticationInteractor.authenticationChallengeResult)
+            val authResult by collectLastValue(authenticationInteractor.onAuthenticationResult)
             lockDeviceAndOpenPinBouncer()
 
             FakeAuthenticationRepository.DEFAULT_PIN.forEach(underTest::onPinButtonClicked)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryRepositoryTest.kt
index d3049d9..565049b 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryRepositoryTest.kt
@@ -99,7 +99,7 @@
         }
 
     @Test
-    fun reportSuccessfulAuthentication_shouldUpdateIsUnlocked() =
+    fun reportSuccessfulAuthentication_updatesIsUnlocked() =
         testScope.runTest {
             val isUnlocked by collectLastValue(underTest.isUnlocked)
             assertThat(isUnlocked).isFalse()
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractorTest.kt
index 910097e..ea19cb7 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractorTest.kt
@@ -19,6 +19,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.authentication.data.repository.FakeAuthenticationRepository
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.coroutines.collectValues
@@ -346,12 +347,14 @@
         }
 
     @Test
-    fun successfulAuthenticationChallengeAttempt_updatedIsUnlockedState() =
+    fun successfulAuthenticationChallengeAttempt_updatesIsUnlockedState() =
         testScope.runTest {
             val isUnlocked by collectLastValue(underTest.isUnlocked)
+            utils.authenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.Pin)
+            utils.deviceEntryRepository.setLockscreenEnabled(true)
             assertThat(isUnlocked).isFalse()
 
-            utils.authenticationRepository.reportAuthenticationAttempt(true)
+            authenticationInteractor.authenticate(FakeAuthenticationRepository.DEFAULT_PIN)
 
             assertThat(isUnlocked).isTrue()
         }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
index c110de9..224903f 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
@@ -265,8 +265,8 @@
                 falsingCollector = utils.falsingCollector(),
                 powerInteractor = powerInteractor,
                 bouncerInteractor = bouncerInteractor,
-                simBouncerInteractor = utils.simBouncerInteractor,
-                authenticationInteractor = utils.authenticationInteractor()
+                simBouncerInteractor = dagger.Lazy { utils.simBouncerInteractor },
+                authenticationInteractor = dagger.Lazy { utils.authenticationInteractor() }
             )
         startable.start()
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt
index 61d55f0..2e4986d 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt
@@ -90,8 +90,8 @@
             falsingCollector = falsingCollector,
             powerInteractor = powerInteractor,
             bouncerInteractor = bouncerInteractor,
-            simBouncerInteractor = utils.simBouncerInteractor,
-            authenticationInteractor = authenticationInteractor,
+            simBouncerInteractor = dagger.Lazy { utils.simBouncerInteractor },
+            authenticationInteractor = dagger.Lazy { authenticationInteractor },
         )
 
     @Test
diff --git a/packages/SystemUI/res/layout/widget_picker.xml b/packages/SystemUI/res/layout/widget_picker.xml
deleted file mode 100644
index 21dc224..0000000
--- a/packages/SystemUI/res/layout/widget_picker.xml
+++ /dev/null
@@ -1,30 +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.
-  -->
-
-<HorizontalScrollView
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    android:layout_width="match_parent"
-    android:layout_height="match_parent">
-
-    <LinearLayout
-        android:id="@+id/widgets_container"
-        android:layout_width="wrap_content"
-        android:layout_height="match_parent"
-        android:gravity="center_vertical"
-        android:orientation="horizontal">
-    </LinearLayout>
-
-</HorizontalScrollView>
diff --git a/packages/SystemUI/src/com/android/systemui/authentication/data/repository/AuthenticationRepository.kt b/packages/SystemUI/src/com/android/systemui/authentication/data/repository/AuthenticationRepository.kt
index fda23b7f..f2b55f4 100644
--- a/packages/SystemUI/src/com/android/systemui/authentication/data/repository/AuthenticationRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/authentication/data/repository/AuthenticationRepository.kt
@@ -24,13 +24,13 @@
 import com.android.internal.widget.LockPatternUtils
 import com.android.internal.widget.LockscreenCredential
 import com.android.keyguard.KeyguardSecurityModel
-import com.android.systemui.authentication.shared.model.AuthenticationLockoutModel
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
 import com.android.systemui.authentication.shared.model.AuthenticationResultModel
 import com.android.systemui.broadcast.BroadcastDispatcher
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.scene.shared.flag.SceneContainerFlags
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionsRepository
 import com.android.systemui.user.data.repository.UserRepository
 import com.android.systemui.util.kotlin.pairwise
@@ -43,9 +43,7 @@
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.SharedFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
@@ -60,12 +58,6 @@
 /** Defines interface for classes that can access authentication-related application state. */
 interface AuthenticationRepository {
     /**
-     * Emits the result whenever a PIN/Pattern/Password security challenge is attempted by the user
-     * in order to unlock the device.
-     */
-    val authenticationChallengeResult: SharedFlow<Boolean>
-
-    /**
      * The exact length a PIN should be for us to enable PIN length hinting.
      *
      * A PIN that's shorter or longer than this is not eligible for the UI to render hints showing
@@ -80,16 +72,6 @@
     val isPatternVisible: StateFlow<Boolean>
 
     /**
-     * The current authentication lockout (aka "throttling") state, set when the user has to wait
-     * before being able to try another authentication attempt. `null` indicates throttling isn't
-     * active.
-     */
-    val lockout: MutableStateFlow<AuthenticationLockoutModel?>
-
-    /** Whether throttling has occurred at least once since the last successful authentication. */
-    val hasLockoutOccurred: MutableStateFlow<Boolean>
-
-    /**
      * Whether the auto confirm feature is enabled for the currently-selected user.
      *
      * Note that the length of the PIN is also important to take into consideration, please see
@@ -98,6 +80,28 @@
     val isAutoConfirmFeatureEnabled: StateFlow<Boolean>
 
     /**
+     * The number of failed authentication attempts for the selected user since their last
+     * successful authentication.
+     */
+    val failedAuthenticationAttempts: StateFlow<Int>
+
+    /**
+     * Timestamp for when the current lockout (aka "throttling") will end, allowing the user to
+     * attempt authentication again. Returns `null` if no lockout is active.
+     *
+     * Note that the value is in milliseconds and matches [SystemClock.elapsedRealtime].
+     *
+     * Also note that the value may change when the selected user is changed.
+     */
+    val lockoutEndTimestamp: Long?
+
+    /**
+     * Whether lockout has occurred at least once since the last successful authentication of any
+     * user.
+     */
+    val hasLockoutOccurred: StateFlow<Boolean>
+
+    /**
      * The currently-configured authentication method. This determines how the authentication
      * challenge needs to be completed in order to unlock an otherwise locked device.
      *
@@ -142,23 +146,6 @@
     /** Reports that the user has entered a temporary device lockout (throttling). */
     suspend fun reportLockoutStarted(durationMs: Int)
 
-    /** Returns the current number of failed authentication attempts. */
-    suspend fun getFailedAuthenticationAttemptCount(): Int
-
-    /**
-     * Returns the timestamp for when the current lockout will end, allowing the user to attempt
-     * authentication again.
-     *
-     * Note that this is in milliseconds and it matches [SystemClock.elapsedRealtime].
-     */
-    suspend fun getLockoutEndTimestamp(): Long
-
-    /**
-     * Sets the lockout timeout duration (time during which the user should not be allowed to
-     * attempt authentication).
-     */
-    suspend fun setLockoutDuration(durationMs: Int)
-
     /**
      * Checks the given [LockscreenCredential] to see if it's correct, returning an
      * [AuthenticationResultModel] representing what happened.
@@ -172,6 +159,8 @@
 constructor(
     @Application private val applicationScope: CoroutineScope,
     @Background private val backgroundDispatcher: CoroutineDispatcher,
+    flags: SceneContainerFlags,
+    private val clock: SystemClock,
     private val getSecurityMode: Function<Int, KeyguardSecurityModel.SecurityMode>,
     private val userRepository: UserRepository,
     private val lockPatternUtils: LockPatternUtils,
@@ -179,8 +168,6 @@
     mobileConnectionsRepository: MobileConnectionsRepository,
 ) : AuthenticationRepository {
 
-    override val authenticationChallengeResult = MutableSharedFlow<Boolean>()
-
     override val hintedPinLength: Int = 6
 
     override val isPatternVisible: StateFlow<Boolean> =
@@ -189,10 +176,6 @@
             getFreshValue = lockPatternUtils::isVisiblePatternEnabled,
         )
 
-    override val lockout: MutableStateFlow<AuthenticationLockoutModel?> = MutableStateFlow(null)
-
-    override val hasLockoutOccurred: MutableStateFlow<Boolean> = MutableStateFlow(false)
-
     override val isAutoConfirmFeatureEnabled: StateFlow<Boolean> =
         refreshingFlow(
             initialValue = false,
@@ -234,6 +217,31 @@
             getFreshValue = { userId -> lockPatternUtils.isPinEnhancedPrivacyEnabled(userId) },
         )
 
+    private val _failedAuthenticationAttempts = MutableStateFlow(0)
+    override val failedAuthenticationAttempts: StateFlow<Int> =
+        _failedAuthenticationAttempts.asStateFlow()
+
+    override val lockoutEndTimestamp: Long?
+        get() =
+            lockPatternUtils.getLockoutAttemptDeadline(selectedUserId).takeIf {
+                clock.elapsedRealtime() < it
+            }
+
+    private val _hasLockoutOccurred = MutableStateFlow(false)
+    override val hasLockoutOccurred: StateFlow<Boolean> = _hasLockoutOccurred.asStateFlow()
+
+    init {
+        if (flags.isEnabled()) {
+            // Hydrate failedAuthenticationAttempts initially and whenever the selected user
+            // changes.
+            applicationScope.launch {
+                userRepository.selectedUserInfo.collect {
+                    _failedAuthenticationAttempts.value = getFailedAuthenticationAttemptCount()
+                }
+            }
+        }
+    }
+
     override suspend fun getAuthenticationMethod(): AuthenticationMethodModel {
         return withContext(backgroundDispatcher) {
             blockingAuthenticationMethodInternal(selectedUserId)
@@ -248,35 +256,20 @@
         withContext(backgroundDispatcher) {
             if (isSuccessful) {
                 lockPatternUtils.reportSuccessfulPasswordAttempt(selectedUserId)
+                _hasLockoutOccurred.value = false
             } else {
                 lockPatternUtils.reportFailedPasswordAttempt(selectedUserId)
             }
-            authenticationChallengeResult.emit(isSuccessful)
+            _failedAuthenticationAttempts.value = getFailedAuthenticationAttemptCount()
         }
     }
 
     override suspend fun reportLockoutStarted(durationMs: Int) {
-        return withContext(backgroundDispatcher) {
+        lockPatternUtils.setLockoutAttemptDeadline(selectedUserId, durationMs)
+        withContext(backgroundDispatcher) {
             lockPatternUtils.reportPasswordLockout(durationMs, selectedUserId)
         }
-    }
-
-    override suspend fun getFailedAuthenticationAttemptCount(): Int {
-        return withContext(backgroundDispatcher) {
-            lockPatternUtils.getCurrentFailedPasswordAttempts(selectedUserId)
-        }
-    }
-
-    override suspend fun getLockoutEndTimestamp(): Long {
-        return withContext(backgroundDispatcher) {
-            lockPatternUtils.getLockoutAttemptDeadline(selectedUserId)
-        }
-    }
-
-    override suspend fun setLockoutDuration(durationMs: Int) {
-        withContext(backgroundDispatcher) {
-            lockPatternUtils.setLockoutAttemptDeadline(selectedUserId, durationMs)
-        }
+        _hasLockoutOccurred.value = true
     }
 
     override suspend fun checkCredential(
@@ -292,6 +285,12 @@
         }
     }
 
+    private suspend fun getFailedAuthenticationAttemptCount(): Int {
+        return withContext(backgroundDispatcher) {
+            lockPatternUtils.getCurrentFailedPasswordAttempts(selectedUserId)
+        }
+    }
+
     private val selectedUserId: Int
         get() = userRepository.getSelectedUserInfo().id
 
diff --git a/packages/SystemUI/src/com/android/systemui/authentication/domain/interactor/AuthenticationInteractor.kt b/packages/SystemUI/src/com/android/systemui/authentication/domain/interactor/AuthenticationInteractor.kt
index 797154e..c85ffe6 100644
--- a/packages/SystemUI/src/com/android/systemui/authentication/domain/interactor/AuthenticationInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/authentication/domain/interactor/AuthenticationInteractor.kt
@@ -16,36 +16,25 @@
 
 package com.android.systemui.authentication.domain.interactor
 
-import com.android.app.tracing.TraceUtils.Companion.withContext
 import com.android.internal.widget.LockPatternView
 import com.android.internal.widget.LockscreenCredential
 import com.android.systemui.authentication.data.repository.AuthenticationRepository
-import com.android.systemui.authentication.shared.model.AuthenticationLockoutModel
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
 import com.android.systemui.authentication.shared.model.AuthenticationPatternCoordinate
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
-import com.android.systemui.dagger.qualifiers.Background
-import com.android.systemui.user.data.repository.UserRepository
 import com.android.systemui.util.time.SystemClock
 import javax.inject.Inject
-import kotlin.math.ceil
-import kotlin.math.max
-import kotlin.time.Duration.Companion.seconds
-import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Job
-import kotlinx.coroutines.async
-import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.SharedFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asSharedFlow
 import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
-import kotlinx.coroutines.launch
 
 /**
  * Hosts application business logic related to user authentication.
@@ -59,10 +48,7 @@
 @Inject
 constructor(
     @Application private val applicationScope: CoroutineScope,
-    @Background private val backgroundDispatcher: CoroutineDispatcher,
     private val repository: AuthenticationRepository,
-    private val userRepository: UserRepository,
-    private val clock: SystemClock,
 ) {
     /**
      * The currently-configured authentication method. This determines how the authentication
@@ -85,13 +71,6 @@
     val authenticationMethod: Flow<AuthenticationMethodModel> = repository.authenticationMethod
 
     /**
-     * The current authentication lockout (aka "throttling") state, set when the user has to wait
-     * before being able to try another authentication attempt. `null` indicates lockout isn't
-     * active.
-     */
-    val lockout: StateFlow<AuthenticationLockoutModel?> = repository.lockout
-
-    /**
      * Whether the auto confirm feature is enabled for the currently-selected user.
      *
      * Note that the length of the PIN is also important to take into consideration, please see
@@ -130,26 +109,35 @@
     /** Whether the pattern should be visible for the currently-selected user. */
     val isPatternVisible: StateFlow<Boolean> = repository.isPatternVisible
 
+    private val _onAuthenticationResult = MutableSharedFlow<Boolean>()
     /**
      * Emits the outcome (successful or unsuccessful) whenever a PIN/Pattern/Password security
      * challenge is attempted by the user in order to unlock the device.
      */
-    val authenticationChallengeResult: SharedFlow<Boolean> =
-        repository.authenticationChallengeResult
+    val onAuthenticationResult: SharedFlow<Boolean> = _onAuthenticationResult.asSharedFlow()
 
     /** Whether the "enhanced PIN privacy" setting is enabled for the current user. */
     val isPinEnhancedPrivacyEnabled: StateFlow<Boolean> = repository.isPinEnhancedPrivacyEnabled
 
-    private var lockoutCountdownJob: Job? = null
+    /**
+     * The number of failed authentication attempts for the selected user since the last successful
+     * authentication.
+     */
+    val failedAuthenticationAttempts: StateFlow<Int> = repository.failedAuthenticationAttempts
 
-    init {
-        applicationScope.launch {
-            userRepository.selectedUserInfo
-                .map { it.id }
-                .distinctUntilChanged()
-                .collect { onSelectedUserChanged() }
-        }
-    }
+    /**
+     * Timestamp for when the current lockout (aka "throttling") will end, allowing the user to
+     * attempt authentication again. Returns `null` if no lockout is active.
+     *
+     * To be notified whenever a lockout is started, the caller should subscribe to
+     * [onAuthenticationResult].
+     *
+     * Note that the value is in milliseconds and matches [SystemClock.elapsedRealtime].
+     *
+     * Also note that the value may change when the selected user is changed.
+     */
+    val lockoutEndTimestamp: Long?
+        get() = repository.lockoutEndTimestamp
 
     /**
      * Returns the currently-configured authentication method. This determines how the
@@ -190,7 +178,7 @@
         val skipCheck =
             when {
                 // Lockout is active, the UI layer should not have called this; skip the attempt.
-                lockout.value != null -> true
+                repository.lockoutEndTimestamp != null -> true
                 // The input is too short; skip the attempt.
                 input.isTooShort(authMethod) -> true
                 // Auto-confirm attempt when the feature is not enabled; skip the attempt.
@@ -211,27 +199,16 @@
         credential.zeroize()
 
         if (authenticationResult.isSuccessful || !tryAutoConfirm) {
-            repository.reportAuthenticationAttempt(
-                isSuccessful = authenticationResult.isSuccessful,
-            )
+            repository.reportAuthenticationAttempt(authenticationResult.isSuccessful)
         }
 
         // Check if lockout should start and, if so, kick off the countdown:
         if (!authenticationResult.isSuccessful && authenticationResult.lockoutDurationMs > 0) {
-            repository.apply {
-                setLockoutDuration(durationMs = authenticationResult.lockoutDurationMs)
-                reportLockoutStarted(durationMs = authenticationResult.lockoutDurationMs)
-                hasLockoutOccurred.value = true
-            }
-            startLockoutCountdown()
+            repository.reportLockoutStarted(authenticationResult.lockoutDurationMs)
         }
 
-        if (authenticationResult.isSuccessful) {
-            // Since authentication succeeded, refresh lockout to make sure the state is completely
-            // reflecting the upstream source of truth.
-            refreshLockout()
-
-            repository.hasLockoutOccurred.value = false
+        if (authenticationResult.isSuccessful || !tryAutoConfirm) {
+            _onAuthenticationResult.emit(authenticationResult.isSuccessful)
         }
 
         return if (authenticationResult.isSuccessful) {
@@ -249,54 +226,6 @@
         }
     }
 
-    /** Starts refreshing the lockout state every second. */
-    private suspend fun startLockoutCountdown() {
-        cancelLockoutCountdown()
-        lockoutCountdownJob =
-            applicationScope.launch {
-                while (refreshLockout()) {
-                    delay(1.seconds.inWholeMilliseconds)
-                }
-            }
-    }
-
-    /** Cancels any lockout state countdown started in [startLockoutCountdown]. */
-    private fun cancelLockoutCountdown() {
-        lockoutCountdownJob?.cancel()
-        lockoutCountdownJob = null
-    }
-
-    /** Notifies that the currently-selected user has changed. */
-    private suspend fun onSelectedUserChanged() {
-        cancelLockoutCountdown()
-        if (refreshLockout()) {
-            startLockoutCountdown()
-        }
-    }
-
-    /**
-     * Refreshes the lockout state, hydrating the repository with the latest state.
-     *
-     * @return Whether lockout is active or not.
-     */
-    private suspend fun refreshLockout(): Boolean {
-        withContext("$TAG#refreshLockout", backgroundDispatcher) {
-            val failedAttemptCount = async { repository.getFailedAuthenticationAttemptCount() }
-            val deadline = async { repository.getLockoutEndTimestamp() }
-            val remainingMs = max(0, deadline.await() - clock.elapsedRealtime())
-            repository.lockout.value =
-                if (remainingMs > 0) {
-                    AuthenticationLockoutModel(
-                        failedAttemptCount = failedAttemptCount.await(),
-                        remainingSeconds = ceil(remainingMs / 1000f).toInt(),
-                    )
-                } else {
-                    null // Lockout ended.
-                }
-        }
-        return repository.lockout.value != null
-    }
-
     private fun AuthenticationMethodModel.createCredential(
         input: List<Any>
     ): LockscreenCredential? {
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractor.kt b/packages/SystemUI/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractor.kt
index 724c0fe..1095abe 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/domain/interactor/BouncerInteractor.kt
@@ -19,7 +19,6 @@
 import android.content.Context
 import com.android.systemui.authentication.domain.interactor.AuthenticationInteractor
 import com.android.systemui.authentication.domain.interactor.AuthenticationResult
-import com.android.systemui.authentication.shared.model.AuthenticationLockoutModel
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
 import com.android.systemui.bouncer.data.repository.BouncerRepository
 import com.android.systemui.classifier.FalsingClassifier
@@ -29,17 +28,15 @@
 import com.android.systemui.keyguard.domain.interactor.KeyguardFaceAuthInteractor
 import com.android.systemui.power.domain.interactor.PowerInteractor
 import com.android.systemui.res.R
-import com.android.systemui.scene.shared.flag.SceneContainerFlags
-import com.android.systemui.util.kotlin.pairwise
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.async
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.SharedFlow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.flow.filter
+import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.launch
 
 /** Encapsulates business logic and application state accessing use-cases. */
@@ -52,33 +49,12 @@
     private val repository: BouncerRepository,
     private val authenticationInteractor: AuthenticationInteractor,
     private val keyguardFaceAuthInteractor: KeyguardFaceAuthInteractor,
-    flags: SceneContainerFlags,
     private val falsingInteractor: FalsingInteractor,
     private val powerInteractor: PowerInteractor,
     private val simBouncerInteractor: SimBouncerInteractor,
 ) {
-
-    /** The user-facing message to show in the bouncer. */
-    val message: StateFlow<String?> =
-        combine(repository.message, authenticationInteractor.lockout) { message, lockout ->
-                messageOrLockoutMessage(message, lockout)
-            }
-            .stateIn(
-                scope = applicationScope,
-                started = SharingStarted.WhileSubscribed(),
-                initialValue =
-                    messageOrLockoutMessage(
-                        repository.message.value,
-                        authenticationInteractor.lockout.value,
-                    )
-            )
-
-    /**
-     * The current authentication lockout (aka "throttling") state, set when the user has to wait
-     * before being able to try another authentication attempt. `null` indicates lockout isn't
-     * active.
-     */
-    val lockout: StateFlow<AuthenticationLockoutModel?> = authenticationInteractor.lockout
+    /** The user-facing message to show in the bouncer when lockout is not active. */
+    val message: StateFlow<String?> = repository.message
 
     /** Whether the auto confirm feature is enabled for the currently-selected user. */
     val isAutoConfirmEnabled: StateFlow<Boolean> = authenticationInteractor.isAutoConfirmEnabled
@@ -101,18 +77,13 @@
     /** Emits a [Unit] each time the IME (keyboard) is hidden by the user. */
     val onImeHiddenByUser: SharedFlow<Unit> = _onImeHiddenByUser
 
-    init {
-        if (flags.isEnabled()) {
-            // Clear the message if moved from locked-out to no-longer locked-out.
-            applicationScope.launch {
-                lockout.pairwise().collect { (previous, current) ->
-                    if (previous != null && current == null) {
-                        clearMessage()
-                    }
-                }
+    /** Emits a [Unit] each time a lockout is started for the selected user. */
+    val onLockoutStarted: Flow<Unit> =
+        authenticationInteractor.onAuthenticationResult
+            .filter { successfullyAuthenticated ->
+                !successfullyAuthenticated && authenticationInteractor.lockoutEndTimestamp != null
             }
-        }
-    }
+            .map {}
 
     /** Notifies that the user has places down a pointer, not necessarily dragging just yet. */
     fun onDown() {
@@ -188,7 +159,7 @@
         }
 
         if (authenticationInteractor.getAuthenticationMethod() == AuthenticationMethodModel.Sim) {
-            // We authenticate sim in SimInteractor
+            // SIM is authenticated in SimBouncerInteractor.
             return AuthenticationResult.SKIPPED
         }
 
@@ -196,18 +167,20 @@
         // view-models, whose lifecycle (and thus scope) is shorter than this interactor.
         // This allows the task to continue running properly even when the calling scope has been
         // cancelled.
-        return applicationScope
-            .async {
-                val authResult = authenticationInteractor.authenticate(input, tryAutoConfirm)
-                if (
-                    authResult == AuthenticationResult.FAILED ||
-                        (authResult == AuthenticationResult.SKIPPED && !tryAutoConfirm)
-                ) {
-                    showErrorMessage()
-                }
-                authResult
-            }
-            .await()
+        val authResult =
+            applicationScope
+                .async { authenticationInteractor.authenticate(input, tryAutoConfirm) }
+                .await()
+
+        if (authenticationInteractor.lockoutEndTimestamp != null) {
+            clearMessage()
+        } else if (
+            authResult == AuthenticationResult.FAILED ||
+                (authResult == AuthenticationResult.SKIPPED && !tryAutoConfirm)
+        ) {
+            showErrorMessage()
+        }
+        return authResult
     }
 
     /**
@@ -250,19 +223,4 @@
             else -> ""
         }
     }
-
-    private fun messageOrLockoutMessage(
-        message: String?,
-        lockoutModel: AuthenticationLockoutModel?,
-    ): String {
-        return when {
-            lockoutModel != null ->
-                applicationContext.getString(
-                    com.android.internal.R.string.lockscreen_too_many_failed_attempts_countdown,
-                    lockoutModel.remainingSeconds,
-                )
-            message != null -> message
-            else -> ""
-        }
-    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModel.kt b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModel.kt
index 4b14343..be6cf85 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerViewModel.kt
@@ -19,6 +19,7 @@
 import android.content.Context
 import android.graphics.Bitmap
 import androidx.core.graphics.drawable.toBitmap
+import com.android.internal.R
 import com.android.systemui.authentication.domain.interactor.AuthenticationInteractor
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
 import com.android.systemui.bouncer.domain.interactor.BouncerActionButtonInteractor
@@ -34,19 +35,24 @@
 import com.android.systemui.user.ui.viewmodel.UserActionViewModel
 import com.android.systemui.user.ui.viewmodel.UserSwitcherViewModel
 import com.android.systemui.user.ui.viewmodel.UserViewModel
+import com.android.systemui.util.time.SystemClock
 import dagger.Module
 import dagger.Provides
+import kotlin.math.ceil
+import kotlin.math.max
+import kotlin.time.Duration.Companion.seconds
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.cancel
+import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.job
@@ -58,13 +64,14 @@
     @Application private val applicationScope: CoroutineScope,
     @Main private val mainDispatcher: CoroutineDispatcher,
     private val bouncerInteractor: BouncerInteractor,
-    authenticationInteractor: AuthenticationInteractor,
+    private val authenticationInteractor: AuthenticationInteractor,
     flags: SceneContainerFlags,
     selectedUser: Flow<UserViewModel>,
     users: Flow<List<UserViewModel>>,
     userSwitcherMenu: Flow<List<UserActionViewModel>>,
     actionButtonInteractor: BouncerActionButtonInteractor,
     private val simBouncerInteractor: SimBouncerInteractor,
+    private val clock: SystemClock,
 ) {
     val selectedUserImage: StateFlow<Bitmap?> =
         selectedUser
@@ -104,15 +111,6 @@
     val isUserSwitcherVisible: Boolean
         get() = bouncerInteractor.isUserSwitcherVisible
 
-    private val isInputEnabled: StateFlow<Boolean> =
-        bouncerInteractor.lockout
-            .map { it == null }
-            .stateIn(
-                scope = applicationScope,
-                started = SharingStarted.WhileSubscribed(),
-                initialValue = bouncerInteractor.lockout.value == null,
-            )
-
     // Handle to the scope of the child ViewModel (stored in [authMethod]).
     private var childViewModelScope: CoroutineScope? = null
     private val _dialogMessage = MutableStateFlow<String?>(null)
@@ -138,19 +136,22 @@
      */
     val dialogMessage: StateFlow<String?> = _dialogMessage.asStateFlow()
 
+    /**
+     * A message shown when the user has attempted the wrong credential too many times and now must
+     * wait a while before attempting to authenticate again.
+     *
+     * This is updated every second (countdown) during the lockout duration. When lockout is not
+     * active, this is `null` and no lockout message should be shown.
+     */
+    private val lockoutMessage = MutableStateFlow<String?>(null)
+
     /** The user-facing message to show in the bouncer. */
     val message: StateFlow<MessageViewModel> =
-        combine(bouncerInteractor.message, bouncerInteractor.lockout) { message, lockout ->
-                toMessageViewModel(message, isLockedOut = lockout != null)
-            }
+        combine(bouncerInteractor.message, lockoutMessage) { _, _ -> createMessageViewModel() }
             .stateIn(
                 scope = applicationScope,
                 started = SharingStarted.WhileSubscribed(),
-                initialValue =
-                    toMessageViewModel(
-                        message = bouncerInteractor.message.value,
-                        isLockedOut = bouncerInteractor.lockout.value != null,
-                    ),
+                initialValue = createMessageViewModel(),
             )
 
     /**
@@ -194,24 +195,29 @@
                 initialValue = isFoldSplitRequired(authMethodViewModel.value),
             )
 
+    private val isInputEnabled: StateFlow<Boolean> =
+        lockoutMessage
+            .map { it == null }
+            .stateIn(
+                scope = applicationScope,
+                started = SharingStarted.WhileSubscribed(),
+                initialValue = authenticationInteractor.lockoutEndTimestamp == null,
+            )
+
+    private var lockoutCountdownJob: Job? = null
+
     init {
         if (flags.isEnabled()) {
             applicationScope.launch {
-                combine(bouncerInteractor.lockout, authMethodViewModel) {
-                        lockout,
-                        authMethodViewModel ->
-                        if (lockout != null && authMethodViewModel != null) {
-                            applicationContext.getString(
-                                authMethodViewModel.lockoutMessageId,
-                                lockout.failedAttemptCount,
-                                lockout.remainingSeconds,
-                            )
-                        } else {
-                            null
-                        }
-                    }
-                    .distinctUntilChanged()
-                    .collect { dialogMessage -> _dialogMessage.value = dialogMessage }
+                bouncerInteractor.onLockoutStarted.collect {
+                    showLockoutDialog()
+                    startLockoutCountdown()
+                }
+            }
+
+            applicationScope.launch {
+                // Update the lockout countdown whenever the selected user is switched.
+                selectedUser.collect { startLockoutCountdown() }
             }
         }
     }
@@ -221,6 +227,48 @@
         _dialogMessage.value = null
     }
 
+    private fun showLockoutDialog() {
+        applicationScope.launch {
+            val failedAttempts = authenticationInteractor.failedAuthenticationAttempts.value
+            _dialogMessage.value =
+                authMethodViewModel.value?.lockoutMessageId?.let { messageId ->
+                    applicationContext.getString(
+                        messageId,
+                        failedAttempts,
+                        remainingLockoutSeconds()
+                    )
+                }
+        }
+    }
+
+    /** Shows the countdown message and refreshes it every second. */
+    private fun startLockoutCountdown() {
+        lockoutCountdownJob?.cancel()
+        lockoutCountdownJob =
+            applicationScope.launch {
+                do {
+                    val remainingSeconds = remainingLockoutSeconds()
+                    lockoutMessage.value =
+                        if (remainingSeconds > 0) {
+                            applicationContext.getString(
+                                R.string.lockscreen_too_many_failed_attempts_countdown,
+                                remainingSeconds,
+                            )
+                        } else {
+                            null
+                        }
+                    delay(1.seconds)
+                } while (remainingSeconds > 0)
+                lockoutCountdownJob = null
+            }
+    }
+
+    private fun remainingLockoutSeconds(): Int {
+        val endTimestampMs = authenticationInteractor.lockoutEndTimestamp ?: 0
+        val remainingMs = max(0, endTimestampMs - clock.elapsedRealtime())
+        return ceil(remainingMs / 1000f).toInt()
+    }
+
     private fun isSideBySideSupported(authMethod: AuthMethodBouncerViewModel?): Boolean {
         return isUserSwitcherVisible || authMethod !is PasswordBouncerViewModel
     }
@@ -229,12 +277,11 @@
         return authMethod !is PasswordBouncerViewModel
     }
 
-    private fun toMessageViewModel(
-        message: String?,
-        isLockedOut: Boolean,
-    ): MessageViewModel {
+    private fun createMessageViewModel(): MessageViewModel {
+        val isLockedOut = lockoutMessage.value != null
         return MessageViewModel(
-            text = message ?: "",
+            // A lockout message takes precedence over the non-lockout message.
+            text = lockoutMessage.value ?: bouncerInteractor.message.value ?: "",
             isUpdateAnimated = !isLockedOut,
         )
     }
@@ -328,6 +375,7 @@
         userSwitcherViewModel: UserSwitcherViewModel,
         actionButtonInteractor: BouncerActionButtonInteractor,
         simBouncerInteractor: SimBouncerInteractor,
+        clock: SystemClock,
     ): BouncerViewModel {
         return BouncerViewModel(
             applicationContext = applicationContext,
@@ -341,6 +389,7 @@
             userSwitcherMenu = userSwitcherViewModel.menu,
             actionButtonInteractor = actionButtonInteractor,
             simBouncerInteractor = simBouncerInteractor,
+            clock = clock,
         )
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/PasswordBouncerViewModel.kt b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/PasswordBouncerViewModel.kt
index b682717..8e14778 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/PasswordBouncerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/PasswordBouncerViewModel.kt
@@ -56,13 +56,11 @@
 
     /** Whether the UI should request focus on the text field element. */
     val isTextFieldFocusRequested =
-        combine(interactor.lockout, isTextFieldFocused) { throttling, hasFocus ->
-                throttling == null && !hasFocus
-            }
+        combine(isInputEnabled, isTextFieldFocused) { hasInput, hasFocus -> hasInput && !hasFocus }
             .stateIn(
                 scope = viewModelScope,
                 started = SharingStarted.WhileSubscribed(),
-                initialValue = interactor.lockout.value == null && !isTextFieldFocused.value,
+                initialValue = isInputEnabled.value && !isTextFieldFocused.value,
             )
 
     override fun onHidden() {
@@ -104,7 +102,7 @@
      * hidden.
      */
     suspend fun onImeVisibilityChanged(isVisible: Boolean) {
-        if (isImeVisible && !isVisible && interactor.lockout.value == null) {
+        if (isImeVisible && !isVisible && isInputEnabled.value) {
             interactor.onImeHiddenByUser()
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt b/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt
index 887b18c..0a13e48 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt
@@ -16,8 +16,9 @@
 
 package com.android.systemui.communal.widgets
 
-import android.appwidget.AppWidgetProviderInfo
+import android.content.ComponentName
 import android.content.Intent
+import android.content.pm.PackageManager
 import android.os.Bundle
 import android.os.RemoteException
 import android.util.Log
@@ -39,10 +40,8 @@
     private var windowManagerService: IWindowManager? = null,
 ) : ComponentActivity() {
     companion object {
-        /**
-         * Intent extra name for the {@link AppWidgetProviderInfo} of a widget to add to hub mode.
-         */
-        const val ADD_WIDGET_INFO = "add_widget_info"
+        private const val EXTRA_FILTER_STRATEGY = "filter_strategy"
+        private const val FILTER_STRATEGY_GLANCEABLE_HUB = 1
         private const val TAG = "EditWidgetsActivity"
     }
 
@@ -51,13 +50,8 @@
             when (result.resultCode) {
                 RESULT_OK -> {
                     result.data
-                        ?.let {
-                            it.getParcelableExtra(
-                                ADD_WIDGET_INFO,
-                                AppWidgetProviderInfo::class.java
-                            )
-                        }
-                        ?.let { communalInteractor.addWidget(it.provider, 0) }
+                        ?.getParcelableExtra(Intent.EXTRA_COMPONENT_NAME, ComponentName::class.java)
+                        ?.let { communalInteractor.addWidget(it, 0) }
                         ?: run { Log.w(TAG, "No AppWidgetProviderInfo found in result.") }
                 }
                 else ->
@@ -71,13 +65,35 @@
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
 
+        setShowWhenLocked(true)
+
         setCommunalEditWidgetActivityContent(
             activity = this,
             viewModel = communalViewModel,
             onOpenWidgetPicker = {
-                addWidgetActivityLauncher.launch(
-                    Intent(applicationContext, WidgetPickerActivity::class.java)
-                )
+                val localPackageManager: PackageManager = getPackageManager()
+                val intent =
+                    Intent(Intent.ACTION_MAIN).also { it.addCategory(Intent.CATEGORY_HOME) }
+                localPackageManager
+                    .resolveActivity(intent, PackageManager.MATCH_DEFAULT_ONLY)
+                    ?.activityInfo
+                    ?.packageName
+                    ?.let { packageName ->
+                        try {
+                            addWidgetActivityLauncher.launch(
+                                Intent(Intent.ACTION_PICK).also {
+                                    it.setPackage(packageName)
+                                    it.putExtra(
+                                        EXTRA_FILTER_STRATEGY,
+                                        FILTER_STRATEGY_GLANCEABLE_HUB
+                                    )
+                                }
+                            )
+                        } catch (e: Exception) {
+                            Log.e(TAG, "Failed to launch widget picker activity", e)
+                        }
+                    }
+                    ?: run { Log.e(TAG, "Couldn't resolve launcher package name") }
             },
             onEditDone = {
                 try {
diff --git a/packages/SystemUI/src/com/android/systemui/communal/widgets/WidgetPickerActivity.kt b/packages/SystemUI/src/com/android/systemui/communal/widgets/WidgetPickerActivity.kt
deleted file mode 100644
index a26afc8..0000000
--- a/packages/SystemUI/src/com/android/systemui/communal/widgets/WidgetPickerActivity.kt
+++ /dev/null
@@ -1,104 +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.systemui.communal.widgets
-
-import android.appwidget.AppWidgetManager
-import android.appwidget.AppWidgetProviderInfo
-import android.content.Intent
-import android.graphics.Color
-import android.os.Bundle
-import android.util.Log
-import android.view.ViewGroup
-import android.widget.ImageView
-import android.widget.LinearLayout
-import androidx.activity.ComponentActivity
-import androidx.core.view.setMargins
-import androidx.core.view.setPadding
-import com.android.systemui.res.R
-import javax.inject.Inject
-
-/**
- * An Activity responsible for displaying a list of widgets to add to the hub mode grid. This is
- * essentially a placeholder until Launcher's widget picker can be used.
- */
-class WidgetPickerActivity
-@Inject
-constructor(
-    private val appWidgetManager: AppWidgetManager,
-) : ComponentActivity() {
-
-    override fun onCreate(savedInstanceState: Bundle?) {
-        super.onCreate(savedInstanceState)
-
-        setContentView(R.layout.widget_picker)
-        loadWidgets()
-    }
-
-    private fun loadWidgets() {
-        val containerView: ViewGroup? = findViewById(R.id.widgets_container)
-        containerView?.apply {
-            try {
-                appWidgetManager
-                    .getInstalledProviders(AppWidgetProviderInfo.WIDGET_CATEGORY_KEYGUARD)
-                    ?.stream()
-                    ?.forEach { widgetInfo ->
-                        val activity = this@WidgetPickerActivity
-                        (widgetInfo.loadPreviewImage(activity, 0)
-                                ?: widgetInfo.loadIcon(activity, 0))
-                            ?.let {
-                                addView(
-                                    ImageView(activity).also { v ->
-                                        v.setImageDrawable(it)
-                                        v.setBackgroundColor(WIDGET_PREVIEW_BACKGROUND_COLOR)
-                                        v.setPadding(WIDGET_PREVIEW_PADDING)
-                                        v.layoutParams =
-                                            LinearLayout.LayoutParams(
-                                                    WIDGET_PREVIEW_SIZE,
-                                                    WIDGET_PREVIEW_SIZE
-                                                )
-                                                .also { lp ->
-                                                    lp.setMargins(WIDGET_PREVIEW_MARGINS)
-                                                }
-                                        v.setOnClickListener {
-                                            setResult(
-                                                RESULT_OK,
-                                                Intent()
-                                                    .putExtra(
-                                                        EditWidgetsActivity.ADD_WIDGET_INFO,
-                                                        widgetInfo
-                                                    )
-                                            )
-                                            finish()
-                                        }
-                                    }
-                                )
-                            }
-                    }
-            } catch (e: RuntimeException) {
-                Log.e(TAG, "Exception fetching widget providers", e)
-            }
-        }
-    }
-
-    companion object {
-        private const val WIDGET_PREVIEW_SIZE = 600
-        private const val WIDGET_PREVIEW_MARGINS = 32
-        private const val WIDGET_PREVIEW_PADDING = 32
-        private val WIDGET_PREVIEW_BACKGROUND_COLOR = Color.rgb(216, 225, 220)
-        private const val TAG = "WidgetPickerActivity"
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/DefaultActivityBinder.java b/packages/SystemUI/src/com/android/systemui/dagger/DefaultActivityBinder.java
index 4b27af1..9afd5ed 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/DefaultActivityBinder.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/DefaultActivityBinder.java
@@ -20,7 +20,6 @@
 
 import com.android.systemui.ForegroundServicesDialog;
 import com.android.systemui.communal.widgets.EditWidgetsActivity;
-import com.android.systemui.communal.widgets.WidgetPickerActivity;
 import com.android.systemui.contrast.ContrastDialogActivity;
 import com.android.systemui.keyguard.WorkLockActivity;
 import com.android.systemui.people.PeopleSpaceActivity;
@@ -158,12 +157,6 @@
     @ClassKey(EditWidgetsActivity.class)
     public abstract Activity bindEditWidgetsActivity(EditWidgetsActivity activity);
 
-    /** Inject into WidgetPickerActivity. */
-    @Binds
-    @IntoMap
-    @ClassKey(WidgetPickerActivity.class)
-    public abstract Activity bindWidgetPickerActivity(WidgetPickerActivity activity);
-
     /** Inject into SwitchToManagedProfileForCallActivity. */
     @Binds
     @IntoMap
diff --git a/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractor.kt b/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractor.kt
index 47be8ab..f6a9570 100644
--- a/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractor.kt
@@ -159,7 +159,7 @@
     fun attemptDeviceEntry() {
         // TODO (b/307768356),
         //       1. Check if the device is already authenticated by trust agent/passive biometrics
-        //       2. show SPFS/UDFPS bouncer if it is available AlternateBouncerInteractor.show
+        //       2. Show SPFS/UDFPS bouncer if it is available AlternateBouncerInteractor.show
         //       3. For face auth only setups trigger face auth, delay transitioning to bouncer for
         //          a small amount of time.
         //       4. Transition to bouncer scene
@@ -197,8 +197,8 @@
     init {
         if (flags.isEnabled()) {
             applicationScope.launch {
-                authenticationInteractor.authenticationChallengeResult.collectLatest { successful ->
-                    if (successful) {
+                authenticationInteractor.onAuthenticationResult.collectLatest { isSuccessful ->
+                    if (isSuccessful) {
                         repository.reportSuccessfulAuthentication()
                     }
                 }
diff --git a/packages/SystemUI/src/com/android/systemui/flags/FlagDependenciesBase.kt b/packages/SystemUI/src/com/android/systemui/flags/FlagDependenciesBase.kt
index 6560ee3..14fda5e 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/FlagDependenciesBase.kt
+++ b/packages/SystemUI/src/com/android/systemui/flags/FlagDependenciesBase.kt
@@ -73,9 +73,19 @@
     ) {
         val isMet = !alphaEnabled || betaEnabled
         override fun toString(): String {
-            val isMetBullet = if (isMet) "+" else "-"
-            return "$isMetBullet $alphaName ($alphaEnabled) DEPENDS ON $betaName ($betaEnabled)"
+            val prefix =
+                when {
+                    !isMet -> "  [NOT MET]"
+                    alphaEnabled -> "      [met]"
+                    betaEnabled -> "    [ready]"
+                    else -> "[not ready]"
+                }
+            val alphaState = if (alphaEnabled) "enabled" else "disabled"
+            val betaState = if (betaEnabled) "enabled" else "disabled"
+            return "$prefix $alphaName ($alphaState) DEPENDS ON $betaName ($betaState)"
         }
+        /** Used whe posting a notification of unmet dependencies */
+        fun shortUnmetString(): String = "$alphaName DEPENDS ON $betaName"
     }
 
     protected infix fun UnreleasedFlag.dependsOn(other: UnreleasedFlag) =
@@ -124,7 +134,7 @@
         unmet: List<FlagDependenciesBase.Dependency>
     ) {
         val title = "Invalid flag dependencies: ${unmet.size} of ${all.size}"
-        val details = unmet.joinToString("\n")
+        val details = unmet.joinToString("\n") { it.shortUnmetString() }
         Log.e("FlagDependencies", "$title:\n$details")
         val channel = NotificationChannel("FLAGS", "Flags", NotificationManager.IMPORTANCE_DEFAULT)
         val notification =
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
index 47518bb..5abb4dd 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
@@ -48,6 +48,7 @@
 import com.android.systemui.util.asIndenting
 import com.android.systemui.util.printSection
 import com.android.systemui.util.println
+import dagger.Lazy
 import java.io.PrintWriter
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
@@ -80,8 +81,8 @@
     private val sceneLogger: SceneLogger,
     @FalsingCollectorActual private val falsingCollector: FalsingCollector,
     private val powerInteractor: PowerInteractor,
-    private val simBouncerInteractor: SimBouncerInteractor,
-    private val authenticationInteractor: AuthenticationInteractor,
+    private val simBouncerInteractor: Lazy<SimBouncerInteractor>,
+    private val authenticationInteractor: Lazy<AuthenticationInteractor>,
 ) : CoreStartable {
 
     override fun start() {
@@ -152,7 +153,7 @@
             }
         }
         applicationScope.launch {
-            simBouncerInteractor.isAnySimSecure.collect { isAnySimLocked ->
+            simBouncerInteractor.get().isAnySimSecure.collect { isAnySimLocked ->
                 val canSwipeToEnter = deviceEntryInteractor.canSwipeToEnter.value
                 val isUnlocked = deviceEntryInteractor.isUnlocked.value
 
@@ -166,15 +167,17 @@
                     isUnlocked && canSwipeToEnter == false -> {
                         switchToScene(
                             targetSceneKey = SceneKey.Gone,
-                            loggingReason = "All SIM cards unlocked and device already" +
-                                " unlocked and lockscreen doesn't require a swipe to dismiss."
+                            loggingReason =
+                                "All SIM cards unlocked and device already" +
+                                    " unlocked and lockscreen doesn't require a swipe to dismiss."
                         )
                     }
                     else -> {
                         switchToScene(
                             targetSceneKey = SceneKey.Lockscreen,
-                            loggingReason = "All SIM cards unlocked and device still locked" +
-                                " or lockscreen still requires a swipe to dismiss."
+                            loggingReason =
+                                "All SIM cards unlocked and device still locked" +
+                                    " or lockscreen still requires a swipe to dismiss."
                         )
                     }
                 }
@@ -262,7 +265,7 @@
                                     " to swipe up on lockscreen to enter.",
                         )
                     } else if (
-                        authenticationInteractor.getAuthenticationMethod() ==
+                        authenticationInteractor.get().getAuthenticationMethod() ==
                             AuthenticationMethodModel.Sim
                     ) {
                         switchToScene(
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractor.kt
index 30e2f0e0..9215568 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractor.kt
@@ -21,6 +21,7 @@
 import com.android.systemui.statusbar.data.repository.NotificationListenerSettingsRepository
 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.ActiveNotificationModel
 import com.android.wm.shell.bubbles.Bubbles
 import java.util.Optional
@@ -37,10 +38,12 @@
 constructor(
     private val activeNotificationsInteractor: ActiveNotificationsInteractor,
     private val bubbles: Optional<Bubbles>,
+    private val headsUpNotificationIconInteractor: HeadsUpNotificationIconInteractor,
     private val keyguardViewStateRepository: NotificationsKeyguardViewStateRepository,
 ) {
     /** Returns a subset of all active notifications based on the supplied filtration parameters. */
     fun filteredNotifSet(
+        forceShowHeadsUp: Boolean = false,
         showAmbient: Boolean = true,
         showLowPriority: Boolean = true,
         showDismissed: Boolean = true,
@@ -49,18 +52,21 @@
     ): Flow<Set<ActiveNotificationModel>> {
         return combine(
             activeNotificationsInteractor.topLevelRepresentativeNotifications,
+            headsUpNotificationIconInteractor.isolatedNotification,
             keyguardViewStateRepository.areNotificationsFullyHidden,
-        ) { notifications, notifsFullyHidden ->
+        ) { notifications, isolatedNotifKey, notifsFullyHidden ->
             notifications
                 .asSequence()
                 .filter { model: ActiveNotificationModel ->
                     shouldShowNotificationIcon(
                         model = model,
+                        forceShowHeadsUp = forceShowHeadsUp,
                         showAmbient = showAmbient,
                         showLowPriority = showLowPriority,
                         showDismissed = showDismissed,
                         showRepliedMessages = showRepliedMessages,
                         showPulsing = showPulsing,
+                        isolatedNotifKey = isolatedNotifKey,
                         notifsFullyHidden = notifsFullyHidden,
                     )
                 }
@@ -70,14 +76,17 @@
 
     private fun shouldShowNotificationIcon(
         model: ActiveNotificationModel,
+        forceShowHeadsUp: Boolean,
         showAmbient: Boolean,
         showLowPriority: Boolean,
         showDismissed: Boolean,
         showRepliedMessages: Boolean,
         showPulsing: Boolean,
+        isolatedNotifKey: String?,
         notifsFullyHidden: Boolean,
     ): Boolean {
         return when {
+            forceShowHeadsUp && model.key == isolatedNotifKey -> true
             !showAmbient && model.isAmbient -> false
             !showLowPriority && model.isSilent -> false
             !showDismissed && model.isRowDismissed -> false
@@ -118,6 +127,7 @@
     val statusBarNotifs: Flow<Set<ActiveNotificationModel>> =
         settingsRepository.showSilentStatusIcons.flatMapLatest { showSilentIcons ->
             iconsInteractor.filteredNotifSet(
+                forceShowHeadsUp = true,
                 showAmbient = false,
                 showLowPriority = showSilentIcons,
                 showDismissed = false,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerStatusBarViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerStatusBarViewModel.kt
index 6e5ac47..d00cd1f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerStatusBarViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerStatusBarViewModel.kt
@@ -96,8 +96,8 @@
                     iconsViewData.visibleIcons.firstOrNull { it.notifKey == isolatedNotif }
                 }
             }
-            .pairwise(initialValue = null)
             .distinctUntilChanged()
+            .pairwise(initialValue = null)
             .sample(shadeInteractor.shadeExpansion) { (prev, iconInfo), shadeExpansion ->
                 val animate =
                     when {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
index 5872840..31ca106 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
@@ -513,15 +513,13 @@
     private void setImageViewAnimationRunning(ImageView imageView, boolean running) {
         if (imageView != null) {
             Drawable drawable = imageView.getDrawable();
-            if (drawable instanceof AnimationDrawable) {
-                AnimationDrawable animationDrawable = (AnimationDrawable) drawable;
+            if (drawable instanceof AnimationDrawable animationDrawable) {
                 if (running) {
                     animationDrawable.start();
                 } else {
                     animationDrawable.stop();
                 }
-            } else if (drawable instanceof AnimatedVectorDrawable) {
-                AnimatedVectorDrawable animationDrawable = (AnimatedVectorDrawable) drawable;
+            } else if (drawable instanceof AnimatedVectorDrawable animationDrawable) {
                 if (running) {
                     animationDrawable.start();
                 } else {
@@ -3439,8 +3437,7 @@
 
     @Override
     protected boolean childNeedsClipping(View child) {
-        if (child instanceof NotificationContentView) {
-            NotificationContentView contentView = (NotificationContentView) child;
+        if (child instanceof NotificationContentView contentView) {
             if (isClippingNeeded()) {
                 return true;
             } else if (hasRoundedCorner()
@@ -3522,8 +3519,7 @@
 
         @Override
         public void applyToView(View view) {
-            if (view instanceof ExpandableNotificationRow) {
-                ExpandableNotificationRow row = (ExpandableNotificationRow) view;
+            if (view instanceof ExpandableNotificationRow row) {
                 if (row.isExpandAnimationRunning()) {
                     return;
                 }
@@ -3543,8 +3539,7 @@
         @Override
         protected void onYTranslationAnimationFinished(View view) {
             super.onYTranslationAnimationFinished(view);
-            if (view instanceof ExpandableNotificationRow) {
-                ExpandableNotificationRow row = (ExpandableNotificationRow) view;
+            if (view instanceof ExpandableNotificationRow row) {
                 if (row.isHeadsUpAnimatingAway()) {
                     row.setHeadsUpAnimatingAway(false);
                 }
@@ -3553,8 +3548,7 @@
 
         @Override
         public void animateTo(View child, AnimationProperties properties) {
-            if (child instanceof ExpandableNotificationRow) {
-                ExpandableNotificationRow row = (ExpandableNotificationRow) child;
+            if (child instanceof ExpandableNotificationRow row) {
                 if (row.isExpandAnimationRunning()) {
                     return;
                 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableView.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableView.java
index 49674d6..c4d266e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableView.java
@@ -676,8 +676,7 @@
         mViewState.headsUpIsVisible = false;
 
         // handling reset for child notifications
-        if (this instanceof ExpandableNotificationRow) {
-            ExpandableNotificationRow row = (ExpandableNotificationRow) this;
+        if (this instanceof ExpandableNotificationRow row) {
             List<ExpandableNotificationRow> children = row.getAttachedChildren();
             if (row.isSummaryWithChildren() && children != null) {
                 for (ExpandableNotificationRow childRow : children) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationChildrenContainer.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationChildrenContainer.java
index 0236fc2..45b9c26 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationChildrenContainer.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationChildrenContainer.java
@@ -869,8 +869,7 @@
         Path clipPath = mChildClipPath;
         if (clipPath != null) {
             final float translation;
-            if (child instanceof ExpandableNotificationRow) {
-                ExpandableNotificationRow notificationRow = (ExpandableNotificationRow) child;
+            if (child instanceof ExpandableNotificationRow notificationRow) {
                 translation = notificationRow.getTranslation();
             } else {
                 translation = child.getTranslationX();
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
index 62fdc29..0f640c9 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
@@ -993,8 +993,7 @@
         for (int i = 0; i < childCount; i++) {
             View child = getChildAt(i);
             if (child.getVisibility() != View.GONE
-                    && child instanceof ExpandableNotificationRow) {
-                ExpandableNotificationRow row = (ExpandableNotificationRow) child;
+                    && child instanceof ExpandableNotificationRow row) {
                 if ((row.isPinned() || row.isHeadsUpAnimatingAway()) && row.getTranslation() < 0
                         && row.getProvider().shouldShowGutsOnSnapOpen()) {
                     top = Math.min(top, row.getTranslationY());
@@ -1129,10 +1128,9 @@
             for (int i = 0; i < n; i++) {
                 View view = getChildAt(i);
                 if (view.getVisibility() == View.GONE
-                        || !(view instanceof ExpandableNotificationRow)) {
+                        || !(view instanceof ExpandableNotificationRow row)) {
                     continue;
                 }
-                ExpandableNotificationRow row = (ExpandableNotificationRow) view;
                 currentIndex++;
                 boolean beforeSpeedBump;
                 if (mHighPriorityBeforeSpeedBump) {
@@ -1769,16 +1767,14 @@
     }
 
     public static boolean isPinnedHeadsUp(View v) {
-        if (v instanceof ExpandableNotificationRow) {
-            ExpandableNotificationRow row = (ExpandableNotificationRow) v;
+        if (v instanceof ExpandableNotificationRow row) {
             return row.isHeadsUp() && row.isPinned();
         }
         return false;
     }
 
     private boolean isHeadsUp(View v) {
-        if (v instanceof ExpandableNotificationRow) {
-            ExpandableNotificationRow row = (ExpandableNotificationRow) v;
+        if (v instanceof ExpandableNotificationRow row) {
             return row.isHeadsUp();
         }
         return false;
@@ -1820,8 +1816,7 @@
 
             if ((bottom - top >= mMinInteractionHeight || !requireMinHeight)
                     && touchY >= top && touchY <= bottom && touchX >= left && touchX <= right) {
-                if (slidingChild instanceof ExpandableNotificationRow) {
-                    ExpandableNotificationRow row = (ExpandableNotificationRow) slidingChild;
+                if (slidingChild instanceof ExpandableNotificationRow row) {
                     NotificationEntry entry = row.getEntry();
                     if (!mIsExpanded && row.isHeadsUp() && row.isPinned()
                             && mTopHeadsUpEntry.getRow() != row
@@ -2364,8 +2359,7 @@
             float rowTranslation = child.getTranslationY();
             if (rowTranslation >= translationY) {
                 return child;
-            } else if (!ignoreChildren && child instanceof ExpandableNotificationRow) {
-                ExpandableNotificationRow row = (ExpandableNotificationRow) child;
+            } else if (!ignoreChildren && child instanceof ExpandableNotificationRow row) {
                 if (row.isSummaryWithChildren() && row.areChildrenExpanded()) {
                     List<ExpandableNotificationRow> notificationChildren =
                             row.getAttachedChildren();
@@ -2886,8 +2880,7 @@
     }
 
     private void focusNextViewIfFocused(View view) {
-        if (view instanceof ExpandableNotificationRow) {
-            ExpandableNotificationRow row = (ExpandableNotificationRow) view;
+        if (view instanceof ExpandableNotificationRow row) {
             if (row.shouldRefocusOnDismiss()) {
                 View nextView = row.getChildAfterViewWhenDismissed();
                 if (nextView == null) {
@@ -3035,8 +3028,7 @@
     }
 
     private int getIntrinsicHeight(View view) {
-        if (view instanceof ExpandableView) {
-            ExpandableView expandableView = (ExpandableView) view;
+        if (view instanceof ExpandableView expandableView) {
             return expandableView.getIntrinsicHeight();
         }
         return view.getHeight();
@@ -3126,8 +3118,7 @@
         generateAddAnimation(child, false /* fromMoreCard */);
         updateAnimationState(child);
         updateChronometerForChild(child);
-        if (child instanceof ExpandableNotificationRow) {
-            ExpandableNotificationRow row = (ExpandableNotificationRow) child;
+        if (child instanceof ExpandableNotificationRow row) {
             row.setDismissUsingRowTranslationX(mDismissUsingRowTranslationX);
 
         }
@@ -3196,8 +3187,7 @@
     }
 
     private void updateAnimationState(boolean running, View child) {
-        if (child instanceof ExpandableNotificationRow) {
-            ExpandableNotificationRow row = (ExpandableNotificationRow) child;
+        if (child instanceof ExpandableNotificationRow row) {
             row.setAnimationRunning(running);
         }
     }
@@ -3394,8 +3384,7 @@
             // we need to know the view after this one
             float removedTranslation = child.getTranslationY();
             boolean ignoreChildren = true;
-            if (child instanceof ExpandableNotificationRow) {
-                ExpandableNotificationRow row = (ExpandableNotificationRow) child;
+            if (child instanceof ExpandableNotificationRow row) {
                 if (row.isRemoved() && row.wasChildInGroupWhenRemoved()) {
                     removedTranslation = row.getTranslationWhenRemoved();
                     ignoreChildren = false;
@@ -3437,8 +3426,7 @@
     private void generatePositionChangeEvents() {
         for (ExpandableView child : mChildrenChangingPositions) {
             Integer duration = null;
-            if (child instanceof ExpandableNotificationRow) {
-                ExpandableNotificationRow row = (ExpandableNotificationRow) child;
+            if (child instanceof ExpandableNotificationRow row) {
                 if (row.getEntry().isMarkedForUserTriggeredMovement()) {
                     duration = StackStateAnimator.ANIMATION_DURATION_PRIORITY_CHANGE;
                     row.getEntry().markForUserTriggeredMovement(false);
@@ -4122,8 +4110,7 @@
     private void clearUserLockedViews() {
         for (int i = 0; i < getChildCount(); i++) {
             ExpandableView child = getChildAtIndex(i);
-            if (child instanceof ExpandableNotificationRow) {
-                ExpandableNotificationRow row = (ExpandableNotificationRow) child;
+            if (child instanceof ExpandableNotificationRow row) {
                 row.setUserLocked(false);
             }
         }
@@ -4137,8 +4124,7 @@
         );
         for (int i = 0; i < getChildCount(); i++) {
             ExpandableView child = getChildAtIndex(i);
-            if (child instanceof ExpandableNotificationRow) {
-                ExpandableNotificationRow row = (ExpandableNotificationRow) child;
+            if (child instanceof ExpandableNotificationRow row) {
                 clearTemporaryViewsInGroup(
                         /* viewGroup = */ row.getChildrenContainer(),
                         /* reason = */ "clearTemporaryViewsInGroup(row.getChildrenContainer())"
@@ -4223,8 +4209,7 @@
     }
 
     void updateChronometerForChild(View child) {
-        if (child instanceof ExpandableNotificationRow) {
-            ExpandableNotificationRow row = (ExpandableNotificationRow) child;
+        if (child instanceof ExpandableNotificationRow row) {
             row.setChronometerRunning(mIsExpanded);
         }
     }
@@ -4263,8 +4248,7 @@
     }
 
     private void updateScrollPositionOnExpandInBottom(ExpandableView view) {
-        if (view instanceof ExpandableNotificationRow && !onKeyguard()) {
-            ExpandableNotificationRow row = (ExpandableNotificationRow) view;
+        if (view instanceof ExpandableNotificationRow row && !onKeyguard()) {
             // TODO: once we're recycling this will need to check the adapter position of the child
             if (row.isUserLocked() && row != getFirstChildNotGone()) {
                 if (row.isSummaryWithChildren()) {
@@ -4323,8 +4307,7 @@
     private void clearHeadsUpDisappearRunning() {
         for (int i = 0; i < getChildCount(); i++) {
             View view = getChildAt(i);
-            if (view instanceof ExpandableNotificationRow) {
-                ExpandableNotificationRow row = (ExpandableNotificationRow) view;
+            if (view instanceof ExpandableNotificationRow row) {
                 row.setHeadsUpAnimatingAway(false);
                 if (row.isSummaryWithChildren()) {
                     for (ExpandableNotificationRow child : row.getAttachedChildren()) {
@@ -5208,8 +5191,7 @@
                     }
                     View swipedView = mSwipeHelper.getSwipedView();
                     pw.println("Swiped view: " + swipedView);
-                    if (swipedView instanceof ExpandableView) {
-                        ExpandableView expandableView = (ExpandableView) swipedView;
+                    if (swipedView instanceof ExpandableView expandableView) {
                         expandableView.dump(pw, args);
                     }
                 });
@@ -5292,8 +5274,7 @@
         if (view instanceof SectionHeaderView && silentSectionWillBeGone) {
             return true;
         }
-        if (view instanceof ExpandableNotificationRow) {
-            ExpandableNotificationRow row = (ExpandableNotificationRow) view;
+        if (view instanceof ExpandableNotificationRow row) {
             if (isVisible(row) && includeChildInClearAll(row, selection)) {
                 return true;
             }
@@ -5319,9 +5300,7 @@
             if (shouldHideParent(view, selection)) {
                 viewsToHide.add(view);
             }
-            if (view instanceof ExpandableNotificationRow) {
-                ExpandableNotificationRow parent = (ExpandableNotificationRow) view;
-
+            if (view instanceof ExpandableNotificationRow parent) {
                 if (isChildrenVisible(parent)) {
                     for (ExpandableNotificationRow child : parent.getAttachedChildren()) {
                         if (isVisible(child) && includeChildInClearAll(child, selection)) {
@@ -5341,10 +5320,9 @@
 
         for (int i = 0; i < childCount; i++) {
             final View view = getChildAt(i);
-            if (!(view instanceof ExpandableNotificationRow)) {
+            if (!(view instanceof ExpandableNotificationRow parent)) {
                 continue;
             }
-            ExpandableNotificationRow parent = (ExpandableNotificationRow) view;
             if (includeChildInClearAll(parent, selection)) {
                 viewsToRemove.add(parent);
             }
@@ -5983,8 +5961,7 @@
         for (int i = 0; i < getChildCount(); i++) {
             View child = getChildAt(i);
             mSwipeHelper.forceResetSwipeState(child);
-            if (child instanceof ExpandableNotificationRow) {
-                ExpandableNotificationRow childRow = (ExpandableNotificationRow) child;
+            if (child instanceof ExpandableNotificationRow childRow) {
                 List<ExpandableNotificationRow> grandchildren = childRow.getAttachedChildren();
                 if (grandchildren != null) {
                     for (ExpandableNotificationRow grandchild : grandchildren) {
@@ -6270,8 +6247,7 @@
     }
 
     static boolean canChildBeDismissed(View v) {
-        if (v instanceof ExpandableNotificationRow) {
-            ExpandableNotificationRow row = (ExpandableNotificationRow) v;
+        if (v instanceof ExpandableNotificationRow row) {
             if (row.areGutsExposed() || !row.getEntry().hasFinishedInitialization()) {
                 return false;
             }
@@ -6281,8 +6257,7 @@
     }
 
     static boolean canChildBeCleared(View v) {
-        if (v instanceof ExpandableNotificationRow) {
-            ExpandableNotificationRow row = (ExpandableNotificationRow) v;
+        if (v instanceof ExpandableNotificationRow row) {
             if (row.areGutsExposed() || !row.getEntry().hasFinishedInitialization()) {
                 return false;
             }
@@ -6377,8 +6352,7 @@
         /* Only ever called as a consequence of an expansion gesture in the shade. */
         @Override
         public void setUserExpandedChild(View v, boolean userExpanded) {
-            if (v instanceof ExpandableNotificationRow) {
-                ExpandableNotificationRow row = (ExpandableNotificationRow) v;
+            if (v instanceof ExpandableNotificationRow row) {
                 if (userExpanded && onKeyguard()) {
                     // Due to a race when locking the screen while touching, a notification may be
                     // expanded even after we went back to keyguard. An example of this happens if
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
index 7c7d943..6f5058c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
@@ -405,8 +405,7 @@
             if (!mAllowLongPress) {
                 return;
             }
-            if (view instanceof ExpandableNotificationRow) {
-                ExpandableNotificationRow row = (ExpandableNotificationRow) view;
+            if (view instanceof ExpandableNotificationRow row) {
                 mMetricsLogger.write(row.getEntry().getSbn().getLogMaker()
                         .setCategory(MetricsEvent.ACTION_TOUCH_GEAR)
                         .setType(MetricsEvent.TYPE_ACTION)
@@ -426,8 +425,7 @@
 
         @Override
         public void onMenuShown(View row) {
-            if (row instanceof ExpandableNotificationRow) {
-                ExpandableNotificationRow notificationRow = (ExpandableNotificationRow) row;
+            if (row instanceof ExpandableNotificationRow notificationRow) {
                 mMetricsLogger.write(notificationRow.getEntry().getSbn().getLogMaker()
                         .setCategory(MetricsEvent.ACTION_REVEAL_GEAR)
                         .setType(MetricsEvent.TYPE_ACTION));
@@ -492,10 +490,9 @@
                  */
                 @Override
                 public void onChildDismissed(View view) {
-                    if (!(view instanceof ActivatableNotificationView)) {
+                    if (!(view instanceof ActivatableNotificationView row)) {
                         return;
                     }
-                    ActivatableNotificationView row = (ActivatableNotificationView) view;
                     if (!row.isDismissed()) {
                         handleChildViewDismissed(view);
                     }
@@ -519,8 +516,7 @@
                     if (mView.getClearAllInProgress()) {
                         return;
                     }
-                    if (view instanceof ExpandableNotificationRow) {
-                        ExpandableNotificationRow row = (ExpandableNotificationRow) view;
+                    if (view instanceof ExpandableNotificationRow row) {
                         if (row.isHeadsUp()) {
                             mHeadsUpManager.addSwipedOutNotification(
                                     row.getEntry().getSbn().getKey());
@@ -551,8 +547,7 @@
                             ev.getY(),
                             true /* requireMinHeight */,
                             false /* ignoreDecors */);
-                    if (child instanceof ExpandableNotificationRow) {
-                        ExpandableNotificationRow row = (ExpandableNotificationRow) child;
+                    if (child instanceof ExpandableNotificationRow row) {
                         ExpandableNotificationRow parent = row.getNotificationParent();
                         if (parent != null && parent.areChildrenExpanded()
                                 && (parent.areGutsExposed()
@@ -582,8 +577,7 @@
                 @Override
                 public void onChildSnappedBack(View animView, float targetLeft) {
                     mView.onSwipeEnd();
-                    if (animView instanceof ExpandableNotificationRow) {
-                        ExpandableNotificationRow row = (ExpandableNotificationRow) animView;
+                    if (animView instanceof ExpandableNotificationRow row) {
                         if (row.isPinned() && !canChildBeDismissed(row)
                                 && row.getEntry().getSbn().getNotification().fullScreenIntent
                                 == null) {
@@ -1980,8 +1974,7 @@
 
             // Check if we need to clear any snooze leavebehinds
             if (guts != null && !NotificationSwipeHelper.isTouchInView(ev, guts)
-                    && guts.getGutsContent() instanceof NotificationSnooze) {
-                NotificationSnooze ns = (NotificationSnooze) guts.getGutsContent();
+                    && guts.getGutsContent() instanceof NotificationSnooze ns) {
                 if ((ns.isExpanded() && isCancelOrUp)
                         || (!horizontalSwipeWantsIt && scrollerWantsIt)) {
                     // If the leavebehind is expanded we clear it on the next up event, otherwise we
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithm.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithm.java
index c4e6b909..664a6b6 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithm.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithm.java
@@ -254,8 +254,7 @@
 
     public static void logView(View view, String s) {
         String viewString = "";
-        if (view instanceof ExpandableNotificationRow) {
-            ExpandableNotificationRow row = ((ExpandableNotificationRow) view);
+        if (view instanceof ExpandableNotificationRow row) {
             if (row.getEntry() == null) {
                 viewString = "ExpandableNotificationRow has null NotificationEntry";
             } else {
@@ -289,8 +288,7 @@
         int childCount = algorithmState.visibleChildren.size();
         for (int i = 0; i < childCount; i++) {
             ExpandableView v = algorithmState.visibleChildren.get(i);
-            if (v instanceof ExpandableNotificationRow) {
-                ExpandableNotificationRow row = (ExpandableNotificationRow) v;
+            if (v instanceof ExpandableNotificationRow row) {
                 row.updateChildrenStates();
             }
         }
@@ -401,8 +399,7 @@
                     continue;
                 }
                 notGoneIndex = updateNotGoneIndex(state, notGoneIndex, v);
-                if (v instanceof ExpandableNotificationRow) {
-                    ExpandableNotificationRow row = (ExpandableNotificationRow) v;
+                if (v instanceof ExpandableNotificationRow row) {
 
                     // handle the notGoneIndex for the children as well
                     List<ExpandableNotificationRow> children = row.getAttachedChildren();
@@ -533,10 +530,9 @@
     private boolean hasNonClearableNotifs(StackScrollAlgorithmState algorithmState) {
         for (int i = 0; i < algorithmState.visibleChildren.size(); i++) {
             View child = algorithmState.visibleChildren.get(i);
-            if (!(child instanceof ExpandableNotificationRow)) {
+            if (!(child instanceof ExpandableNotificationRow row)) {
                 continue;
             }
-            final ExpandableNotificationRow row = (ExpandableNotificationRow) child;
             if (!row.canViewBeCleared()) {
                 return true;
             }
@@ -740,10 +736,9 @@
         ExpandableNotificationRow pulsingRow = null;
         for (int i = 0; i < childCount; i++) {
             View child = algorithmState.visibleChildren.get(i);
-            if (!(child instanceof ExpandableNotificationRow)) {
+            if (!(child instanceof ExpandableNotificationRow row)) {
                 continue;
             }
-            ExpandableNotificationRow row = (ExpandableNotificationRow) child;
             if (!row.showingPulsing() || (i == 0 && ambientState.isPulseExpanding())) {
                 continue;
             }
@@ -785,10 +780,9 @@
         ExpandableNotificationRow topHeadsUpEntry = null;
         for (int i = 0; i < childCount; i++) {
             View child = algorithmState.visibleChildren.get(i);
-            if (!(child instanceof ExpandableNotificationRow)) {
+            if (!(child instanceof ExpandableNotificationRow row)) {
                 continue;
             }
-            ExpandableNotificationRow row = (ExpandableNotificationRow) child;
             if (!(row.isHeadsUp() || row.isHeadsUpAnimatingAway())) {
                 continue;
             }
@@ -940,8 +934,7 @@
     }
 
     protected int getMaxAllowedChildHeight(View child) {
-        if (child instanceof ExpandableView) {
-            ExpandableView expandableView = (ExpandableView) child;
+        if (child instanceof ExpandableView expandableView) {
             return expandableView.getIntrinsicHeight();
         }
         return child == null ? mCollapsedSize : child.getHeight();
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/RemoteInputView.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/RemoteInputView.java
index 4864fb8..5bced93 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/RemoteInputView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/RemoteInputView.java
@@ -303,8 +303,7 @@
                     mEntry.mRemoteEditImeVisible = editTextRootWindowInsets != null
                             && editTextRootWindowInsets.isVisible(WindowInsets.Type.ime());
                     if (!mEntry.mRemoteEditImeVisible && !mEditText.mShowImeOnInputConnection) {
-                        // Pass null to ensure all inputs are cleared for this entry b/227115380
-                            mController.removeRemoteInput(mEntry, null,
+                            mController.removeRemoteInput(mEntry, mToken,
                                     /* reason= */"RemoteInputView$WindowInsetAnimation#onEnd");
                     }
                 }
@@ -536,6 +535,11 @@
         if (mEntry.getRow().isChangingPosition() || isTemporarilyDetached()) {
             return;
         }
+        // RemoteInputView can be detached from window before IME close event in some cases like
+        // remote input view removal with notification update. As a result of this, RemoteInputView
+        // will stop ime animation updates, which results in never removing remote input. That's why
+        // we have to set mRemoteEditImeAnimatingAway false on detach to remove remote input.
+        mEntry.mRemoteEditImeAnimatingAway = false;
         mController.removeRemoteInput(mEntry, mToken,
                 /* reason= */"RemoteInputView#onDetachedFromWindow");
         mController.removeSpinning(mEntry.getKey(), mToken);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/biometrics/data/repository/BiometricStatusRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/biometrics/data/repository/BiometricStatusRepositoryTest.kt
index 27d93eb..8f0e910 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/biometrics/data/repository/BiometricStatusRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/biometrics/data/repository/BiometricStatusRepositoryTest.kt
@@ -24,7 +24,6 @@
 import android.hardware.biometrics.BiometricRequestConstants.REASON_AUTH_SETTINGS
 import android.hardware.biometrics.BiometricRequestConstants.REASON_ENROLL_ENROLLING
 import android.hardware.biometrics.BiometricRequestConstants.REASON_ENROLL_FIND_SENSOR
-import android.platform.test.annotations.RequiresFlagsEnabled
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.biometrics.shared.model.AuthenticationReason
@@ -48,7 +47,6 @@
 import org.mockito.junit.MockitoJUnit
 import org.mockito.junit.MockitoRule
 
-@RequiresFlagsEnabled(FLAG_SIDEFPS_CONTROLLER_REFACTOR)
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(JUnit4::class)
@@ -62,6 +60,7 @@
 
     @Before
     fun setUp() {
+        mSetFlagsRule.enableFlags(FLAG_SIDEFPS_CONTROLLER_REFACTOR)
         underTest = BiometricStatusRepositoryImpl(testScope.backgroundScope, biometricManager)
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/biometrics/domain/interactor/BiometricStatusInteractorImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/biometrics/domain/interactor/BiometricStatusInteractorImplTest.kt
index 6978923..d7b7d79 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/biometrics/domain/interactor/BiometricStatusInteractorImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/biometrics/domain/interactor/BiometricStatusInteractorImplTest.kt
@@ -19,7 +19,6 @@
 import android.app.ActivityManager
 import android.app.ActivityTaskManager
 import android.content.ComponentName
-import android.platform.test.annotations.RequiresFlagsEnabled
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.biometrics.data.repository.FakeBiometricStatusRepository
@@ -44,7 +43,6 @@
 import org.mockito.junit.MockitoJUnit
 import org.mockito.junit.MockitoRule
 
-@RequiresFlagsEnabled(FLAG_SIDEFPS_CONTROLLER_REFACTOR)
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(JUnit4::class)
@@ -59,6 +57,7 @@
 
     @Before
     fun setup() {
+        mSetFlagsRule.enableFlags(FLAG_SIDEFPS_CONTROLLER_REFACTOR)
         biometricStatusRepository = FakeBiometricStatusRepository()
         underTest = BiometricStatusInteractorImpl(activityTaskManager, biometricStatusRepository)
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/DeviceEntrySideFpsOverlayInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/DeviceEntrySideFpsOverlayInteractorTest.kt
index 70d3f81..0616a34 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/DeviceEntrySideFpsOverlayInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/DeviceEntrySideFpsOverlayInteractorTest.kt
@@ -17,7 +17,6 @@
 package com.android.systemui.keyguard.domain.interactor
 
 import android.os.Handler
-import android.platform.test.annotations.RequiresFlagsEnabled
 import androidx.test.filters.SmallTest
 import com.android.keyguard.KeyguardSecurityModel
 import com.android.keyguard.KeyguardUpdateMonitor
@@ -58,7 +57,6 @@
 import org.mockito.junit.MockitoRule
 
 @OptIn(ExperimentalCoroutinesApi::class)
-@RequiresFlagsEnabled(FLAG_SIDEFPS_CONTROLLER_REFACTOR)
 @SmallTest
 @RunWith(JUnit4::class)
 class DeviceEntrySideFpsOverlayInteractorTest : SysuiTestCase() {
@@ -80,6 +78,7 @@
 
     @Before
     fun setup() {
+        mSetFlagsRule.enableFlags(FLAG_SIDEFPS_CONTROLLER_REFACTOR)
         primaryBouncerInteractor =
             PrimaryBouncerInteractor(
                 bouncerRepository,
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 360a373..47feccf 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
@@ -30,6 +30,7 @@
 import com.android.systemui.statusbar.notification.data.repository.ActiveNotificationListRepository
 import com.android.systemui.statusbar.notification.data.repository.ActiveNotificationsStore
 import com.android.systemui.statusbar.notification.data.repository.FakeNotificationsKeyguardViewStateRepository
+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
@@ -264,6 +265,7 @@
     interface TestComponent : SysUITestComponent<StatusBarNotificationIconsInteractor> {
 
         val activeNotificationListRepository: ActiveNotificationListRepository
+        val headsUpIconsInteractor: HeadsUpNotificationIconInteractor
         val keyguardViewStateRepository: FakeNotificationsKeyguardViewStateRepository
         val notificationListenerSettingsRepository: NotificationListenerSettingsRepository
 
@@ -336,6 +338,14 @@
                 .comparingElementsUsing(byIsLastMessageFromReply)
                 .doesNotContain(true)
         }
+
+    @Test
+    fun filteredEntrySet_includesIsolatedIcon() =
+        testComponent.runTest {
+            val filteredSet by collectLastValue(underTest.statusBarNotifs)
+            headsUpIconsInteractor.setIsolatedIconNotificationKey("notif5")
+            assertThat(filteredSet).comparingElementsUsing(byKey).contains("notif5")
+        }
 }
 
 private val testIcons =
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerStatusBarViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerStatusBarViewModelTest.kt
index 349a35eb..c40401f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerStatusBarViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerStatusBarViewModelTest.kt
@@ -399,4 +399,29 @@
 
             assertThat(isolatedIcon?.value?.notifKey).isEqualTo("notif1")
         }
+
+    @Test
+    fun isolatedIcon_lastMessageIsFromReply_notNull() =
+        testComponent.runTest {
+            val icon: Icon = mock()
+            headsUpViewStateRepository.isolatedNotification.value = "notif1"
+            activeNotificationsRepository.activeNotifications.value =
+                ActiveNotificationsStore.Builder()
+                    .apply {
+                        addIndividualNotif(
+                            activeNotificationModel(
+                                key = "notif1",
+                                groupKey = "group",
+                                statusBarIcon = icon,
+                                isLastMessageFromReply = true,
+                            )
+                        )
+                    }
+                    .build()
+
+            val isolatedIcon by collectLastValue(underTest.isolatedIcon)
+            runCurrent()
+
+            assertThat(isolatedIcon?.value?.notifKey).isEqualTo("notif1")
+        }
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/authentication/data/repository/FakeAuthenticationRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/authentication/data/repository/FakeAuthenticationRepository.kt
index 7c5696c..8486508 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/authentication/data/repository/FakeAuthenticationRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/authentication/data/repository/FakeAuthenticationRepository.kt
@@ -20,7 +20,6 @@
 import com.android.internal.widget.LockPatternView
 import com.android.internal.widget.LockscreenCredential
 import com.android.keyguard.KeyguardSecurityModel.SecurityMode
-import com.android.systemui.authentication.shared.model.AuthenticationLockoutModel
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
 import com.android.systemui.authentication.shared.model.AuthenticationPatternCoordinate
 import com.android.systemui.authentication.shared.model.AuthenticationResultModel
@@ -29,7 +28,6 @@
 import dagger.Module
 import dagger.Provides
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
@@ -40,15 +38,11 @@
     private val currentTime: () -> Long,
 ) : AuthenticationRepository {
 
-    override val authenticationChallengeResult = MutableSharedFlow<Boolean>()
-
     override val hintedPinLength: Int = HINTING_PIN_LENGTH
 
     private val _isPatternVisible = MutableStateFlow(true)
     override val isPatternVisible: StateFlow<Boolean> = _isPatternVisible.asStateFlow()
 
-    override val lockout: MutableStateFlow<AuthenticationLockoutModel?> = MutableStateFlow(null)
-
     override val hasLockoutOccurred = MutableStateFlow(false)
 
     private val _isAutoConfirmFeatureEnabled = MutableStateFlow(false)
@@ -68,8 +62,6 @@
     override val isPinEnhancedPrivacyEnabled: StateFlow<Boolean> =
         _isPinEnhancedPrivacyEnabled.asStateFlow()
 
-    private var failedAttemptCount = 0
-    private var lockoutEndTimestamp = 0L
     private var credentialOverride: List<Any>? = null
     private var securityMode: SecurityMode = DEFAULT_AUTHENTICATION_METHOD.toSecurityMode()
 
@@ -89,11 +81,27 @@
     }
 
     override suspend fun reportAuthenticationAttempt(isSuccessful: Boolean) {
-        failedAttemptCount = if (isSuccessful) 0 else failedAttemptCount + 1
-        authenticationChallengeResult.emit(isSuccessful)
+        if (isSuccessful) {
+            _failedAuthenticationAttempts.value = 0
+            _lockoutEndTimestamp = null
+            hasLockoutOccurred.value = false
+            lockoutStartedReportCount = 0
+        } else {
+            _failedAuthenticationAttempts.value++
+        }
     }
 
+    private var _failedAuthenticationAttempts = MutableStateFlow(0)
+    override val failedAuthenticationAttempts: StateFlow<Int> =
+        _failedAuthenticationAttempts.asStateFlow()
+
+    private var _lockoutEndTimestamp: Long? = null
+    override val lockoutEndTimestamp: Long?
+        get() = if (currentTime() < (_lockoutEndTimestamp ?: 0)) _lockoutEndTimestamp else null
+
     override suspend fun reportLockoutStarted(durationMs: Int) {
+        _lockoutEndTimestamp = (currentTime() + durationMs).takeIf { durationMs > 0 }
+        hasLockoutOccurred.value = true
         lockoutStartedReportCount++
     }
 
@@ -101,25 +109,10 @@
         return (credentialOverride ?: DEFAULT_PIN).size
     }
 
-    override suspend fun getFailedAuthenticationAttemptCount(): Int {
-        return failedAttemptCount
-    }
-
-    override suspend fun getLockoutEndTimestamp(): Long {
-        return lockoutEndTimestamp
-    }
-
     fun setAutoConfirmFeatureEnabled(isEnabled: Boolean) {
         _isAutoConfirmFeatureEnabled.value = isEnabled
     }
 
-    override suspend fun setLockoutDuration(durationMs: Int) {
-        lockoutEndTimestamp = if (durationMs > 0) currentTime() + durationMs else 0
-        if (durationMs > 0) {
-            hasLockoutOccurred.value = true
-        }
-    }
-
     override suspend fun checkCredential(
         credential: LockscreenCredential
     ): AuthenticationResultModel {
@@ -136,8 +129,8 @@
                 else -> error("Unexpected credential type ${credential.type}!")
             }
 
-        return if (isSuccessful || failedAttemptCount < MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT - 1) {
-            hasLockoutOccurred.value = false
+        val failedAttempts = _failedAuthenticationAttempts.value
+        return if (isSuccessful || failedAttempts < MAX_FAILED_AUTH_TRIES_BEFORE_LOCKOUT - 1) {
             AuthenticationResultModel(
                 isSuccessful = isSuccessful,
                 lockoutDurationMs = 0,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/authentication/domain/interactor/AuthenticationInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/authentication/domain/interactor/AuthenticationInteractorKosmos.kt
index 060ca4c..05cb059 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/authentication/domain/interactor/AuthenticationInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/authentication/domain/interactor/AuthenticationInteractorKosmos.kt
@@ -19,17 +19,11 @@
 import com.android.systemui.authentication.data.repository.authenticationRepository
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.applicationCoroutineScope
-import com.android.systemui.kosmos.testDispatcher
-import com.android.systemui.user.data.repository.userRepository
-import com.android.systemui.util.time.fakeSystemClock
 
 val Kosmos.authenticationInteractor by
     Kosmos.Fixture {
         AuthenticationInteractor(
             applicationScope = applicationCoroutineScope,
             repository = authenticationRepository,
-            backgroundDispatcher = testDispatcher,
-            userRepository = userRepository,
-            clock = fakeSystemClock,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneTestUtils.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneTestUtils.kt
index 0b41926..25b97b3 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneTestUtils.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneTestUtils.kt
@@ -91,7 +91,6 @@
 import com.android.systemui.telephony.data.repository.TelephonyRepository
 import com.android.systemui.telephony.domain.interactor.TelephonyInteractor
 import com.android.systemui.user.data.repository.FakeUserRepository
-import com.android.systemui.user.data.repository.UserRepository
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor
 import com.android.systemui.user.ui.viewmodel.UserActionViewModel
 import com.android.systemui.user.ui.viewmodel.UserViewModel
@@ -174,7 +173,7 @@
             mobileConnectionsRepository = mobileConnectionsRepository,
         )
 
-    val userRepository: UserRepository by lazy {
+    val userRepository: FakeUserRepository by lazy {
         FakeUserRepository().apply {
             val users = listOf(UserInfo(/* id=  */ 0, "name", /* flags= */ 0))
             setUserInfos(users)
@@ -236,9 +235,6 @@
         return AuthenticationInteractor(
             applicationScope = applicationScope(),
             repository = repository,
-            backgroundDispatcher = testDispatcher,
-            userRepository = userRepository,
-            clock = mock { whenever(elapsedRealtime()).thenAnswer { testScope.currentTime } }
         )
     }
 
@@ -274,7 +270,6 @@
             repository = bouncerRepository,
             authenticationInteractor = authenticationInteractor,
             keyguardFaceAuthInteractor = keyguardFaceAuthInteractor,
-            flags = sceneContainerFlags,
             falsingInteractor = falsingInteractor(),
             powerInteractor = powerInteractor(),
             simBouncerInteractor = simBouncerInteractor,
@@ -312,6 +307,7 @@
             userSwitcherMenu = flowOf(createMenuActions()),
             actionButtonInteractor = actionButtonInteractor,
             simBouncerInteractor = simBouncerInteractor,
+            clock = mock { whenever(elapsedRealtime()).thenAnswer { testScope.currentTime } },
         )
     }
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractorKosmos.kt
index 5c8fe0d..2d2f546 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractorKosmos.kt
@@ -24,6 +24,7 @@
 import com.android.systemui.statusbar.data.repository.notificationListenerSettingsRepository
 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.wm.shell.bubbles.bubblesOptional
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 
@@ -43,6 +44,7 @@
     NotificationIconsInteractor(
         activeNotificationsInteractor = activeNotificationsInteractor,
         bubbles = bubblesOptional,
+        headsUpNotificationIconInteractor = headsUpNotificationIconInteractor,
         keyguardViewStateRepository = notificationsKeyguardViewStateRepository,
     )
 }
diff --git a/services/companion/java/com/android/server/companion/transport/CompanionTransportManager.java b/services/companion/java/com/android/server/companion/transport/CompanionTransportManager.java
index df74770..62c6703 100644
--- a/services/companion/java/com/android/server/companion/transport/CompanionTransportManager.java
+++ b/services/companion/java/com/android/server/companion/transport/CompanionTransportManager.java
@@ -130,7 +130,7 @@
         synchronized (mTransports) {
             for (int i = 0; i < associationIds.length; i++) {
                 if (mTransports.contains(associationIds[i])) {
-                    mTransports.get(associationIds[i]).requestForResponse(message, data);
+                    mTransports.get(associationIds[i]).sendMessage(message, data);
                 }
             }
         }
@@ -220,7 +220,7 @@
             if (transport == null) {
                 return CompletableFuture.failedFuture(new IOException("Missing transport"));
             }
-            return transport.requestForResponse(MESSAGE_REQUEST_PERMISSION_RESTORE, data);
+            return transport.sendMessage(MESSAGE_REQUEST_PERMISSION_RESTORE, data);
         }
     }
 
diff --git a/services/companion/java/com/android/server/companion/transport/Transport.java b/services/companion/java/com/android/server/companion/transport/Transport.java
index 32d4061..22b18ac 100644
--- a/services/companion/java/com/android/server/companion/transport/Transport.java
+++ b/services/companion/java/com/android/server/companion/transport/Transport.java
@@ -16,6 +16,9 @@
 
 package com.android.server.companion.transport;
 
+import static android.companion.CompanionDeviceManager.MESSAGE_ONEWAY_FROM_WEARABLE;
+import static android.companion.CompanionDeviceManager.MESSAGE_ONEWAY_PING;
+import static android.companion.CompanionDeviceManager.MESSAGE_ONEWAY_TO_WEARABLE;
 import static android.companion.CompanionDeviceManager.MESSAGE_REQUEST_CONTEXT_SYNC;
 import static android.companion.CompanionDeviceManager.MESSAGE_REQUEST_PERMISSION_RESTORE;
 import static android.companion.CompanionDeviceManager.MESSAGE_REQUEST_PING;
@@ -80,6 +83,10 @@
         return (message & 0xFF000000) == 0x33000000;
     }
 
+    private static boolean isOneway(int message) {
+        return (message & 0xFF000000) == 0x43000000;
+    }
+
     @GuardedBy("mPendingRequests")
     protected final SparseArray<CompletableFuture<byte[]>> mPendingRequests =
             new SparseArray<>();
@@ -134,6 +141,42 @@
     protected abstract void sendMessage(int message, int sequence, @NonNull byte[] data)
             throws IOException;
 
+    /**
+     * Send a message using this transport. If the message was a request, then the returned Future
+     * object will complete successfully only if the remote device both received and processed it
+     * as expected. If the message was a send-and-forget type message, then the Future object will
+     * resolve successfully immediately (with null) upon sending the message.
+     *
+     * @param message the message type
+     * @param data the message payload
+     * @return Future object containing the result of the sent message.
+     */
+    public Future<byte[]> sendMessage(int message, byte[] data) {
+        final CompletableFuture<byte[]> pending = new CompletableFuture<>();
+        if (isOneway(message)) {
+            return sendAndForget(message, data);
+        } else if (isRequest(message)) {
+            return requestForResponse(message, data);
+        } else {
+            Slog.w(TAG, "Failed to send message 0x" + Integer.toHexString(message));
+            pending.completeExceptionally(new IllegalArgumentException(
+                    "The message being sent must be either a one-way or a request."
+            ));
+        }
+        return pending;
+    }
+
+    /**
+     * @deprecated Method was renamed to sendMessage(int, byte[]) to support both
+     * send-and-forget type messages as well as wait-for-response type messages.
+     *
+     * @param message request message type
+     * @param data the message payload
+     * @return future object containing the result of the request.
+     *
+     * @see #sendMessage(int, byte[])
+     */
+    @Deprecated
     public Future<byte[]> requestForResponse(int message, byte[] data) {
         if (DEBUG) Slog.d(TAG, "Requesting for response");
         final int sequence = mNextSequence.incrementAndGet();
@@ -154,6 +197,20 @@
         return pending;
     }
 
+    private Future<byte[]> sendAndForget(int message, byte[]data) {
+        if (DEBUG) Slog.d(TAG, "Sending a one-way message");
+        final CompletableFuture<byte[]> pending = new CompletableFuture<>();
+
+        try {
+            sendMessage(message, -1, data);
+            pending.complete(null);
+        } catch (IOException e) {
+            pending.completeExceptionally(e);
+        }
+
+        return pending;
+    }
+
     protected final void handleMessage(int message, int sequence, @NonNull byte[] data)
             throws IOException {
         if (DEBUG) {
@@ -162,7 +219,9 @@
                     + " from association " + mAssociationId);
         }
 
-        if (isRequest(message)) {
+        if (isOneway(message)) {
+            processOneway(message, data);
+        } else if (isRequest(message)) {
             try {
                 processRequest(message, sequence, data);
             } catch (IOException e) {
@@ -175,6 +234,21 @@
         }
     }
 
+    private void processOneway(int message, byte[] data) {
+        switch (message) {
+            case MESSAGE_ONEWAY_PING:
+            case MESSAGE_ONEWAY_FROM_WEARABLE:
+            case MESSAGE_ONEWAY_TO_WEARABLE: {
+                callback(message, data);
+                break;
+            }
+            default: {
+                Slog.w(TAG, "Ignoring unknown message 0x" + Integer.toHexString(message));
+                break;
+            }
+        }
+    }
+
     private void processRequest(int message, int sequence, byte[] data)
             throws IOException {
         switch (message) {
diff --git a/services/core/java/com/android/server/display/AutomaticBrightnessController.java b/services/core/java/com/android/server/display/AutomaticBrightnessController.java
index 2314bb7..3024dd2 100644
--- a/services/core/java/com/android/server/display/AutomaticBrightnessController.java
+++ b/services/core/java/com/android/server/display/AutomaticBrightnessController.java
@@ -16,6 +16,8 @@
 
 package com.android.server.display;
 
+import static com.android.server.display.config.DisplayBrightnessMappingConfig.autoBrightnessModeToString;
+
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -72,12 +74,14 @@
     @IntDef(prefix = { "AUTO_BRIGHTNESS_MODE_" }, value = {
             AUTO_BRIGHTNESS_MODE_DEFAULT,
             AUTO_BRIGHTNESS_MODE_IDLE,
+            AUTO_BRIGHTNESS_MODE_DOZE
     })
     @Retention(RetentionPolicy.SOURCE)
     public @interface AutomaticBrightnessMode{}
 
     public static final int AUTO_BRIGHTNESS_MODE_DEFAULT = 0;
     public static final int AUTO_BRIGHTNESS_MODE_IDLE = 1;
+    public static final int AUTO_BRIGHTNESS_MODE_DOZE = 2;
 
     // How long the current sensor reading is assumed to be valid beyond the current time.
     // This provides a bit of prediction, as well as ensures that the weight for the last sample is
@@ -616,12 +620,13 @@
         pw.println("  mPendingForegroundAppPackageName=" + mPendingForegroundAppPackageName);
         pw.println("  mForegroundAppCategory=" + mForegroundAppCategory);
         pw.println("  mPendingForegroundAppCategory=" + mPendingForegroundAppCategory);
-        pw.println("  Current mode=" + mCurrentBrightnessMapper.getMode());
+        pw.println("  Current mode="
+                + autoBrightnessModeToString(mCurrentBrightnessMapper.getMode()));
 
         pw.println();
         for (int i = 0; i < mBrightnessMappingStrategyMap.size(); i++) {
-            pw.println("  Mapper for mode " + modeToString(mBrightnessMappingStrategyMap.keyAt(i))
-                    + "=");
+            pw.println("  Mapper for mode "
+                    + autoBrightnessModeToString(mBrightnessMappingStrategyMap.keyAt(i)) + "=");
             mBrightnessMappingStrategyMap.valueAt(i).dump(pw,
                     mBrightnessRangeController.getNormalBrightnessMax());
         }
@@ -1224,14 +1229,6 @@
         }
     }
 
-    private String modeToString(@AutomaticBrightnessMode int mode) {
-        return switch (mode) {
-            case AUTO_BRIGHTNESS_MODE_DEFAULT -> "default";
-            case AUTO_BRIGHTNESS_MODE_IDLE -> "idle";
-            default -> Integer.toString(mode);
-        };
-    }
-
     private class ShortTermModel {
         // When the short term model is invalidated, we don't necessarily reset it (i.e. clear the
         // user's adjustment) immediately, but wait for a drastic enough change in the ambient
diff --git a/services/core/java/com/android/server/display/BrightnessMappingStrategy.java b/services/core/java/com/android/server/display/BrightnessMappingStrategy.java
index f2ffd4d..6a4b00f 100644
--- a/services/core/java/com/android/server/display/BrightnessMappingStrategy.java
+++ b/services/core/java/com/android/server/display/BrightnessMappingStrategy.java
@@ -19,6 +19,7 @@
 import static android.text.TextUtils.formatSimple;
 
 import static com.android.server.display.AutomaticBrightnessController.AUTO_BRIGHTNESS_MODE_DEFAULT;
+import static com.android.server.display.AutomaticBrightnessController.AUTO_BRIGHTNESS_MODE_DOZE;
 import static com.android.server.display.AutomaticBrightnessController.AUTO_BRIGHTNESS_MODE_IDLE;
 
 import android.annotation.Nullable;
@@ -98,8 +99,8 @@
         switch (mode) {
             case AUTO_BRIGHTNESS_MODE_DEFAULT -> {
                 brightnessLevelsNits = displayDeviceConfig.getAutoBrightnessBrighteningLevelsNits();
-                luxLevels = displayDeviceConfig.getAutoBrightnessBrighteningLevelsLux();
-                brightnessLevels = displayDeviceConfig.getAutoBrightnessBrighteningLevels();
+                luxLevels = displayDeviceConfig.getAutoBrightnessBrighteningLevelsLux(mode);
+                brightnessLevels = displayDeviceConfig.getAutoBrightnessBrighteningLevels(mode);
             }
             case AUTO_BRIGHTNESS_MODE_IDLE -> {
                 brightnessLevelsNits = getFloatArray(resources.obtainTypedArray(
@@ -107,6 +108,10 @@
                 luxLevels = getLuxLevels(resources.getIntArray(
                         com.android.internal.R.array.config_autoBrightnessLevelsIdle));
             }
+            case AUTO_BRIGHTNESS_MODE_DOZE -> {
+                luxLevels = displayDeviceConfig.getAutoBrightnessBrighteningLevelsLux(mode);
+                brightnessLevels = displayDeviceConfig.getAutoBrightnessBrighteningLevels(mode);
+            }
         }
 
         // Display independent, mode independent values
diff --git a/services/core/java/com/android/server/display/DisplayDeviceConfig.java b/services/core/java/com/android/server/display/DisplayDeviceConfig.java
index 7d22a87..a6f42d7 100644
--- a/services/core/java/com/android/server/display/DisplayDeviceConfig.java
+++ b/services/core/java/com/android/server/display/DisplayDeviceConfig.java
@@ -1591,25 +1591,29 @@
     }
 
     /**
-     * @return The default auto-brightness brightening ambient lux levels
+     * @param mode The auto-brightness mode
+     * @return The default auto-brightness brightening ambient lux levels for the specified mode
+     * and the normal brightness preset
      */
-    public float[] getAutoBrightnessBrighteningLevelsLux() {
+    public float[] getAutoBrightnessBrighteningLevelsLux(
+            @AutomaticBrightnessController.AutomaticBrightnessMode int mode) {
         if (mDisplayBrightnessMapping == null) {
             return null;
         }
-        return mDisplayBrightnessMapping.getLuxArray();
+        return mDisplayBrightnessMapping.getLuxArray(mode);
     }
 
     /**
      * @param mode The auto-brightness mode
-     * @param setting The brightness setting
-     * @return Auto brightness brightening ambient lux levels for the specified mode and setting
+     * @param preset The brightness preset. Presets are used on devices that allow users to choose
+     *               from a set of predefined options in display auto-brightness settings.
+     * @return Auto brightness brightening ambient lux levels for the specified mode and preset
      */
-    public float[] getAutoBrightnessBrighteningLevelsLux(String mode, String setting) {
+    public float[] getAutoBrightnessBrighteningLevelsLux(String mode, String preset) {
         if (mDisplayBrightnessMapping == null) {
             return null;
         }
-        return mDisplayBrightnessMapping.getLuxArray(mode, setting);
+        return mDisplayBrightnessMapping.getLuxArray(mode, preset);
     }
 
     /**
@@ -1623,25 +1627,29 @@
     }
 
     /**
-     * @return The default auto-brightness brightening levels
+     * @param mode The auto-brightness mode
+     * @return The default auto-brightness brightening levels for the specified mode and the normal
+     * brightness preset
      */
-    public float[] getAutoBrightnessBrighteningLevels() {
+    public float[] getAutoBrightnessBrighteningLevels(
+            @AutomaticBrightnessController.AutomaticBrightnessMode int mode) {
         if (mDisplayBrightnessMapping == null) {
             return null;
         }
-        return mDisplayBrightnessMapping.getBrightnessArray();
+        return mDisplayBrightnessMapping.getBrightnessArray(mode);
     }
 
     /**
      * @param mode The auto-brightness mode
-     * @param setting The brightness setting
-     * @return Auto brightness brightening backlight levels for the specified mode and setting
+     * @param preset The brightness preset. Presets are used on devices that allow users to choose
+     *               from a set of predefined options in display auto-brightness settings.
+     * @return Auto brightness brightening backlight levels for the specified mode and preset
      */
-    public float[] getAutoBrightnessBrighteningLevels(String mode, String setting) {
+    public float[] getAutoBrightnessBrighteningLevels(String mode, String preset) {
         if (mDisplayBrightnessMapping == null) {
             return null;
         }
-        return mDisplayBrightnessMapping.getBrightnessArray(mode, setting);
+        return mDisplayBrightnessMapping.getBrightnessArray(mode, preset);
     }
 
     /**
diff --git a/services/core/java/com/android/server/display/DisplayPowerController2.java b/services/core/java/com/android/server/display/DisplayPowerController2.java
index 6d09cc9..c088a6d 100644
--- a/services/core/java/com/android/server/display/DisplayPowerController2.java
+++ b/services/core/java/com/android/server/display/DisplayPowerController2.java
@@ -17,6 +17,7 @@
 package com.android.server.display;
 
 import static com.android.server.display.AutomaticBrightnessController.AUTO_BRIGHTNESS_MODE_DEFAULT;
+import static com.android.server.display.AutomaticBrightnessController.AUTO_BRIGHTNESS_MODE_DOZE;
 import static com.android.server.display.AutomaticBrightnessController.AUTO_BRIGHTNESS_MODE_IDLE;
 
 import android.animation.Animator;
@@ -1006,6 +1007,13 @@
             }
         }
 
+        BrightnessMappingStrategy dozeModeBrightnessMapper =
+                BrightnessMappingStrategy.create(resources, mDisplayDeviceConfig,
+                        AUTO_BRIGHTNESS_MODE_DOZE, mDisplayWhiteBalanceController);
+        if (mFlags.areAutoBrightnessModesEnabled() && dozeModeBrightnessMapper != null) {
+            brightnessMappers.put(AUTO_BRIGHTNESS_MODE_DOZE, dozeModeBrightnessMapper);
+        }
+
         float userLux = BrightnessMappingStrategy.INVALID_LUX;
         float userNits = BrightnessMappingStrategy.INVALID_NITS;
         if (mAutomaticBrightnessController != null) {
@@ -1349,6 +1357,13 @@
         animateScreenStateChange(state, mDisplayStateController.shouldPerformScreenOffTransition());
         state = mPowerState.getScreenState();
 
+        // Switch to doze auto-brightness mode if needed
+        if (mFlags.areAutoBrightnessModesEnabled() && mAutomaticBrightnessController != null
+                && !mAutomaticBrightnessController.isInIdleMode()) {
+            setAutomaticScreenBrightnessMode(Display.isDozeState(state)
+                    ? AUTO_BRIGHTNESS_MODE_DOZE : AUTO_BRIGHTNESS_MODE_DEFAULT);
+        }
+
         final boolean userSetBrightnessChanged = mDisplayBrightnessController
                 .updateUserSetScreenBrightness();
 
diff --git a/services/core/java/com/android/server/display/config/DisplayBrightnessMappingConfig.java b/services/core/java/com/android/server/display/config/DisplayBrightnessMappingConfig.java
index 2162850..8f12329 100644
--- a/services/core/java/com/android/server/display/config/DisplayBrightnessMappingConfig.java
+++ b/services/core/java/com/android/server/display/config/DisplayBrightnessMappingConfig.java
@@ -16,11 +16,16 @@
 
 package com.android.server.display.config;
 
+import static com.android.server.display.AutomaticBrightnessController.AUTO_BRIGHTNESS_MODE_DEFAULT;
+import static com.android.server.display.AutomaticBrightnessController.AUTO_BRIGHTNESS_MODE_DOZE;
+import static com.android.server.display.AutomaticBrightnessController.AUTO_BRIGHTNESS_MODE_IDLE;
+
 import android.content.Context;
 import android.os.PowerManager;
 import android.util.Spline;
 
 import com.android.internal.display.BrightnessSynchronizer;
+import com.android.server.display.AutomaticBrightnessController;
 import com.android.server.display.DisplayDeviceConfig;
 import com.android.server.display.feature.DisplayManagerFlags;
 
@@ -33,7 +38,9 @@
  */
 public class DisplayBrightnessMappingConfig {
 
-    private static final String DEFAULT_BRIGHTNESS_MAPPING_KEY = "default_normal";
+    private static final String DEFAULT_BRIGHTNESS_PRESET_NAME = "normal";
+    private static final String DEFAULT_BRIGHTNESS_MAPPING_KEY =
+            AutoBrightnessModeName._default.getRawName() + "_" + DEFAULT_BRIGHTNESS_PRESET_NAME;
 
     /**
      * Array of desired screen brightness in nits corresponding to the lux values
@@ -45,19 +52,22 @@
 
     /**
      * Map of arrays of desired screen brightness corresponding to the lux values
-     * in mBrightnessLevelsLuxMap, indexed by the auto-brightness mode and the brightness setting.
+     * in mBrightnessLevelsLuxMap, indexed by the auto-brightness mode and the brightness preset.
      * The brightness values must be non-negative and non-decreasing. They must be between
      * {@link PowerManager.BRIGHTNESS_MIN} and {@link PowerManager.BRIGHTNESS_MAX}.
      *
-     * The keys are a concatenation of the auto-brightness mode and the brightness setting
-     * separated by an underscore, e.g. default_normal, default_dim, default_bright, doze_normal,
-     * doze_dim, doze_bright.
+     * The keys are a concatenation of the auto-brightness mode and the brightness preset separated
+     * by an underscore, e.g. default_normal, default_dim, default_bright, doze_normal, doze_dim,
+     * doze_bright.
+     *
+     * The presets are used on devices that allow users to choose from a set of predefined options
+     * in display auto-brightness settings.
      */
     private final Map<String, float[]> mBrightnessLevelsMap = new HashMap<>();
 
     /**
      * Map of arrays of light sensor lux values to define our levels for auto-brightness support,
-     * indexed by the auto-brightness mode and the brightness setting.
+     * indexed by the auto-brightness mode and the brightness preset.
      *
      * The first lux value in every array is always 0.
      *
@@ -69,9 +79,12 @@
      * Spline interpolation is used to determine the auto-brightness values for lux levels between
      * these control points.
      *
-     * The keys are a concatenation of the auto-brightness mode and the brightness setting
-     * separated by an underscore, e.g. default_normal, default_dim, default_bright, doze_normal,
-     * doze_dim, doze_bright.
+     * The keys are a concatenation of the auto-brightness mode and the brightness preset separated
+     * by an underscore, e.g. default_normal, default_dim, default_bright, doze_normal, doze_dim,
+     * doze_bright.
+     *
+     * The presets are used on devices that allow users to choose from a set of predefined options
+     * in display auto-brightness settings.
      */
     private final Map<String, float[]> mBrightnessLevelsLuxMap = new HashMap<>();
 
@@ -138,19 +151,23 @@
     }
 
     /**
-     * @return The default auto-brightness brightening ambient lux levels
+     * @param mode The auto-brightness mode
+     * @return The default auto-brightness brightening ambient lux levels for the specified mode
+     * and the normal brightness preset
      */
-    public float[] getLuxArray() {
-        return mBrightnessLevelsLuxMap.get(DEFAULT_BRIGHTNESS_MAPPING_KEY);
+    public float[] getLuxArray(@AutomaticBrightnessController.AutomaticBrightnessMode int mode) {
+        return mBrightnessLevelsLuxMap.get(
+                autoBrightnessModeToString(mode) + "_" + DEFAULT_BRIGHTNESS_PRESET_NAME);
     }
 
     /**
      * @param mode The auto-brightness mode
-     * @param setting The brightness setting
-     * @return Auto brightness brightening ambient lux levels for the specified mode and setting
+     * @param preset The brightness preset. Presets are used on devices that allow users to choose
+     *               from a set of predefined options in display auto-brightness settings.
+     * @return Auto brightness brightening ambient lux levels for the specified mode and preset
      */
-    public float[] getLuxArray(String mode, String setting) {
-        return mBrightnessLevelsLuxMap.get(mode + "_" + setting);
+    public float[] getLuxArray(String mode, String preset) {
+        return mBrightnessLevelsLuxMap.get(mode + "_" + preset);
     }
 
     /**
@@ -161,19 +178,24 @@
     }
 
     /**
-     * @return The default auto-brightness brightening levels
+     * @param mode The auto-brightness mode
+     * @return The default auto-brightness brightening levels for the specified mode and the normal
+     * brightness preset
      */
-    public float[] getBrightnessArray() {
-        return mBrightnessLevelsMap.get(DEFAULT_BRIGHTNESS_MAPPING_KEY);
+    public float[] getBrightnessArray(
+            @AutomaticBrightnessController.AutomaticBrightnessMode int mode) {
+        return mBrightnessLevelsMap.get(
+                autoBrightnessModeToString(mode) + "_" + DEFAULT_BRIGHTNESS_PRESET_NAME);
     }
 
     /**
      * @param mode The auto-brightness mode
-     * @param setting The brightness setting
-     * @return Auto brightness brightening ambient lux levels for the specified mode and setting
+     * @param preset The brightness preset. Presets are used on devices that allow users to choose
+     *               from a set of predefined options in display auto-brightness settings.
+     * @return Auto brightness brightening ambient lux levels for the specified mode and preset
      */
-    public float[] getBrightnessArray(String mode, String setting) {
-        return mBrightnessLevelsMap.get(mode + "_" + setting);
+    public float[] getBrightnessArray(String mode, String preset) {
+        return mBrightnessLevelsMap.get(mode + "_" + preset);
     }
 
     @Override
@@ -205,6 +227,26 @@
                 + ", mBrightnessLevelsMap= " + brightnessLevelsMapString;
     }
 
+    /**
+     * @param mode The auto-brightness mode
+     * @return The string representing the mode
+     */
+    public static String autoBrightnessModeToString(
+            @AutomaticBrightnessController.AutomaticBrightnessMode int mode) {
+        switch (mode) {
+            case AUTO_BRIGHTNESS_MODE_DEFAULT -> {
+                return AutoBrightnessModeName._default.getRawName();
+            }
+            case AUTO_BRIGHTNESS_MODE_IDLE -> {
+                return AutoBrightnessModeName.idle.getRawName();
+            }
+            case AUTO_BRIGHTNESS_MODE_DOZE -> {
+                return AutoBrightnessModeName.doze.getRawName();
+            }
+            default -> throw new IllegalArgumentException("Unknown auto-brightness mode: " + mode);
+        }
+    }
+
     private float[] brightnessArrayIntToFloat(int[] brightnessInt,
             Spline backlightToBrightnessSpline) {
         float[] brightnessFloat = new float[brightnessInt.length];
diff --git a/services/core/java/com/android/server/location/gnss/GnssNetworkConnectivityHandler.java b/services/core/java/com/android/server/location/gnss/GnssNetworkConnectivityHandler.java
index 78c8cde..403b421 100644
--- a/services/core/java/com/android/server/location/gnss/GnssNetworkConnectivityHandler.java
+++ b/services/core/java/com/android/server/location/gnss/GnssNetworkConnectivityHandler.java
@@ -19,6 +19,7 @@
 import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED;
 
 import android.content.Context;
+import android.location.flags.Flags;
 import android.net.ConnectivityManager;
 import android.net.LinkAddress;
 import android.net.LinkProperties;
@@ -48,6 +49,7 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Handles network connection requests and network state change updates for AGPS data download.
@@ -91,6 +93,10 @@
     // network with SUPL connectivity or report an error.
     private static final int SUPL_NETWORK_REQUEST_TIMEOUT_MILLIS = 20 * 1000;
 
+    // If the chipset does not request to release a SUPL connection before the specified timeout in
+    // milliseconds, the connection will be automatically released.
+    private static final long SUPL_CONNECTION_TIMEOUT_MILLIS = TimeUnit.MINUTES.toMillis(1);
+
     private static final int HASH_MAP_INITIAL_CAPACITY_TO_TRACK_CONNECTED_NETWORKS = 5;
 
     // Keeps track of networks and their state as notified by the network request callbacks.
@@ -121,6 +127,8 @@
     private static final long WAKELOCK_TIMEOUT_MILLIS = 60 * 1000;
     private final PowerManager.WakeLock mWakeLock;
 
+    private final Object mSuplConnectionReleaseOnTimeoutToken = new Object();
+
     /**
      * Network attributes needed when updating HAL about network connectivity status changes.
      */
@@ -609,6 +617,13 @@
                     mSuplConnectivityCallback,
                     mHandler,
                     SUPL_NETWORK_REQUEST_TIMEOUT_MILLIS);
+            if (Flags.releaseSuplConnectionOnTimeout()) {
+                // Schedule to release the SUPL connection after timeout
+                mHandler.removeCallbacksAndMessages(mSuplConnectionReleaseOnTimeoutToken);
+                mHandler.postDelayed(() -> handleReleaseSuplConnection(GPS_RELEASE_AGPS_DATA_CONN),
+                        mSuplConnectionReleaseOnTimeoutToken,
+                        SUPL_CONNECTION_TIMEOUT_MILLIS);
+            }
         } catch (RuntimeException e) {
             Log.e(TAG, "Failed to request network.", e);
             mSuplConnectivityCallback = null;
@@ -639,6 +654,10 @@
             Log.d(TAG, message);
         }
 
+        if (Flags.releaseSuplConnectionOnTimeout()) {
+            // Remove pending task to avoid releasing an incorrect connection
+            mHandler.removeCallbacksAndMessages(mSuplConnectionReleaseOnTimeoutToken);
+        }
         if (mAGpsDataConnectionState == AGPS_DATA_CONNECTION_CLOSED) {
             return;
         }
diff --git a/services/core/xsd/display-device-config/display-device-config.xsd b/services/core/xsd/display-device-config/display-device-config.xsd
index adbd3c9..8078745 100644
--- a/services/core/xsd/display-device-config/display-device-config.xsd
+++ b/services/core/xsd/display-device-config/display-device-config.xsd
@@ -631,7 +631,7 @@
             <xs:annotation name="nonnull"/>
             <xs:annotation name="final"/>
         </xs:element>
-        <xs:element name="mode" type="xs:string" minOccurs="0"/>
+        <xs:element name="mode" type="AutoBrightnessModeName" minOccurs="0"/>
         <xs:element name="setting" type="xs:string" minOccurs="0"/>
     </xs:complexType>
 
@@ -765,4 +765,14 @@
             </xs:element>
         </xs:sequence>
     </xs:complexType>
+
+    <!-- Predefined type names as defined by
+    AutomaticBrightnessController.AutomaticBrightnessMode -->
+    <xs:simpleType  name="AutoBrightnessModeName">
+        <xs:restriction base="xs:string">
+            <xs:enumeration value="default"/>
+            <xs:enumeration value="idle"/>
+            <xs:enumeration value="doze"/>
+        </xs:restriction>
+    </xs:simpleType>
 </xs:schema>
diff --git a/services/core/xsd/display-device-config/schema/current.txt b/services/core/xsd/display-device-config/schema/current.txt
index 98c95ed..91172a3 100644
--- a/services/core/xsd/display-device-config/schema/current.txt
+++ b/services/core/xsd/display-device-config/schema/current.txt
@@ -16,6 +16,13 @@
     method public void setEnabled(boolean);
   }
 
+  public enum AutoBrightnessModeName {
+    method public String getRawName();
+    enum_constant public static final com.android.server.display.config.AutoBrightnessModeName _default;
+    enum_constant public static final com.android.server.display.config.AutoBrightnessModeName doze;
+    enum_constant public static final com.android.server.display.config.AutoBrightnessModeName idle;
+  }
+
   public class BlockingZoneConfig {
     ctor public BlockingZoneConfig();
     method public final com.android.server.display.config.BlockingZoneThreshold getBlockingZoneThreshold();
@@ -219,10 +226,10 @@
   public class LuxToBrightnessMapping {
     ctor public LuxToBrightnessMapping();
     method @NonNull public final com.android.server.display.config.NonNegativeFloatToFloatMap getMap();
-    method public String getMode();
+    method public com.android.server.display.config.AutoBrightnessModeName getMode();
     method public String getSetting();
     method public final void setMap(@NonNull com.android.server.display.config.NonNegativeFloatToFloatMap);
-    method public void setMode(String);
+    method public void setMode(com.android.server.display.config.AutoBrightnessModeName);
     method public void setSetting(String);
   }
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/BrightnessMappingStrategyTest.java b/services/tests/displayservicetests/src/com/android/server/display/BrightnessMappingStrategyTest.java
index c5a1ba1..f4eaa5b 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/BrightnessMappingStrategyTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/BrightnessMappingStrategyTest.java
@@ -695,9 +695,11 @@
             mDdc = mock(DisplayDeviceConfig.class);
             when(mDdc.getNits()).thenReturn(DISPLAY_RANGE_NITS);
             when(mDdc.getBrightness()).thenReturn(DISPLAY_LEVELS_RANGE_BACKLIGHT_FLOAT);
-            when(mDdc.getAutoBrightnessBrighteningLevelsLux()).thenReturn(LUX_LEVELS);
+            when(mDdc.getAutoBrightnessBrighteningLevelsLux(AUTO_BRIGHTNESS_MODE_DEFAULT))
+                    .thenReturn(LUX_LEVELS);
             when(mDdc.getAutoBrightnessBrighteningLevelsNits()).thenReturn(EMPTY_FLOAT_ARRAY);
-            when(mDdc.getAutoBrightnessBrighteningLevels()).thenReturn(EMPTY_FLOAT_ARRAY);
+            when(mDdc.getAutoBrightnessBrighteningLevels(AUTO_BRIGHTNESS_MODE_DEFAULT))
+                    .thenReturn(EMPTY_FLOAT_ARRAY);
         }
 
         DdcBuilder setNitsRange(float[] nitsArray) {
@@ -711,7 +713,8 @@
         }
 
         DdcBuilder setAutoBrightnessLevelsLux(float[] luxLevels) {
-            when(mDdc.getAutoBrightnessBrighteningLevelsLux()).thenReturn(luxLevels);
+            when(mDdc.getAutoBrightnessBrighteningLevelsLux(AUTO_BRIGHTNESS_MODE_DEFAULT))
+                    .thenReturn(luxLevels);
             return this;
         }
 
@@ -721,7 +724,8 @@
         }
 
         DdcBuilder setAutoBrightnessLevels(float[] brightnessLevels) {
-            when(mDdc.getAutoBrightnessBrighteningLevels()).thenReturn(brightnessLevels);
+            when(mDdc.getAutoBrightnessBrighteningLevels(AUTO_BRIGHTNESS_MODE_DEFAULT))
+                    .thenReturn(brightnessLevels);
             return this;
         }
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/DisplayDeviceConfigTest.java b/services/tests/displayservicetests/src/com/android/server/display/DisplayDeviceConfigTest.java
index a4c15b5..61c6076 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/DisplayDeviceConfigTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/DisplayDeviceConfigTest.java
@@ -18,6 +18,7 @@
 
 
 import static com.android.internal.display.BrightnessSynchronizer.brightnessIntToFloat;
+import static com.android.server.display.AutomaticBrightnessController.AUTO_BRIGHTNESS_MODE_DEFAULT;
 import static com.android.server.display.config.SensorData.SupportedMode;
 import static com.android.server.display.utils.DeviceConfigParsingUtils.ambientBrightnessThresholdsIntToFloat;
 import static com.android.server.display.utils.DeviceConfigParsingUtils.displayBrightnessThresholdsIntToFloat;
@@ -605,13 +606,13 @@
 
     private void verifyConfigValuesFromConfigResource() {
         assertNull(mDisplayDeviceConfig.getName());
-        assertArrayEquals(mDisplayDeviceConfig.getAutoBrightnessBrighteningLevelsNits(), new
-                float[]{2.0f, 200.0f, 600.0f}, ZERO_DELTA);
-        assertArrayEquals(mDisplayDeviceConfig.getAutoBrightnessBrighteningLevelsLux(), new
-                float[]{0.0f, 110.0f, 500.0f}, ZERO_DELTA);
-        assertArrayEquals(mDisplayDeviceConfig.getAutoBrightnessBrighteningLevels(), new
-                float[]{brightnessIntToFloat(50), brightnessIntToFloat(100),
-                brightnessIntToFloat(150)}, SMALL_DELTA);
+        assertArrayEquals(mDisplayDeviceConfig.getAutoBrightnessBrighteningLevelsNits(),
+                new float[]{2.0f, 200.0f, 600.0f}, ZERO_DELTA);
+        assertArrayEquals(mDisplayDeviceConfig.getAutoBrightnessBrighteningLevelsLux(
+                AUTO_BRIGHTNESS_MODE_DEFAULT), new float[]{0.0f, 110.0f, 500.0f}, ZERO_DELTA);
+        assertArrayEquals(mDisplayDeviceConfig.getAutoBrightnessBrighteningLevels(
+                AUTO_BRIGHTNESS_MODE_DEFAULT), new float[]{brightnessIntToFloat(50),
+                brightnessIntToFloat(100), brightnessIntToFloat(150)}, SMALL_DELTA);
 
         // Test thresholds
         assertEquals(0, mDisplayDeviceConfig.getAmbientLuxBrighteningMinThreshold(), ZERO_DELTA);
@@ -737,9 +738,11 @@
                 getValidProxSensor(), /* includeIdleMode= */ false));
 
         assertArrayEquals(new float[]{0.0f, 80},
-                mDisplayDeviceConfig.getAutoBrightnessBrighteningLevelsLux(), ZERO_DELTA);
+                mDisplayDeviceConfig.getAutoBrightnessBrighteningLevelsLux(
+                        AUTO_BRIGHTNESS_MODE_DEFAULT), ZERO_DELTA);
         assertArrayEquals(new float[]{0.2f, 0.3f},
-                mDisplayDeviceConfig.getAutoBrightnessBrighteningLevels(), SMALL_DELTA);
+                mDisplayDeviceConfig.getAutoBrightnessBrighteningLevels(
+                        AUTO_BRIGHTNESS_MODE_DEFAULT), SMALL_DELTA);
 
         assertArrayEquals(new float[]{0.0f, 90},
                 mDisplayDeviceConfig.getAutoBrightnessBrighteningLevelsLux("default", "dim"),
@@ -772,9 +775,11 @@
 
         assertArrayEquals(new float[]{brightnessIntToFloat(50), brightnessIntToFloat(100),
                         brightnessIntToFloat(150)},
-                mDisplayDeviceConfig.getAutoBrightnessBrighteningLevels(), SMALL_DELTA);
+                mDisplayDeviceConfig.getAutoBrightnessBrighteningLevels(
+                        AUTO_BRIGHTNESS_MODE_DEFAULT), SMALL_DELTA);
         assertArrayEquals(new float[]{0, 110, 500},
-                mDisplayDeviceConfig.getAutoBrightnessBrighteningLevelsLux(), ZERO_DELTA);
+                mDisplayDeviceConfig.getAutoBrightnessBrighteningLevelsLux(
+                        AUTO_BRIGHTNESS_MODE_DEFAULT), ZERO_DELTA);
         assertArrayEquals(new float[]{2, 200, 600},
                 mDisplayDeviceConfig.getAutoBrightnessBrighteningLevelsNits(), SMALL_DELTA);
     }
diff --git a/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerController2Test.java b/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerController2Test.java
index 02bd35a..ffdc8b4 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerController2Test.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerController2Test.java
@@ -18,6 +18,8 @@
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doAnswer;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
+import static com.android.server.display.AutomaticBrightnessController.AUTO_BRIGHTNESS_MODE_DEFAULT;
+import static com.android.server.display.AutomaticBrightnessController.AUTO_BRIGHTNESS_MODE_DOZE;
 import static com.android.server.display.AutomaticBrightnessController.AUTO_BRIGHTNESS_MODE_IDLE;
 
 import static org.junit.Assert.assertNotNull;
@@ -1568,6 +1570,56 @@
                 eq(BRIGHTNESS_RAMP_RATE_FAST_INCREASE), eq(false));
     }
 
+    @Test
+    public void testSwitchToDozeAutoBrightnessMode() {
+        when(mDisplayManagerFlagsMock.areAutoBrightnessModesEnabled()).thenReturn(true);
+        when(mHolder.displayPowerState.getScreenState()).thenReturn(Display.STATE_DOZE);
+
+        DisplayPowerRequest dpr = new DisplayPowerRequest();
+        dpr.policy = DisplayPowerRequest.POLICY_DOZE;
+        mHolder.dpc.requestPowerState(dpr, /* waitForNegativeProximity= */ false);
+        advanceTime(1); // Run updatePowerState
+
+        // One triggered by handleBrightnessModeChange, another triggered by requestPowerState
+        verify(mHolder.automaticBrightnessController, times(2))
+                .switchMode(AUTO_BRIGHTNESS_MODE_DOZE);
+
+        // Back to default mode
+        when(mHolder.displayPowerState.getScreenState()).thenReturn(Display.STATE_ON);
+        dpr.policy = DisplayPowerRequest.POLICY_BRIGHT;
+        mHolder.dpc.requestPowerState(dpr, /* waitForNegativeProximity= */ false);
+        advanceTime(1); // Run updatePowerState
+
+        verify(mHolder.automaticBrightnessController).switchMode(AUTO_BRIGHTNESS_MODE_DEFAULT);
+    }
+
+    @Test
+    public void testDoesNotSwitchFromIdleToDozeAutoBrightnessMode() {
+        when(mDisplayManagerFlagsMock.areAutoBrightnessModesEnabled()).thenReturn(true);
+        when(mHolder.displayPowerState.getScreenState()).thenReturn(Display.STATE_DOZE);
+        when(mHolder.automaticBrightnessController.isInIdleMode()).thenReturn(true);
+
+        DisplayPowerRequest dpr = new DisplayPowerRequest();
+        mHolder.dpc.requestPowerState(dpr, /* waitForNegativeProximity= */ false);
+        advanceTime(1); // Run updatePowerState
+
+        verify(mHolder.automaticBrightnessController, never())
+                .switchMode(AUTO_BRIGHTNESS_MODE_DOZE);
+    }
+
+    @Test
+    public void testDoesNotSwitchDozeAutoBrightnessModeIfFeatureFlagOff() {
+        when(mDisplayManagerFlagsMock.areAutoBrightnessModesEnabled()).thenReturn(false);
+        when(mHolder.displayPowerState.getScreenState()).thenReturn(Display.STATE_DOZE);
+
+        DisplayPowerRequest dpr = new DisplayPowerRequest();
+        mHolder.dpc.requestPowerState(dpr, /* waitForNegativeProximity= */ false);
+        advanceTime(1); // Run updatePowerState
+
+        verify(mHolder.automaticBrightnessController, never())
+                .switchMode(AUTO_BRIGHTNESS_MODE_DOZE);
+    }
+
     /**
      * Creates a mock and registers it to {@link LocalServices}.
      */
diff --git a/services/tests/mockingservicestests/src/com/android/server/job/controllers/FlexibilityControllerTest.java b/services/tests/mockingservicestests/src/com/android/server/job/controllers/FlexibilityControllerTest.java
index d24500d..650c473 100644
--- a/services/tests/mockingservicestests/src/com/android/server/job/controllers/FlexibilityControllerTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/job/controllers/FlexibilityControllerTest.java
@@ -28,8 +28,6 @@
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doAnswer;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
-import static com.android.server.job.JobSchedulerService.ACTIVE_INDEX;
-import static com.android.server.job.JobSchedulerService.EXEMPTED_INDEX;
 import static com.android.server.job.controllers.FlexibilityController.FcConfig.DEFAULT_FALLBACK_FLEXIBILITY_DEADLINE_MS;
 import static com.android.server.job.controllers.FlexibilityController.FcConfig.DEFAULT_UNSEEN_CONSTRAINT_GRACE_PERIOD_MS;
 import static com.android.server.job.controllers.FlexibilityController.FcConfig.KEY_APPLIED_CONSTRAINTS;
@@ -214,7 +212,6 @@
         JobStatus js = JobStatus.createFromJobInfo(
                 jobInfo, 1000, SOURCE_PACKAGE, SOURCE_USER_ID, "FCTest", testTag);
         js.enqueueTime = FROZEN_TIME;
-        js.setStandbyBucket(ACTIVE_INDEX);
         if (js.hasFlexibilityConstraint()) {
             js.setNumAppliedFlexibleConstraints(Integer.bitCount(
                     mFlexibilityController.getRelevantAppliedConstraintsLocked(js)));
@@ -850,75 +847,14 @@
     }
 
     @Test
-    public void testAllowlistedAppBypass() {
-        JobStatus jsHigh = createJobStatus("testAllowlistedAppBypass",
-                createJob(0).setPriority(JobInfo.PRIORITY_HIGH));
-        JobStatus jsDefault = createJobStatus("testAllowlistedAppBypass",
-                createJob(0).setPriority(JobInfo.PRIORITY_DEFAULT));
-        JobStatus jsLow = createJobStatus("testAllowlistedAppBypass",
-                createJob(0).setPriority(JobInfo.PRIORITY_LOW));
-        JobStatus jsMin = createJobStatus("testAllowlistedAppBypass",
-                createJob(0).setPriority(JobInfo.PRIORITY_MIN));
-        jsHigh.setStandbyBucket(EXEMPTED_INDEX);
-        jsDefault.setStandbyBucket(EXEMPTED_INDEX);
-        jsLow.setStandbyBucket(EXEMPTED_INDEX);
-        jsMin.setStandbyBucket(EXEMPTED_INDEX);
-
-        synchronized (mFlexibilityController.mLock) {
-            assertTrue(mFlexibilityController.isFlexibilitySatisfiedLocked(jsHigh));
-            assertTrue(mFlexibilityController.isFlexibilitySatisfiedLocked(jsDefault));
-            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsLow));
-            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsMin));
-        }
-    }
-
-    @Test
-    public void testForegroundAppBypass() {
-        JobStatus jsHigh = createJobStatus("testAllowlistedAppBypass",
-                createJob(0).setPriority(JobInfo.PRIORITY_HIGH));
-        JobStatus jsDefault = createJobStatus("testAllowlistedAppBypass",
-                createJob(0).setPriority(JobInfo.PRIORITY_DEFAULT));
-        JobStatus jsLow = createJobStatus("testAllowlistedAppBypass",
-                createJob(0).setPriority(JobInfo.PRIORITY_LOW));
-        JobStatus jsMin = createJobStatus("testAllowlistedAppBypass",
-                createJob(0).setPriority(JobInfo.PRIORITY_MIN));
-
-        when(mJobSchedulerService.getUidBias(mSourceUid)).thenReturn(JobInfo.BIAS_DEFAULT);
-        synchronized (mFlexibilityController.mLock) {
-            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsHigh));
-            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsDefault));
-            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsLow));
-            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsMin));
-        }
-
-        when(mJobSchedulerService.getUidBias(mSourceUid))
-                .thenReturn(JobInfo.BIAS_BOUND_FOREGROUND_SERVICE);
-        synchronized (mFlexibilityController.mLock) {
-            assertTrue(mFlexibilityController.isFlexibilitySatisfiedLocked(jsHigh));
-            assertTrue(mFlexibilityController.isFlexibilitySatisfiedLocked(jsDefault));
-            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsLow));
-            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsMin));
-        }
-
-        when(mJobSchedulerService.getUidBias(mSourceUid))
-                .thenReturn(JobInfo.BIAS_FOREGROUND_SERVICE);
-        synchronized (mFlexibilityController.mLock) {
-            assertTrue(mFlexibilityController.isFlexibilitySatisfiedLocked(jsHigh));
-            assertTrue(mFlexibilityController.isFlexibilitySatisfiedLocked(jsDefault));
-            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsLow));
-            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsMin));
-        }
-    }
-
-    @Test
     public void testTopAppBypass() {
-        JobInfo.Builder jb = createJob(0).setPriority(JobInfo.PRIORITY_MIN);
+        JobInfo.Builder jb = createJob(0);
         JobStatus js = createJobStatus("testTopAppBypass", jb);
         mJobStore.add(js);
 
         // Needed because if before and after Uid bias is the same, nothing happens.
         when(mJobSchedulerService.getUidBias(mSourceUid))
-                .thenReturn(JobInfo.BIAS_DEFAULT);
+                .thenReturn(JobInfo.BIAS_FOREGROUND_SERVICE);
 
         synchronized (mFlexibilityController.mLock) {
             mFlexibilityController.maybeStartTrackingJobLocked(js, null);
@@ -929,7 +865,7 @@
             assertTrue(mFlexibilityController.isFlexibilitySatisfiedLocked(js));
             assertTrue(js.isConstraintSatisfied(CONSTRAINT_FLEXIBLE));
 
-            setUidBias(mSourceUid, JobInfo.BIAS_SYNC_INITIALIZATION);
+            setUidBias(mSourceUid, JobInfo.BIAS_FOREGROUND_SERVICE);
 
             assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(js));
             assertFalse(js.isConstraintSatisfied(CONSTRAINT_FLEXIBLE));
diff --git a/services/tests/servicestests/src/com/android/server/biometrics/sensors/SensorOverlaysTest.java b/services/tests/servicestests/src/com/android/server/biometrics/sensors/SensorOverlaysTest.java
index 94cb860..74f8f08 100644
--- a/services/tests/servicestests/src/com/android/server/biometrics/sensors/SensorOverlaysTest.java
+++ b/services/tests/servicestests/src/com/android/server/biometrics/sensors/SensorOverlaysTest.java
@@ -29,7 +29,7 @@
 import android.hardware.fingerprint.ISidefpsController;
 import android.hardware.fingerprint.IUdfpsOverlayController;
 import android.platform.test.annotations.Presubmit;
-import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.flag.junit.SetFlagsRule;
 
 import androidx.test.filters.SmallTest;
 
@@ -43,10 +43,10 @@
 import java.util.ArrayList;
 import java.util.List;
 
-@RequiresFlagsDisabled(FLAG_SIDEFPS_CONTROLLER_REFACTOR)
 @Presubmit
 @SmallTest
 public class SensorOverlaysTest {
+    @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
 
     private static final int SENSOR_ID = 11;
     private static final long REQUEST_ID = 8;
@@ -59,6 +59,7 @@
 
     @Before
     public void setup() {
+        mSetFlagsRule.disableFlags(FLAG_SIDEFPS_CONTROLLER_REFACTOR);
         when(mAcquisitionClient.getRequestId()).thenReturn(REQUEST_ID);
         when(mAcquisitionClient.hasRequestId()).thenReturn(true);
     }
diff --git a/services/tests/servicestests/src/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintAuthenticationClientTest.java b/services/tests/servicestests/src/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintAuthenticationClientTest.java
index c24227f..774ea5b 100644
--- a/services/tests/servicestests/src/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintAuthenticationClientTest.java
+++ b/services/tests/servicestests/src/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintAuthenticationClientTest.java
@@ -161,6 +161,7 @@
 
     @Before
     public void setup() {
+        mSetFlagsRule.disableFlags(FLAG_SIDEFPS_CONTROLLER_REFACTOR);
         mContext.addMockSystemService(BiometricManager.class, mBiometricManager);
         when(mBiometricContext.getAuthSessionCoordinator()).thenReturn(mAuthSessionCoordinator);
         when(mBiometricLogger.getAmbientLightProbe(anyBoolean())).thenAnswer(i ->
diff --git a/services/tests/servicestests/src/com/android/server/usage/AppStandbyControllerTests.java b/services/tests/servicestests/src/com/android/server/usage/AppStandbyControllerTests.java
index 8487903..89056cc 100644
--- a/services/tests/servicestests/src/com/android/server/usage/AppStandbyControllerTests.java
+++ b/services/tests/servicestests/src/com/android/server/usage/AppStandbyControllerTests.java
@@ -64,17 +64,18 @@
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
-import static org.mockito.Matchers.eq;
 import static org.mockito.Matchers.intThat;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
 
+import android.annotation.DurationMillisLong;
 import android.annotation.NonNull;
 import android.app.ActivityManager;
 import android.app.usage.AppStandbyInfo;
 import android.app.usage.UsageEvents;
+import android.app.usage.UsageStatsManager;
 import android.app.usage.UsageStatsManagerInternal;
 import android.appwidget.AppWidgetManager;
 import android.content.Context;
@@ -99,7 +100,6 @@
 import android.view.Display;
 
 import androidx.test.InstrumentationRegistry;
-import androidx.test.filters.FlakyTest;
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
@@ -110,6 +110,7 @@
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -125,6 +126,7 @@
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.function.BooleanSupplier;
 import java.util.stream.Collectors;
 
 /**
@@ -175,6 +177,9 @@
     private static final int ASSERT_RETRY_ATTEMPTS = 20;
     private static final int ASSERT_RETRY_DELAY_MILLISECONDS = 500;
 
+    @DurationMillisLong
+    private static final long FLUSH_TIMEOUT_MILLISECONDS = 5_000;
+
     /** Mock variable used in {@link MyInjector#isPackageInstalled(String, int, int)} */
     private static boolean isPackageInstalled = true;
 
@@ -563,6 +568,7 @@
         mInjector = new MyInjector(myContext, Looper.getMainLooper());
         mController = setupController();
         setupInitialUsageHistory();
+        flushHandler(mController);
     }
 
     @After
@@ -571,12 +577,9 @@
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testBoundWidgetPackageExempt() throws Exception {
         assumeTrue(mInjector.getContext().getSystemService(AppWidgetManager.class) != null);
-        assertEquals(STANDBY_BUCKET_ACTIVE,
-                mController.getAppStandbyBucket(PACKAGE_EXEMPTED_1, USER_ID,
-                        mInjector.mElapsedRealtime, false));
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_EXEMPTED_1);
     }
 
     @Test
@@ -654,7 +657,6 @@
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testIsAppIdle_Charging() throws Exception {
         TestParoleListener paroleListener = new TestParoleListener();
         mController.addListener(paroleListener);
@@ -662,7 +664,7 @@
         setChargingState(mController, false);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_FORCED_BY_SYSTEM);
-        assertEquals(STANDBY_BUCKET_RARE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
         assertTrue(mController.isAppIdleFiltered(PACKAGE_1, UID_1, USER_ID, 0));
         assertTrue(mController.isAppIdleFiltered(PACKAGE_1, UID_1, USER_ID, false));
         assertFalse(mController.isInParole());
@@ -671,7 +673,7 @@
         setChargingState(mController, true);
         paroleListener.awaitOnLatch(2000);
         assertTrue(paroleListener.getParoleState());
-        assertEquals(STANDBY_BUCKET_RARE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
         assertFalse(mController.isAppIdleFiltered(PACKAGE_1, UID_1, USER_ID, 0));
         assertFalse(mController.isAppIdleFiltered(PACKAGE_1, UID_1, USER_ID, false));
         assertTrue(mController.isInParole());
@@ -680,14 +682,13 @@
         setChargingState(mController, false);
         paroleListener.awaitOnLatch(2000);
         assertFalse(paroleListener.getParoleState());
-        assertEquals(STANDBY_BUCKET_RARE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
         assertTrue(mController.isAppIdleFiltered(PACKAGE_1, UID_1, USER_ID, 0));
         assertTrue(mController.isAppIdleFiltered(PACKAGE_1, UID_1, USER_ID, false));
         assertFalse(mController.isInParole());
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testIsAppIdle_Enabled() throws Exception {
         setChargingState(mController, false);
         TestParoleListener paroleListener = new TestParoleListener();
@@ -696,7 +697,7 @@
         setAppIdleEnabled(mController, true);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_FORCED_BY_SYSTEM);
-        assertEquals(STANDBY_BUCKET_RARE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
         assertTrue(mController.isAppIdleFiltered(PACKAGE_1, UID_1, USER_ID, 0));
         assertTrue(mController.isAppIdleFiltered(PACKAGE_1, UID_1, USER_ID, false));
 
@@ -711,7 +712,7 @@
         setAppIdleEnabled(mController, true);
         paroleListener.awaitOnLatch(2000);
         assertFalse(paroleListener.getParoleState());
-        assertEquals(STANDBY_BUCKET_RARE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
         assertTrue(mController.isAppIdleFiltered(PACKAGE_1, UID_1, USER_ID, 0));
         assertTrue(mController.isAppIdleFiltered(PACKAGE_1, UID_1, USER_ID, false));
     }
@@ -723,6 +724,7 @@
     private void assertTimeout(AppStandbyController controller, long elapsedTime, int bucket,
             int userId) {
         mInjector.mElapsedRealtime = elapsedTime;
+        flushHandler(controller);
         controller.checkIdleStates(userId);
         assertEquals(bucket,
                 controller.getAppStandbyBucket(PACKAGE_1, userId, mInjector.mElapsedRealtime,
@@ -744,50 +746,85 @@
     }
 
     private int getStandbyBucket(int userId, AppStandbyController controller, String packageName) {
+        flushHandler(controller);
         return controller.getAppStandbyBucket(packageName, userId, mInjector.mElapsedRealtime,
                 true);
     }
 
+    private List<AppStandbyInfo> getStandbyBuckets(int userId) {
+        flushHandler(mController);
+        return mController.getAppStandbyBuckets(userId);
+    }
+
     private int getStandbyBucketReason(String packageName) {
+        flushHandler(mController);
         return mController.getAppStandbyBucketReason(packageName, USER_ID,
                 mInjector.mElapsedRealtime);
     }
 
-    private void assertBucket(int bucket) throws InterruptedException {
-        assertBucket(bucket, PACKAGE_1);
+    private void waitAndAssertBucket(int bucket, String pkg) {
+        waitAndAssertBucket(mController, bucket, pkg);
     }
 
-    private void assertBucket(int bucket, String pkg) throws InterruptedException {
-        int retries = 0;
-        do {
-            if (bucket == getStandbyBucket(mController, pkg)) {
-                // Success
-                return;
-            }
-            Thread.sleep(ASSERT_RETRY_DELAY_MILLISECONDS);
-            retries++;
-        } while(retries < ASSERT_RETRY_ATTEMPTS);
-        // try one last time
-        assertEquals(bucket, getStandbyBucket(mController, pkg));
+    private void waitAndAssertBucket(AppStandbyController controller, int bucket, String pkg) {
+        StringBuilder sb = new StringBuilder();
+        sb.append(pkg);
+        sb.append(" was not in the ");
+        sb.append(UsageStatsManager.standbyBucketToString(bucket));
+        sb.append(" (");
+        sb.append(bucket);
+        sb.append(") bucket.");
+        waitAndAssertBucket(sb.toString(), controller, bucket, pkg);
     }
 
-    private void assertNotBucket(int bucket) throws InterruptedException {
-        final String pkg = PACKAGE_1;
+    private void waitAndAssertBucket(String msg, int bucket, String pkg) {
+        waitAndAssertBucket(msg, mController, bucket, pkg);
+    }
+
+    private void waitAndAssertBucket(String msg, AppStandbyController controller, int bucket,
+            String pkg) {
+        waitAndAssertBucket(msg, controller, bucket, USER_ID, pkg);
+    }
+
+    private void waitAndAssertBucket(String msg, AppStandbyController controller, int bucket,
+            int userId,
+            String pkg) {
+        waitUntil(() -> bucket == getStandbyBucket(userId, controller, pkg));
+        assertEquals(msg, bucket, getStandbyBucket(userId, controller, pkg));
+    }
+
+    private void waitAndAssertNotBucket(int bucket, String pkg) {
+        waitAndAssertNotBucket(mController, bucket, pkg);
+    }
+
+    private void waitAndAssertNotBucket(AppStandbyController controller, int bucket, String pkg) {
+        waitUntil(() -> bucket != getStandbyBucket(controller, pkg));
+        assertNotEquals(bucket, getStandbyBucket(controller, pkg));
+    }
+
+    private void waitAndAssertLastNoteEvent(int event) {
+        waitUntil(() -> {
+            flushHandler(mController);
+            return event == mInjector.mLastNoteEvent;
+        });
+        assertEquals(event, mInjector.mLastNoteEvent);
+    }
+
+    // Waits until condition is true or times out.
+    private void waitUntil(BooleanSupplier resultSupplier) {
         int retries = 0;
         do {
-            if (bucket != getStandbyBucket(mController, pkg)) {
-                // Success
-                return;
+            if (resultSupplier.getAsBoolean()) return;
+            try {
+                Thread.sleep(ASSERT_RETRY_DELAY_MILLISECONDS);
+            } catch (InterruptedException ie) {
+                // Do nothing
             }
-            Thread.sleep(ASSERT_RETRY_DELAY_MILLISECONDS);
             retries++;
-        } while(retries < ASSERT_RETRY_ATTEMPTS);
-        // try one last time
-        assertNotEquals(bucket, getStandbyBucket(mController, pkg));
+        } while (retries < ASSERT_RETRY_ATTEMPTS);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testBuckets() throws Exception {
         assertTimeout(mController, 0, STANDBY_BUCKET_NEVER);
 
@@ -820,14 +857,13 @@
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testSetAppStandbyBucket() throws Exception {
         // For a known package, standby bucket should be set properly
         reportEvent(mController, USER_INTERACTION, 0, PACKAGE_1);
         mInjector.mElapsedRealtime = HOUR_MS;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_ACTIVE,
                 REASON_MAIN_TIMEOUT);
-        assertEquals(STANDBY_BUCKET_ACTIVE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // For an unknown package, standby bucket should not be set, hence NEVER is returned
         // Ensure the unknown package is not already in history by removing it
@@ -836,21 +872,20 @@
         mController.setAppStandbyBucket(PACKAGE_UNKNOWN, USER_ID, STANDBY_BUCKET_ACTIVE,
                 REASON_MAIN_TIMEOUT);
         isPackageInstalled = true; // Reset mocked variable for other tests
-        assertEquals(STANDBY_BUCKET_NEVER, getStandbyBucket(mController, PACKAGE_UNKNOWN));
+        waitAndAssertBucket(STANDBY_BUCKET_NEVER, PACKAGE_UNKNOWN);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testAppStandbyBucketOnInstallAndUninstall() throws Exception {
         // On package install, standby bucket should be ACTIVE
         reportEvent(mController, USER_INTERACTION, 0, PACKAGE_UNKNOWN);
-        assertEquals(STANDBY_BUCKET_ACTIVE, getStandbyBucket(mController, PACKAGE_UNKNOWN));
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_UNKNOWN);
 
         // On uninstall, package should not exist in history and should return a NEVER bucket
         mController.clearAppIdleForPackage(PACKAGE_UNKNOWN, USER_ID);
-        assertEquals(STANDBY_BUCKET_NEVER, getStandbyBucket(mController, PACKAGE_UNKNOWN));
+        waitAndAssertBucket(STANDBY_BUCKET_NEVER, PACKAGE_UNKNOWN);
         // Ensure uninstalled app is not in history
-        List<AppStandbyInfo> buckets = mController.getAppStandbyBuckets(USER_ID);
+        List<AppStandbyInfo> buckets = getStandbyBuckets(USER_ID);
         for(AppStandbyInfo bucket : buckets) {
             if (bucket.mPackageName.equals(PACKAGE_UNKNOWN)) {
                 fail("packageName found in app idle history after uninstall.");
@@ -859,7 +894,6 @@
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testScreenTimeAndBuckets() throws Exception {
         mInjector.setDisplayOn(false);
 
@@ -876,22 +910,21 @@
         // RARE bucket, should fail because the screen wasn't ON.
         mInjector.mElapsedRealtime = RARE_THRESHOLD + 1;
         mController.checkIdleStates(USER_ID);
-        assertNotEquals(STANDBY_BUCKET_RARE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertNotBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
 
         mInjector.setDisplayOn(true);
         assertTimeout(mController, RARE_THRESHOLD + 2 * HOUR_MS + 1, STANDBY_BUCKET_RARE);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testForcedIdle() throws Exception {
         mController.forceIdleState(PACKAGE_1, USER_ID, true);
-        assertEquals(STANDBY_BUCKET_RARE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
         assertTrue(mController.isAppIdleFiltered(PACKAGE_1, UID_1, USER_ID, 0));
 
         mController.forceIdleState(PACKAGE_1, USER_ID, false);
-        assertEquals(STANDBY_BUCKET_ACTIVE, mController.getAppStandbyBucket(PACKAGE_1, USER_ID, 0,
-                true));
+
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
         assertFalse(mController.isAppIdleFiltered(PACKAGE_1, UID_1, USER_ID, 0));
     }
 
@@ -901,15 +934,15 @@
                 .onPropertiesChanged(mInjector.getDeviceConfigProperties());
 
         reportEvent(mController, USER_INTERACTION, 0, PACKAGE_1);
-        assertEquals(STANDBY_BUCKET_ACTIVE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
         mInjector.mElapsedRealtime = 1;
         rearmQuotaBumpLatch(PACKAGE_1, USER_ID);
         reportEvent(mController, NOTIFICATION_SEEN, mInjector.mElapsedRealtime, PACKAGE_1);
-        assertEquals(STANDBY_BUCKET_ACTIVE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         mController.forceIdleState(PACKAGE_1, USER_ID, true);
         reportEvent(mController, NOTIFICATION_SEEN, mInjector.mElapsedRealtime, PACKAGE_1);
-        assertEquals(STANDBY_BUCKET_WORKING_SET, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_WORKING_SET, PACKAGE_1);
         assertFalse(mQuotaBumpLatch.await(1, TimeUnit.SECONDS));
     }
 
@@ -917,9 +950,10 @@
     public void testNotificationEvent_bucketPromotion_changePromotedBucket() throws Exception {
         mInjector.mPropertiesChangedListener
                 .onPropertiesChanged(mInjector.getDeviceConfigProperties());
+        mInjector.mElapsedRealtime += RARE_THRESHOLD + 1;
         mController.forceIdleState(PACKAGE_1, USER_ID, true);
         reportEvent(mController, NOTIFICATION_SEEN, mInjector.mElapsedRealtime, PACKAGE_1);
-        assertEquals(STANDBY_BUCKET_WORKING_SET, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_WORKING_SET, PACKAGE_1);
 
         // TODO: Avoid hardcoding these string constants.
         mInjector.mSettingsBuilder.setInt("notification_seen_promoted_bucket",
@@ -928,11 +962,10 @@
                 mInjector.getDeviceConfigProperties());
         mController.forceIdleState(PACKAGE_1, USER_ID, true);
         reportEvent(mController, NOTIFICATION_SEEN, mInjector.mElapsedRealtime, PACKAGE_1);
-        assertEquals(STANDBY_BUCKET_FREQUENT, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_FREQUENT, PACKAGE_1);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testNotificationEvent_quotaBump() throws Exception {
         mInjector.mSettingsBuilder
                 .setBoolean("trigger_quota_bump_on_notification_seen", true);
@@ -942,7 +975,7 @@
                 .onPropertiesChanged(mInjector.getDeviceConfigProperties());
 
         reportEvent(mController, USER_INTERACTION, 0, PACKAGE_1);
-        assertEquals(STANDBY_BUCKET_ACTIVE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
         mInjector.mElapsedRealtime = RARE_THRESHOLD * 2;
         setAndAssertBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE, REASON_MAIN_FORCED_BY_SYSTEM);
 
@@ -951,83 +984,80 @@
 
         reportEvent(mController, NOTIFICATION_SEEN, mInjector.mElapsedRealtime, PACKAGE_1);
         assertTrue(mQuotaBumpLatch.await(1, TimeUnit.SECONDS));
-        assertEquals(STANDBY_BUCKET_RARE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testSlicePinnedEvent() throws Exception {
         reportEvent(mController, USER_INTERACTION, 0, PACKAGE_1);
-        assertEquals(STANDBY_BUCKET_ACTIVE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
         mInjector.mElapsedRealtime = 1;
         reportEvent(mController, SLICE_PINNED, mInjector.mElapsedRealtime, PACKAGE_1);
-        assertEquals(STANDBY_BUCKET_ACTIVE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         mController.forceIdleState(PACKAGE_1, USER_ID, true);
         reportEvent(mController, SLICE_PINNED, mInjector.mElapsedRealtime, PACKAGE_1);
-        assertEquals(STANDBY_BUCKET_WORKING_SET, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_WORKING_SET, PACKAGE_1);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testSlicePinnedPrivEvent() throws Exception {
         mController.forceIdleState(PACKAGE_1, USER_ID, true);
         reportEvent(mController, SLICE_PINNED_PRIV, mInjector.mElapsedRealtime, PACKAGE_1);
-        assertEquals(STANDBY_BUCKET_ACTIVE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testPredictionTimedOut() throws Exception {
         // Set it to timeout or usage, so that prediction can override it
         mInjector.mElapsedRealtime = HOUR_MS;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_TIMEOUT);
-        assertEquals(STANDBY_BUCKET_RARE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
 
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_ACTIVE,
                 REASON_MAIN_PREDICTED);
-        assertEquals(STANDBY_BUCKET_ACTIVE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // Fast forward 12 hours
         mInjector.mElapsedRealtime += WORKING_SET_THRESHOLD;
         mController.checkIdleStates(USER_ID);
         // Should still be in predicted bucket, since prediction timeout is 1 day since prediction
-        assertEquals(STANDBY_BUCKET_ACTIVE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
         // Fast forward two more hours
         mInjector.mElapsedRealtime += 2 * HOUR_MS;
         mController.checkIdleStates(USER_ID);
         // Should have now applied prediction timeout
-        assertEquals(STANDBY_BUCKET_WORKING_SET, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_WORKING_SET, PACKAGE_1);
 
         // Fast forward RARE bucket
         mInjector.mElapsedRealtime += RARE_THRESHOLD;
         mController.checkIdleStates(USER_ID);
         // Should continue to apply prediction timeout
-        assertEquals(STANDBY_BUCKET_RARE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
+    @Ignore("b/317086276")
     public void testOverrides() throws Exception {
         // Can force to NEVER
         mInjector.mElapsedRealtime = HOUR_MS;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_NEVER,
                 REASON_MAIN_FORCED_BY_USER);
-        assertEquals(STANDBY_BUCKET_NEVER, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_NEVER, PACKAGE_1);
 
         // Prediction can't override FORCED reasons
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_FORCED_BY_SYSTEM);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_WORKING_SET,
                 REASON_MAIN_PREDICTED);
-        assertEquals(STANDBY_BUCKET_RARE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
 
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_FREQUENT,
                 REASON_MAIN_FORCED_BY_USER);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_WORKING_SET,
                 REASON_MAIN_PREDICTED);
-        assertEquals(STANDBY_BUCKET_FREQUENT, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_FREQUENT, PACKAGE_1);
 
         // Prediction can't override NEVER
         mInjector.mElapsedRealtime = 2 * HOUR_MS;
@@ -1035,115 +1065,114 @@
                 REASON_MAIN_DEFAULT);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_ACTIVE,
                 REASON_MAIN_PREDICTED);
-        assertEquals(STANDBY_BUCKET_NEVER, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_NEVER, PACKAGE_1);
 
         // Prediction can't set to NEVER
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_ACTIVE,
                 REASON_MAIN_USAGE);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_NEVER,
                 REASON_MAIN_PREDICTED);
-        assertEquals(STANDBY_BUCKET_ACTIVE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // Prediction can't remove from RESTRICTED
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_USER);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_WORKING_SET,
                 REASON_MAIN_PREDICTED);
-        assertEquals(STANDBY_BUCKET_RESTRICTED, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
 
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_SYSTEM);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_WORKING_SET,
                 REASON_MAIN_PREDICTED);
-        assertEquals(STANDBY_BUCKET_RESTRICTED, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
 
         // Force from user can remove from RESTRICTED
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_USER);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_WORKING_SET,
                 REASON_MAIN_FORCED_BY_USER);
-        assertEquals(STANDBY_BUCKET_WORKING_SET, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_WORKING_SET, PACKAGE_1);
 
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_SYSTEM);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_WORKING_SET,
                 REASON_MAIN_FORCED_BY_USER);
-        assertEquals(STANDBY_BUCKET_WORKING_SET, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_WORKING_SET, PACKAGE_1);
 
         // Force from system can remove from RESTRICTED if it was put it in due to system
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_SYSTEM);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_WORKING_SET,
                 REASON_MAIN_FORCED_BY_SYSTEM);
-        assertEquals(STANDBY_BUCKET_WORKING_SET, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_WORKING_SET, PACKAGE_1);
 
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_USER);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_WORKING_SET,
                 REASON_MAIN_FORCED_BY_SYSTEM);
-        assertEquals(STANDBY_BUCKET_RESTRICTED, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
 
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_PREDICTED);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_WORKING_SET,
                 REASON_MAIN_FORCED_BY_SYSTEM);
-        assertEquals(STANDBY_BUCKET_RESTRICTED, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
 
         // Non-user usage can't remove from RESTRICTED
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_SYSTEM);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_WORKING_SET,
                 REASON_MAIN_USAGE);
-        assertEquals(STANDBY_BUCKET_RESTRICTED, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_WORKING_SET,
                 REASON_MAIN_USAGE | REASON_SUB_USAGE_SYSTEM_INTERACTION);
-        assertEquals(STANDBY_BUCKET_RESTRICTED, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_WORKING_SET,
                 REASON_MAIN_USAGE | REASON_SUB_USAGE_SYNC_ADAPTER);
-        assertEquals(STANDBY_BUCKET_RESTRICTED, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_WORKING_SET,
                 REASON_MAIN_USAGE | REASON_SUB_USAGE_EXEMPTED_SYNC_SCHEDULED_NON_DOZE);
-        assertEquals(STANDBY_BUCKET_RESTRICTED, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
 
         // Explicit user usage can remove from RESTRICTED
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_USER);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_WORKING_SET,
                 REASON_MAIN_USAGE | REASON_SUB_USAGE_USER_INTERACTION);
-        assertEquals(STANDBY_BUCKET_WORKING_SET, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_WORKING_SET, PACKAGE_1);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_SYSTEM);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_ACTIVE,
                 REASON_MAIN_USAGE | REASON_SUB_USAGE_MOVE_TO_FOREGROUND);
-        assertEquals(STANDBY_BUCKET_ACTIVE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testTimeout() throws Exception {
         reportEvent(mController, USER_INTERACTION, 0, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         mInjector.mElapsedRealtime = 2000;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_FREQUENT,
                 REASON_MAIN_PREDICTED);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // bucketing works after timeout
         mInjector.mElapsedRealtime = mController.mPredictionTimeoutMillis - 100;
         mController.checkIdleStates(USER_ID);
         // Use recent prediction
-        assertBucket(STANDBY_BUCKET_FREQUENT);
+        waitAndAssertBucket(STANDBY_BUCKET_FREQUENT, PACKAGE_1);
+        waitAndAssertBucket(STANDBY_BUCKET_FREQUENT, PACKAGE_1);
 
         // Way past prediction timeout, use system thresholds
         mInjector.mElapsedRealtime = RARE_THRESHOLD;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_RARE);
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
     }
 
     /** Test that timeouts still work properly even if invalid configuration values are set. */
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testTimeout_InvalidThresholds() throws Exception {
         mInjector.mSettingsBuilder
                 .setLong("screen_threshold_active", -1)
@@ -1161,19 +1190,19 @@
 
         reportEvent(mController, USER_INTERACTION, 0, PACKAGE_1);
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         mInjector.mElapsedRealtime = HOUR_MS;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_FREQUENT);
+        waitAndAssertBucket(STANDBY_BUCKET_FREQUENT, PACKAGE_1);
 
         mInjector.mElapsedRealtime = 2 * HOUR_MS;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_RARE);
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
 
         mInjector.mElapsedRealtime = 4 * HOUR_MS;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
     }
 
     /**
@@ -1181,74 +1210,72 @@
      * timeout has passed.
      */
     @Test
-    @FlakyTest(bugId = 185169504)
+    @Ignore("b/317086276")
     public void testTimeoutBeforeRestricted() throws Exception {
         reportEvent(mController, USER_INTERACTION, 0, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         mInjector.mElapsedRealtime += WORKING_SET_THRESHOLD;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_SYSTEM);
         // Bucket shouldn't change
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // bucketing works after timeout
         mInjector.mElapsedRealtime += DAY_MS;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_SYSTEM);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
 
         // Way past all timeouts. Make sure timeout processing doesn't raise bucket.
         mInjector.mElapsedRealtime += RESTRICTED_THRESHOLD * 4;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
     }
 
     /**
      * Test that an app is put into the RESTRICTED bucket after enough time has passed.
      */
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testRestrictedDelay() throws Exception {
         reportEvent(mController, USER_INTERACTION, 0, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         mInjector.mElapsedRealtime += mInjector.getAutoRestrictedBucketDelayMs() - 5000;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_SYSTEM);
         // Bucket shouldn't change
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // bucketing works after timeout
         mInjector.mElapsedRealtime += 6000;
 
         Thread.sleep(6000);
         // Enough time has passed. The app should automatically be put into the RESTRICTED bucket.
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
     }
 
     /**
      * Test that an app is put into the RESTRICTED bucket after enough time has passed.
      */
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testRestrictedDelay_DelayChange() throws Exception {
         reportEvent(mController, USER_INTERACTION, 0, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         mInjector.mAutoRestrictedBucketDelayMs = 2 * HOUR_MS;
         mInjector.mElapsedRealtime += 2 * HOUR_MS - 5000;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_SYSTEM);
         // Bucket shouldn't change
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // bucketing works after timeout
         mInjector.mElapsedRealtime += 6000;
 
         Thread.sleep(6000);
         // Enough time has passed. The app should automatically be put into the RESTRICTED bucket.
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
     }
 
     /**
@@ -1256,36 +1283,35 @@
      * a low bucket after the RESTRICTED timeout.
      */
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testRestrictedTimeoutOverridesRestoredLowBucketPrediction() throws Exception {
         reportEvent(mController, USER_INTERACTION, mInjector.mElapsedRealtime, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // Predict to RARE Not long enough to time out into RESTRICTED.
         mInjector.mElapsedRealtime += RARE_THRESHOLD;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_PREDICTED);
-        assertBucket(STANDBY_BUCKET_RARE);
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
 
         // Add a short timeout event
         mInjector.mElapsedRealtime += 1000;
         reportEvent(mController, SYSTEM_INTERACTION, mInjector.mElapsedRealtime, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
         mInjector.mElapsedRealtime += 1000;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // Long enough that it could have timed out into RESTRICTED. Instead of reverting to
         // predicted RARE, should go into RESTRICTED
         mInjector.mElapsedRealtime += RESTRICTED_THRESHOLD * 4;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
 
         // Ensure that prediction can still raise it out despite this override.
         mInjector.mElapsedRealtime += 1;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_ACTIVE,
                 REASON_MAIN_PREDICTED);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
     }
 
     /**
@@ -1293,7 +1319,6 @@
      * a low bucket after the RESTRICTED timeout.
      */
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testRestrictedTimeoutOverridesPredictionLowBucket() throws Exception {
         reportEvent(mController, USER_INTERACTION, mInjector.mElapsedRealtime, PACKAGE_1);
 
@@ -1301,7 +1326,7 @@
         mInjector.mElapsedRealtime += RARE_THRESHOLD;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_PREDICTED);
-        assertBucket(STANDBY_BUCKET_RARE);
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
 
         mInjector.mElapsedRealtime += 1;
         reportEvent(mController, USER_INTERACTION, mInjector.mElapsedRealtime, PACKAGE_1);
@@ -1310,10 +1335,10 @@
         mInjector.mElapsedRealtime += RESTRICTED_THRESHOLD * 4;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_ACTIVE,
                 REASON_MAIN_PREDICTED);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_PREDICTED);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
     }
 
     /**
@@ -1321,261 +1346,250 @@
      * interaction.
      */
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testSystemInteractionOverridesRestrictedTimeout() throws Exception {
         reportEvent(mController, USER_INTERACTION, mInjector.mElapsedRealtime, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // Long enough that it could have timed out into RESTRICTED.
         mInjector.mElapsedRealtime += RESTRICTED_THRESHOLD * 4;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
 
         // Report system interaction.
         mInjector.mElapsedRealtime += 1000;
         reportEvent(mController, SYSTEM_INTERACTION, mInjector.mElapsedRealtime, PACKAGE_1);
 
         // Ensure that it's raised out of RESTRICTED for the system interaction elevation duration.
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
         mInjector.mElapsedRealtime += 1000;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // Elevation duration over. Should fall back down.
         mInjector.mElapsedRealtime += 10 * MINUTE_MS;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testPredictionRaiseFromRestrictedTimeout_highBucket() throws Exception {
         reportEvent(mController, USER_INTERACTION, mInjector.mElapsedRealtime, PACKAGE_1);
 
         // Way past all timeouts. App times out into RESTRICTED bucket.
         mInjector.mElapsedRealtime += RESTRICTED_THRESHOLD * 4;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
 
         // Since the app timed out into RESTRICTED, prediction should be able to remove from the
         // bucket.
         mInjector.mElapsedRealtime += RESTRICTED_THRESHOLD;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_ACTIVE,
                 REASON_MAIN_PREDICTED);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testPredictionRaiseFromRestrictedTimeout_lowBucket() throws Exception {
         reportEvent(mController, USER_INTERACTION, mInjector.mElapsedRealtime, PACKAGE_1);
 
         // Way past all timeouts. App times out into RESTRICTED bucket.
         mInjector.mElapsedRealtime += RESTRICTED_THRESHOLD * 4;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
 
         // Prediction into a low bucket means no expectation of the app being used, so we shouldn't
         // elevate the app from RESTRICTED.
         mInjector.mElapsedRealtime += RESTRICTED_THRESHOLD;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_PREDICTED);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testCascadingTimeouts() throws Exception {
         mInjector.mPropertiesChangedListener
                 .onPropertiesChanged(mInjector.getDeviceConfigProperties());
 
         reportEvent(mController, USER_INTERACTION, 0, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         reportEvent(mController, NOTIFICATION_SEEN, 1000, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_WORKING_SET,
                 REASON_MAIN_PREDICTED);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         mInjector.mElapsedRealtime = 2000 + mController.mStrongUsageTimeoutMillis;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_FREQUENT,
                 REASON_MAIN_PREDICTED);
-        assertBucket(STANDBY_BUCKET_WORKING_SET);
+        waitAndAssertBucket(STANDBY_BUCKET_WORKING_SET, PACKAGE_1);
 
         mInjector.mElapsedRealtime = 2000 + mController.mNotificationSeenTimeoutMillis;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_FREQUENT,
                 REASON_MAIN_PREDICTED);
-        assertBucket(STANDBY_BUCKET_FREQUENT);
+        waitAndAssertBucket(STANDBY_BUCKET_FREQUENT, PACKAGE_1);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testOverlappingTimeouts() throws Exception {
         mInjector.mPropertiesChangedListener
                 .onPropertiesChanged(mInjector.getDeviceConfigProperties());
 
         reportEvent(mController, USER_INTERACTION, 0, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         reportEvent(mController, NOTIFICATION_SEEN, 1000, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // Overlapping USER_INTERACTION before previous one times out
         reportEvent(mController, USER_INTERACTION, mController.mStrongUsageTimeoutMillis - 1000,
                 PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // Still in ACTIVE after first USER_INTERACTION times out
         mInjector.mElapsedRealtime = mController.mStrongUsageTimeoutMillis + 1000;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_FREQUENT,
                 REASON_MAIN_PREDICTED);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // Both timed out, so NOTIFICATION_SEEN timeout should be effective
         mInjector.mElapsedRealtime = mController.mStrongUsageTimeoutMillis * 2 + 2000;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_FREQUENT,
                 REASON_MAIN_PREDICTED);
-        assertBucket(STANDBY_BUCKET_WORKING_SET);
+        waitAndAssertBucket(STANDBY_BUCKET_WORKING_SET, PACKAGE_1);
 
         mInjector.mElapsedRealtime = mController.mNotificationSeenTimeoutMillis + 2000;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_PREDICTED);
-        assertBucket(STANDBY_BUCKET_RARE);
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testSystemInteractionTimeout() throws Exception {
         reportEvent(mController, USER_INTERACTION, 0, PACKAGE_1);
         // Fast forward to RARE
         mInjector.mElapsedRealtime = RARE_THRESHOLD + 100;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_RARE);
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
 
         // Trigger a SYSTEM_INTERACTION and verify bucket
         reportEvent(mController, SYSTEM_INTERACTION, mInjector.mElapsedRealtime, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // Verify it's still in ACTIVE close to end of timeout
         mInjector.mElapsedRealtime += mController.mSystemInteractionTimeoutMillis - 100;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // Verify bucket moves to RARE after timeout
         mInjector.mElapsedRealtime += 200;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_RARE);
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testInitialForegroundServiceTimeout() throws Exception {
         mInjector.mElapsedRealtime = 1 * RARE_THRESHOLD + 100;
         // Make sure app is in NEVER bucket
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_NEVER,
                 REASON_MAIN_FORCED_BY_USER);
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_NEVER);
+        waitAndAssertBucket(STANDBY_BUCKET_NEVER, PACKAGE_1);
 
         mInjector.mElapsedRealtime += 100;
 
         // Trigger a FOREGROUND_SERVICE_START and verify bucket
         reportEvent(mController, FOREGROUND_SERVICE_START, mInjector.mElapsedRealtime, PACKAGE_1);
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // Verify it's still in ACTIVE close to end of timeout
         mInjector.mElapsedRealtime += mController.mInitialForegroundServiceStartTimeoutMillis - 100;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // Verify bucket moves to RARE after timeout
         mInjector.mElapsedRealtime += 200;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_RARE);
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
 
         // Trigger a FOREGROUND_SERVICE_START again
         reportEvent(mController, FOREGROUND_SERVICE_START, mInjector.mElapsedRealtime, PACKAGE_1);
         mController.checkIdleStates(USER_ID);
         // Bucket should not be immediately elevated on subsequent service starts
-        assertBucket(STANDBY_BUCKET_RARE);
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testPredictionNotOverridden() throws Exception {
         mInjector.mPropertiesChangedListener
                 .onPropertiesChanged(mInjector.getDeviceConfigProperties());
 
         reportEvent(mController, USER_INTERACTION, 0, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         mInjector.mElapsedRealtime = WORKING_SET_THRESHOLD - 1000;
         reportEvent(mController, NOTIFICATION_SEEN, mInjector.mElapsedRealtime, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // Falls back to WORKING_SET
         mInjector.mElapsedRealtime += 5000;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_WORKING_SET);
+        waitAndAssertBucket(STANDBY_BUCKET_WORKING_SET, PACKAGE_1);
 
         // Predict to ACTIVE
         mInjector.mElapsedRealtime += 1000;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_ACTIVE,
                 REASON_MAIN_PREDICTED);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // CheckIdleStates should not change the prediction
         mInjector.mElapsedRealtime += 1000;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testPredictionStrikesBack() throws Exception {
         reportEvent(mController, USER_INTERACTION, 0, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // Predict to FREQUENT
         mInjector.mElapsedRealtime = RARE_THRESHOLD;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_FREQUENT,
                 REASON_MAIN_PREDICTED);
-        assertBucket(STANDBY_BUCKET_FREQUENT);
+        waitAndAssertBucket(STANDBY_BUCKET_FREQUENT, PACKAGE_1);
 
         // Add a short timeout event
         mInjector.mElapsedRealtime += 1000;
         reportEvent(mController, SYSTEM_INTERACTION, mInjector.mElapsedRealtime, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
         mInjector.mElapsedRealtime += 1000;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_ACTIVE);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         // Verify it reverted to predicted
         mInjector.mElapsedRealtime += WORKING_SET_THRESHOLD / 2;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_FREQUENT);
+        waitAndAssertBucket(STANDBY_BUCKET_FREQUENT, PACKAGE_1);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testSystemForcedFlags_NotAddedForUserForce() throws Exception {
         final int expectedReason = REASON_MAIN_FORCED_BY_USER;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_USER);
-        assertEquals(STANDBY_BUCKET_RESTRICTED, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         assertEquals(expectedReason, getStandbyBucketReason(PACKAGE_1));
 
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_SYSTEM | REASON_SUB_FORCED_SYSTEM_FLAG_ABUSE);
-        assertEquals(STANDBY_BUCKET_RESTRICTED, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         assertEquals(expectedReason, getStandbyBucketReason(PACKAGE_1));
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testSystemForcedFlags_AddedForSystemForce() throws Exception {
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_ACTIVE,
                 REASON_MAIN_DEFAULT);
@@ -1584,13 +1598,13 @@
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_SYSTEM
                         | REASON_SUB_FORCED_SYSTEM_FLAG_BACKGROUND_RESOURCE_USAGE);
-        assertEquals(STANDBY_BUCKET_RESTRICTED, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         assertEquals(REASON_MAIN_FORCED_BY_SYSTEM
                         | REASON_SUB_FORCED_SYSTEM_FLAG_BACKGROUND_RESOURCE_USAGE,
                 getStandbyBucketReason(PACKAGE_1));
 
         mController.restrictApp(PACKAGE_1, USER_ID, REASON_SUB_FORCED_SYSTEM_FLAG_ABUSE);
-        assertEquals(STANDBY_BUCKET_RESTRICTED, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         // Flags should be combined
         assertEquals(REASON_MAIN_FORCED_BY_SYSTEM
                 | REASON_SUB_FORCED_SYSTEM_FLAG_BACKGROUND_RESOURCE_USAGE
@@ -1598,7 +1612,6 @@
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testSystemForcedFlags_SystemForceChangesBuckets() throws Exception {
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_ACTIVE,
                 REASON_MAIN_DEFAULT);
@@ -1607,14 +1620,14 @@
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_FREQUENT,
                 REASON_MAIN_FORCED_BY_SYSTEM
                         | REASON_SUB_FORCED_SYSTEM_FLAG_BACKGROUND_RESOURCE_USAGE);
-        assertEquals(STANDBY_BUCKET_FREQUENT, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_FREQUENT, PACKAGE_1);
         assertEquals(REASON_MAIN_FORCED_BY_SYSTEM
                         | REASON_SUB_FORCED_SYSTEM_FLAG_BACKGROUND_RESOURCE_USAGE,
                 getStandbyBucketReason(PACKAGE_1));
 
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_FREQUENT,
                 REASON_MAIN_FORCED_BY_SYSTEM | REASON_SUB_FORCED_SYSTEM_FLAG_ABUSE);
-        assertEquals(STANDBY_BUCKET_FREQUENT, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_FREQUENT, PACKAGE_1);
         // Flags should be combined
         assertEquals(REASON_MAIN_FORCED_BY_SYSTEM
                         | REASON_SUB_FORCED_SYSTEM_FLAG_BACKGROUND_RESOURCE_USAGE
@@ -1623,20 +1636,19 @@
 
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_FORCED_BY_SYSTEM | REASON_SUB_FORCED_SYSTEM_FLAG_BUGGY);
-        assertEquals(STANDBY_BUCKET_RARE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
         // Flags should be combined
         assertEquals(REASON_MAIN_FORCED_BY_SYSTEM | REASON_SUB_FORCED_SYSTEM_FLAG_BUGGY,
                 getStandbyBucketReason(PACKAGE_1));
 
         mController.restrictApp(PACKAGE_1, USER_ID, REASON_SUB_FORCED_SYSTEM_FLAG_UNDEFINED);
-        assertEquals(STANDBY_BUCKET_RESTRICTED, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         // Flags should not be combined since the bucket changed.
         assertEquals(REASON_MAIN_FORCED_BY_SYSTEM | REASON_SUB_FORCED_SYSTEM_FLAG_UNDEFINED,
                 getStandbyBucketReason(PACKAGE_1));
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testRestrictApp_MainReason() throws Exception {
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_ACTIVE,
                 REASON_MAIN_DEFAULT);
@@ -1644,11 +1656,11 @@
 
         mController.restrictApp(PACKAGE_1, USER_ID, REASON_MAIN_PREDICTED, 0);
         // Call should be ignored.
-        assertEquals(STANDBY_BUCKET_ACTIVE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         mController.restrictApp(PACKAGE_1, USER_ID, REASON_MAIN_FORCED_BY_USER, 0);
         // Call should go through
-        assertEquals(STANDBY_BUCKET_RESTRICTED, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
     }
 
     @Test
@@ -1724,15 +1736,15 @@
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testUserInteraction_CrossProfile() throws Exception {
         mInjector.mRunningUsers = new int[] {USER_ID, USER_ID2, USER_ID3};
         mInjector.mCrossProfileTargets = Arrays.asList(USER_HANDLE_USER2);
         reportEvent(mController, USER_INTERACTION, 0, PACKAGE_1);
-        assertEquals("Cross profile connected package bucket should be elevated on usage",
-                STANDBY_BUCKET_ACTIVE, getStandbyBucket(USER_ID2, mController, PACKAGE_1));
-        assertEquals("Not Cross profile connected package bucket should not be elevated on usage",
-                STANDBY_BUCKET_NEVER, getStandbyBucket(USER_ID3, mController, PACKAGE_1));
+        waitAndAssertBucket("Cross profile connected package bucket should be elevated on usage",
+                mController, STANDBY_BUCKET_ACTIVE, USER_ID2, PACKAGE_1);
+        waitAndAssertBucket(
+                "Not Cross profile connected package bucket should not be elevated on usage",
+                mController, STANDBY_BUCKET_NEVER, USER_ID3, PACKAGE_1);
 
         assertTimeout(mController, WORKING_SET_THRESHOLD - 1, STANDBY_BUCKET_ACTIVE, USER_ID);
         assertTimeout(mController, WORKING_SET_THRESHOLD - 1, STANDBY_BUCKET_ACTIVE, USER_ID2);
@@ -1742,51 +1754,50 @@
 
         mInjector.mCrossProfileTargets = Collections.emptyList();
         reportEvent(mController, USER_INTERACTION, 0, PACKAGE_1);
-        assertEquals("No longer cross profile connected package bucket should not be "
-                        + "elevated on usage",
-                STANDBY_BUCKET_WORKING_SET, getStandbyBucket(USER_ID2, mController, PACKAGE_1));
+        waitAndAssertBucket("No longer cross profile connected package bucket should not be "
+                        + "elevated on usage", mController, STANDBY_BUCKET_WORKING_SET, USER_ID2,
+                PACKAGE_1);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testUnexemptedSyncScheduled() throws Exception {
         rearmLatch(PACKAGE_1);
         mController.addListener(mListener);
-        assertEquals("Test package did not start in the Never bucket", STANDBY_BUCKET_NEVER,
-                getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket("Test package did not start in the Never bucket", STANDBY_BUCKET_NEVER,
+                PACKAGE_1);
 
         mController.postReportSyncScheduled(PACKAGE_1, USER_ID, false);
         mStateChangedLatch.await(1000, TimeUnit.MILLISECONDS);
-        assertEquals("Unexempted sync scheduled should bring the package out of the Never bucket",
-                STANDBY_BUCKET_WORKING_SET, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket(
+                "Unexempted sync scheduled should bring the package out of the Never bucket",
+                STANDBY_BUCKET_WORKING_SET, PACKAGE_1);
 
         setAndAssertBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE, REASON_MAIN_FORCED_BY_SYSTEM);
 
         rearmLatch(PACKAGE_1);
         mController.postReportSyncScheduled(PACKAGE_1, USER_ID, false);
         mStateChangedLatch.await(1000, TimeUnit.MILLISECONDS);
-        assertEquals("Unexempted sync scheduled should not elevate a non Never bucket",
-                STANDBY_BUCKET_RARE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket("Unexempted sync scheduled should not elevate a non Never bucket",
+                STANDBY_BUCKET_RARE, PACKAGE_1);
     }
 
     @Test
-    @FlakyTest(bugId = 185169504)
     public void testExemptedSyncScheduled() throws Exception {
         setAndAssertBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE, REASON_MAIN_FORCED_BY_SYSTEM);
         mInjector.mDeviceIdleMode = true;
         rearmLatch(PACKAGE_1);
         mController.postReportSyncScheduled(PACKAGE_1, USER_ID, true);
         mStateChangedLatch.await(1000, TimeUnit.MILLISECONDS);
-        assertEquals("Exempted sync scheduled in doze should set bucket to working set",
-                STANDBY_BUCKET_WORKING_SET, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket("Exempted sync scheduled in doze should set bucket to working set",
+                STANDBY_BUCKET_WORKING_SET, PACKAGE_1);
 
         setAndAssertBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE, REASON_MAIN_FORCED_BY_SYSTEM);
         mInjector.mDeviceIdleMode = false;
         rearmLatch(PACKAGE_1);
         mController.postReportSyncScheduled(PACKAGE_1, USER_ID, true);
         mStateChangedLatch.await(1000, TimeUnit.MILLISECONDS);
-        assertEquals("Exempted sync scheduled while not in doze should set bucket to active",
-                STANDBY_BUCKET_ACTIVE, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket("Exempted sync scheduled while not in doze should set bucket to active",
+                STANDBY_BUCKET_ACTIVE, PACKAGE_1);
     }
 
     @Test
@@ -1796,14 +1807,14 @@
         reportEvent(mController, USER_INTERACTION, mInjector.mElapsedRealtime, PACKAGE_1);
         mInjector.mElapsedRealtime += RESTRICTED_THRESHOLD * 4;
         mController.checkIdleStates(USER_ID);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
 
         mController.maybeUnrestrictBuggyApp(PACKAGE_1, USER_ID);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         mController.maybeUnrestrictBuggyApp(PACKAGE_1, USER_ID2);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         mController.maybeUnrestrictBuggyApp("com.random.package", USER_ID);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
 
         // Updates shouldn't change bucket if the app was forced by the system for a non-buggy
         // reason.
@@ -1814,11 +1825,11 @@
                         | REASON_SUB_FORCED_SYSTEM_FLAG_BACKGROUND_RESOURCE_USAGE);
 
         mController.maybeUnrestrictBuggyApp(PACKAGE_1, USER_ID);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         mController.maybeUnrestrictBuggyApp(PACKAGE_1, USER_ID2);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         mController.maybeUnrestrictBuggyApp("com.random.package", USER_ID);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
 
         // Updates should change bucket if the app was forced by the system for a buggy reason.
         reportEvent(mController, USER_INTERACTION, mInjector.mElapsedRealtime, PACKAGE_1);
@@ -1827,11 +1838,11 @@
                 REASON_MAIN_FORCED_BY_SYSTEM | REASON_SUB_FORCED_SYSTEM_FLAG_BUGGY);
 
         mController.maybeUnrestrictBuggyApp(PACKAGE_1, USER_ID2);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         mController.maybeUnrestrictBuggyApp("com.random.package", USER_ID);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         mController.maybeUnrestrictBuggyApp(PACKAGE_1, USER_ID);
-        assertNotEquals(STANDBY_BUCKET_RESTRICTED, getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertNotBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
 
         // Updates shouldn't change bucket if the app was forced by the system for more than just
         // a buggy reason.
@@ -1842,13 +1853,13 @@
                         | REASON_SUB_FORCED_SYSTEM_FLAG_BUGGY);
 
         mController.maybeUnrestrictBuggyApp(PACKAGE_1, USER_ID);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         assertEquals(REASON_MAIN_FORCED_BY_SYSTEM | REASON_SUB_FORCED_SYSTEM_FLAG_ABUSE,
                 getStandbyBucketReason(PACKAGE_1));
         mController.maybeUnrestrictBuggyApp(PACKAGE_1, USER_ID2);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         mController.maybeUnrestrictBuggyApp("com.random.package", USER_ID);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
 
         // Updates shouldn't change bucket if the app was forced by the user.
         reportEvent(mController, USER_INTERACTION, mInjector.mElapsedRealtime, PACKAGE_1);
@@ -1857,11 +1868,11 @@
                 REASON_MAIN_FORCED_BY_USER);
 
         mController.maybeUnrestrictBuggyApp(PACKAGE_1, USER_ID);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         mController.maybeUnrestrictBuggyApp(PACKAGE_1, USER_ID2);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
         mController.maybeUnrestrictBuggyApp("com.random.package", USER_ID);
-        assertBucket(STANDBY_BUCKET_RESTRICTED);
+        waitAndAssertBucket(STANDBY_BUCKET_RESTRICTED, PACKAGE_1);
     }
 
     @Test
@@ -1876,37 +1887,37 @@
 
         mController.setAppStandbyBucket(PACKAGE_SYSTEM_HEADFULL, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_TIMEOUT);
-        assertBucket(STANDBY_BUCKET_RARE, PACKAGE_SYSTEM_HEADFULL);
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_SYSTEM_HEADFULL);
 
         // Make sure headless system apps don't get lowered.
         mController.setAppStandbyBucket(PACKAGE_SYSTEM_HEADLESS, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_TIMEOUT);
-        assertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_SYSTEM_HEADLESS);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_SYSTEM_HEADLESS);
 
         // Package 1 doesn't have activities and is headless, but is not a system app, so it can
         // be lowered.
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_TIMEOUT);
-        assertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
     }
 
     @Test
     public void testWellbeingAppElevated() throws Exception {
         reportEvent(mController, USER_INTERACTION, mInjector.mElapsedRealtime, PACKAGE_WELLBEING);
-        assertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_WELLBEING);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_WELLBEING);
         reportEvent(mController, USER_INTERACTION, mInjector.mElapsedRealtime, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
         mInjector.mElapsedRealtime += RESTRICTED_THRESHOLD;
 
         // Make sure the default wellbeing app does not get lowered below WORKING_SET.
         mController.setAppStandbyBucket(PACKAGE_WELLBEING, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_TIMEOUT);
-        assertBucket(STANDBY_BUCKET_WORKING_SET, PACKAGE_WELLBEING);
+        waitAndAssertBucket(STANDBY_BUCKET_WORKING_SET, PACKAGE_WELLBEING);
 
         // A non default wellbeing app should be able to fall lower than WORKING_SET.
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_TIMEOUT);
-        assertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_1);
     }
 
     @Test
@@ -1914,22 +1925,22 @@
         mInjector.mClockApps.add(Pair.create(PACKAGE_1, UID_1));
 
         reportEvent(mController, USER_INTERACTION, mInjector.mElapsedRealtime, PACKAGE_1);
-        assertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_1);
 
         reportEvent(mController, USER_INTERACTION, mInjector.mElapsedRealtime, PACKAGE_2);
-        assertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_2);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_2);
 
         mInjector.mElapsedRealtime += RESTRICTED_THRESHOLD;
 
         // Make sure a clock app does not get lowered below WORKING_SET.
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_TIMEOUT);
-        assertBucket(STANDBY_BUCKET_WORKING_SET, PACKAGE_1);
+        waitAndAssertBucket(STANDBY_BUCKET_WORKING_SET, PACKAGE_1);
 
         // A non clock app should be able to fall lower than WORKING_SET.
         mController.setAppStandbyBucket(PACKAGE_2, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_TIMEOUT);
-        assertBucket(STANDBY_BUCKET_RARE, PACKAGE_2);
+        waitAndAssertBucket(STANDBY_BUCKET_RARE, PACKAGE_2);
     }
 
     @Test
@@ -2067,13 +2078,13 @@
     public void testBackgroundLocationBucket() throws Exception {
         reportEvent(mController, USER_INTERACTION, mInjector.mElapsedRealtime,
                 PACKAGE_BACKGROUND_LOCATION);
-        assertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_BACKGROUND_LOCATION);
+        waitAndAssertBucket(STANDBY_BUCKET_ACTIVE, PACKAGE_BACKGROUND_LOCATION);
 
         mInjector.mElapsedRealtime += RESTRICTED_THRESHOLD;
         // Make sure PACKAGE_BACKGROUND_LOCATION does not get lowered than STANDBY_BUCKET_FREQUENT.
         mController.setAppStandbyBucket(PACKAGE_BACKGROUND_LOCATION, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_TIMEOUT);
-        assertBucket(STANDBY_BUCKET_FREQUENT, PACKAGE_BACKGROUND_LOCATION);
+        waitAndAssertBucket(STANDBY_BUCKET_FREQUENT, PACKAGE_BACKGROUND_LOCATION);
     }
 
     @Test
@@ -2083,41 +2094,41 @@
 
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_FORCED_BY_USER);
-        assertEquals(BatteryStats.HistoryItem.EVENT_PACKAGE_INACTIVE, mInjector.mLastNoteEvent);
+        waitAndAssertLastNoteEvent(BatteryStats.HistoryItem.EVENT_PACKAGE_INACTIVE);
 
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_ACTIVE,
                 REASON_MAIN_FORCED_BY_USER);
-        assertEquals(BatteryStats.HistoryItem.EVENT_PACKAGE_ACTIVE, mInjector.mLastNoteEvent);
+        waitAndAssertLastNoteEvent(BatteryStats.HistoryItem.EVENT_PACKAGE_ACTIVE);
 
         // Since we're staying on the PACKAGE_ACTIVE side, noteEvent shouldn't be called.
         // Reset the last event to confirm the method isn't called.
         mInjector.mLastNoteEvent = BatteryStats.HistoryItem.EVENT_NONE;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_WORKING_SET,
                 REASON_MAIN_FORCED_BY_USER);
-        assertEquals(BatteryStats.HistoryItem.EVENT_NONE, mInjector.mLastNoteEvent);
+        waitAndAssertLastNoteEvent(BatteryStats.HistoryItem.EVENT_NONE);
 
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RARE,
                 REASON_MAIN_FORCED_BY_USER);
-        assertEquals(BatteryStats.HistoryItem.EVENT_PACKAGE_INACTIVE, mInjector.mLastNoteEvent);
+        waitAndAssertLastNoteEvent(BatteryStats.HistoryItem.EVENT_PACKAGE_INACTIVE);
 
         // Since we're staying on the PACKAGE_ACTIVE side, noteEvent shouldn't be called.
         // Reset the last event to confirm the method isn't called.
         mInjector.mLastNoteEvent = BatteryStats.HistoryItem.EVENT_NONE;
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_USER);
-        assertEquals(BatteryStats.HistoryItem.EVENT_NONE, mInjector.mLastNoteEvent);
+        waitAndAssertLastNoteEvent(BatteryStats.HistoryItem.EVENT_NONE);
 
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_FREQUENT,
                 REASON_MAIN_FORCED_BY_USER);
-        assertEquals(BatteryStats.HistoryItem.EVENT_PACKAGE_ACTIVE, mInjector.mLastNoteEvent);
+        waitAndAssertLastNoteEvent(BatteryStats.HistoryItem.EVENT_PACKAGE_ACTIVE);
 
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_RESTRICTED,
                 REASON_MAIN_FORCED_BY_USER);
-        assertEquals(BatteryStats.HistoryItem.EVENT_PACKAGE_INACTIVE, mInjector.mLastNoteEvent);
+        waitAndAssertLastNoteEvent(BatteryStats.HistoryItem.EVENT_PACKAGE_INACTIVE);
 
         mController.setAppStandbyBucket(PACKAGE_1, USER_ID, STANDBY_BUCKET_EXEMPTED,
                 REASON_MAIN_FORCED_BY_USER);
-        assertEquals(BatteryStats.HistoryItem.EVENT_PACKAGE_ACTIVE, mInjector.mLastNoteEvent);
+        waitAndAssertLastNoteEvent(BatteryStats.HistoryItem.EVENT_PACKAGE_ACTIVE);
     }
 
     private String getAdminAppsStr(int userId) {
@@ -2187,8 +2198,7 @@
         rearmLatch(pkg);
         mController.setAppStandbyBucket(pkg, user, bucket, reason);
         mStateChangedLatch.await(1, TimeUnit.SECONDS);
-        assertEquals("Failed to set package bucket", bucket,
-                getStandbyBucket(mController, PACKAGE_1));
+        waitAndAssertBucket("Failed to set package bucket", bucket, PACKAGE_1);
     }
 
     private void rearmLatch(String pkgName) {
@@ -2205,4 +2215,12 @@
         mLatchUserId = userId;
         mQuotaBumpLatch = new CountDownLatch(1);
     }
+
+    private void flushHandler(AppStandbyController controller) {
+        assertTrue("Failed to flush handler!", controller.flushHandler(FLUSH_TIMEOUT_MILLISECONDS));
+        // Some AppStandbyController handler messages queue another handler message. Flush again
+        // to catch those as well.
+        assertTrue("Failed to flush handler (the second time)!",
+                controller.flushHandler(FLUSH_TIMEOUT_MILLISECONDS));
+    }
 }
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
index e004ca0..a0e49a6 100755
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
@@ -297,6 +297,7 @@
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.ClassRule;
+import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TestRule;
@@ -13737,6 +13738,7 @@
     }
 
     @Test
+    @Ignore("b/316989461")
     public void cancelNotificationsFromListener_rapidClear_oldNew_cancelOne()
             throws RemoteException {
         mSetFlagsRule.enableFlags(android.view.contentprotection.flags.Flags
@@ -13766,6 +13768,7 @@
     }
 
     @Test
+    @Ignore("b/316989461")
     public void cancelNotificationsFromListener_rapidClear_old_cancelOne() throws RemoteException {
         mSetFlagsRule.enableFlags(android.view.contentprotection.flags.Flags
                 .FLAG_RAPID_CLEAR_NOTIFICATIONS_BY_LISTENER_APP_OP_ENABLED);
@@ -13793,6 +13796,7 @@
     }
 
     @Test
+    @Ignore("b/316989461")
     public void cancelNotificationsFromListener_rapidClear_oldNew_cancelOne_flagDisabled()
             throws RemoteException {
         mSetFlagsRule.disableFlags(android.view.contentprotection.flags.Flags
@@ -13823,6 +13827,7 @@
     }
 
     @Test
+    @Ignore("b/316989461")
     public void cancelNotificationsFromListener_rapidClear_oldNew_cancelAll()
             throws RemoteException {
         mSetFlagsRule.enableFlags(android.view.contentprotection.flags.Flags
@@ -13851,6 +13856,7 @@
     }
 
     @Test
+    @Ignore("b/316989461")
     public void cancelNotificationsFromListener_rapidClear_old_cancelAll() throws RemoteException {
         mSetFlagsRule.enableFlags(android.view.contentprotection.flags.Flags
                 .FLAG_RAPID_CLEAR_NOTIFICATIONS_BY_LISTENER_APP_OP_ENABLED);
@@ -13877,6 +13883,7 @@
     }
 
     @Test
+    @Ignore("b/316989461")
     public void cancelNotificationsFromListener_rapidClear_oldNew_cancelAll_flagDisabled()
             throws RemoteException {
         mSetFlagsRule.disableFlags(android.view.contentprotection.flags.Flags
diff --git a/services/tests/wmtests/src/com/android/server/wm/SurfaceControlViewHostTests.java b/services/tests/wmtests/src/com/android/server/wm/SurfaceControlViewHostTests.java
index ef427bb..a8b2178 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SurfaceControlViewHostTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SurfaceControlViewHostTests.java
@@ -16,7 +16,7 @@
 
 package com.android.server.wm;
 
-import static android.server.wm.CtsWindowInfoUtils.dumpWindowsOnScreen;
+import static android.server.wm.CtsWindowInfoUtils.assertAndDumpWindowState;
 import static android.server.wm.CtsWindowInfoUtils.waitForWindowFocus;
 import static android.server.wm.CtsWindowInfoUtils.waitForWindowVisible;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION;
@@ -129,37 +129,22 @@
             mScvh2.setView(mView2, lp2);
         });
 
-        boolean wasVisible = waitForWindowVisible(mView1);
-        if (!wasVisible) {
-            dumpWindowsOnScreen(TAG, "requestFocusWithMultipleWindows");
-        }
-        assertTrue("Failed to wait for view1", wasVisible);
-
-        wasVisible = waitForWindowVisible(mView2);
-        if (!wasVisible) {
-            dumpWindowsOnScreen(TAG, "requestFocusWithMultipleWindows-not visible");
-        }
-        assertTrue("Failed to wait for view2", wasVisible);
+        assertAndDumpWindowState(TAG, "Failed to wait for view1", waitForWindowVisible(mView1));
+        assertAndDumpWindowState(TAG, "Failed to wait for view2", waitForWindowVisible(mView2));
 
         IWindow window = IWindow.Stub.asInterface(mSurfaceView.getWindowToken());
 
         WindowManagerGlobal.getWindowSession().grantEmbeddedWindowFocus(window,
                 mScvh1.getInputTransferToken(), true);
 
-        boolean gainedFocus = waitForWindowFocus(mView1, true);
-        if (!gainedFocus) {
-            dumpWindowsOnScreen(TAG, "requestFocusWithMultipleWindows-view1 not focus");
-        }
-        assertTrue("Failed to gain focus for view1", gainedFocus);
+        assertAndDumpWindowState(TAG, "Failed to wait for view1 focus",
+                waitForWindowFocus(mView1, true));
 
         WindowManagerGlobal.getWindowSession().grantEmbeddedWindowFocus(window,
                 mScvh2.getInputTransferToken(), true);
 
-        gainedFocus = waitForWindowFocus(mView2, true);
-        if (!gainedFocus) {
-            dumpWindowsOnScreen(TAG, "requestFocusWithMultipleWindows-view2 not focus");
-        }
-        assertTrue("Failed to gain focus for view2", gainedFocus);
+        assertAndDumpWindowState(TAG, "Failed to wait for view2 focus",
+                waitForWindowFocus(mView2, true));
     }
 
     private static class TestWindowlessWindowManager extends WindowlessWindowManager {
diff --git a/services/tests/wmtests/src/com/android/server/wm/TrustedOverlayTests.java b/services/tests/wmtests/src/com/android/server/wm/TrustedOverlayTests.java
index ac49839..6a15b05 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TrustedOverlayTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TrustedOverlayTests.java
@@ -16,6 +16,7 @@
 
 package com.android.server.wm;
 
+import static android.server.wm.CtsWindowInfoUtils.assertAndDumpWindowState;
 import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_TRUSTED_OVERLAY;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_PANEL;
 
@@ -138,11 +139,8 @@
                     return false;
                 }, TIMEOUT_S, TimeUnit.SECONDS);
 
-        if (!foundTrusted[0]) {
-            CtsWindowInfoUtils.dumpWindowsOnScreen(TAG, mName.getMethodName());
-        }
-
-        assertTrue("Failed to find window or was not marked trusted", foundTrusted[0]);
+        assertAndDumpWindowState(TAG, "Failed to find window or was not marked trusted",
+                foundTrusted[0]);
     }
 
     private void testTrustedOverlayChildHelper(boolean expectedTrustedChild)
diff --git a/telecomm/java/android/telecom/Call.java b/telecomm/java/android/telecom/Call.java
index def52a5..874c10c 100644
--- a/telecomm/java/android/telecom/Call.java
+++ b/telecomm/java/android/telecom/Call.java
@@ -210,100 +210,6 @@
             "android.telecom.extra.SILENT_RINGING_REQUESTED";
 
     /**
-     * Call event sent from a {@link Call} via {@link #sendCallEvent(String, Bundle)} to inform
-     * Telecom that the user has requested that the current {@link Call} should be handed over
-     * to another {@link ConnectionService}.
-     * <p>
-     * The caller must specify the {@link #EXTRA_HANDOVER_PHONE_ACCOUNT_HANDLE} to indicate to
-     * Telecom which {@link PhoneAccountHandle} the {@link Call} should be handed over to.
-     * @hide
-     * @deprecated Use {@link Call#handoverTo(PhoneAccountHandle, int, Bundle)} and its associated
-     * APIs instead.
-     */
-    public static final String EVENT_REQUEST_HANDOVER =
-            "android.telecom.event.REQUEST_HANDOVER";
-
-    /**
-     * Extra key used with the {@link #EVENT_REQUEST_HANDOVER} call event.  Specifies the
-     * {@link PhoneAccountHandle} to which a call should be handed over to.
-     * @hide
-     * @deprecated Use {@link Call#handoverTo(PhoneAccountHandle, int, Bundle)} and its associated
-     * APIs instead.
-     */
-    public static final String EXTRA_HANDOVER_PHONE_ACCOUNT_HANDLE =
-            "android.telecom.extra.HANDOVER_PHONE_ACCOUNT_HANDLE";
-
-    /**
-     * Integer extra key used with the {@link #EVENT_REQUEST_HANDOVER} call event.  Specifies the
-     * video state of the call when it is handed over to the new {@link PhoneAccount}.
-     * <p>
-     * Valid values: {@link VideoProfile#STATE_AUDIO_ONLY},
-     * {@link VideoProfile#STATE_BIDIRECTIONAL}, {@link VideoProfile#STATE_RX_ENABLED}, and
-     * {@link VideoProfile#STATE_TX_ENABLED}.
-     * @hide
-     * @deprecated Use {@link Call#handoverTo(PhoneAccountHandle, int, Bundle)} and its associated
-     * APIs instead.
-     */
-    public static final String EXTRA_HANDOVER_VIDEO_STATE =
-            "android.telecom.extra.HANDOVER_VIDEO_STATE";
-
-    /**
-     * Extra key used with the {@link #EVENT_REQUEST_HANDOVER} call event.  Used by the
-     * {@link InCallService} initiating a handover to provide a {@link Bundle} with extra
-     * information to the handover {@link ConnectionService} specified by
-     * {@link #EXTRA_HANDOVER_PHONE_ACCOUNT_HANDLE}.
-     * <p>
-     * This {@link Bundle} is not interpreted by Telecom, but passed as-is to the
-     * {@link ConnectionService} via the request extras when
-     * {@link ConnectionService#onCreateOutgoingConnection(PhoneAccountHandle, ConnectionRequest)}
-     * is called to initate the handover.
-     * @hide
-     * @deprecated Use {@link Call#handoverTo(PhoneAccountHandle, int, Bundle)} and its associated
-     * APIs instead.
-     */
-    public static final String EXTRA_HANDOVER_EXTRAS = "android.telecom.extra.HANDOVER_EXTRAS";
-
-    /**
-     * Call event sent from Telecom to the handover {@link ConnectionService} via
-     * {@link Connection#onCallEvent(String, Bundle)} to inform a {@link Connection} that a handover
-     * to the {@link ConnectionService} has completed successfully.
-     * <p>
-     * A handover is initiated with the {@link #EVENT_REQUEST_HANDOVER} call event.
-     * @hide
-     * @deprecated Use {@link Call#handoverTo(PhoneAccountHandle, int, Bundle)} and its associated
-     * APIs instead.
-     */
-    public static final String EVENT_HANDOVER_COMPLETE =
-            "android.telecom.event.HANDOVER_COMPLETE";
-
-    /**
-     * Call event sent from Telecom to the handover destination {@link ConnectionService} via
-     * {@link Connection#onCallEvent(String, Bundle)} to inform the handover destination that the
-     * source connection has disconnected.  The {@link Bundle} parameter for the call event will be
-     * {@code null}.
-     * <p>
-     * A handover is initiated with the {@link #EVENT_REQUEST_HANDOVER} call event.
-     * @hide
-     * @deprecated Use {@link Call#handoverTo(PhoneAccountHandle, int, Bundle)} and its associated
-     * APIs instead.
-     */
-    public static final String EVENT_HANDOVER_SOURCE_DISCONNECTED =
-            "android.telecom.event.HANDOVER_SOURCE_DISCONNECTED";
-
-    /**
-     * Call event sent from Telecom to the handover {@link ConnectionService} via
-     * {@link Connection#onCallEvent(String, Bundle)} to inform a {@link Connection} that a handover
-     * to the {@link ConnectionService} has failed.
-     * <p>
-     * A handover is initiated with the {@link #EVENT_REQUEST_HANDOVER} call event.
-     * @hide
-     * @deprecated Use {@link Call#handoverTo(PhoneAccountHandle, int, Bundle)} and its associated
-     * APIs instead.
-     */
-    public static final String EVENT_HANDOVER_FAILED =
-            "android.telecom.event.HANDOVER_FAILED";
-
-    /**
      * Event reported from the Telecom stack to report an in-call diagnostic message which the
      * dialer app may opt to display to the user.  A diagnostic message is used to communicate
      * scenarios the device has detected which may impact the quality of the ongoing call.
diff --git a/telecomm/java/android/telecom/Connection.java b/telecomm/java/android/telecom/Connection.java
index 4a541da..ee9bf898 100644
--- a/telecomm/java/android/telecom/Connection.java
+++ b/telecomm/java/android/telecom/Connection.java
@@ -961,28 +961,6 @@
             "android.telecom.event.CALL_REMOTELY_UNHELD";
 
     /**
-     * Connection event used to inform an {@link InCallService} which initiated a call handover via
-     * {@link Call#EVENT_REQUEST_HANDOVER} that the handover from this {@link Connection} has
-     * successfully completed.
-     * @hide
-     * @deprecated Use {@link Call#handoverTo(PhoneAccountHandle, int, Bundle)} and its associated
-     * APIs instead.
-     */
-    public static final String EVENT_HANDOVER_COMPLETE =
-            "android.telecom.event.HANDOVER_COMPLETE";
-
-    /**
-     * Connection event used to inform an {@link InCallService} which initiated a call handover via
-     * {@link Call#EVENT_REQUEST_HANDOVER} that the handover from this {@link Connection} has failed
-     * to complete.
-     * @hide
-     * @deprecated Use {@link Call#handoverTo(PhoneAccountHandle, int, Bundle)} and its associated
-     * APIs instead.
-     */
-    public static final String EVENT_HANDOVER_FAILED =
-            "android.telecom.event.HANDOVER_FAILED";
-
-    /**
      * String Connection extra key used to store SIP invite fields for an incoming call for IMS call
      */
     public static final String EXTRA_SIP_INVITE = "android.telecom.extra.SIP_INVITE";