diff --git a/core/tests/batterystatstests/BatteryUsageStatsProtoTests/Android.bp b/core/tests/batterystatstests/BatteryUsageStatsProtoTests/Android.bp
deleted file mode 100644
index 1fb5f2c..0000000
--- a/core/tests/batterystatstests/BatteryUsageStatsProtoTests/Android.bp
+++ /dev/null
@@ -1,30 +0,0 @@
-package {
-    // See: http://go/android-license-faq
-    // A large-scale-change added 'default_applicable_licenses' to import
-    // all of the 'license_kinds' from "frameworks_base_license"
-    // to get the below license kinds:
-    //   SPDX-license-identifier-Apache-2.0
-    default_applicable_licenses: ["frameworks_base_license"],
-}
-
-android_test {
-    name: "BatteryUsageStatsProtoTests",
-    srcs: ["src/**/*.java"],
-
-    static_libs: [
-        "androidx.test.rules",
-        "junit",
-        "mockito-target-minus-junit4",
-        "platform-test-annotations",
-        "platformprotosnano",
-        "statsdprotolite",
-        "truth",
-    ],
-
-    libs: ["android.test.runner"],
-
-    platform_apis: true,
-    certificate: "platform",
-
-    test_suites: ["device-tests"],
-}
diff --git a/core/tests/batterystatstests/BatteryUsageStatsProtoTests/AndroidManifest.xml b/core/tests/batterystatstests/BatteryUsageStatsProtoTests/AndroidManifest.xml
deleted file mode 100644
index 9128dca..0000000
--- a/core/tests/batterystatstests/BatteryUsageStatsProtoTests/AndroidManifest.xml
+++ /dev/null
@@ -1,28 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!--
-  ~ Copyright (C) 2021 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.
-  -->
-
-<manifest xmlns:android="http://schemas.android.com/apk/res/android"
-          package="com.android.frameworks.core.batteryusagestatsprototests">
-
-    <uses-permission android:name="android.permission.BATTERY_STATS"/>
-
-    <instrumentation
-        android:name="androidx.test.runner.AndroidJUnitRunner"
-        android:targetPackage="com.android.frameworks.core.batteryusagestatsprototests"
-        android:label="BatteryUsageStats Proto Tests" />
-
-</manifest>
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index 67985ef..092ee16 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -102,6 +102,7 @@
 
 import com.android.internal.R;
 import com.android.internal.annotations.GuardedBy;
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.app.IBatteryStats;
 import com.android.internal.os.BinderCallsStats;
 import com.android.internal.os.Clock;
@@ -1191,7 +1192,7 @@
                                     .setMinConsumedPowerThreshold(minConsumedPowerThreshold)
                                     .build();
                     bus = getBatteryUsageStats(List.of(query)).get(0);
-                    return StatsPerUidLogger.logStats(bus, data);
+                    return new StatsPerUidLogger(new FrameworkStatsLogger()).logStats(bus, data);
                 }
                 default:
                     throw new UnsupportedOperationException("Unknown tagId=" + atomTag);
@@ -1204,7 +1205,35 @@
         }
     }
 
