Merge "Factor out common PowerStatsCollector / PowerStats HAL integration code" into main

26 files changed, 364 insertions, 524 deletions

diff --git a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
index 47ae9d0ff37c..680b1acedf78 100644
--- a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
+++ b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
@@ -179,7 +179,6 @@ import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.locks.ReentrantLock;
-import java.util.function.IntSupplier;
 import java.util.function.LongSupplier;
 import java.util.function.Supplier;
 
@@ -2027,7 +2026,7 @@ public class BatteryStatsImpl extends BatteryStats {
         void setContext(Context context) {
             mPackageManager = context.getPackageManager();
             mConsumedEnergyRetriever = new PowerStatsCollector.ConsumedEnergyRetrieverImpl(
-                    LocalServices.getService(PowerStatsInternal.class));
+                    LocalServices.getService(PowerStatsInternal.class), () -> mBatteryVoltageMv);
             mNetworkStatsManager = context.getSystemService(NetworkStatsManager.class);
             mTelephonyManager =
                     (TelephonyManager) context.getSystemService(Context.TELEPHONY_SERVICE);
@@ -2082,11 +2081,6 @@ public class BatteryStatsImpl extends BatteryStats {
         }
 
         @Override
-        public IntSupplier getVoltageSupplier() {
-            return () -> mBatteryVoltageMv;
-        }
-
-        @Override
         public ScreenPowerStatsCollector.ScreenUsageTimeRetriever getScreenUsageTimeRetriever() {
             return mScreenUsageTimeRetriever;
         }
diff --git a/services/core/java/com/android/server/power/stats/BluetoothPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/BluetoothPowerStatsCollector.java
index 622cc1630e69..539c41537f8e 100644
--- a/services/core/java/com/android/server/power/stats/BluetoothPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/BluetoothPowerStatsCollector.java
@@ -30,12 +30,10 @@ import com.android.internal.os.Clock;
 import com.android.internal.os.PowerStats;
 import com.android.server.power.stats.format.BluetoothPowerStatsLayout;
 
-import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
-import java.util.function.IntSupplier;
 
 public class BluetoothPowerStatsCollector extends PowerStatsCollector {
     private static final String TAG = "BluetoothPowerStatsCollector";
@@ -62,7 +60,6 @@ public class BluetoothPowerStatsCollector extends PowerStatsCollector {
         long getPowerStatsCollectionThrottlePeriod(String powerComponentName);
         PackageManager getPackageManager();
         ConsumedEnergyRetriever getConsumedEnergyRetriever();
-        IntSupplier getVoltageSupplier();
         BluetoothStatsRetriever getBluetoothStatsRetriever();
     }
 
@@ -74,10 +71,7 @@ public class BluetoothPowerStatsCollector extends PowerStatsCollector {
     private long[] mDeviceStats;
     private BluetoothStatsRetriever mBluetoothStatsRetriever;
     private ConsumedEnergyRetriever mConsumedEnergyRetriever;
-    private IntSupplier mVoltageSupplier;
-    private int[] mEnergyConsumerIds = new int[0];
-    private long[] mLastConsumedEnergyUws;
-    private int mLastVoltageMv;
+    private ConsumedEnergyHelper mConsumedEnergyHelper;
 
     private long mLastRxTime;
     private long mLastTxTime;
@@ -123,15 +117,11 @@ public class BluetoothPowerStatsCollector extends PowerStatsCollector {
             return false;
         }
 
-        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
-        mVoltageSupplier = mInjector.getVoltageSupplier();
         mBluetoothStatsRetriever = mInjector.getBluetoothStatsRetriever();
-        mEnergyConsumerIds =
-                mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.BLUETOOTH);
-        mLastConsumedEnergyUws = new long[mEnergyConsumerIds.length];
-        Arrays.fill(mLastConsumedEnergyUws, ENERGY_UNSPECIFIED);
-
-        mLayout = new BluetoothPowerStatsLayout(mEnergyConsumerIds.length);
+        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
+        mConsumedEnergyHelper = new ConsumedEnergyHelper(mConsumedEnergyRetriever,
+                EnergyConsumerType.BLUETOOTH);
+        mLayout = new BluetoothPowerStatsLayout(mConsumedEnergyHelper.getEnergyConsumerCount());
 
         PersistableBundle extras = new PersistableBundle();
         mLayout.toExtras(extras);
@@ -157,9 +147,7 @@ public class BluetoothPowerStatsCollector extends PowerStatsCollector {
         collectBluetoothActivityInfo();
         collectBluetoothScanStats();
 
-        if (mEnergyConsumerIds.length != 0) {
-            collectEnergyConsumers();
-        }
+        mConsumedEnergyHelper.collectConsumedEnergy(mPowerStats, mLayout);
 
         return mPowerStats;
     }
@@ -291,34 +279,6 @@ public class BluetoothPowerStatsCollector extends PowerStatsCollector {
         mLayout.setDeviceScanTime(mDeviceStats, totalScanTime);
     }
 
-    private void collectEnergyConsumers() {
-        int voltageMv = mVoltageSupplier.getAsInt();
-        if (voltageMv <= 0) {
-            Slog.wtf(TAG, "Unexpected battery voltage (" + voltageMv
-                    + " mV) when querying energy consumers");
-            return;
-        }
-
-        int averageVoltage = mLastVoltageMv != 0 ? (mLastVoltageMv + voltageMv) / 2 : voltageMv;
-        mLastVoltageMv = voltageMv;
-
-        long[] energyUws = mConsumedEnergyRetriever.getConsumedEnergyUws(mEnergyConsumerIds);
-        if (energyUws == null) {
-            return;
-        }
-
-        for (int i = energyUws.length - 1; i >= 0; i--) {
-            long energyDelta = mLastConsumedEnergyUws[i] != ENERGY_UNSPECIFIED
-                    ? energyUws[i] - mLastConsumedEnergyUws[i] : 0;
-            if (energyDelta < 0) {
-                // Likely, restart of powerstats HAL
-                energyDelta = 0;
-            }
-            mLayout.setConsumedEnergy(mPowerStats.stats, i, uJtoUc(energyDelta, averageVoltage));
-            mLastConsumedEnergyUws[i] = energyUws[i];
-        }
-    }
-
     @Override
     protected void onUidRemoved(int uid) {
         super.onUidRemoved(uid);
diff --git a/services/core/java/com/android/server/power/stats/CameraPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/CameraPowerStatsCollector.java
index 7ded031354ac..0f70064c1227 100644
--- a/services/core/java/com/android/server/power/stats/CameraPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/CameraPowerStatsCollector.java
@@ -26,7 +26,6 @@ public class CameraPowerStatsCollector extends EnergyConsumerPowerStatsCollector
     public CameraPowerStatsCollector(Injector injector) {
         super(injector, BatteryConsumer.POWER_COMPONENT_CAMERA,
                 BatteryConsumer.powerComponentIdToString(BatteryConsumer.POWER_COMPONENT_CAMERA),
-                EnergyConsumerType.CAMERA, /* energy consumer name */ null,
-                new BinaryStatePowerStatsLayout());
+                EnergyConsumerType.CAMERA, new BinaryStatePowerStatsLayout());
     }
 }
diff --git a/services/core/java/com/android/server/power/stats/CpuPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/CpuPowerStatsCollector.java
index 7b006f0af84a..128f14a31898 100644
--- a/services/core/java/com/android/server/power/stats/CpuPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/CpuPowerStatsCollector.java
@@ -34,9 +34,7 @@ import com.android.internal.os.PowerStats;
 import com.android.server.power.stats.format.CpuPowerStatsLayout;
 
 import java.io.PrintWriter;
-import java.util.Arrays;
 import java.util.Locale;
-import java.util.function.IntSupplier;
 
 /**
  * Collects snapshots of power-related system statistics.
@@ -47,7 +45,6 @@ import java.util.function.IntSupplier;
 public class CpuPowerStatsCollector extends PowerStatsCollector {
     private static final String TAG = "CpuPowerStatsCollector";
     private static final long NANOS_PER_MILLIS = 1000000;
-    private static final long ENERGY_UNSPECIFIED = -1;
     private static final int DEFAULT_CPU_POWER_BRACKETS = 3;
     private static final int DEFAULT_CPU_POWER_BRACKETS_PER_ENERGY_CONSUMER = 2;
 
@@ -59,7 +56,6 @@ public class CpuPowerStatsCollector extends PowerStatsCollector {
         PowerProfile getPowerProfile();
         KernelCpuStatsReader getKernelCpuStatsReader();
         ConsumedEnergyRetriever getConsumedEnergyRetriever();
-        IntSupplier getVoltageSupplier();
         long getPowerStatsCollectionThrottlePeriod(String powerComponentName);
 
         default int getDefaultCpuPowerBrackets() {
@@ -77,8 +73,7 @@ public class CpuPowerStatsCollector extends PowerStatsCollector {
     private CpuScalingPolicies mCpuScalingPolicies;
     private PowerProfile mPowerProfile;
     private KernelCpuStatsReader mKernelCpuStatsReader;
-    private ConsumedEnergyRetriever mConsumedEnergyRetriever;
-    private IntSupplier mVoltageSupplier;
+    private ConsumedEnergyHelper mConsumedEnergyHelper;
     private int mDefaultCpuPowerBrackets;
     private int mDefaultCpuPowerBracketsPerEnergyConsumer;
     private long[] mCpuTimeByScalingStep;
@@ -86,15 +81,12 @@ public class CpuPowerStatsCollector extends PowerStatsCollector {
     private long[] mTempUidStats;
     private final SparseArray<UidStats> mUidStats = new SparseArray<>();
     private boolean mIsPerUidTimeInStateSupported;
-    private int[] mCpuEnergyConsumerIds = new int[0];
     private PowerStats.Descriptor mPowerStatsDescriptor;
     // Reusable instance
     private PowerStats mCpuPowerStats;
     private CpuPowerStatsLayout mLayout;
     private long mLastUpdateTimestampNanos;
     private long mLastUpdateUptimeMillis;
-    private int mLastVoltageMv;
-    private long[] mLastConsumedEnergyUws;
 
     CpuPowerStatsCollector(Injector injector) {
         super(injector.getHandler(), injector.getPowerStatsCollectionThrottlePeriod(
@@ -116,24 +108,21 @@ public class CpuPowerStatsCollector extends PowerStatsCollector {
         mCpuScalingPolicies = mInjector.getCpuScalingPolicies();
         mPowerProfile = mInjector.getPowerProfile();
         mKernelCpuStatsReader = mInjector.getKernelCpuStatsReader();
-        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
-        mVoltageSupplier = mInjector.getVoltageSupplier();
         mDefaultCpuPowerBrackets = mInjector.getDefaultCpuPowerBrackets();
         mDefaultCpuPowerBracketsPerEnergyConsumer =
                 mInjector.getDefaultCpuPowerBracketsPerEnergyConsumer();
 
         mIsPerUidTimeInStateSupported = mKernelCpuStatsReader.isSupportedFeature();
-        mCpuEnergyConsumerIds =
-                mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.CPU_CLUSTER);
-        mLastConsumedEnergyUws = new long[mCpuEnergyConsumerIds.length];
-        Arrays.fill(mLastConsumedEnergyUws, ENERGY_UNSPECIFIED);
+
+        mConsumedEnergyHelper = new ConsumedEnergyHelper(mInjector.getConsumedEnergyRetriever(),
+                EnergyConsumerType.CPU_CLUSTER);
 
         int cpuScalingStepCount = mCpuScalingPolicies.getScalingStepCount();
         mCpuTimeByScalingStep = new long[cpuScalingStepCount];
         mTempCpuTimeByScalingStep = new long[cpuScalingStepCount];
         int[] scalingStepToPowerBracketMap = initPowerBrackets();
 
-        mLayout = new CpuPowerStatsLayout(mCpuEnergyConsumerIds.length,
+        mLayout = new CpuPowerStatsLayout(mConsumedEnergyHelper.getEnergyConsumerCount(),
                 mCpuScalingPolicies.getPolicies().length, scalingStepToPowerBracketMap);
 
         PersistableBundle extras = new PersistableBundle();
@@ -153,16 +142,18 @@ public class CpuPowerStatsCollector extends PowerStatsCollector {
     private int[] initPowerBrackets() {
         if (mPowerProfile.getCpuPowerBracketCount() != PowerProfile.POWER_BRACKETS_UNSPECIFIED) {
             return initPowerBracketsFromPowerProfile();
-        } else if (mCpuEnergyConsumerIds.length == 0 || mCpuEnergyConsumerIds.length == 1) {
+        } else if (mConsumedEnergyHelper.getEnergyConsumerCount() == 0
+                || mConsumedEnergyHelper.getEnergyConsumerCount() == 1) {
             return initDefaultPowerBrackets(mDefaultCpuPowerBrackets);
-        } else if (mCpuScalingPolicies.getPolicies().length == mCpuEnergyConsumerIds.length) {
+        } else if (mCpuScalingPolicies.getPolicies().length
+                == mConsumedEnergyHelper.getEnergyConsumerCount()) {
             return initPowerBracketsByCluster(mDefaultCpuPowerBracketsPerEnergyConsumer);
         } else {
             Slog.i(TAG, "Assigning a single power brackets to each CPU_CLUSTER energy consumer."
                         + " Number of CPU clusters ("
                         + mCpuScalingPolicies.getPolicies().length
                         + ") does not match the number of energy consumers ("
-                        + mCpuEnergyConsumerIds.length + "). "
+                        + mConsumedEnergyHelper.getEnergyConsumerCount() + "). "
                         + " Using default power bucket assignment.");
             return initDefaultPowerBrackets(mDefaultCpuPowerBrackets);
         }
@@ -363,41 +354,11 @@ public class CpuPowerStatsCollector extends PowerStatsCollector {
         }
         mLayout.setUsageDuration(mCpuPowerStats.stats, uptimeDelta);
 
-        if (mCpuEnergyConsumerIds.length != 0) {
-            collectEnergyConsumers();
-        }
+        mConsumedEnergyHelper.collectConsumedEnergy(mCpuPowerStats, mLayout);
 
         return mCpuPowerStats;
     }
 
-    private void collectEnergyConsumers() {
-        int voltageMv = mVoltageSupplier.getAsInt();
-        if (voltageMv <= 0) {
-            Slog.wtf(TAG, "Unexpected battery voltage (" + voltageMv
-                          + " mV) when querying energy consumers");
-            return;
-        }
-
-        int averageVoltage = mLastVoltageMv != 0 ? (mLastVoltageMv + voltageMv) / 2 : voltageMv;
-        mLastVoltageMv = voltageMv;
-
-        long[] energyUws = mConsumedEnergyRetriever.getConsumedEnergyUws(mCpuEnergyConsumerIds);
-        if (energyUws == null) {
-            return;
-        }
-
-        for (int i = energyUws.length - 1; i >= 0; i--) {
-            long energyDelta = mLastConsumedEnergyUws[i] != ENERGY_UNSPECIFIED
-                    ? energyUws[i] - mLastConsumedEnergyUws[i] : 0;
-            if (energyDelta < 0) {
-                // Likely, restart of powerstats HAL
-                energyDelta = 0;
-            }
-            mLayout.setConsumedEnergy(mCpuPowerStats.stats, i, uJtoUc(energyDelta, averageVoltage));
-            mLastConsumedEnergyUws[i] = energyUws[i];
-        }
-    }
-
     @VisibleForNative
     interface KernelCpuStatsCallback {
         @Keep // Called from native
diff --git a/services/core/java/com/android/server/power/stats/CustomEnergyConsumerPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/CustomEnergyConsumerPowerStatsCollector.java
index 7d7eb29971c9..c1c7e12f2cab 100644
--- a/services/core/java/com/android/server/power/stats/CustomEnergyConsumerPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/CustomEnergyConsumerPowerStatsCollector.java
@@ -49,7 +49,8 @@ public class CustomEnergyConsumerPowerStatsCollector extends PowerStatsCollector
         for (int i = 0; i < energyConsumerIds.length; i++) {
             String name = retriever.getEnergyConsumerName(energyConsumerIds[i]);
             EnergyConsumerPowerStatsCollector collector = new EnergyConsumerPowerStatsCollector(
-                    mInjector, powerComponentId++, name, energyConsumerIds[i], sLayout);
+                    mInjector, powerComponentId++, name, EnergyConsumerType.OTHER,
+                    energyConsumerIds[i], sLayout);
             collector.setEnabled(true);
             collector.addConsumer(this::deliverStats);
             mCollectors.add(collector);
diff --git a/services/core/java/com/android/server/power/stats/EnergyConsumerPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/EnergyConsumerPowerStatsCollector.java
index c5229b90cff2..1d2e38849708 100644
--- a/services/core/java/com/android/server/power/stats/EnergyConsumerPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/EnergyConsumerPowerStatsCollector.java
@@ -16,24 +16,15 @@
 
 package com.android.server.power.stats;
 
-import android.hardware.power.stats.EnergyConsumerAttribution;
-import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.Handler;
 import android.os.PersistableBundle;
-import android.util.Slog;
-import android.util.SparseLongArray;
 
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerStats;
 import com.android.server.power.stats.format.EnergyConsumerPowerStatsLayout;
 
-import java.util.function.IntSupplier;
-
 public class EnergyConsumerPowerStatsCollector extends PowerStatsCollector {
-    private static final String TAG = "EnergyConsumerPowerStatsCollector";
-
-    private static final long ENERGY_UNSPECIFIED = -1;
 
     public interface Injector {
         Handler getHandler();
@@ -41,54 +32,41 @@ public class EnergyConsumerPowerStatsCollector extends PowerStatsCollector {
         PowerStatsUidResolver getUidResolver();
         long getPowerStatsCollectionThrottlePeriod(String powerComponentName);
         ConsumedEnergyRetriever getConsumedEnergyRetriever();
-        IntSupplier getVoltageSupplier();
     }
 
+    private static final int UNSPECIFIED = -1;
+
     private final Injector mInjector;
     private final int mPowerComponentId;
     private final String mPowerComponentName;
+    private final int mEnergyConsumerId;
     private final int mEnergyConsumerType;
-    private final String mEnergyConsumerName;
 
     private final EnergyConsumerPowerStatsLayout mLayout;
     private boolean mIsInitialized;
 
     private PowerStats mPowerStats;
-    private ConsumedEnergyRetriever mConsumedEnergyRetriever;
-    private IntSupplier mVoltageSupplier;
-    private int[] mEnergyConsumerIds;
-    private long mLastConsumedEnergyUws = ENERGY_UNSPECIFIED;
-    private SparseLongArray mLastConsumerEnergyPerUid = new SparseLongArray();
-    private int mLastVoltageMv;
+    private ConsumedEnergyHelper mConsumedEnergyHelper;
     private long mLastUpdateTimestamp;
-    private boolean mFirstCollection = true;
 
     EnergyConsumerPowerStatsCollector(Injector injector, int powerComponentId,
             String powerComponentName, @EnergyConsumerType int energyConsumerType,
-            String energyConsumerName, EnergyConsumerPowerStatsLayout statsLayout) {
-        super(injector.getHandler(),
-                injector.getPowerStatsCollectionThrottlePeriod(powerComponentName),
-                injector.getUidResolver(), injector.getClock());
-        mInjector = injector;
-        mPowerComponentId = powerComponentId;
-        mPowerComponentName = powerComponentName;
-        mEnergyConsumerType = energyConsumerType;
-        mEnergyConsumerName = energyConsumerName;
-        mLayout = statsLayout;
+            EnergyConsumerPowerStatsLayout statsLayout) {
+        this(injector, powerComponentId, powerComponentName, energyConsumerType, UNSPECIFIED,
+                statsLayout);
     }
 
     EnergyConsumerPowerStatsCollector(Injector injector, int powerComponentId,
-            String powerComponentName, int energyConsumerId,
-            EnergyConsumerPowerStatsLayout statsLayout) {
+            String powerComponentName, @EnergyConsumerType int energyConsumerType,
+            int energyConsumerId, EnergyConsumerPowerStatsLayout statsLayout) {
         super(injector.getHandler(),
                 injector.getPowerStatsCollectionThrottlePeriod(powerComponentName),
                 injector.getUidResolver(), injector.getClock());
         mInjector = injector;
         mPowerComponentId = powerComponentId;
         mPowerComponentName = powerComponentName;
-        mEnergyConsumerIds = new int[]{energyConsumerId};
-        mEnergyConsumerType = EnergyConsumerType.OTHER;
-        mEnergyConsumerName = null;
+        mEnergyConsumerId = energyConsumerId;
+        mEnergyConsumerType = energyConsumerType;
         mLayout = statsLayout;
     }
 
@@ -101,12 +79,14 @@ public class EnergyConsumerPowerStatsCollector extends PowerStatsCollector {
             return false;
         }
 
-        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
-        mVoltageSupplier = mInjector.getVoltageSupplier();
-        if (mEnergyConsumerIds == null) {
-            mEnergyConsumerIds = mConsumedEnergyRetriever.getEnergyConsumerIds(mEnergyConsumerType,
-                    mEnergyConsumerName);
+        if (mEnergyConsumerId != UNSPECIFIED) {
+            mConsumedEnergyHelper = new ConsumedEnergyHelper(mInjector.getConsumedEnergyRetriever(),
+                    mEnergyConsumerId, true /* perUidAttribution */);
+        } else {
+            mConsumedEnergyHelper = new ConsumedEnergyHelper(mInjector.getConsumedEnergyRetriever(),
+                    mEnergyConsumerType);
         }
