Merge "Revert^2 "Add exemptions to flex policy."" into main
diff --git a/apex/jobscheduler/framework/java/com/android/server/DeviceIdleInternal.java b/apex/jobscheduler/framework/java/com/android/server/DeviceIdleInternal.java
index caf7e7f..1fc888b 100644
--- a/apex/jobscheduler/framework/java/com/android/server/DeviceIdleInternal.java
+++ b/apex/jobscheduler/framework/java/com/android/server/DeviceIdleInternal.java
@@ -16,6 +16,7 @@
 
 package com.android.server;
 
+import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.os.PowerExemptionManager;
 import android.os.PowerExemptionManager.ReasonCode;
@@ -77,6 +78,9 @@
 
     int[] getPowerSaveTempWhitelistAppIds();
 
+    @NonNull
+    String[] getFullPowerWhitelistExceptIdle();
+
     /**
      * Listener to be notified when DeviceIdleController determines that the device has moved or is
      * stationary.
diff --git a/apex/jobscheduler/service/java/com/android/server/DeviceIdleController.java b/apex/jobscheduler/service/java/com/android/server/DeviceIdleController.java
index 6383ed8..31214cb 100644
--- a/apex/jobscheduler/service/java/com/android/server/DeviceIdleController.java
+++ b/apex/jobscheduler/service/java/com/android/server/DeviceIdleController.java
@@ -2374,6 +2374,11 @@
             return DeviceIdleController.this.isAppOnWhitelistInternal(appid);
         }
 
+        @Override
+        public String[] getFullPowerWhitelistExceptIdle() {
+            return DeviceIdleController.this.getFullPowerWhitelistInternalUnchecked();
+        }
+
         /**
          * Returns the array of app ids whitelisted by user. Take care not to
          * modify this, as it is a reference to the original copy. But the reference
@@ -3100,10 +3105,14 @@
     }
 
     private String[] getFullPowerWhitelistInternal(final int callingUid, final int callingUserId) {
-        final String[] apps;
+        return ArrayUtils.filter(getFullPowerWhitelistInternalUnchecked(), String[]::new,
+                (pkg) -> !mPackageManagerInternal.filterAppAccess(pkg, callingUid, callingUserId));
+    }
+
+    private String[] getFullPowerWhitelistInternalUnchecked() {
         synchronized (this) {
             int size = mPowerSaveWhitelistApps.size() + mPowerSaveWhitelistUserApps.size();
-            apps = new String[size];
+            final String[] apps = new String[size];
             int cur = 0;
             for (int i = 0; i < mPowerSaveWhitelistApps.size(); i++) {
                 apps[cur] = mPowerSaveWhitelistApps.keyAt(i);
@@ -3113,9 +3122,8 @@
                 apps[cur] = mPowerSaveWhitelistUserApps.keyAt(i);
                 cur++;
             }
+            return apps;
         }
-        return ArrayUtils.filter(apps, String[]::new,
-                (pkg) -> !mPackageManagerInternal.filterAppAccess(pkg, callingUid, callingUserId));
     }
 
     public boolean isPowerSaveWhitelistExceptIdleAppInternal(String packageName) {
diff --git a/apex/jobscheduler/service/java/com/android/server/job/controllers/BackgroundJobsController.java b/apex/jobscheduler/service/java/com/android/server/job/controllers/BackgroundJobsController.java
index e3ba50d..e3ac780 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/controllers/BackgroundJobsController.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/controllers/BackgroundJobsController.java
@@ -318,6 +318,10 @@
 
         try {
             final boolean isStopped = mPackageManagerInternal.isPackageStopped(packageName, uid);
+            if (DEBUG) {
+                Slog.d(TAG,
+                        "Pulled stopped state of " + packageName + " (" + uid + "): " + isStopped);
+            }
             mPackageStoppedState.add(uid, packageName, isStopped);
             return isStopped;
         } catch (PackageManager.NameNotFoundException e) {
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 0e67b9a..2c9af67 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
@@ -30,11 +30,15 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.app.job.JobInfo;
+import android.content.BroadcastReceiver;
 import android.content.Context;
+import android.content.Intent;
+import android.content.IntentFilter;
 import android.content.pm.PackageManager;
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
+import android.os.PowerManager;
 import android.os.UserHandle;
 import android.provider.DeviceConfig;
 import android.util.ArraySet;
@@ -48,6 +52,8 @@
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.AppSchedulingModuleThread;
+import com.android.server.DeviceIdleInternal;
+import com.android.server.LocalServices;
 import com.android.server.job.JobSchedulerService;
 import com.android.server.utils.AlarmQueue;
 
@@ -127,6 +133,19 @@
     @GuardedBy("mLock")
     private final SparseLongArray mLastSeenConstraintTimesElapsed = new SparseLongArray();
 
+    private DeviceIdleInternal mDeviceIdleInternal;
+    private final ArraySet<String> mPowerAllowlistedApps = new ArraySet<>();
+
+    private final BroadcastReceiver mBroadcastReceiver = new BroadcastReceiver() {
+        @Override
+        public void onReceive(Context context, Intent intent) {
+            switch (intent.getAction()) {
+                case PowerManager.ACTION_POWER_SAVE_WHITELIST_CHANGED:
+                    mHandler.post(FlexibilityController.this::updatePowerAllowlistCache);
+                    break;
+            }
+        }
+    };
     @VisibleForTesting
     @GuardedBy("mLock")
     final FlexibilityTracker mFlexibilityTracker;
@@ -180,8 +199,16 @@
                 }
             };
 
-    private static final int MSG_UPDATE_JOBS = 0;
-    private static final int MSG_UPDATE_JOB = 1;
+    private static final int MSG_CHECK_ALL_JOBS = 0;
+    /** Check the jobs in {@link #mJobsToCheck} */
+    private static final int MSG_CHECK_JOBS = 1;
+    /** Check the jobs of packages in {@link #mPackagesToCheck} */
+    private static final int MSG_CHECK_PACKAGES = 2;
+
+    @GuardedBy("mLock")
+    private final ArraySet<JobStatus> mJobsToCheck = new ArraySet<>();
+    @GuardedBy("mLock")
+    private final ArraySet<String> mPackagesToCheck = new ArraySet<>();
 
     public FlexibilityController(
             JobSchedulerService service, PrefetchController prefetchController) {
@@ -204,6 +231,16 @@
         mPercentToDropConstraints =
                 mFcConfig.DEFAULT_PERCENT_TO_DROP_FLEXIBLE_CONSTRAINTS;
         mPrefetchController = prefetchController;
+
+        if (mFlexibilityEnabled) {
+            registerBroadcastReceiver();
+        }
+    }
+
+    @Override
+    public void onSystemServicesReady() {
+        mDeviceIdleInternal = LocalServices.getService(DeviceIdleInternal.class);
+        mHandler.post(FlexibilityController.this::updatePowerAllowlistCache);
     }
 
     @Override
