Fix device-level wakelock battery attribution

Bug: 187844042
Test: atest FrameworksCoreTests:WakelockPowerCalculatorTest
Change-Id: I4e14d962db37adeafb793022b50da0065628238c
diff --git a/core/java/com/android/internal/os/WakelockPowerCalculator.java b/core/java/com/android/internal/os/WakelockPowerCalculator.java
index d594107..e0ef129 100644
--- a/core/java/com/android/internal/os/WakelockPowerCalculator.java
+++ b/core/java/com/android/internal/os/WakelockPowerCalculator.java
@@ -75,22 +75,29 @@
         // this remainder to the OS, if possible.
         calculateRemaining(result, batteryStats, rawRealtimeUs, rawUptimeUs,
                 BatteryStats.STATS_SINCE_CHARGED, osPowerMah, osDurationMs, totalAppDurationMs);
+        final double remainingPowerMah = result.powerMah;
         if (osBatteryConsumer != null) {
             osBatteryConsumer.setUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_WAKELOCK,
                     result.durationMs)
-                    .setConsumedPower(BatteryConsumer.POWER_COMPONENT_WAKELOCK, result.powerMah);
+                    .setConsumedPower(BatteryConsumer.POWER_COMPONENT_WAKELOCK, remainingPowerMah);
         }
 
-        final long wakeTimeMillis =
-                calculateWakeTimeMillis(batteryStats, rawRealtimeUs, rawUptimeUs);
-        final double powerMah = mPowerEstimator.calculatePower(wakeTimeMillis);
+        long wakeTimeMs = calculateWakeTimeMillis(batteryStats, rawRealtimeUs, rawUptimeUs);
+        if (wakeTimeMs < 0) {
+            wakeTimeMs = 0;
+        }
         builder.getAggregateBatteryConsumerBuilder(
                 BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE)
-                .setUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_WAKELOCK, wakeTimeMillis)
-                .setConsumedPower(BatteryConsumer.POWER_COMPONENT_WAKELOCK, powerMah);
+                .setUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_WAKELOCK,
+                        wakeTimeMs)
+                .setConsumedPower(BatteryConsumer.POWER_COMPONENT_WAKELOCK,
+                        appPowerMah + remainingPowerMah);
         builder.getAggregateBatteryConsumerBuilder(
                 BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_ALL_APPS)
-                .setConsumedPower(BatteryConsumer.POWER_COMPONENT_WAKELOCK, appPowerMah);
+                .setUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_WAKELOCK,
+                        totalAppDurationMs)
+                .setConsumedPower(BatteryConsumer.POWER_COMPONENT_WAKELOCK,
+                        appPowerMah);
     }
 
     @Override
@@ -167,9 +174,16 @@
             }
             result.durationMs = osDurationMs + wakeTimeMillis;
             result.powerMah = osPowerMah + power;
+        } else {
+            result.durationMs = 0;
+            result.powerMah = 0;
         }
     }
 
+    /**
+     * Return on-battery/screen-off time.  May be negative if the screen-on time exceeds
+     * the on-battery time.
+     */
     private long calculateWakeTimeMillis(BatteryStats batteryStats, long rawRealtimeUs,
             long rawUptimeUs) {
         final long batteryUptimeUs = batteryStats.getBatteryUptime(rawUptimeUs);
diff --git a/core/tests/coretests/src/com/android/internal/os/WakelockPowerCalculatorTest.java b/core/tests/coretests/src/com/android/internal/os/WakelockPowerCalculatorTest.java
index 82830f2..a7f4fb3 100644
--- a/core/tests/coretests/src/com/android/internal/os/WakelockPowerCalculatorTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/WakelockPowerCalculatorTest.java
@@ -79,8 +79,8 @@
         assertThat(deviceConsumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_WAKELOCK))
                 .isWithin(PRECISION).of(0.6);
 
-        BatteryConsumer appConsumer = mStatsRule.getDeviceBatteryConsumer();
+        BatteryConsumer appConsumer = mStatsRule.getAppsBatteryConsumer();
         assertThat(appConsumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_WAKELOCK))
-                .isWithin(PRECISION).of(0.6);
+                .isWithin(PRECISION).of(0.1);
     }
 }