Merge changes from topic "wifi-power-stats" into main

* changes:
  Enable PowerStatsExporter for WiFi
  Introduce PowerStatsProcessor for WiFi
  Introduce PowerStatsCollector for WiFi
diff --git a/core/java/android/os/connectivity/WifiActivityEnergyInfo.java b/core/java/android/os/connectivity/WifiActivityEnergyInfo.java
index ad74a9f..2ceb1c7 100644
--- a/core/java/android/os/connectivity/WifiActivityEnergyInfo.java
+++ b/core/java/android/os/connectivity/WifiActivityEnergyInfo.java
@@ -37,8 +37,10 @@
  * real-time.
  * @hide
  */
+@android.ravenwood.annotation.RavenwoodKeepWholeClass
 @SystemApi
 public final class WifiActivityEnergyInfo implements Parcelable {
+    private static final long DEFERRED_ENERGY_ESTIMATE = -1;
     @ElapsedRealtimeLong
     private final long mTimeSinceBootMillis;
     @StackState
@@ -52,7 +54,7 @@
     @IntRange(from = 0)
     private final long mControllerIdleDurationMillis;
     @IntRange(from = 0)
-    private final long mControllerEnergyUsedMicroJoules;
+    private long mControllerEnergyUsedMicroJoules;
 
     /** @hide */
     @Retention(RetentionPolicy.SOURCE)
@@ -99,9 +101,10 @@
                 rxDurationMillis,
                 scanDurationMillis,
                 idleDurationMillis,
-                calculateEnergyMicroJoules(txDurationMillis, rxDurationMillis, idleDurationMillis));
+                DEFERRED_ENERGY_ESTIMATE);
     }
 
+    @android.ravenwood.annotation.RavenwoodReplace
     private static long calculateEnergyMicroJoules(
             long txDurationMillis, long rxDurationMillis, long idleDurationMillis) {
         final Context context = ActivityThread.currentActivityThread().getSystemContext();
@@ -125,6 +128,11 @@
                 * voltage);
     }
 
+    private static long calculateEnergyMicroJoules$ravenwood(long txDurationMillis,
+            long rxDurationMillis, long idleDurationMillis) {
+        return 0;
+    }
+
     /** @hide */
     public WifiActivityEnergyInfo(
             @ElapsedRealtimeLong long timeSinceBootMillis,
@@ -152,7 +160,7 @@
                 + " mControllerRxDurationMillis=" + mControllerRxDurationMillis
                 + " mControllerScanDurationMillis=" + mControllerScanDurationMillis
                 + " mControllerIdleDurationMillis=" + mControllerIdleDurationMillis
-                + " mControllerEnergyUsedMicroJoules=" + mControllerEnergyUsedMicroJoules
+                + " mControllerEnergyUsedMicroJoules=" + getControllerEnergyUsedMicroJoules()
                 + " }";
     }
 
@@ -231,6 +239,11 @@
     /** Get the energy consumed by Wifi, in microjoules. */
     @IntRange(from = 0)
     public long getControllerEnergyUsedMicroJoules() {
+        if (mControllerEnergyUsedMicroJoules == DEFERRED_ENERGY_ESTIMATE) {
+            mControllerEnergyUsedMicroJoules = calculateEnergyMicroJoules(
+                    mControllerTxDurationMillis, mControllerRxDurationMillis,
+                    mControllerIdleDurationMillis);
+        }
         return mControllerEnergyUsedMicroJoules;
     }
 
diff --git a/core/res/res/values/config_battery_stats.xml b/core/res/res/values/config_battery_stats.xml
index ae47899..8d97362 100644
--- a/core/res/res/values/config_battery_stats.xml
+++ b/core/res/res/values/config_battery_stats.xml
@@ -35,6 +35,9 @@
     <!-- Mobile Radio power stats collection throttle period in milliseconds. -->
     <integer name="config_defaultPowerStatsThrottlePeriodMobileRadio">3600000</integer>
 
+    <!-- Mobile Radio power stats collection throttle period in milliseconds. -->
+    <integer name="config_defaultPowerStatsThrottlePeriodWifi">3600000</integer>
+
     <!-- PowerStats aggregation period in milliseconds. This is the interval at which the power
     stats aggregation procedure is performed and the results stored in PowerStatsStore. -->
     <integer name="config_powerStatsAggregationPeriod">14400000</integer>
diff --git a/core/res/res/values/symbols.xml b/core/res/res/values/symbols.xml
index fbc8811..624026c 100644
--- a/core/res/res/values/symbols.xml
+++ b/core/res/res/values/symbols.xml
@@ -5218,6 +5218,7 @@
   <java-symbol type="bool" name="config_batteryStatsResetOnUnplugAfterSignificantCharge" />
   <java-symbol type="integer" name="config_defaultPowerStatsThrottlePeriodCpu" />
   <java-symbol type="integer" name="config_defaultPowerStatsThrottlePeriodMobileRadio" />
+  <java-symbol type="integer" name="config_defaultPowerStatsThrottlePeriodWifi" />
   <java-symbol type="integer" name="config_powerStatsAggregationPeriod" />
   <java-symbol type="integer" name="config_aggregatedPowerStatsSpanDuration" />
 
diff --git a/ravenwood/texts/ravenwood-annotation-allowed-classes.txt b/ravenwood/texts/ravenwood-annotation-allowed-classes.txt
index 9b4d378..243e224 100644
--- a/ravenwood/texts/ravenwood-annotation-allowed-classes.txt
+++ b/ravenwood/texts/ravenwood-annotation-allowed-classes.txt
@@ -266,6 +266,8 @@
 android.telephony.ModemActivityInfo
 android.telephony.ServiceState
 
+android.os.connectivity.WifiActivityEnergyInfo
+
 com.android.server.LocalServices
 
 com.android.internal.util.BitUtils
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index f98799d..4f84149 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -133,6 +133,7 @@
 import com.android.server.power.stats.PowerStatsStore;
 import com.android.server.power.stats.PowerStatsUidResolver;
 import com.android.server.power.stats.SystemServerCpuThreadReader.SystemServiceCpuThreadTimes;
+import com.android.server.power.stats.WifiPowerStatsProcessor;
 import com.android.server.power.stats.wakeups.CpuWakeupStats;
 
 import java.io.File;
@@ -412,6 +413,8 @@
                 com.android.internal.R.integer.config_defaultPowerStatsThrottlePeriodCpu);
         final long powerStatsThrottlePeriodMobileRadio = context.getResources().getInteger(
                 com.android.internal.R.integer.config_defaultPowerStatsThrottlePeriodMobileRadio);
+        final long powerStatsThrottlePeriodWifi = context.getResources().getInteger(
+                com.android.internal.R.integer.config_defaultPowerStatsThrottlePeriodWifi);
         mBatteryStatsConfig =
                 new BatteryStatsImpl.BatteryStatsConfig.Builder()
                         .setResetOnUnplugHighBatteryLevel(resetOnUnplugHighBatteryLevel)
@@ -422,6 +425,9 @@
                         .setPowerStatsThrottlePeriodMillis(
                                 BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
                                 powerStatsThrottlePeriodMobileRadio)
+                        .setPowerStatsThrottlePeriodMillis(
+                                BatteryConsumer.POWER_COMPONENT_WIFI,
+                                powerStatsThrottlePeriodWifi)
                         .build();
         mPowerStatsUidResolver = new PowerStatsUidResolver();
         mStats = new BatteryStatsImpl(mBatteryStatsConfig, Clock.SYSTEM_CLOCK, mMonotonicClock,
@@ -480,6 +486,7 @@
                         AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
                 .setProcessor(
                         new CpuPowerStatsProcessor(mPowerProfile, mCpuScalingPolicies));
+
         config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
                 .trackDeviceStates(
                         AggregatedPowerStatsConfig.STATE_POWER,
@@ -490,9 +497,21 @@
                         AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
                 .setProcessor(
                         new MobileRadioPowerStatsProcessor(mPowerProfile));
+
         config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_PHONE,
                         BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
                 .setProcessor(new PhoneCallPowerStatsProcessor());
+
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_WIFI)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
+                .setProcessor(
+                        new WifiPowerStatsProcessor(mPowerProfile));
         return config;
     }
 
@@ -518,14 +537,22 @@
     public void systemServicesReady() {
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_CPU,
                 Flags.streamlinedBatteryStats());
-        mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
-                Flags.streamlinedConnectivityBatteryStats());
         mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
                 BatteryConsumer.POWER_COMPONENT_CPU,
                 Flags.streamlinedBatteryStats());
+
+        mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
+                Flags.streamlinedConnectivityBatteryStats());
         mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
                 BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
                 Flags.streamlinedConnectivityBatteryStats());
+
+        mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_WIFI,
+                Flags.streamlinedConnectivityBatteryStats());
+        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+                BatteryConsumer.POWER_COMPONENT_WIFI,
+                Flags.streamlinedConnectivityBatteryStats());
+
         mWorker.systemServicesReady();
         mStats.systemServicesReady(mContext);
         mCpuWakeupStats.systemServicesReady();
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 54cb9c9..49c4000 100644
--- a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
+++ b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
@@ -19,6 +19,7 @@
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
 import static android.os.BatteryStats.Uid.NUM_PROCESS_STATE;
+import static android.os.BatteryStats.Uid.NUM_WIFI_BATCHED_SCAN_BINS;
 import static android.os.BatteryStatsManager.NUM_WIFI_STATES;
 import static android.os.BatteryStatsManager.NUM_WIFI_SUPPL_STATES;
 
@@ -292,7 +293,25 @@
     private int[] mCpuPowerBracketMap;
     private final CpuPowerStatsCollector mCpuPowerStatsCollector;
     private final MobileRadioPowerStatsCollector mMobileRadioPowerStatsCollector;
+    private final WifiPowerStatsCollector mWifiPowerStatsCollector;
     private final SparseBooleanArray mPowerStatsCollectorEnabled = new SparseBooleanArray();
+    private final WifiPowerStatsCollector.WifiStatsRetriever mWifiStatsRetriever =
+            new WifiPowerStatsCollector.WifiStatsRetriever() {
+                @Override
+                public void retrieveWifiScanTimes(Callback callback) {
+                    synchronized (BatteryStatsImpl.this) {
+                        retrieveWifiScanTimesLocked(callback);
+                    }
+                }
+
+                @Override
+                public long getWifiActiveDuration() {
+                    synchronized (BatteryStatsImpl.this) {
+                        return getGlobalWifiRunningTime(mClock.elapsedRealtime() * 1000,
+                                STATS_SINCE_CHARGED) / 1000;
+                    }
+                }
+            };
 
     public LongSparseArray<SamplingTimer> getKernelMemoryStats() {
         return mKernelMemoryStats;
@@ -501,6 +520,8 @@
                         TimeUnit.MINUTES.toMillis(1));
                 setPowerStatsThrottlePeriodMillis(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
                         TimeUnit.HOURS.toMillis(1));
