From 2c5045d41a797507a24b68e5393f57bee96f1336 Mon Sep 17 00:00:00 2001
From: Kweku Adams <kwekua@google.com>
Date: Thu, 19 Nov 2020 10:14:58 -0800
Subject: [PATCH] Explicitly ensure the standby bucket is updated in testing.

Also add a missing switch case.

Bug: 171305774
Test: atest --rerun-until-failure 5 com.android.server.job.controllers.QuotaControllerTest
Change-Id: I8fe79032da6c4f7d386729ce0cebc96f52805c95
---
 .../job/controllers/QuotaController.java      | 69 ++++++++++---------
 .../job/controllers/QuotaControllerTest.java  |  3 +
 2 files changed, 40 insertions(+), 32 deletions(-)

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 7d7de3be8249c..1d72b42d824d3 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
@@ -1679,42 +1679,47 @@ public final class QuotaController extends StateController {
             // Update job bookkeeping out of band.
             JobSchedulerBackgroundThread.getHandler().post(() -> {
                 final int bucketIndex = JobSchedulerService.standbyBucketToBucketIndex(bucket);
-                if (DEBUG) {
-                    Slog.i(TAG, "Moving pkg " + string(userId, packageName) + " to bucketIndex "
-                            + bucketIndex);
-                }
-                List<JobStatus> restrictedChanges = new ArrayList<>();
-                synchronized (mLock) {
-                    ArraySet<JobStatus> jobs = mTrackedJobs.get(userId, packageName);
-                    if (jobs == null || jobs.size() == 0) {
-                        return;
-                    }
-                    for (int i = jobs.size() - 1; i >= 0; i--) {
-                        JobStatus js = jobs.valueAt(i);
-                        // Effective standby bucket can change after this in some situations so
-                        // use the real bucket so that the job is tracked by the controllers.
-                        if ((bucketIndex == RESTRICTED_INDEX
-                                || js.getStandbyBucket() == RESTRICTED_INDEX)
-                                && bucketIndex != js.getStandbyBucket()) {
-                            restrictedChanges.add(js);
-                        }
-                        js.setStandbyBucket(bucketIndex);
-                    }
-                    Timer timer = mPkgTimers.get(userId, packageName);
-                    if (timer != null && timer.isActive()) {
-                        timer.rescheduleCutoff();
-                    }
-                    if (maybeUpdateConstraintForPkgLocked(userId, packageName)) {
-                        mStateChangedListener.onControllerStateChanged();
-                    }
-                }
-                if (restrictedChanges.size() > 0) {
-                    mStateChangedListener.onRestrictedBucketChanged(restrictedChanges);
-                }
+                updateStandbyBucket(userId, packageName, bucketIndex);
             });
         }
     }
 
+    @VisibleForTesting
+    void updateStandbyBucket(
+            final int userId, final @NonNull String packageName, final int bucketIndex) {
+        if (DEBUG) {
+            Slog.i(TAG, "Moving pkg " + string(userId, packageName)
+                    + " to bucketIndex " + bucketIndex);
+        }
+        List<JobStatus> restrictedChanges = new ArrayList<>();
+        synchronized (mLock) {
+            ArraySet<JobStatus> jobs = mTrackedJobs.get(userId, packageName);
+            if (jobs == null || jobs.size() == 0) {
+                return;
+            }
+            for (int i = jobs.size() - 1; i >= 0; i--) {
+                JobStatus js = jobs.valueAt(i);
+                // Effective standby bucket can change after this in some situations so
+                // use the real bucket so that the job is tracked by the controllers.
+                if ((bucketIndex == RESTRICTED_INDEX || js.getStandbyBucket() == RESTRICTED_INDEX)
+                        && bucketIndex != js.getStandbyBucket()) {
+                    restrictedChanges.add(js);
+                }
+                js.setStandbyBucket(bucketIndex);
+            }
+            Timer timer = mPkgTimers.get(userId, packageName);
+            if (timer != null && timer.isActive()) {
+                timer.rescheduleCutoff();
+            }
+            if (maybeUpdateConstraintForPkgLocked(userId, packageName)) {
+                mStateChangedListener.onControllerStateChanged();
+            }
+        }
+        if (restrictedChanges.size() > 0) {
+            mStateChangedListener.onRestrictedBucketChanged(restrictedChanges);
+        }
+    }
+
     private final class DeleteTimingSessionsFunctor implements Consumer<List<TimingSession>> {
         private final Predicate<TimingSession> mTooOld = new Predicate<TimingSession>() {
             public boolean test(TimingSession ts) {
diff --git a/services/tests/mockingservicestests/src/com/android/server/job/controllers/QuotaControllerTest.java b/services/tests/mockingservicestests/src/com/android/server/job/controllers/QuotaControllerTest.java
index 77fef1274415a..8795d7711c633 100644
--- a/services/tests/mockingservicestests/src/com/android/server/job/controllers/QuotaControllerTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/job/controllers/QuotaControllerTest.java
@@ -284,6 +284,8 @@ public class QuotaControllerTest {
                 return UsageStatsManager.STANDBY_BUCKET_FREQUENT;
             case RARE_INDEX:
                 return UsageStatsManager.STANDBY_BUCKET_RARE;
+            case RESTRICTED_INDEX:
+                return UsageStatsManager.STANDBY_BUCKET_RESTRICTED;
             default:
                 return UsageStatsManager.STANDBY_BUCKET_NEVER;
         }
@@ -292,6 +294,7 @@ public class QuotaControllerTest {
     private void setStandbyBucket(int bucketIndex) {
         when(mUsageStatsManager.getAppStandbyBucket(eq(SOURCE_PACKAGE), eq(SOURCE_USER_ID),
                 anyLong())).thenReturn(bucketIndexToUsageStatsBucket(bucketIndex));
+        mQuotaController.updateStandbyBucket(SOURCE_USER_ID, SOURCE_PACKAGE, bucketIndex);
     }
 
     private void setStandbyBucket(int bucketIndex, JobStatus... jobs) {