Merge "Make WakelockStatsFrameworkEvents test work with ravenwood" into main
diff --git a/services/core/java/com/android/server/power/stats/WakelockStatsFrameworkEvents.java b/services/core/java/com/android/server/power/stats/WakelockStatsFrameworkEvents.java
index c9693bd..f387fec 100644
--- a/services/core/java/com/android/server/power/stats/WakelockStatsFrameworkEvents.java
+++ b/services/core/java/com/android/server/power/stats/WakelockStatsFrameworkEvents.java
@@ -150,26 +150,14 @@
     }
 
     @VisibleForTesting
-    public boolean inOverflow() {
-        synchronized (mLock) {
-            return inOverflowLocked();
-        }
-    }
-
     @GuardedBy("mLock")
-    private boolean inOverflowLocked() {
+    public boolean inOverflow() {
         return mWakeLockStats.size() >= SUMMARY_THRESHOLD;
     }
 
     @VisibleForTesting
-    public boolean inHardCap() {
-        synchronized (mLock) {
-            return inHardCapLocked();
-        }
-    }
-
     @GuardedBy("mLock")
-    private boolean inHardCapLocked() {
+    public boolean inHardCap() {
         return mWakeLockStats.size() >= MAX_WAKELOCK_DIMENSIONS;
     }
 
@@ -189,9 +177,9 @@
                 long wakeLockDur = eventUptimeMillis - data.acquireUptimeMillis;
 
                 // Rewrite key if in an overflow state.
-                if (inOverflowLocked() && !mWakeLockStats.containsKey(key)) {
+                if (inOverflow() && !mWakeLockStats.containsKey(key)) {
                     key.setOverflow();
-                    if (inHardCapLocked() && !mWakeLockStats.containsKey(key)) {
+                    if (inHardCap() && !mWakeLockStats.containsKey(key)) {
                         key.setHardCap();
                     }
                 }
@@ -207,12 +195,41 @@
         }
     }
 
-    public List<StatsEvent> pullFrameworkWakelockInfoAtoms() {
-        return pullFrameworkWakelockInfoAtoms(SystemClock.uptimeMillis());
+    // Shim interface for testing.
+    @VisibleForTesting
+    public interface EventLogger {
+        void logResult(
+                int uid, String tag, int wakeLockLevel, long uptimeMillis, long completedCount);
     }
 
-    public List<StatsEvent> pullFrameworkWakelockInfoAtoms(long nowMillis) {
+    public List<StatsEvent> pullFrameworkWakelockInfoAtoms() {
         List<StatsEvent> result = new ArrayList<>();
+        EventLogger logger =
+                new EventLogger() {
+                    public void logResult(
+                            int uid,
+                            String tag,
+                            int wakeLockLevel,
+                            long uptimeMillis,
+                            long completedCount) {
+                        StatsEvent event =
+                                StatsEvent.newBuilder()
+                                        .setAtomId(FrameworkStatsLog.FRAMEWORK_WAKELOCK_INFO)
+                                        .writeInt(uid)
+                                        .writeString(tag)
+                                        .writeInt(wakeLockLevel)
+                                        .writeLong(uptimeMillis)
+                                        .writeLong(completedCount)
+                                        .build();
+                        result.add(event);
+                    }
+                };
+        pullFrameworkWakelockInfoAtoms(SystemClock.uptimeMillis(), logger);
+        return result;
+    }
+
+    @VisibleForTesting
+    public void pullFrameworkWakelockInfoAtoms(long nowMillis, EventLogger logger) {
         HashSet<WakeLockKey> keys = new HashSet<>();
 
         // Used to collect open WakeLocks when in an overflow state.
@@ -223,13 +240,13 @@
 
             // If we are in an overflow state, an open wakelock may have a new key
             // that needs to be summarized.
-            if (inOverflowLocked()) {
+            if (inOverflow()) {
                 for (WakeLockKey key : mOpenWakeLocks.keySet()) {
                     if (!mWakeLockStats.containsKey(key)) {
                         WakeLockData data = mOpenWakeLocks.get(key);
 
                         key.setOverflow();
-                        if (inHardCapLocked() && !mWakeLockStats.containsKey(key)) {
+                        if (inHardCap() && !mWakeLockStats.containsKey(key)) {
                             key.setHardCap();
                         }
                         keys.add(key);
@@ -257,20 +274,14 @@
 
                 stats.uptimeMillis += openWakeLockUptime + extraTime.uptimeMillis;
 
-                StatsEvent event =
-                        StatsEvent.newBuilder()
-                                .setAtomId(FrameworkStatsLog.FRAMEWORK_WAKELOCK_INFO)
-                                .writeInt(key.getUid())
-                                .writeString(key.getTag())
-                                .writeInt(key.getPowerManagerWakeLockLevel())
-                                .writeLong(stats.uptimeMillis)
-                                .writeLong(stats.completedCount)
-                                .build();
-                result.add(event);
+                logger.logResult(
+                        key.getUid(),
+                        key.getTag(),
+                        key.getPowerManagerWakeLockLevel(),
+                        stats.uptimeMillis,
+                        stats.completedCount);
             }
         }
-
-        return result;
     }
 
     private static final String TAG = "BatteryStatsPulledMetrics";
diff --git a/services/tests/powerstatstests/Android.bp b/services/tests/powerstatstests/Android.bp
index d9e071f..d6ca10a 100644
--- a/services/tests/powerstatstests/Android.bp
+++ b/services/tests/powerstatstests/Android.bp
@@ -27,9 +27,6 @@
         "servicestests-utils",
         "platform-test-annotations",
         "flag-junit",
-        "statsdprotolite",
-        "StatsdTestUtils",
-        "platformprotoslite",
     ],
 
     libs: [
@@ -77,10 +74,6 @@
         "src/com/android/server/power/stats/format/*.java",
         "src/com/android/server/power/stats/processor/*.java",
     ],
-    // TODO(b/372292543): Enable this test.
-    exclude_srcs: [
-        "src/com/android/server/power/stats/WakelockStatsFrameworkEventsTest.java",
-    ],
     java_resources: [
         "res/xml/power_profile*.xml",
     ],
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockPowerStatsCollectorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockPowerStatsCollectorTest.java
index cef3fdd..03491bc 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockPowerStatsCollectorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockPowerStatsCollectorTest.java
@@ -25,7 +25,6 @@
 import android.content.Context;
 import android.os.Process;
 import android.platform.test.annotations.DisableFlags;
-import android.platform.test.annotations.DisabledOnRavenwood;
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.platform.test.ravenwood.RavenwoodConfig;
 import android.platform.test.ravenwood.RavenwoodConfig.Config;
@@ -76,7 +75,6 @@
 
     @Test
     @DisableFlags(Flags.FLAG_FRAMEWORK_WAKELOCK_INFO)
-    @DisabledOnRavenwood(reason = "b/372292543 temporary disable")
     public void collectStats() {
         PowerStatsCollector powerStatsCollector = mBatteryStats.getPowerStatsCollector(
                 POWER_COMPONENT_WAKELOCK);
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockStatsFrameworkEventsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockStatsFrameworkEventsTest.java
index cb644db..1fe3f58 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockStatsFrameworkEventsTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockStatsFrameworkEventsTest.java
@@ -19,41 +19,25 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
-import android.os.WakeLockLevelEnum;
-import android.util.StatsEvent;
-import android.util.StatsEventTestUtils;
+import android.os.nano.OsProtoEnums;
 
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
-import com.android.os.AtomsProto;
-import com.android.os.framework.FrameworkExtensionAtoms;
-import com.android.os.framework.FrameworkExtensionAtoms.FrameworkWakelockInfo;
-
-import com.google.protobuf.CodedInputStream;
-import com.google.protobuf.CodedOutputStream;
-import com.google.protobuf.ExtensionRegistryLite;
-
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
 import java.util.ArrayList;
-import java.util.List;
 
 @RunWith(AndroidJUnit4.class)
 @SmallTest
 public class WakelockStatsFrameworkEventsTest {
     private WakelockStatsFrameworkEvents mEvents;
-    private ExtensionRegistryLite mRegistry;
 
     @Before
     public void setup() {
         mEvents = new WakelockStatsFrameworkEvents();
-        mRegistry = ExtensionRegistryLite.newInstance();
-        FrameworkExtensionAtoms.registerAllExtensions(mRegistry);
     }
 
     private static final int UID_1 = 1;
@@ -62,8 +46,8 @@
     private static final String TAG_1 = "TAG1";
     private static final String TAG_2 = "TAG2";
 
-    private static final WakeLockLevelEnum WAKELOCK_TYPE_1 = WakeLockLevelEnum.PARTIAL_WAKE_LOCK;
-    private static final WakeLockLevelEnum WAKELOCK_TYPE_2 = WakeLockLevelEnum.DOZE_WAKE_LOCK;
+    private static final int WAKELOCK_TYPE_1 = OsProtoEnums.PARTIAL_WAKE_LOCK;
+    private static final int WAKELOCK_TYPE_2 = OsProtoEnums.DOZE_WAKE_LOCK;
 
     private static final long TS_1 = 1000;
     private static final long TS_2 = 2000;
@@ -71,19 +55,37 @@
     private static final long TS_4 = 4000;
     private static final long TS_5 = 5000;
 
+    // Mirrors com.android.os.framework.FrameworkWakelockInfo proto.
+    private static class WakelockInfo {
+        public int uid;
+        public String tag;
+        public int type;
+        public long uptimeMillis;
+        public long completedCount;
+
+        WakelockInfo(int uid, String tag, int type, long uptimeMillis, long completedCount) {
+            this.uid = uid;
+            this.tag = tag;
+            this.type = type;
+            this.uptimeMillis = uptimeMillis;
+            this.completedCount = completedCount;
+        }
+    }
+
     // Assumes that mEvents is empty.
+    @SuppressWarnings("GuardedBy")
     private void makeMetricsAlmostOverflow() throws Exception {
         for (int i = 0; i < mEvents.SUMMARY_THRESHOLD - 1; i++) {
             String tag = "forceOverflow" + i;
-            mEvents.noteStartWakeLock(UID_1, tag, WAKELOCK_TYPE_1.getNumber(), TS_1);
-            mEvents.noteStopWakeLock(UID_1, tag, WAKELOCK_TYPE_1.getNumber(), TS_2);
+            mEvents.noteStartWakeLock(UID_1, tag, WAKELOCK_TYPE_1, TS_1);
+            mEvents.noteStopWakeLock(UID_1, tag, WAKELOCK_TYPE_1, TS_2);
         }
 
         assertFalse("not overflow", mEvents.inOverflow());
-        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
-        FrameworkWakelockInfo notOverflowInfo =
+        ArrayList<WakelockInfo> info = pullResults(TS_4);
+        WakelockInfo notOverflowInfo =
                 info.stream()
-                        .filter(i -> i.getAttributionTag().equals(mEvents.OVERFLOW_TAG))
+                        .filter(i -> i.tag.equals(mEvents.OVERFLOW_TAG))
                         .findFirst()
                         .orElse(null);
 
@@ -91,38 +93,36 @@
 
         // Add one more to hit an overflow state.
         String lastTag = "forceOverflowLast";
-        mEvents.noteStartWakeLock(UID_1, lastTag, WAKELOCK_TYPE_2.getNumber(), TS_1);
-        mEvents.noteStopWakeLock(UID_1, lastTag, WAKELOCK_TYPE_2.getNumber(), TS_2);
+        mEvents.noteStartWakeLock(UID_1, lastTag, WAKELOCK_TYPE_2, TS_1);
+        mEvents.noteStopWakeLock(UID_1, lastTag, WAKELOCK_TYPE_2, TS_2);
 
         assertTrue("overflow", mEvents.inOverflow());
         info = pullResults(TS_4);
 
-        FrameworkWakelockInfo tag1Info =
-                info.stream()
-                        .filter(i -> i.getAttributionTag().equals(lastTag))
-                        .findFirst()
-                        .orElse(null);
+        WakelockInfo tag1Info =
+                info.stream().filter(i -> i.tag.equals(lastTag)).findFirst().orElse(null);
 
         assertTrue("lastTag found", tag1Info != null);
-        assertEquals("uid", UID_1, tag1Info.getAttributionUid());
-        assertEquals("tag", lastTag, tag1Info.getAttributionTag());
-        assertEquals("type", WAKELOCK_TYPE_2, tag1Info.getType());
-        assertEquals("duration", TS_2 - TS_1, tag1Info.getUptimeMillis());
-        assertEquals("count", 1, tag1Info.getCompletedCount());
+        assertEquals("uid", UID_1, tag1Info.uid);
+        assertEquals("tag", lastTag, tag1Info.tag);
+        assertEquals("type", WAKELOCK_TYPE_2, tag1Info.type);
+        assertEquals("duration", TS_2 - TS_1, tag1Info.uptimeMillis);
+        assertEquals("count", 1, tag1Info.completedCount);
     }
 
     // Assumes that mEvents is empty.
+    @SuppressWarnings("GuardedBy")
     private void makeMetricsAlmostHardCap() throws Exception {
         for (int i = 0; i < mEvents.MAX_WAKELOCK_DIMENSIONS - 1; i++) {
-            mEvents.noteStartWakeLock(i /* uid */, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
-            mEvents.noteStopWakeLock(i /* uid */, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_2);
+            mEvents.noteStartWakeLock(i /* uid */, TAG_1, WAKELOCK_TYPE_1, TS_1);
+            mEvents.noteStopWakeLock(i /* uid */, TAG_1, WAKELOCK_TYPE_1, TS_2);
         }
 
         assertFalse("not hard capped", mEvents.inHardCap());
-        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
-        FrameworkWakelockInfo notOverflowInfo =
+        ArrayList<WakelockInfo> info = pullResults(TS_4);
+        WakelockInfo notOverflowInfo =
                 info.stream()
-                        .filter(i -> i.getAttributionTag().equals(mEvents.HARD_CAP_TAG))
+                        .filter(i -> i.tag.equals(mEvents.HARD_CAP_TAG))
                         .findFirst()
                         .orElse(null);
 
@@ -130,98 +130,86 @@
 
         // Add one more to hit an hardcap state.
         int hardCapUid = mEvents.MAX_WAKELOCK_DIMENSIONS;
-        mEvents.noteStartWakeLock(hardCapUid, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_1);
-        mEvents.noteStopWakeLock(hardCapUid, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_2);
+        mEvents.noteStartWakeLock(hardCapUid, TAG_2, WAKELOCK_TYPE_2, TS_1);
+        mEvents.noteStopWakeLock(hardCapUid, TAG_2, WAKELOCK_TYPE_2, TS_2);
 
         assertTrue("hard capped", mEvents.inHardCap());
         info = pullResults(TS_4);
 
-        FrameworkWakelockInfo tag2Info =
-                info.stream()
-                        .filter(i -> i.getAttributionUid() == hardCapUid)
-                        .findFirst()
-                        .orElse(null);
+        WakelockInfo tag2Info =
+                info.stream().filter(i -> i.uid == hardCapUid).findFirst().orElse(null);
 
         assertTrue("hardCapUid found", tag2Info != null);
-        assertEquals("uid", hardCapUid, tag2Info.getAttributionUid());
-        assertEquals("tag", mEvents.OVERFLOW_TAG, tag2Info.getAttributionTag());
-        assertEquals(
-                "type", WakeLockLevelEnum.forNumber(mEvents.OVERFLOW_LEVEL), tag2Info.getType());
-        assertEquals("duration", TS_2 - TS_1, tag2Info.getUptimeMillis());
-        assertEquals("count", 1, tag2Info.getCompletedCount());
+        assertEquals("uid", hardCapUid, tag2Info.uid);
+        assertEquals("tag", mEvents.OVERFLOW_TAG, tag2Info.tag);
+        assertEquals("type", mEvents.OVERFLOW_LEVEL, tag2Info.type);
+        assertEquals("duration", TS_2 - TS_1, tag2Info.uptimeMillis);
+        assertEquals("count", 1, tag2Info.completedCount);
     }
 
-    private ArrayList<FrameworkWakelockInfo> pullResults(long timestamp) throws Exception {
-        ArrayList<FrameworkWakelockInfo> result = new ArrayList<>();
-        List<StatsEvent> events = mEvents.pullFrameworkWakelockInfoAtoms(timestamp);
-
-        for (StatsEvent e : events) {
-            // The returned atom does not have external extensions registered.
-            // So we serialize and then deserialize with extensions registered.
-            AtomsProto.Atom atom = StatsEventTestUtils.convertToAtom(e);
-
-            ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
-            CodedOutputStream codedos = CodedOutputStream.newInstance(outputStream);
-            atom.writeTo(codedos);
-            codedos.flush();
-
-            ByteArrayInputStream inputStream = new ByteArrayInputStream(outputStream.toByteArray());
-            CodedInputStream codedis = CodedInputStream.newInstance(inputStream);
-            AtomsProto.Atom atomWithExtensions = AtomsProto.Atom.parseFrom(codedis, mRegistry);
-
-            assertTrue(
-                    atomWithExtensions.hasExtension(FrameworkExtensionAtoms.frameworkWakelockInfo));
-            FrameworkWakelockInfo info =
-                    atomWithExtensions.getExtension(FrameworkExtensionAtoms.frameworkWakelockInfo);
-            result.add(info);
-        }
-
-        return result;
+    private ArrayList<WakelockInfo> pullResults(long timestamp) {
+        ArrayList<WakelockInfo> results = new ArrayList<>();
+        WakelockStatsFrameworkEvents.EventLogger logger =
+                new WakelockStatsFrameworkEvents.EventLogger() {
+                    public void logResult(
+                            int uid,
+                            String tag,
+                            int wakeLockLevel,
+                            long uptimeMillis,
+                            long completedCount) {
+                        WakelockInfo info =
+                                new WakelockInfo(
+                                        uid, tag, wakeLockLevel, uptimeMillis, completedCount);
+                        results.add(info);
+                    }
+                };
+        mEvents.pullFrameworkWakelockInfoAtoms(timestamp, logger);
+        return results;
     }
 
     @Test
-    public void singleWakelock() throws Exception {
-        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
-        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_2);
+    public void singleWakelock() {
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1, TS_1);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1, TS_2);
 
-        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_3);
+        ArrayList<WakelockInfo> info = pullResults(TS_3);
 
         assertEquals("size", 1, info.size());
-        assertEquals("uid", UID_1, info.get(0).getAttributionUid());
-        assertEquals("tag", TAG_1, info.get(0).getAttributionTag());
-        assertEquals("type", WAKELOCK_TYPE_1, info.get(0).getType());
-        assertEquals("duration", TS_2 - TS_1, info.get(0).getUptimeMillis());
-        assertEquals("count", 1, info.get(0).getCompletedCount());
+        assertEquals("uid", UID_1, info.get(0).uid);
+        assertEquals("tag", TAG_1, info.get(0).tag);
+        assertEquals("type", WAKELOCK_TYPE_1, info.get(0).type);
+        assertEquals("duration", TS_2 - TS_1, info.get(0).uptimeMillis);
+        assertEquals("count", 1, info.get(0).completedCount);
     }
 
     @Test
     public void wakelockOpen() throws Exception {
-        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1, TS_1);
 
-        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_3);
+        ArrayList<WakelockInfo> info = pullResults(TS_3);
 
         assertEquals("size", 1, info.size());
-        assertEquals("uid", UID_1, info.get(0).getAttributionUid());
-        assertEquals("tag", TAG_1, info.get(0).getAttributionTag());
-        assertEquals("type", WAKELOCK_TYPE_1, info.get(0).getType());
-        assertEquals("duration", TS_3 - TS_1, info.get(0).getUptimeMillis());
-        assertEquals("count", 0, info.get(0).getCompletedCount());
+        assertEquals("uid", UID_1, info.get(0).uid);
+        assertEquals("tag", TAG_1, info.get(0).tag);
+        assertEquals("type", WAKELOCK_TYPE_1, info.get(0).type);
+        assertEquals("duration", TS_3 - TS_1, info.get(0).uptimeMillis);
+        assertEquals("count", 0, info.get(0).completedCount);
     }
 
     @Test
     public void wakelockOpenOverlap() throws Exception {
-        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
-        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_2);
-        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_3);
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1, TS_1);
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1, TS_2);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1, TS_3);
 
-        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
+        ArrayList<WakelockInfo> info = pullResults(TS_4);
 
         assertEquals("size", 1, info.size());
-        assertEquals("uid", UID_1, info.get(0).getAttributionUid());
-        assertEquals("tag", TAG_1, info.get(0).getAttributionTag());
-        assertEquals("type", WAKELOCK_TYPE_1, info.get(0).getType());
-        assertEquals("duration", TS_4 - TS_1, info.get(0).getUptimeMillis());
-        assertEquals("count", 0, info.get(0).getCompletedCount());
+        assertEquals("uid", UID_1, info.get(0).uid);
+        assertEquals("tag", TAG_1, info.get(0).tag);
+        assertEquals("type", WAKELOCK_TYPE_1, info.get(0).type);
+        assertEquals("duration", TS_4 - TS_1, info.get(0).uptimeMillis);
+        assertEquals("count", 0, info.get(0).completedCount);
     }
 
     @Test
