Delete WrappedSysUiStatsEvent in favor of new StatsEventTestUtils

The custom WrappedSysUiStatsEvent was used to enable testing of pulled atoms, which otherwise were only available in a previously-uninspectable StatsEvent format. Now that there is an official test utility for converting StatsEvent objects to the associated Atom (which we can read), we can inspect the contents of generated StatsEvent directly, removing the need for wrapping the builder and maintaining separate custom testing code.

Also suppresses lint warning on ZenModeHelperTest so we don't have to lock the locks to set up mConfig.

Bug: 161127189
Test: atest ZenModeHelperTest, PreferencesHelperTest

Change-Id: Ibb647adc635fcebf628c12aa0c47392a7cdc3424

7 files changed, 249 insertions, 372 deletions

diff --git a/services/core/java/com/android/server/notification/NotificationManagerService.java b/services/core/java/com/android/server/notification/NotificationManagerService.java
index 6f0a4b48c09e..be6133b5b6ae 100755
--- a/services/core/java/com/android/server/notification/NotificationManagerService.java
+++ b/services/core/java/com/android/server/notification/NotificationManagerService.java
@@ -2328,8 +2328,7 @@ public class NotificationManagerService extends SystemService {
         mRankingHandler = rankingHandler;
         mConditionProviders = conditionProviders;
         mZenModeHelper = new ZenModeHelper(getContext(), mHandler.getLooper(), mConditionProviders,
-                new SysUiStatsEvent.BuilderFactory(), flagResolver,
-                new ZenModeEventLogger(mPackageManagerClient));
+                flagResolver, new ZenModeEventLogger(mPackageManagerClient));
         mZenModeHelper.addCallback(new ZenModeHelper.Callback() {
             @Override
             public void onConfigChanged() {
@@ -2390,7 +2389,6 @@ public class NotificationManagerService extends SystemService {
                 mNotificationChannelLogger,
                 mAppOps,
                 mUserProfiles,
-                new SysUiStatsEvent.BuilderFactory(),
                 mShowReviewPermissionsNotification);
         mRankingHelper = new RankingHelper(getContext(),
                 mRankingHandler,
diff --git a/services/core/java/com/android/server/notification/PreferencesHelper.java b/services/core/java/com/android/server/notification/PreferencesHelper.java
index 0e37f101ce70..b132a23b575b 100644
--- a/services/core/java/com/android/server/notification/PreferencesHelper.java
+++ b/services/core/java/com/android/server/notification/PreferencesHelper.java
@@ -25,7 +25,6 @@ import static android.app.NotificationManager.IMPORTANCE_DEFAULT;
 import static android.app.NotificationManager.IMPORTANCE_MAX;
 import static android.app.NotificationManager.IMPORTANCE_NONE;
 import static android.app.NotificationManager.IMPORTANCE_UNSPECIFIED;
-import static android.util.StatsLog.ANNOTATION_ID_IS_UID;
 
 import static com.android.internal.util.FrameworkStatsLog.PACKAGE_NOTIFICATION_CHANNEL_GROUP_PREFERENCES;
 import static com.android.internal.util.FrameworkStatsLog.PACKAGE_NOTIFICATION_CHANNEL_PREFERENCES;
@@ -78,6 +77,7 @@ import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.config.sysui.SystemUiSystemPropertiesFlags;
 import com.android.internal.config.sysui.SystemUiSystemPropertiesFlags.NotificationFlags;
 import com.android.internal.logging.MetricsLogger;
+import com.android.internal.util.FrameworkStatsLog;
 import com.android.internal.util.Preconditions;
 import com.android.internal.util.XmlUtils;
 import com.android.modules.utils.TypedXmlPullParser;
@@ -169,7 +169,6 @@ public class PreferencesHelper implements RankingConfig {
      * fields.
      */
     private static final int DEFAULT_LOCKED_APP_FIELDS = 0;
-    private final SysUiStatsEvent.BuilderFactory mStatsEventBuilderFactory;
 
     /**
      * All user-lockable fields for a given application.
@@ -208,7 +207,6 @@ public class PreferencesHelper implements RankingConfig {
             ZenModeHelper zenHelper, PermissionHelper permHelper, PermissionManager permManager,
             NotificationChannelLogger notificationChannelLogger,
             AppOpsManager appOpsManager, ManagedServices.UserProfiles userProfiles,
-            SysUiStatsEvent.BuilderFactory statsEventBuilderFactory,
             boolean showReviewPermissionsNotification) {
         mContext = context;
         mZenModeHelper = zenHelper;
@@ -219,7 +217,6 @@ public class PreferencesHelper implements RankingConfig {
         mNotificationChannelLogger = notificationChannelLogger;
         mAppOps = appOpsManager;
         mUserProfiles = userProfiles;
-        mStatsEventBuilderFactory = statsEventBuilderFactory;
         mShowReviewPermissionsNotification = showReviewPermissionsNotification;
 
         XML_VERSION = 4;
@@ -2190,11 +2187,7 @@ public class PreferencesHelper implements RankingConfig {
                     break;
                 }
                 pulledEvents++;
-                SysUiStatsEvent.Builder event = mStatsEventBuilderFactory.newBuilder()
-                        .setAtomId(PACKAGE_NOTIFICATION_PREFERENCES);
                 final PackagePreferences r = mPackagePreferences.valueAt(i);
-                event.writeInt(r.uid);
-                event.addBooleanAnnotation(ANNOTATION_ID_IS_UID, true);
 
                 // collect whether this package's importance info was user-set for later, if needed
                 // before the migration is enabled, this will simply default to false in all cases.
@@ -2214,15 +2207,7 @@ public class PreferencesHelper implements RankingConfig {
 
                     pkgsWithPermissionsToHandle.remove(key);
                 }
-                event.writeInt(importance);
 
-                event.writeInt(r.visibility);
-                event.writeInt(r.lockedAppFields);
-
-                // optional bool user_set_importance = 5;
-                event.writeBoolean(importanceIsUserSet);
-
-                // optional FsiState fsi_state = 6;
                 final boolean isStickyHunFlagEnabled = SystemUiSystemPropertiesFlags.getResolver()
                         .isEnabled(NotificationFlags.SHOW_STICKY_HUN_FOR_DENIED_FSI);
 
@@ -2232,20 +2217,23 @@ public class PreferencesHelper implements RankingConfig {
                 final int fsiState = getFsiState(r.pkg, r.uid, requestedFSIPermission,
                         isStickyHunFlagEnabled);
 
-                event.writeInt(fsiState);
-
-                // optional bool is_fsi_permission_user_set = 7;
                 final int currentPermissionFlags = mPm.getPermissionFlags(
                         android.Manifest.permission.USE_FULL_SCREEN_INTENT, r.pkg,
                         UserHandle.getUserHandleForUid(r.uid));
 
-                final boolean isUserSet =
+                final boolean fsiIsUserSet =
                         isFsiPermissionUserSet(r.pkg, r.uid, fsiState, currentPermissionFlags,
                                 isStickyHunFlagEnabled);
 
-                event.writeBoolean(isUserSet);
-
-                events.add(event.build());
+                events.add(FrameworkStatsLog.buildStatsEvent(
+                        PACKAGE_NOTIFICATION_PREFERENCES,
+                        /* optional int32 uid = 1 [(is_uid) = true] */ r.uid,
+                        /* optional int32 importance = 2 */ importance,
+                        /* optional int32 visibility = 3 */ r.visibility,
+                        /* optional int32 user_locked_fields = 4 */ r.lockedAppFields,
+                        /* optional bool user_set_importance = 5 */ importanceIsUserSet,
+                        /* optional FsiState fsi_state = 6 */ fsiState,
+                        /* optional bool is_fsi_permission_user_set = 7 */ fsiIsUserSet));
             }
         }
 
@@ -2256,18 +2244,18 @@ public class PreferencesHelper implements RankingConfig {
                     break;
                 }
                 pulledEvents++;
-                SysUiStatsEvent.Builder event = mStatsEventBuilderFactory.newBuilder()
-                        .setAtomId(PACKAGE_NOTIFICATION_PREFERENCES);
-                event.writeInt(p.first);
-                event.addBooleanAnnotation(ANNOTATION_ID_IS_UID, true);
-                event.writeInt(pkgPermissions.get(p).first ? IMPORTANCE_DEFAULT : IMPORTANCE_NONE);
-
-                // fill out the rest of the fields with default values so as not to confuse the
-                // builder
-                event.writeInt(DEFAULT_VISIBILITY);
-                event.writeInt(DEFAULT_LOCKED_APP_FIELDS);
-                event.writeBoolean(pkgPermissions.get(p).second); // user_set_importance field
-                events.add(event.build());
+                // Because all fields are required in FrameworkStatsLog.buildStatsEvent, we have
+                // to fill in default values for all the unspecified fields.
+                events.add(FrameworkStatsLog.buildStatsEvent(
+                        PACKAGE_NOTIFICATION_PREFERENCES,
+                        /* optional int32 uid = 1 [(is_uid) = true] */ p.first,
+                        /* optional int32 importance = 2 */ pkgPermissions.get(p).first
+                                ? IMPORTANCE_DEFAULT : IMPORTANCE_NONE,
+                        /* optional int32 visibility = 3 */ DEFAULT_VISIBILITY,
+                        /* optional int32 user_locked_fields = 4 */ DEFAULT_LOCKED_APP_FIELDS,
+                        /* optional bool user_set_importance = 5 */ pkgPermissions.get(p).second,
+                        /* optional FsiState fsi_state = 6 */ 0,
+                        /* optional bool is_fsi_permission_user_set = 7 */ false));
             }
         }
     }