+                setPowerStatsThrottlePeriodMillis(BatteryConsumer.POWER_COMPONENT_WIFI,
+                        TimeUnit.HOURS.toMillis(1));
             }
 
             /**
@@ -1885,11 +1906,12 @@
     }
 
     private class PowerStatsCollectorInjector implements CpuPowerStatsCollector.Injector,
-            MobileRadioPowerStatsCollector.Injector {
+            MobileRadioPowerStatsCollector.Injector, WifiPowerStatsCollector.Injector {
         private PackageManager mPackageManager;
         private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
         private NetworkStatsManager mNetworkStatsManager;
         private TelephonyManager mTelephonyManager;
+        private WifiManager mWifiManager;
 
         void setContext(Context context) {
             mPackageManager = context.getPackageManager();
@@ -1897,6 +1919,7 @@
                     LocalServices.getService(PowerStatsInternal.class));
             mNetworkStatsManager = context.getSystemService(NetworkStatsManager.class);
             mTelephonyManager = context.getSystemService(TelephonyManager.class);
+            mWifiManager = context.getSystemService(WifiManager.class);
         }
 
         @Override
@@ -1950,11 +1973,26 @@
         }
 
         @Override
+        public Supplier<NetworkStats> getWifiNetworkStatsSupplier() {
+            return () -> readWifiNetworkStatsLocked(mNetworkStatsManager);
+        }
+
+        @Override
+        public WifiPowerStatsCollector.WifiStatsRetriever getWifiStatsRetriever() {
+            return mWifiStatsRetriever;
+        }
+
+        @Override
         public TelephonyManager getTelephonyManager() {
             return mTelephonyManager;
         }
 
         @Override
+        public WifiManager getWifiManager() {
+            return mWifiManager;
+        }
+
+        @Override
         public LongSupplier getCallDurationSupplier() {
             return () -> mPhoneOnTimer.getTotalTimeLocked(mClock.elapsedRealtime() * 1000,
                     STATS_SINCE_CHARGED);
@@ -6354,7 +6392,11 @@
                     HistoryItem.STATE2_WIFI_ON_FLAG);
             mWifiOn = true;
             mWifiOnTimer.startRunningLocked(elapsedRealtimeMs);
-            scheduleSyncExternalStatsLocked("wifi-off", ExternalStatsSync.UPDATE_WIFI);
+            if (mWifiPowerStatsCollector.isEnabled()) {
+                mWifiPowerStatsCollector.schedule();
+            } else {
+                scheduleSyncExternalStatsLocked("wifi-off", ExternalStatsSync.UPDATE_WIFI);
+            }
         }
     }
 
@@ -6365,7 +6407,11 @@
                     HistoryItem.STATE2_WIFI_ON_FLAG);
             mWifiOn = false;
             mWifiOnTimer.stopRunningLocked(elapsedRealtimeMs);
-            scheduleSyncExternalStatsLocked("wifi-on", ExternalStatsSync.UPDATE_WIFI);
+            if (mWifiPowerStatsCollector.isEnabled()) {
+                mWifiPowerStatsCollector.schedule();
+            } else {
+                scheduleSyncExternalStatsLocked("wifi-on", ExternalStatsSync.UPDATE_WIFI);
+            }
         }
     }
 
@@ -6757,8 +6803,11 @@
                             .noteWifiRunningLocked(elapsedRealtimeMs);
                 }
             }
-
-            scheduleSyncExternalStatsLocked("wifi-running", ExternalStatsSync.UPDATE_WIFI);
+            if (mWifiPowerStatsCollector.isEnabled()) {
+                mWifiPowerStatsCollector.schedule();
+            } else {
+                scheduleSyncExternalStatsLocked("wifi-running", ExternalStatsSync.UPDATE_WIFI);
+            }
         } else {
             Log.w(TAG, "noteWifiRunningLocked -- called while WIFI running");
         }
@@ -6827,7 +6876,11 @@
                 }
             }
 
-            scheduleSyncExternalStatsLocked("wifi-stopped", ExternalStatsSync.UPDATE_WIFI);
+            if (mWifiPowerStatsCollector.isEnabled()) {
+                mWifiPowerStatsCollector.schedule();
+            } else {
+                scheduleSyncExternalStatsLocked("wifi-stopped", ExternalStatsSync.UPDATE_WIFI);
+            }
         } else {
             Log.w(TAG, "noteWifiStoppedLocked -- called while WIFI not running");
         }
@@ -6842,7 +6895,11 @@
             }
             mWifiState = wifiState;
             mWifiStateTimer[wifiState].startRunningLocked(elapsedRealtimeMs);
-            scheduleSyncExternalStatsLocked("wifi-state", ExternalStatsSync.UPDATE_WIFI);
+            if (mWifiPowerStatsCollector.isEnabled()) {
+                mWifiPowerStatsCollector.schedule();
+            } else {
+                scheduleSyncExternalStatsLocked("wifi-state", ExternalStatsSync.UPDATE_WIFI);
+            }
         }
     }
 
@@ -6965,6 +7022,25 @@
                 .noteWifiBatchedScanStoppedLocked(elapsedRealtimeMs);
     }
 
+    private void retrieveWifiScanTimesLocked(
+            WifiPowerStatsCollector.WifiStatsRetriever.Callback callback) {
+        long elapsedTimeUs = mClock.elapsedRealtime() * 1000;
+        for (int i = mUidStats.size() - 1; i >= 0; i--) {
+            int uid = mUidStats.keyAt(i);
+            Uid uidStats = mUidStats.valueAt(i);
+            long scanTimeUs = uidStats.getWifiScanTime(elapsedTimeUs, STATS_SINCE_CHARGED);
+            long batchScanTimeUs = 0;
+            for (int bucket = 0; bucket < NUM_WIFI_BATCHED_SCAN_BINS; bucket++) {
+                batchScanTimeUs += uidStats.getWifiBatchedScanTime(bucket, elapsedTimeUs,
+                        STATS_SINCE_CHARGED);
+            }
+            if (scanTimeUs != 0 || batchScanTimeUs != 0) {
+                callback.onWifiScanTime(uid, (scanTimeUs + 500) / 1000,
+                        (batchScanTimeUs + 500) / 1000);
+            }
+        }
+    }
+
     private int mWifiMulticastNesting = 0;
 
     @GuardedBy("this")
@@ -11101,6 +11177,11 @@
                 BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO));
         mMobileRadioPowerStatsCollector.addConsumer(this::recordPowerStats);
 
+        mWifiPowerStatsCollector = new WifiPowerStatsCollector(
+                mPowerStatsCollectorInjector, mBatteryStatsConfig.getPowerStatsThrottlePeriod(
+                BatteryConsumer.POWER_COMPONENT_WIFI));
+        mWifiPowerStatsCollector.addConsumer(this::recordPowerStats);
+
         mStartCount++;
         initTimersAndCounters();
         mOnBattery = mOnBatteryInternal = false;
@@ -12095,10 +12176,10 @@
                 }
             }
             if (lastEntry != null) {
-                delta.mRxBytes = entry.getRxBytes() - lastEntry.getRxBytes();
-                delta.mRxPackets = entry.getRxPackets() - lastEntry.getRxPackets();
-                delta.mTxBytes = entry.getTxBytes() - lastEntry.getTxBytes();
-                delta.mTxPackets = entry.getTxPackets() - lastEntry.getTxPackets();
+                delta.mRxBytes = Math.max(0, entry.getRxBytes() - lastEntry.getRxBytes());
+                delta.mRxPackets = Math.max(0, entry.getRxPackets() - lastEntry.getRxPackets());
+                delta.mTxBytes = Math.max(0, entry.getTxBytes() - lastEntry.getTxBytes());
+                delta.mTxPackets = Math.max(0, entry.getTxPackets() - lastEntry.getTxPackets());
             } else {
                 delta.mRxBytes = entry.getRxBytes();
                 delta.mRxPackets = entry.getRxPackets();
@@ -12119,6 +12200,10 @@
     public void updateWifiState(@Nullable final WifiActivityEnergyInfo info,
             final long consumedChargeUC, long elapsedRealtimeMs, long uptimeMs,
             @NonNull NetworkStatsManager networkStatsManager) {
+        if (mWifiPowerStatsCollector.isEnabled()) {
+            return;
+        }
+
         if (DEBUG_ENERGY) {
             synchronized (mWifiNetworkLock) {
                 Slog.d(TAG, "Updating wifi stats: " + Arrays.toString(mWifiIfaces));
@@ -14507,6 +14592,10 @@
                 mPowerStatsCollectorEnabled.get(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO));
         mMobileRadioPowerStatsCollector.schedule();
 
+        mWifiPowerStatsCollector.setEnabled(
+                mPowerStatsCollectorEnabled.get(BatteryConsumer.POWER_COMPONENT_WIFI));
+        mWifiPowerStatsCollector.schedule();
+
         mSystemReady = true;
     }
 
@@ -14521,6 +14610,8 @@
                 return mCpuPowerStatsCollector;
             case BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO:
                 return mMobileRadioPowerStatsCollector;
+            case BatteryConsumer.POWER_COMPONENT_WIFI:
+                return mWifiPowerStatsCollector;
         }
         return null;
     }
@@ -16056,6 +16147,7 @@
     public void schedulePowerStatsSampleCollection() {
         mCpuPowerStatsCollector.forceSchedule();
         mMobileRadioPowerStatsCollector.forceSchedule();
+        mWifiPowerStatsCollector.forceSchedule();
     }
 
     /**
@@ -16074,6 +16166,7 @@
     public void dumpStatsSample(PrintWriter pw) {
         mCpuPowerStatsCollector.collectAndDump(pw);
         mMobileRadioPowerStatsCollector.collectAndDump(pw);
+        mWifiPowerStatsCollector.collectAndDump(pw);
     }
 
     private final Runnable mWriteAsyncRunnable = () -> {
diff --git a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
index 97f0986..0d5eabc 100644
--- a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
+++ b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
@@ -87,7 +87,9 @@
                         mPowerCalculators.add(new PhonePowerCalculator(mPowerProfile));
                     }
                 }
-                mPowerCalculators.add(new WifiPowerCalculator(mPowerProfile));
+                if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_WIFI)) {
+                    mPowerCalculators.add(new WifiPowerCalculator(mPowerProfile));
+                }
                 mPowerCalculators.add(new BluetoothPowerCalculator(mPowerProfile));
                 mPowerCalculators.add(new SensorPowerCalculator(
                         mContext.getSystemService(SensorManager.class)));
diff --git a/services/core/java/com/android/server/power/stats/WifiPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/WifiPowerStatsCollector.java
new file mode 100644
index 0000000..6321053
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/WifiPowerStatsCollector.java
@@ -0,0 +1,328 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.server.power.stats;
+
+import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerType;
+import android.net.NetworkStats;
+import android.net.wifi.WifiManager;
+import android.os.BatteryConsumer;
+import android.os.Handler;
+import android.os.PersistableBundle;
+import android.os.connectivity.WifiActivityEnergyInfo;
+import android.util.Slog;
+import android.util.SparseArray;
+
+import com.android.internal.os.Clock;
+import com.android.internal.os.PowerStats;
+
+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 {
+    private static final String TAG = "WifiPowerStatsCollector";
+
+    private static final long WIFI_ACTIVITY_REQUEST_TIMEOUT = 20000;
+
+    private static final long ENERGY_UNSPECIFIED = -1;
+
+    interface WifiStatsRetriever {
+        interface Callback {
+            void onWifiScanTime(int uid, long scanTimeMs, long batchScanTimeMs);
+        }
+
+        void retrieveWifiScanTimes(Callback callback);
+        long getWifiActiveDuration();
+    }
+
+    interface Injector {
+        Handler getHandler();
+        Clock getClock();
+        PowerStatsUidResolver getUidResolver();
+        PackageManager getPackageManager();
+        ConsumedEnergyRetriever getConsumedEnergyRetriever();
+        IntSupplier getVoltageSupplier();
+        Supplier<NetworkStats> getWifiNetworkStatsSupplier();
+        WifiManager getWifiManager();
+        WifiStatsRetriever getWifiStatsRetriever();
+    }
+
+    private final Injector mInjector;
+
+    private WifiPowerStatsLayout mLayout;
+    private boolean mIsInitialized;
+    private boolean mPowerReportingSupported;
+
+    private PowerStats mPowerStats;
+    private long[] mDeviceStats;
+    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 WifiActivityEnergyInfo mLastWifiActivityInfo =
+            new WifiActivityEnergyInfo(0, 0, 0, 0, 0, 0);
+    private NetworkStats mLastNetworkStats;
+    private long[] mLastConsumedEnergyUws;
+    private int mLastVoltageMv;
+
+    private static class WifiScanTimes {
+        public long basicScanTimeMs;
+        public long batchedScanTimeMs;
+    }
+    private final WifiScanTimes mScanTimes = new WifiScanTimes();
+    private final SparseArray<WifiScanTimes> mLastScanTimes = new SparseArray<>();
+    private long mLastWifiActiveDuration;
+
+    public WifiPowerStatsCollector(Injector injector, long throttlePeriodMs) {
+        super(injector.getHandler(), throttlePeriodMs, injector.getUidResolver(),
+                injector.getClock());
+        mInjector = injector;
+    }
+
+    @Override
+    public void setEnabled(boolean enabled) {
+        if (enabled) {
+            PackageManager packageManager = mInjector.getPackageManager();
+            super.setEnabled(packageManager != null
+                    && packageManager.hasSystemFeature(PackageManager.FEATURE_WIFI));
+        } else {
+            super.setEnabled(false);
+        }
+    }
+
+    private boolean ensureInitialized() {
+        if (mIsInitialized) {
+            return true;
+        }
+
+        if (!isEnabled()) {
+            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);
+
+        mLayout = new WifiPowerStatsLayout();
+        mLayout.addDeviceWifiActivity(mPowerReportingSupported);
+        mLayout.addDeviceSectionEnergyConsumers(mEnergyConsumerIds.length);
+        mLayout.addUidNetworkStats();
+        mLayout.addDeviceSectionUsageDuration();
+        mLayout.addDeviceSectionPowerEstimate();
+        mLayout.addUidSectionPowerEstimate();
+
+        PersistableBundle extras = new PersistableBundle();
+        mLayout.toExtras(extras);
+        PowerStats.Descriptor powerStatsDescriptor = new PowerStats.Descriptor(
+                BatteryConsumer.POWER_COMPONENT_WIFI, mLayout.getDeviceStatsArrayLength(),
+                null, 0, mLayout.getUidStatsArrayLength(),
+                extras);
+        mPowerStats = new PowerStats(powerStatsDescriptor);
+        mDeviceStats = mPowerStats.stats;
+
+        mIsInitialized = true;
+        return true;
+    }
+
+    @Override
+    protected PowerStats collectStats() {
+        if (!ensureInitialized()) {
+            return null;
+        }
+
+        if (mPowerReportingSupported) {
+            collectWifiActivityInfo();
+        } else {
+            collectWifiActivityStats();
+        }
+        collectNetworkStats();
+        collectWifiScanTime();
+
+        if (mEnergyConsumerIds.length != 0) {
+            collectEnergyConsumers();
+        }
+
+        return mPowerStats;
+    }
+
+    private void collectWifiActivityInfo() {
+        CompletableFuture<WifiActivityEnergyInfo> immediateFuture = new CompletableFuture<>();
+        mWifiManager.getWifiActivityEnergyInfoAsync(Runnable::run,
+                immediateFuture::complete);
+
+        WifiActivityEnergyInfo activityInfo;
+        try {
+            activityInfo = immediateFuture.get(WIFI_ACTIVITY_REQUEST_TIMEOUT,
+                    TimeUnit.MILLISECONDS);
+        } catch (Exception e) {
+            Slog.e(TAG, "Cannot acquire WifiActivityEnergyInfo", e);
+            activityInfo = null;
+        }
+
+        if (activityInfo == null) {
+            return;
+        }
+
+        long rxDuration = activityInfo.getControllerRxDurationMillis()
+                - mLastWifiActivityInfo.getControllerRxDurationMillis();
+        long txDuration = activityInfo.getControllerTxDurationMillis()
+                - mLastWifiActivityInfo.getControllerTxDurationMillis();
+        long scanDuration = activityInfo.getControllerScanDurationMillis()
+                - mLastWifiActivityInfo.getControllerScanDurationMillis();
+        long idleDuration = activityInfo.getControllerIdleDurationMillis()
+                - mLastWifiActivityInfo.getControllerIdleDurationMillis();
+
+        mLayout.setDeviceRxTime(mDeviceStats, rxDuration);
+        mLayout.setDeviceTxTime(mDeviceStats, txDuration);
+        mLayout.setDeviceScanTime(mDeviceStats, scanDuration);
+        mLayout.setDeviceIdleTime(mDeviceStats, idleDuration);
+
+        mPowerStats.durationMs = rxDuration + txDuration + scanDuration + idleDuration;
+
+        mLastWifiActivityInfo = activityInfo;
+    }
+
+    private void collectWifiActivityStats() {
+        long duration = mWifiStatsRetriever.getWifiActiveDuration();
+        mLayout.setDeviceActiveTime(mDeviceStats, Math.max(0, duration - mLastWifiActiveDuration));
+        mLastWifiActiveDuration = duration;
+        mPowerStats.durationMs = duration;
+    }
+
+    private void collectNetworkStats() {
+        mPowerStats.uidStats.clear();
+
+        NetworkStats networkStats = mNetworkStatsSupplier.get();
+        if (networkStats == null) {
+            return;
+        }
+
+        List<BatteryStatsImpl.NetworkStatsDelta> delta =
+                BatteryStatsImpl.computeDelta(networkStats, mLastNetworkStats);
+        mLastNetworkStats = networkStats;
+        for (int i = delta.size() - 1; i >= 0; i--) {
+            BatteryStatsImpl.NetworkStatsDelta uidDelta = delta.get(i);
+            long rxBytes = uidDelta.getRxBytes();
+            long txBytes = uidDelta.getTxBytes();
+            long rxPackets = uidDelta.getRxPackets();
+            long txPackets = uidDelta.getTxPackets();
+            if (rxBytes == 0 && txBytes == 0 && rxPackets == 0 && txPackets == 0) {
+                continue;
+            }
+
+            int uid = mUidResolver.mapUid(uidDelta.getUid());
+            long[] stats = mPowerStats.uidStats.get(uid);
+            if (stats == null) {
+                stats = new long[mLayout.getUidStatsArrayLength()];
+                mPowerStats.uidStats.put(uid, stats);
+                mLayout.setUidRxBytes(stats, rxBytes);
+                mLayout.setUidTxBytes(stats, txBytes);
+                mLayout.setUidRxPackets(stats, rxPackets);
+                mLayout.setUidTxPackets(stats, txPackets);
+            } else {
+                mLayout.setUidRxBytes(stats, mLayout.getUidRxBytes(stats) + rxBytes);
+                mLayout.setUidTxBytes(stats, mLayout.getUidTxBytes(stats) + txBytes);
+                mLayout.setUidRxPackets(stats, mLayout.getUidRxPackets(stats) + rxPackets);
+                mLayout.setUidTxPackets(stats, mLayout.getUidTxPackets(stats) + txPackets);
+            }
+        }
+    }
+
+    private void collectWifiScanTime() {
+        mScanTimes.basicScanTimeMs = 0;
+        mScanTimes.batchedScanTimeMs = 0;
+        mWifiStatsRetriever.retrieveWifiScanTimes((uid, scanTimeMs, batchScanTimeMs) -> {
+            WifiScanTimes lastScanTimes = mLastScanTimes.get(uid);
+            if (lastScanTimes == null) {
+                lastScanTimes = new WifiScanTimes();
+                mLastScanTimes.put(uid, lastScanTimes);
+            }
+
+            long scanTimeDelta = Math.max(0, scanTimeMs - lastScanTimes.basicScanTimeMs);
+            long batchScanTimeDelta = Math.max(0,
+                    batchScanTimeMs - lastScanTimes.batchedScanTimeMs);
+            if (scanTimeDelta != 0 || batchScanTimeDelta != 0) {
+                mScanTimes.basicScanTimeMs += scanTimeDelta;
+                mScanTimes.batchedScanTimeMs += batchScanTimeDelta;
+                uid = mUidResolver.mapUid(uid);
+                long[] stats = mPowerStats.uidStats.get(uid);
+                if (stats == null) {
+                    stats = new long[mLayout.getUidStatsArrayLength()];
+                    mPowerStats.uidStats.put(uid, stats);
+                    mLayout.setUidScanTime(stats, scanTimeDelta);
+                    mLayout.setUidBatchScanTime(stats, batchScanTimeDelta);
+                } else {
+                    mLayout.setUidScanTime(stats, mLayout.getUidScanTime(stats) + scanTimeDelta);
+                    mLayout.setUidBatchScanTime(stats,
+                            mLayout.getUidBatchedScanTime(stats) + batchScanTimeDelta);
+                }
+            }
+            lastScanTimes.basicScanTimeMs = scanTimeMs;
+            lastScanTimes.batchedScanTimeMs = batchScanTimeMs;
+        });
+
+        mLayout.setDeviceBasicScanTime(mDeviceStats, mScanTimes.basicScanTimeMs);
+        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);
+        mLastScanTimes.remove(uid);
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/WifiPowerStatsLayout.java b/services/core/java/com/android/server/power/stats/WifiPowerStatsLayout.java
new file mode 100644
index 0000000..0fa6ec6
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/WifiPowerStatsLayout.java
@@ -0,0 +1,241 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.power.stats;
+
+import android.annotation.NonNull;
+import android.os.PersistableBundle;
+
+import com.android.internal.os.PowerStats;
+
+public class WifiPowerStatsLayout extends PowerStatsLayout {
+    private static final String TAG = "WifiPowerStatsLayout";
+    private static final int UNSPECIFIED = -1;
+    private static final String EXTRA_POWER_REPORTING_SUPPORTED = "prs";
+    private static final String EXTRA_DEVICE_RX_TIME_POSITION = "dt-rx";
+    private static final String EXTRA_DEVICE_TX_TIME_POSITION = "dt-tx";
+    private static final String EXTRA_DEVICE_SCAN_TIME_POSITION = "dt-scan";
+    private static final String EXTRA_DEVICE_BASIC_SCAN_TIME_POSITION = "dt-basic-scan";
+    private static final String EXTRA_DEVICE_BATCHED_SCAN_TIME_POSITION = "dt-batch-scan";
+    private static final String EXTRA_DEVICE_IDLE_TIME_POSITION = "dt-idle";
+    private static final String EXTRA_DEVICE_ACTIVE_TIME_POSITION = "dt-on";
+    private static final String EXTRA_UID_RX_BYTES_POSITION = "urxb";
+    private static final String EXTRA_UID_TX_BYTES_POSITION = "utxb";
+    private static final String EXTRA_UID_RX_PACKETS_POSITION = "urxp";
+    private static final String EXTRA_UID_TX_PACKETS_POSITION = "utxp";
+    private static final String EXTRA_UID_SCAN_TIME_POSITION = "ut-scan";
+    private static final String EXTRA_UID_BATCH_SCAN_TIME_POSITION = "ut-bscan";
+
+    private boolean mPowerReportingSupported;
+    private int mDeviceRxTimePosition;
+    private int mDeviceTxTimePosition;
+    private int mDeviceIdleTimePosition;
+    private int mDeviceScanTimePosition;
+    private int mDeviceBasicScanTimePosition;
+    private int mDeviceBatchedScanTimePosition;
+    private int mDeviceActiveTimePosition;
+    private int mUidRxBytesPosition;
+    private int mUidTxBytesPosition;
+    private int mUidRxPacketsPosition;
+    private int mUidTxPacketsPosition;
+    private int mUidScanTimePosition;
+    private int mUidBatchScanTimePosition;
+
+    WifiPowerStatsLayout() {
+    }
+
+    WifiPowerStatsLayout(@NonNull PowerStats.Descriptor descriptor) {
+        super(descriptor);
+    }
+
+    void addDeviceWifiActivity(boolean powerReportingSupported) {
+        mPowerReportingSupported = powerReportingSupported;
+        if (mPowerReportingSupported) {
+            mDeviceActiveTimePosition = UNSPECIFIED;
+            mDeviceRxTimePosition = addDeviceSection(1);
+            mDeviceTxTimePosition = addDeviceSection(1);
+            mDeviceIdleTimePosition = addDeviceSection(1);
+            mDeviceScanTimePosition = addDeviceSection(1);
+        } else {
+            mDeviceActiveTimePosition = addDeviceSection(1);
+            mDeviceRxTimePosition = UNSPECIFIED;
+            mDeviceTxTimePosition = UNSPECIFIED;
+            mDeviceIdleTimePosition = UNSPECIFIED;
+            mDeviceScanTimePosition = UNSPECIFIED;
+        }
+        mDeviceBasicScanTimePosition = addDeviceSection(1);
+        mDeviceBatchedScanTimePosition = addDeviceSection(1);
+    }
+
+    void addUidNetworkStats() {
+        mUidRxBytesPosition = addUidSection(1);
+        mUidTxBytesPosition = addUidSection(1);
+        mUidRxPacketsPosition = addUidSection(1);
+        mUidTxPacketsPosition = addUidSection(1);
+        mUidScanTimePosition = addUidSection(1);
+        mUidBatchScanTimePosition = addUidSection(1);
+    }
+
+    public boolean isPowerReportingSupported() {
+        return mPowerReportingSupported;
+    }
+
+    public void setDeviceRxTime(long[] stats, long durationMillis) {
+        stats[mDeviceRxTimePosition] = durationMillis;
+    }
+
+    public long getDeviceRxTime(long[] stats) {
+        return stats[mDeviceRxTimePosition];
+    }
+
+    public void setDeviceTxTime(long[] stats, long durationMillis) {
+        stats[mDeviceTxTimePosition] = durationMillis;
+    }
+
+    public long getDeviceTxTime(long[] stats) {
+        return stats[mDeviceTxTimePosition];
+    }
+
+    public void setDeviceScanTime(long[] stats, long durationMillis) {
+        stats[mDeviceScanTimePosition] = durationMillis;
+    }
+
+    public long getDeviceScanTime(long[] stats) {
+        return stats[mDeviceScanTimePosition];
+    }
+
+    public void setDeviceBasicScanTime(long[] stats, long durationMillis) {
+        stats[mDeviceBasicScanTimePosition] = durationMillis;
+    }
+
+    public long getDeviceBasicScanTime(long[] stats) {
+        return stats[mDeviceBasicScanTimePosition];
+    }
+
+    public void setDeviceBatchedScanTime(long[] stats, long durationMillis) {
+        stats[mDeviceBatchedScanTimePosition] = durationMillis;
+    }
+
+    public long getDeviceBatchedScanTime(long[] stats) {
+        return stats[mDeviceBatchedScanTimePosition];
+    }
+
+    public void setDeviceIdleTime(long[] stats, long durationMillis) {
+        stats[mDeviceIdleTimePosition] = durationMillis;
+    }
+
+    public long getDeviceIdleTime(long[] stats) {
+        return stats[mDeviceIdleTimePosition];
+    }
+
+    public void setDeviceActiveTime(long[] stats, long durationMillis) {
+        stats[mDeviceActiveTimePosition] = durationMillis;
+    }
+
+    public long getDeviceActiveTime(long[] stats) {
+        return stats[mDeviceActiveTimePosition];
+    }
+
+    public void setUidRxBytes(long[] stats, long count) {
+        stats[mUidRxBytesPosition] = count;
+    }
+
+    public long getUidRxBytes(long[] stats) {
+        return stats[mUidRxBytesPosition];
+    }
+
+    public void setUidTxBytes(long[] stats, long count) {
+        stats[mUidTxBytesPosition] = count;
+    }
+
+    public long getUidTxBytes(long[] stats) {
+        return stats[mUidTxBytesPosition];
+    }
+
+    public void setUidRxPackets(long[] stats, long count) {
+        stats[mUidRxPacketsPosition] = count;
+    }
+
+    public long getUidRxPackets(long[] stats) {
+        return stats[mUidRxPacketsPosition];
+    }
+
+    public void setUidTxPackets(long[] stats, long count) {
+        stats[mUidTxPacketsPosition] = count;
+    }
+
+    public long getUidTxPackets(long[] stats) {
+        return stats[mUidTxPacketsPosition];
+    }
+
+    public void setUidScanTime(long[] stats, long count) {
+        stats[mUidScanTimePosition] = count;
+    }
+
+    public long getUidScanTime(long[] stats) {
+        return stats[mUidScanTimePosition];
+    }
+
+    public void setUidBatchScanTime(long[] stats, long count) {
+        stats[mUidBatchScanTimePosition] = count;
+    }
+
+    public long getUidBatchedScanTime(long[] stats) {
+        return stats[mUidBatchScanTimePosition];
+    }
+
+    /**
+     * Copies the elements of the stats array layout into <code>extras</code>
+     */
+    public void toExtras(PersistableBundle extras) {
+        super.toExtras(extras);
+        extras.putBoolean(EXTRA_POWER_REPORTING_SUPPORTED, mPowerReportingSupported);
+        extras.putInt(EXTRA_DEVICE_RX_TIME_POSITION, mDeviceRxTimePosition);
+        extras.putInt(EXTRA_DEVICE_TX_TIME_POSITION, mDeviceTxTimePosition);
+        extras.putInt(EXTRA_DEVICE_SCAN_TIME_POSITION, mDeviceScanTimePosition);
+        extras.putInt(EXTRA_DEVICE_BASIC_SCAN_TIME_POSITION, mDeviceBasicScanTimePosition);
+        extras.putInt(EXTRA_DEVICE_BATCHED_SCAN_TIME_POSITION, mDeviceBatchedScanTimePosition);
+        extras.putInt(EXTRA_DEVICE_IDLE_TIME_POSITION, mDeviceIdleTimePosition);
+        extras.putInt(EXTRA_DEVICE_ACTIVE_TIME_POSITION, mDeviceActiveTimePosition);
+        extras.putInt(EXTRA_UID_RX_BYTES_POSITION, mUidRxBytesPosition);
+        extras.putInt(EXTRA_UID_TX_BYTES_POSITION, mUidTxBytesPosition);
+        extras.putInt(EXTRA_UID_RX_PACKETS_POSITION, mUidRxPacketsPosition);
+        extras.putInt(EXTRA_UID_TX_PACKETS_POSITION, mUidTxPacketsPosition);
+        extras.putInt(EXTRA_UID_SCAN_TIME_POSITION, mUidScanTimePosition);
+        extras.putInt(EXTRA_UID_BATCH_SCAN_TIME_POSITION, mUidBatchScanTimePosition);
+    }
+
+    /**
+     * Retrieves elements of the stats array layout from <code>extras</code>
+     */
+    public void fromExtras(PersistableBundle extras) {
+        super.fromExtras(extras);
+        mPowerReportingSupported = extras.getBoolean(EXTRA_POWER_REPORTING_SUPPORTED);
+        mDeviceRxTimePosition = extras.getInt(EXTRA_DEVICE_RX_TIME_POSITION);
+        mDeviceTxTimePosition = extras.getInt(EXTRA_DEVICE_TX_TIME_POSITION);
+        mDeviceScanTimePosition = extras.getInt(EXTRA_DEVICE_SCAN_TIME_POSITION);
+        mDeviceBasicScanTimePosition = extras.getInt(EXTRA_DEVICE_BASIC_SCAN_TIME_POSITION);
+        mDeviceBatchedScanTimePosition = extras.getInt(EXTRA_DEVICE_BATCHED_SCAN_TIME_POSITION);
+        mDeviceIdleTimePosition = extras.getInt(EXTRA_DEVICE_IDLE_TIME_POSITION);
+        mDeviceActiveTimePosition = extras.getInt(EXTRA_DEVICE_ACTIVE_TIME_POSITION);
+        mUidRxBytesPosition = extras.getInt(EXTRA_UID_RX_BYTES_POSITION);
+        mUidTxBytesPosition = extras.getInt(EXTRA_UID_TX_BYTES_POSITION);
+        mUidRxPacketsPosition = extras.getInt(EXTRA_UID_RX_PACKETS_POSITION);
+        mUidTxPacketsPosition = extras.getInt(EXTRA_UID_TX_PACKETS_POSITION);
+        mUidScanTimePosition = extras.getInt(EXTRA_UID_SCAN_TIME_POSITION);
+        mUidBatchScanTimePosition = extras.getInt(EXTRA_UID_BATCH_SCAN_TIME_POSITION);
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/WifiPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/WifiPowerStatsProcessor.java
new file mode 100644
index 0000000..5e9cc40
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/WifiPowerStatsProcessor.java
@@ -0,0 +1,425 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.power.stats;
+
+import android.util.Slog;
+
+import com.android.internal.os.PowerProfile;
+import com.android.internal.os.PowerStats;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+public class WifiPowerStatsProcessor extends PowerStatsProcessor {
+    private static final String TAG = "WifiPowerStatsProcessor";
+    private static final boolean DEBUG = false;
+
+    private final UsageBasedPowerEstimator mRxPowerEstimator;
+    private final UsageBasedPowerEstimator mTxPowerEstimator;
+    private final UsageBasedPowerEstimator mIdlePowerEstimator;
+
+    private final UsageBasedPowerEstimator mActivePowerEstimator;
+    private final UsageBasedPowerEstimator mScanPowerEstimator;
+    private final UsageBasedPowerEstimator mBatchedScanPowerEstimator;
+
+    private PowerStats.Descriptor mLastUsedDescriptor;
+    private WifiPowerStatsLayout mStatsLayout;
+    // Sequence of steps for power estimation and intermediate results.
+    private PowerEstimationPlan mPlan;
+
+    private long[] mTmpDeviceStatsArray;
+    private long[] mTmpUidStatsArray;
+    private boolean mHasWifiPowerController;
+
+    public WifiPowerStatsProcessor(PowerProfile powerProfile) {
+        mRxPowerEstimator = new UsageBasedPowerEstimator(
+                powerProfile.getAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_RX));
+        mTxPowerEstimator = new UsageBasedPowerEstimator(
+                powerProfile.getAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_TX));
+        mIdlePowerEstimator = new UsageBasedPowerEstimator(
+                powerProfile.getAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_IDLE));
+        mActivePowerEstimator = new UsageBasedPowerEstimator(
+                powerProfile.getAveragePower(PowerProfile.POWER_WIFI_ACTIVE));
+        mScanPowerEstimator = new UsageBasedPowerEstimator(
+                powerProfile.getAveragePower(PowerProfile.POWER_WIFI_SCAN));
+        mBatchedScanPowerEstimator = new UsageBasedPowerEstimator(
+                powerProfile.getAveragePower(PowerProfile.POWER_WIFI_BATCHED_SCAN));
+    }
+
+    private static class Intermediates {
+        /**
+         * Estimated power for the RX state.
+         */
+        public double rxPower;
+        /**
+         * Estimated power for the TX state.
+         */
+        public double txPower;
+        /**
+         * Estimated power in the SCAN state
+         */
+        public double scanPower;
+        /**
+         * Estimated power for IDLE, SCAN states.
+         */
+        public double idlePower;
+        /**
+         * Number of received packets
+         */
+        public long rxPackets;
+        /**
+         * Number of transmitted packets
+         */
+        public long txPackets;
+        /**
+         * Total duration of unbatched scans across all UIDs.
+         */
+        public long basicScanDuration;
+        /**
+         * Estimated power in the unbatched SCAN state
+         */
+        public double basicScanPower;
+        /**
+         * Total duration of batched scans across all UIDs.
+         */
+        public long batchedScanDuration;
+        /**
+         * Estimated power in the BATCHED SCAN state
+         */
+        public double batchedScanPower;
+        /**
+         * Estimated total power when active; used only in the absence of WiFiManager power
+         * reporting.
+         */
+        public double activePower;
+        /**
+         * Measured consumed energy from power monitoring hardware (micro-coulombs)
+         */
+        public long consumedEnergy;
+    }
+
+    @Override
+    void finish(PowerComponentAggregatedPowerStats stats) {
+        if (stats.getPowerStatsDescriptor() == null) {
+            return;
+        }
+
+        unpackPowerStatsDescriptor(stats.getPowerStatsDescriptor());
+
+        if (mPlan == null) {
+            mPlan = new PowerEstimationPlan(stats.getConfig());
+        }
+
+        for (int i = mPlan.deviceStateEstimations.size() - 1; i >= 0; i--) {
+            DeviceStateEstimation estimation = mPlan.deviceStateEstimations.get(i);
+            Intermediates intermediates = new Intermediates();
+            estimation.intermediates = intermediates;
+            computeDevicePowerEstimates(stats, estimation.stateValues, intermediates);
+        }
+
+        double ratio = 1.0;
+        if (mStatsLayout.getEnergyConsumerCount() != 0) {
+            ratio = computeEstimateAdjustmentRatioUsingConsumedEnergy();
+            if (ratio != 1) {
+                for (int i = mPlan.deviceStateEstimations.size() - 1; i >= 0; i--) {
+                    DeviceStateEstimation estimation = mPlan.deviceStateEstimations.get(i);
+                    adjustDevicePowerEstimates(stats, estimation.stateValues,
+                            (Intermediates) estimation.intermediates, ratio);
+                }
+            }
+        }
+
+        combineDeviceStateEstimates();
+
+        ArrayList<Integer> uids = new ArrayList<>();
+        stats.collectUids(uids);
+        if (!uids.isEmpty()) {
+            for (int uid : uids) {
+                for (int i = 0; i < mPlan.uidStateEstimates.size(); i++) {
+                    computeUidActivityTotals(stats, uid, mPlan.uidStateEstimates.get(i));
+                }
+            }
+
+            for (int uid : uids) {
+                for (int i = 0; i < mPlan.uidStateEstimates.size(); i++) {
+                    computeUidPowerEstimates(stats, uid, mPlan.uidStateEstimates.get(i));
+                }
+            }
+        }
+        mPlan.resetIntermediates();
+    }
+
+    private void unpackPowerStatsDescriptor(PowerStats.Descriptor descriptor) {
+        if (descriptor.equals(mLastUsedDescriptor)) {
+            return;
+        }
+
+        mLastUsedDescriptor = descriptor;
+        mStatsLayout = new WifiPowerStatsLayout(descriptor);
+        mTmpDeviceStatsArray = new long[descriptor.statsArrayLength];
+        mTmpUidStatsArray = new long[descriptor.uidStatsArrayLength];
+        mHasWifiPowerController = mStatsLayout.isPowerReportingSupported();
+    }
+
+    /**
+     * Compute power estimates using the power profile.
+     */
+    private void computeDevicePowerEstimates(PowerComponentAggregatedPowerStats stats,
+            int[] deviceStates, Intermediates intermediates) {
+        if (!stats.getDeviceStats(mTmpDeviceStatsArray, deviceStates)) {
+            return;
+        }
+
+        for (int i = mStatsLayout.getEnergyConsumerCount() - 1; i >= 0; i--) {
+            intermediates.consumedEnergy += mStatsLayout.getConsumedEnergy(mTmpDeviceStatsArray, i);
+        }
+
+        intermediates.basicScanDuration =
+                mStatsLayout.getDeviceBasicScanTime(mTmpDeviceStatsArray);
+        intermediates.batchedScanDuration =
+                mStatsLayout.getDeviceBatchedScanTime(mTmpDeviceStatsArray);
+        if (mHasWifiPowerController) {
+            intermediates.rxPower = mRxPowerEstimator.calculatePower(
+                    mStatsLayout.getDeviceRxTime(mTmpDeviceStatsArray));
+            intermediates.txPower = mTxPowerEstimator.calculatePower(
+                    mStatsLayout.getDeviceTxTime(mTmpDeviceStatsArray));
+            intermediates.scanPower = mScanPowerEstimator.calculatePower(
+                    mStatsLayout.getDeviceScanTime(mTmpDeviceStatsArray));
+            intermediates.idlePower = mIdlePowerEstimator.calculatePower(
+                    mStatsLayout.getDeviceIdleTime(mTmpDeviceStatsArray));
+            mStatsLayout.setDevicePowerEstimate(mTmpDeviceStatsArray,
+                    intermediates.rxPower + intermediates.txPower + intermediates.scanPower
+                            + intermediates.idlePower);
+        } else {
+            intermediates.activePower = mActivePowerEstimator.calculatePower(
+                    mStatsLayout.getDeviceActiveTime(mTmpDeviceStatsArray));
+            intermediates.basicScanPower =
+                    mScanPowerEstimator.calculatePower(intermediates.basicScanDuration);
+            intermediates.batchedScanPower =
+                    mBatchedScanPowerEstimator.calculatePower(intermediates.batchedScanDuration);
+            mStatsLayout.setDevicePowerEstimate(mTmpDeviceStatsArray,
+                    intermediates.activePower + intermediates.basicScanPower
+                            + intermediates.batchedScanPower);
+        }
+
+        stats.setDeviceStats(deviceStates, mTmpDeviceStatsArray);
+    }
+
+    /**
+     * Compute an adjustment ratio using the total power estimated using the power profile
+     * and the total power measured by hardware.
+     */
+    private double computeEstimateAdjustmentRatioUsingConsumedEnergy() {
+        long totalConsumedEnergy = 0;
+        double totalPower = 0;
+
+        for (int i = mPlan.deviceStateEstimations.size() - 1; i >= 0; i--) {
+            Intermediates intermediates =
+                    (Intermediates) mPlan.deviceStateEstimations.get(i).intermediates;
+            if (mHasWifiPowerController) {
+                totalPower += intermediates.rxPower + intermediates.txPower
+                        + intermediates.scanPower + intermediates.idlePower;
+            } else {
+                totalPower += intermediates.activePower + intermediates.basicScanPower
+                        + intermediates.batchedScanPower;
+            }
+            totalConsumedEnergy += intermediates.consumedEnergy;
+        }
+
+        if (totalPower == 0) {
+            return 1;
+        }
+
+        return uCtoMah(totalConsumedEnergy) / totalPower;
+    }
+
+    /**
+     * Uniformly apply the same adjustment to all power estimates in order to ensure that the total
+     * estimated power matches the measured consumed power.  We are not claiming that all
+     * averages captured in the power profile have to be off by the same percentage in reality.
+     */
+    private void adjustDevicePowerEstimates(PowerComponentAggregatedPowerStats stats,
+            int[] deviceStates, Intermediates intermediates, double ratio) {
+        double adjutedPower;
+        if (mHasWifiPowerController) {
+            intermediates.rxPower *= ratio;
+            intermediates.txPower *= ratio;
+            intermediates.scanPower *= ratio;
+            intermediates.idlePower *= ratio;
+            adjutedPower = intermediates.rxPower + intermediates.txPower + intermediates.scanPower
+                    + intermediates.idlePower;
+        } else {
+            intermediates.activePower *= ratio;
+            intermediates.basicScanPower *= ratio;
+            intermediates.batchedScanPower *= ratio;
+            adjutedPower = intermediates.activePower + intermediates.basicScanPower
+                    + intermediates.batchedScanPower;
+        }
+
+        if (!stats.getDeviceStats(mTmpDeviceStatsArray, deviceStates)) {
+            return;
+        }
+
+        mStatsLayout.setDevicePowerEstimate(mTmpDeviceStatsArray, adjutedPower);
+        stats.setDeviceStats(deviceStates, mTmpDeviceStatsArray);
+    }
+
+    /**
+     * Combine power estimates before distributing them proportionally to UIDs.
+     */
+    private void combineDeviceStateEstimates() {
+        for (int i = mPlan.combinedDeviceStateEstimations.size() - 1; i >= 0; i--) {
+            CombinedDeviceStateEstimate cdse = mPlan.combinedDeviceStateEstimations.get(i);
+            Intermediates
+                    cdseIntermediates = new Intermediates();
+            cdse.intermediates = cdseIntermediates;
+            List<DeviceStateEstimation> deviceStateEstimations = cdse.deviceStateEstimations;
+            for (int j = deviceStateEstimations.size() - 1; j >= 0; j--) {
+                DeviceStateEstimation dse = deviceStateEstimations.get(j);
+                Intermediates intermediates = (Intermediates) dse.intermediates;
+                if (mHasWifiPowerController) {
+                    cdseIntermediates.rxPower += intermediates.rxPower;
+                    cdseIntermediates.txPower += intermediates.txPower;
+                    cdseIntermediates.scanPower += intermediates.scanPower;
+                    cdseIntermediates.idlePower += intermediates.idlePower;
+                } else {
+                    cdseIntermediates.activePower += intermediates.activePower;
+                    cdseIntermediates.basicScanPower += intermediates.basicScanPower;
+                    cdseIntermediates.batchedScanPower += intermediates.batchedScanPower;
+                }
+                cdseIntermediates.basicScanDuration += intermediates.basicScanDuration;
+                cdseIntermediates.batchedScanDuration += intermediates.batchedScanDuration;
+                cdseIntermediates.consumedEnergy += intermediates.consumedEnergy;
+            }
+        }
+    }
+
+    private void computeUidActivityTotals(PowerComponentAggregatedPowerStats stats, int uid,
+            UidStateEstimate uidStateEstimate) {
+        Intermediates intermediates =
+                (Intermediates) uidStateEstimate.combinedDeviceStateEstimate.intermediates;
+        for (UidStateProportionalEstimate proportionalEstimate :
+                uidStateEstimate.proportionalEstimates) {
+            if (!stats.getUidStats(mTmpUidStatsArray, uid, proportionalEstimate.stateValues)) {
+                continue;
+            }
+
+            intermediates.rxPackets += mStatsLayout.getUidRxPackets(mTmpUidStatsArray);
+            intermediates.txPackets += mStatsLayout.getUidTxPackets(mTmpUidStatsArray);
+        }
+    }
+
+    private void computeUidPowerEstimates(PowerComponentAggregatedPowerStats stats, int uid,
+            UidStateEstimate uidStateEstimate) {
+        Intermediates intermediates =
+                (Intermediates) uidStateEstimate.combinedDeviceStateEstimate.intermediates;
+        for (UidStateProportionalEstimate proportionalEstimate :
+                uidStateEstimate.proportionalEstimates) {
+            if (!stats.getUidStats(mTmpUidStatsArray, uid, proportionalEstimate.stateValues)) {
+                continue;
+            }
+
+            double power = 0;
+            if (mHasWifiPowerController) {
+                if (intermediates.rxPackets != 0) {
+                    power += intermediates.rxPower * mStatsLayout.getUidRxPackets(mTmpUidStatsArray)
+                            / intermediates.rxPackets;
+                }
+                if (intermediates.txPackets != 0) {
+                    power += intermediates.txPower * mStatsLayout.getUidTxPackets(mTmpUidStatsArray)
+                            / intermediates.txPackets;
+                }
+                long totalScanDuration =
+                        intermediates.basicScanDuration + intermediates.batchedScanDuration;
+                if (totalScanDuration != 0) {
+                    long scanDuration = mStatsLayout.getUidScanTime(mTmpUidStatsArray)
+                            + mStatsLayout.getUidBatchedScanTime(mTmpUidStatsArray);
+                    power += intermediates.scanPower * scanDuration / totalScanDuration;
+                }
+            } else {
+                long totalPackets = intermediates.rxPackets + intermediates.txPackets;
+                if (totalPackets != 0) {
+                    long packets = mStatsLayout.getUidRxPackets(mTmpUidStatsArray)
+                            + mStatsLayout.getUidTxPackets(mTmpUidStatsArray);
+                    power += intermediates.activePower * packets / totalPackets;
+                }
+
+                if (intermediates.basicScanDuration != 0) {
+                    long scanDuration = mStatsLayout.getUidScanTime(mTmpUidStatsArray);
+                    power += intermediates.basicScanPower * scanDuration
+                            / intermediates.basicScanDuration;
+                }
+
+                if (intermediates.batchedScanDuration != 0) {
+                    long batchedScanDuration = mStatsLayout.getUidBatchedScanTime(
+                            mTmpUidStatsArray);
+                    power += intermediates.batchedScanPower * batchedScanDuration
+                            / intermediates.batchedScanDuration;
+                }
+            }
+            mStatsLayout.setUidPowerEstimate(mTmpUidStatsArray, power);
+            stats.setUidStats(uid, proportionalEstimate.stateValues, mTmpUidStatsArray);
+
+            if (DEBUG) {
+                Slog.d(TAG, "UID: " + uid
+                        + " states: " + Arrays.toString(proportionalEstimate.stateValues)
+                        + " stats: " + Arrays.toString(mTmpUidStatsArray)
+                        + " rx: " + mStatsLayout.getUidRxPackets(mTmpUidStatsArray)
+                        + " rx-power: " + intermediates.rxPower
+                        + " rx-packets: " + intermediates.rxPackets
+                        + " tx: " + mStatsLayout.getUidTxPackets(mTmpUidStatsArray)
+                        + " tx-power: " + intermediates.txPower
+                        + " tx-packets: " + intermediates.txPackets
+                        + " power: " + power);
+            }
+        }
+    }
+
+    @Override
+    String deviceStatsToString(PowerStats.Descriptor descriptor, long[] stats) {
+        unpackPowerStatsDescriptor(descriptor);
+        if (mHasWifiPowerController) {
+            return "rx: " + mStatsLayout.getDeviceRxTime(stats)
+                    + " tx: " + mStatsLayout.getDeviceTxTime(stats)
+                    + " scan: " + mStatsLayout.getDeviceScanTime(stats)
+                    + " idle: " + mStatsLayout.getDeviceIdleTime(stats)
+                    + " power: " + mStatsLayout.getDevicePowerEstimate(stats);
+        } else {
+            return "active: " + mStatsLayout.getDeviceActiveTime(stats)
+                    + " scan: " + mStatsLayout.getDeviceBasicScanTime(stats)
+                    + " batched-scan: " + mStatsLayout.getDeviceBatchedScanTime(stats)
+                    + " power: " + mStatsLayout.getDevicePowerEstimate(stats);
+        }
+    }
+
+    @Override
+    String stateStatsToString(PowerStats.Descriptor descriptor, int key, long[] stats) {
+        // Unsupported for this power component
+        return null;
+    }
+
+    @Override
+    String uidStatsToString(PowerStats.Descriptor descriptor, long[] stats) {
+        unpackPowerStatsDescriptor(descriptor);
+        return "rx: " + mStatsLayout.getUidRxPackets(stats)
+                + " tx: " + mStatsLayout.getUidTxPackets(stats)
+                + " scan: " + mStatsLayout.getUidScanTime(stats)
+                + " batched-scan: " + mStatsLayout.getUidBatchedScanTime(stats)
+                + " power: " + mStatsLayout.getUidPowerEstimate(stats);
+    }
+}
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
new file mode 100644
index 0000000..8b1d423
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsCollectorTest.java
@@ -0,0 +1,416 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.power.stats;
+
+import static android.net.NetworkStats.DEFAULT_NETWORK_NO;
+import static android.net.NetworkStats.METERED_NO;
+import static android.net.NetworkStats.ROAMING_NO;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.when;
+
+import android.content.Context;
+import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerType;
+import android.net.NetworkStats;
+import android.net.wifi.WifiManager;
+import android.os.BatteryConsumer;
+import android.os.BatteryStatsManager;
+import android.os.Handler;
+import android.os.WorkSource;
+import android.os.connectivity.WifiActivityEnergyInfo;
+import android.platform.test.ravenwood.RavenwoodRule;
+import android.util.IndentingPrintWriter;
+import android.util.SparseArray;
+
+import com.android.internal.os.Clock;
+import com.android.internal.os.PowerStats;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.Mock;
+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 {
+    private static final int APP_UID1 = 42;
+    private static final int APP_UID2 = 24;
+    private static final int APP_UID3 = 44;
+    private static final int ISOLATED_UID = 99123;
+
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
+
+    @Rule(order = 1)
+    public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
+            .setPowerStatsThrottlePeriodMillis(BatteryConsumer.POWER_COMPONENT_WIFI, 1000);
+
+    private MockBatteryStatsImpl mBatteryStats;
+
+    private final MockClock mClock = mStatsRule.getMockClock();
+
+    @Mock
+    private Context mContext;
+    @Mock
+    private PackageManager mPackageManager;
+    @Mock
+    private WifiManager mWifiManager;
+    @Mock
+    private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
+    @Mock
+    private Supplier<NetworkStats> mNetworkStatsSupplier;
+    @Mock
+    private PowerStatsUidResolver mPowerStatsUidResolver;
+
+    private NetworkStats mNetworkStats;
+    private List<NetworkStats.Entry> mNetworkStatsEntries;
+
+    private static class ScanTimes {
+        public long scanTimeMs;
+        public long batchScanTimeMs;
+    }
+
+    private final SparseArray<ScanTimes> mScanTimes = new SparseArray<>();
+    private long mWifiActiveDuration;
+
+    private final WifiPowerStatsCollector.WifiStatsRetriever mWifiStatsRetriever =
+            new WifiPowerStatsCollector.WifiStatsRetriever() {
+        @Override
+        public void retrieveWifiScanTimes(Callback callback) {
+            for (int i = 0; i < mScanTimes.size(); i++) {
+                int uid = mScanTimes.keyAt(i);
+                ScanTimes scanTimes = mScanTimes.valueAt(i);
+                callback.onWifiScanTime(uid, scanTimes.scanTimeMs, scanTimes.batchScanTimeMs);
+            }
+        }
+
+        @Override
+        public long getWifiActiveDuration() {
+            return mWifiActiveDuration;
+        }
+    };
+
+    private final List<PowerStats> mRecordedPowerStats = new ArrayList<>();
+
+    private WifiPowerStatsCollector.Injector mInjector = new WifiPowerStatsCollector.Injector() {
+        @Override
+        public Handler getHandler() {
+            return mStatsRule.getHandler();
+        }
+
+        @Override
+        public Clock getClock() {
+            return mStatsRule.getMockClock();
+        }
+
+        @Override
+        public PowerStatsUidResolver getUidResolver() {
+            return mPowerStatsUidResolver;
+        }
+
+        @Override
+        public PackageManager getPackageManager() {
+            return mPackageManager;
+        }
+
+        @Override
+        public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
+            return mConsumedEnergyRetriever;
+        }
+
+        @Override
+        public IntSupplier getVoltageSupplier() {
+            return () -> 3500;
+        }
+
+        @Override
+        public Supplier<NetworkStats> getWifiNetworkStatsSupplier() {
+            return mNetworkStatsSupplier;
+        }
+
+        @Override
+        public WifiPowerStatsCollector.WifiStatsRetriever getWifiStatsRetriever() {
+            return mWifiStatsRetriever;
+        }
+
+        @Override
+        public WifiManager getWifiManager() {
+            return mWifiManager;
+        }
+    };
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+        when(mContext.getPackageManager()).thenReturn(mPackageManager);
+        when(mPackageManager.hasSystemFeature(PackageManager.FEATURE_WIFI)).thenReturn(true);
+        when(mPowerStatsUidResolver.mapUid(anyInt())).thenAnswer(invocation -> {
+            int uid = invocation.getArgument(0);
+            if (uid == ISOLATED_UID) {
+                return APP_UID2;
+            } else {
+                return uid;
+            }
+        });
+        mBatteryStats = mStatsRule.getBatteryStats();
+    }
+
+    @SuppressWarnings("GuardedBy")
+    @Test
+    public void triggering() throws Throwable {
+        PowerStatsCollector collector = mBatteryStats.getPowerStatsCollector(
+                BatteryConsumer.POWER_COMPONENT_WIFI);
+        collector.addConsumer(mRecordedPowerStats::add);
+
+        mBatteryStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_WIFI, true);
+
+        mockWifiActivityInfo(1000, 2000, 3000, 600, 100);
+
+        // This should trigger a sample collection to establish a baseline
+        mBatteryStats.onSystemReady(mContext);
+
+        mStatsRule.waitForBackgroundThread();
+        assertThat(mRecordedPowerStats).hasSize(1);
+
+        mRecordedPowerStats.clear();
+        mStatsRule.setTime(20000, 20000);
+        mBatteryStats.noteWifiOnLocked(mClock.realtime, mClock.uptime);
+        mStatsRule.waitForBackgroundThread();
+        assertThat(mRecordedPowerStats).hasSize(1);
+
+        mRecordedPowerStats.clear();
+        mStatsRule.setTime(40000, 40000);
+        mBatteryStats.noteWifiOffLocked(mClock.realtime, mClock.uptime);
+        mStatsRule.waitForBackgroundThread();
+        assertThat(mRecordedPowerStats).hasSize(1);
+
+        mRecordedPowerStats.clear();
+        mStatsRule.setTime(50000, 50000);
+        mBatteryStats.noteWifiRunningLocked(new WorkSource(APP_UID1), mClock.realtime,
+                mClock.uptime);
+        mStatsRule.waitForBackgroundThread();
+        assertThat(mRecordedPowerStats).hasSize(1);
+
+        mRecordedPowerStats.clear();
+        mStatsRule.setTime(60000, 60000);
+        mBatteryStats.noteWifiStoppedLocked(new WorkSource(APP_UID1), mClock.realtime,
+                mClock.uptime);
+        mStatsRule.waitForBackgroundThread();
+        assertThat(mRecordedPowerStats).hasSize(1);
+
+        mRecordedPowerStats.clear();
+        mStatsRule.setTime(70000, 70000);
+        mBatteryStats.noteWifiStateLocked(BatteryStatsManager.WIFI_STATE_ON_CONNECTED_STA,
+                "mywyfy", mClock.realtime);
+        mStatsRule.waitForBackgroundThread();
+        assertThat(mRecordedPowerStats).hasSize(1);
+    }
+
+    @Test
+    public void collectStats_powerReportingSupported() throws Throwable {
+        PowerStats powerStats = collectPowerStats(true);
+        assertThat(powerStats.durationMs).isEqualTo(7500);
+
+        PowerStats.Descriptor descriptor = powerStats.descriptor;
+        WifiPowerStatsLayout layout = new WifiPowerStatsLayout(descriptor);
+        assertThat(layout.isPowerReportingSupported()).isTrue();
+        assertThat(layout.getDeviceRxTime(powerStats.stats)).isEqualTo(6000);
+        assertThat(layout.getDeviceTxTime(powerStats.stats)).isEqualTo(1000);
+        assertThat(layout.getDeviceScanTime(powerStats.stats)).isEqualTo(200);
+        assertThat(layout.getDeviceIdleTime(powerStats.stats)).isEqualTo(300);
+        assertThat(layout.getConsumedEnergy(powerStats.stats, 0))
+                .isEqualTo((64321 - 10000) * 1000 / 3500);
+
+        verifyUidStats(powerStats);
+    }
+
+    @Test
+    public void collectStats_powerReportingUnsupported() {
+        PowerStats powerStats = collectPowerStats(false);
+        assertThat(powerStats.durationMs).isEqualTo(13200);
+
+        PowerStats.Descriptor descriptor = powerStats.descriptor;
+        WifiPowerStatsLayout layout = new WifiPowerStatsLayout(descriptor);
+        assertThat(layout.isPowerReportingSupported()).isFalse();
+        assertThat(layout.getDeviceActiveTime(powerStats.stats)).isEqualTo(7500);
+        assertThat(layout.getDeviceBasicScanTime(powerStats.stats)).isEqualTo(234 + 100 + 300);
+        assertThat(layout.getDeviceBatchedScanTime(powerStats.stats)).isEqualTo(345 + 200 + 400);
+        assertThat(layout.getConsumedEnergy(powerStats.stats, 0))
+                .isEqualTo((64321 - 10000) * 1000 / 3500);
+
+        verifyUidStats(powerStats);
+    }
+
+    private void verifyUidStats(PowerStats powerStats) {
+        WifiPowerStatsLayout layout = new WifiPowerStatsLayout(powerStats.descriptor);
+        assertThat(powerStats.uidStats.size()).isEqualTo(2);
+        long[] actual1 = powerStats.uidStats.get(APP_UID1);
+        assertThat(layout.getUidRxBytes(actual1)).isEqualTo(1000);
+        assertThat(layout.getUidTxBytes(actual1)).isEqualTo(2000);
+        assertThat(layout.getUidRxPackets(actual1)).isEqualTo(100);
+        assertThat(layout.getUidTxPackets(actual1)).isEqualTo(200);
+        assertThat(layout.getUidScanTime(actual1)).isEqualTo(234);
+        assertThat(layout.getUidBatchedScanTime(actual1)).isEqualTo(345);
+
+        // Combines APP_UID2 and ISOLATED_UID
+        long[] actual2 = powerStats.uidStats.get(APP_UID2);
+        assertThat(layout.getUidRxBytes(actual2)).isEqualTo(6000);
+        assertThat(layout.getUidTxBytes(actual2)).isEqualTo(3000);
+        assertThat(layout.getUidRxPackets(actual2)).isEqualTo(60);
+        assertThat(layout.getUidTxPackets(actual2)).isEqualTo(30);
+        assertThat(layout.getUidScanTime(actual2)).isEqualTo(100 + 300);
+        assertThat(layout.getUidBatchedScanTime(actual2)).isEqualTo(200 + 400);
+
+        assertThat(powerStats.uidStats.get(ISOLATED_UID)).isNull();
+        assertThat(powerStats.uidStats.get(APP_UID3)).isNull();
+    }
+
+    @Test
+    public void dump() throws Throwable {
+        PowerStats powerStats = collectPowerStats(true);
+        StringWriter sw = new StringWriter();
+        IndentingPrintWriter pw = new IndentingPrintWriter(sw);
+        powerStats.dump(pw);
+        pw.flush();
+        String dump = sw.toString();
+        assertThat(dump).contains("duration=7500");
+        assertThat(dump).contains(
+                "stats=[6000, 1000, 300, 200, 634, 945, " + ((64321 - 10000) * 1000 / 3500)
+                        + ", 0, 0]");
+        assertThat(dump).contains("UID 24: [6000, 3000, 60, 30, 400, 600, 0]");
+        assertThat(dump).contains("UID 42: [1000, 2000, 100, 200, 234, 345, 0]");
+    }
+
+    private PowerStats collectPowerStats(boolean hasPowerReporting) {
+        when(mWifiManager.isEnhancedPowerReportingSupported()).thenReturn(hasPowerReporting);
+
+        WifiPowerStatsCollector collector = new WifiPowerStatsCollector(mInjector, 0);
+        collector.setEnabled(true);
+
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.WIFI))
+                .thenReturn(new int[]{777});
+
+        if (hasPowerReporting) {
+            mockWifiActivityInfo(1000, 600, 100, 2000, 3000);
+        } else {
+            mWifiActiveDuration = 5700;
+        }
+        mockNetworkStats(1000);
+        mockNetworkStatsEntry(APP_UID1, 4321, 321, 1234, 23);
+        mockNetworkStatsEntry(APP_UID2, 4000, 40, 2000, 20);
+        mockNetworkStatsEntry(ISOLATED_UID, 2000, 20, 1000, 10);
+        mockNetworkStatsEntry(APP_UID3, 314, 281, 314, 281);
+        mockWifiScanTimes(APP_UID1, 1000, 2000);
+        mockWifiScanTimes(APP_UID2, 3000, 4000);
+        mockWifiScanTimes(ISOLATED_UID, 5000, 6000);
+
+        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
+                .thenReturn(new long[]{10000});
+
+        collector.collectStats();
+
+        if (hasPowerReporting) {
+            mockWifiActivityInfo(1100, 6600, 1100, 2200, 3300);
+        } else {
+            mWifiActiveDuration = 13200;
+        }
+        mockNetworkStats(1100);
+        mockNetworkStatsEntry(APP_UID1, 5321, 421, 3234, 223);
+        mockNetworkStatsEntry(APP_UID2, 8000, 80, 4000, 40);
+        mockNetworkStatsEntry(ISOLATED_UID, 4000, 40, 2000, 20);
+        mockNetworkStatsEntry(APP_UID3, 314, 281, 314, 281);    // Unchanged
+        mockWifiScanTimes(APP_UID1, 1234, 2345);
+        mockWifiScanTimes(APP_UID2, 3100, 4200);
+        mockWifiScanTimes(ISOLATED_UID, 5300, 6400);
+
+        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
+                .thenReturn(new long[]{64321});
+
+        mStatsRule.setTime(20000, 20000);
+        return collector.collectStats();
+    }
+
+    private void mockWifiActivityInfo(long timestamp, long rxTimeMs, long txTimeMs, int scanTimeMs,
+            int idleTimeMs) {
+        int stackState = 0;
+        WifiActivityEnergyInfo info = new WifiActivityEnergyInfo(timestamp, stackState, txTimeMs,
+                rxTimeMs, scanTimeMs, idleTimeMs);
+        doAnswer(invocation -> {
+            WifiManager.OnWifiActivityEnergyInfoListener listener = invocation.getArgument(1);
+            listener.onWifiActivityEnergyInfo(info);
+            return null;
+        }).when(mWifiManager).getWifiActivityEnergyInfoAsync(any(), any());
+    }
+
+    private void mockNetworkStats(long elapsedRealtime) {
+        if (RavenwoodRule.isOnRavenwood()) {
+            mNetworkStats = mock(NetworkStats.class);
+            ArrayList<NetworkStats.Entry> networkStatsEntries = new ArrayList<>();
+            when(mNetworkStats.iterator()).thenAnswer(inv -> networkStatsEntries.iterator());
+            mNetworkStatsEntries = networkStatsEntries;
+        } else {
+            mNetworkStats = new NetworkStats(elapsedRealtime, 1);
+        }
+        when(mNetworkStatsSupplier.get()).thenReturn(mNetworkStats);
+    }
+
+    private void mockNetworkStatsEntry(int uid, long rxBytes, long rxPackets, long txBytes,
+            long txPackets) {
+        if (RavenwoodRule.isOnRavenwood()) {
+            NetworkStats.Entry entry = mock(NetworkStats.Entry.class);
+            when(entry.getUid()).thenReturn(uid);
+            when(entry.getMetered()).thenReturn(METERED_NO);
+            when(entry.getRoaming()).thenReturn(ROAMING_NO);
+            when(entry.getDefaultNetwork()).thenReturn(DEFAULT_NETWORK_NO);
+            when(entry.getRxBytes()).thenReturn(rxBytes);
+            when(entry.getRxPackets()).thenReturn(rxPackets);
+            when(entry.getTxBytes()).thenReturn(txBytes);
+            when(entry.getTxPackets()).thenReturn(txPackets);
+            when(entry.getOperations()).thenReturn(100L);
+            mNetworkStatsEntries.add(entry);
+        } else {
+            mNetworkStats = mNetworkStats
+                    .addEntry(new NetworkStats.Entry("wifi", uid, 0, 0,
+                            METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, rxBytes, rxPackets,
+                            txBytes, txPackets, 100));
+            reset(mNetworkStatsSupplier);
+            when(mNetworkStatsSupplier.get()).thenReturn(mNetworkStats);
+        }
+    }
+
+    private void mockWifiScanTimes(int uid, long scanTimeMs, long batchScanTimeMs) {
+        ScanTimes scanTimes = new ScanTimes();
+        scanTimes.scanTimeMs = scanTimeMs;
+        scanTimes.batchScanTimeMs = batchScanTimeMs;
+        mScanTimes.put(uid, scanTimes);
+    }
+}
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsProcessorTest.java
new file mode 100644
index 0000000..257a1a6
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsProcessorTest.java
@@ -0,0 +1,592 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.server.power.stats;
+
+import static android.net.NetworkStats.DEFAULT_NETWORK_NO;
+import static android.net.NetworkStats.METERED_NO;
+import static android.net.NetworkStats.ROAMING_NO;
+import static android.os.BatteryConsumer.PROCESS_STATE_BACKGROUND;
+import static android.os.BatteryConsumer.PROCESS_STATE_CACHED;
+import static android.os.BatteryConsumer.PROCESS_STATE_FOREGROUND;
+import static android.os.BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE;
+
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_SCREEN;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import android.annotation.Nullable;
+import android.content.Context;
+import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerType;
+import android.net.NetworkStats;
+import android.net.wifi.WifiManager;
+import android.os.BatteryConsumer;
+import android.os.Handler;
+import android.os.Process;
+import android.os.connectivity.WifiActivityEnergyInfo;
+import android.platform.test.ravenwood.RavenwoodRule;
+import android.util.SparseArray;
+
+import com.android.internal.os.Clock;
+import com.android.internal.os.PowerProfile;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+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 {
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
+
+    private static final double PRECISION = 0.00001;
+    private static final int APP_UID1 = Process.FIRST_APPLICATION_UID + 42;
+    private static final int APP_UID2 = Process.FIRST_APPLICATION_UID + 101;
+    private static final int WIFI_ENERGY_CONSUMER_ID = 1;
+    private static final int VOLTAGE_MV = 3500;
+
+    @Rule(order = 1)
+    public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
+            .setAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_IDLE, 360.0)
+            .setAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_RX, 480.0)
+            .setAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_TX, 720.0)
+            .setAveragePower(PowerProfile.POWER_WIFI_ACTIVE, 360.0)
+            .setAveragePower(PowerProfile.POWER_WIFI_SCAN, 480.0)
+            .setAveragePower(PowerProfile.POWER_WIFI_BATCHED_SCAN, 720.0)
+            .initMeasuredEnergyStatsLocked();
+
+    @Mock
+    private Context mContext;
+    @Mock
+    private PowerStatsUidResolver mPowerStatsUidResolver;
+    @Mock
+    private PackageManager mPackageManager;
+    @Mock
+    private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
+    @Mock
+    private Supplier<NetworkStats> mNetworkStatsSupplier;
+    @Mock
+    private WifiManager mWifiManager;
+
+    private static class ScanTimes {
+        public long scanTimeMs;
+        public long batchScanTimeMs;
+    }
+
+    private final SparseArray<ScanTimes> mScanTimes = new SparseArray<>();
+    private long mWifiActiveDuration;
+
+    private final WifiPowerStatsCollector.WifiStatsRetriever mWifiStatsRetriever =
+            new WifiPowerStatsCollector.WifiStatsRetriever() {
+        @Override
+        public void retrieveWifiScanTimes(Callback callback) {
+            for (int i = 0; i < mScanTimes.size(); i++) {
+                int uid = mScanTimes.keyAt(i);
+                ScanTimes scanTimes = mScanTimes.valueAt(i);
+                callback.onWifiScanTime(uid, scanTimes.scanTimeMs, scanTimes.batchScanTimeMs);
+            }
+        }
+
+        @Override
+        public long getWifiActiveDuration() {
+            return mWifiActiveDuration;
+        }
+    };
+
+    private final WifiPowerStatsCollector.Injector mInjector =
+            new WifiPowerStatsCollector.Injector() {
+                @Override
+                public Handler getHandler() {
+                    return mStatsRule.getHandler();
+                }
+
+                @Override
+                public Clock getClock() {
+                    return mStatsRule.getMockClock();
+                }
+
+                @Override
+                public PowerStatsUidResolver getUidResolver() {
+                    return mPowerStatsUidResolver;
+                }
+
+                @Override
+                public PackageManager getPackageManager() {
+                    return mPackageManager;
+                }
+
+                @Override
+                public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
+                    return mConsumedEnergyRetriever;
+                }
+
+                @Override
+                public IntSupplier getVoltageSupplier() {
+                    return () -> VOLTAGE_MV;
+                }
+
+                @Override
+                public Supplier<NetworkStats> getWifiNetworkStatsSupplier() {
+                    return mNetworkStatsSupplier;
+                }
+
+                @Override
+                public WifiManager getWifiManager() {
+                    return mWifiManager;
+                }
+
+                @Override
+                public WifiPowerStatsCollector.WifiStatsRetriever getWifiStatsRetriever() {
+                    return mWifiStatsRetriever;
+                }
+            };
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+
+        when(mContext.getPackageManager()).thenReturn(mPackageManager);
+        when(mPackageManager.hasSystemFeature(PackageManager.FEATURE_WIFI)).thenReturn(true);
+        when(mPowerStatsUidResolver.mapUid(anyInt()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+    }
+
+    @Test
+    public void powerProfileModel_powerController() {
+        when(mWifiManager.isEnhancedPowerReportingSupported()).thenReturn(true);
+
+        // No power monitoring hardware
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.WIFI))
+                .thenReturn(new int[0]);
+
+        WifiPowerStatsProcessor processor =
+                new WifiPowerStatsProcessor(mStatsRule.getPowerProfile());
+
+        PowerComponentAggregatedPowerStats aggregatedStats = createAggregatedPowerStats(processor);
+
+        WifiPowerStatsCollector collector = new WifiPowerStatsCollector(mInjector, 0);
+        collector.setEnabled(true);
+
+        // Initial empty WifiActivityEnergyInfo.
+        mockWifiActivityEnergyInfo(new WifiActivityEnergyInfo(0L,
+                WifiActivityEnergyInfo.STACK_STATE_INVALID, 0L, 0L, 0L, 0L));
+
+        // Establish a baseline
+        aggregatedStats.addPowerStats(collector.collectStats(), 0);
+
+        // Turn the screen off after 2.5 seconds
+        aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_OTHER, 2500);
+        aggregatedStats.setUidState(APP_UID1, STATE_PROCESS_STATE, PROCESS_STATE_BACKGROUND, 2500);
+        aggregatedStats.setUidState(APP_UID1, STATE_PROCESS_STATE, PROCESS_STATE_FOREGROUND_SERVICE,
+                5000);
+
+        // Note application network activity
+        NetworkStats networkStats = mockNetworkStats(10000, 1,
+                mockNetworkStatsEntry("wifi", APP_UID1, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 10000, 1500, 20000, 300, 100),
+                mockNetworkStatsEntry("wifi", APP_UID2, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 5000, 500, 3000, 100, 111));
+        when(mNetworkStatsSupplier.get()).thenReturn(networkStats);
+
+        mockWifiScanTimes(APP_UID1, 300, 400);
+        mockWifiScanTimes(APP_UID2, 100, 200);
+
+        mockWifiActivityEnergyInfo(new WifiActivityEnergyInfo(10000,
+                WifiActivityEnergyInfo.STACK_STATE_STATE_ACTIVE, 2000, 3000, 100, 600));
+
+        mStatsRule.setTime(10_000, 10_000);
+
+        aggregatedStats.addPowerStats(collector.collectStats(), 10_000);
+
+        processor.finish(aggregatedStats);
+
+        WifiPowerStatsLayout statsLayout =
+                new WifiPowerStatsLayout(aggregatedStats.getPowerStatsDescriptor());
+
+        // RX power = 'rx-duration * PowerProfile[wifi.controller.rx]`
+        //        RX power = 3000 * 480 = 1440000 mA-ms = 0.4 mAh
+        // TX power = 'tx-duration * PowerProfile[wifi.controller.tx]`
+        //        TX power = 2000 * 720 = 1440000 mA-ms = 0.4 mAh
+        // Scan power = 'scan-duration * PowerProfile[wifi.scan]`
+        //        Scan power = 100 * 480 = 48000 mA-ms = 0.013333 mAh
+        // Idle power = 'idle-duration * PowerProfile[wifi.idle]`
+        //        Idle power = 600 * 360 = 216000 mA-ms = 0.06 mAh
+        // Total power = RX + TX + Scan + Idle = 0.873333
+        // Screen-on  - 25%
+        // Screen-off - 75%
+        double expectedPower = 0.873333;
+        long[] deviceStats = new long[aggregatedStats.getPowerStatsDescriptor().statsArrayLength];
+        aggregatedStats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_ON));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(expectedPower * 0.25);
+
+        aggregatedStats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_OTHER));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(expectedPower * 0.75);
+
+        // UID1 =
+        //     (1500 / 2000) * 0.4        // rx
+        //     + (300 / 400) * 0.4        // tx
+        //     + (700 / 1000) * 0.013333  // scan (basic + batched)
+        //   = 0.609333 mAh
+        double expectedPower1 = 0.609333;
+        long[] uidStats = new long[aggregatedStats.getPowerStatsDescriptor().uidStatsArrayLength];
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_FOREGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_BACKGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_FOREGROUND_SERVICE));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.5);
+
+        // UID2 =
+        //     (500 / 2000) * 0.4         // rx
+        //     + (100 / 400) * 0.4        // tx
+        //     + (300 / 1000) * 0.013333  // scan (basic + batched)
+        //   = 0.204 mAh
+        double expectedPower2 = 0.204;
+        aggregatedStats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower2 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower2 * 0.75);
+    }
+
+    @Test
+    public void consumedEnergyModel_powerController() {
+        when(mWifiManager.isEnhancedPowerReportingSupported()).thenReturn(true);
+
+        // PowerStats hardware is available
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.WIFI))
+                .thenReturn(new int[] {WIFI_ENERGY_CONSUMER_ID});
+
+        WifiPowerStatsProcessor processor =
+                new WifiPowerStatsProcessor(mStatsRule.getPowerProfile());
+
+        PowerComponentAggregatedPowerStats aggregatedStats = createAggregatedPowerStats(processor);
+
+        WifiPowerStatsCollector collector = new WifiPowerStatsCollector(mInjector, 0);
+        collector.setEnabled(true);
+
+        // Initial empty WifiActivityEnergyInfo.
+        mockWifiActivityEnergyInfo(new WifiActivityEnergyInfo(0L,
+                WifiActivityEnergyInfo.STACK_STATE_INVALID, 0L, 0L, 0L, 0L));
+
+        when(mConsumedEnergyRetriever.getConsumedEnergyUws(
+                new int[]{WIFI_ENERGY_CONSUMER_ID}))
+                .thenReturn(new long[]{0});
+
+        // Establish a baseline
+        aggregatedStats.addPowerStats(collector.collectStats(), 0);
+
+        // Turn the screen off after 2.5 seconds
+        aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_OTHER, 2500);
+        aggregatedStats.setUidState(APP_UID1, STATE_PROCESS_STATE, PROCESS_STATE_BACKGROUND, 2500);
+        aggregatedStats.setUidState(APP_UID1, STATE_PROCESS_STATE, PROCESS_STATE_FOREGROUND_SERVICE,
+                5000);
+
+        // Note application network activity
+        NetworkStats networkStats = mockNetworkStats(10000, 1,
+                mockNetworkStatsEntry("wifi", APP_UID1, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 10000, 1500, 20000, 300, 100),
+                mockNetworkStatsEntry("wifi", APP_UID2, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 5000, 500, 3000, 100, 111));
+        when(mNetworkStatsSupplier.get()).thenReturn(networkStats);
+
+        mockWifiScanTimes(APP_UID1, 300, 400);
+        mockWifiScanTimes(APP_UID2, 100, 200);
+
+        mockWifiActivityEnergyInfo(new WifiActivityEnergyInfo(10000,
+                WifiActivityEnergyInfo.STACK_STATE_STATE_ACTIVE, 2000, 3000, 100, 600));
+
+        mStatsRule.setTime(10_000, 10_000);
+
+        // 10 mAh represented as microWattSeconds
+        long energyUws = 10 * 3600 * VOLTAGE_MV;
+        when(mConsumedEnergyRetriever.getConsumedEnergyUws(
+                new int[]{WIFI_ENERGY_CONSUMER_ID})).thenReturn(new long[]{energyUws});
+
+        aggregatedStats.addPowerStats(collector.collectStats(), 10_000);
+
+        processor.finish(aggregatedStats);
+
+        WifiPowerStatsLayout statsLayout =
+                new WifiPowerStatsLayout(aggregatedStats.getPowerStatsDescriptor());
+
+        // All estimates are computed as in the #powerProfileModel_powerController test,
+        // except they are all scaled by the same ratio to ensure that the total estimated
+        // energy is equal to the measured energy
+        double expectedPower = 10;
+        long[] deviceStats = new long[aggregatedStats.getPowerStatsDescriptor().statsArrayLength];
+        aggregatedStats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_ON));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(expectedPower * 0.25);
+
+        aggregatedStats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_OTHER));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(expectedPower * 0.75);
+
+        // UID1
+        //   0.609333           // power profile model estimate
+        //   0.873333           // power profile model estimate for total power
+        //   10                 // total consumed energy
+        //   = 0.609333 * (10 / 0.873333) = 6.9771
+        double expectedPower1 = 6.9771;
+        long[] uidStats = new long[aggregatedStats.getPowerStatsDescriptor().uidStatsArrayLength];
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_FOREGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_BACKGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_FOREGROUND_SERVICE));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.5);
+
+        // UID2
+        //   0.204              // power profile model estimate
+        //   0.873333           // power profile model estimate for total power
+        //   10                 // total consumed energy
+        //   = 0.204 * (10 / 0.873333) = 2.33588
+        double expectedPower2 = 2.33588;
+        aggregatedStats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower2 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower2 * 0.75);
+    }
+
+    @Test
+    public void powerProfileModel_noPowerController() {
+        when(mWifiManager.isEnhancedPowerReportingSupported()).thenReturn(false);
+
+        // No power monitoring hardware
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.WIFI))
+                .thenReturn(new int[0]);
+
+        WifiPowerStatsProcessor processor =
+                new WifiPowerStatsProcessor(mStatsRule.getPowerProfile());
+
+        PowerComponentAggregatedPowerStats aggregatedStats = createAggregatedPowerStats(processor);
+
+        WifiPowerStatsCollector collector = new WifiPowerStatsCollector(mInjector, 0);
+        collector.setEnabled(true);
+
+        // Establish a baseline
+        aggregatedStats.addPowerStats(collector.collectStats(), 0);
+
+        // Turn the screen off after 2.5 seconds
+        aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_OTHER, 2500);
+        aggregatedStats.setUidState(APP_UID1, STATE_PROCESS_STATE, PROCESS_STATE_BACKGROUND, 2500);
+        aggregatedStats.setUidState(APP_UID1, STATE_PROCESS_STATE, PROCESS_STATE_FOREGROUND_SERVICE,
+                5000);
+
+        // Note application network activity
+        NetworkStats networkStats = mockNetworkStats(10000, 1,
+                mockNetworkStatsEntry("wifi", APP_UID1, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 10000, 1500, 20000, 300, 100),
+                mockNetworkStatsEntry("wifi", APP_UID2, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 5000, 500, 3000, 100, 111));
+        when(mNetworkStatsSupplier.get()).thenReturn(networkStats);
+
+        mScanTimes.clear();
+        mWifiActiveDuration = 8000;
+        mockWifiScanTimes(APP_UID1, 300, 400);
+        mockWifiScanTimes(APP_UID2, 100, 200);
+
+        mStatsRule.setTime(10_000, 10_000);
+
+        aggregatedStats.addPowerStats(collector.collectStats(), 10_000);
+
+        processor.finish(aggregatedStats);
+
+        WifiPowerStatsLayout statsLayout =
+                new WifiPowerStatsLayout(aggregatedStats.getPowerStatsDescriptor());
+
+        // Total active power = 'active-duration * PowerProfile[wifi.on]`
+        //        active = 8000 * 360 = 2880000 mA-ms = 0.8 mAh
+        // UID1 rxPackets + txPackets = 1800
+        // UID2 rxPackets + txPackets = 600
+        // Total rx+tx packets = 2400
+        // Total scan power = `scan-duration * PowerProfile[wifi.scan]`
+        //        scan = (100 + 300) * 480 = 192000 mA-ms = 0.05333 mAh
+        // Total batch scan power = `(200 + 400) * PowerProfile[wifi.batchedscan]`
+        //        bscan = (200 + 400) * 720 = 432000 mA-ms = 0.12 mAh
+        //
+        // Expected power = active + scan + bscan = 0.97333
+        double expectedPower = 0.97333;
+        long[] deviceStats = new long[aggregatedStats.getPowerStatsDescriptor().statsArrayLength];
+        aggregatedStats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_ON));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(expectedPower * 0.25);
+
+        aggregatedStats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_OTHER));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(expectedPower * 0.75);
+
+        // UID1 =
+        //     (1800 / 2400) * 0.8      // active
+        //     + (300 / 400) * 0.05333  // scan
+        //     + (400 / 600) * 0.12     // batched scan
+        //   = 0.72 mAh
+        double expectedPower1 = 0.72;
+        long[] uidStats = new long[aggregatedStats.getPowerStatsDescriptor().uidStatsArrayLength];
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_FOREGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_BACKGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_FOREGROUND_SERVICE));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.5);
+
+        // UID2 =
+        //     (600 / 2400) * 0.8       // active
+        //     + (100 / 400) * 0.05333  // scan
+        //     + (200 / 600) * 0.12     // batched scan
+        //   = 0.253333 mAh
+        double expectedPower2 = 0.25333;
+        aggregatedStats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower2 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower2 * 0.75);
+    }
+
+    private static PowerComponentAggregatedPowerStats createAggregatedPowerStats(
+            WifiPowerStatsProcessor processor) {
+        AggregatedPowerStatsConfig.PowerComponent config =
+                new AggregatedPowerStatsConfig.PowerComponent(BatteryConsumer.POWER_COMPONENT_WIFI)
+                        .trackDeviceStates(STATE_POWER, STATE_SCREEN)
+                        .trackUidStates(STATE_POWER, STATE_SCREEN, STATE_PROCESS_STATE)
+                        .setProcessor(processor);
+
+        PowerComponentAggregatedPowerStats aggregatedStats =
+                new PowerComponentAggregatedPowerStats(
+                        new AggregatedPowerStats(mock(AggregatedPowerStatsConfig.class)), config);
+
+        aggregatedStats.setState(STATE_POWER, POWER_STATE_OTHER, 0);
+        aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_ON, 0);
+        aggregatedStats.setUidState(APP_UID1, STATE_PROCESS_STATE, PROCESS_STATE_FOREGROUND, 0);
+        aggregatedStats.setUidState(APP_UID2, STATE_PROCESS_STATE, PROCESS_STATE_CACHED, 0);
+
+        return aggregatedStats;
+    }
+
+    private int[] states(int... states) {
+        return states;
+    }
+
+    private void mockWifiActivityEnergyInfo(WifiActivityEnergyInfo waei) {
+        doAnswer(invocation -> {
+            WifiManager.OnWifiActivityEnergyInfoListener
+                    listener = invocation.getArgument(1);
+            listener.onWifiActivityEnergyInfo(waei);
+            return null;
+        }).when(mWifiManager).getWifiActivityEnergyInfoAsync(any(), any());
+    }
+
+    private NetworkStats mockNetworkStats(int elapsedTime, int initialSize,
+            NetworkStats.Entry... entries) {
+        NetworkStats stats;
+        if (RavenwoodRule.isOnRavenwood()) {
+            stats = mock(NetworkStats.class);
+            when(stats.iterator()).thenAnswer(inv -> List.of(entries).iterator());
+        } else {
+            stats = new NetworkStats(elapsedTime, initialSize);
+            for (NetworkStats.Entry entry : entries) {
+                stats = stats.addEntry(entry);
+            }
+        }
+        return stats;
+    }
+
+    private static NetworkStats.Entry mockNetworkStatsEntry(@Nullable String iface, int uid,
+            int set, int tag, int metered, int roaming, int defaultNetwork, long rxBytes,
+            long rxPackets, long txBytes, long txPackets, long operations) {
+        if (RavenwoodRule.isOnRavenwood()) {
+            NetworkStats.Entry entry = mock(NetworkStats.Entry.class);
+            when(entry.getUid()).thenReturn(uid);
+            when(entry.getMetered()).thenReturn(metered);
+            when(entry.getRoaming()).thenReturn(roaming);
+            when(entry.getDefaultNetwork()).thenReturn(defaultNetwork);
+            when(entry.getRxBytes()).thenReturn(rxBytes);
+            when(entry.getRxPackets()).thenReturn(rxPackets);
+            when(entry.getTxBytes()).thenReturn(txBytes);
+            when(entry.getTxPackets()).thenReturn(txPackets);
+            when(entry.getOperations()).thenReturn(operations);
+            return entry;
+        } else {
+            return new NetworkStats.Entry(iface, uid, set, tag, metered,
+                    roaming, defaultNetwork, rxBytes, rxPackets, txBytes, txPackets, operations);
+        }
+    }
+
+    private void mockWifiScanTimes(int uid, long scanTimeMs, long batchScanTimeMs) {
+        ScanTimes scanTimes = new ScanTimes();
+        scanTimes.scanTimeMs = scanTimeMs;
+        scanTimes.batchScanTimeMs = batchScanTimeMs;
+        mScanTimes.put(uid, scanTimes);
+    }
+}