@@ -229,23 +217,20 @@
         makeMetricsAlmostOverflow();
 
         // This one gets tagged as an overflow.
-        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_1);
-        mEvents.noteStopWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_2);
+        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2, TS_1);
+        mEvents.noteStopWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2, TS_2);
 
-        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
-        FrameworkWakelockInfo overflowInfo =
+        ArrayList<WakelockInfo> info = pullResults(TS_4);
+        WakelockInfo overflowInfo =
                 info.stream()
-                        .filter(i -> i.getAttributionTag().equals(mEvents.OVERFLOW_TAG))
+                        .filter(i -> i.tag.equals(mEvents.OVERFLOW_TAG))
                         .findFirst()
                         .orElse(null);
 
-        assertEquals("uid", UID_1, overflowInfo.getAttributionUid());
-        assertEquals(
-                "type",
-                WakeLockLevelEnum.forNumber(mEvents.OVERFLOW_LEVEL),
-                overflowInfo.getType());
-        assertEquals("duration", TS_2 - TS_1, overflowInfo.getUptimeMillis());
-        assertEquals("count", 1, overflowInfo.getCompletedCount());
+        assertEquals("uid", UID_1, overflowInfo.uid);
+        assertEquals("type", mEvents.OVERFLOW_LEVEL, overflowInfo.type);
+        assertEquals("duration", TS_2 - TS_1, overflowInfo.uptimeMillis);
+        assertEquals("count", 1, overflowInfo.completedCount);
     }
 
     @Test