@@ -241,6 +278,7 @@
             mFlexibilityAlarmQueue.removeAlarmForKey(js);
             mFlexibilityTracker.remove(js);
         }
+        mJobsToCheck.remove(js);
     }
 
     @Override
@@ -266,7 +304,14 @@
     @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)
+                // For apps in the power allowlist, automatically exclude DEFAULT+ priority jobs.
+                || (js.getEffectivePriority() >= JobInfo.PRIORITY_DEFAULT
+                        && mPowerAllowlistedApps.contains(js.getSourcePackageName()))
                 || hasEnoughSatisfiedConstraintsLocked(js)
                 || mService.isCurrentlyRunningLocked(js);
     }
@@ -371,7 +416,7 @@
 
                 // 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_UPDATE_JOBS).sendToTarget();
+                mHandler.obtainMessage(MSG_CHECK_ALL_JOBS).sendToTarget();
             }
         }
     }
@@ -485,7 +530,9 @@
     @Override
     @GuardedBy("mLock")
     public void onUidBiasChangedLocked(int uid, int prevBias, int newBias) {
-        if (prevBias != JobInfo.BIAS_TOP_APP && newBias != JobInfo.BIAS_TOP_APP) {
+        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.
             return;
         }
         final long nowElapsed = sElapsedRealtimeClock.millis();
@@ -557,6 +604,39 @@
         mFcConfig.processConstantLocked(properties, key);
     }
 