-    private static class StatsPerUidLogger {
+    public static class FrameworkStatsLogger {
+        /**
+         * Wrapper for the FrameworkStatsLog.buildStatsEvent method that makes it easier
+         * for mocking.
+         */
+        @VisibleForTesting
+        public StatsEvent buildStatsEvent(long sessionStartTs, long sessionEndTs,
+                long sessionDuration, int sessionDischargePercentage, long sessionDischargeDuration,
+                int uid, @BatteryConsumer.ProcessState int processState, long timeInStateMillis,
+                String powerComponentName, float totalConsumedPowerMah, float powerComponentMah,
+                long powerComponentDurationMillis) {
+            return FrameworkStatsLog.buildStatsEvent(
+                    FrameworkStatsLog.BATTERY_USAGE_STATS_PER_UID,
+                    sessionStartTs,
+                    sessionEndTs,
+                    sessionDuration,
+                    sessionDischargePercentage,
+                    sessionDischargeDuration,
+                    uid,
+                    processState,
+                    timeInStateMillis,
+                    powerComponentName,
+                    totalConsumedPowerMah,
+                    powerComponentMah,
+                    powerComponentDurationMillis);
+        }
+    }
+
+    public static class StatsPerUidLogger {
 
         private static final int STATSD_METRIC_MAX_DIMENSIONS_COUNT = 3000;
 
@@ -1224,7 +1253,18 @@
                 long dischargeDuration) {}
         ;
 
-        static int logStats(BatteryUsageStats bus, List<StatsEvent> data) {
+        private final FrameworkStatsLogger mFrameworkStatsLogger;
+
+        public StatsPerUidLogger(FrameworkStatsLogger frameworkStatsLogger) {
+            mFrameworkStatsLogger = frameworkStatsLogger;
+        }
+
+        /**
+         * Generates StatsEvents for the supplied battery usage stats and adds them to
+         * the supplied list.
+         */
+        @VisibleForTesting
+        public int logStats(BatteryUsageStats bus, List<StatsEvent> data) {
             final SessionInfo sessionInfo =
                     new SessionInfo(
                             bus.getStatsStartTimestamp(),
@@ -1340,7 +1380,7 @@
             return StatsManager.PULL_SUCCESS;
         }
 
-        private static boolean addStatsForPredefinedComponent(
+        private boolean addStatsForPredefinedComponent(
                 List<StatsEvent> data,
                 SessionInfo sessionInfo,
                 int uid,
@@ -1380,7 +1420,7 @@
                     powerComponentDurationMillis);
         }
 
-        private static boolean addStatsForCustomComponent(
+        private boolean addStatsForCustomComponent(
                 List<StatsEvent> data,
                 SessionInfo sessionInfo,
                 int uid,
@@ -1422,7 +1462,7 @@
          * Returns true on success and false if reached max atoms capacity and no more atoms should
          * be added
          */
-        private static boolean addStatsAtom(
+        private boolean addStatsAtom(
                 List<StatsEvent> data,
                 SessionInfo sessionInfo,
                 int uid,
@@ -1432,9 +1472,7 @@
                 float totalConsumedPowerMah,
                 float powerComponentMah,
                 long powerComponentDurationMillis) {
-            data.add(
-                    FrameworkStatsLog.buildStatsEvent(
-                            FrameworkStatsLog.BATTERY_USAGE_STATS_PER_UID,
+            data.add(mFrameworkStatsLogger.buildStatsEvent(
                             sessionInfo.startTs(),
                             sessionInfo.endTs(),
                             sessionInfo.duration(),
diff --git a/services/tests/powerstatstests/Android.bp b/services/tests/powerstatstests/Android.bp
index f2b4136..b2a5b02 100644
--- a/services/tests/powerstatstests/Android.bp
+++ b/services/tests/powerstatstests/Android.bp
@@ -59,6 +59,7 @@
     name: "PowerStatsTestsRavenwood",
     static_libs: [
         "services.core",
+        "platformprotosnano",
         "coretests-aidl",
         "ravenwood-junit",
         "truth",
diff --git a/core/tests/batterystatstests/BatteryUsageStatsProtoTests/src/com/android/internal/os/BatteryUsageStatsPulledTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsAtomTest.java
similarity index 73%
rename from core/tests/batterystatstests/BatteryUsageStatsProtoTests/src/com/android/internal/os/BatteryUsageStatsPulledTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsAtomTest.java
index ac1f7d0..62efbc3 100644
--- a/core/tests/batterystatstests/BatteryUsageStatsProtoTests/src/com/android/internal/os/BatteryUsageStatsPulledTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsAtomTest.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.internal.os;
+package com.android.server.power.stats;
 
 import static android.os.BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE;
 
@@ -23,39 +23,262 @@
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.fail;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 
 import android.os.AggregateBatteryConsumer;
 import android.os.BatteryConsumer;
 import android.os.BatteryUsageStats;
+import android.os.Process;
 import android.os.UidBatteryConsumer;
 import android.os.nano.BatteryUsageStatsAtomsProto;
 import android.os.nano.BatteryUsageStatsAtomsProto.BatteryConsumerData.PowerComponentUsage;
+import android.platform.test.ravenwood.RavenwoodRule;
+import android.util.StatsEvent;
 
 import androidx.test.filters.SmallTest;
 
+import com.android.server.am.BatteryStatsService;
+
 import com.google.protobuf.nano.InvalidProtocolBufferNanoException;
 
+import org.junit.Rule;
 import org.junit.Test;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 
-
 @SmallTest
-public class BatteryUsageStatsPulledTest {
+public class BatteryUsageStatsAtomTest {
+
+    @Rule
+    public final RavenwoodRule mRavenwood = new RavenwoodRule();
 
     private static final int UID_0 = 1000;
     private static final int UID_1 = 2000;
     private static final int UID_2 = 3000;
     private static final int UID_3 = 4000;
-    private static final int[] UID_USAGE_TIME_PROCESS_STATES = {
-            BatteryConsumer.PROCESS_STATE_FOREGROUND,
-            BatteryConsumer.PROCESS_STATE_BACKGROUND,
-            BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE
-    };
 
     @Test
-    public void testGetStatsProto() {
+    public void testAtom_BatteryUsageStatsPerUid() {
+        final BatteryUsageStats bus = buildBatteryUsageStats();
+        BatteryStatsService.FrameworkStatsLogger statsLogger =
+                mock(BatteryStatsService.FrameworkStatsLogger.class);
+
+        List<StatsEvent> actual = new ArrayList<>();
+        new BatteryStatsService.StatsPerUidLogger(statsLogger).logStats(bus, actual);
+
+        // Device-wide totals
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                Process.INVALID_UID,
+                BatteryConsumer.PROCESS_STATE_UNSPECIFIED,
+                0L,
+                "cpu",
+                30000.0f,
+                20100.0f,
+                20300L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                Process.INVALID_UID,
+                BatteryConsumer.PROCESS_STATE_UNSPECIFIED,
+                0L,
+                "camera",
+                30000.0f,
+                20150.0f,
+                0L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                Process.INVALID_UID,
+                BatteryConsumer.PROCESS_STATE_UNSPECIFIED,
+                0L,
+                "CustomConsumer1",
+                30000.0f,
+                20200.0f,
+                20400L
+        );
+
+        // Per-proc state estimates for UID_0
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_UNSPECIFIED,
+                0L,
+                "screen",
+                1650.0f,
+                300.0f,
+                0L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_UNSPECIFIED,
+                0L,
+                "cpu",
+                1650.0f,
+                400.0f,
+                600L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND,
+                1000L,
+                "cpu",
+                1650.0f,
+                9100.0f,
+                8100L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_BACKGROUND,
+                2000L,
+                "cpu",
+                1650.0f,
+                9200.0f,
+                8200L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE,
+                0L,
+                "cpu",
+                1650.0f,
+                9300.0f,
+                8400L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_CACHED,
+                0L,
+                "cpu",
+                1650.0f,
+                9400.0f,
+                0L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND,
+                1000L,
+                "CustomConsumer1",
+                1650.0f,
+                450.0f,
+                0L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_BACKGROUND,
+                2000L,
+                "CustomConsumer1",
+                1650.0f,
+                450.0f,
+                0L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND,
+                1000L,
+                "CustomConsumer2",
+                1650.0f,
+                500.0f,
+                800L
+        );
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_0,
+                BatteryConsumer.PROCESS_STATE_BACKGROUND,
+                2000L,
+                "CustomConsumer2",
+                1650.0f,
+                500.0f,
+                800L
+        );
+
+        // Nothing for UID_1, because its power consumption is 0
+
+        // Only "screen" is populated for UID_2
+        verify(statsLogger).buildStatsEvent(
+                1000L,
+                20000L,
+                10000L,
+                20,
+                1234L,
+                UID_2,
+                BatteryConsumer.PROCESS_STATE_UNSPECIFIED,
+                0L,
+                "screen",
+                766.0f,
+                766.0f,
+                0L
+        );
+
+        verifyNoMoreInteractions(statsLogger);
+    }
+
+    @Test
+    public void testAtom_BatteryUsageStatsAtomsProto() {
         final BatteryUsageStats bus = buildBatteryUsageStats();
         final byte[] bytes = bus.getStatsProto();
         BatteryUsageStatsAtomsProto proto;
@@ -68,9 +291,7 @@
 
         assertEquals(bus.getStatsStartTimestamp(), proto.sessionStartMillis);
         assertEquals(bus.getStatsEndTimestamp(), proto.sessionEndMillis);
-        assertEquals(
-                bus.getStatsEndTimestamp() - bus.getStatsStartTimestamp(),
-                proto.sessionDurationMillis);
+        assertEquals(10000, proto.sessionDurationMillis);
         assertEquals(bus.getDischargePercentage(), proto.sessionDischargePercentage);
         assertEquals(bus.getDischargeDurationMs(), proto.dischargeDurationMillis);
 
@@ -90,8 +311,8 @@
         final List<android.os.UidBatteryConsumer> uidConsumers = bus.getUidBatteryConsumers();
         uidConsumers.sort((a, b) -> a.getUid() - b.getUid());
 
-        final BatteryUsageStatsAtomsProto.UidBatteryConsumer[] uidConsumersProto
-                = proto.uidBatteryConsumers;
+        final BatteryUsageStatsAtomsProto.UidBatteryConsumer[] uidConsumersProto =
+                proto.uidBatteryConsumers;
         Arrays.sort(uidConsumersProto, (a, b) -> a.uid - b.uid);
 
         // UID_0 - After sorting, UID_0 should be in position 0 for both data structures
@@ -186,6 +407,12 @@
         }
     }
 
+    private static final int[] UID_USAGE_TIME_PROCESS_STATES = {
+            BatteryConsumer.PROCESS_STATE_FOREGROUND,
+            BatteryConsumer.PROCESS_STATE_BACKGROUND,
+            BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE
+    };
+
     private void assertSameUidBatteryConsumer(
             android.os.UidBatteryConsumer uidConsumer,
             BatteryUsageStatsAtomsProto.UidBatteryConsumer uidConsumerProto,
@@ -195,10 +422,10 @@
         assertEquals("Uid consumers had mismatched uids", uid, uidConsumer.getUid());
 
         assertEquals("For uid " + uid,
-                uidConsumer.getTimeInStateMs(android.os.UidBatteryConsumer.STATE_FOREGROUND),
+                uidConsumer.getTimeInProcessStateMs(BatteryConsumer.PROCESS_STATE_FOREGROUND),
                 uidConsumerProto.timeInForegroundMillis);
         assertEquals("For uid " + uid,
-                uidConsumer.getTimeInStateMs(android.os.UidBatteryConsumer.STATE_BACKGROUND),
+                uidConsumer.getTimeInProcessStateMs(BatteryConsumer.PROCESS_STATE_BACKGROUND),
                 uidConsumerProto.timeInBackgroundMillis);
         for (int processState : UID_USAGE_TIME_PROCESS_STATES) {
             final long timeInStateMillis = uidConsumer.getTimeInProcessStateMs(processState);
@@ -265,7 +492,9 @@
                         .setDischargePercentage(20)
                         .setDischargedPowerRange(1000, 2000)
                         .setDischargeDurationMs(1234)
-                        .setStatsStartTimestamp(1000);
+                        .setStatsStartTimestamp(1000)
+                        .setStatsEndTimestamp(20000)
+                        .setStatsDuration(10000);
         final UidBatteryConsumer.Builder uidBuilder = builder
                 .getOrCreateUidBatteryConsumerBuilder(UID_0)
                 .setPackageWithHighestDrain("myPackage0")