@@ -253,22 +238,19 @@
         makeMetricsAlmostOverflow();
 
         // This is the open wakelock that overflows.
-        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_1);
+        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2, TS_1);
 
-        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
-        FrameworkWakelockInfo overflowInfo =
+        ArrayList<WakelockInfo> info = pullResults(TS_4);
+        WakelockInfo overflowInfo =
                 info.stream()
-                        .filter(i -> i.getAttributionTag().equals(mEvents.OVERFLOW_TAG))
+                        .filter(i -> i.tag.equals(mEvents.OVERFLOW_TAG))
                         .findFirst()
                         .orElse(null);
 
-        assertEquals("uid", UID_1, overflowInfo.getAttributionUid());
-        assertEquals(
-                "type",
-                WakeLockLevelEnum.forNumber(mEvents.OVERFLOW_LEVEL),
-                overflowInfo.getType());
-        assertEquals("duration", (TS_4 - TS_1), overflowInfo.getUptimeMillis());
-        assertEquals("count", 0, overflowInfo.getCompletedCount());
+        assertEquals("uid", UID_1, overflowInfo.uid);
+        assertEquals("type", mEvents.OVERFLOW_LEVEL, overflowInfo.type);
+        assertEquals("duration", (TS_4 - TS_1), overflowInfo.uptimeMillis);
+        assertEquals("count", 0, overflowInfo.completedCount);
     }
 
     @Test
