Remove POWER_COMPONENT_USAGE and TIME_COMPONENT_USAGE

... and use specific power components, such as POWER_COMPONENT_SCREEN

Bug: 184397335
Test: atest FrameworksCoreTests:com.android.internal.os.BatteryStatsTests
Test: mp :BatteryStatsViewer && adb shell am start -n com.android.frameworks.core.batterystatsviewer/.BatteryStatsViewerActivity
Change-Id: I551b2d37d2da77e21e622961f5290b55fb24da61

14 files changed, 72 insertions, 51 deletions

diff --git a/core/java/android/os/BatteryConsumer.java b/core/java/android/os/BatteryConsumer.java
index b558edcded47..ba63ba4521df 100644
--- a/core/java/android/os/BatteryConsumer.java
+++ b/core/java/android/os/BatteryConsumer.java
@@ -36,7 +36,7 @@ public abstract class BatteryConsumer {
      * @hide
      */
     @IntDef(prefix = {"POWER_COMPONENT_"}, value = {
-            POWER_COMPONENT_USAGE,
+            POWER_COMPONENT_SCREEN,
             POWER_COMPONENT_CPU,
             POWER_COMPONENT_BLUETOOTH,
             POWER_COMPONENT_CAMERA,
@@ -49,14 +49,16 @@ public abstract class BatteryConsumer {
             POWER_COMPONENT_GNSS,
             POWER_COMPONENT_WIFI,
             POWER_COMPONENT_WAKELOCK,
-            POWER_COMPONENT_SCREEN,
+            POWER_COMPONENT_MEMORY,
+            POWER_COMPONENT_PHONE,
+            POWER_COMPONENT_IDLE,
             POWER_COMPONENT_REATTRIBUTED_TO_OTHER_CONSUMERS,
     })
     @Retention(RetentionPolicy.SOURCE)
     public static @interface PowerComponent {
     }
 
-    public static final int POWER_COMPONENT_USAGE = 0;
+    public static final int POWER_COMPONENT_SCREEN = 0;
     public static final int POWER_COMPONENT_CPU = 1;
     public static final int POWER_COMPONENT_BLUETOOTH = 2;
     public static final int POWER_COMPONENT_CAMERA = 3;
@@ -69,13 +71,15 @@ public abstract class BatteryConsumer {
     public static final int POWER_COMPONENT_GNSS = 10;
     public static final int POWER_COMPONENT_WIFI = 11;
     public static final int POWER_COMPONENT_WAKELOCK = 12;
-    public static final int POWER_COMPONENT_SCREEN = 13;
+    public static final int POWER_COMPONENT_MEMORY = 13;
+    public static final int POWER_COMPONENT_PHONE = 13;
+    public static final int POWER_COMPONENT_IDLE = 15;
     // Power that is re-attributed to other battery consumers. For example, for System Server
     // this represents the power attributed to apps requesting system services.
     // The value should be negative or zero.
-    public static final int POWER_COMPONENT_REATTRIBUTED_TO_OTHER_CONSUMERS = 14;
+    public static final int POWER_COMPONENT_REATTRIBUTED_TO_OTHER_CONSUMERS = 16;
 
-    public static final int POWER_COMPONENT_COUNT = 15;
+    public static final int POWER_COMPONENT_COUNT = 17;
 
     public static final int FIRST_CUSTOM_POWER_COMPONENT_ID = 1000;
     public static final int LAST_CUSTOM_POWER_COMPONENT_ID = 9999;
@@ -87,7 +91,7 @@ public abstract class BatteryConsumer {
      * @hide
      */
     @IntDef(prefix = {"TIME_COMPONENT_"}, value = {
-            TIME_COMPONENT_USAGE,
+            TIME_COMPONENT_SCREEN,
             TIME_COMPONENT_CPU,
             TIME_COMPONENT_CPU_FOREGROUND,
             TIME_COMPONENT_BLUETOOTH,
@@ -98,13 +102,15 @@ public abstract class BatteryConsumer {
             TIME_COMPONENT_GNSS,
             TIME_COMPONENT_WIFI,
             TIME_COMPONENT_WAKELOCK,
-            TIME_COMPONENT_SCREEN,
+            TIME_COMPONENT_MEMORY,
+            TIME_COMPONENT_PHONE,
+            TIME_COMPONENT_IDLE,
     })
     @Retention(RetentionPolicy.SOURCE)
     public static @interface TimeComponent {
     }
 
-    public static final int TIME_COMPONENT_USAGE = 0;
+    public static final int TIME_COMPONENT_SCREEN = 0;
     public static final int TIME_COMPONENT_CPU = 1;
     public static final int TIME_COMPONENT_CPU_FOREGROUND = 2;
     public static final int TIME_COMPONENT_BLUETOOTH = 3;
@@ -117,9 +123,11 @@ public abstract class BatteryConsumer {
     public static final int TIME_COMPONENT_GNSS = 10;
     public static final int TIME_COMPONENT_WIFI = 11;
     public static final int TIME_COMPONENT_WAKELOCK = 12;
-    public static final int TIME_COMPONENT_SCREEN = 13;
+    public static final int TIME_COMPONENT_MEMORY = 13;
+    public static final int TIME_COMPONENT_PHONE = 14;
+    public static final int TIME_COMPONENT_IDLE = 15;
 
-    public static final int TIME_COMPONENT_COUNT = 14;
+    public static final int TIME_COMPONENT_COUNT = 16;
 
     public static final int FIRST_CUSTOM_TIME_COMPONENT_ID = 1000;
     public static final int LAST_CUSTOM_TIME_COMPONENT_ID = 9999;
@@ -148,8 +156,7 @@ public abstract class BatteryConsumer {
      */
     public static final int POWER_MODEL_MEASURED_ENERGY = 1;
 
-    private final PowerComponents mPowerComponents;
-    private String[] mCustomPowerComponentNames;
+    protected final PowerComponents mPowerComponents;
 
     protected BatteryConsumer(@NonNull PowerComponents powerComponents) {
         mPowerComponents = powerComponents;
diff --git a/core/java/android/os/PowerComponents.java b/core/java/android/os/PowerComponents.java
index 2fea4acf3876..a0a41f4d6630 100644
--- a/core/java/android/os/PowerComponents.java
+++ b/core/java/android/os/PowerComponents.java
@@ -192,6 +192,19 @@ class PowerComponents {
     }
 
     /**
+     * Returns the largest usage duration among all time components.
+     */
+    public long getMaxComponentUsageDurationMillis() {
+        long max = 0;
+        for (int i = mTimeComponentsMs.length - 1; i >= 0; i--) {
+            if (mTimeComponentsMs[i] > max) {
+                max = mTimeComponentsMs[i];
+            }
+        }
+        return max;
+    }
+
+    /**
      * Builder for PowerComponents.
      */
     static final class Builder {
@@ -312,10 +325,10 @@ class PowerComponents {
         }
 
         public void addPowerAndDuration(Builder other) {
-            for (int i = 0; i < mPowerComponentsMah.length; i++) {
+            for (int i = mPowerComponentsMah.length - 1; i >= 0; i--) {
                 mPowerComponentsMah[i] += other.mPowerComponentsMah[i];
             }
-            for (int i = 0; i < mTimeComponentsMs.length; i++) {
+            for (int i = mTimeComponentsMs.length - 1; i >= 0; i--) {
                 mTimeComponentsMs[i] += other.mTimeComponentsMs[i];
             }
         }
diff --git a/core/java/android/os/SystemBatteryConsumer.java b/core/java/android/os/SystemBatteryConsumer.java
index 70b0532049fb..1327978ad0ba 100644
--- a/core/java/android/os/SystemBatteryConsumer.java
+++ b/core/java/android/os/SystemBatteryConsumer.java
@@ -104,6 +104,13 @@ public class SystemBatteryConsumer extends BatteryConsumer implements Parcelable
     }
 
     /**
+     * Returns the amount of time this consumer was operating.
+     */
+    public long getUsageDurationMillis() {
+        return mPowerComponents.getMaxComponentUsageDurationMillis();
+    }
+
+    /**
      * Writes the contents into a Parcel.
      */
     @Override
diff --git a/core/java/com/android/internal/os/AmbientDisplayPowerCalculator.java b/core/java/com/android/internal/os/AmbientDisplayPowerCalculator.java
index 2f4958223270..aa416c568dbc 100644
--- a/core/java/com/android/internal/os/AmbientDisplayPowerCalculator.java
+++ b/core/java/com/android/internal/os/AmbientDisplayPowerCalculator.java
@@ -52,8 +52,8 @@ public class AmbientDisplayPowerCalculator extends PowerCalculator {
                 measuredEnergyUC, mPowerEstimator, durationMs);
         builder.getOrCreateSystemBatteryConsumerBuilder(
                         SystemBatteryConsumer.DRAIN_TYPE_AMBIENT_DISPLAY)
-                .setConsumedPower(BatteryConsumer.POWER_COMPONENT_USAGE, powerMah, powerModel)
-                .setUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_USAGE, durationMs);
+                .setConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN, powerMah, powerModel)
+                .setUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_SCREEN, durationMs);
     }
 
     /**
diff --git a/core/java/com/android/internal/os/IdlePowerCalculator.java b/core/java/com/android/internal/os/IdlePowerCalculator.java
index 4a4991bd3966..5cb54bd7fbc9 100644
--- a/core/java/com/android/internal/os/IdlePowerCalculator.java
+++ b/core/java/com/android/internal/os/IdlePowerCalculator.java
@@ -54,8 +54,8 @@ public class IdlePowerCalculator extends PowerCalculator {
                 BatteryStats.STATS_SINCE_CHARGED);
         if (mPowerMah != 0) {
             builder.getOrCreateSystemBatteryConsumerBuilder(SystemBatteryConsumer.DRAIN_TYPE_IDLE)
-                    .setConsumedPower(BatteryConsumer.POWER_COMPONENT_USAGE, mPowerMah)
-                    .setUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_USAGE, mDurationMs);
+                    .setConsumedPower(BatteryConsumer.POWER_COMPONENT_IDLE, mPowerMah)
+                    .setUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_IDLE, mDurationMs);
         }
     }
 
diff --git a/core/java/com/android/internal/os/MemoryPowerCalculator.java b/core/java/com/android/internal/os/MemoryPowerCalculator.java
index 21dcce9fb64f..9ec40c6f6eed 100644
--- a/core/java/com/android/internal/os/MemoryPowerCalculator.java
+++ b/core/java/com/android/internal/os/MemoryPowerCalculator.java
@@ -32,8 +32,8 @@ public class MemoryPowerCalculator extends PowerCalculator {
         final double powerMah = calculatePower(batteryStats, rawRealtimeUs,
                 BatteryStats.STATS_SINCE_CHARGED);
         builder.getOrCreateSystemBatteryConsumerBuilder(SystemBatteryConsumer.DRAIN_TYPE_MEMORY)
-                .setUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_USAGE, durationMs)
-                .setConsumedPower(BatteryConsumer.POWER_COMPONENT_USAGE, powerMah);
+                .setUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_MEMORY, durationMs)
+                .setConsumedPower(BatteryConsumer.POWER_COMPONENT_MEMORY, powerMah);
     }
 
     @Override
diff --git a/core/java/com/android/internal/os/PhonePowerCalculator.java b/core/java/com/android/internal/os/PhonePowerCalculator.java
index 362ca0761b2c..6f279d99a5c9 100644
--- a/core/java/com/android/internal/os/PhonePowerCalculator.java
+++ b/core/java/com/android/internal/os/PhonePowerCalculator.java
@@ -45,8 +45,8 @@ public class PhonePowerCalculator extends PowerCalculator {
         final double phoneOnPower = mPowerEstimator.calculatePower(phoneOnTimeMs);
         if (phoneOnPower != 0) {
             builder.getOrCreateSystemBatteryConsumerBuilder(SystemBatteryConsumer.DRAIN_TYPE_PHONE)
-                    .setConsumedPower(BatteryConsumer.POWER_COMPONENT_USAGE, phoneOnPower)
-                    .setUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_USAGE, phoneOnTimeMs);
+                    .setConsumedPower(BatteryConsumer.POWER_COMPONENT_PHONE, phoneOnPower)
+                    .setUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_PHONE, phoneOnTimeMs);
         }
     }
 
diff --git a/core/java/com/android/internal/os/ScreenPowerCalculator.java b/core/java/com/android/internal/os/ScreenPowerCalculator.java
index 0267def918c2..0743c897a91f 100644
--- a/core/java/com/android/internal/os/ScreenPowerCalculator.java
+++ b/core/java/com/android/internal/os/ScreenPowerCalculator.java
@@ -96,9 +96,9 @@ public class ScreenPowerCalculator extends PowerCalculator {
         }
 
         builder.getOrCreateSystemBatteryConsumerBuilder(SystemBatteryConsumer.DRAIN_TYPE_SCREEN)
-                .setUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_USAGE,
+                .setUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_SCREEN,
                         totalPowerAndDuration.durationMs)
-                .setConsumedPower(BatteryConsumer.POWER_COMPONENT_USAGE,
+                .setConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN,
                         Math.max(totalPowerAndDuration.powerMah, totalAppPower), powerModel)
                 .setPowerConsumedByApps(totalAppPower);
     }
diff --git a/core/tests/batterystatstests/BatteryStatsViewer/src/com/android/frameworks/core/batterystatsviewer/BatteryConsumerData.java b/core/tests/batterystatstests/BatteryStatsViewer/src/com/android/frameworks/core/batterystatsviewer/BatteryConsumerData.java
index 78569bdb4adc..8df3de3db668 100644
--- a/core/tests/batterystatstests/BatteryStatsViewer/src/com/android/frameworks/core/batterystatsviewer/BatteryConsumerData.java
+++ b/core/tests/batterystatstests/BatteryStatsViewer/src/com/android/frameworks/core/batterystatsviewer/BatteryConsumerData.java
@@ -213,14 +213,6 @@ public class BatteryConsumerData {
         }
         addEntry("System services", EntryType.POWER,
                 requestedBatterySipper.systemServiceCpuPowerMah, totalSystemServiceCpuPowerMah);
-        if (requestedBatteryConsumer != null) {
-            addEntry("Usage", EntryType.POWER,
-                    requestedBatteryConsumer.getConsumedPower(
-                            BatteryConsumer.POWER_COMPONENT_USAGE), totalUsagePowerMah);
-        } else {
-            addEntry("Usage (sipper)", EntryType.POWER,
-                    requestedBatterySipper.usagePowerMah, totalUsagePowerMah);
-        }
         addEntry("Wake lock", EntryType.POWER,
                 requestedBatterySipper.wakeLockPowerMah, totalWakeLockPowerMah);
         addEntry("Mobile radio", EntryType.POWER,
diff --git a/core/tests/coretests/src/com/android/internal/os/AmbientDisplayPowerCalculatorTest.java b/core/tests/coretests/src/com/android/internal/os/AmbientDisplayPowerCalculatorTest.java
index cf47efddddd9..c63ec45ed5a9 100644
--- a/core/tests/coretests/src/com/android/internal/os/AmbientDisplayPowerCalculatorTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/AmbientDisplayPowerCalculatorTest.java
@@ -65,12 +65,12 @@ public class AmbientDisplayPowerCalculatorTest {
         SystemBatteryConsumer consumer =
                 mStatsRule.getSystemBatteryConsumer(
                         SystemBatteryConsumer.DRAIN_TYPE_AMBIENT_DISPLAY);
-        assertThat(consumer.getUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_USAGE))
+        assertThat(consumer.getUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_SCREEN))
                 .isEqualTo(90 * MINUTE_IN_MS);
         // 100,000,00 uC / 1000 (micro-/milli-) / 360 (seconds/hour) = 27.777778 mAh
-        assertThat(consumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_USAGE))
+        assertThat(consumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
                 .isWithin(PRECISION).of(27.777778);
-        assertThat(consumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_USAGE))
+        assertThat(consumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
                 .isEqualTo(BatteryConsumer.POWER_MODEL_MEASURED_ENERGY);
     }
 
@@ -91,11 +91,11 @@ public class AmbientDisplayPowerCalculatorTest {
         SystemBatteryConsumer consumer =
                 mStatsRule.getSystemBatteryConsumer(
                         SystemBatteryConsumer.DRAIN_TYPE_AMBIENT_DISPLAY);
-        assertThat(consumer.getUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_USAGE))
+        assertThat(consumer.getUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_SCREEN))
                 .isEqualTo(90 * MINUTE_IN_MS);
-        assertThat(consumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_USAGE))
+        assertThat(consumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
                 .isWithin(PRECISION).of(15.0);
-        assertThat(consumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_USAGE))
+        assertThat(consumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
                 .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
     }
 }
diff --git a/core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsTest.java b/core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsTest.java
index 0dd19776bbb3..b25359970c5a 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsTest.java
@@ -77,7 +77,7 @@ public class BatteryUsageStatsTest {
                 .setTimeInStateMs(UidBatteryConsumer.STATE_FOREGROUND, 1000)
                 .setTimeInStateMs(UidBatteryConsumer.STATE_BACKGROUND, 2000)
                 .setConsumedPower(
-                        BatteryConsumer.POWER_COMPONENT_USAGE, 300)
+                        BatteryConsumer.POWER_COMPONENT_SCREEN, 300)
                 .setConsumedPower(
                         BatteryConsumer.POWER_COMPONENT_CPU, 400)
                 .setConsumedPowerForCustomComponent(
@@ -123,7 +123,7 @@ public class BatteryUsageStatsTest {
                 assertThat(uidBatteryConsumer.getTimeInStateMs(
                         UidBatteryConsumer.STATE_BACKGROUND)).isEqualTo(2000);
                 assertThat(uidBatteryConsumer.getConsumedPower(
-                        BatteryConsumer.POWER_COMPONENT_USAGE)).isEqualTo(300);
+                        BatteryConsumer.POWER_COMPONENT_SCREEN)).isEqualTo(300);
                 assertThat(uidBatteryConsumer.getConsumedPower(
                         BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(400);
                 assertThat(uidBatteryConsumer.getConsumedPowerForCustomComponent(
@@ -157,6 +157,8 @@ public class BatteryUsageStatsTest {
                         BatteryConsumer.FIRST_CUSTOM_TIME_COMPONENT_ID)).isEqualTo(10400);
                 assertThat(systemBatteryConsumer.getConsumedPower()).isEqualTo(20300);
                 assertThat(systemBatteryConsumer.getPowerConsumedByApps()).isEqualTo(20000);
+                assertThat(systemBatteryConsumer.getUsageDurationMillis())
+                        .isEqualTo(10400); // max
                 assertThat(systemBatteryConsumer.getCustomPowerComponentCount()).isEqualTo(1);
                 assertThat(systemBatteryConsumer.getCustomPowerComponentName(
                         BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo("FOO");
diff --git a/core/tests/coretests/src/com/android/internal/os/IdlePowerCalculatorTest.java b/core/tests/coretests/src/com/android/internal/os/IdlePowerCalculatorTest.java
index 781e72560279..a9800b7e6b75 100644
--- a/core/tests/coretests/src/com/android/internal/os/IdlePowerCalculatorTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/IdlePowerCalculatorTest.java
@@ -48,9 +48,9 @@ public class IdlePowerCalculatorTest {
 
         SystemBatteryConsumer consumer =
                 mStatsRule.getSystemBatteryConsumer(SystemBatteryConsumer.DRAIN_TYPE_IDLE);
-        assertThat(consumer.getUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_USAGE))
+        assertThat(consumer.getUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_IDLE))
                 .isEqualTo(3000);
-        assertThat(consumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_USAGE))
+        assertThat(consumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_IDLE))
                 .isWithin(PRECISION).of(0.7);
     }
 }
diff --git a/core/tests/coretests/src/com/android/internal/os/MemoryPowerCalculatorTest.java b/core/tests/coretests/src/com/android/internal/os/MemoryPowerCalculatorTest.java
index 8f21503a6d77..71dbcdbf7f46 100644
--- a/core/tests/coretests/src/com/android/internal/os/MemoryPowerCalculatorTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/MemoryPowerCalculatorTest.java
@@ -55,9 +55,9 @@ public class MemoryPowerCalculatorTest {
 
         SystemBatteryConsumer consumer =
                 mStatsRule.getSystemBatteryConsumer(SystemBatteryConsumer.DRAIN_TYPE_MEMORY);
-        assertThat(consumer.getUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_USAGE))
+        assertThat(consumer.getUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_MEMORY))
                 .isEqualTo(3000);
-        assertThat(consumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_USAGE))
+        assertThat(consumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_MEMORY))
                 .isWithin(PRECISION).of(0.7);
     }
 }
diff --git a/core/tests/coretests/src/com/android/internal/os/ScreenPowerCalculatorTest.java b/core/tests/coretests/src/com/android/internal/os/ScreenPowerCalculatorTest.java
index 9cd6ea8a6a3b..7d829e44d9ec 100644
--- a/core/tests/coretests/src/com/android/internal/os/ScreenPowerCalculatorTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/ScreenPowerCalculatorTest.java
@@ -84,13 +84,13 @@ public class ScreenPowerCalculatorTest {
 
         SystemBatteryConsumer consumer =
                 mStatsRule.getSystemBatteryConsumer(SystemBatteryConsumer.DRAIN_TYPE_SCREEN);
-        assertThat(consumer.getUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_USAGE))
+        assertThat(consumer.getUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_SCREEN))
                 .isEqualTo(80 * MINUTE_IN_MS);
 
         // 600000000 uAs * (1 mA / 1000 uA) * (1 h / 3600 s)  = 166.66666 mAh
-        assertThat(consumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_USAGE))
+        assertThat(consumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
                 .isWithin(PRECISION).of(166.66666);
-        assertThat(consumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_USAGE))
+        assertThat(consumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
                 .isEqualTo(BatteryConsumer.POWER_MODEL_MEASURED_ENERGY);
         assertThat(consumer.getConsumedPower())
                 .isWithin(PRECISION).of(166.66666);
@@ -153,11 +153,11 @@ public class ScreenPowerCalculatorTest {
 
         SystemBatteryConsumer consumer =
                 mStatsRule.getSystemBatteryConsumer(SystemBatteryConsumer.DRAIN_TYPE_SCREEN);
-        assertThat(consumer.getUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_USAGE))
+        assertThat(consumer.getUsageDurationMillis(BatteryConsumer.TIME_COMPONENT_SCREEN))
                 .isEqualTo(80 * MINUTE_IN_MS);
-        assertThat(consumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_USAGE))
+        assertThat(consumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
                 .isWithin(PRECISION).of(92.0);
-        assertThat(consumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_USAGE))
+        assertThat(consumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
                 .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
         assertThat(consumer.getConsumedPower())
                 .isWithin(PRECISION).of(92.0);