@@ -2288,20 +2276,21 @@ public class PreferencesHelper implements RankingConfig {
                     if (++totalChannelsPulled > NOTIFICATION_CHANNEL_PULL_LIMIT) {
                         break;
                     }
-                    SysUiStatsEvent.Builder event = mStatsEventBuilderFactory.newBuilder()
-                            .setAtomId(PACKAGE_NOTIFICATION_CHANNEL_PREFERENCES);
-                    event.writeInt(r.uid);
-                    event.addBooleanAnnotation(ANNOTATION_ID_IS_UID, true);
-                    event.writeString(channel.getId());
-                    event.writeString(channel.getName().toString());
-                    event.writeString(channel.getDescription());
-                    event.writeInt(channel.getImportance());
-                    event.writeInt(channel.getUserLockedFields());
-                    event.writeBoolean(channel.isDeleted());
-                    event.writeBoolean(channel.getConversationId() != null);
-                    event.writeBoolean(channel.isDemoted());
-                    event.writeBoolean(channel.isImportantConversation());
-                    events.add(event.build());
+                    events.add(FrameworkStatsLog.buildStatsEvent(
+                            PACKAGE_NOTIFICATION_CHANNEL_PREFERENCES,
+                            /* optional int32 uid = 1 [(is_uid) = true] */ r.uid,
+                            /* optional string channel_id = 2 */ channel.getId(),
+                            /* optional string channel_name = 3 */ channel.getName().toString(),
+                            /* optional string description = 4 */ channel.getDescription(),
+                            /* optional int32 importance = 5 */ channel.getImportance(),
+                            /* optional int32 user_locked_fields = 6 */
+                            channel.getUserLockedFields(),
+                            /* optional bool is_deleted = 7 */ channel.isDeleted(),
+                            /* optional bool is_conversation = 8 */
+                            channel.getConversationId() != null,
+                            /* optional bool is_demoted_conversation = 9 */ channel.isDemoted(),
+                            /* optional bool is_important_conversation = 10 */
+                            channel.isImportantConversation()));
                 }
             }
         }
@@ -2323,16 +2312,15 @@ public class PreferencesHelper implements RankingConfig {
                     if (++totalGroupsPulled > NOTIFICATION_CHANNEL_GROUP_PULL_LIMIT) {
                         break;
                     }
-                    SysUiStatsEvent.Builder event = mStatsEventBuilderFactory.newBuilder()
-                            .setAtomId(PACKAGE_NOTIFICATION_CHANNEL_GROUP_PREFERENCES);
-                    event.writeInt(r.uid);
-                    event.addBooleanAnnotation(ANNOTATION_ID_IS_UID, true);
-                    event.writeString(groupChannel.getId());
-                    event.writeString(groupChannel.getName().toString());
-                    event.writeString(groupChannel.getDescription());
-                    event.writeBoolean(groupChannel.isBlocked());
-                    event.writeInt(groupChannel.getUserLockedFields());
-                    events.add(event.build());
+                    events.add(FrameworkStatsLog.buildStatsEvent(
+                            PACKAGE_NOTIFICATION_CHANNEL_GROUP_PREFERENCES,
+                            /* optional int32 uid = 1 [(is_uid) = true] */ r.uid,
+                            /* optional string group_id = 2 */ groupChannel.getId(),
+                            /* optional string group_name = 3 */ groupChannel.getName().toString(),
+                            /* optional string description = 4 */ groupChannel.getDescription(),
+                            /* optional bool is_blocked = 5 */ groupChannel.isBlocked(),
+                            /* optional int32 user_locked_fields = 6 */
+                            groupChannel.getUserLockedFields()));
                 }
             }
         }
diff --git a/services/core/java/com/android/server/notification/ZenModeHelper.java b/services/core/java/com/android/server/notification/ZenModeHelper.java
index 1f5bd3e0cc60..e49074573ed3 100644
--- a/services/core/java/com/android/server/notification/ZenModeHelper.java
+++ b/services/core/java/com/android/server/notification/ZenModeHelper.java
@@ -21,7 +21,6 @@ import static android.app.NotificationManager.AUTOMATIC_RULE_STATUS_ENABLED;
 import static android.app.NotificationManager.AUTOMATIC_RULE_STATUS_REMOVED;
 import static android.app.NotificationManager.Policy.PRIORITY_SENDERS_ANY;
 import static android.service.notification.NotificationServiceProto.ROOT_CONFIG;
-import static android.util.StatsLog.ANNOTATION_ID_IS_UID;
 
 import static com.android.internal.util.FrameworkStatsLog.DND_MODE_RULE;
 
@@ -81,6 +80,7 @@ import com.android.internal.config.sysui.SystemUiSystemPropertiesFlags;
 import com.android.internal.logging.MetricsLogger;
 import com.android.internal.messages.nano.SystemMessageProto.SystemMessage;
 import com.android.internal.notification.SystemNotificationChannels;
+import com.android.internal.util.FrameworkStatsLog;
 import com.android.internal.util.XmlUtils;
 import com.android.modules.utils.TypedXmlPullParser;
 import com.android.modules.utils.TypedXmlSerializer;