@@ -276,23 +258,20 @@
         makeMetricsAlmostHardCap();
 
         // This one gets tagged as a hard cap.
-        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_1);
-        mEvents.noteStopWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_2);
+        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2, TS_1);
+        mEvents.noteStopWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2, TS_2);
 
-        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
-        FrameworkWakelockInfo hardCapInfo =
+        ArrayList<WakelockInfo> info = pullResults(TS_4);
+        WakelockInfo hardCapInfo =
                 info.stream()
-                        .filter(i -> i.getAttributionTag().equals(mEvents.HARD_CAP_TAG))
+                        .filter(i -> i.tag.equals(mEvents.HARD_CAP_TAG))
                         .findFirst()
                         .orElse(null);
 
-        assertEquals("uid", mEvents.HARD_CAP_UID, hardCapInfo.getAttributionUid());
-        assertEquals(
-                "type",
-                WakeLockLevelEnum.forNumber(mEvents.OVERFLOW_LEVEL),
-                hardCapInfo.getType());
-        assertEquals("duration", TS_2 - TS_1, hardCapInfo.getUptimeMillis());
-        assertEquals("count", 1, hardCapInfo.getCompletedCount());
+        assertEquals("uid", mEvents.HARD_CAP_UID, hardCapInfo.uid);
+        assertEquals("type", mEvents.OVERFLOW_LEVEL, hardCapInfo.type);
+        assertEquals("duration", TS_2 - TS_1, hardCapInfo.uptimeMillis);
+        assertEquals("count", 1, hardCapInfo.completedCount);
     }
 
     @Test