+    private void registerBroadcastReceiver() {
+        IntentFilter filter = new IntentFilter(PowerManager.ACTION_POWER_SAVE_WHITELIST_CHANGED);
+        mContext.registerReceiver(mBroadcastReceiver, filter);
+    }
+
+    private void unregisterBroadcastReceiver() {
+        mContext.unregisterReceiver(mBroadcastReceiver);
+    }
+
+    private void updatePowerAllowlistCache() {
+        if (mDeviceIdleInternal == null) {
+            return;
+        }
+
+        // Don't call out to DeviceIdleController with the lock held.
+        final String[] allowlistedPkgs = mDeviceIdleInternal.getFullPowerWhitelistExceptIdle();
+        final ArraySet<String> changedPkgs = new ArraySet<>();
+        synchronized (mLock) {
+            changedPkgs.addAll(mPowerAllowlistedApps);
+            mPowerAllowlistedApps.clear();
+            for (final String pkgName : allowlistedPkgs) {
+                mPowerAllowlistedApps.add(pkgName);
+                if (changedPkgs.contains(pkgName)) {
+                    changedPkgs.remove(pkgName);
+                } else {
+                    changedPkgs.add(pkgName);
+                }
+            }
+            mPackagesToCheck.addAll(changedPkgs);
+            mHandler.sendEmptyMessage(MSG_CHECK_PACKAGES);
+        }
+    }
+
     @VisibleForTesting
     class FlexibilityTracker {
         final ArrayList<ArraySet<JobStatus>> mTrackedJobs;
@@ -710,7 +790,8 @@
                     }
                     mFlexibilityTracker.setNumDroppedFlexibleConstraints(js,
                             js.getNumAppliedFlexibleConstraints());
-                    mHandler.obtainMessage(MSG_UPDATE_JOB, js).sendToTarget();
+                    mJobsToCheck.add(js);
+                    mHandler.sendEmptyMessage(MSG_CHECK_JOBS);
                     return;
                 }
                 if (nextTimeElapsed == NO_LIFECYCLE_END) {
@@ -761,10 +842,12 @@
         @Override
         public void handleMessage(Message msg) {
             switch (msg.what) {
-                case MSG_UPDATE_JOBS:
-                    removeMessages(MSG_UPDATE_JOBS);
+                case MSG_CHECK_ALL_JOBS:
+                    removeMessages(MSG_CHECK_ALL_JOBS);
 
                     synchronized (mLock) {
+                        mJobsToCheck.clear();
+                        mPackagesToCheck.clear();
                         final long nowElapsed = sElapsedRealtimeClock.millis();
                         final ArraySet<JobStatus> changedJobs = new ArraySet<>();
 
@@ -790,19 +873,50 @@
                     }
                     break;
 
-                case MSG_UPDATE_JOB:
+                case MSG_CHECK_JOBS:
                     synchronized (mLock) {
-                        final JobStatus js = (JobStatus) msg.obj;
-                        if (DEBUG) {
-                            Slog.d("blah", "Checking on " + js.toShortString());
-                        }
                         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);
+                        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);
+                            }
+                        }
+
+                        mJobsToCheck.clear();
+                        if (changedJobs.size() > 0) {
+                            mStateChangedListener.onControllerStateChanged(changedJobs);
+                        }
+                    }
+                    break;
+
+                case MSG_CHECK_PACKAGES:
+                    synchronized (mLock) {
+                        final long nowElapsed = sElapsedRealtimeClock.millis();
+                        final ArraySet<JobStatus> changedJobs = new ArraySet<>();
+
+                        mService.getJobStore().forEachJob(
+                                (js) -> mPackagesToCheck.contains(js.getSourcePackageName())
+                                        || mPackagesToCheck.contains(js.getCallingPackageName()),
+                                (js) -> {
+                                    if (DEBUG) {
+                                        Slog.d(TAG, "Checking on " + js.toShortString());
+                                    }
+                                    if (js.setFlexibilityConstraintSatisfied(
+                                            nowElapsed, isFlexibilitySatisfiedLocked(js))) {
+                                        changedJobs.add(js);
+                                    }
+                                });
+
+                        mPackagesToCheck.clear();
+                        if (changedJobs.size() > 0) {
+                            mStateChangedListener.onControllerStateChanged(changedJobs);
                         }
                     }
                     break;
