Merge "Logs BundlePreferences on pkg data pull" into main

5 files changed, 351 insertions, 10 deletions

diff --git a/services/core/java/com/android/server/notification/NotificationManagerService.java b/services/core/java/com/android/server/notification/NotificationManagerService.java
index 341038f878d9..05aa4134cbd5 100644
--- a/services/core/java/com/android/server/notification/NotificationManagerService.java
+++ b/services/core/java/com/android/server/notification/NotificationManagerService.java
@@ -158,6 +158,7 @@ import static android.view.WindowManager.LayoutParams.TYPE_TOAST;
 import static android.view.contentprotection.flags.Flags.rapidClearNotificationsByListenerAppOpEnabled;
 
 import static com.android.internal.util.FrameworkStatsLog.DND_MODE_RULE;
+import static com.android.internal.util.FrameworkStatsLog.NOTIFICATION_BUNDLE_PREFERENCES;
 import static com.android.internal.util.FrameworkStatsLog.PACKAGE_NOTIFICATION_CHANNEL_GROUP_PREFERENCES;
 import static com.android.internal.util.FrameworkStatsLog.PACKAGE_NOTIFICATION_CHANNEL_PREFERENCES;
 import static com.android.internal.util.FrameworkStatsLog.PACKAGE_NOTIFICATION_PREFERENCES;
@@ -362,6 +363,7 @@ import com.android.server.lights.LightsManager;
 import com.android.server.notification.GroupHelper.NotificationAttributes;
 import com.android.server.notification.ManagedServices.ManagedServiceInfo;
 import com.android.server.notification.ManagedServices.UserProfiles;
+import com.android.server.notification.NotificationRecordLogger.NotificationPullStatsEvent;
 import com.android.server.notification.NotificationRecordLogger.NotificationReportedEvent;
 import com.android.server.notification.toast.CustomToastRecord;
 import com.android.server.notification.toast.TextToastRecord;