@@ -300,134 +279,123 @@
         makeMetricsAlmostHardCap();
 
         // This is the open wakelock that overflows.
-        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_1);
+        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2, TS_1);
 
-        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
-        FrameworkWakelockInfo hardCapInfo =
+        ArrayList<WakelockInfo> info = pullResults(TS_4);
+        WakelockInfo hardCapInfo =
                 info.stream()
-                        .filter(i -> i.getAttributionTag().equals(mEvents.HARD_CAP_TAG))
+                        .filter(i -> i.tag.equals(mEvents.HARD_CAP_TAG))
                         .findFirst()
                         .orElse(null);
 
-        assertEquals("uid", mEvents.HARD_CAP_UID, hardCapInfo.getAttributionUid());
-        assertEquals(
-                "type",
-                WakeLockLevelEnum.forNumber(mEvents.OVERFLOW_LEVEL),
-                hardCapInfo.getType());
-        assertEquals("duration", (TS_4 - TS_1), hardCapInfo.getUptimeMillis());
-        assertEquals("count", 0, hardCapInfo.getCompletedCount());
+        assertEquals("uid", mEvents.HARD_CAP_UID, hardCapInfo.uid);
+        assertEquals("type", mEvents.OVERFLOW_LEVEL, hardCapInfo.type);
+        assertEquals("duration", (TS_4 - TS_1), hardCapInfo.uptimeMillis);
+        assertEquals("count", 0, hardCapInfo.completedCount);
     }
 
     @Test
     public void overlappingWakelocks() throws Exception {
-        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
-        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_2);
-        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_3);
-        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_4);
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1, TS_1);
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1, TS_2);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1, TS_3);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1, TS_4);
 