+
         PersistableBundle extras = new PersistableBundle();
         mLayout.toExtras(extras);
         PowerStats.Descriptor powerStatsDescriptor = new PowerStats.Descriptor(
@@ -125,85 +105,15 @@ public class EnergyConsumerPowerStatsCollector extends PowerStatsCollector {
             return null;
         }
 
-        if (mEnergyConsumerIds.length == 0) {
-            return null;
-        }
-
-        int voltageMv = mVoltageSupplier.getAsInt();
-        int averageVoltage = mLastVoltageMv != 0 ? (mLastVoltageMv + voltageMv) / 2 : voltageMv;
-        if (averageVoltage <= 0) {
-            Slog.wtf(TAG, "Unexpected battery voltage (" + voltageMv
-                    + " mV) when querying energy consumers");
-            return null;
-        }
-
-        mLastVoltageMv = voltageMv;
-
-        EnergyConsumerResult[] energy =
-                    mConsumedEnergyRetriever.getConsumedEnergy(mEnergyConsumerIds);
-        long consumedEnergy = 0;
-        if (energy != null) {
-            for (int i = energy.length - 1; i >= 0; i--) {
-                if (energy[i].energyUWs != ENERGY_UNSPECIFIED) {
-                    consumedEnergy += energy[i].energyUWs;
-                }
-            }
-        }
-
-        long energyDelta = mLastConsumedEnergyUws != ENERGY_UNSPECIFIED
-                ? consumedEnergy - mLastConsumedEnergyUws : 0;
-        mLastConsumedEnergyUws = consumedEnergy;
-        if (energyDelta < 0) {
-            // Likely, restart of powerstats HAL
-            energyDelta = 0;
-        }
+        mPowerStats.uidStats.clear();
 
-        if (energyDelta == 0 && !mFirstCollection) {
+        if (!mConsumedEnergyHelper.collectConsumedEnergy(mPowerStats, mLayout)) {
             return null;
         }
 
-        mLayout.setConsumedEnergy(mPowerStats.stats, 0, uJtoUc(energyDelta, averageVoltage));
-
-        mPowerStats.uidStats.clear();
-        if (energy != null) {
-            for (int i = energy.length - 1; i >= 0; i--) {
-                EnergyConsumerAttribution[] perUid = energy[i].attribution;
-                if (perUid == null) {
-                    continue;
-                }
-
-                for (EnergyConsumerAttribution attribution : perUid) {
-                    int uid = mUidResolver.mapUid(attribution.uid);
-                    long lastEnergy = mLastConsumerEnergyPerUid.get(uid, ENERGY_UNSPECIFIED);
-                    mLastConsumerEnergyPerUid.put(uid, attribution.energyUWs);
-                    if (lastEnergy == ENERGY_UNSPECIFIED) {
-                        continue;
-                    }
-                    long deltaEnergy = attribution.energyUWs - lastEnergy;
-                    if (deltaEnergy <= 0) {
-                        continue;
-                    }
-                    long[] uidStats = mPowerStats.uidStats.get(uid);
-                    if (uidStats == null) {
-                        uidStats = new long[mLayout.getUidStatsArrayLength()];
-                        mPowerStats.uidStats.put(uid, uidStats);
-                    }
-
-                    mLayout.setUidConsumedEnergy(uidStats, 0,
-                            mLayout.getUidConsumedEnergy(uidStats, 0)
-                                    + uJtoUc(deltaEnergy, averageVoltage));
-                }
-            }
-        }
         long timestamp = mClock.elapsedRealtime();
         mPowerStats.durationMs = timestamp - mLastUpdateTimestamp;
         mLastUpdateTimestamp = timestamp;
-        mFirstCollection = false;
         return mPowerStats;
     }
-
-    @Override
-    protected void onUidRemoved(int uid) {
-        mLastConsumerEnergyPerUid.delete(uid);
-    }
 }