@@ -2856,6 +2858,7 @@ public class NotificationManagerService extends SystemService {
             mStatsManager.clearPullAtomCallback(PACKAGE_NOTIFICATION_PREFERENCES);
             mStatsManager.clearPullAtomCallback(PACKAGE_NOTIFICATION_CHANNEL_PREFERENCES);
             mStatsManager.clearPullAtomCallback(PACKAGE_NOTIFICATION_CHANNEL_GROUP_PREFERENCES);
+            mStatsManager.clearPullAtomCallback(NOTIFICATION_BUNDLE_PREFERENCES);
             mStatsManager.clearPullAtomCallback(DND_MODE_RULE);
         }
         if (mAppOps != null) {
@@ -2960,6 +2963,12 @@ public class NotificationManagerService extends SystemService {
                 ConcurrentUtils.DIRECT_EXECUTOR,
                 mPullAtomCallback
         );
+        mStatsManager.setPullAtomCallback(
+                NOTIFICATION_BUNDLE_PREFERENCES,
+                null, // use default PullAtomMetadata values
+                ConcurrentUtils.DIRECT_EXECUTOR,
+                mPullAtomCallback
+        );
     }
 
     private class StatsPullAtomCallbackImpl implements StatsManager.StatsPullAtomCallback {
@@ -2969,6 +2978,7 @@ public class NotificationManagerService extends SystemService {
                 case PACKAGE_NOTIFICATION_PREFERENCES:
                 case PACKAGE_NOTIFICATION_CHANNEL_PREFERENCES:
                 case PACKAGE_NOTIFICATION_CHANNEL_GROUP_PREFERENCES:
+                case NOTIFICATION_BUNDLE_PREFERENCES:
                 case DND_MODE_RULE:
                     return pullNotificationStates(atomTag, data);
                 default:
@@ -2980,8 +2990,15 @@ public class NotificationManagerService extends SystemService {
     private int pullNotificationStates(int atomTag, List<StatsEvent> data) {
         switch(atomTag) {
             case PACKAGE_NOTIFICATION_PREFERENCES:
-                mPreferencesHelper.pullPackagePreferencesStats(data,
-                        getAllUsersNotificationPermissions());
+                if (notificationClassificationUi()) {
+                    Set<String> pkgs = mAssistants.getPackagesWithKeyTypeAdjustmentSettings();
+                    mPreferencesHelper.pullPackagePreferencesStats(data,
+                            getAllUsersNotificationPermissions(),
+                            getPackageSpecificAdjustmentKeyTypes(pkgs));
+                } else {
+                    mPreferencesHelper.pullPackagePreferencesStats(data,
+                            getAllUsersNotificationPermissions());
+                }
                 break;
             case PACKAGE_NOTIFICATION_CHANNEL_PREFERENCES:
                 mPreferencesHelper.pullPackageChannelPreferencesStats(data);
@@ -2989,6 +3006,11 @@ public class NotificationManagerService extends SystemService {
             case PACKAGE_NOTIFICATION_CHANNEL_GROUP_PREFERENCES:
                 mPreferencesHelper.pullPackageChannelGroupPreferencesStats(data);
                 break;
+            case NOTIFICATION_BUNDLE_PREFERENCES:
+                if (notificationClassification() && notificationClassificationUi()) {
+                    mAssistants.pullBundlePreferencesStats(data);
+                }
+                break;
             case DND_MODE_RULE:
                 mZenModeHelper.pullRules(data);
                 break;
@@ -7481,6 +7503,24 @@ public class NotificationManagerService extends SystemService {
         return allPermissions;
     }
 
+    @VisibleForTesting
+    @FlaggedApi(android.app.Flags.FLAG_NOTIFICATION_CLASSIFICATION_UI)
+    protected @NonNull Map<String, Set<Integer>> getPackageSpecificAdjustmentKeyTypes(
+            Set<String> pkgs) {
+        ArrayMap<String, Set<Integer>> pkgToAllowedTypes = new ArrayMap<>();
+        for (String pkg : pkgs) {
+            int[] allowedTypesArray = mAssistants.getAllowedAdjustmentKeyTypesForPackage(pkg);
+            if (allowedTypesArray != null) {
+                Set<Integer> allowedTypes = new ArraySet<Integer>();
+                for (int i : allowedTypesArray) {
+                    allowedTypes.add(i);
+                }
+                pkgToAllowedTypes.append(pkg, allowedTypes);
+            }
+        }
+        return pkgToAllowedTypes;
+    }
+
     private void dumpJson(PrintWriter pw, @NonNull DumpFilter filter,
             ArrayMap<Pair<Integer, String>, Pair<Boolean, Boolean>> pkgPermissions) {
         JSONObject dump = new JSONObject();
@@ -12056,6 +12096,22 @@ public class NotificationManagerService extends SystemService {
         }
 
         @FlaggedApi(android.app.Flags.FLAG_NOTIFICATION_CLASSIFICATION_UI)
+        protected @NonNull Set<String> getPackagesWithKeyTypeAdjustmentSettings() {
+            if (notificationClassificationUi()) {
+                Set<String> packagesWithModifications = new ArraySet<String>();
+                synchronized (mLock) {
+                    for (String pkg : mClassificationTypePackagesEnabledTypes.keySet()) {
+                        if (mClassificationTypePackagesEnabledTypes.get(pkg) != null) {
+                            packagesWithModifications.add(pkg);
+                        }
+                    }
+                }
+                return packagesWithModifications;
+            }
+            return new ArraySet<String>();
+        }
+
+        @FlaggedApi(android.app.Flags.FLAG_NOTIFICATION_CLASSIFICATION_UI)
         protected @NonNull int[] getAllowedAdjustmentKeyTypesForPackage(String pkg) {
             synchronized (mLock) {
                 if (notificationClassificationUi()) {
@@ -12656,6 +12712,32 @@ public class NotificationManagerService extends SystemService {
                 Slog.e(TAG, "unable to notify assistant (capabilities): " + info, ex);
             }
         }
+
+        /**
+         * Fills out {@link BundlePreferences} proto and wraps it in a {@link StatsEvent}.
+         */
+        @FlaggedApi(android.service.notification.Flags.FLAG_NOTIFICATION_CLASSIFICATION)
+        protected void pullBundlePreferencesStats(List<StatsEvent> events) {
+            boolean bundlesAllowed = true;
+            synchronized (mLock) {
+                List<String> unsupportedAdjustments = new ArrayList(
+                        mNasUnsupported.getOrDefault(
+                                UserHandle.getUserId(Binder.getCallingUid()),
+                                new HashSet<>())
+                );
+                bundlesAllowed = !unsupportedAdjustments.contains(Adjustment.KEY_TYPE);
+            }
+
+            int[] allowedBundleTypes = getAllowedAdjustmentKeyTypes();
+
+            events.add(FrameworkStatsLog.buildStatsEvent(
+                    NOTIFICATION_BUNDLE_PREFERENCES,
+                    /* optional int32 event_id = 1 */
+                    NotificationPullStatsEvent.NOTIFICATION_BUNDLE_PREFERENCES_PULLED.getId(),
+                    /* optional bool bundles_allowed = 2 */ bundlesAllowed,
+                    /* repeated android.stats.notification.BundleTypes allowed_bundle_types = 3 */
+                    allowedBundleTypes));
+        }
     }
 
     /**
diff --git a/services/core/java/com/android/server/notification/NotificationRecordLogger.java b/services/core/java/com/android/server/notification/NotificationRecordLogger.java
index 3943aa583fee..6c0035b82a86 100644
--- a/services/core/java/com/android/server/notification/NotificationRecordLogger.java
+++ b/services/core/java/com/android/server/notification/NotificationRecordLogger.java
@@ -32,8 +32,6 @@ import android.service.notification.NotificationListenerService;
 import android.service.notification.NotificationStats;
 import android.util.Log;
 
-import com.android.internal.config.sysui.SystemUiSystemPropertiesFlags;
-import com.android.internal.config.sysui.SystemUiSystemPropertiesFlags.NotificationFlags;
 import com.android.internal.logging.InstanceId;
 import com.android.internal.logging.UiEvent;
 import com.android.internal.logging.UiEventLogger;
@@ -368,6 +366,19 @@ interface NotificationRecordLogger {
         }
     }
 
+    enum NotificationPullStatsEvent implements UiEventLogger.UiEventEnum {
+        @UiEvent(doc = "Notification Bundle Preferences pulled.")
+        NOTIFICATION_BUNDLE_PREFERENCES_PULLED(2072);
+
+        private final int mId;
+        NotificationPullStatsEvent(int id) {
+            mId = id;
+        }
+        @Override public int getId() {
+            return mId;
+        }
+    }
+
     /**
      * A helper for extracting logging information from one or two NotificationRecords.
      */
diff --git a/services/core/java/com/android/server/notification/PreferencesHelper.java b/services/core/java/com/android/server/notification/PreferencesHelper.java
index 36eabae69b22..45b155049c72 100644
--- a/services/core/java/com/android/server/notification/PreferencesHelper.java
+++ b/services/core/java/com/android/server/notification/PreferencesHelper.java
@@ -16,6 +16,7 @@
 
 package com.android.server.notification;
 
+import static android.app.Flags.notificationClassificationUi;
 import static android.app.AppOpsManager.OP_SYSTEM_ALERT_WINDOW;
 import static android.app.NotificationChannel.DEFAULT_CHANNEL_ID;
 import static android.app.NotificationChannel.NEWS_ID;
@@ -2523,6 +2524,25 @@ public class PreferencesHelper implements RankingConfig {
      */
     public void pullPackagePreferencesStats(List<StatsEvent> events,
             ArrayMap<Pair<Integer, String>, Pair<Boolean, Boolean>> pkgPermissions) {
+        pullPackagePreferencesStats(events, pkgPermissions, new ArrayMap<String, Set<Integer>>());
+    }
+
+
+    /**
+     * Fills out {@link PackageNotificationPreferences} proto and wraps it in a {@link StatsEvent}.
+     * @param events Newly filled out StatsEvent protos are added to this list as output.
+     * @param pkgPermissions Maps from a pair representing a uid and package to a pair of booleans,
+     *                       where the first represents whether the notification permission was
+     *                       granted to that package, and the second represents whether the
+     *                       permission was user-set.
+     * @param pkgAdjustmentKeyTypes A map of package names that are not allowed to have their
+     *                                 notifications classified into differently typed notification
+     *                                 channels, and the channels that they're allowed to be
+     *                                 classified into.
+     */
+    public void pullPackagePreferencesStats(List<StatsEvent> events,
+            ArrayMap<Pair<Integer, String>, Pair<Boolean, Boolean>> pkgPermissions,
+            @NonNull Map<String, Set<Integer>> pkgAdjustmentKeyTypes) {
         Set<Pair<Integer, String>> pkgsWithPermissionsToHandle = null;
         if (pkgPermissions != null) {
             pkgsWithPermissionsToHandle = pkgPermissions.keySet();
@@ -2568,6 +2588,14 @@ public class PreferencesHelper implements RankingConfig {
                         isFsiPermissionUserSet(r.pkg, r.uid, fsiState,
                                 currentPermissionFlags);
 
+                if (!notificationClassificationUi()
+                        && pkgAdjustmentKeyTypes.keySet().size() > 0) {
+                    Slog.w(TAG, "Pkg adjustment types improperly allowed without flag set");
+                }
+
+                int[] allowedBundleTypes =
+                        getAllowedTypesForPackage(pkgAdjustmentKeyTypes, r.pkg);
+
                 events.add(FrameworkStatsLog.buildStatsEvent(
                         PACKAGE_NOTIFICATION_PREFERENCES,
                         /* optional int32 uid = 1 [(is_uid) = true] */ r.uid,
@@ -2576,7 +2604,9 @@ public class PreferencesHelper implements RankingConfig {
                         /* 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));
+                        /* optional bool is_fsi_permission_user_set = 7 */ fsiIsUserSet,
+                        /* repeated int32 allowed_bundle_types = 8 */ allowedBundleTypes
+                ));
             }
         }
 
@@ -2587,6 +2617,10 @@ public class PreferencesHelper implements RankingConfig {
                     break;
                 }
                 pulledEvents++;
+
+                int[] allowedBundleTypes =
+                        getAllowedTypesForPackage(pkgAdjustmentKeyTypes, p.second);
+
                 // Because all fields are required in FrameworkStatsLog.buildStatsEvent, we have
                 // to fill in default values for all the unspecified fields.
                 events.add(FrameworkStatsLog.buildStatsEvent(
@@ -2598,9 +2632,29 @@ public class PreferencesHelper implements RankingConfig {
                         /* 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));
+                        /* optional bool is_fsi_permission_user_set = 7 */ false,
+                        /* repeated BundleTypes allowed_bundle_types = 8 */ allowedBundleTypes));
+            }
+        }
+    }
+
+    private int[] getAllowedTypesForPackage(@NonNull
+                                            Map<String, Set<Integer>> pkgAdjustmentKeyTypes,
+                                            String pkg) {
+        int[] allowedBundleTypes = new int[]{};
+        if (notificationClassificationUi()) {
+            if (pkgAdjustmentKeyTypes.containsKey(pkg)) {
+                // Convert from set to int[]
+                Set<Integer> types = pkgAdjustmentKeyTypes.get(pkg);
+                allowedBundleTypes = new int[types.size()];
+                int i = 0;
+                for (int val : types) {
+                    allowedBundleTypes[i] = val;
+                    i++;
+                }
             }
         }
+        return allowedBundleTypes;
     }
 
     /**
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/NotificationAssistantsTest.java b/services/tests/uiservicestests/src/com/android/server/notification/NotificationAssistantsTest.java
index 4f5cdb73edd2..1df8e3deb84b 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationAssistantsTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationAssistantsTest.java
@@ -64,6 +64,8 @@ import android.testing.TestableContext;
 import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.IntArray;
+import android.util.StatsEvent;
+import android.util.StatsEventTestUtils;
 import android.util.Xml;
 
 import androidx.test.runner.AndroidJUnit4;
@@ -71,9 +73,17 @@ import androidx.test.runner.AndroidJUnit4;
 import com.android.internal.util.CollectionUtils;
 import com.android.modules.utils.TypedXmlPullParser;
 import com.android.modules.utils.TypedXmlSerializer;
+import com.android.os.AtomsProto;
+import com.android.os.notification.NotificationBundlePreferences;
+import com.android.os.notification.NotificationExtensionAtoms;
+import com.android.os.notification.NotificationProtoEnums;
 import com.android.server.UiServiceTestCase;
 import com.android.server.notification.NotificationManagerService.NotificationAssistants;
 
+import com.google.protobuf.CodedInputStream;
+import com.google.protobuf.CodedOutputStream;
+import com.google.protobuf.ExtensionRegistryLite;
+
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -120,6 +130,8 @@ public class NotificationAssistantsTest extends UiServiceTestCase {
 
     ComponentName mCn = new ComponentName("a", "b");
 
+    private ExtensionRegistryLite mRegistry;
+
 
     // Helper function to hold mApproved lock, avoid GuardedBy lint errors
     private boolean isUserSetServicesEmpty(NotificationAssistants assistant, int userId) {
@@ -204,6 +216,8 @@ public class NotificationAssistantsTest extends UiServiceTestCase {
         when(mUserProfiles.getCurrentProfileIds()).thenReturn(profileIds);
         when(mNm.isNASMigrationDone(anyInt())).thenReturn(true);
         when(mNm.canUseManagedServices(any(), anyInt(), any())).thenReturn(true);
+        mRegistry = ExtensionRegistryLite.newInstance();
+        NotificationExtensionAtoms.registerAllExtensions(mRegistry);
     }
 
     @Test
@@ -749,6 +763,28 @@ public class NotificationAssistantsTest extends UiServiceTestCase {
     }
 
     @Test
+    @EnableFlags({android.service.notification.Flags.FLAG_NOTIFICATION_CLASSIFICATION,
+            android.app.Flags.FLAG_NOTIFICATION_CLASSIFICATION_UI})
+    public void testGetPackagesWithKeyTypeAdjustmentSettings() throws Exception {
+        String pkg = "my.package";
+        String pkg2 = "my.package.2";
+        setDefaultAllowedAdjustmentKeyTypes(mAssistants);
+        assertThat(mAssistants.isTypeAdjustmentAllowedForPackage(pkg, TYPE_PROMOTION)).isTrue();
+        assertThat(mAssistants.getPackagesWithKeyTypeAdjustmentSettings()).isEmpty();
+
+        mAssistants.setAssistantAdjustmentKeyTypeStateForPackage(pkg, TYPE_PROMOTION, true);
+        assertThat(mAssistants.getPackagesWithKeyTypeAdjustmentSettings())
+                .containsExactly(pkg);
+        mAssistants.setAssistantAdjustmentKeyTypeStateForPackage(pkg, TYPE_PROMOTION, false);
+        assertThat(mAssistants.getPackagesWithKeyTypeAdjustmentSettings())
+                .containsExactly(pkg);
+        mAssistants.setAssistantAdjustmentKeyTypeStateForPackage(pkg2, TYPE_NEWS, true);
+        mAssistants.setAssistantAdjustmentKeyTypeStateForPackage(pkg2, TYPE_PROMOTION, false);
+        assertThat(mAssistants.getPackagesWithKeyTypeAdjustmentSettings())
+                .containsExactly(pkg, pkg2);
+    }
+
+    @Test
     @EnableFlags(android.app.Flags.FLAG_NOTIFICATION_CLASSIFICATION_UI)
     public void testSetAssistantAdjustmentKeyTypeStateForPackage_usesGlobalDefault() {
         String pkg = "my.package";
@@ -892,4 +928,88 @@ public class NotificationAssistantsTest extends UiServiceTestCase {
         assertThat(mAssistants.getAllowedAdjustmentKeyTypesForPackage(pkg)).asList()
                 .containsExactly(TYPE_NEWS, TYPE_PROMOTION, TYPE_SOCIAL_MEDIA);
     }
+
+    @Test
+    @SuppressWarnings("GuardedBy")
+    @EnableFlags({android.service.notification.Flags.FLAG_NOTIFICATION_CLASSIFICATION,
+            android.app.Flags.FLAG_NOTIFICATION_CLASSIFICATION_UI})
+    public void testPullBundlePreferencesStats_fillsOutStatsEvent()
+            throws Exception {
+        // Create the current user and enable the package
+        int userId = ActivityManager.getCurrentUser();
+        mAssistants.loadDefaultsFromConfig(true);
+        mAssistants.setPackageOrComponentEnabled(mCn.flattenToString(), userId, true,
+                true, true);
+        ManagedServices.ManagedServiceInfo info =
+                mAssistants.new ManagedServiceInfo(null, mCn, userId, false, null, 35, 2345256);
+
+        // Ensure bundling is enabled
+        mAssistants.setAdjustmentTypeSupportedState(info, Adjustment.KEY_TYPE, true);
+        // Enable these specific bundle types
+        mAssistants.setAssistantAdjustmentKeyTypeState(TYPE_PROMOTION, false);
+        mAssistants.setAssistantAdjustmentKeyTypeState(TYPE_NEWS, true);
+        mAssistants.setAssistantAdjustmentKeyTypeState(TYPE_CONTENT_RECOMMENDATION, true);
+
+        // When pullBundlePreferencesStats is run with the given preferences
+        ArrayList<StatsEvent> events = new ArrayList<>();
+        mAssistants.pullBundlePreferencesStats(events);
+
+        // The StatsEvent is filled out with the expected NotificationBundlePreferences values.
+        assertThat(events.size()).isEqualTo(1);
+        AtomsProto.Atom atom = StatsEventTestUtils.convertToAtom(events.get(0));
+
+        // The returned atom does not have external extensions registered.
+        // So we serialize and then deserialize with extensions registered.
+        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);
+        atom = AtomsProto.Atom.parseFrom(codedis, mRegistry);
+        assertTrue(atom.hasExtension(NotificationExtensionAtoms.notificationBundlePreferences));
+        NotificationBundlePreferences p =
+                atom.getExtension(NotificationExtensionAtoms.notificationBundlePreferences);
+        assertThat(p.getBundlesAllowed()).isTrue();
+        assertThat(p.getAllowedBundleTypes(0).getNumber())
+                .isEqualTo(NotificationProtoEnums.TYPE_NEWS);
+        assertThat(p.getAllowedBundleTypes(1).getNumber())
+                .isEqualTo(NotificationProtoEnums.TYPE_CONTENT_RECOMMENDATION);
+
+        // Disable the top-level bundling setting
+        mAssistants.setAdjustmentTypeSupportedState(info, Adjustment.KEY_TYPE, false);
+        // Enable these specific bundle types
+        mAssistants.setAssistantAdjustmentKeyTypeState(TYPE_PROMOTION, true);
+        mAssistants.setAssistantAdjustmentKeyTypeState(TYPE_NEWS, false);
+        mAssistants.setAssistantAdjustmentKeyTypeState(TYPE_CONTENT_RECOMMENDATION, true);
+
+        ArrayList<StatsEvent> eventsDisabled = new ArrayList<>();
+        mAssistants.pullBundlePreferencesStats(eventsDisabled);
+
+        // The StatsEvent is filled out with the expected NotificationBundlePreferences values.
+        assertThat(eventsDisabled.size()).isEqualTo(1);
+        AtomsProto.Atom atomDisabled = StatsEventTestUtils.convertToAtom(eventsDisabled.get(0));
+
+        // The returned atom does not have external extensions registered.
+        // So we serialize and then deserialize with extensions registered.
+        outputStream = new ByteArrayOutputStream();
+        codedos = CodedOutputStream.newInstance(outputStream);
+        atomDisabled.writeTo(codedos);
+        codedos.flush();
+
+        inputStream = new ByteArrayInputStream(outputStream.toByteArray());
+        codedis = CodedInputStream.newInstance(inputStream);
+        atomDisabled = AtomsProto.Atom.parseFrom(codedis, mRegistry);
+        assertTrue(atomDisabled.hasExtension(NotificationExtensionAtoms
+                .notificationBundlePreferences));
+
+        NotificationBundlePreferences p2 =
+                atomDisabled.getExtension(NotificationExtensionAtoms.notificationBundlePreferences);
+        assertThat(p2.getBundlesAllowed()).isFalse();
+        assertThat(p2.getAllowedBundleTypes(0).getNumber())
+                .isEqualTo(NotificationProtoEnums.TYPE_PROMOTION);
+        assertThat(p2.getAllowedBundleTypes(1).getNumber())
+                .isEqualTo(NotificationProtoEnums.TYPE_CONTENT_RECOMMENDATION);
+    }
 }
\ No newline at end of file
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 f41805d40b0d..4bc286117ac6 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/PreferencesHelperTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/PreferencesHelperTest.java
@@ -19,6 +19,7 @@ import static android.app.AppOpsManager.MODE_ALLOWED;
 import static android.app.AppOpsManager.MODE_DEFAULT;
 import static android.app.AppOpsManager.OP_SYSTEM_ALERT_WINDOW;
 import static android.app.Flags.FLAG_MODES_UI;
+import static android.app.Flags.FLAG_NOTIFICATION_CLASSIFICATION_UI;
 import static android.app.Notification.VISIBILITY_PRIVATE;
 import static android.app.Notification.VISIBILITY_SECRET;
 import static android.app.NotificationChannel.ALLOW_BUBBLE_ON;
@@ -161,6 +162,7 @@ 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.os.notification.NotificationProtoEnums;
 import com.android.server.UiServiceTestCase;
 import com.android.server.notification.PermissionHelper.PackagePermission;
 import com.android.server.uri.UriGrantsManagerInternal;
@@ -256,7 +258,8 @@ public class PreferencesHelperTest extends UiServiceTestCase {
     public static List<FlagsParameterization> getParams() {
         return FlagsParameterization.allCombinationsOf(
                 android.app.Flags.FLAG_API_RICH_ONGOING,
-                FLAG_NOTIFICATION_CLASSIFICATION, FLAG_MODES_UI);
+                FLAG_NOTIFICATION_CLASSIFICATION, FLAG_NOTIFICATION_CLASSIFICATION_UI,
+                FLAG_MODES_UI);
     }
 
     public PreferencesHelperTest(FlagsParameterization flags) {
@@ -6135,6 +6138,7 @@ public class PreferencesHelperTest extends UiServiceTestCase {
     }
 
     @Test
+    @DisableFlags({FLAG_NOTIFICATION_CLASSIFICATION_UI})
     public void testPullPackagePreferencesStats_postPermissionMigration()
             throws InvalidProtocolBufferException {
         // make sure there's at least one channel for each package we want to test
@@ -6155,6 +6159,11 @@ public class PreferencesHelperTest extends UiServiceTestCase {
         mHelper.canShowBadge(PKG_O, UID_O);
         mHelper.canShowBadge(PKG_P, UID_P);
 
+        ArrayList<StatsEvent> events = new ArrayList<>();
+
+        mHelper.pullPackagePreferencesStats(events, appPermissions,
+                new ArrayMap<String, Set<Integer>>());
+
         // expected output. format: uid -> importance, as only uid (and not package name)
         // is in PackageNotificationPreferences
         ArrayMap<Integer, Pair<Integer, Boolean>> expected = new ArrayMap<>();
@@ -6162,9 +6171,6 @@ public class PreferencesHelperTest extends UiServiceTestCase {
         expected.put(UID_O, new Pair<>(IMPORTANCE_NONE, true));         // banned by permissions
         expected.put(UID_P, new Pair<>(IMPORTANCE_UNSPECIFIED, false)); // default: unspecified
 
-        ArrayList<StatsEvent> events = new ArrayList<>();
-        mHelper.pullPackagePreferencesStats(events, appPermissions);
-
         assertEquals("total number of packages", 3, events.size());
         for (StatsEvent ev : events) {
             AtomsProto.Atom atom = StatsEventTestUtils.convertToAtom(ev);
@@ -6180,6 +6186,74 @@ public class PreferencesHelperTest extends UiServiceTestCase {
     }
 
     @Test
+    @EnableFlags({FLAG_NOTIFICATION_CLASSIFICATION, FLAG_NOTIFICATION_CLASSIFICATION_UI})
+    public void testPullPackagePreferencesStats_createsExpectedStatsEvents()
+            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,
+                UID_N_MR1, false);
+        NotificationChannel channelB = new NotificationChannel("b", "b", IMPORTANCE_DEFAULT);
+        mHelper.createNotificationChannel(PKG_O, UID_O, channelB, true, false, UID_O, false);
+        NotificationChannel channelC = new NotificationChannel("c", "c", IMPORTANCE_DEFAULT);
+        mHelper.createNotificationChannel(PKG_P, UID_P, channelC, true, false, UID_P, false);
+
+        // build a collection of app permissions that should be passed in and used
+        ArrayMap<Pair<Integer, String>, Pair<Boolean, Boolean>> pkgPermissions = new ArrayMap<>();
+        pkgPermissions.put(new Pair<>(UID_N_MR1, PKG_N_MR1), new Pair<>(true, false));
+        pkgPermissions.put(new Pair<>(UID_O, PKG_O), new Pair<>(false, true)); // in local prefs
+
+        // local preferences
+        mHelper.canShowBadge(PKG_O, UID_O);
+        mHelper.canShowBadge(PKG_P, UID_P);
+
+        // Sets bundles_allowed to true for these packages.
+        ArrayMap<String, Set<Integer>> packageSpecificAdjustmentKeyTypes = new ArrayMap<>();
+        Set<Integer> nMr1BundlesSet = new ArraySet<Integer>();
+        nMr1BundlesSet.add(TYPE_NEWS);
+        nMr1BundlesSet.add(TYPE_SOCIAL_MEDIA);
+        packageSpecificAdjustmentKeyTypes.put(PKG_N_MR1, nMr1BundlesSet);
+        Set<Integer> pBundlesSet = new ArraySet<Integer>();
+        packageSpecificAdjustmentKeyTypes.put(PKG_P, pBundlesSet);
+
+        ArrayList<StatsEvent> events = new ArrayList<>();
+
+        mHelper.pullPackagePreferencesStats(events, pkgPermissions,
+                packageSpecificAdjustmentKeyTypes);
+
+        assertEquals("total number of packages", 3, events.size());
+
+        AtomsProto.Atom atom0 = StatsEventTestUtils.convertToAtom(events.get(0));
+        assertTrue(atom0.hasPackageNotificationPreferences());
+        PackageNotificationPreferences p0 = atom0.getPackageNotificationPreferences();
+        assertThat(p0.getUid()).isEqualTo(UID_O);
+        assertThat(p0.getImportance()).isEqualTo(IMPORTANCE_NONE); // banned by permissions
+        assertThat(p0.getUserSetImportance()).isTrue();
+        assertThat(p0.getAllowedBundleTypesList()).hasSize(0);
+
+        AtomsProto.Atom atom1 = StatsEventTestUtils.convertToAtom(events.get(1));
+        assertTrue(atom1.hasPackageNotificationPreferences());
+        PackageNotificationPreferences p1 = atom1.getPackageNotificationPreferences();
+        assertThat(p1.getUid()).isEqualTo(UID_N_MR1);
+        assertThat(p1.getImportance()).isEqualTo(IMPORTANCE_DEFAULT);
+        assertThat(p1.getUserSetImportance()).isFalse();
+        assertThat(p1.getAllowedBundleTypesList()).hasSize(2);
+
+        assertThat(p1.getAllowedBundleTypes(0).getNumber())
+                .isEqualTo(NotificationProtoEnums.TYPE_SOCIAL_MEDIA);
+        assertThat(p1.getAllowedBundleTypes(1).getNumber())
+                .isEqualTo(NotificationProtoEnums.TYPE_NEWS);
+
+        AtomsProto.Atom atom2 = StatsEventTestUtils.convertToAtom(events.get(2));
+        assertTrue(atom2.hasPackageNotificationPreferences());
+        PackageNotificationPreferences p2 = atom2.getPackageNotificationPreferences();
+        assertThat(p2.getUid()).isEqualTo(UID_P);
+        assertThat(p2.getImportance()).isEqualTo(IMPORTANCE_UNSPECIFIED); // default: unspecified
+        assertThat(p2.getUserSetImportance()).isFalse();
+        assertThat(p2.getAllowedBundleTypesList()).hasSize(0);
+    }
+
+    @Test
     public void testUnlockNotificationChannelImportance() {
         NotificationChannel channel = new NotificationChannel("a", "a", IMPORTANCE_LOW);
         mHelper.createNotificationChannel(PKG_O, UID_O, channel, true, false, UID_O, false);