@@ -882,10 +996,12 @@
                             mFlexibilityEnabled = true;
                             mPrefetchController
                                     .registerPrefetchChangedListener(mPrefetchChangedListener);
+                            registerBroadcastReceiver();
                         } else {
                             mFlexibilityEnabled = false;
                             mPrefetchController
                                     .unRegisterPrefetchChangedListener(mPrefetchChangedListener);
+                            unregisterBroadcastReceiver();
                         }
                     }
                     break;
@@ -985,7 +1101,14 @@
             pw.println(":");
             pw.increaseIndent();
 
-            pw.print(KEY_APPLIED_CONSTRAINTS, APPLIED_CONSTRAINTS).println();
+            pw.print(KEY_APPLIED_CONSTRAINTS, APPLIED_CONSTRAINTS);
+            pw.print("(");
+            if (APPLIED_CONSTRAINTS != 0) {
+                JobStatus.dumpConstraints(pw, APPLIED_CONSTRAINTS);
+            } else {
+                pw.print("nothing");
+            }
+            pw.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,
@@ -1044,6 +1167,10 @@
         pw.decreaseIndent();
 
         pw.println();
+        pw.print("Power allowlisted packages: ");
+        pw.println(mPowerAllowlistedApps);
+
+        pw.println();
         mFlexibilityTracker.dump(pw, predicate);
         pw.println();
         mFlexibilityAlarmQueue.dump(pw);
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 650c473..116d5db 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
@@ -26,15 +26,18 @@
 import static android.text.format.DateUtils.MINUTE_IN_MILLIS;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doAnswer;
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.doNothing;
 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.FLEXIBLE_CONSTRAINTS;
 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;
 import static com.android.server.job.controllers.FlexibilityController.FcConfig.KEY_DEADLINE_PROXIMITY_LIMIT;
 import static com.android.server.job.controllers.FlexibilityController.FcConfig.KEY_FALLBACK_FLEXIBILITY_DEADLINE;
 import static com.android.server.job.controllers.FlexibilityController.FcConfig.KEY_PERCENTS_TO_DROP_NUM_FLEXIBLE_CONSTRAINTS;
-import static com.android.server.job.controllers.FlexibilityController.FLEXIBLE_CONSTRAINTS;
 import static com.android.server.job.controllers.FlexibilityController.SYSTEM_WIDE_FLEXIBLE_CONSTRAINTS;
 import static com.android.server.job.controllers.JobStatus.CONSTRAINT_BATTERY_NOT_LOW;
 import static com.android.server.job.controllers.JobStatus.CONSTRAINT_CHARGING;
@@ -50,24 +53,32 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.app.AlarmManager;
 import android.app.AppGlobals;
 import android.app.job.JobInfo;
+import android.content.BroadcastReceiver;
 import android.content.ComponentName;
 import android.content.Context;
+import android.content.Intent;
 import android.content.pm.PackageManager;
 import android.content.pm.PackageManagerInternal;
 import android.net.NetworkRequest;
 import android.os.Looper;
+import android.os.PowerManager;
 import android.provider.DeviceConfig;
 import android.util.ArraySet;
+import android.util.EmptyArray;
 
 import com.android.server.AppSchedulingModuleThread;