diff --git a/services/core/java/com/android/server/power/stats/GnssPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/GnssPowerStatsCollector.java
index ad5fe0b0064f..c1b4c31b1b77 100644
--- a/services/core/java/com/android/server/power/stats/GnssPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/GnssPowerStatsCollector.java
@@ -26,7 +26,6 @@ public class GnssPowerStatsCollector extends EnergyConsumerPowerStatsCollector {
     public GnssPowerStatsCollector(Injector injector) {
         super(injector, BatteryConsumer.POWER_COMPONENT_GNSS,
                 BatteryConsumer.powerComponentIdToString(BatteryConsumer.POWER_COMPONENT_GNSS),
-                EnergyConsumerType.GNSS, /* energy consumer name */ null,
-                new GnssPowerStatsLayout());
+                EnergyConsumerType.GNSS, new GnssPowerStatsLayout());
     }
 }
diff --git a/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsCollector.java
index 506c072dd0ed..cbd6fab2a9f7 100644
--- a/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsCollector.java
@@ -41,7 +41,6 @@ import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
-import java.util.function.IntSupplier;
 import java.util.function.LongSupplier;
 import java.util.function.Supplier;
 
@@ -57,8 +56,6 @@ public class MobileRadioPowerStatsCollector extends PowerStatsCollector {
 
     private static final long MODEM_ACTIVITY_REQUEST_TIMEOUT = 20000;
 
-    private static final long ENERGY_UNSPECIFIED = -1;
-
     @VisibleForTesting
     @AccessNetworkConstants.RadioAccessNetworkType
     static final int[] NETWORK_TYPES = {
@@ -85,7 +82,6 @@ public class MobileRadioPowerStatsCollector extends PowerStatsCollector {
         long getPowerStatsCollectionThrottlePeriod(String powerComponentName);
         PackageManager getPackageManager();
         ConsumedEnergyRetriever getConsumedEnergyRetriever();
-        IntSupplier getVoltageSupplier();
         Supplier<NetworkStats> getMobileNetworkStatsSupplier();
         TelephonyManager getTelephonyManager();
         LongSupplier getCallDurationSupplier();
@@ -104,14 +100,10 @@ public class MobileRadioPowerStatsCollector extends PowerStatsCollector {
     private LongSupplier mCallDurationSupplier;
     private LongSupplier mScanDurationSupplier;
     private volatile Supplier<NetworkStats> mNetworkStatsSupplier;
-    private ConsumedEnergyRetriever mConsumedEnergyRetriever;
-    private IntSupplier mVoltageSupplier;
-    private int[] mEnergyConsumerIds = new int[0];
+    private ConsumedEnergyHelper mConsumedEnergyHelper;
     private long mLastUpdateTimestampMillis;
     private ModemActivityInfo mLastModemActivityInfo;
     private NetworkStats mLastNetworkStats;
-    private long[] mLastConsumedEnergyUws;
-    private int mLastVoltageMv;
     private long mLastCallDuration;
     private long mLastScanDuration;
 
@@ -145,20 +137,15 @@ public class MobileRadioPowerStatsCollector extends PowerStatsCollector {
             return false;
         }
 
-        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
-        mVoltageSupplier = mInjector.getVoltageSupplier();
-
         mTelephonyManager = mInjector.getTelephonyManager();
         mNetworkStatsSupplier = mInjector.getMobileNetworkStatsSupplier();
         mCallDurationSupplier = mInjector.getCallDurationSupplier();
         mScanDurationSupplier = mInjector.getPhoneSignalScanDurationSupplier();
 
-        mEnergyConsumerIds = mConsumedEnergyRetriever.getEnergyConsumerIds(
+        mConsumedEnergyHelper = new ConsumedEnergyHelper(mInjector.getConsumedEnergyRetriever(),
                 EnergyConsumerType.MOBILE_RADIO);
-        mLastConsumedEnergyUws = new long[mEnergyConsumerIds.length];
-        Arrays.fill(mLastConsumedEnergyUws, ENERGY_UNSPECIFIED);
 
-        mLayout = new MobileRadioPowerStatsLayout(mEnergyConsumerIds.length);
+        mLayout = new MobileRadioPowerStatsLayout(mConsumedEnergyHelper.getEnergyConsumerCount());
 
         SparseArray<String> stateLabels = new SparseArray<>();
         for (int rat = 0; rat < BatteryStats.RADIO_ACCESS_TECHNOLOGY_COUNT; rat++) {
@@ -204,9 +191,8 @@ public class MobileRadioPowerStatsCollector extends PowerStatsCollector {
 
         ModemActivityInfo modemActivityDelta = collectModemActivityInfo();
         List<BatteryStatsImpl.NetworkStatsDelta> networkStatsDeltas = collectNetworkStats();
-        if (mEnergyConsumerIds.length != 0) {
-            collectEnergyConsumers();
-        }
+
+        mConsumedEnergyHelper.collectConsumedEnergy(mPowerStats, mLayout);
 
         if (mPowerStats.durationMs == 0) {
             setTimestamp(mClock.elapsedRealtime());
@@ -340,34 +326,6 @@ public class MobileRadioPowerStatsCollector extends PowerStatsCollector {
         return delta;
     }
 
-    private void collectEnergyConsumers() {
-        int voltageMv = mVoltageSupplier.getAsInt();
-        if (voltageMv <= 0) {
-            Slog.wtf(TAG, "Unexpected battery voltage (" + voltageMv
-                    + " mV) when querying energy consumers");
-            return;
-        }
-
-        int averageVoltage = mLastVoltageMv != 0 ? (mLastVoltageMv + voltageMv) / 2 : voltageMv;
-        mLastVoltageMv = voltageMv;
-
-        long[] energyUws = mConsumedEnergyRetriever.getConsumedEnergyUws(mEnergyConsumerIds);
-        if (energyUws == null) {
-            return;
-        }
-
-        for (int i = energyUws.length - 1; i >= 0; i--) {
-            long energyDelta = mLastConsumedEnergyUws[i] != ENERGY_UNSPECIFIED
-                    ? energyUws[i] - mLastConsumedEnergyUws[i] : 0;
-            if (energyDelta < 0) {
-                // Likely, restart of powerstats HAL
-                energyDelta = 0;
-            }
-            mLayout.setConsumedEnergy(mPowerStats.stats, i, uJtoUc(energyDelta, averageVoltage));
-            mLastConsumedEnergyUws[i] = energyUws[i];
-        }
-    }
-
     private void setTimestamp(long timestamp) {
         mPowerStats.durationMs = Math.max(timestamp - mLastUpdateTimestampMillis, 0);
         mLastUpdateTimestampMillis = timestamp;
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsCollector.java b/services/core/java/com/android/server/power/stats/PowerStatsCollector.java
index 7f220320edc5..291f0e38908e 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsCollector.java
@@ -19,6 +19,7 @@ package com.android.server.power.stats;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.hardware.power.stats.EnergyConsumer;
+import android.hardware.power.stats.EnergyConsumerAttribution;
 import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.ConditionVariable;
@@ -26,13 +27,16 @@ import android.os.Handler;
 import android.power.PowerStatsInternal;
 import android.util.IndentingPrintWriter;
 import android.util.Slog;
+import android.util.SparseLongArray;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.format.PowerStatsLayout;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.List;
@@ -41,6 +45,7 @@ import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.function.Consumer;
+import java.util.function.IntSupplier;
 
 /**
  * Collects snapshots of power-related system statistics.
@@ -53,6 +58,8 @@ public abstract class PowerStatsCollector {
     private static final String TAG = "PowerStatsCollector";
     private static final int MILLIVOLTS_PER_VOLT = 1000;
     private static final long POWER_STATS_ENERGY_CONSUMERS_TIMEOUT = 20000;
+    private static final long ENERGY_UNSPECIFIED = -1;
+
     private final Handler mHandler;
     protected final PowerStatsUidResolver mUidResolver;
     protected final Clock mClock;
@@ -236,45 +243,30 @@ public abstract class PowerStatsCollector {
     }
 
     public interface ConsumedEnergyRetriever {
+
         @NonNull
-        int[] getEnergyConsumerIds(@EnergyConsumerType int energyConsumerType, String name);
+        int[] getEnergyConsumerIds(@EnergyConsumerType int energyConsumerType);
 
         String getEnergyConsumerName(int energyConsumerId);
 
         @Nullable
         EnergyConsumerResult[] getConsumedEnergy(int[] energyConsumerIds);
 
-        @Nullable
-        default long[] getConsumedEnergyUws(int[] energyConsumerIds) {
-            EnergyConsumerResult[] results = getConsumedEnergy(energyConsumerIds);
-            if (results == null) {
-                return null;
-            }
-
-            long[] energy = new long[energyConsumerIds.length];
-            for (int i = 0; i < energyConsumerIds.length; i++) {
-                int id = energyConsumerIds[i];
-                for (EnergyConsumerResult result : results) {
-                    if (result.id == id) {
-                        energy[i] = result.energyUWs;
-                        break;
-                    }
-                }
-            }
-            return energy;
-        }
-
-        default int[] getEnergyConsumerIds(@EnergyConsumerType int energyConsumerType) {
-            return getEnergyConsumerIds(energyConsumerType, null);
-        }
+        /**
+         * Returns the last known battery/charger voltage in milli-volts.
+         */
+        int getVoltageMv();
     }
 
     static class ConsumedEnergyRetrieverImpl implements ConsumedEnergyRetriever {
         private final PowerStatsInternal mPowerStatsInternal;
+        private final IntSupplier mVoltageSupplier;
         private EnergyConsumer[] mEnergyConsumers;
 
-        ConsumedEnergyRetrieverImpl(PowerStatsInternal powerStatsInternal) {
+        ConsumedEnergyRetrieverImpl(PowerStatsInternal powerStatsInternal,
+                IntSupplier voltageSupplier) {
             mPowerStatsInternal = powerStatsInternal;
+            mVoltageSupplier = voltageSupplier;
         }
 
         private void ensureEnergyConsumers() {
@@ -293,8 +285,9 @@ public abstract class PowerStatsCollector {
             }
         }
 
+        @NonNull
         @Override
-        public int[] getEnergyConsumerIds(int energyConsumerType, String name) {
+        public int[] getEnergyConsumerIds(int energyConsumerType) {
             ensureEnergyConsumers();
 
             if (mEnergyConsumers.length == 0) {
@@ -303,8 +296,7 @@ public abstract class PowerStatsCollector {
 
             List<EnergyConsumer> energyConsumers = new ArrayList<>();
             for (EnergyConsumer energyConsumer : mEnergyConsumers) {
-                if (energyConsumer.type == energyConsumerType
-                        && (name == null || name.equals(energyConsumer.name))) {
+                if (energyConsumer.type == energyConsumerType) {
                     energyConsumers.add(energyConsumer);
                 }
             }
@@ -335,6 +327,11 @@ public abstract class PowerStatsCollector {
         }
 
         @Override
+        public int getVoltageMv() {
+            return mVoltageSupplier.getAsInt();
+        }
+
+        @Override
         public String getEnergyConsumerName(int energyConsumerId) {
             ensureEnergyConsumers();
 
@@ -368,4 +365,152 @@ public abstract class PowerStatsCollector {
             return sb.toString();
         }
     }
+
+    class ConsumedEnergyHelper implements PowerStatsUidResolver.Listener {
+        private final ConsumedEnergyRetriever mConsumedEnergyRetriever;
+        private final @EnergyConsumerType int mEnergyConsumerType;
+        private final boolean mPerUidAttributionSupported;
+
+        private boolean mIsInitialized;
+        private boolean mFirstCollection = true;
+        private int[] mEnergyConsumerIds;
+        private long[] mLastConsumedEnergyUws;
+        private final SparseLongArray mLastConsumerEnergyPerUid;
+        private int mLastVoltageMv;
+
+        ConsumedEnergyHelper(ConsumedEnergyRetriever consumedEnergyRetriever,
+                @EnergyConsumerType int energyConsumerType) {
+            mConsumedEnergyRetriever = consumedEnergyRetriever;
+            mEnergyConsumerType = energyConsumerType;
+            mPerUidAttributionSupported = false;
+            mLastConsumerEnergyPerUid = null;
+        }
+
+        ConsumedEnergyHelper(ConsumedEnergyRetriever consumedEnergyRetriever,
+                int energyConsumerId, boolean perUidAttributionSupported) {
+            mConsumedEnergyRetriever = consumedEnergyRetriever;
+            mEnergyConsumerType = EnergyConsumerType.OTHER;
+            mEnergyConsumerIds = new int[]{energyConsumerId};
+            mPerUidAttributionSupported = perUidAttributionSupported;
+            mLastConsumerEnergyPerUid = mPerUidAttributionSupported ? new SparseLongArray() : null;
+        }
+
+        private void ensureInitialized() {
+            if (!mIsInitialized) {
+                if (mEnergyConsumerIds == null) {
+                    mEnergyConsumerIds = mConsumedEnergyRetriever.getEnergyConsumerIds(
+                            mEnergyConsumerType);
+                }
+                mLastConsumedEnergyUws = new long[mEnergyConsumerIds.length];
+                Arrays.fill(mLastConsumedEnergyUws, ENERGY_UNSPECIFIED);
+                mUidResolver.addListener(this);
+                mIsInitialized = true;
+            }
+        }
+
+        int getEnergyConsumerCount() {
+            ensureInitialized();
+            return mEnergyConsumerIds.length;
+        }
+
+        boolean collectConsumedEnergy(PowerStats powerStats, PowerStatsLayout layout) {
+            ensureInitialized();
+
+            if (mEnergyConsumerIds.length == 0) {
+                return false;
+            }
+
+            int voltageMv = mConsumedEnergyRetriever.getVoltageMv();
+            int averageVoltage = mLastVoltageMv != 0 ? (mLastVoltageMv + voltageMv) / 2 : voltageMv;
+            if (averageVoltage <= 0) {
+                Slog.wtf(TAG, "Unexpected battery voltage (" + voltageMv
+                        + " mV) when querying energy consumers");
+                return false;
+            }
+
+            mLastVoltageMv = voltageMv;
+
+            EnergyConsumerResult[] energy =
+                    mConsumedEnergyRetriever.getConsumedEnergy(mEnergyConsumerIds);
+            System.out.println("mEnergyConsumerIds = " + Arrays.toString(mEnergyConsumerIds) + " "
+                    + "energy = "
+                    + Arrays.toString(energy));
+            if (energy == null) {
+                return false;
+            }
+
+            for (int i = 0; i < mEnergyConsumerIds.length; i++) {
+                populatePowerStats(powerStats, layout, energy, i, averageVoltage);
+            }
+            mFirstCollection = false;
+            return true;
+        }
+
+        private void populatePowerStats(PowerStats powerStats, PowerStatsLayout layout,
+                @NonNull EnergyConsumerResult[] energy, int energyConsumerIndex,
+                int averageVoltage) {
+            long consumedEnergy = energy[energyConsumerIndex].energyUWs;
+            long energyDelta = mLastConsumedEnergyUws[energyConsumerIndex] != ENERGY_UNSPECIFIED
+                    ? consumedEnergy - mLastConsumedEnergyUws[energyConsumerIndex] : 0;
+            mLastConsumedEnergyUws[energyConsumerIndex] = consumedEnergy;
+            if (energyDelta < 0) {
+                // Likely, restart of powerstats HAL
+                energyDelta = 0;
+            }
+
+            if (energyDelta == 0 && !mFirstCollection) {
+                return;
+            }
+
+            layout.setConsumedEnergy(powerStats.stats, energyConsumerIndex,
+                    uJtoUc(energyDelta, averageVoltage));
+
+            if (!mPerUidAttributionSupported) {
+                return;
+            }
+
+            EnergyConsumerAttribution[] perUid = energy[energyConsumerIndex].attribution;
+            if (perUid == null) {
+                return;
+            }
+
+            for (EnergyConsumerAttribution attribution : perUid) {
+                int uid = mUidResolver.mapUid(attribution.uid);
+                long lastEnergy = mLastConsumerEnergyPerUid.get(uid, ENERGY_UNSPECIFIED);
+                mLastConsumerEnergyPerUid.put(uid, attribution.energyUWs);
+                if (lastEnergy == ENERGY_UNSPECIFIED) {
+                    continue;
+                }
+                long deltaEnergy = attribution.energyUWs - lastEnergy;
+                if (deltaEnergy <= 0) {
+                    continue;
+                }
+
+                long[] uidStats = powerStats.uidStats.get(uid);
+                if (uidStats == null) {
+                    uidStats = new long[layout.getUidStatsArrayLength()];
+                    powerStats.uidStats.put(uid, uidStats);
+                }
+
+                layout.setUidConsumedEnergy(uidStats, energyConsumerIndex,
+                        layout.getUidConsumedEnergy(uidStats, energyConsumerIndex)
+                                + uJtoUc(deltaEnergy, averageVoltage));
+            }
+        }
+
+        @Override
+        public void onAfterIsolatedUidRemoved(int isolatedUid, int parentUid) {
+            if (mLastConsumerEnergyPerUid != null) {
+                mHandler.post(() -> mLastConsumerEnergyPerUid.delete(isolatedUid));
+            }
+        }
+
+        @Override
+        public void onIsolatedUidAdded(int isolatedUid, int parentUid) {
+        }
+
+        @Override
+        public void onBeforeIsolatedUidRemoved(int isolatedUid, int parentUid) {
+        }
+    }
 }
diff --git a/services/core/java/com/android/server/power/stats/ScreenPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/ScreenPowerStatsCollector.java
index 4b5d89907f84..8371e6681747 100644
--- a/services/core/java/com/android/server/power/stats/ScreenPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/ScreenPowerStatsCollector.java
@@ -21,16 +21,12 @@ import android.os.BatteryConsumer;
 import android.os.BatteryStats;
 import android.os.Handler;
 import android.os.PersistableBundle;
-import android.util.Slog;
 import android.util.SparseLongArray;
 
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerStats;
 import com.android.server.power.stats.format.ScreenPowerStatsLayout;
 
-import java.util.Arrays;
-import java.util.function.IntSupplier;
-
 public class ScreenPowerStatsCollector extends PowerStatsCollector {
     private static final String TAG = "ScreenPowerStatsCollector";
 
@@ -52,24 +48,17 @@ public class ScreenPowerStatsCollector extends PowerStatsCollector {
         PowerStatsUidResolver getUidResolver();
         long getPowerStatsCollectionThrottlePeriod(String powerComponentName);
         ConsumedEnergyRetriever getConsumedEnergyRetriever();
-        IntSupplier getVoltageSupplier();
         ScreenUsageTimeRetriever getScreenUsageTimeRetriever();
         int getDisplayCount();
     }
 
-    private static final long ENERGY_UNSPECIFIED = -1;
-
     private final Injector mInjector;
     private boolean mIsInitialized;
     private ScreenPowerStatsLayout mLayout;
     private int mDisplayCount;
     private PowerStats mPowerStats;
-    private ConsumedEnergyRetriever mConsumedEnergyRetriever;
-    private IntSupplier mVoltageSupplier;
+    private ConsumedEnergyHelper mConsumedEnergyHelper;
     private ScreenUsageTimeRetriever mScreenUsageTimeRetriever;
-    private int[] mEnergyConsumerIds = new int[0];
-    private long[] mLastConsumedEnergyUws;
-    private int mLastVoltageMv;
     private boolean mFirstSample = true;
     private long[] mLastScreenOnTime;
     private long[][] mLastBrightnessLevelTime;
@@ -96,15 +85,11 @@ public class ScreenPowerStatsCollector extends PowerStatsCollector {
         }
 
         mDisplayCount = mInjector.getDisplayCount();
-        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
-        mVoltageSupplier = mInjector.getVoltageSupplier();
         mScreenUsageTimeRetriever = mInjector.getScreenUsageTimeRetriever();
-        mEnergyConsumerIds = mConsumedEnergyRetriever.getEnergyConsumerIds(
-                EnergyConsumerType.DISPLAY);
-        mLastConsumedEnergyUws = new long[mEnergyConsumerIds.length];
-        Arrays.fill(mLastConsumedEnergyUws, ENERGY_UNSPECIFIED);
 
-        mLayout = new ScreenPowerStatsLayout(mEnergyConsumerIds.length,
+        mConsumedEnergyHelper = new ConsumedEnergyHelper(mInjector.getConsumedEnergyRetriever(),
+                EnergyConsumerType.DISPLAY);
+        mLayout = new ScreenPowerStatsLayout(mConsumedEnergyHelper.getEnergyConsumerCount(),
                 mInjector.getDisplayCount());
 
         PersistableBundle extras = new PersistableBundle();
@@ -130,9 +115,7 @@ public class ScreenPowerStatsCollector extends PowerStatsCollector {
             return null;
         }
 
-        if (mEnergyConsumerIds.length != 0) {
-            collectEnergyConsumers();
-        }
+        mConsumedEnergyHelper.collectConsumedEnergy(mPowerStats, mLayout);
 
         for (int display = 0; display < mDisplayCount; display++) {
             long screenOnTimeMs = mScreenUsageTimeRetriever.getScreenOnTimeMs(display);
@@ -188,34 +171,6 @@ public class ScreenPowerStatsCollector extends PowerStatsCollector {
         return mPowerStats;
     }
 
-    private void collectEnergyConsumers() {
-        int voltageMv = mVoltageSupplier.getAsInt();
-        if (voltageMv <= 0) {
-            Slog.wtf(TAG, "Unexpected battery voltage (" + voltageMv
-                    + " mV) when querying energy consumers");
-            return;
-        }
-
-        int averageVoltage = mLastVoltageMv != 0 ? (mLastVoltageMv + voltageMv) / 2 : voltageMv;
-        mLastVoltageMv = voltageMv;
-
-        long[] energyUws = mConsumedEnergyRetriever.getConsumedEnergyUws(mEnergyConsumerIds);
-        if (energyUws == null) {
-            return;
-        }
-
-        for (int i = energyUws.length - 1; i >= 0; i--) {
-            long energyDelta = mLastConsumedEnergyUws[i] != ENERGY_UNSPECIFIED
-                    ? energyUws[i] - mLastConsumedEnergyUws[i] : 0;
-            if (energyDelta < 0) {
-                // Likely, restart of powerstats HAL
-                energyDelta = 0;
-            }
-            mLayout.setConsumedEnergy(mPowerStats.stats, i, uJtoUc(energyDelta, averageVoltage));
-            mLastConsumedEnergyUws[i] = energyUws[i];
-        }
-    }
-
     @Override
     protected void onUidRemoved(int uid) {
         mLastTopActivityTime.delete(uid);
diff --git a/services/core/java/com/android/server/power/stats/WifiPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/WifiPowerStatsCollector.java
index 45bf5739cddb..7a84b05823a4 100644
--- a/services/core/java/com/android/server/power/stats/WifiPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/WifiPowerStatsCollector.java
@@ -30,11 +30,9 @@ import com.android.internal.os.Clock;
 import com.android.internal.os.PowerStats;
 import com.android.server.power.stats.format.WifiPowerStatsLayout;
 
-import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
-import java.util.function.IntSupplier;
 import java.util.function.Supplier;
 
 public class WifiPowerStatsCollector extends PowerStatsCollector {
@@ -42,8 +40,6 @@ public class WifiPowerStatsCollector extends PowerStatsCollector {
 
     private static final long WIFI_ACTIVITY_REQUEST_TIMEOUT = 20000;
 
-    private static final long ENERGY_UNSPECIFIED = -1;
-
     interface Observer {
         void onWifiPowerStatsRetrieved(WifiActivityEnergyInfo info,
                 List<BatteryStatsImpl.NetworkStatsDelta> delta, long elapsedRealtimeMs,
@@ -66,7 +62,6 @@ public class WifiPowerStatsCollector extends PowerStatsCollector {
         long getPowerStatsCollectionThrottlePeriod(String powerComponentName);
         PackageManager getPackageManager();
         ConsumedEnergyRetriever getConsumedEnergyRetriever();
-        IntSupplier getVoltageSupplier();
         Supplier<NetworkStats> getWifiNetworkStatsSupplier();
         WifiManager getWifiManager();
         WifiStatsRetriever getWifiStatsRetriever();
@@ -84,13 +79,9 @@ public class WifiPowerStatsCollector extends PowerStatsCollector {
     private volatile WifiManager mWifiManager;
     private volatile Supplier<NetworkStats> mNetworkStatsSupplier;
     private volatile WifiStatsRetriever mWifiStatsRetriever;
-    private ConsumedEnergyRetriever mConsumedEnergyRetriever;
-    private IntSupplier mVoltageSupplier;
-    private int[] mEnergyConsumerIds = new int[0];
+    private ConsumedEnergyHelper mConsumedEnergyHelper;
     private WifiActivityEnergyInfo mLastWifiActivityInfo;
     private NetworkStats mLastNetworkStats;
-    private long[] mLastConsumedEnergyUws;
-    private int mLastVoltageMv;
 
     private static class WifiScanTimes {
         public long basicScanTimeMs;
@@ -129,19 +120,17 @@ public class WifiPowerStatsCollector extends PowerStatsCollector {
             return false;
         }
 
-        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
-        mVoltageSupplier = mInjector.getVoltageSupplier();
         mWifiManager = mInjector.getWifiManager();
         mNetworkStatsSupplier = mInjector.getWifiNetworkStatsSupplier();
         mWifiStatsRetriever = mInjector.getWifiStatsRetriever();
         mPowerReportingSupported =
                 mWifiManager != null && mWifiManager.isEnhancedPowerReportingSupported();
 
-        mEnergyConsumerIds = mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.WIFI);
-        mLastConsumedEnergyUws = new long[mEnergyConsumerIds.length];
-        Arrays.fill(mLastConsumedEnergyUws, ENERGY_UNSPECIFIED);
+        mConsumedEnergyHelper = new ConsumedEnergyHelper(mInjector.getConsumedEnergyRetriever(),
+                EnergyConsumerType.WIFI);
 
-        mLayout = new WifiPowerStatsLayout(mEnergyConsumerIds.length, mPowerReportingSupported);
+        mLayout = new WifiPowerStatsLayout(mConsumedEnergyHelper.getEnergyConsumerCount(),
+                mPowerReportingSupported);
 
         PersistableBundle extras = new PersistableBundle();
         mLayout.toExtras(extras);
@@ -171,9 +160,7 @@ public class WifiPowerStatsCollector extends PowerStatsCollector {
         List<BatteryStatsImpl.NetworkStatsDelta> networkStatsDeltas = collectNetworkStats();
         collectWifiScanTime();
 
-        if (mEnergyConsumerIds.length != 0) {
-            collectEnergyConsumers();
-        }
+        mConsumedEnergyHelper.collectConsumedEnergy(mPowerStats, mLayout);
 
         if (mObserver != null) {
             mObserver.onWifiPowerStatsRetrieved(activityInfo, networkStatsDeltas,
@@ -313,34 +300,6 @@ public class WifiPowerStatsCollector extends PowerStatsCollector {
         mLayout.setDeviceBatchedScanTime(mDeviceStats, mScanTimes.batchedScanTimeMs);
     }
 
-    private void collectEnergyConsumers() {
-        int voltageMv = mVoltageSupplier.getAsInt();
-        if (voltageMv <= 0) {
-            Slog.wtf(TAG, "Unexpected battery voltage (" + voltageMv
-                    + " mV) when querying energy consumers");
-            return;
-        }
-
-        int averageVoltage = mLastVoltageMv != 0 ? (mLastVoltageMv + voltageMv) / 2 : voltageMv;
-        mLastVoltageMv = voltageMv;
-
-        long[] energyUws = mConsumedEnergyRetriever.getConsumedEnergyUws(mEnergyConsumerIds);
-        if (energyUws == null) {
-            return;
-        }
-
-        for (int i = energyUws.length - 1; i >= 0; i--) {
-            long energyDelta = mLastConsumedEnergyUws[i] != ENERGY_UNSPECIFIED
-                    ? energyUws[i] - mLastConsumedEnergyUws[i] : 0;
-            if (energyDelta < 0) {
-                // Likely, restart of powerstats HAL
-                energyDelta = 0;
-            }
-            mLayout.setConsumedEnergy(mPowerStats.stats, i, uJtoUc(energyDelta, averageVoltage));
-            mLastConsumedEnergyUws[i] = energyUws[i];
-        }
-    }
-
     @Override
     protected void onUidRemoved(int uid) {
         super.onUidRemoved(uid);
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BluetoothPowerStatsCollectorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BluetoothPowerStatsCollectorTest.java
index cfcbfca71b3b..e392c5d190f7 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BluetoothPowerStatsCollectorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BluetoothPowerStatsCollectorTest.java
@@ -27,6 +27,7 @@ import android.bluetooth.BluetoothAdapter;
 import android.bluetooth.UidTraffic;
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.BatteryConsumer;
 import android.os.Handler;
@@ -49,7 +50,6 @@ import java.io.StringWriter;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.Executor;
-import java.util.function.IntSupplier;
 
 public class BluetoothPowerStatsCollectorTest {
     private static final int APP_UID1 = 42;
@@ -133,11 +133,6 @@ public class BluetoothPowerStatsCollectorTest {
                 }
 
                 @Override
-                public IntSupplier getVoltageSupplier() {
-                    return () -> 3500;
-                }
-
-                @Override
                 public BluetoothPowerStatsCollector.BluetoothStatsRetriever
                         getBluetoothStatsRetriever() {
                     return mBluetoothStatsRetriever;
@@ -233,6 +228,7 @@ public class BluetoothPowerStatsCollectorTest {
         BluetoothPowerStatsCollector collector = new BluetoothPowerStatsCollector(mInjector);
         collector.setEnabled(true);
 
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(3500);
         when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.BLUETOOTH))
                 .thenReturn(new int[]{777});
 
@@ -243,8 +239,7 @@ public class BluetoothPowerStatsCollectorTest {
 
         mUidScanTimes.put(APP_UID1, 100);
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
-                .thenReturn(new long[]{10000});
+        mockConsumedEnergy(777, 10000);
 
         // Establish a baseline
         collector.collectStats();
@@ -259,13 +254,19 @@ public class BluetoothPowerStatsCollectorTest {
         mUidScanTimes.put(APP_UID2, 300);
         mUidScanTimes.put(ISOLATED_UID, 400);
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
-                .thenReturn(new long[]{64321});
+        mockConsumedEnergy(777, 64321);
 
         mStatsRule.setTime(20000, 20000);
         return collector.collectStats();
     }
 
+    private void mockConsumedEnergy(int consumerId, long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        when(mConsumedEnergyRetriever.getConsumedEnergy(eq(new int[]{consumerId})))
+                .thenReturn(new EnergyConsumerResult[]{ecr});
+    }
+
     private BluetoothActivityEnergyInfo mockBluetoothActivityEnergyInfo(long timestamp,
             long rxTimeMs, long txTimeMs, long idleTimeMs, UidTraffic... uidTraffic) {
         if (RavenwoodRule.isOnRavenwood()) {
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsCollectorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsCollectorTest.java
index d995227cfec3..1fea46256b12 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsCollectorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsCollectorTest.java
@@ -26,6 +26,7 @@ import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.when;
 
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.BatteryConsumer;
 import android.os.ConditionVariable;
@@ -54,7 +55,6 @@ import org.xmlpull.v1.XmlPullParserException;
 
 import java.io.IOException;
 import java.io.StringWriter;
-import java.util.function.IntSupplier;
 
 @RunWith(AndroidJUnit4.class)
 @SmallTest
@@ -125,11 +125,6 @@ public class CpuPowerStatsCollectorTest {
         }
 
         @Override
-        public IntSupplier getVoltageSupplier() {
-            return () -> 3500;
-        }
-
-        @Override
         public long getPowerStatsCollectionThrottlePeriod(String powerComponentName) {
             return 0;
         }
@@ -151,7 +146,9 @@ public class CpuPowerStatsCollectorTest {
         mHandlerThread.start();
         mHandler = mHandlerThread.getThreadHandler();
         when(mMockKernelCpuStatsReader.isSupportedFeature()).thenReturn(true);
-        when(mConsumedEnergyRetriever.getEnergyConsumerIds(anyInt())).thenReturn(new int[0]);
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(anyInt()))
+                .thenReturn(new int[0]);
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(3500);
         mUidResolver.noteIsolatedUidAdded(ISOLATED_UID, UID_2);
     }
 
@@ -459,11 +456,20 @@ public class CpuPowerStatsCollectorTest {
 
     private void mockEnergyConsumers() {
         reset(mConsumedEnergyRetriever);
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(3500);
         when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.CPU_CLUSTER))
                 .thenReturn(new int[]{1, 2});
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{1, 2})))
-                .thenReturn(new long[]{1000, 2000})
-                .thenReturn(new long[]{1500, 2700});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(eq(new int[]{1, 2})))
+                .thenReturn(new EnergyConsumerResult[]{
+                        mockEnergyConsumer(1000), mockEnergyConsumer(2000)})
+                .thenReturn(new EnergyConsumerResult[]{
+                        mockEnergyConsumer(1500), mockEnergyConsumer(2700)});
+    }
+
+    private EnergyConsumerResult mockEnergyConsumer(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
     }
 
     private static int[] getScalingStepToPowerBracketMap(CpuPowerStatsCollector collector) {
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerStatsCollectorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerStatsCollectorTest.java
index 2e4f32b45c4a..00b911bc1ae7 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerStatsCollectorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerStatsCollectorTest.java
@@ -33,6 +33,7 @@ import static org.mockito.Mockito.when;
 import android.annotation.NonNull;
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.net.NetworkStats;
 import android.os.BatteryConsumer;
@@ -64,7 +65,6 @@ import java.lang.reflect.Field;
 import java.lang.reflect.Modifier;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.function.IntSupplier;
 import java.util.function.LongSupplier;
 import java.util.function.Supplier;
 
@@ -140,11 +140,6 @@ public class MobileRadioPowerStatsCollectorTest {
         }
 
         @Override
-        public IntSupplier getVoltageSupplier() {
-            return () -> 3500;
-        }
-
-        @Override
         public Supplier<NetworkStats> getMobileNetworkStatsSupplier() {
             return mNetworkStatsSupplier;
         }
@@ -179,6 +174,7 @@ public class MobileRadioPowerStatsCollectorTest {
                 return uid;
             }
         });
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(3500);
         mBatteryStats = mStatsRule.getBatteryStats();
     }
 
@@ -415,8 +411,8 @@ public class MobileRadioPowerStatsCollectorTest {
                 4321, 321, 1234, 23,
                 4000, 40, 2000, 20);
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
-                .thenReturn(new long[]{10000});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(eq(new int[]{777})))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(10000)});
 
         when(mCallDurationSupplier.getAsLong()).thenReturn(10000L);
         when(mScanDurationSupplier.getAsLong()).thenReturn(20000L);
@@ -438,8 +434,8 @@ public class MobileRadioPowerStatsCollectorTest {
                 5321, 421, 3234, 223,
                 8000, 80, 4000, 40);
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
-                .thenReturn(new long[]{64321});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(eq(new int[]{777})))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(64321)});
         when(mCallDurationSupplier.getAsLong()).thenReturn(50000L);
         when(mScanDurationSupplier.getAsLong()).thenReturn(80000L);
 
@@ -447,6 +443,12 @@ public class MobileRadioPowerStatsCollectorTest {
         return collector.collectStats();
     }
 
+    private EnergyConsumerResult mockEnergyConsumer(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
+    }
+
     private void mockModemActivityInfo(long timestamp, int sleepTimeMs, int idleTimeMs,
             int networkType1, int freqRange1, int rxTimeMs1, @NonNull int[] txTimeMs1,
             int networkType2, int freqRange2, int rxTimeMs2, @NonNull int[] txTimeMs2) {
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsCollectorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsCollectorTest.java
index 89d6c1c66ca3..a04f7212204d 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsCollectorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsCollectorTest.java
@@ -114,13 +114,15 @@ public class PowerStatsCollectorTest {
         mockEnergyConsumers(powerStatsInternal);
 
         PowerStatsCollector.ConsumedEnergyRetrieverImpl retriever =
-                new PowerStatsCollector.ConsumedEnergyRetrieverImpl(powerStatsInternal);
+                new PowerStatsCollector.ConsumedEnergyRetrieverImpl(powerStatsInternal, ()-> 3500);
         int[] energyConsumerIds = retriever.getEnergyConsumerIds(EnergyConsumerType.CPU_CLUSTER);
         assertThat(energyConsumerIds).isEqualTo(new int[]{1, 2});
-        long[] energy = retriever.getConsumedEnergyUws(energyConsumerIds);
-        assertThat(energy).isEqualTo(new long[]{1000, 2000});
-        energy = retriever.getConsumedEnergyUws(energyConsumerIds);
-        assertThat(energy).isEqualTo(new long[]{1500, 2700});
+        EnergyConsumerResult[] energy = retriever.getConsumedEnergy(energyConsumerIds);
+        assertThat(energy[0].energyUWs).isEqualTo(1000);
+        assertThat(energy[1].energyUWs).isEqualTo(2000);
+        energy = retriever.getConsumedEnergy(energyConsumerIds);
+        assertThat(energy[0].energyUWs).isEqualTo(1500);
+        assertThat(energy[1].energyUWs).isEqualTo(2700);
     }
 
     @SuppressWarnings("unchecked")
@@ -176,4 +178,11 @@ public class PowerStatsCollectorTest {
                 .thenReturn(future1)
                 .thenReturn(future2);
     }
+
+    private EnergyConsumerResult mockEnergyConsumerResult(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
+    }
+
 }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsCollectorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsCollectorTest.java
index 8fd3f4dd00a1..8c09d1dbfc9d 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsCollectorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsCollectorTest.java
@@ -23,6 +23,7 @@ import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.when;
 
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.BatteryConsumer;
 import android.os.BatteryStats;
@@ -40,8 +41,6 @@ import org.junit.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import java.util.function.IntSupplier;
-
 public class ScreenPowerStatsCollectorTest {
     private static final int APP_UID1 = 42;
     private static final int APP_UID2 = 24;
@@ -90,11 +89,6 @@ public class ScreenPowerStatsCollectorTest {
         }
 
         @Override
-        public IntSupplier getVoltageSupplier() {
-            return () -> 3500;
-        }
-
-        @Override
         public int getDisplayCount() {
             return 2;
         }
@@ -116,6 +110,7 @@ public class ScreenPowerStatsCollectorTest {
                 return uid;
             }
         });
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(3500);
     }
 
     @Test
@@ -126,8 +121,8 @@ public class ScreenPowerStatsCollectorTest {
         // Establish a baseline
         when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.DISPLAY))
                 .thenReturn(new int[]{77});
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(new int[]{77}))
-                .thenReturn(new long[]{10_000});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(new int[]{77}))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(10_000)});
 
         doAnswer(inv -> {
             ScreenPowerStatsCollector.ScreenUsageTimeRetriever.Callback callback =
@@ -140,8 +135,8 @@ public class ScreenPowerStatsCollectorTest {
 
         collector.collectStats();
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(new int[]{77}))
-                .thenReturn(new long[]{45_000});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(new int[]{77}))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(45_000)});
         when(mScreenUsageTimeRetriever.getScreenOnTimeMs(0))
                 .thenReturn(60_000L);
         when(mScreenUsageTimeRetriever.getBrightnessLevelTimeMs(0,
@@ -204,4 +199,10 @@ public class ScreenPowerStatsCollectorTest {
         assertThat(layout.getUidTopActivityDuration(powerStats.uidStats.get(APP_UID2)))
                 .isEqualTo(10000);
     }
+
+    private EnergyConsumerResult mockEnergyConsumer(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
+    }
 }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsCollectorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsCollectorTest.java
index f887d83f9f60..8b5e6ee9cf89 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsCollectorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsCollectorTest.java
@@ -32,6 +32,7 @@ import static org.mockito.Mockito.when;
 
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.net.NetworkStats;
 import android.net.wifi.WifiManager;
@@ -58,7 +59,6 @@ import org.mockito.MockitoAnnotations;
 import java.io.StringWriter;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.function.IntSupplier;
 import java.util.function.Supplier;
 
 public class WifiPowerStatsCollectorTest {
@@ -155,11 +155,6 @@ public class WifiPowerStatsCollectorTest {
         }
 
         @Override
-        public IntSupplier getVoltageSupplier() {
-            return () -> 3500;
-        }
-
-        @Override
         public Supplier<NetworkStats> getWifiNetworkStatsSupplier() {
             return mNetworkStatsSupplier;
         }
@@ -369,6 +364,7 @@ public class WifiPowerStatsCollectorTest {
         WifiPowerStatsCollector collector = new WifiPowerStatsCollector(mInjector, null);
         collector.setEnabled(true);
 
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(3500);
         when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.WIFI))
                 .thenReturn(new int[]{777});
 
@@ -386,8 +382,8 @@ public class WifiPowerStatsCollectorTest {
         mockWifiScanTimes(APP_UID2, 3000, 4000);
         mockWifiScanTimes(ISOLATED_UID, 5000, 6000);
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
-                .thenReturn(new long[]{10000});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(eq(new int[]{777})))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(10_000)});
 
         collector.collectStats();
 
@@ -405,13 +401,19 @@ public class WifiPowerStatsCollectorTest {
         mockWifiScanTimes(APP_UID2, 3100, 4200);
         mockWifiScanTimes(ISOLATED_UID, 5300, 6400);
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
-                .thenReturn(new long[]{64321});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(eq(new int[]{777})))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(64_321)});
 
         mStatsRule.setTime(20000, 20000);
         return collector.collectStats();
     }
 
+    private EnergyConsumerResult mockEnergyConsumer(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
+    }
+
     private void mockWifiActivityInfo(long timestamp, long rxTimeMs, long txTimeMs, int scanTimeMs,
             int idleTimeMs) {
         int stackState = 0;
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/AmbientDisplayPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/AmbientDisplayPowerStatsProcessorTest.java
index 6477b26b90e8..21e615f8c740 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/AmbientDisplayPowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/AmbientDisplayPowerStatsProcessorTest.java
@@ -48,8 +48,6 @@ import org.junit.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import java.util.function.IntSupplier;
-
 public class AmbientDisplayPowerStatsProcessorTest {
 
     @Rule(order = 0)
@@ -99,11 +97,6 @@ public class AmbientDisplayPowerStatsProcessorTest {
                 }
 
                 @Override
-                public IntSupplier getVoltageSupplier() {
-                    return () -> VOLTAGE_MV;
-                }
-
-                @Override
                 public int getDisplayCount() {
                     return 2;
                 }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/BluetoothPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/BluetoothPowerStatsProcessorTest.java
index 371c6c925b87..6dfc22077357 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/BluetoothPowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/BluetoothPowerStatsProcessorTest.java
@@ -38,6 +38,7 @@ import android.bluetooth.BluetoothAdapter;
 import android.bluetooth.UidTraffic;
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.BatteryConsumer;
 import android.os.Handler;
@@ -63,7 +64,6 @@ import org.mockito.MockitoAnnotations;
 
 import java.util.List;
 import java.util.concurrent.Executor;
-import java.util.function.IntSupplier;
 import java.util.function.Supplier;
 
 public class BluetoothPowerStatsProcessorTest {
@@ -148,11 +148,6 @@ public class BluetoothPowerStatsProcessorTest {
                 }
 
                 @Override
-                public IntSupplier getVoltageSupplier() {
-                    return () -> VOLTAGE_MV;
-                }
-
-                @Override
                 public BluetoothStatsRetriever getBluetoothStatsRetriever() {
                     return mBluetoothStatsRetriever;
                 }
@@ -368,7 +363,8 @@ public class BluetoothPowerStatsProcessorTest {
 
     @Test
     public void consumedEnergyModel() {
-        // No power monitoring hardware
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(VOLTAGE_MV);
+        // Power monitoring hardware exists
         when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.BLUETOOTH))
                 .thenReturn(new int[]{BLUETOOTH_ENERGY_CONSUMER_ID});
 
@@ -383,8 +379,8 @@ public class BluetoothPowerStatsProcessorTest {
 
         mUidScanTimes.put(APP_UID1, 100);
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(
-                new int[]{BLUETOOTH_ENERGY_CONSUMER_ID})).thenReturn(new long[]{0});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(new int[]{BLUETOOTH_ENERGY_CONSUMER_ID}))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(0)});
 
         aggregatedStats.start(0);
 
@@ -409,8 +405,8 @@ public class BluetoothPowerStatsProcessorTest {
 
         // 10 mAh represented as microWattSeconds
         long energyUws = 10 * 3600 * VOLTAGE_MV;
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(
-                new int[]{BLUETOOTH_ENERGY_CONSUMER_ID})).thenReturn(new long[]{energyUws});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(new int[]{BLUETOOTH_ENERGY_CONSUMER_ID}))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(energyUws)});
 
         aggregatedStats.addPowerStats(collector.collectStats(), 10_000);
 
@@ -495,6 +491,12 @@ public class BluetoothPowerStatsProcessorTest {
         return states;
     }
 
+    private EnergyConsumerResult mockEnergyConsumer(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
+    }
+
     private BluetoothActivityEnergyInfo mockBluetoothActivityEnergyInfo(long timestamp,
             long rxTimeMs, long txTimeMs, long idleTimeMs, UidTraffic... uidTraffic) {
         if (RavenwoodRule.isOnRavenwood()) {
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/CameraPowerStatsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/CameraPowerStatsTest.java
index 4c60e0af4ba6..0afcbf15415c 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/CameraPowerStatsTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/CameraPowerStatsTest.java
@@ -30,7 +30,6 @@ import static com.android.server.power.stats.processor.AggregatedPowerStatsConfi
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.when;
 
@@ -59,7 +58,6 @@ import org.junit.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import java.util.function.IntSupplier;
 import java.util.function.Supplier;
 
 public class CameraPowerStatsTest {
@@ -109,11 +107,6 @@ public class CameraPowerStatsTest {
                 public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
                     return mConsumedEnergyRetriever;
                 }
-
-                @Override
-                public IntSupplier getVoltageSupplier() {
-                    return () -> VOLTAGE_MV;
-                }
             };
 
     private MonotonicClock mMonotonicClock;
@@ -126,8 +119,9 @@ public class CameraPowerStatsTest {
 
     @Test
     public void energyConsumerModel() {
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(VOLTAGE_MV);
         when(mConsumedEnergyRetriever
-                .getEnergyConsumerIds(eq((int) EnergyConsumerType.CAMERA), any()))
+                .getEnergyConsumerIds(eq((int) EnergyConsumerType.CAMERA)))
                 .thenReturn(new int[]{ENERGY_CONSUMER_ID});
 
         PowerComponentAggregatedPowerStats stats = createAggregatedPowerStats(
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/CustomEnergyConsumerPowerStatsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/CustomEnergyConsumerPowerStatsTest.java
index 1c03b56614c8..42baba765276 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/CustomEnergyConsumerPowerStatsTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/CustomEnergyConsumerPowerStatsTest.java
@@ -60,7 +60,6 @@ import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
-import java.util.function.IntSupplier;
 
 public class CustomEnergyConsumerPowerStatsTest {
     @Rule(order = 0)
@@ -111,11 +110,6 @@ public class CustomEnergyConsumerPowerStatsTest {
                 public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
                     return mConsumedEnergyRetriever;
                 }
-
-                @Override
-                public IntSupplier getVoltageSupplier() {
-                    return () -> VOLTAGE_MV;
-                }
             };
 
 
@@ -243,6 +237,7 @@ public class CustomEnergyConsumerPowerStatsTest {
 
     private List<PowerStats> collectPowerStats(long timestamp, int chargeUc1, int chargeUc2,
             EnergyConsumerAttribution... attributions2) throws Exception {
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(VOLTAGE_MV);
         when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.OTHER))
                 .thenReturn(new int[]{ENERGY_CONSUMER_ID1, ENERGY_CONSUMER_ID2});
         when(mConsumedEnergyRetriever.getEnergyConsumerName(ENERGY_CONSUMER_ID1))
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/GnssPowerStatsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/GnssPowerStatsTest.java
index 2440287739bc..e6207d48d8a0 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/GnssPowerStatsTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/GnssPowerStatsTest.java
@@ -30,7 +30,6 @@ import static com.android.server.power.stats.processor.AggregatedPowerStatsConfi
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.when;
 
@@ -60,7 +59,6 @@ import org.junit.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import java.util.function.IntSupplier;
 import java.util.function.Supplier;
 
 public class GnssPowerStatsTest {
@@ -112,11 +110,6 @@ public class GnssPowerStatsTest {
                 public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
                     return mConsumedEnergyRetriever;
                 }
-
-                @Override
-                public IntSupplier getVoltageSupplier() {
-                    return () -> VOLTAGE_MV;
-                }
             };
 
     private MonotonicClock mMonotonicClock;
@@ -133,7 +126,7 @@ public class GnssPowerStatsTest {
     public void powerProfileModel() {
         // ODPM unsupported
         when(mConsumedEnergyRetriever
-                .getEnergyConsumerIds(eq((int) EnergyConsumerType.GNSS), any()))
+                .getEnergyConsumerIds(eq((int) EnergyConsumerType.GNSS)))
                 .thenReturn(new int[0]);
 
         PowerComponentAggregatedPowerStats stats = createAggregatedPowerStats(
@@ -223,7 +216,7 @@ public class GnssPowerStatsTest {
     public void initialStateGnssOn() {
         // ODPM unsupported
         when(mConsumedEnergyRetriever
-                .getEnergyConsumerIds(eq((int) EnergyConsumerType.GNSS), any()))
+                .getEnergyConsumerIds(eq((int) EnergyConsumerType.GNSS)))
                 .thenReturn(new int[0]);
 
         PowerComponentAggregatedPowerStats stats = createAggregatedPowerStats(
@@ -300,8 +293,9 @@ public class GnssPowerStatsTest {
 
     @Test
     public void energyConsumerModel() {
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(VOLTAGE_MV);
         when(mConsumedEnergyRetriever
-                .getEnergyConsumerIds(eq((int) EnergyConsumerType.GNSS), any()))
+                .getEnergyConsumerIds(eq((int) EnergyConsumerType.GNSS)))
                 .thenReturn(new int[]{ENERGY_CONSUMER_ID});
 
         PowerComponentAggregatedPowerStats stats = createAggregatedPowerStats(
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/MobileRadioPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/MobileRadioPowerStatsProcessorTest.java
index 33f407444ce0..80358c5768f1 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/MobileRadioPowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/MobileRadioPowerStatsProcessorTest.java
@@ -41,6 +41,7 @@ import static org.mockito.Mockito.when;
 import android.annotation.Nullable;
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.net.NetworkStats;
 import android.os.BatteryConsumer;
@@ -66,7 +67,6 @@ import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 import java.util.List;
-import java.util.function.IntSupplier;
 import java.util.function.LongSupplier;
 import java.util.function.Supplier;
 
@@ -134,11 +134,6 @@ public class MobileRadioPowerStatsProcessorTest {
                 }
 
                 @Override
-                public IntSupplier getVoltageSupplier() {
-                    return () -> VOLTAGE_MV;
-                }
-
-                @Override
                 public Supplier<NetworkStats> getMobileNetworkStatsSupplier() {
                     return mNetworkStatsSupplier;
                 }
@@ -407,6 +402,7 @@ public class MobileRadioPowerStatsProcessorTest {
 
     private PowerComponentAggregatedPowerStats prepareAggregatedStats_energyConsumerModel() {
         // PowerStats hardware is available
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(VOLTAGE_MV);
         when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.MOBILE_RADIO))
                 .thenReturn(new int[] {MOBILE_RADIO_ENERGY_CONSUMER_ID});
 
@@ -435,9 +431,9 @@ public class MobileRadioPowerStatsProcessorTest {
         // Initial empty ModemActivityInfo.
         mockModemActivityInfo(new ModemActivityInfo(0L, 0L, 0L, new int[5], 0L));
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(
+        when(mConsumedEnergyRetriever.getConsumedEnergy(
                 new int[]{MOBILE_RADIO_ENERGY_CONSUMER_ID}))
-                .thenReturn(new long[]{0});
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(0)});
 
         aggregatedStats.start(0);
 
@@ -466,8 +462,8 @@ public class MobileRadioPowerStatsProcessorTest {
         mStatsRule.setTime(10_000, 10_000);
 
         long energyUws = 10_000_000L * VOLTAGE_MV / 1000L;
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(
-                new int[]{MOBILE_RADIO_ENERGY_CONSUMER_ID})).thenReturn(new long[]{energyUws});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(new int[]{MOBILE_RADIO_ENERGY_CONSUMER_ID}))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(energyUws)});
 
         when(mCallDurationSupplier.getAsLong()).thenReturn(200L);
         when(mScanDurationSupplier.getAsLong()).thenReturn(5555L);
@@ -484,6 +480,12 @@ public class MobileRadioPowerStatsProcessorTest {
         return states;
     }
 
+    private EnergyConsumerResult mockEnergyConsumer(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
+    }
+
     private void mockModemActivityInfo(ModemActivityInfo emptyMai) {
         doAnswer(invocation -> {
             OutcomeReceiver<ModemActivityInfo, TelephonyManager.ModemActivityInfoException>
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/PhoneCallPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/PhoneCallPowerStatsProcessorTest.java
index 2744ceb2e7e7..535f2da603b8 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/PhoneCallPowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/PhoneCallPowerStatsProcessorTest.java
@@ -53,7 +53,6 @@ import org.junit.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import java.util.function.IntSupplier;
 import java.util.function.LongSupplier;
 import java.util.function.Supplier;
 
@@ -118,11 +117,6 @@ public class PhoneCallPowerStatsProcessorTest {
                 }
 
                 @Override
-                public IntSupplier getVoltageSupplier() {
-                    return () -> VOLTAGE_MV;
-                }
-
-                @Override
                 public Supplier<NetworkStats> getMobileNetworkStatsSupplier() {
                     return mNetworkStatsSupplier;
                 }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/ScreenPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/ScreenPowerStatsProcessorTest.java
index 412f0c8fe737..185216583f1b 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/ScreenPowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/ScreenPowerStatsProcessorTest.java
@@ -29,6 +29,7 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.when;
 
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.BatteryConsumer;
 import android.os.BatteryStats;
@@ -52,7 +53,6 @@ import org.junit.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import java.util.function.IntSupplier;
 import java.util.function.Supplier;
 
 public class ScreenPowerStatsProcessorTest {
@@ -110,11 +110,6 @@ public class ScreenPowerStatsProcessorTest {
         }
 
         @Override
-        public IntSupplier getVoltageSupplier() {
-            return () -> VOLTAGE_MV;
-        }
-
-        @Override
         public int getDisplayCount() {
             return 2;
         }
@@ -128,6 +123,7 @@ public class ScreenPowerStatsProcessorTest {
     @Before
     public void setup() {
         MockitoAnnotations.initMocks(this);
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(VOLTAGE_MV);
     }
 
     @Test
@@ -181,8 +177,8 @@ public class ScreenPowerStatsProcessorTest {
         if (energyConsumer) {
             when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.DISPLAY))
                     .thenReturn(new int[]{77});
-            when(mConsumedEnergyRetriever.getConsumedEnergyUws(new int[]{77}))
-                    .thenReturn(new long[]{10_000});
+            when(mConsumedEnergyRetriever.getConsumedEnergy(new int[]{77}))
+                    .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(10_000)});
         } else {
             when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.DISPLAY))
                     .thenReturn(new int[0]);
@@ -204,8 +200,8 @@ public class ScreenPowerStatsProcessorTest {
         if (energyConsumer) {
             // 400 mAh represented as microWattSeconds
             long energyUws = 400L * 3600 * VOLTAGE_MV;
-            when(mConsumedEnergyRetriever.getConsumedEnergyUws(new int[]{77}))
-                    .thenReturn(new long[]{10_000 + energyUws});
+            when(mConsumedEnergyRetriever.getConsumedEnergy(new int[]{77}))
+                    .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(10_000 + energyUws)});
         }
 
         when(mScreenUsageTimeRetriever.getScreenOnTimeMs(0))
@@ -262,6 +258,12 @@ public class ScreenPowerStatsProcessorTest {
         return aggregatedStats;
     }
 
+    private EnergyConsumerResult mockEnergyConsumer(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
+    }
+
     private void assertDevicePowerEstimate(
             PowerComponentAggregatedPowerStats aggregatedStats,
             int powerState, int screenState, double expectedScreenPowerEstimate,
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/WifiPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/WifiPowerStatsProcessorTest.java
index 43023f8b9ce7..baf468e3f7de 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/WifiPowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/WifiPowerStatsProcessorTest.java
@@ -41,6 +41,7 @@ import static org.mockito.Mockito.when;
 import android.annotation.Nullable;
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.net.NetworkStats;
 import android.net.wifi.WifiManager;
@@ -68,7 +69,6 @@ import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 import java.util.List;
-import java.util.function.IntSupplier;
 import java.util.function.Supplier;
 
 public class WifiPowerStatsProcessorTest {
@@ -165,11 +165,6 @@ public class WifiPowerStatsProcessorTest {
                 }
 
                 @Override
-                public IntSupplier getVoltageSupplier() {
-                    return () -> VOLTAGE_MV;
-                }
-
-                @Override
                 public Supplier<NetworkStats> getWifiNetworkStatsSupplier() {
                     return mNetworkStatsSupplier;
                 }
@@ -314,6 +309,7 @@ public class WifiPowerStatsProcessorTest {
         when(mWifiManager.isEnhancedPowerReportingSupported()).thenReturn(true);
 
         // PowerStats hardware is available
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(VOLTAGE_MV);
         when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.WIFI))
                 .thenReturn(new int[] {WIFI_ENERGY_CONSUMER_ID});
 
@@ -327,9 +323,9 @@ public class WifiPowerStatsProcessorTest {
         mockWifiActivityEnergyInfo(new WifiActivityEnergyInfo(0L,
                 WifiActivityEnergyInfo.STACK_STATE_INVALID, 0L, 0L, 0L, 0L));
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(
+        when(mConsumedEnergyRetriever.getConsumedEnergy(
                 new int[]{WIFI_ENERGY_CONSUMER_ID}))
-                .thenReturn(new long[]{0});
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(0)});
 
         aggregatedStats.start(0);
 
@@ -360,8 +356,8 @@ public class WifiPowerStatsProcessorTest {
 
         // 10 mAh represented as microWattSeconds
         long energyUws = 10 * 3600 * VOLTAGE_MV;
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(
-                new int[]{WIFI_ENERGY_CONSUMER_ID})).thenReturn(new long[]{energyUws});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(new int[]{WIFI_ENERGY_CONSUMER_ID}))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(energyUws)});
 
         aggregatedStats.addPowerStats(collector.collectStats(), 10_000);
 
@@ -552,6 +548,12 @@ public class WifiPowerStatsProcessorTest {
         return states;
     }
 
+    private EnergyConsumerResult mockEnergyConsumer(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
+    }
+
     private void mockWifiActivityEnergyInfo(WifiActivityEnergyInfo waei) {
         doAnswer(invocation -> {
             WifiManager.OnWifiActivityEnergyInfoListener