@@ -116,7 +116,6 @@ public class ZenModeHelper {
     private final SettingsObserver mSettingsObserver;
     private final AppOpsManager mAppOps;
     private final NotificationManager mNotificationManager;
-    private final SysUiStatsEvent.BuilderFactory mStatsEventBuilderFactory;
     private ZenModeConfig mDefaultConfig;
     private final ArrayList<Callback> mCallbacks = new ArrayList<Callback>();
     private final ZenModeFiltering mFiltering;
@@ -152,7 +151,6 @@ public class ZenModeHelper {
     private String[] mPriorityOnlyDndExemptPackages;
 
     public ZenModeHelper(Context context, Looper looper, ConditionProviders conditionProviders,
-            SysUiStatsEvent.BuilderFactory statsEventBuilderFactory,
             SystemUiSystemPropertiesFlags.FlagResolver flagResolver,
             ZenModeEventLogger zenModeEventLogger) {
         mContext = context;
@@ -174,7 +172,6 @@ public class ZenModeHelper {
         mFiltering = new ZenModeFiltering(mContext);
         mConditions = new ZenModeConditions(this, conditionProviders);
         mServiceConfig = conditionProviders.getConfig();
-        mStatsEventBuilderFactory = statsEventBuilderFactory;
         mFlagResolver = flagResolver;
         mZenModeEventLogger = zenModeEventLogger;
     }
@@ -1314,17 +1311,14 @@ public class ZenModeHelper {
             for (int i = 0; i < numConfigs; i++) {
                 final int user = mConfigs.keyAt(i);
                 final ZenModeConfig config = mConfigs.valueAt(i);
-                SysUiStatsEvent.Builder data = mStatsEventBuilderFactory.newBuilder()
-                        .setAtomId(DND_MODE_RULE)
-                        .writeInt(user)
-                        .writeBoolean(config.manualRule != null) // enabled
-                        .writeBoolean(config.areChannelsBypassingDnd)
-                        .writeInt(ROOT_CONFIG)
-                        .writeString("") // name, empty for root config
-                        .writeInt(Process.SYSTEM_UID) // system owns root config
-                        .addBooleanAnnotation(ANNOTATION_ID_IS_UID, true)
-                        .writeByteArray(config.toZenPolicy().toProto());
-                events.add(data.build());
+                events.add(FrameworkStatsLog.buildStatsEvent(DND_MODE_RULE,
+                        /* optional int32 user = 1 */ user,
+                        /* optional bool enabled = 2 */ config.manualRule != null,
+                        /* optional bool channels_bypassing = 3 */ config.areChannelsBypassingDnd,
+                        /* optional LoggedZenMode zen_mode = 4 */ ROOT_CONFIG,
+                        /* optional string id = 5 */ "", // empty for root config
+                        /* optional int32 uid = 6 */ Process.SYSTEM_UID, // system owns root config
+                        /* optional DNDPolicyProto policy = 7 */ config.toZenPolicy().toProto()));
                 if (config.manualRule != null) {
                     ruleToProtoLocked(user, config.manualRule, true, events);
                 }
@@ -1355,21 +1349,18 @@ public class ZenModeHelper {
         }
 
         SysUiStatsEvent.Builder data;
-        data = mStatsEventBuilderFactory.newBuilder()
-                .setAtomId(DND_MODE_RULE)
-                .writeInt(user)
-                .writeBoolean(rule.enabled)
-                .writeBoolean(false) // channels_bypassing unused for rules
-                .writeInt(rule.zenMode)
-                .writeString(id)
-                .writeInt(getPackageUid(pkg, user))
-                .addBooleanAnnotation(ANNOTATION_ID_IS_UID, true);
         byte[] policyProto = new byte[]{};
         if (rule.zenPolicy != null) {
             policyProto = rule.zenPolicy.toProto();
         }
-        data.writeByteArray(policyProto);
-        events.add(data.build());
+        events.add(FrameworkStatsLog.buildStatsEvent(DND_MODE_RULE,
+                /* optional int32 user = 1 */ user,
+                /* optional bool enabled = 2 */ rule.enabled,
+                /* optional bool channels_bypassing = 3 */ false, // unused for rules
+                /* optional android.stats.dnd.ZenMode zen_mode = 4 */ rule.zenMode,
+                /* optional string id = 5 */ id,
+                /* optional int32 uid = 6 */ getPackageUid(pkg, user),
+                /* optional DNDPolicyProto policy = 7 */ policyProto));
     }
 
     private int getPackageUid(String pkg, int user) {
diff --git a/services/tests/uiservicestests/Android.bp b/services/tests/uiservicestests/Android.bp
index 4b658958cf18..2dacda0af7f4 100644
--- a/services/tests/uiservicestests/Android.bp
+++ b/services/tests/uiservicestests/Android.bp
@@ -34,6 +34,7 @@ android_test {
         "platform-test-annotations",
         "platformprotosnano",
         "statsdprotolite",
+        "StatsdTestUtils",
         "hamcrest-library",
         "servicestests-utils",
         "testables",
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/PreferencesHelperTest.java b/services/tests/uiservicestests/src/com/android/server/notification/PreferencesHelperTest.java
index 9b745f5aaf4f..020afdbce987 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/PreferencesHelperTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/PreferencesHelperTest.java
@@ -44,22 +44,11 @@ import static android.app.NotificationManager.VISIBILITY_NO_OVERRIDE;
 import static android.media.AudioAttributes.CONTENT_TYPE_SONIFICATION;
 import static android.media.AudioAttributes.USAGE_NOTIFICATION;
 import static android.os.UserHandle.USER_SYSTEM;
-import static android.util.StatsLog.ANNOTATION_ID_IS_UID;
 
 import static com.android.internal.config.sysui.SystemUiSystemPropertiesFlags.NotificationFlags.PROPAGATE_CHANNEL_UPDATES_TO_CONVERSATIONS;
-import static com.android.internal.util.FrameworkStatsLog.PACKAGE_NOTIFICATION_CHANNEL_PREFERENCES;
-import static com.android.internal.util.FrameworkStatsLog.PACKAGE_NOTIFICATION_PREFERENCES;
 import static com.android.internal.util.FrameworkStatsLog.PACKAGE_NOTIFICATION_PREFERENCES__FSI_STATE__DENIED;
 import static com.android.internal.util.FrameworkStatsLog.PACKAGE_NOTIFICATION_PREFERENCES__FSI_STATE__GRANTED;
 import static com.android.internal.util.FrameworkStatsLog.PACKAGE_NOTIFICATION_PREFERENCES__FSI_STATE__NOT_REQUESTED;
-import static com.android.os.AtomsProto.PackageNotificationChannelPreferences.CHANNEL_ID_FIELD_NUMBER;
-import static com.android.os.AtomsProto.PackageNotificationChannelPreferences.CHANNEL_NAME_FIELD_NUMBER;
-import static com.android.os.AtomsProto.PackageNotificationChannelPreferences.IMPORTANCE_FIELD_NUMBER;
-import static com.android.os.AtomsProto.PackageNotificationChannelPreferences.IS_CONVERSATION_FIELD_NUMBER;
-import static com.android.os.AtomsProto.PackageNotificationChannelPreferences.IS_DELETED_FIELD_NUMBER;
-import static com.android.os.AtomsProto.PackageNotificationChannelPreferences.IS_DEMOTED_CONVERSATION_FIELD_NUMBER;
-import static com.android.os.AtomsProto.PackageNotificationChannelPreferences.IS_IMPORTANT_CONVERSATION_FIELD_NUMBER;
-import static com.android.os.AtomsProto.PackageNotificationChannelPreferences.UID_FIELD_NUMBER;
 import static com.android.server.notification.NotificationChannelLogger.NotificationChannelEvent.NOTIFICATION_CHANNEL_UPDATED_BY_USER;
 import static com.android.server.notification.PreferencesHelper.DEFAULT_BUBBLE_PREFERENCE;
 import static com.android.server.notification.PreferencesHelper.NOTIFICATION_CHANNEL_COUNT_LIMIT;
@@ -134,6 +123,7 @@ import android.util.ArraySet;
 import android.util.IntArray;
 import android.util.Pair;
 import android.util.StatsEvent;
+import android.util.StatsEventTestUtils;
 import android.util.Xml;
 import android.util.proto.ProtoOutputStream;
 
@@ -144,11 +134,14 @@ import com.android.internal.config.sysui.SystemUiSystemPropertiesFlags;
 import com.android.internal.config.sysui.TestableFlagResolver;
 import com.android.modules.utils.TypedXmlPullParser;
 import com.android.modules.utils.TypedXmlSerializer;
+import com.android.os.AtomsProto;
+import com.android.os.AtomsProto.PackageNotificationChannelPreferences;
 import com.android.os.AtomsProto.PackageNotificationPreferences;
 import com.android.server.UiServiceTestCase;
 import com.android.server.notification.PermissionHelper.PackagePermission;
 
 import com.google.common.collect.ImmutableList;
+import com.google.protobuf.InvalidProtocolBufferException;
 
 import org.json.JSONArray;
 import org.json.JSONObject;
@@ -218,7 +211,6 @@ public class PreferencesHelperTest extends UiServiceTestCase {
     private PreferencesHelper mHelper;
     private AudioAttributes mAudioAttributes;
     private NotificationChannelLoggerFake mLogger = new NotificationChannelLoggerFake();
-    private WrappedSysUiStatsEvent.WrappedBuilderFactory mStatsEventBuilderFactory;
 
     @Before
     public void setUp() throws Exception {
@@ -331,11 +323,9 @@ public class PreferencesHelperTest extends UiServiceTestCase {
 
         when(mUserProfiles.getCurrentProfileIds()).thenReturn(IntArray.wrap(new int[] {0}));
 
-        mStatsEventBuilderFactory = new WrappedSysUiStatsEvent.WrappedBuilderFactory();
-
         mHelper = new PreferencesHelper(getContext(), mPm, mHandler, mMockZenModeHelper,
                 mPermissionHelper, mPermissionManager, mLogger, mAppOpsManager, mUserProfiles,
-                mStatsEventBuilderFactory, false);
+                false);
         resetZenModeHelper();
 
         mAudioAttributes = new AudioAttributes.Builder()
@@ -683,7 +673,7 @@ public class PreferencesHelperTest extends UiServiceTestCase {
     public void testReadXml_oldXml_migrates() throws Exception {
         mHelper = new PreferencesHelper(getContext(), mPm, mHandler, mMockZenModeHelper,
                 mPermissionHelper, mPermissionManager, mLogger, mAppOpsManager, mUserProfiles,
-                mStatsEventBuilderFactory, /* showReviewPermissionsNotification= */ true);
+                /* showReviewPermissionsNotification= */ true);
 
         String xml = "<ranking version=\"2\">\n"
                 + "<package name=\"" + PKG_N_MR1 + "\" uid=\"" + UID_N_MR1
@@ -825,7 +815,7 @@ public class PreferencesHelperTest extends UiServiceTestCase {
     public void testReadXml_newXml_noMigration_showPermissionNotification() throws Exception {
         mHelper = new PreferencesHelper(getContext(), mPm, mHandler, mMockZenModeHelper,
                 mPermissionHelper, mPermissionManager, mLogger, mAppOpsManager, mUserProfiles,
-                mStatsEventBuilderFactory, /* showReviewPermissionsNotification= */ true);
+                /* showReviewPermissionsNotification= */ true);
 
         String xml = "<ranking version=\"3\">\n"
                 + "<package name=\"" + PKG_N_MR1 + "\" show_badge=\"true\">\n"
@@ -883,7 +873,7 @@ public class PreferencesHelperTest extends UiServiceTestCase {
     public void testReadXml_newXml_permissionNotificationOff() throws Exception {
         mHelper = new PreferencesHelper(getContext(), mPm, mHandler, mMockZenModeHelper,
                 mPermissionHelper, mPermissionManager, mLogger, mAppOpsManager, mUserProfiles,
-                mStatsEventBuilderFactory, /* showReviewPermissionsNotification= */ false);
+                /* showReviewPermissionsNotification= */ false);
 
         String xml = "<ranking version=\"3\">\n"
                 + "<package name=\"" + PKG_N_MR1 + "\" show_badge=\"true\">\n"
@@ -941,7 +931,7 @@ public class PreferencesHelperTest extends UiServiceTestCase {
     public void testReadXml_newXml_noMigration_noPermissionNotification() throws Exception {
         mHelper = new PreferencesHelper(getContext(), mPm, mHandler, mMockZenModeHelper,
                 mPermissionHelper, mPermissionManager, mLogger, mAppOpsManager, mUserProfiles,
-                mStatsEventBuilderFactory, /* showReviewPermissionsNotification= */ true);
+                /* showReviewPermissionsNotification= */ true);
 
         String xml = "<ranking version=\"4\">\n"
                 + "<package name=\"" + PKG_N_MR1 + "\" show_badge=\"true\">\n"
@@ -1521,7 +1511,7 @@ public class PreferencesHelperTest extends UiServiceTestCase {
 
         mHelper = new PreferencesHelper(mContext, mPm, mHandler, mMockZenModeHelper,
                 mPermissionHelper, mPermissionManager, mLogger, mAppOpsManager, mUserProfiles,
-                mStatsEventBuilderFactory, false);
+                false);
 
         NotificationChannel channel =
                 new NotificationChannel("id", "name", IMPORTANCE_LOW);
@@ -5392,7 +5382,7 @@ public class PreferencesHelperTest extends UiServiceTestCase {
     }
 
     @Test
-    public void testPullPackageChannelPreferencesStats() {
+    public void testPullPackageChannelPreferencesStats() throws InvalidProtocolBufferException {
         String channelId = "parent";
         String name = "messages";
         NotificationChannel fodderA = new NotificationChannel("a", "a", IMPORTANCE_LOW);
@@ -5406,25 +5396,40 @@ public class PreferencesHelperTest extends UiServiceTestCase {
         ArrayList<StatsEvent> events = new ArrayList<>();
         mHelper.pullPackageChannelPreferencesStats(events);
 
-        int found = 0;
-        for (WrappedSysUiStatsEvent.WrappedBuilder builder : mStatsEventBuilderFactory.builders) {
-            if (builder.getAtomId() == PACKAGE_NOTIFICATION_CHANNEL_PREFERENCES
-                    && channelId.equals(builder.getValue(CHANNEL_ID_FIELD_NUMBER))) {
-                ++found;
-                assertEquals("uid", UID_O, builder.getValue(UID_FIELD_NUMBER));
-                assertTrue("uid annotation", builder.getBooleanAnnotation(
-                        UID_FIELD_NUMBER, ANNOTATION_ID_IS_UID));
-                assertEquals("importance", IMPORTANCE_DEFAULT, builder.getValue(
-                        IMPORTANCE_FIELD_NUMBER));
-                assertEquals("name", name, builder.getValue(CHANNEL_NAME_FIELD_NUMBER));
-                assertFalse("isconv", builder.getBoolean(IS_CONVERSATION_FIELD_NUMBER));
-                assertFalse("deleted", builder.getBoolean(IS_DELETED_FIELD_NUMBER));
+        // number of channels with preferences should be 3 total
+        assertEquals("expected number of events", 3, events.size());
+        for (StatsEvent ev : events) {
+            // all of these events should be of PackageNotificationChannelPreferences type,
+            // and therefore we expect the atom to have this field.
+            AtomsProto.Atom atom = StatsEventTestUtils.convertToAtom(ev);
+            assertTrue(atom.hasPackageNotificationChannelPreferences());
+            PackageNotificationChannelPreferences p =
+                    atom.getPackageNotificationChannelPreferences();
+
+            // uid is shared across all channels; conversation & deleted are not set in any of
+            // these channels; beyond that check individual channel properties
+            assertEquals("uid", UID_O, p.getUid());
+            assertFalse("is conversation", p.getIsConversation());
+            assertFalse("is deleted", p.getIsDeleted());
+
+            String eventChannelId = p.getChannelId();
+            if (eventChannelId.equals(channelId)) {
+                assertEquals("channel name", name, p.getChannelName());
+                assertEquals("importance", IMPORTANCE_DEFAULT, p.getImportance());
+                assertFalse("is conversation", p.getIsConversation());
+            } else if (eventChannelId.equals("a")) {
+                assertEquals("channel name", "a", p.getChannelName());
+                assertEquals("importance", IMPORTANCE_LOW, p.getImportance());
+            } else { // b
+                assertEquals("channel name", "b", p.getChannelName());
+                assertEquals("importance", IMPORTANCE_HIGH, p.getImportance());
             }
         }
     }
 
     @Test
-    public void testPullPackageChannelPreferencesStats_one_to_one() {
+    public void testPullPackageChannelPreferencesStats_one_to_one()
+            throws InvalidProtocolBufferException {
         NotificationChannel channelA = new NotificationChannel("a", "a", IMPORTANCE_LOW);
         mHelper.createNotificationChannel(PKG_O, UID_O, channelA, true, false, UID_O, false);
         NotificationChannel channelB = new NotificationChannel("b", "b", IMPORTANCE_LOW);
@@ -5437,19 +5442,22 @@ public class PreferencesHelperTest extends UiServiceTestCase {
         ArrayList<StatsEvent> events = new ArrayList<>();
         mHelper.pullPackageChannelPreferencesStats(events);
 
-        int found = 0;
-        for (WrappedSysUiStatsEvent.WrappedBuilder builder : mStatsEventBuilderFactory.builders) {
-            if (builder.getAtomId() == PACKAGE_NOTIFICATION_CHANNEL_PREFERENCES) {
-                Object id = builder.getValue(CHANNEL_ID_FIELD_NUMBER);
-                assertTrue("missing channel in the output", channels.contains(id));
-                channels.remove(id);
-            }
+        assertEquals("total events", 3, events.size());
+        for (StatsEvent ev : events) {
+            AtomsProto.Atom atom = StatsEventTestUtils.convertToAtom(ev);
+            assertTrue(atom.hasPackageNotificationChannelPreferences());
+            PackageNotificationChannelPreferences p =
+                    atom.getPackageNotificationChannelPreferences();
+            String id = p.getChannelId();
+            assertTrue("missing channel in the output", channels.contains(id));
+            channels.remove(id);
         }
         assertTrue("unexpected channel in output", channels.isEmpty());
     }
 
     @Test
-    public void testPullPackageChannelPreferencesStats_conversation() {
+    public void testPullPackageChannelPreferencesStats_conversation()
+            throws InvalidProtocolBufferException {
         String conversationId = "friend";
 
         NotificationChannel parent =
@@ -5467,21 +5475,25 @@ public class PreferencesHelperTest extends UiServiceTestCase {
         ArrayList<StatsEvent> events = new ArrayList<>();
         mHelper.pullPackageChannelPreferencesStats(events);
 
-        for (WrappedSysUiStatsEvent.WrappedBuilder builder : mStatsEventBuilderFactory.builders) {
-            if (builder.getAtomId() == PACKAGE_NOTIFICATION_CHANNEL_PREFERENCES
-                    && channelId.equals(builder.getValue(CHANNEL_ID_FIELD_NUMBER))) {
-                assertTrue("isConveration should be true", builder.getBoolean(
-                        IS_CONVERSATION_FIELD_NUMBER));
-                assertFalse("not demoted", builder.getBoolean(
-                        IS_DEMOTED_CONVERSATION_FIELD_NUMBER));
-                assertFalse("not important", builder.getBoolean(
-                        IS_IMPORTANT_CONVERSATION_FIELD_NUMBER));
+        // In this case, we want to check the properties of the conversation channel (not parent)
+        assertEquals("total events", 2, events.size());
+        for (StatsEvent ev : events) {
+            AtomsProto.Atom atom = StatsEventTestUtils.convertToAtom(ev);
+            assertTrue(atom.hasPackageNotificationChannelPreferences());
+            PackageNotificationChannelPreferences p =
+                    atom.getPackageNotificationChannelPreferences();
+
+            if (channelId.equals(p.getChannelId())) {
+                assertTrue("isConversation should be true", p.getIsConversation());
+                assertFalse("not demoted", p.getIsDemotedConversation());
+                assertFalse("not important", p.getIsImportantConversation());
             }
         }
     }
 
     @Test
-    public void testPullPackageChannelPreferencesStats_conversation_demoted() {
+    public void testPullPackageChannelPreferencesStats_conversation_demoted()
+            throws InvalidProtocolBufferException {
         NotificationChannel parent =
                 new NotificationChannel("parent", "messages", IMPORTANCE_DEFAULT);
         mHelper.createNotificationChannel(PKG_O, UID_O, parent, true, false, UID_O, false);
@@ -5496,21 +5508,23 @@ public class PreferencesHelperTest extends UiServiceTestCase {
         ArrayList<StatsEvent> events = new ArrayList<>();
         mHelper.pullPackageChannelPreferencesStats(events);
 
-        for (WrappedSysUiStatsEvent.WrappedBuilder builder : mStatsEventBuilderFactory.builders) {
-            if (builder.getAtomId() == PACKAGE_NOTIFICATION_CHANNEL_PREFERENCES
-                    && channelId.equals(builder.getValue(CHANNEL_ID_FIELD_NUMBER))) {
-                assertTrue("isConveration should be true", builder.getBoolean(
-                        IS_CONVERSATION_FIELD_NUMBER));
-                assertTrue("is demoted", builder.getBoolean(
-                        IS_DEMOTED_CONVERSATION_FIELD_NUMBER));
-                assertFalse("not important", builder.getBoolean(
-                        IS_IMPORTANT_CONVERSATION_FIELD_NUMBER));
+        assertEquals("total events", 2, events.size());
+        for (StatsEvent ev : events) {
+            AtomsProto.Atom atom = StatsEventTestUtils.convertToAtom(ev);
+            assertTrue(atom.hasPackageNotificationChannelPreferences());
+            PackageNotificationChannelPreferences p =
+                    atom.getPackageNotificationChannelPreferences();
+            if (channelId.equals(p.getChannelId())) {
+                assertTrue("isConversation should be true", p.getIsConversation());
+                assertTrue("is demoted", p.getIsDemotedConversation());
+                assertFalse("not important", p.getIsImportantConversation());
             }
         }
     }
 
     @Test
-    public void testPullPackageChannelPreferencesStats_conversation_priority() {
+    public void testPullPackageChannelPreferencesStats_conversation_priority()
+            throws InvalidProtocolBufferException {
         NotificationChannel parent =
                 new NotificationChannel("parent", "messages", IMPORTANCE_DEFAULT);
         mHelper.createNotificationChannel(PKG_O, UID_O, parent, true, false, UID_O, false);
@@ -5525,21 +5539,23 @@ public class PreferencesHelperTest extends UiServiceTestCase {
         ArrayList<StatsEvent> events = new ArrayList<>();
         mHelper.pullPackageChannelPreferencesStats(events);
 
-        for (WrappedSysUiStatsEvent.WrappedBuilder builder : mStatsEventBuilderFactory.builders) {
-            if (builder.getAtomId() == PACKAGE_NOTIFICATION_CHANNEL_PREFERENCES
-                    && channelId.equals(builder.getValue(CHANNEL_ID_FIELD_NUMBER))) {
-                assertTrue("isConveration should be true", builder.getBoolean(
-                        IS_CONVERSATION_FIELD_NUMBER));
-                assertFalse("not demoted", builder.getBoolean(
-                        IS_DEMOTED_CONVERSATION_FIELD_NUMBER));
-                assertTrue("is important", builder.getBoolean(
-                        IS_IMPORTANT_CONVERSATION_FIELD_NUMBER));
+        assertEquals("total events", 2, events.size());
+        for (StatsEvent ev : events) {
+            AtomsProto.Atom atom = StatsEventTestUtils.convertToAtom(ev);
+            assertTrue(atom.hasPackageNotificationChannelPreferences());
+            PackageNotificationChannelPreferences p =
+                    atom.getPackageNotificationChannelPreferences();
+            if (channelId.equals(p.getChannelId())) {
+                assertTrue("isConversation should be true", p.getIsConversation());
+                assertFalse("not demoted", p.getIsDemotedConversation());
+                assertTrue("is important", p.getIsImportantConversation());
             }
         }
     }
 
     @Test
-    public void testPullPackagePreferencesStats_postPermissionMigration() {
+    public void testPullPackagePreferencesStats_postPermissionMigration()
+            throws InvalidProtocolBufferException {
         // make sure there's at least one channel for each package we want to test
         NotificationChannel channelA = new NotificationChannel("a", "a", IMPORTANCE_DEFAULT);
         mHelper.createNotificationChannel(PKG_N_MR1, UID_N_MR1, channelA, true, false,
@@ -5568,23 +5584,18 @@ public class PreferencesHelperTest extends UiServiceTestCase {
         ArrayList<StatsEvent> events = new ArrayList<>();
         mHelper.pullPackagePreferencesStats(events, appPermissions);
 
-        int found = 0;
-        for (WrappedSysUiStatsEvent.WrappedBuilder builder : mStatsEventBuilderFactory.builders) {
-            if (builder.getAtomId() == PACKAGE_NOTIFICATION_PREFERENCES) {
-                ++found;
-                int uid = builder.getInt(PackageNotificationPreferences.UID_FIELD_NUMBER);
-                boolean userSet = builder.getBoolean(
-                        PackageNotificationPreferences.USER_SET_IMPORTANCE_FIELD_NUMBER);
-
-                // if it's one of the expected ids, then make sure the importance matches
-                assertTrue(expected.containsKey(uid));
-                assertThat(expected.get(uid).first).isEqualTo(
-                        builder.getInt(PackageNotificationPreferences.IMPORTANCE_FIELD_NUMBER));
-                assertThat(expected.get(uid).second).isEqualTo(userSet);
-            }
+        assertEquals("total number of packages", 3, events.size());
+        for (StatsEvent ev : events) {
+            AtomsProto.Atom atom = StatsEventTestUtils.convertToAtom(ev);
+            assertTrue(atom.hasPackageNotificationPreferences());
+            PackageNotificationPreferences p = atom.getPackageNotificationPreferences();
+            int uid = p.getUid();
+
+            // if it's one of the expected ids, then make sure the importance matches
+            assertTrue(expected.containsKey(uid));
+            assertThat(expected.get(uid).first).isEqualTo(p.getImportance());
+            assertThat(expected.get(uid).second).isEqualTo(p.getUserSetImportance());
         }
-        // should have at least one entry for each of the packages we expected to see
-        assertThat(found).isAtLeast(3);
     }
 
     @Test
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/WrappedSysUiStatsEvent.java b/services/tests/uiservicestests/src/com/android/server/notification/WrappedSysUiStatsEvent.java
deleted file mode 100644
index 89adc724f600..000000000000
--- a/services/tests/uiservicestests/src/com/android/server/notification/WrappedSysUiStatsEvent.java
+++ /dev/null
@@ -1,143 +0,0 @@
-/*
- * Copyright (C) 2020 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.notification;
-
-import android.util.StatsEvent;
-
-import com.android.server.notification.SysUiStatsEvent.Builder;
-
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-
-/**
- * Wrapper for SysUiStatsEvent that implements validation.
- */
-public class WrappedSysUiStatsEvent {
-
-    static class WrappedBuilder extends Builder {
-        private ArrayList<Object> mValues;
-        private HashMap<Integer, HashMap<Byte, Object>> mAnnotations;
-        private int mAtomId;
-        private int mLastIndex;
-        private boolean mBuilt;
-
-        WrappedBuilder(StatsEvent.Builder builder) {
-            super(builder);
-            mValues = new ArrayList<>();
-            mAnnotations = new HashMap<>();
-            mValues.add(0); // proto fields are 1-based
-        }
-
-        @Override
-        public Builder setAtomId(int atomId) {
-            mAtomId = atomId;
-            super.setAtomId(atomId);
-            return this;
-        }
-
-        @Override
-        public Builder writeInt(int value) {
-            addValue(Integer.valueOf(value));
-            super.writeInt(value);
-            return this;
-        }
-
-        @Override
-        public Builder addBooleanAnnotation(byte annotation, boolean value) {
-            addAnnotation(annotation, Boolean.valueOf(value));
-            super.addBooleanAnnotation(annotation, value);
-            return this;
-        }
-
-        @Override
-        public Builder writeString(String value) {
-            addValue(value);
-            super.writeString(value);
-            return this;
-        }
-
-        @Override
-        public Builder writeBoolean(boolean value) {
-            addValue(Boolean.valueOf(value));
-            super.writeBoolean(value);
-            return this;
-        }
-
-        @Override
-        public StatsEvent build() {
-            mBuilt = true;
-            return super.build();
-        }
-
-        public Object getValue(int index) {
-            return index < mValues.size() ? mValues.get(index) : null;
-        }
-
-        /** useful to make assertTrue() statements more readable. */
-        public boolean getBoolean(int index) {
-            return (Boolean) mValues.get(index);
-        }
-
-        /** useful to make assertTrue() statements more readable. */
-        public int getInt(int index) {
-            return (Integer) mValues.get(index);
-        }
-
-        /** useful to make assertTrue() statements more readable. */
-        public String getString(int index) {
-            return (String) mValues.get(index);
-        }
-
-        private void addValue(Object value) {
-            mLastIndex = mValues.size();
-            mValues.add(value);
-        }
-
-        private void addAnnotation(byte annotation, Object value) {
-            Integer key = Integer.valueOf(mLastIndex);
-            if (!mAnnotations.containsKey(key)) {
-                mAnnotations.put(key, new HashMap<>());
-            }
-            mAnnotations.get(key).put(Byte.valueOf(annotation), value);
-        }
-
-        public boolean getBooleanAnnotation(int i, byte a) {
-            return ((Boolean) mAnnotations.get(Integer.valueOf(i)).get(Byte.valueOf(a)))
-                    .booleanValue();
-        }
-
-        public int getAtomId() {
-            return mAtomId;
-        }
-    }
-
-    static class WrappedBuilderFactory extends SysUiStatsEvent.BuilderFactory {
-        public List<WrappedBuilder> builders;
-
-        WrappedBuilderFactory() {
-            builders = new ArrayList<>();
-        }
-
-        @Override
-        Builder newBuilder() {
-            WrappedBuilder b = new WrappedBuilder(StatsEvent.newBuilder());
-            builders.add(b);
-            return b;
-        }
-    }
-}
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java b/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java
index 3ee75de23fdb..e540068e6bc6 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java
@@ -33,18 +33,12 @@ import static android.app.NotificationManager.Policy.SUPPRESSED_EFFECT_BADGE;
 import static android.app.NotificationManager.Policy.SUPPRESSED_EFFECT_FULL_SCREEN_INTENT;
 import static android.app.NotificationManager.Policy.SUPPRESSED_EFFECT_LIGHTS;
 import static android.app.NotificationManager.Policy.SUPPRESSED_EFFECT_PEEK;
+import static android.provider.Settings.Global.ZEN_MODE_ALARMS;
 import static android.provider.Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS;
 import static android.service.notification.Condition.STATE_FALSE;
 import static android.service.notification.Condition.STATE_TRUE;
-import static android.util.StatsLog.ANNOTATION_ID_IS_UID;
 
 import static com.android.internal.config.sysui.SystemUiSystemPropertiesFlags.NotificationFlags.LOG_DND_STATE_EVENTS;
-import static com.android.internal.util.FrameworkStatsLog.DND_MODE_RULE;
-import static com.android.os.dnd.DNDModeProto.CHANNELS_BYPASSING_FIELD_NUMBER;
-import static com.android.os.dnd.DNDModeProto.ENABLED_FIELD_NUMBER;
-import static com.android.os.dnd.DNDModeProto.ID_FIELD_NUMBER;
-import static com.android.os.dnd.DNDModeProto.UID_FIELD_NUMBER;
-import static com.android.os.dnd.DNDModeProto.ZEN_MODE_FIELD_NUMBER;
 import static com.android.os.dnd.DNDProtoEnums.PEOPLE_STARRED;
 import static com.android.os.dnd.DNDProtoEnums.ROOT_CONFIG;
 import static com.android.os.dnd.DNDProtoEnums.STATE_ALLOW;
@@ -73,6 +67,7 @@ import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.annotation.SuppressLint;
 import android.app.AppGlobals;
 import android.app.AppOpsManager;
 import android.app.AutomaticZenRule;
@@ -106,6 +101,7 @@ import android.testing.TestableLooper;
 import android.util.ArrayMap;
 import android.util.Log;
 import android.util.StatsEvent;
+import android.util.StatsEventTestUtils;
 import android.util.Xml;
 
 import com.android.internal.R;
@@ -113,12 +109,15 @@ import com.android.internal.config.sysui.TestableFlagResolver;
 import com.android.internal.messages.nano.SystemMessageProto.SystemMessage;
 import com.android.modules.utils.TypedXmlPullParser;
 import com.android.modules.utils.TypedXmlSerializer;
+import com.android.os.AtomsProto;
+import com.android.os.dnd.DNDModeProto;
 import com.android.os.dnd.DNDPolicyProto;
 import com.android.os.dnd.DNDProtoEnums;
 import com.android.server.UiServiceTestCase;
 import com.android.server.notification.ManagedServices.UserProfiles;
 
 import com.google.common.collect.ImmutableList;
+import com.google.protobuf.InvalidProtocolBufferException;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -140,13 +139,13 @@ import java.util.List;
 import java.util.Objects;
 
 @SmallTest
+@SuppressLint("GuardedBy") // It's ok for this test to access guarded methods from the service.
 @RunWith(AndroidTestingRunner.class)
 @TestableLooper.RunWithLooper
 public class ZenModeHelperTest extends UiServiceTestCase {
 
     private static final String EVENTS_DEFAULT_RULE_ID = "EVENTS_DEFAULT_RULE";
     private static final String SCHEDULE_DEFAULT_RULE_ID = "EVERY_NIGHT_DEFAULT_RULE";
-    private static final int ZEN_MODE_FOR_TESTING = 99;
     private static final String CUSTOM_PKG_NAME = "not.android";
     private static final int CUSTOM_PKG_UID = 1;
     private static final String CUSTOM_RULE_ID = "custom_rule";
@@ -159,7 +158,6 @@ public class ZenModeHelperTest extends UiServiceTestCase {
     private ZenModeHelper mZenModeHelper;
     private ContentResolver mContentResolver;
     @Mock AppOpsManager mAppOps;
-    private WrappedSysUiStatsEvent.WrappedBuilderFactory mStatsEventBuilderFactory;
     TestableFlagResolver mTestFlagResolver = new TestableFlagResolver();
     ZenModeEventLoggerFake mZenModeEventLogger;
 
@@ -178,7 +176,6 @@ public class ZenModeHelperTest extends UiServiceTestCase {
             Log.d("ZenModeHelperTest", "Couldn't mock default zen mode config xml file err=" +
                     e.toString());
         }
-        mStatsEventBuilderFactory = new WrappedSysUiStatsEvent.WrappedBuilderFactory();
 
         when(mContext.getSystemService(AppOpsManager.class)).thenReturn(mAppOps);
         when(mContext.getSystemService(NotificationManager.class)).thenReturn(mNotificationManager);
@@ -188,8 +185,7 @@ public class ZenModeHelperTest extends UiServiceTestCase {
         mConditionProviders.addSystemProvider(new ScheduleConditionProvider());
         mZenModeEventLogger = new ZenModeEventLoggerFake(mPackageManager);
         mZenModeHelper = new ZenModeHelper(mContext, mTestableLooper.getLooper(),
-                mConditionProviders, mStatsEventBuilderFactory, mTestFlagResolver,
-                mZenModeEventLogger);
+                mConditionProviders, mTestFlagResolver, mZenModeEventLogger);
 
         ResolveInfo ri = new ResolveInfo();
         ri.activityInfo = new ActivityInfo();
@@ -911,37 +907,35 @@ public class ZenModeHelperTest extends UiServiceTestCase {
     }
 
     @Test
-    public void testProto() {
+    public void testProto() throws InvalidProtocolBufferException {
         mZenModeHelper.mZenMode = ZEN_MODE_IMPORTANT_INTERRUPTIONS;
         // existence of manual rule means it should be in output
         mZenModeHelper.mConfig.manualRule = new ZenModeConfig.ZenRule();
         mZenModeHelper.mConfig.manualRule.pkg = "android";  // system
+        mZenModeHelper.mConfig.automaticRules = new ArrayMap<>(); // no automatic rules
 
-        int n = mZenModeHelper.mConfig.automaticRules.size();
-        List<String> ids = new ArrayList<>(n);
-        for (ZenModeConfig.ZenRule rule : mZenModeHelper.mConfig.automaticRules.values()) {
-            ids.add(rule.id);
-        }
+        List<String> ids = new ArrayList<>();
         ids.add(ZenModeConfig.MANUAL_RULE_ID);
         ids.add(""); // for ROOT_CONFIG, logged with empty string as id
 
         List<StatsEvent> events = new LinkedList<>();
         mZenModeHelper.pullRules(events);
-        assertEquals(n + 2, events.size());  // automatic rules + manual rule + root config
-        for (WrappedSysUiStatsEvent.WrappedBuilder builder : mStatsEventBuilderFactory.builders) {
-            if (builder.getAtomId() == DND_MODE_RULE) {
-                if (builder.getInt(ZEN_MODE_FIELD_NUMBER) == ROOT_CONFIG) {
-                    assertTrue(builder.getBoolean(ENABLED_FIELD_NUMBER));
-                    assertFalse(builder.getBoolean(CHANNELS_BYPASSING_FIELD_NUMBER));
-                }
-                assertEquals(Process.SYSTEM_UID, builder.getInt(UID_FIELD_NUMBER));
-                assertTrue(builder.getBooleanAnnotation(UID_FIELD_NUMBER, ANNOTATION_ID_IS_UID));
-                String name = (String) builder.getValue(ID_FIELD_NUMBER);
-                assertTrue("unexpected rule id", ids.contains(name));
-                ids.remove(name);
-            } else {
-                fail("unexpected atom id: " + builder.getAtomId());
+        assertEquals(2, events.size());  // manual rule + root config
+        for (StatsEvent ev : events) {
+            AtomsProto.Atom atom = StatsEventTestUtils.convertToAtom(ev);
+            assertTrue(atom.hasDndModeRule());
+            DNDModeProto cfg = atom.getDndModeRule();
+            // Additional check for ID to clearly identify the root config because there's some
+            // odd behavior in the test util around enum value of 0 (the usual default, but not in
+            // this case).
+            if (cfg.getZenMode().getNumber() == ROOT_CONFIG && cfg.getId().equals("")) {
+                assertTrue(cfg.getEnabled());
+                assertFalse(cfg.getChannelsBypassing());
             }
+            assertEquals(Process.SYSTEM_UID, cfg.getUid());
+            String name = cfg.getId();
+            assertTrue("unexpected rule id", ids.contains(name));
+            ids.remove(name);
         }
         assertEquals("extra rule in output", 0, ids.size());
     }
@@ -949,28 +943,61 @@ public class ZenModeHelperTest extends UiServiceTestCase {
     @Test
     public void testProtoWithAutoRule() throws Exception {
         setupZenConfig();
-        // one enabled automatic rule
-        mZenModeHelper.mConfig.automaticRules = getCustomAutomaticRules(ZEN_MODE_FOR_TESTING);
+        // one enabled automatic rule. we use a non-usual zen mode value (though it has to be
+        // a real one in the enum because non-valid enum values are reverted to default).
+        mZenModeHelper.mConfig.automaticRules = getCustomAutomaticRules(ZEN_MODE_ALARMS);
 
         List<StatsEvent> events = new LinkedList<>();
         mZenModeHelper.pullRules(events);
 
         boolean foundCustomEvent = false;
-        for (WrappedSysUiStatsEvent.WrappedBuilder builder : mStatsEventBuilderFactory.builders) {
-            if (builder.getAtomId() == DND_MODE_RULE) {
-                if (ZEN_MODE_FOR_TESTING == builder.getInt(ZEN_MODE_FIELD_NUMBER)) {
-                    foundCustomEvent = true;
-                    assertEquals(CUSTOM_PKG_UID, builder.getInt(UID_FIELD_NUMBER));
-                    assertTrue(builder.getBoolean(ENABLED_FIELD_NUMBER));
-                }
-            } else {
-                fail("unexpected atom id: " + builder.getAtomId());
+        for (StatsEvent ev : events) {
+            AtomsProto.Atom atom = StatsEventTestUtils.convertToAtom(ev);
+            assertTrue(atom.hasDndModeRule());
+            DNDModeProto cfg = atom.getDndModeRule();
+            if (cfg.getZenMode().getNumber() == ZEN_MODE_ALARMS) {
+                foundCustomEvent = true;
+                assertEquals(CUSTOM_PKG_UID, cfg.getUid());
+                assertTrue(cfg.getEnabled());
             }
         }
         assertTrue("couldn't find custom rule", foundCustomEvent);
     }
 
     @Test
+    public void testProtoWithDefaultAutoRules() throws Exception {
+        setupZenConfig();
+        // clear the automatic rules so we can reset to only the default rules
+        mZenModeHelper.mConfig.automaticRules = new ArrayMap<>();
+
+        // read in XML to restore the default rules
+        ByteArrayOutputStream baos = writeXmlAndPurge(5);
+        TypedXmlPullParser parser = Xml.newFastPullParser();
+        parser.setInput(new BufferedInputStream(
+                new ByteArrayInputStream(baos.toByteArray())), null);
+        parser.nextTag();
+        mZenModeHelper.readXml(parser, false, UserHandle.USER_ALL);
+        List<StatsEvent> events = new LinkedList<>();
+        mZenModeHelper.pullRules(events);
+
+        // list for tracking which ids we've seen in the pulled atom output
+        List<String> ids = new ArrayList<>();
+        ids.addAll(ZenModeConfig.DEFAULT_RULE_IDS);
+        ids.add("");  // empty string for root config
+
+        for (StatsEvent ev : events) {
+            AtomsProto.Atom atom = StatsEventTestUtils.convertToAtom(ev);
+            assertTrue(atom.hasDndModeRule());
+            DNDModeProto cfg = atom.getDndModeRule();
+            if (!ids.contains(cfg.getId())) {
+                fail("unexpected ID found: " + cfg.getId());
+            }
+            ids.remove(cfg.getId());
+        }
+        assertEquals("default ID(s) not found", 0, ids.size());
+    }
+
+    @Test
     public void ruleUidsCached() throws Exception {
         setupZenConfig();
         // one enabled automatic rule
@@ -1019,10 +1046,11 @@ public class ZenModeHelperTest extends UiServiceTestCase {
         List<StatsEvent> events = new LinkedList<>();
         mZenModeHelper.pullRules(events);
 
-        boolean foundCustomEvent = false;
-        for (WrappedSysUiStatsEvent.WrappedBuilder builder : mStatsEventBuilderFactory.builders) {
-            if (builder.getAtomId() == DND_MODE_RULE
-                    && "customRule".equals(builder.getString(ID_FIELD_NUMBER))) {
+        for (StatsEvent ev : events) {
+            AtomsProto.Atom atom = StatsEventTestUtils.convertToAtom(ev);
+            assertTrue(atom.hasDndModeRule());
+            DNDModeProto cfg = atom.getDndModeRule();
+            if ("customRule".equals(cfg.getId())) {
                 fail("non-default IDs should be redacted");
             }
         }
@@ -1039,14 +1067,17 @@ public class ZenModeHelperTest extends UiServiceTestCase {
         mZenModeHelper.pullRules(events);
 
         boolean foundManualRule = false;
-        for (WrappedSysUiStatsEvent.WrappedBuilder builder : mStatsEventBuilderFactory.builders) {
-            if (builder.getAtomId() == DND_MODE_RULE
-                    && ZenModeConfig.MANUAL_RULE_ID.equals(builder.getString(ID_FIELD_NUMBER))) {
-                assertEquals(0, builder.getInt(UID_FIELD_NUMBER));
+        for (StatsEvent ev : events) {
+            AtomsProto.Atom atom = StatsEventTestUtils.convertToAtom(ev);
+            assertTrue(atom.hasDndModeRule());
+            DNDModeProto cfg = atom.getDndModeRule();
+            if (ZenModeConfig.MANUAL_RULE_ID.equals(cfg.getId())) {
+                assertEquals(0, cfg.getUid());
                 foundManualRule = true;
             }
         }
-        assertTrue("couldn't find manual rule", foundManualRule);    }
+        assertTrue("couldn't find manual rule", foundManualRule);
+    }
 
     @Test
     public void testWriteXml_onlyBackupsTargetUser() throws Exception {