+import com.android.server.DeviceIdleInternal;
 import com.android.server.LocalServices;
 import com.android.server.job.JobSchedulerService;
 import com.android.server.job.JobStore;
@@ -77,6 +88,7 @@
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.ArgumentCaptor;
 import org.mockito.ArgumentMatchers;
 import org.mockito.Mock;
 import org.mockito.MockitoSession;
@@ -95,6 +107,7 @@
     private static final long FROZEN_TIME = 100L;
 
     private MockitoSession mMockingSession;
+    private BroadcastReceiver mBroadcastReceiver;
     private FlexibilityController mFlexibilityController;
     private DeviceConfig.Properties.Builder mDeviceConfigPropertiesBuilder;
     private JobStore mJobStore;
@@ -106,6 +119,8 @@
     @Mock
     private Context mContext;
     @Mock
+    private DeviceIdleInternal mDeviceIdleInternal;
+    @Mock
     private JobSchedulerService mJobSchedulerService;
     @Mock
     private PrefetchController mPrefetchController;
@@ -128,10 +143,13 @@
         // Called in FlexibilityController constructor.
         when(mContext.getMainLooper()).thenReturn(Looper.getMainLooper());
         when(mContext.getSystemService(Context.ALARM_SERVICE)).thenReturn(mAlarmManager);
+        doNothing().when(mAlarmManager).setExact(anyInt(), anyLong(), anyString(), any(), any());
         when(mContext.getPackageManager()).thenReturn(mPackageManager);
         when(mPackageManager.hasSystemFeature(
                 PackageManager.FEATURE_AUTOMOTIVE)).thenReturn(false);
         when(mPackageManager.hasSystemFeature(PackageManager.FEATURE_EMBEDDED)).thenReturn(false);
+        doReturn(mDeviceIdleInternal)
+                .when(() -> LocalServices.getService(DeviceIdleInternal.class));
         // Used in FlexibilityController.FcConstants.
         doAnswer((Answer<Void>) invocationOnMock -> null)
                 .when(() -> DeviceConfig.addOnPropertiesChangedListener(
@@ -146,7 +164,7 @@
                         eq(DeviceConfig.NAMESPACE_JOB_SCHEDULER), ArgumentMatchers.<String>any()));
         //used to get jobs by UID
         mJobStore = JobStore.initAndGetForTesting(mContext, mContext.getFilesDir());
-        when(mJobSchedulerService.getJobStore()).thenReturn(mJobStore);
+        doReturn(mJobStore).when(mJobSchedulerService).getJobStore();
         // Used in JobStatus.
         doReturn(mock(PackageManagerInternal.class))
                 .when(() -> LocalServices.getService(PackageManagerInternal.class));
@@ -156,6 +174,8 @@
         JobSchedulerService.sElapsedRealtimeClock =
                 Clock.fixed(Instant.ofEpochMilli(FROZEN_TIME), ZoneOffset.UTC);
         // Initialize real objects.
+        ArgumentCaptor<BroadcastReceiver> receiverCaptor =
+                ArgumentCaptor.forClass(BroadcastReceiver.class);
         mFlexibilityController = new FlexibilityController(mJobSchedulerService,
                 mPrefetchController);
         mFcConfig = mFlexibilityController.getFcConfig();
@@ -166,6 +186,11 @@
         setDeviceConfigLong(KEY_DEADLINE_PROXIMITY_LIMIT, 0L);
         setDeviceConfigInt(KEY_APPLIED_CONSTRAINTS, FLEXIBLE_CONSTRAINTS);
         waitForQuietModuleThread();
+
+        verify(mContext).registerReceiver(receiverCaptor.capture(),
+                ArgumentMatchers.argThat(filter ->
+                        filter.hasAction(PowerManager.ACTION_POWER_SAVE_WHITELIST_CHANGED)));
+        mBroadcastReceiver = receiverCaptor.getValue();
     }
 
     @After
@@ -212,6 +237,7 @@
         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)));
