diff --git a/core/tests/batterystatstests/BatteryUsageStatsProtoTests/Android.bp b/core/tests/batterystatstests/BatteryUsageStatsProtoTests/Android.bp
new file mode 100644
index 0000000..1fb5f2c
--- /dev/null
+++ b/core/tests/batterystatstests/BatteryUsageStatsProtoTests/Android.bp
@@ -0,0 +1,30 @@
+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
new file mode 100644
index 0000000..9128dca
--- /dev/null
+++ b/core/tests/batterystatstests/BatteryUsageStatsProtoTests/AndroidManifest.xml
@@ -0,0 +1,28 @@
+<?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/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsAtomTest.java b/core/tests/batterystatstests/BatteryUsageStatsProtoTests/src/com/android/internal/os/BatteryUsageStatsPulledTest.java
similarity index 73%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsAtomTest.java
rename to core/tests/batterystatstests/BatteryUsageStatsProtoTests/src/com/android/internal/os/BatteryUsageStatsPulledTest.java
index 62efbc3..ac1f7d0 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsAtomTest.java
+++ b/core/tests/batterystatstests/BatteryUsageStatsProtoTests/src/com/android/internal/os/BatteryUsageStatsPulledTest.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.server.power.stats;
+package com.android.internal.os;
 
 import static android.os.BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE;
 
@@ -23,262 +23,39 @@
 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 BatteryUsageStatsAtomTest {
 
-    @Rule
-    public final RavenwoodRule mRavenwood = new RavenwoodRule();
+@SmallTest
+public class BatteryUsageStatsPulledTest {
 
     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 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() {
+    public void testGetStatsProto() {
         final BatteryUsageStats bus = buildBatteryUsageStats();
         final byte[] bytes = bus.getStatsProto();
         BatteryUsageStatsAtomsProto proto;
@@ -291,7 +68,9 @@
 
         assertEquals(bus.getStatsStartTimestamp(), proto.sessionStartMillis);
         assertEquals(bus.getStatsEndTimestamp(), proto.sessionEndMillis);
-        assertEquals(10000, proto.sessionDurationMillis);
+        assertEquals(
+                bus.getStatsEndTimestamp() - bus.getStatsStartTimestamp(),
+                proto.sessionDurationMillis);
         assertEquals(bus.getDischargePercentage(), proto.sessionDischargePercentage);
         assertEquals(bus.getDischargeDurationMs(), proto.dischargeDurationMillis);
 
@@ -311,8 +90,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
@@ -407,12 +186,6 @@
         }
     }
 
-    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,
@@ -422,10 +195,10 @@
         assertEquals("Uid consumers had mismatched uids", uid, uidConsumer.getUid());
 
         assertEquals("For uid " + uid,
-                uidConsumer.getTimeInProcessStateMs(BatteryConsumer.PROCESS_STATE_FOREGROUND),
+                uidConsumer.getTimeInStateMs(android.os.UidBatteryConsumer.STATE_FOREGROUND),
                 uidConsumerProto.timeInForegroundMillis);
         assertEquals("For uid " + uid,
-                uidConsumer.getTimeInProcessStateMs(BatteryConsumer.PROCESS_STATE_BACKGROUND),
+                uidConsumer.getTimeInStateMs(android.os.UidBatteryConsumer.STATE_BACKGROUND),
                 uidConsumerProto.timeInBackgroundMillis);
         for (int processState : UID_USAGE_TIME_PROCESS_STATES) {
             final long timeInStateMillis = uidConsumer.getTimeInProcessStateMs(processState);
@@ -492,9 +265,7 @@
                         .setDischargePercentage(20)
                         .setDischargedPowerRange(1000, 2000)
                         .setDischargeDurationMs(1234)
-                        .setStatsStartTimestamp(1000)
-                        .setStatsEndTimestamp(20000)
-                        .setStatsDuration(10000);
+                        .setStatsStartTimestamp(1000);
         final UidBatteryConsumer.Builder uidBuilder = builder
                 .getOrCreateUidBatteryConsumerBuilder(UID_0)
                 .setPackageWithHighestDrain("myPackage0")
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index 092ee16..67985ef 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -102,7 +102,6 @@
 
 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;
@@ -1192,7 +1191,7 @@
                                     .setMinConsumedPowerThreshold(minConsumedPowerThreshold)
                                     .build();
                     bus = getBatteryUsageStats(List.of(query)).get(0);
-                    return new StatsPerUidLogger(new FrameworkStatsLogger()).logStats(bus, data);
+                    return StatsPerUidLogger.logStats(bus, data);
                 }
                 default:
                     throw new UnsupportedOperationException("Unknown tagId=" + atomTag);
@@ -1205,35 +1204,7 @@
         }
     }
 
-    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 class StatsPerUidLogger {
 
         private static final int STATSD_METRIC_MAX_DIMENSIONS_COUNT = 3000;
 
@@ -1253,18 +1224,7 @@
                 long dischargeDuration) {}
         ;
 
-        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) {
+        static int logStats(BatteryUsageStats bus, List<StatsEvent> data) {
             final SessionInfo sessionInfo =
                     new SessionInfo(
                             bus.getStatsStartTimestamp(),
@@ -1380,7 +1340,7 @@
             return StatsManager.PULL_SUCCESS;
         }
 
-        private boolean addStatsForPredefinedComponent(
+        private static boolean addStatsForPredefinedComponent(
                 List<StatsEvent> data,
                 SessionInfo sessionInfo,
                 int uid,
@@ -1420,7 +1380,7 @@
                     powerComponentDurationMillis);
         }
 
-        private boolean addStatsForCustomComponent(
+        private static boolean addStatsForCustomComponent(
                 List<StatsEvent> data,
                 SessionInfo sessionInfo,
                 int uid,
@@ -1462,7 +1422,7 @@
          * Returns true on success and false if reached max atoms capacity and no more atoms should
          * be added
          */
-        private boolean addStatsAtom(
+        private static boolean addStatsAtom(
                 List<StatsEvent> data,
                 SessionInfo sessionInfo,
                 int uid,
@@ -1472,7 +1432,9 @@
                 float totalConsumedPowerMah,
                 float powerComponentMah,
                 long powerComponentDurationMillis) {
-            data.add(mFrameworkStatsLogger.buildStatsEvent(
+            data.add(
+                    FrameworkStatsLog.buildStatsEvent(
+                            FrameworkStatsLog.BATTERY_USAGE_STATS_PER_UID,
                             sessionInfo.startTs(),
                             sessionInfo.endTs(),
                             sessionInfo.duration(),
diff --git a/services/tests/powerstatstests/Android.bp b/services/tests/powerstatstests/Android.bp
index b2a5b02..f2b4136 100644
--- a/services/tests/powerstatstests/Android.bp
+++ b/services/tests/powerstatstests/Android.bp
@@ -59,7 +59,6 @@
     name: "PowerStatsTestsRavenwood",
     static_libs: [
         "services.core",
-        "platformprotosnano",
         "coretests-aidl",
         "ravenwood-junit",
         "truth",