-        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_5);
+        ArrayList<WakelockInfo> info = pullResults(TS_5);
 
         assertEquals("size", 1, info.size());
-        assertEquals("uid", UID_1, info.get(0).getAttributionUid());
-        assertEquals("tag", TAG_1, info.get(0).getAttributionTag());
-        assertEquals("type", WAKELOCK_TYPE_1, info.get(0).getType());
-        assertEquals("duration", TS_4 - TS_1, info.get(0).getUptimeMillis());
-        assertEquals("count", 1, info.get(0).getCompletedCount());
+        assertEquals("uid", UID_1, info.get(0).uid);
+        assertEquals("tag", TAG_1, info.get(0).tag);
+        assertEquals("type", WAKELOCK_TYPE_1, info.get(0).type);
+        assertEquals("duration", TS_4 - TS_1, info.get(0).uptimeMillis);
+        assertEquals("count", 1, info.get(0).completedCount);
     }
 
     @Test
     public void diffUid() throws Exception {
-        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
-        mEvents.noteStartWakeLock(UID_2, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_2);
-        mEvents.noteStopWakeLock(UID_2, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_3);
-        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_4);
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1, TS_1);
+        mEvents.noteStartWakeLock(UID_2, TAG_1, WAKELOCK_TYPE_1, TS_2);
+        mEvents.noteStopWakeLock(UID_2, TAG_1, WAKELOCK_TYPE_1, TS_3);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1, TS_4);
 