@@ -598,10 +624,10 @@
     @Test
     public void testGetLifeCycleBeginningElapsedLocked_Prefetch() {
         // prefetch with lifecycle
-        when(mPrefetchController.getLaunchTimeThresholdMs()).thenReturn(700L);
+        doReturn(700L).when(mPrefetchController).getLaunchTimeThresholdMs();
         JobInfo.Builder jb = createJob(0).setPrefetch(true);
         JobStatus js = createJobStatus("time", jb);
-        when(mPrefetchController.getNextEstimatedLaunchTimeLocked(js)).thenReturn(900L);
+        doReturn(900L).when(mPrefetchController).getNextEstimatedLaunchTimeLocked(js);
         assertEquals(900L - 700L, mFlexibilityController.getLifeCycleBeginningElapsedLocked(js));
         // prefetch with enqueue
         jb = createJob(0).setPrefetch(true);
@@ -616,7 +642,7 @@
         // prefetch without estimate
         mFlexibilityController.mPrefetchLifeCycleStart
                 .add(js.getUserId(), js.getSourcePackageName(), 500L);
-        when(mPrefetchController.getNextEstimatedLaunchTimeLocked(js)).thenReturn(Long.MAX_VALUE);
+        doReturn(Long.MAX_VALUE).when(mPrefetchController).getNextEstimatedLaunchTimeLocked(js);
         jb = createJob(0).setPrefetch(true);
         js = createJobStatus("time", jb);
         assertEquals(500L, mFlexibilityController.getLifeCycleBeginningElapsedLocked(js));
@@ -642,12 +668,12 @@
         // prefetch no estimate
         JobInfo.Builder jb = createJob(0).setPrefetch(true);
         JobStatus js = createJobStatus("time", jb);
-        when(mPrefetchController.getNextEstimatedLaunchTimeLocked(js)).thenReturn(Long.MAX_VALUE);
+        doReturn(Long.MAX_VALUE).when(mPrefetchController).getNextEstimatedLaunchTimeLocked(js);
         assertEquals(Long.MAX_VALUE, mFlexibilityController.getLifeCycleEndElapsedLocked(js, 0));
         // prefetch with estimate
         jb = createJob(0).setPrefetch(true);
         js = createJobStatus("time", jb);
-        when(mPrefetchController.getNextEstimatedLaunchTimeLocked(js)).thenReturn(1000L);
+        doReturn(1000L).when(mPrefetchController).getNextEstimatedLaunchTimeLocked(js);
         assertEquals(1000L, mFlexibilityController.getLifeCycleEndElapsedLocked(js, 0));
     }
 
@@ -696,7 +722,7 @@
         // Stop satisfied constraints from causing a false positive.
         js.setNumAppliedFlexibleConstraints(100);
         synchronized (mFlexibilityController.mLock) {
-            when(mJobSchedulerService.isCurrentlyRunningLocked(js)).thenReturn(true);
+            doReturn(true).when(mJobSchedulerService).isCurrentlyRunningLocked(js);
             assertTrue(mFlexibilityController.isFlexibilitySatisfiedLocked(js));
         }
     }
@@ -847,14 +873,85 @@
     }
 
     @Test
+    public void testAllowlistedAppBypass() {
+        setPowerWhitelistExceptIdle();
+        mFlexibilityController.onSystemServicesReady();
+
+        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);
+
+        setPowerWhitelistExceptIdle();
+        synchronized (mFlexibilityController.mLock) {
+            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsHigh));
+            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsDefault));
+            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsLow));
+            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsMin));
+        }
+
+        setPowerWhitelistExceptIdle(SOURCE_PACKAGE);
+        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));
+
+        doReturn(JobInfo.BIAS_DEFAULT).when(mJobSchedulerService).getUidBias(mSourceUid);
+        synchronized (mFlexibilityController.mLock) {
+            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsHigh));
+            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsDefault));
+            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsLow));
+            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsMin));
+        }
+
+        setUidBias(mSourceUid, JobInfo.BIAS_BOUND_FOREGROUND_SERVICE);
+        synchronized (mFlexibilityController.mLock) {
+            assertTrue(mFlexibilityController.isFlexibilitySatisfiedLocked(jsHigh));
+            assertTrue(mFlexibilityController.isFlexibilitySatisfiedLocked(jsDefault));
+            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsLow));
+            assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(jsMin));
+        }
+
+        setUidBias(mSourceUid, 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);
+        JobInfo.Builder jb = createJob(0).setPriority(JobInfo.PRIORITY_MIN);
         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_FOREGROUND_SERVICE);
