Add BatteryUsageStats.add(BatteryUsageStats)

Bug: 187223764
Test: atest FrameworksCoreTests:com.android.internal.os.BatteryUsageStatsTest
Change-Id: I1e5060d726fd0f4634fd658483e80a9ab02e24fb
Merged-In: I1e5060d726fd0f4634fd658483e80a9ab02e24fb
(cherry picked from commit 63d9aecdbb)

core/java/android/os/AggregateBatteryConsumer.java

@@ -90,6 +90,14 @@ public final class AggregateBatteryConsumer extends BatteryConsumer implements P
             return this;
         }
 
+        /**
+         * Adds power and usage duration from the supplied AggregateBatteryConsumer.
+         */
+        public void add(AggregateBatteryConsumer aggregateBatteryConsumer) {
+            mConsumedPowerMah += aggregateBatteryConsumer.mConsumedPowerMah;
+            mPowerComponentsBuilder.addPowerAndDuration(aggregateBatteryConsumer.mPowerComponents);
+        }
+
         /**
          * Creates a read-only object out of the Builder values.
          */

core/java/android/os/BatteryUsageStats.java

@@ -30,6 +30,7 @@ import java.io.PrintWriter;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Comparator;
 import java.util.List;
 
@@ -96,11 +97,7 @@ public final class BatteryUsageStats implements Parcelable {
     private BatteryUsageStats(@NonNull Builder builder) {
         mStatsStartTimestampMs = builder.mStatsStartTimestampMs;
         mStatsEndTimestampMs = builder.mStatsEndTimestampMs;
-        if (builder.mStatsDurationMs != -1) {
+        mStatsDurationMs = builder.getStatsDuration();
-            mStatsDurationMs = builder.mStatsDurationMs;
-        } else {
-            mStatsDurationMs = mStatsEndTimestampMs - mStatsStartTimestampMs;
-        }
         mBatteryCapacityMah = builder.mBatteryCapacityMah;
         mDischargePercentage = builder.mDischargePercentage;
         mDischargedPowerLowerBound = builder.mDischargedPowerLowerBoundMah;
@@ -658,6 +655,14 @@ public final class BatteryUsageStats implements Parcelable {
             return this;
         }
 
+        private long getStatsDuration() {
+            if (mStatsDurationMs != -1) {
+                return mStatsDurationMs;
+            } else {
+                return mStatsEndTimestampMs - mStatsStartTimestampMs;
+            }
+        }
+
         /**
          * Sets the battery discharge amount since BatteryStats reset as percentage of the full
          * charge.
@@ -737,6 +742,22 @@ public final class BatteryUsageStats implements Parcelable {
             return builder;
         }
 
+        /**
+         * Creates or returns a UidBatteryConsumer, which represents battery attribution
+         * data for an individual UID. This version of the method is not suitable for use
+         * with PowerCalculators.
+         */
+        @NonNull
+        public UidBatteryConsumer.Builder getOrCreateUidBatteryConsumerBuilder(int uid) {
+            UidBatteryConsumer.Builder builder = mUidBatteryConsumerBuilders.get(uid);
+            if (builder == null) {
+                builder = new UidBatteryConsumer.Builder(mCustomPowerComponentNames,
+                        mIncludePowerModels, uid);
+                mUidBatteryConsumerBuilders.put(uid, builder);
+            }
+            return builder;
+        }
+
         /**
          * Creates or returns a UserBatteryConsumer, which represents battery attribution
          * data for an individual {@link UserHandle}.
@@ -756,5 +777,59 @@ public final class BatteryUsageStats implements Parcelable {
         public SparseArray<UidBatteryConsumer.Builder> getUidBatteryConsumerBuilders() {
             return mUidBatteryConsumerBuilders;
         }
+
+        /**
+         * Adds battery usage stats from another snapshots. The two snapshots are assumed to be
+         * non-overlapping, meaning that the power consumption estimates and session durations
+         * can be simply summed across the two snapshots.  This remains true even if the timestamps
+         * seem to indicate that the sessions are in fact overlapping: timestamps may be off as a
+         * result of realtime clock adjustments by the user or the system.
+         */
+        @NonNull
+        public Builder add(BatteryUsageStats stats) {
+            if (!Arrays.equals(mCustomPowerComponentNames, stats.mCustomPowerComponentNames)) {
+                throw new IllegalArgumentException(
+                        "BatteryUsageStats have different custom power components");
+            }
+
+            if (mUserBatteryConsumerBuilders.size() != 0
+                    || !stats.getUserBatteryConsumers().isEmpty()) {
+                throw new UnsupportedOperationException(
+                        "Combining UserBatteryConsumers is not supported");
+            }
+
+            mDischargedPowerLowerBoundMah += stats.mDischargedPowerLowerBound;
+            mDischargedPowerUpperBoundMah += stats.mDischargedPowerUpperBound;
+            mDischargePercentage += stats.mDischargePercentage;
+
+            mStatsDurationMs = getStatsDuration() + stats.getStatsDuration();
+
+            if (mStatsStartTimestampMs == 0
+                    || stats.mStatsStartTimestampMs < mStatsStartTimestampMs) {
+                mStatsStartTimestampMs = stats.mStatsStartTimestampMs;
+            }
+
+            final boolean addingLaterSnapshot = stats.mStatsEndTimestampMs > mStatsEndTimestampMs;
+            if (addingLaterSnapshot) {
+                mStatsEndTimestampMs = stats.mStatsEndTimestampMs;
+            }
+
+            for (int scope = 0; scope < AGGREGATE_BATTERY_CONSUMER_SCOPE_COUNT; scope++) {
+                getAggregateBatteryConsumerBuilder(scope)
+                        .add(stats.mAggregateBatteryConsumers[scope]);
+            }
+
+            for (UidBatteryConsumer consumer : stats.getUidBatteryConsumers()) {
+                getOrCreateUidBatteryConsumerBuilder(consumer.getUid()).add(consumer);
+            }
+
+            if (addingLaterSnapshot) {
+                mBatteryCapacityMah = stats.mBatteryCapacityMah;
+                mBatteryTimeRemainingMs = stats.mBatteryTimeRemainingMs;
+                mChargeTimeRemainingMs = stats.mChargeTimeRemainingMs;
+            }
+
+            return this;
+        }
     }
 }