-        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_5);
+        ArrayList<WakelockInfo> info = pullResults(TS_5);
         assertEquals("size", 2, info.size());
 
-        FrameworkWakelockInfo uid1Info =
-                info.stream().filter(i -> i.getAttributionUid() == UID_1).findFirst().orElse(null);
+        WakelockInfo uid1Info = info.stream().filter(i -> i.uid == UID_1).findFirst().orElse(null);
 
         assertTrue("UID_1 found", uid1Info != null);
-        assertEquals("uid", UID_1, uid1Info.getAttributionUid());
-        assertEquals("tag", TAG_1, uid1Info.getAttributionTag());
-        assertEquals("type", WAKELOCK_TYPE_1, uid1Info.getType());
-        assertEquals("duration", TS_4 - TS_1, uid1Info.getUptimeMillis());
-        assertEquals("count", 1, uid1Info.getCompletedCount());
+        assertEquals("uid", UID_1, uid1Info.uid);
+        assertEquals("tag", TAG_1, uid1Info.tag);
+        assertEquals("type", WAKELOCK_TYPE_1, uid1Info.type);
+        assertEquals("duration", TS_4 - TS_1, uid1Info.uptimeMillis);
+        assertEquals("count", 1, uid1Info.completedCount);
 
-        FrameworkWakelockInfo uid2Info =
-                info.stream().filter(i -> i.getAttributionUid() == UID_2).findFirst().orElse(null);
+        WakelockInfo uid2Info = info.stream().filter(i -> i.uid == UID_2).findFirst().orElse(null);
         assertTrue("UID_2 found", uid2Info != null);
-        assertEquals("uid", UID_2, uid2Info.getAttributionUid());
-        assertEquals("tag", TAG_1, uid2Info.getAttributionTag());
-        assertEquals("type", WAKELOCK_TYPE_1, uid2Info.getType());
-        assertEquals("duration", TS_3 - TS_2, uid2Info.getUptimeMillis());
-        assertEquals("count", 1, uid2Info.getCompletedCount());
+        assertEquals("uid", UID_2, uid2Info.uid);
+        assertEquals("tag", TAG_1, uid2Info.tag);
+        assertEquals("type", WAKELOCK_TYPE_1, uid2Info.type);
+        assertEquals("duration", TS_3 - TS_2, uid2Info.uptimeMillis);
+        assertEquals("count", 1, uid2Info.completedCount);
     }
 
     @Test
     public void diffTag() throws Exception {
-        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
-        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_1.getNumber(), TS_2);
-        mEvents.noteStopWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_1.getNumber(), TS_3);
-        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_4);
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1, TS_1);
+        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_1, TS_2);
+        mEvents.noteStopWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_1, TS_3);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1, TS_4);
 
-        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_5);
+        ArrayList<WakelockInfo> info = pullResults(TS_5);
         assertEquals("size", 2, info.size());
 
-        FrameworkWakelockInfo uid1Info =
-                info.stream()
-                        .filter(i -> i.getAttributionTag().equals(TAG_1))
-                        .findFirst()
-                        .orElse(null);
+        WakelockInfo uid1Info =
+                info.stream().filter(i -> i.tag.equals(TAG_1)).findFirst().orElse(null);
 
         assertTrue("TAG_1 found", uid1Info != null);
