Merge "Fix support for BatteryUsageStatsQuery.accumulated().setMaxAge()" into main

7 files changed, 180 insertions, 100 deletions

diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index c27126a01a32..aea24d978bee 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -1101,6 +1101,9 @@ public final class BatteryStatsService extends IBatteryStats.Stub
 
     /** StatsPullAtomCallback for pulling BatteryUsageStats data. */
     private class StatsPullAtomCallbackImpl implements StatsManager.StatsPullAtomCallback {
+        private static final long BATTERY_USAGE_STATS_PER_UID_MAX_STATS_AGE =
+                TimeUnit.HOURS.toMillis(2);
+
         @Override
         public int onPullAtom(int atomTag, List<StatsEvent> data) {
             final BatteryUsageStats bus;
@@ -1168,7 +1171,8 @@ public final class BatteryStatsService extends IBatteryStats.Stub
                             .setMinConsumedPowerThreshold(minConsumedPowerThreshold);
 
                     if (isBatteryUsageStatsAccumulationSupported()) {
-                        query.accumulated();
+                        query.accumulated()
+                                .setMaxStatsAgeMs(BATTERY_USAGE_STATS_PER_UID_MAX_STATS_AGE);
                     }
 
                     bus = getBatteryUsageStats(List.of(query.build())).get(0);
diff --git a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
index 63e8d9973237..8c588b4c9b98 100644
--- a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
+++ b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
@@ -195,23 +195,22 @@ public class BatteryUsageStatsProvider {
             mLastAccumulationMonotonicHistorySize = historySize;
         }
 
-        handler.post(() -> accumulateBatteryUsageStats(stats));
+        // No need to store the accumulated stats asynchronously, as the entire accumulation
+        // operation is async
+        handler.post(() -> accumulateBatteryUsageStats(stats, false));
     }
 
     /**
      * Computes BatteryUsageStats for the period since the last accumulated stats were stored,
-     * adds them to the accumulated stats and saves the result.
+     * adds them to the accumulated stats and asynchronously saves the result.
      */
     public void accumulateBatteryUsageStats(BatteryStatsImpl stats) {
-        AccumulatedBatteryUsageStats accumulatedStats = loadAccumulatedBatteryUsageStats();
+        accumulateBatteryUsageStats(stats, true);
+    }
 
-        final BatteryUsageStatsQuery query = new BatteryUsageStatsQuery.Builder()
-                .setMaxStatsAgeMs(0)
-                .includeProcessStateData()
-                .includePowerStateData()
-                .includeScreenStateData()
-                .build();
-        updateAccumulatedBatteryUsageStats(accumulatedStats, stats, query);
+    private void accumulateBatteryUsageStats(BatteryStatsImpl stats, boolean storeAsync) {
+        AccumulatedBatteryUsageStats accumulatedStats = loadAccumulatedBatteryUsageStats();
+        updateAccumulatedBatteryUsageStats(accumulatedStats, stats);
 
         PowerStatsSpan powerStatsSpan = new PowerStatsSpan(AccumulatedBatteryUsageStatsSection.ID);
         powerStatsSpan.addSection(
@@ -220,8 +219,13 @@ public class BatteryUsageStatsProvider {
                 accumulatedStats.startWallClockTime,
                 accumulatedStats.endMonotonicTime - accumulatedStats.startMonotonicTime);
         mMonotonicClock.write();
-        mPowerStatsStore.storePowerStatsSpanAsync(powerStatsSpan,
-                accumulatedStats.builder::discard);
+        if (storeAsync) {
+            mPowerStatsStore.storePowerStatsSpanAsync(powerStatsSpan,
+                    accumulatedStats.builder::discard);
+        } else {
+            mPowerStatsStore.storePowerStatsSpan(powerStatsSpan);
+            accumulatedStats.builder.discard();
+        }
     }
 
     /**
@@ -269,7 +273,7 @@ public class BatteryUsageStatsProvider {
         BatteryUsageStats batteryUsageStats;
         if ((query.getFlags()
                 & BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_ACCUMULATED) != 0) {
-            batteryUsageStats = getAccumulatedBatteryUsageStats(stats, query, currentTimeMs);
+            batteryUsageStats = getAccumulatedBatteryUsageStats(stats, query);
         } else if (query.getAggregatedToTimestamp() == 0) {
             BatteryUsageStats.Builder builder = computeBatteryUsageStats(stats, query,
                     query.getMonotonicStartTime(),
@@ -288,9 +292,13 @@ public class BatteryUsageStatsProvider {
     }
 
     private BatteryUsageStats getAccumulatedBatteryUsageStats(BatteryStatsImpl stats,
-            BatteryUsageStatsQuery query, long currentTimeMs) {
+            BatteryUsageStatsQuery query) {
         AccumulatedBatteryUsageStats accumulatedStats = loadAccumulatedBatteryUsageStats();
-        updateAccumulatedBatteryUsageStats(accumulatedStats, stats, query);
+        if (accumulatedStats.endMonotonicTime == MonotonicClock.UNDEFINED
+                || mMonotonicClock.monotonicTime() - accumulatedStats.endMonotonicTime
+                > query.getMaxStatsAge()) {
+            updateAccumulatedBatteryUsageStats(accumulatedStats, stats);
+        }
         return accumulatedStats.builder.build();
     }
 
@@ -321,7 +329,7 @@ public class BatteryUsageStatsProvider {
     }
 
     private void updateAccumulatedBatteryUsageStats(AccumulatedBatteryUsageStats accumulatedStats,
-            BatteryStatsImpl stats, BatteryUsageStatsQuery query) {
+            BatteryStatsImpl stats) {
         long startMonotonicTime = accumulatedStats.endMonotonicTime;
         if (startMonotonicTime == MonotonicClock.UNDEFINED) {
             startMonotonicTime = stats.getMonotonicStartTime();
@@ -333,6 +341,7 @@ public class BatteryUsageStatsProvider {
             accumulatedStats.builder = new BatteryUsageStats.Builder(
                     stats.getCustomEnergyConsumerNames(), true, true, true, 0);
             accumulatedStats.startWallClockTime = stats.getStartClockTime();
+            accumulatedStats.startMonotonicTime = stats.getMonotonicStartTime();
             accumulatedStats.builder.setStatsStartTimestamp(accumulatedStats.startWallClockTime);
         }
 
@@ -342,7 +351,7 @@ public class BatteryUsageStatsProvider {
         accumulatedStats.builder.setStatsDuration(endWallClockTime - startMonotonicTime);
 
         mPowerAttributor.estimatePowerConsumption(accumulatedStats.builder, stats.getHistory(),
-                startMonotonicTime, MonotonicClock.UNDEFINED);
+                startMonotonicTime, endMonotonicTime);
 
         populateGeneralInfo(accumulatedStats.builder, stats);
     }
diff --git a/services/core/java/com/android/server/power/stats/processor/PowerStatsAggregator.java b/services/core/java/com/android/server/power/stats/processor/PowerStatsAggregator.java
index dcdd3bd8b3fa..2609cf7aebfc 100644
--- a/services/core/java/com/android/server/power/stats/processor/PowerStatsAggregator.java
+++ b/services/core/java/com/android/server/power/stats/processor/PowerStatsAggregator.java
@@ -23,6 +23,7 @@ import android.util.SparseBooleanArray;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.os.BatteryStatsHistory;
 import com.android.internal.os.BatteryStatsHistoryIterator;
+import com.android.internal.os.MonotonicClock;
 
 import java.util.function.Consumer;
 
@@ -169,6 +170,9 @@ public class PowerStatsAggregator {
                     }
                 }
             }
+            if (endTimeMs != MonotonicClock.UNDEFINED) {
+                lastTime = endTimeMs;
+            }
             if (lastTime > baseTime) {
                 mStats.setDuration(lastTime - baseTime);
                 mStats.finish(lastTime);
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java
index 709f83ba907d..73dcfe77e67f 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java
@@ -36,6 +36,7 @@ import android.os.BatteryStats;
 import android.os.BatteryUsageStats;
 import android.os.BatteryUsageStatsQuery;
 import android.os.ConditionVariable;
+import android.os.Handler;
 import android.os.Parcel;
 import android.os.Process;
 import android.os.UidBatteryConsumer;
@@ -81,8 +82,9 @@ public class BatteryUsageStatsProviderTest {
                     .setAveragePower(PowerProfile.POWER_BATTERY_CAPACITY, 4000.0);
 
     private MockClock mMockClock = mStatsRule.getMockClock();
-    private MonotonicClock mMonotonicClock = new MonotonicClock(666777, mMockClock);
+    private MonotonicClock mMonotonicClock = mStatsRule.getMonotonicClock();
     private Context mContext;
+    private PowerStatsStore mPowerStatsStore;
 
     @Before
     public void setup() throws IOException {
@@ -93,6 +95,9 @@ public class BatteryUsageStatsProviderTest {
         } else {
             mContext = InstrumentationRegistry.getContext();
         }
+        mPowerStatsStore = spy(new PowerStatsStore(
+                new File(mStatsRule.getHistoryDir(), getClass().getSimpleName()),
+                mStatsRule.getHandler()));
     }
 
     @Test
@@ -274,10 +279,7 @@ public class BatteryUsageStatsProviderTest {
         powerAttributor.setPowerComponentSupported(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT,
                 true);
 
-        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
-                powerAttributor, mStatsRule.getPowerProfile(),
-                mStatsRule.getCpuScalingPolicies(), mock(PowerStatsStore.class), 0, mMockClock,
-                mMonotonicClock);
+        BatteryUsageStatsProvider provider = createBatteryUsageStatsProvider(0);
 
         return provider.getBatteryUsageStats(batteryStats, BatteryUsageStatsQuery.DEFAULT);
     }
@@ -331,30 +333,30 @@ public class BatteryUsageStatsProviderTest {
         BatteryStats.HistoryItem item;
 
         assertThat(item = iterator.next()).isNotNull();
-        assertHistoryItem(item,
+        assertHistoryItem(batteryStats, item,
                 BatteryStats.HistoryItem.CMD_RESET, BatteryStats.HistoryItem.EVENT_NONE,
                 null, 0, 3_600_000, 90, 1_000_000);
 
         assertThat(item = iterator.next()).isNotNull();
-        assertHistoryItem(item,
+        assertHistoryItem(batteryStats, item,
                 BatteryStats.HistoryItem.CMD_UPDATE, BatteryStats.HistoryItem.EVENT_NONE,
                 null, 0, 3_600_000, 90, 1_000_000);
         assertThat(item.states & BatteryStats.HistoryItem.STATE_CPU_RUNNING_FLAG).isNotEqualTo(0);
 
         assertThat(item = iterator.next()).isNotNull();
-        assertHistoryItem(item,
+        assertHistoryItem(batteryStats, item,
                 BatteryStats.HistoryItem.CMD_UPDATE, BatteryStats.HistoryItem.EVENT_NONE,
                 null, 0, 3_600_000, 90, 2_000_000);
         assertThat(item.states & BatteryStats.HistoryItem.STATE_CPU_RUNNING_FLAG).isEqualTo(0);
 
         assertThat(item = iterator.next()).isNotNull();
-        assertHistoryItem(item,
+        assertHistoryItem(batteryStats, item,
                 BatteryStats.HistoryItem.CMD_UPDATE,
                 BatteryStats.HistoryItem.EVENT_ALARM | BatteryStats.HistoryItem.EVENT_FLAG_START,
                 "foo", APP_UID, 3_600_000, 90, 3_000_000);
 
         assertThat(item = iterator.next()).isNotNull();
-        assertHistoryItem(item,
+        assertHistoryItem(batteryStats, item,
                 BatteryStats.HistoryItem.CMD_UPDATE,
                 BatteryStats.HistoryItem.EVENT_ALARM | BatteryStats.HistoryItem.EVENT_FLAG_FINISH,
                 "foo", APP_UID, 3_600_000, 90, 3_001_000);
@@ -441,14 +443,15 @@ public class BatteryUsageStatsProviderTest {
             assertThat(item.eventTag.string).startsWith(uid + " ");
             assertThat(item.batteryChargeUah).isEqualTo(3_600_000);
             assertThat(item.batteryLevel).isEqualTo(90);
-            assertThat(item.time).isEqualTo((long) 1_000_000);
+            assertThat(item.time).isEqualTo(batteryStats.getMonotonicStartTime() + 1_000_000);
         }
 
         assertThat(expectedUid).isEqualTo(200);
     }
 
-    private void assertHistoryItem(BatteryStats.HistoryItem item, int command, int eventCode,
-            String tag, int uid, int batteryChargeUah, int batteryLevel, long elapsedTimeMs) {
+    private void assertHistoryItem(MockBatteryStatsImpl batteryStats, BatteryStats.HistoryItem item,
+            int command, int eventCode, String tag, int uid, int batteryChargeUah, int batteryLevel,
+            long elapsedTimeMs) {
         assertThat(item.cmd).isEqualTo(command);
         assertThat(item.eventCode).isEqualTo(eventCode);
         if (tag == null) {
@@ -460,7 +463,7 @@ public class BatteryUsageStatsProviderTest {
         assertThat(item.batteryChargeUah).isEqualTo(batteryChargeUah);
         assertThat(item.batteryLevel).isEqualTo(batteryLevel);
 
-        assertThat(item.time).isEqualTo(elapsedTimeMs);
+        assertThat(item.time).isEqualTo(batteryStats.getMonotonicStartTime() + elapsedTimeMs);
     }
 
     @Test
@@ -566,38 +569,66 @@ public class BatteryUsageStatsProviderTest {
 
         assertThat(stats.getStatsStartTimestamp()).isEqualTo(5 * MINUTE_IN_MS);
         assertThat(stats.getStatsEndTimestamp()).isEqualTo(55 * MINUTE_IN_MS);
-        assertThat(stats.getAggregateBatteryConsumer(
-                        BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE)
-                .getConsumedPower(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT))
-                .isWithin(0.0001)
-                .of(180.0);  // 360 mA * 0.5 hours
-        assertThat(stats.getAggregateBatteryConsumer(
-                        BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE)
-                .getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT))
-                .isEqualTo((10 + 20) * MINUTE_IN_MS);
-        final UidBatteryConsumer uidBatteryConsumer = stats.getUidBatteryConsumers().stream()
-                .filter(uid -> uid.getUid() == APP_UID).findFirst().get();
-        assertThat(uidBatteryConsumer
-                .getConsumedPower(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT))
-                .isWithin(0.1)
-                .of(180.0);
+        assertBatteryConsumer(stats, 180.0, (10 + 20) * MINUTE_IN_MS);
+        assertBatteryConsumer(stats, APP_UID, 180.0, (10 + 20) * MINUTE_IN_MS);
 
         stats.close();
     }
 
     @Test
     public void accumulateBatteryUsageStats() throws Throwable {
-        accumulateBatteryUsageStats(10000000, 1);
+        MockBatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
+        accumulateBatteryUsageStats(batteryStats, 10000000, 0);
         // Accumulate every 200 bytes of battery history
-        accumulateBatteryUsageStats(200, 2);
-        accumulateBatteryUsageStats(50, 5);
+        accumulateBatteryUsageStats(batteryStats, 200, 2);
+        accumulateBatteryUsageStats(batteryStats, 50, 4);
         // Accumulate on every invocation of accumulateBatteryUsageStats
-        accumulateBatteryUsageStats(0, 7);
+        accumulateBatteryUsageStats(batteryStats, 0, 7);
     }
 
-    private void accumulateBatteryUsageStats(int accumulatedBatteryUsageStatsSpanSize,
-            int expectedNumberOfUpdates) throws Throwable {
-        BatteryStatsImpl batteryStats = spy(mStatsRule.getBatteryStats());
+    @Test
+    public void getAccumulatedBatteryUsageStats() throws Throwable {
+        MockBatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
+
+        // Only accumulate the first 25 minutes
+        accumulateBatteryUsageStats(batteryStats, 200, 1);
+
+        BatteryUsageStatsProvider batteryUsageStatsProvider = createBatteryUsageStatsProvider(200);
+
+        // At this point the last stored accumulated stats are `115 - 30 = 85` minutes old
+        BatteryUsageStats stats = batteryUsageStatsProvider.getBatteryUsageStats(batteryStats,
+                new BatteryUsageStatsQuery.Builder()
+                        .accumulated()
+                        .setMaxStatsAgeMs(90 * MINUTE_IN_MS)
+                        .build());
+
+        assertThat(stats.getStatsStartTimestamp()).isEqualTo(5 * MINUTE_IN_MS);
+        assertThat(stats.getStatsEndTimestamp()).isEqualTo(30 * MINUTE_IN_MS);
+        assertBatteryConsumer(stats, 60.0, 10 * MINUTE_IN_MS);
+        assertBatteryConsumer(stats, APP_UID, 60.0, 10 * MINUTE_IN_MS);
+
+        stats.close();
+
+        // Now force the usage stats to catch up to the current time
+        stats = batteryUsageStatsProvider.getBatteryUsageStats(batteryStats,
+                new BatteryUsageStatsQuery.Builder()
+                        .accumulated()
+                        .setMaxStatsAgeMs(5 * MINUTE_IN_MS)
+                        .build());
+
+        assertThat(stats.getStatsStartTimestamp()).isEqualTo(5 * MINUTE_IN_MS);
+        assertThat(stats.getStatsEndTimestamp()).isEqualTo(115 * MINUTE_IN_MS);
+        assertBatteryConsumer(stats, 360.0, 60 * MINUTE_IN_MS);
+        assertBatteryConsumer(stats, APP_UID, 360.0, 60 * MINUTE_IN_MS);
+
+        stats.close();
+    }
+
+    private void accumulateBatteryUsageStats(MockBatteryStatsImpl batteryStatsImpl,
+            int accumulatedBatteryUsageStatsSpanSize, int expectedNumberOfUpdates)
+            throws Throwable {
+        Handler handler = mStatsRule.getHandler();
+        MockBatteryStatsImpl batteryStats = spy(batteryStatsImpl);
         // Note - these two are in microseconds
         when(batteryStats.computeBatteryTimeRemaining(anyLong())).thenReturn(111_000L);
         when(batteryStats.computeChargeTimeRemaining(anyLong())).thenReturn(777_000L);
@@ -610,82 +641,76 @@ public class BatteryUsageStatsProviderTest {
             batteryStats.resetAllStatsAndHistoryLocked(BatteryStatsImpl.RESET_REASON_ADB_COMMAND);
         }
 
-        PowerStatsStore powerStatsStore = spy(new PowerStatsStore(
-                new File(mStatsRule.getHistoryDir(), getClass().getSimpleName()),
-                mStatsRule.getHandler()));
-        powerStatsStore.reset();
+        mPowerStatsStore.reset();
 
         int[] count = new int[1];
         doAnswer(inv -> {
             count[0]++;
-            return null;
-        }).when(powerStatsStore).storePowerStatsSpan(any(PowerStatsSpan.class));
+            return inv.callRealMethod();
+        }).when(mPowerStatsStore).storePowerStatsSpan(any(PowerStatsSpan.class));
 
-        MultiStatePowerAttributor powerAttributor = new MultiStatePowerAttributor(mContext,
-                powerStatsStore, mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies(),
-                () -> 3500);
-        for (int powerComponentId = 0; powerComponentId < BatteryConsumer.POWER_COMPONENT_COUNT;
-                powerComponentId++) {
-            powerAttributor.setPowerComponentSupported(powerComponentId, true);
-        }
-        powerAttributor.setPowerComponentSupported(BatteryConsumer.POWER_COMPONENT_ANY, true);
-
-        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
-                powerAttributor, mStatsRule.getPowerProfile(),
-                mStatsRule.getCpuScalingPolicies(), powerStatsStore,
-                accumulatedBatteryUsageStatsSpanSize, mMockClock, mMonotonicClock);
+        BatteryUsageStatsProvider batteryUsageStatsProvider = createBatteryUsageStatsProvider(
+                accumulatedBatteryUsageStatsSpanSize);
 
-        provider.accumulateBatteryUsageStatsAsync(batteryStats, mStatsRule.getHandler());
+        batteryUsageStatsProvider.accumulateBatteryUsageStatsAsync(batteryStats, handler);
 
+        setTime(10 * MINUTE_IN_MS);
         synchronized (batteryStats) {
             batteryStats.noteFlashlightOnLocked(APP_UID,
                     10 * MINUTE_IN_MS, 10 * MINUTE_IN_MS);
         }
 
-        provider.accumulateBatteryUsageStatsAsync(batteryStats, mStatsRule.getHandler());
+        batteryUsageStatsProvider.accumulateBatteryUsageStatsAsync(batteryStats, handler);
 
+        setTime(20 * MINUTE_IN_MS);
         synchronized (batteryStats) {
             batteryStats.noteFlashlightOffLocked(APP_UID,
                     20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
         }
 
-        provider.accumulateBatteryUsageStatsAsync(batteryStats, mStatsRule.getHandler());
+        batteryUsageStatsProvider.accumulateBatteryUsageStatsAsync(batteryStats, handler);
 
+        setTime(30 * MINUTE_IN_MS);
         synchronized (batteryStats) {
             batteryStats.noteFlashlightOnLocked(APP_UID,
                     30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS);
         }
 
-        provider.accumulateBatteryUsageStatsAsync(batteryStats, mStatsRule.getHandler());
+        batteryUsageStatsProvider.accumulateBatteryUsageStatsAsync(batteryStats, handler);
 
+        // Make sure the accumulated stats are computed and saved before generating more history
+        mStatsRule.waitForBackgroundThread();
+
+        setTime(50 * MINUTE_IN_MS);
         synchronized (batteryStats) {
             batteryStats.noteFlashlightOffLocked(APP_UID,
                     50 * MINUTE_IN_MS, 50 * MINUTE_IN_MS);
         }
         setTime(55 * MINUTE_IN_MS);
 
-        provider.accumulateBatteryUsageStatsAsync(batteryStats, mStatsRule.getHandler());
+        batteryUsageStatsProvider.accumulateBatteryUsageStatsAsync(batteryStats, handler);
 
         // This section has not been saved yet, but should be added to the accumulated totals
+        setTime(80 * MINUTE_IN_MS);
         synchronized (batteryStats) {
             batteryStats.noteFlashlightOnLocked(APP_UID,
                     80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
         }
 
-        provider.accumulateBatteryUsageStatsAsync(batteryStats, mStatsRule.getHandler());
+        batteryUsageStatsProvider.accumulateBatteryUsageStatsAsync(batteryStats, handler);
 
+        setTime(110 * MINUTE_IN_MS);
         synchronized (batteryStats) {
             batteryStats.noteFlashlightOffLocked(APP_UID,
                     110 * MINUTE_IN_MS, 110 * MINUTE_IN_MS);
         }
         setTime(115 * MINUTE_IN_MS);
 
-        // Pick up the remainder of battery history that has not yet been accumulated
-        provider.accumulateBatteryUsageStats(batteryStats);
+        batteryUsageStatsProvider.accumulateBatteryUsageStatsAsync(batteryStats, handler);
 
         mStatsRule.waitForBackgroundThread();
 
-        BatteryUsageStats stats = provider.getBatteryUsageStats(batteryStats,
+        BatteryUsageStats stats = batteryUsageStatsProvider.getBatteryUsageStats(batteryStats,
                 new BatteryUsageStatsQuery.Builder().accumulated().build());
 
         assertThat(stats.getStatsStartTimestamp()).isEqualTo(5 * MINUTE_IN_MS);
@@ -696,29 +721,55 @@ public class BatteryUsageStatsProviderTest {
         assertThat(stats.getBatteryCapacity()).isEqualTo(4000);  // from PowerProfile
 
         // Total: 10 + 20 + 30 = 60
-        assertThat(stats.getAggregateBatteryConsumer(
-                        BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE)
+        assertBatteryConsumer(stats, 360.0, 60 * MINUTE_IN_MS);
+        assertBatteryConsumer(stats, APP_UID, 360.0, 60 * MINUTE_IN_MS);
+        stats.close();
+
+        mStatsRule.waitForBackgroundThread();
+
+        assertThat(count[0]).isEqualTo(expectedNumberOfUpdates);
+    }
+
+    private BatteryUsageStatsProvider createBatteryUsageStatsProvider(
+            int accumulatedBatteryUsageStatsSpanSize) {
+        MultiStatePowerAttributor powerAttributor = new MultiStatePowerAttributor(mContext,
+                mPowerStatsStore, mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies(),
+                () -> 3500);
+        for (int powerComponentId = 0; powerComponentId < BatteryConsumer.POWER_COMPONENT_COUNT;
+                powerComponentId++) {
+            powerAttributor.setPowerComponentSupported(powerComponentId, true);
+        }
+        powerAttributor.setPowerComponentSupported(BatteryConsumer.POWER_COMPONENT_ANY, true);
+
+        return new BatteryUsageStatsProvider(mContext, powerAttributor,
+                mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies(), mPowerStatsStore,
+                accumulatedBatteryUsageStatsSpanSize, mMockClock, mMonotonicClock);
+    }
+
+    private static void assertBatteryConsumer(BatteryUsageStats stats, double expectedPowerMah,
+            long expectedDurationMs) {
+        AggregateBatteryConsumer aggregatedConsumer = stats.getAggregateBatteryConsumer(
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE);
+        assertThat(aggregatedConsumer
                 .getConsumedPower(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT))
                 .isWithin(0.0001)
-                .of(360.0);  // 360 mA * 1.0 hour
-        assertThat(stats.getAggregateBatteryConsumer(
-                        BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE)
+                .of(expectedPowerMah);
+        assertThat(aggregatedConsumer
                 .getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT))
-                .isEqualTo(60 * MINUTE_IN_MS);
+                .isEqualTo(expectedDurationMs);
+    }
 
+    private static void assertBatteryConsumer(BatteryUsageStats stats, int uid,
+            double expectedPowerMah, long expectedDurationMs) {
         final UidBatteryConsumer uidBatteryConsumer = stats.getUidBatteryConsumers().stream()
-                .filter(uid -> uid.getUid() == APP_UID).findFirst().get();
+                .filter(u -> u.getUid() == uid).findFirst().get();
         assertThat(uidBatteryConsumer
                 .getConsumedPower(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT))
                 .isWithin(0.1)
-                .of(360.0);
+                .of(expectedPowerMah);
         assertThat(uidBatteryConsumer
                 .getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT))
-                .isEqualTo(60 * MINUTE_IN_MS);
-
-        assertThat(count[0]).isEqualTo(expectedNumberOfUpdates);
-
-        stats.close();
+                .isEqualTo(expectedDurationMs);
     }
 
     private void setTime(long timeMs) {
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java
index a3c7ece386c7..9e7e0b646047 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java
@@ -41,6 +41,7 @@ import android.util.SparseArray;
 import android.util.Xml;
 
 import com.android.internal.os.CpuScalingPolicies;
+import com.android.internal.os.MonotonicClock;
 import com.android.internal.os.PowerProfile;
 import com.android.internal.power.EnergyConsumerStats;
 
@@ -65,6 +66,7 @@ public class BatteryUsageStatsRule implements TestRule {
 
     private final PowerProfile mPowerProfile;
     private final MockClock mMockClock = new MockClock();
+    private final MonotonicClock mMonotonicClock = new MonotonicClock(666777, mMockClock);
     private String mTestName;
     private boolean mCreateTempDirectory;
     private File mHistoryDir;
@@ -118,7 +120,7 @@ public class BatteryUsageStatsRule implements TestRule {
             clearDirectory();
         }
         mBatteryStats = new MockBatteryStatsImpl(mBatteryStatsConfigBuilder.build(),
-                mMockClock, mHistoryDir, mHandler, new PowerStatsUidResolver());
+                mMockClock, mMonotonicClock, mHistoryDir, mHandler, new PowerStatsUidResolver());
         mBatteryStats.setPowerProfile(mPowerProfile);
         mBatteryStats.setCpuScalingPolicies(new CpuScalingPolicies(mCpusByPolicy, mFreqsByPolicy));
         synchronized (mBatteryStats) {
@@ -144,6 +146,10 @@ public class BatteryUsageStatsRule implements TestRule {
         return mMockClock;
     }
 
+    public MonotonicClock getMonotonicClock() {
+        return mMonotonicClock;
+    }
+
     public Handler getHandler() {
         return mHandler;
     }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/MockBatteryStatsImpl.java b/services/tests/powerstatstests/src/com/android/server/power/stats/MockBatteryStatsImpl.java
index b374a3202fa2..9a38209a7d17 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/MockBatteryStatsImpl.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/MockBatteryStatsImpl.java
@@ -77,9 +77,15 @@ public class MockBatteryStatsImpl extends BatteryStatsImpl {
                 new PowerStatsUidResolver());
     }
 
-    MockBatteryStatsImpl(BatteryStatsConfig config, Clock clock, File historyDirectory,
-            Handler handler, PowerStatsUidResolver powerStatsUidResolver) {
-        super(config, clock, new MonotonicClock(0, clock), historyDirectory, handler,
+    MockBatteryStatsImpl(BatteryStatsConfig config, Clock clock,
+            File historyDirectory, Handler handler, PowerStatsUidResolver powerStatsUidResolver) {
+        this(config, clock, new MonotonicClock(0, clock), historyDirectory, handler,
+                powerStatsUidResolver);
+    }
+
+    MockBatteryStatsImpl(BatteryStatsConfig config, Clock clock, MonotonicClock monotonicClock,
+            File historyDirectory, Handler handler, PowerStatsUidResolver powerStatsUidResolver) {
+        super(config, clock, monotonicClock, historyDirectory, handler,
                 mock(PlatformIdleStateCallback.class), mock(EnergyStatsRetriever.class),
                 mock(UserInfoProvider.class), mockPowerProfile(),
                 new CpuScalingPolicies(new SparseArray<>(), new SparseArray<>()),
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/PowerStatsExporterTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/PowerStatsExporterTest.java
index 38fe6134d992..d243f92a139f 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/PowerStatsExporterTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/PowerStatsExporterTest.java
@@ -408,7 +408,7 @@ public class PowerStatsExporterTest {
         BatteryUsageStats.Builder builder = new BatteryUsageStats.Builder(
                 new String[]{"cu570m"},
                 /* includeProcessStateData */ true, true, true, /* powerThreshold */ 0);
-        exportAggregatedPowerStats(builder, 3700, 6700);
+        exportAggregatedPowerStats(builder, 3700, 7500);
 
         BatteryUsageStats actual = builder.build();
         String message = "Actual BatteryUsageStats: " + actual;