+                .thenReturn(JobInfo.BIAS_DEFAULT);
 
         synchronized (mFlexibilityController.mLock) {
             mFlexibilityController.maybeStartTrackingJobLocked(js, null);
@@ -865,7 +962,7 @@
             assertTrue(mFlexibilityController.isFlexibilitySatisfiedLocked(js));
             assertTrue(js.isConstraintSatisfied(CONSTRAINT_FLEXIBLE));
 
-            setUidBias(mSourceUid, JobInfo.BIAS_FOREGROUND_SERVICE);
+            setUidBias(mSourceUid, JobInfo.BIAS_SYNC_INITIALIZATION);
 
             assertFalse(mFlexibilityController.isFlexibilitySatisfiedLocked(js));
             assertFalse(js.isConstraintSatisfied(CONSTRAINT_FLEXIBLE));
@@ -1187,9 +1284,9 @@
         JobInfo.Builder jb = createJob(22).setPrefetch(true);
         JobStatus js = createJobStatus("onPrefetchCacheUpdated", jb);
         jobs.add(js);
-        when(mPrefetchController.getLaunchTimeThresholdMs()).thenReturn(7 * HOUR_IN_MILLIS);
-        when(mPrefetchController.getNextEstimatedLaunchTimeLocked(js)).thenReturn(
-                1150L + mFlexibilityController.mConstants.PREFETCH_FORCE_BATCH_RELAX_THRESHOLD_MS);
+        doReturn(7 * HOUR_IN_MILLIS).when(mPrefetchController).getLaunchTimeThresholdMs();
+        doReturn(1150L + mFlexibilityController.mConstants.PREFETCH_FORCE_BATCH_RELAX_THRESHOLD_MS)
+                .when(mPrefetchController).getNextEstimatedLaunchTimeLocked(js);
 
         mFlexibilityController.maybeStartTrackingJobLocked(js, null);
 
@@ -1245,7 +1342,6 @@
         setUidBias(mSourceUid, BIAS_FOREGROUND_SERVICE);
         assertEquals(100L, (long) mFlexibilityController
                 .mPrefetchLifeCycleStart.get(js.getSourceUserId(), js.getSourcePackageName()));
-
     }
 
     @Test
@@ -1259,7 +1355,7 @@
     }
 
     private void runTestUnsupportedDevice(String feature) {
-        when(mPackageManager.hasSystemFeature(feature)).thenReturn(true);
+        doReturn(true).when(mPackageManager).hasSystemFeature(feature);
         mFlexibilityController =
                 new FlexibilityController(mJobSchedulerService, mPrefetchController);
         assertFalse(mFlexibilityController.isEnabled());
@@ -1279,6 +1375,16 @@
         }
     }
 
+    private void setPowerWhitelistExceptIdle(String... packages) {
+        doReturn(packages == null ? EmptyArray.STRING : packages)
+                .when(mDeviceIdleInternal).getFullPowerWhitelistExceptIdle();
+        if (mBroadcastReceiver != null) {
+            mBroadcastReceiver.onReceive(mContext,
+                    new Intent(PowerManager.ACTION_POWER_SAVE_WHITELIST_CHANGED));
+            waitForQuietModuleThread();
+        }
+    }
+
     private void setUidBias(int uid, int bias) {
         int prevBias = mJobSchedulerService.getUidBias(uid);
         doReturn(bias).when(mJobSchedulerService).getUidBias(uid);