-        assertEquals("uid", UID_1, uid1Info.getAttributionUid());
-        assertEquals("tag", TAG_1, uid1Info.getAttributionTag());
-        assertEquals("type", WAKELOCK_TYPE_1, uid1Info.getType());
-        assertEquals("duration", TS_4 - TS_1, uid1Info.getUptimeMillis());
-        assertEquals("count", 1, uid1Info.getCompletedCount());
+        assertEquals("uid", UID_1, uid1Info.uid);
+        assertEquals("tag", TAG_1, uid1Info.tag);
+        assertEquals("type", WAKELOCK_TYPE_1, uid1Info.type);
+        assertEquals("duration", TS_4 - TS_1, uid1Info.uptimeMillis);
+        assertEquals("count", 1, uid1Info.completedCount);
 
-        FrameworkWakelockInfo uid2Info =
-                info.stream()
-                        .filter(i -> i.getAttributionTag().equals(TAG_2))
-                        .findFirst()
-                        .orElse(null);
+        WakelockInfo uid2Info =
+                info.stream().filter(i -> i.tag.equals(TAG_2)).findFirst().orElse(null);
         assertTrue("TAG_2 found", uid2Info != null);
-        assertEquals("uid", UID_1, uid2Info.getAttributionUid());
-        assertEquals("tag", TAG_2, uid2Info.getAttributionTag());
-        assertEquals("type", WAKELOCK_TYPE_1, uid2Info.getType());
-        assertEquals("duration", TS_3 - TS_2, uid2Info.getUptimeMillis());
-        assertEquals("count", 1, uid2Info.getCompletedCount());
+        assertEquals("uid", UID_1, uid2Info.uid);
+        assertEquals("tag", TAG_2, uid2Info.tag);
+        assertEquals("type", WAKELOCK_TYPE_1, uid2Info.type);
+        assertEquals("duration", TS_3 - TS_2, uid2Info.uptimeMillis);
+        assertEquals("count", 1, uid2Info.completedCount);
     }
 
     @Test
     public void diffType() throws Exception {
-        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
-        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_2.getNumber(), TS_2);
-        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_2.getNumber(), TS_3);
-        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_4);
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1, TS_1);
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_2, TS_2);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_2, TS_3);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1, TS_4);
 
-        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_5);
+        ArrayList<WakelockInfo> info = pullResults(TS_5);
         assertEquals("size", 2, info.size());
 
-        FrameworkWakelockInfo uid1Info =
-                info.stream().filter(i -> i.getType() == WAKELOCK_TYPE_1).findFirst().orElse(null);
+        WakelockInfo uid1Info =
+                info.stream().filter(i -> i.type == WAKELOCK_TYPE_1).findFirst().orElse(null);
 
         assertTrue("WAKELOCK_TYPE_1 found", uid1Info != null);
-        assertEquals("uid", UID_1, uid1Info.getAttributionUid());
-        assertEquals("tag", TAG_1, uid1Info.getAttributionTag());
-        assertEquals("type", WAKELOCK_TYPE_1, uid1Info.getType());
-        assertEquals("duration", TS_4 - TS_1, uid1Info.getUptimeMillis());
-        assertEquals("count", 1, uid1Info.getCompletedCount());
+        assertEquals("uid", UID_1, uid1Info.uid);
+        assertEquals("tag", TAG_1, uid1Info.tag);
+        assertEquals("type", WAKELOCK_TYPE_1, uid1Info.type);
+        assertEquals("duration", TS_4 - TS_1, uid1Info.uptimeMillis);
+        assertEquals("count", 1, uid1Info.completedCount);
 
-        FrameworkWakelockInfo uid2Info =
-                info.stream().filter(i -> i.getType() == WAKELOCK_TYPE_2).findFirst().orElse(null);
+        WakelockInfo uid2Info =
+                info.stream().filter(i -> i.type == WAKELOCK_TYPE_2).findFirst().orElse(null);
         assertTrue("WAKELOCK_TYPE_2 found", uid2Info != null);
-        assertEquals("uid", UID_1, uid2Info.getAttributionUid());
-        assertEquals("tag", TAG_1, uid2Info.getAttributionTag());
-        assertEquals("type", WAKELOCK_TYPE_2, uid2Info.getType());
-        assertEquals("duration", TS_3 - TS_2, uid2Info.getUptimeMillis());
-        assertEquals("count", 1, uid2Info.getCompletedCount());
+        assertEquals("uid", UID_1, uid2Info.uid);
+        assertEquals("tag", TAG_1, uid2Info.tag);
+        assertEquals("type", WAKELOCK_TYPE_2, uid2Info.type);
+        assertEquals("duration", TS_3 - TS_2, uid2Info.uptimeMillis);
+        assertEquals("count", 1, uid2Info.completedCount);
     }
 }