core/java/android/os/PowerComponents.java

@@ -24,6 +24,7 @@ import android.util.proto.ProtoOutputStream;
 import com.android.internal.os.PowerCalculator;
 
 import java.io.PrintWriter;
+import java.util.Arrays;
 
 /**
  * Contains details of battery attribution data broken down to individual power drain types
@@ -36,9 +37,12 @@ class PowerComponents {
             - BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID;
 
     private final double mConsumedPowerMah;
+    @NonNull
     private final double[] mPowerComponentsMah;
+    @NonNull
     private final long[] mUsageDurationsMs;
     private final int mCustomPowerComponentCount;
+    @Nullable
     private final byte[] mPowerModels;
     // Not written to Parcel and must be explicitly restored during the parent object's unparceling
     private String[] mCustomPowerComponentNames;
@@ -49,7 +53,7 @@ class PowerComponents {
         mPowerComponentsMah = builder.mPowerComponentsMah;
         mUsageDurationsMs = builder.mUsageDurationsMs;
         mConsumedPowerMah = builder.getTotalPower();
-        mPowerModels = builder.mPowerModels;
+        mPowerModels = builder.getPowerModels();
     }
 
     PowerComponents(@NonNull Parcel source) {
@@ -146,9 +150,13 @@ class PowerComponents {
         }
     }
 
+    public boolean hasPowerModels() {
+        return mPowerModels != null;
+    }
+
     @BatteryConsumer.PowerModel
     int getPowerModel(@BatteryConsumer.PowerComponent int component) {
-        if (mPowerModels == null) {
+        if (!hasPowerModels()) {
             throw new IllegalStateException(
                     "Power model IDs were not requested in the BatteryUsageStatsQuery");
         }
@@ -298,6 +306,8 @@ class PowerComponents {
      * Builder for PowerComponents.
      */
     static final class Builder {
+        private static final byte POWER_MODEL_UNINITIALIZED = -1;
+
         private final double[] mPowerComponentsMah;
         private final String[] mCustomPowerComponentNames;
         private final long[] mUsageDurationsMs;
@@ -311,6 +321,7 @@ class PowerComponents {
             mUsageDurationsMs = new long[powerComponentCount];
             if (includePowerModels) {
                 mPowerModels = new byte[BatteryConsumer.POWER_COMPONENT_COUNT];
+                Arrays.fill(mPowerModels, POWER_MODEL_UNINITIALIZED);
             } else {
                 mPowerModels = null;
             }
@@ -412,12 +423,39 @@ class PowerComponents {
             return this;
         }
 
-        public void addPowerAndDuration(Builder other) {
+        public void addPowerAndDuration(PowerComponents.Builder other) {
+            addPowerAndDuration(other.mPowerComponentsMah, other.mUsageDurationsMs,
+                    other.mPowerModels);
+        }
+
+        public void addPowerAndDuration(PowerComponents other) {
+            addPowerAndDuration(other.mPowerComponentsMah, other.mUsageDurationsMs,
+                    other.mPowerModels);
+        }
+
+        private void addPowerAndDuration(double[] powerComponentsMah,
+                long[] usageDurationsMs, byte[] powerModels) {
+            if (mPowerComponentsMah.length != powerComponentsMah.length) {
+                throw new IllegalArgumentException(
+                        "Number of power components does not match: " + powerComponentsMah.length
+                                + ", expected: " + mPowerComponentsMah.length);
+            }
+
             for (int i = mPowerComponentsMah.length - 1; i >= 0; i--) {
-                mPowerComponentsMah[i] += other.mPowerComponentsMah[i];
+                mPowerComponentsMah[i] += powerComponentsMah[i];
             }
             for (int i = mUsageDurationsMs.length - 1; i >= 0; i--) {
-                mUsageDurationsMs[i] += other.mUsageDurationsMs[i];
+                mUsageDurationsMs[i] += usageDurationsMs[i];
+            }
+            if (mPowerModels != null && powerModels != null) {
+                for (int i = mPowerModels.length - 1; i >= 0; i--) {
+                    if (mPowerModels[i] == POWER_MODEL_UNINITIALIZED) {
+                        mPowerModels[i] = powerModels[i];
+                    } else if (mPowerModels[i] != powerModels[i]
+                            && powerModels[i] != POWER_MODEL_UNINITIALIZED) {
+                        mPowerModels[i] = BatteryConsumer.POWER_MODEL_UNDEFINED;
+                    }
+                }
             }
         }
 
@@ -433,6 +471,19 @@ class PowerComponents {
             return totalPowerMah;
         }
 
+        private byte[] getPowerModels() {
+            if (mPowerModels == null) {
+                return null;
+            }
+
+            byte[] powerModels = new byte[mPowerModels.length];
+            for (int i = mPowerModels.length - 1; i >= 0; i--) {
+                powerModels[i] = mPowerModels[i] != POWER_MODEL_UNINITIALIZED ? mPowerModels[i]
+                        : BatteryConsumer.POWER_MODEL_UNDEFINED;
+            }
+            return powerModels;
+        }
+
         /**
          * Creates a read-only object out of the Builder values.
          */

core/java/android/os/UidBatteryConsumer.java

@@ -19,6 +19,7 @@ package android.os;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.text.TextUtils;
 
 import com.android.internal.os.PowerCalculator;
 
@@ -147,9 +148,10 @@ public final class UidBatteryConsumer extends BatteryConsumer implements Parcela
      * Builder for UidBatteryConsumer.
      */
     public static final class Builder extends BaseBuilder<Builder> {
+        private static final String PACKAGE_NAME_UNINITIALIZED = "";
         private final BatteryStats.Uid mBatteryStatsUid;
         private final int mUid;
-        private String mPackageWithHighestDrain;
+        private String mPackageWithHighestDrain = PACKAGE_NAME_UNINITIALIZED;
         public long mTimeInForegroundMs;
         public long mTimeInBackgroundMs;
         private boolean mExcludeFromBatteryUsageStats;
@@ -161,8 +163,19 @@ public final class UidBatteryConsumer extends BatteryConsumer implements Parcela
             mUid = batteryStatsUid.getUid();
         }
 
+        public Builder(@NonNull String[] customPowerComponentNames, boolean includePowerModels,
+                int uid) {
+            super(customPowerComponentNames, includePowerModels);
+            mBatteryStatsUid = null;
+            mUid = uid;
+        }
+
         @NonNull
         public BatteryStats.Uid getBatteryStatsUid() {
+            if (mBatteryStatsUid == null) {
+                throw new IllegalStateException(
+                        "UidBatteryConsumer.Builder was initialized without a BatteryStats.Uid");
+            }
             return mBatteryStatsUid;
         }
 
@@ -176,7 +189,7 @@ public final class UidBatteryConsumer extends BatteryConsumer implements Parcela
          */
         @NonNull
         public Builder setPackageWithHighestDrain(@Nullable String packageName) {
-            mPackageWithHighestDrain = packageName;
+            mPackageWithHighestDrain = TextUtils.nullIfEmpty(packageName);
             return this;
         }
 
@@ -207,6 +220,30 @@ public final class UidBatteryConsumer extends BatteryConsumer implements Parcela
             return this;
         }
 
+        /**
+         * Adds power and usage duration from the supplied UidBatteryConsumer.
+         */
+        public Builder add(UidBatteryConsumer consumer) {
+            mPowerComponentsBuilder.addPowerAndDuration(consumer.mPowerComponents);
+            mTimeInBackgroundMs += consumer.mTimeInBackgroundMs;
+            mTimeInForegroundMs += consumer.mTimeInForegroundMs;
+
+            if (mPackageWithHighestDrain == PACKAGE_NAME_UNINITIALIZED) {
+                mPackageWithHighestDrain = consumer.mPackageWithHighestDrain;
+            } else if (!TextUtils.equals(mPackageWithHighestDrain,
+                    consumer.mPackageWithHighestDrain)) {
+                // Consider combining two UidBatteryConsumers with this distribution
+                // of power drain between packages:
+                // (package1=100, package2=10) and (package1=100, package2=101).
+                // Since we don't know the actual power distribution between packages at this
+                // point, we have no way to correctly declare package1 as the winner.
+                // The naive logic of picking the consumer with the higher total consumed
+                // power would produce an incorrect result.
+                mPackageWithHighestDrain = null;
+            }
+            return this;
+        }
+
         /**
          * Returns true if this UidBatteryConsumer must be excluded from the
          * BatteryUsageStats.
@@ -220,6 +257,9 @@ public final class UidBatteryConsumer extends BatteryConsumer implements Parcela
          */
         @NonNull
         public UidBatteryConsumer build() {
+            if (mPackageWithHighestDrain == PACKAGE_NAME_UNINITIALIZED) {
+                mPackageWithHighestDrain = null;
+            }
             return new UidBatteryConsumer(this);
         }
     }

core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsTest.java

@@ -16,8 +16,13 @@
 
 package com.android.internal.os;
 
+import static android.os.BatteryConsumer.POWER_MODEL_MEASURED_ENERGY;
+import static android.os.BatteryConsumer.POWER_MODEL_POWER_PROFILE;
+import static android.os.BatteryConsumer.POWER_MODEL_UNDEFINED;
+
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.fail;
 
 import android.os.BatteryConsumer;
@@ -42,15 +47,19 @@ import java.util.Set;
 @RunWith(AndroidJUnit4.class)
 public class BatteryUsageStatsTest {
 
+    private static final int USER_ID = 42;
+    private static final int APP_UID1 = 271;
+    private static final int APP_UID2 = 314;
+
     @Test
     public void testBuilder() {
-        BatteryUsageStats batteryUsageStats = buildBatteryUsageStats().build();
+        BatteryUsageStats batteryUsageStats = buildBatteryUsageStats1(true).build();
-        validateBatteryUsageStats(batteryUsageStats);
+        assertBatteryUsageStats1(batteryUsageStats, true);
     }
 
     @Test
     public void testParcelability() {
-        final BatteryUsageStats outBatteryUsageStats = buildBatteryUsageStats().build();
+        final BatteryUsageStats outBatteryUsageStats = buildBatteryUsageStats1(true).build();
         final Parcel outParcel = Parcel.obtain();
         outParcel.writeParcelable(outBatteryUsageStats, 0);
         final byte[] bytes = outParcel.marshall();
@@ -62,20 +71,20 @@ public class BatteryUsageStatsTest {
         final BatteryUsageStats inBatteryUsageStats =
                 inParcel.readParcelable(getClass().getClassLoader());
         assertThat(inBatteryUsageStats).isNotNull();
-        validateBatteryUsageStats(inBatteryUsageStats);
+        assertBatteryUsageStats1(inBatteryUsageStats, true);
     }
 
 
     @Test
     public void testDefaultSessionDuration() {
         final BatteryUsageStats stats =
-                buildBatteryUsageStats().setStatsDuration(10000).build();
+                buildBatteryUsageStats1(true).setStatsDuration(10000).build();
         assertThat(stats.getStatsDuration()).isEqualTo(10000);
     }
 
     @Test
     public void testDump() {
-        final BatteryUsageStats stats = buildBatteryUsageStats().build();
+        final BatteryUsageStats stats = buildBatteryUsageStats1(true).build();
         final StringWriter out = new StringWriter();
         try (PrintWriter pw = new PrintWriter(out)) {
             stats.dump(pw, "  ");
@@ -87,7 +96,7 @@ public class BatteryUsageStatsTest {
         assertThat(dump).contains("actual drain: 1000-2000");
         assertThat(dump).contains("cpu: 20100 apps: 10100 duration: 20s 300ms");
         assertThat(dump).contains("FOO: 20200 apps: 10200 duration: 20s 400ms");
-        assertThat(dump).contains("UID 2000: 1200 ( screen=300 cpu=400 FOO=500 )");
+        assertThat(dump).contains("UID 271: 1200 ( screen=300 cpu=400 FOO=500 )");
         assertThat(dump).contains("User 42: 30.0 ( cpu=10.0 FOO=20.0 )");
     }
 
@@ -101,154 +110,279 @@ public class BatteryUsageStatsTest {
         assertThat(allNames).hasSize(BatteryConsumer.POWER_COMPONENT_COUNT);
     }
 
-    private BatteryUsageStats.Builder buildBatteryUsageStats() {
+    @Test
-        final MockClocks clocks = new MockClocks();
+    public void testAdd() {
-        final MockBatteryStatsImpl batteryStats = new MockBatteryStatsImpl(clocks);
+        final BatteryUsageStats stats1 = buildBatteryUsageStats1(false).build();
-        final BatteryStatsImpl.Uid batteryStatsUid = batteryStats.getUidStatsLocked(2000);
+        final BatteryUsageStats stats2 = buildBatteryUsageStats2(new String[] {"FOO"}).build();
 
-        final BatteryUsageStats.Builder builder =
+        final BatteryUsageStats sum =
-                new BatteryUsageStats.Builder(new String[]{"FOO"}, true)
+                new BatteryUsageStats.Builder(new String[] {"FOO"}, true)
-                        .setBatteryCapacity(4000)
+                        .add(stats1)
-                        .setDischargePercentage(20)
+                        .add(stats2)
-                        .setDischargedPowerRange(1000, 2000)
+                        .build();
-                        .setStatsStartTimestamp(1000)
-                        .setStatsEndTimestamp(3000);
-        builder.getOrCreateUidBatteryConsumerBuilder(batteryStatsUid)
-                .setPackageWithHighestDrain("foo")
-                .setTimeInStateMs(UidBatteryConsumer.STATE_FOREGROUND, 1000)
-                .setTimeInStateMs(UidBatteryConsumer.STATE_BACKGROUND, 2000)
-                .setConsumedPower(
-                        BatteryConsumer.POWER_COMPONENT_SCREEN, 300)
-                .setConsumedPower(
-                        BatteryConsumer.POWER_COMPONENT_CPU, 400)
-                .setConsumedPowerForCustomComponent(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 500)
-                .setUsageDurationMillis(
-                        BatteryConsumer.POWER_COMPONENT_CPU, 600)
-                .setUsageDurationForCustomComponentMillis(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 800);
 
-        builder.getAggregateBatteryConsumerBuilder(
+        assertBatteryUsageStats(sum, 42345, 50, 2234, 4345, 1000, 5000, 5000);
-                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_ALL_APPS)
-                .setConsumedPower(
-                        BatteryConsumer.POWER_COMPONENT_CPU, 10100)
-                .setConsumedPowerForCustomComponent(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 10200)
-                .setUsageDurationMillis(
-                        BatteryConsumer.POWER_COMPONENT_CPU, 10300)
-                .setUsageDurationForCustomComponentMillis(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 10400);
-
-        builder.getAggregateBatteryConsumerBuilder(
-                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE)
-                .setConsumedPower(30000)
-                .setConsumedPower(
-                        BatteryConsumer.POWER_COMPONENT_CPU, 20100)
-                .setConsumedPowerForCustomComponent(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 20200)
-                .setUsageDurationMillis(
-                        BatteryConsumer.POWER_COMPONENT_CPU, 20300)
-                .setUsageDurationForCustomComponentMillis(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 20400);
-
-        builder.getOrCreateUserBatteryConsumerBuilder(42)
-                .setConsumedPower(
-                        BatteryConsumer.POWER_COMPONENT_CPU, 10)
-                .setConsumedPowerForCustomComponent(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 20)
-                .setUsageDurationMillis(
-                        BatteryConsumer.POWER_COMPONENT_CPU, 30)
-                .setUsageDurationForCustomComponentMillis(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 40);
-
-        return builder;
-    }
-
-    public void validateBatteryUsageStats(BatteryUsageStats batteryUsageStats) {
-        assertThat(batteryUsageStats.getConsumedPower()).isEqualTo(30000);
-        assertThat(batteryUsageStats.getBatteryCapacity()).isEqualTo(4000);
-        assertThat(batteryUsageStats.getDischargePercentage()).isEqualTo(20);
-        assertThat(batteryUsageStats.getDischargedPowerRange().getLower()).isEqualTo(1000);
-        assertThat(batteryUsageStats.getDischargedPowerRange().getUpper()).isEqualTo(2000);
-        assertThat(batteryUsageStats.getStatsStartTimestamp()).isEqualTo(1000);
-        assertThat(batteryUsageStats.getStatsEndTimestamp()).isEqualTo(3000);
-        assertThat(batteryUsageStats.getStatsDuration()).isEqualTo(2000);
 
         final List<UidBatteryConsumer> uidBatteryConsumers =
-                batteryUsageStats.getUidBatteryConsumers();
+                sum.getUidBatteryConsumers();
         for (UidBatteryConsumer uidBatteryConsumer : uidBatteryConsumers) {
-            if (uidBatteryConsumer.getUid() == 2000) {
+            if (uidBatteryConsumer.getUid() == APP_UID1) {
-                assertThat(uidBatteryConsumer.getPackageWithHighestDrain()).isEqualTo("foo");
+                assertUidBatteryConsumer(uidBatteryConsumer, 2124, null,
-                assertThat(uidBatteryConsumer.getTimeInStateMs(
+                        5321, 7432, 423, POWER_MODEL_POWER_PROFILE, 745, POWER_MODEL_UNDEFINED,
-                        UidBatteryConsumer.STATE_FOREGROUND)).isEqualTo(1000);
+                        956, 1167, 1478);
-                assertThat(uidBatteryConsumer.getTimeInStateMs(
+            } else if (uidBatteryConsumer.getUid() == APP_UID2) {
-                        UidBatteryConsumer.STATE_BACKGROUND)).isEqualTo(2000);
+                assertUidBatteryConsumer(uidBatteryConsumer, 1332, "bar",
-                assertThat(uidBatteryConsumer.getConsumedPower(
+                        1111, 2222, 333, POWER_MODEL_POWER_PROFILE, 444, POWER_MODEL_POWER_PROFILE,
-                        BatteryConsumer.POWER_COMPONENT_SCREEN)).isEqualTo(300);
+                        555, 666, 777);
-                assertThat(uidBatteryConsumer.getConsumedPower(
-                        BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(400);
-                assertThat(uidBatteryConsumer.getConsumedPowerForCustomComponent(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(500);
-                assertThat(uidBatteryConsumer.getUsageDurationMillis(
-                        BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(600);
-                assertThat(uidBatteryConsumer.getUsageDurationForCustomComponentMillis(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(800);
-                assertThat(uidBatteryConsumer.getConsumedPower()).isEqualTo(1200);
-                assertThat(uidBatteryConsumer.getCustomPowerComponentCount()).isEqualTo(1);
-                assertThat(uidBatteryConsumer.getCustomPowerComponentName(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo("FOO");
             } else {
                 fail("Unexpected UID " + uidBatteryConsumer.getUid());
             }
         }
 
-        final BatteryConsumer appsBatteryConsumer = batteryUsageStats.getAggregateBatteryConsumer(
+        assertAggregateBatteryConsumer(sum,
-                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_ALL_APPS);
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_ALL_APPS,
-        assertThat(appsBatteryConsumer.getConsumedPower(
+                20223, 20434, 20645, 20856);
-                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(10100);
-        assertThat(appsBatteryConsumer.getConsumedPowerForCustomComponent(
-                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(10200);
-        assertThat(appsBatteryConsumer.getUsageDurationMillis(
-                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(10300);
-        assertThat(appsBatteryConsumer.getUsageDurationForCustomComponentMillis(
-                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(10400);
-        assertThat(appsBatteryConsumer.getCustomPowerComponentCount()).isEqualTo(1);
-        assertThat(appsBatteryConsumer.getCustomPowerComponentName(
-                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo("FOO");
 
-        final BatteryConsumer deviceBatteryConsumer = batteryUsageStats.getAggregateBatteryConsumer(
+        assertAggregateBatteryConsumer(sum,
-                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE);
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE,
-        assertThat(deviceBatteryConsumer.getConsumedPower(
+                40211, 40422, 40633, 40844);
-                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(20100);
+    }
-        assertThat(deviceBatteryConsumer.getConsumedPowerForCustomComponent(
+
-                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(20200);
+    @Test
-        assertThat(deviceBatteryConsumer.getUsageDurationMillis(
+    public void testAdd_customComponentMismatch() {
-                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(20300);
+        final BatteryUsageStats.Builder builder =
-        assertThat(deviceBatteryConsumer.getUsageDurationForCustomComponentMillis(
+                new BatteryUsageStats.Builder(new String[] {"FOO"}, true);
-                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(20400);
+        final BatteryUsageStats stats = buildBatteryUsageStats2(new String[] {"BAR"}).build();
-        assertThat(deviceBatteryConsumer.getCustomPowerComponentCount()).isEqualTo(1);
+
-        assertThat(deviceBatteryConsumer.getCustomPowerComponentName(
+        assertThrows(IllegalArgumentException.class, () -> builder.add(stats));
-                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo("FOO");
+    }
+
+    private BatteryUsageStats.Builder buildBatteryUsageStats1(boolean includeUserBatteryConsumer) {
+        final MockClocks clocks = new MockClocks();
+        final MockBatteryStatsImpl batteryStats = new MockBatteryStatsImpl(clocks);
+
+        final BatteryUsageStats.Builder builder =
+                new BatteryUsageStats.Builder(new String[] {"FOO"}, true)
+                        .setBatteryCapacity(4000)
+                        .setDischargePercentage(20)
+                        .setDischargedPowerRange(1000, 2000)
+                        .setStatsStartTimestamp(1000)
+                        .setStatsEndTimestamp(3000);
+
+        addUidBatteryConsumer(builder, batteryStats, APP_UID1, "foo",
+                1000, 2000,
+                300, POWER_MODEL_POWER_PROFILE, 400, POWER_MODEL_POWER_PROFILE, 500, 600, 800);
+
+        addAggregateBatteryConsumer(builder,
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_ALL_APPS, 0,
+                10100, 10200, 10300, 10400);
+
+        addAggregateBatteryConsumer(builder,
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE, 30000,
+                20100, 20200, 20300, 20400);
+
+
+        if (includeUserBatteryConsumer) {
+            builder.getOrCreateUserBatteryConsumerBuilder(USER_ID)
+                    .setConsumedPower(
+                            BatteryConsumer.POWER_COMPONENT_CPU, 10)
+                    .setConsumedPowerForCustomComponent(
+                            BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 20)
+                    .setUsageDurationMillis(
+                            BatteryConsumer.POWER_COMPONENT_CPU, 30)
+                    .setUsageDurationForCustomComponentMillis(
+                            BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 40);
+        }
+        return builder;
+    }
+
+    private BatteryUsageStats.Builder buildBatteryUsageStats2(String[] customPowerComponentNames) {
+        final MockClocks clocks = new MockClocks();
+        final MockBatteryStatsImpl batteryStats = new MockBatteryStatsImpl(clocks);
+
+        final BatteryUsageStats.Builder builder =
+                new BatteryUsageStats.Builder(customPowerComponentNames, true)
+                        .setDischargePercentage(30)
+                        .setDischargedPowerRange(1234, 2345)
+                        .setStatsStartTimestamp(2000)
+                        .setStatsEndTimestamp(5000);
+
+        addUidBatteryConsumer(builder, batteryStats, APP_UID1, null,
+                4321, 5432,
+                123, POWER_MODEL_POWER_PROFILE, 345, POWER_MODEL_MEASURED_ENERGY, 456, 567, 678);
+
+        addUidBatteryConsumer(builder, batteryStats, APP_UID2, "bar",
+                1111, 2222,
+                333, POWER_MODEL_POWER_PROFILE, 444, POWER_MODEL_POWER_PROFILE, 555, 666, 777);
+
+        addAggregateBatteryConsumer(builder,
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_ALL_APPS, 0,
+                10123, 10234, 10345, 10456);
+
+        addAggregateBatteryConsumer(builder,
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE, 12345,
+                20111, 20222, 20333, 20444);
+
+        return builder;
+    }
+
+    private void addUidBatteryConsumer(BatteryUsageStats.Builder builder,
+            MockBatteryStatsImpl batteryStats, int uid, String packageWithHighestDrain,
+            int timeInStateForeground, int timeInStateBackground, int screenPower,
+            int screenPowerModel, int cpuPower, int cpuPowerModel, int customComponentPower,
+            int cpuDuration, int customComponentDuration) {
+        final BatteryStatsImpl.Uid batteryStatsUid = batteryStats.getUidStatsLocked(uid);
+        builder.getOrCreateUidBatteryConsumerBuilder(batteryStatsUid)
+                .setPackageWithHighestDrain(packageWithHighestDrain)
+                .setTimeInStateMs(UidBatteryConsumer.STATE_FOREGROUND, timeInStateForeground)
+                .setTimeInStateMs(UidBatteryConsumer.STATE_BACKGROUND, timeInStateBackground)
+                .setConsumedPower(
+                        BatteryConsumer.POWER_COMPONENT_SCREEN, screenPower, screenPowerModel)
+                .setConsumedPower(
+                        BatteryConsumer.POWER_COMPONENT_CPU, cpuPower, cpuPowerModel)
+                .setConsumedPowerForCustomComponent(
+                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, customComponentPower)
+                .setUsageDurationMillis(
+                        BatteryConsumer.POWER_COMPONENT_CPU, cpuDuration)
+                .setUsageDurationForCustomComponentMillis(
+                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, customComponentDuration);
+    }
+
+    private void addAggregateBatteryConsumer(BatteryUsageStats.Builder builder, int scope,
+            double consumedPower, int cpuPower, int customComponentPower, int cpuDuration,
+            int customComponentDuration) {
+        builder.getAggregateBatteryConsumerBuilder(scope)
+                .setConsumedPower(consumedPower)
+                .setConsumedPower(
+                        BatteryConsumer.POWER_COMPONENT_CPU, cpuPower)
+                .setConsumedPowerForCustomComponent(
+                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, customComponentPower)
+                .setUsageDurationMillis(
+                        BatteryConsumer.POWER_COMPONENT_CPU, cpuDuration)
+                .setUsageDurationForCustomComponentMillis(
+                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, customComponentDuration);
+    }
+
+    public void assertBatteryUsageStats1(BatteryUsageStats batteryUsageStats,
+            boolean includesUserBatteryConsumers) {
+        assertBatteryUsageStats(batteryUsageStats, 30000, 20, 1000, 2000, 1000, 3000, 2000);
+
+        final List<UidBatteryConsumer> uidBatteryConsumers =
+                batteryUsageStats.getUidBatteryConsumers();
+        assertThat(uidBatteryConsumers).hasSize(1);
+        for (UidBatteryConsumer uidBatteryConsumer : uidBatteryConsumers) {
+            if (uidBatteryConsumer.getUid() == APP_UID1) {
+                assertUidBatteryConsumer(uidBatteryConsumer, 1200, "foo",
+                        1000, 2000, 300, POWER_MODEL_POWER_PROFILE, 400, POWER_MODEL_POWER_PROFILE,
+                        500, 600, 800);
+            } else {
+                fail("Unexpected UID " + uidBatteryConsumer.getUid());
+            }
+        }
 
         final List<UserBatteryConsumer> userBatteryConsumers =
                 batteryUsageStats.getUserBatteryConsumers();
-        for (UserBatteryConsumer userBatteryConsumer : userBatteryConsumers) {
+        if (includesUserBatteryConsumers) {
-            if (userBatteryConsumer.getUserId() == 42) {
+            assertThat(userBatteryConsumers).hasSize(1);
-                assertThat(userBatteryConsumer.getConsumedPower(
+            for (UserBatteryConsumer userBatteryConsumer : userBatteryConsumers) {
-                        BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(10);
+                if (userBatteryConsumer.getUserId() == USER_ID) {
-                assertThat(userBatteryConsumer.getConsumedPowerForCustomComponent(
+                    assertUserBatteryConsumer(userBatteryConsumer, 42, 10, 20, 30, 40);
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(20);
+                } else {
-                assertThat(userBatteryConsumer.getUsageDurationMillis(
+                    fail("Unexpected User ID " + userBatteryConsumer.getUserId());
-                        BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(30);
+                }
-                assertThat(userBatteryConsumer.getUsageDurationForCustomComponentMillis(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(40);
-                assertThat(userBatteryConsumer.getConsumedPower()).isEqualTo(30);
-                assertThat(userBatteryConsumer.getCustomPowerComponentCount()).isEqualTo(1);
-                assertThat(userBatteryConsumer.getCustomPowerComponentName(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo("FOO");
-            } else {
-                fail("Unexpected user ID " + userBatteryConsumer.getUserId());
             }
+        } else {
+            assertThat(userBatteryConsumers).isEmpty();
         }
+
+        assertAggregateBatteryConsumer(batteryUsageStats,
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_ALL_APPS,
+                10100, 10200, 10300, 10400);
+
+        assertAggregateBatteryConsumer(batteryUsageStats,
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE,
+                20100, 20200, 20300, 20400);
+    }
+
+    private void assertBatteryUsageStats(BatteryUsageStats batteryUsageStats, int consumedPower,
+            int dischargePercentage, int dischagePowerLower, int dischargePowerUpper,
+            int statsStartTimestamp, int statsEndTimestamp, int statsDuration) {
+        assertThat(batteryUsageStats.getConsumedPower()).isEqualTo(consumedPower);
+        assertThat(batteryUsageStats.getDischargePercentage()).isEqualTo(dischargePercentage);
+        assertThat(batteryUsageStats.getDischargedPowerRange().getLower()).isEqualTo(
+                dischagePowerLower);
+        assertThat(batteryUsageStats.getDischargedPowerRange().getUpper()).isEqualTo(
+                dischargePowerUpper);
+        assertThat(batteryUsageStats.getStatsStartTimestamp()).isEqualTo(statsStartTimestamp);
+        assertThat(batteryUsageStats.getStatsEndTimestamp()).isEqualTo(statsEndTimestamp);
+        assertThat(batteryUsageStats.getStatsDuration()).isEqualTo(statsDuration);
+    }
+
+    private void assertUidBatteryConsumer(UidBatteryConsumer uidBatteryConsumer,
+            int consumedPower, String packageWithHighestDrain, int timeInStateForeground,
+            int timeInStateBackground, int screenPower, int screenPowerModel, int cpuPower,
+            int cpuPowerModel, int customComponentPower, int cpuDuration,
+            int customComponentDuration) {
+        assertThat(uidBatteryConsumer.getConsumedPower()).isEqualTo(consumedPower);
+        assertThat(uidBatteryConsumer.getPackageWithHighestDrain()).isEqualTo(
+                packageWithHighestDrain);
+        assertThat(uidBatteryConsumer.getTimeInStateMs(
+                UidBatteryConsumer.STATE_FOREGROUND)).isEqualTo(timeInStateForeground);
+        assertThat(uidBatteryConsumer.getTimeInStateMs(
+                UidBatteryConsumer.STATE_BACKGROUND)).isEqualTo(timeInStateBackground);
+        assertThat(uidBatteryConsumer.getConsumedPower(
+                BatteryConsumer.POWER_COMPONENT_SCREEN)).isEqualTo(screenPower);
+        assertThat(uidBatteryConsumer.getPowerModel(
+                BatteryConsumer.POWER_COMPONENT_SCREEN)).isEqualTo(screenPowerModel);
+        assertThat(uidBatteryConsumer.getConsumedPower(
+                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(cpuPower);
+        assertThat(uidBatteryConsumer.getPowerModel(
+                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(cpuPowerModel);
+        assertThat(uidBatteryConsumer.getConsumedPowerForCustomComponent(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(customComponentPower);
+        assertThat(uidBatteryConsumer.getUsageDurationMillis(
+                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(cpuDuration);
+        assertThat(uidBatteryConsumer.getUsageDurationForCustomComponentMillis(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(
+                customComponentDuration);
+        assertThat(uidBatteryConsumer.getCustomPowerComponentCount()).isEqualTo(1);
+        assertThat(uidBatteryConsumer.getCustomPowerComponentName(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo("FOO");
+    }
+
+    private void assertUserBatteryConsumer(UserBatteryConsumer userBatteryConsumer,
+            int userId, int cpuPower, int customComponentPower,
+            int cpuDuration, int customComponentDuration) {
+        assertThat(userBatteryConsumer.getConsumedPower(
+                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(cpuPower);
+        assertThat(userBatteryConsumer.getConsumedPowerForCustomComponent(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(customComponentPower);
+        assertThat(userBatteryConsumer.getUsageDurationMillis(
+                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(cpuDuration);
+        assertThat(userBatteryConsumer.getUsageDurationForCustomComponentMillis(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(
+                customComponentDuration);
+        assertThat(userBatteryConsumer.getCustomPowerComponentCount()).isEqualTo(1);
+        assertThat(userBatteryConsumer.getCustomPowerComponentName(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo("FOO");
+    }
+
+    private void assertAggregateBatteryConsumer(BatteryUsageStats batteryUsageStats,
+            int aggregateBatteryConsumerScopeAllApps, int cpuPower, int customComponentPower,
+            int cpuDuration, int customComponentDuration) {
+        final BatteryConsumer appsBatteryConsumer = batteryUsageStats.getAggregateBatteryConsumer(
+                aggregateBatteryConsumerScopeAllApps);
+        assertThat(appsBatteryConsumer.getConsumedPower(
+                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(cpuPower);
+        assertThat(appsBatteryConsumer.getConsumedPowerForCustomComponent(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(customComponentPower);
+        assertThat(appsBatteryConsumer.getUsageDurationMillis(
+                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(cpuDuration);
+        assertThat(appsBatteryConsumer.getUsageDurationForCustomComponentMillis(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(
+                customComponentDuration);
+        assertThat(appsBatteryConsumer.getCustomPowerComponentCount()).isEqualTo(1);
+        assertThat(appsBatteryConsumer.getCustomPowerComponentName(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo("FOO");
     }
 }