Merge "Correctly propagate origin to ZenModeEventLogger" into main
diff --git a/services/core/java/com/android/server/notification/ZenModeEventLogger.java b/services/core/java/com/android/server/notification/ZenModeEventLogger.java
index 87158cd..df570a0 100644
--- a/services/core/java/com/android/server/notification/ZenModeEventLogger.java
+++ b/services/core/java/com/android/server/notification/ZenModeEventLogger.java
@@ -29,6 +29,7 @@
 import android.os.Process;
 import android.service.notification.DNDPolicyProto;
 import android.service.notification.ZenModeConfig;
+import android.service.notification.ZenModeConfig.ConfigChangeOrigin;
 import android.service.notification.ZenModeDiff;
 import android.service.notification.ZenPolicy;
 import android.util.ArrayMap;
@@ -58,7 +59,7 @@
     // mode change.
     ZenModeEventLogger.ZenStateChanges mChangeState = new ZenModeEventLogger.ZenStateChanges();
 
-    private PackageManager mPm;
+    private final PackageManager mPm;
 
     ZenModeEventLogger(PackageManager pm) {
         mPm = pm;
@@ -97,11 +98,11 @@
      * @param newInfo     ZenModeInfo after this change takes effect
      * @param callingUid  the calling UID associated with the change; may be used to attribute the
      *                    change to a particular package or determine if this is a user action
-     * @param fromSystemOrSystemUi whether the calling UID is either system UID or system UI
+     * @param origin      The origin of the Zen change.
      */
     public final void maybeLogZenChange(ZenModeInfo prevInfo, ZenModeInfo newInfo, int callingUid,
-            boolean fromSystemOrSystemUi) {
-        mChangeState.init(prevInfo, newInfo, callingUid, fromSystemOrSystemUi);
+            @ConfigChangeOrigin int origin) {
+        mChangeState.init(prevInfo, newInfo, callingUid, origin);
         if (mChangeState.shouldLogChanges()) {
             maybeReassignCallingUid();
             logChanges();
@@ -124,7 +125,7 @@
         // We don't consider the manual rule in the old config because if a manual rule is turning
         // off with a call from system, that could easily be a user action to explicitly turn it off
         if (mChangeState.getChangedRuleType() == RULE_TYPE_MANUAL) {
-            if (!mChangeState.mFromSystemOrSystemUi
+            if (!mChangeState.isFromSystemOrSystemUi()
                     || mChangeState.getNewManualRuleEnabler() == null) {
                 return;
             }
@@ -136,7 +137,7 @@
         //   - we've determined it's not a user action
         //   - our current best guess is that the calling uid is system/sysui
         if (mChangeState.getChangedRuleType() == RULE_TYPE_AUTOMATIC) {
-            if (mChangeState.getIsUserAction() || !mChangeState.mFromSystemOrSystemUi) {
+            if (mChangeState.getIsUserAction() || !mChangeState.isFromSystemOrSystemUi()) {
                 return;
             }
 
@@ -221,10 +222,10 @@
         ZenModeConfig mPrevConfig, mNewConfig;
         NotificationManager.Policy mPrevPolicy, mNewPolicy;
         int mCallingUid = Process.INVALID_UID;
-        boolean mFromSystemOrSystemUi = false;
+        @ConfigChangeOrigin int mOrigin = ZenModeConfig.UPDATE_ORIGIN_UNKNOWN;
 
         private void init(ZenModeInfo prevInfo, ZenModeInfo newInfo, int callingUid,
-                boolean fromSystemOrSystemUi) {
+                @ConfigChangeOrigin int origin) {
             // previous & new may be the same -- that would indicate that zen mode hasn't changed.
             mPrevZenMode = prevInfo.mZenMode;
             mNewZenMode = newInfo.mZenMode;
@@ -233,7 +234,7 @@
             mPrevPolicy = prevInfo.mPolicy;
             mNewPolicy = newInfo.mPolicy;
             mCallingUid = callingUid;
-            mFromSystemOrSystemUi = fromSystemOrSystemUi;
+            mOrigin = origin;
         }
 
         /**
@@ -389,12 +390,16 @@
 
         /**
          * Return our best guess as to whether the changes observed are due to a user action.
-         * Note that this won't be 100% accurate as we can't necessarily distinguish between a
-         * system uid call indicating "user interacted with Settings" vs "a system app changed
-         * something automatically".
+         * Note that this (before {@code MODES_API}) won't be 100% accurate as we can't necessarily
+         * distinguish between a system uid call indicating "user interacted with Settings" vs "a
+         * system app changed something automatically".
          */
         boolean getIsUserAction() {
-            // Approach:
+            if (Flags.modesApi()) {
+                return mOrigin == ZenModeConfig.UPDATE_ORIGIN_USER;
+            }
+
+            // Approach for pre-MODES_API:
             //   - if manual rule turned on or off, the calling UID is system, and the new manual
             //     rule does not have an enabler set, guess that this is likely to be a user action.
             //     This may represent a system app turning on DND automatically, but we guess "user"
@@ -419,13 +424,13 @@
             switch (getChangedRuleType()) {
                 case RULE_TYPE_MANUAL:
                     // TODO(b/278888961): Distinguish the automatically-turned-off state
-                    return mFromSystemOrSystemUi && (getNewManualRuleEnabler() == null);
+                    return isFromSystemOrSystemUi() && (getNewManualRuleEnabler() == null);
                 case RULE_TYPE_AUTOMATIC:
                     for (ZenModeDiff.RuleDiff d : getChangedAutomaticRules().values()) {
                         if (d.wasAdded() || d.wasRemoved()) {
                             // If the change comes from system, a rule being added/removed indicates
                             // a likely user action. From an app, it's harder to know for sure.
-                            return mFromSystemOrSystemUi;
+                            return isFromSystemOrSystemUi();
                         }
                         ZenModeDiff.FieldDiff enabled = d.getDiffForField(
                                 ZenModeDiff.RuleDiff.FIELD_ENABLED);
@@ -455,6 +460,13 @@
             return false;
         }
 
+        boolean isFromSystemOrSystemUi() {
+            return mOrigin == ZenModeConfig.UPDATE_ORIGIN_INIT
+                    || mOrigin == ZenModeConfig.UPDATE_ORIGIN_INIT_USER
+                    || mOrigin == ZenModeConfig.UPDATE_ORIGIN_SYSTEM_OR_SYSTEMUI
+                    || mOrigin == ZenModeConfig.UPDATE_ORIGIN_RESTORE_BACKUP;
+        }
+
         /**
          * Get the package UID associated with this change, which is just the calling UID for the
          * relevant method changes. This may get reset by ZenModeEventLogger, which has access to
@@ -612,7 +624,7 @@
             copy.mPrevPolicy = mPrevPolicy.copy();
             copy.mNewPolicy = mNewPolicy.copy();
             copy.mCallingUid = mCallingUid;
-            copy.mFromSystemOrSystemUi = mFromSystemOrSystemUi;
+            copy.mOrigin = mOrigin;
             return copy;
         }
     }
diff --git a/services/core/java/com/android/server/notification/ZenModeHelper.java b/services/core/java/com/android/server/notification/ZenModeHelper.java
index d1de9b0..0a46901 100644
--- a/services/core/java/com/android/server/notification/ZenModeHelper.java
+++ b/services/core/java/com/android/server/notification/ZenModeHelper.java
@@ -1371,12 +1371,8 @@
         if (logZenModeEvents) {
             ZenModeEventLogger.ZenModeInfo newInfo = new ZenModeEventLogger.ZenModeInfo(
                     mZenMode, mConfig, mConsolidatedPolicy);
-            boolean fromSystemOrSystemUi = origin == UPDATE_ORIGIN_SYSTEM_OR_SYSTEMUI
-                    || origin == UPDATE_ORIGIN_INIT
-                    || origin == UPDATE_ORIGIN_INIT_USER
-                    || origin == UPDATE_ORIGIN_RESTORE_BACKUP;
             mZenModeEventLogger.maybeLogZenChange(prevInfo, newInfo, callingUid,
-                    fromSystemOrSystemUi);
+                    origin);
         }
     }
 
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/ZenModeEventLoggerFake.java b/services/tests/uiservicestests/src/com/android/server/notification/ZenModeEventLoggerFake.java
index 4a1435f..1fcee06 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/ZenModeEventLoggerFake.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/ZenModeEventLoggerFake.java
@@ -99,7 +99,7 @@
     public boolean getFromSystemOrSystemUi(int i) throws IllegalArgumentException {
         // While this isn't a logged output value, it's still helpful to check in tests.
         checkInRange(i);
-        return mChanges.get(i).mFromSystemOrSystemUi;
+        return mChanges.get(i).isFromSystemOrSystemUi();
     }
 
     public boolean getIsUserAction(int i) throws IllegalArgumentException {
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 bc63c29e..44f0894 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java
@@ -43,6 +43,8 @@
 import static android.provider.Settings.Global.ZEN_MODE_ALARMS;
 import static android.provider.Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS;
 import static android.provider.Settings.Global.ZEN_MODE_OFF;
+import static android.service.notification.Condition.SOURCE_SCHEDULE;
+import static android.service.notification.Condition.SOURCE_USER_ACTION;
 import static android.service.notification.Condition.STATE_FALSE;
 import static android.service.notification.Condition.STATE_TRUE;
 import static android.service.notification.ZenModeConfig.UPDATE_ORIGIN_APP;
@@ -120,12 +122,13 @@
 import android.service.notification.DeviceEffectsApplier;
 import android.service.notification.ZenDeviceEffects;
 import android.service.notification.ZenModeConfig;
+import android.service.notification.ZenModeConfig.ConfigChangeOrigin;
 import android.service.notification.ZenModeConfig.ScheduleInfo;
 import android.service.notification.ZenModeConfig.ZenRule;
 import android.service.notification.ZenModeDiff;
 import android.service.notification.ZenPolicy;
 import android.test.suitebuilder.annotation.SmallTest;
-import android.testing.AndroidTestingRunner;
+import android.testing.TestWithLooperRule;
 import android.testing.TestableLooper;
 import android.util.ArrayMap;
 import android.util.Log;
@@ -148,6 +151,8 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.truth.Correspondence;
 import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.testing.junit.testparameterinjector.TestParameter;
+import com.google.testing.junit.testparameterinjector.TestParameterInjector;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -173,7 +178,7 @@
 
 @SmallTest
 @SuppressLint("GuardedBy") // It's ok for this test to access guarded methods from the service.
-@RunWith(AndroidTestingRunner.class)
+@RunWith(TestParameterInjector.class)
 @TestableLooper.RunWithLooper
 public class ZenModeHelperTest extends UiServiceTestCase {
 
@@ -215,6 +220,9 @@
     public final SetFlagsRule mSetFlagsRule = new SetFlagsRule(
             SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT);
 
+    @Rule(order = Integer.MAX_VALUE) // set the highest order so it's the innermost rule
+    public TestWithLooperRule mLooperRule = new TestWithLooperRule();
+
     ConditionProviders mConditionProviders;
     @Mock NotificationManager mNotificationManager;
     @Mock PackageManager mPackageManager;
@@ -2341,15 +2349,38 @@
         assertEquals(ZEN_MODE_OFF, mZenModeHelper.mZenMode);
     }
 
+    private enum ModesApiFlag {
+        ENABLED(true, /* originForUserActionInSystemUi= */ UPDATE_ORIGIN_USER),
+        DISABLED(false, /* originForUserActionInSystemUi= */ UPDATE_ORIGIN_SYSTEM_OR_SYSTEMUI);
+
+        private final boolean mEnabled;
+        @ConfigChangeOrigin private final int mOriginForUserActionInSystemUi;
+
+        ModesApiFlag(boolean enabled, @ConfigChangeOrigin int originForUserActionInSystemUi) {
+            this.mEnabled = enabled;
+            this.mOriginForUserActionInSystemUi = originForUserActionInSystemUi;
+        }
+
+        void applyFlag(SetFlagsRule setFlagsRule) {
+            if (mEnabled) {
+                setFlagsRule.enableFlags(Flags.FLAG_MODES_API);
+            } else {
+                setFlagsRule.disableFlags(Flags.FLAG_MODES_API);
+            }
+        }
+    }
+
     @Test
-    public void testZenModeEventLog_setManualZenMode() throws IllegalArgumentException {
+    public void testZenModeEventLog_setManualZenMode(@TestParameter ModesApiFlag modesApiFlag)
+            throws IllegalArgumentException {
+        modesApiFlag.applyFlag(mSetFlagsRule);
         mTestFlagResolver.setFlagOverride(LOG_DND_STATE_EVENTS, true);
         setupZenConfig();
 
         // Turn zen mode on (to important_interruptions)
         // Need to additionally call the looper in order to finish the post-apply-config process
         mZenModeHelper.setManualZenMode(ZEN_MODE_IMPORTANT_INTERRUPTIONS, null,
-                UPDATE_ORIGIN_SYSTEM_OR_SYSTEMUI, "", null, Process.SYSTEM_UID);
+                modesApiFlag.mOriginForUserActionInSystemUi, "", null, Process.SYSTEM_UID);
 
         // Now turn zen mode off, but via a different package UID -- this should get registered as
         // "not an action by the user" because some other app is changing zen mode
@@ -2376,7 +2407,8 @@
         assertEquals(ZEN_MODE_IMPORTANT_INTERRUPTIONS, mZenModeEventLogger.getNewZenMode(0));
         assertEquals(DNDProtoEnums.MANUAL_RULE, mZenModeEventLogger.getChangedRuleType(0));
         assertEquals(1, mZenModeEventLogger.getNumRulesActive(0));
-        assertTrue(mZenModeEventLogger.getFromSystemOrSystemUi(0));
+        assertThat(mZenModeEventLogger.getFromSystemOrSystemUi(0)).isEqualTo(
+                modesApiFlag == ModesApiFlag.DISABLED);
         assertTrue(mZenModeEventLogger.getIsUserAction(0));
         assertEquals(Process.SYSTEM_UID, mZenModeEventLogger.getPackageUid(0));
         checkDndProtoMatchesSetupZenConfig(mZenModeEventLogger.getPolicyProto(0));
@@ -2401,7 +2433,9 @@
     }
 
     @Test
-    public void testZenModeEventLog_automaticRules() throws IllegalArgumentException {
+    public void testZenModeEventLog_automaticRules(@TestParameter ModesApiFlag modesApiFlag)
+            throws IllegalArgumentException {
+        modesApiFlag.applyFlag(mSetFlagsRule);
         mTestFlagResolver.setFlagOverride(LOG_DND_STATE_EVENTS, true);
         setupZenConfig();
 
@@ -2423,8 +2457,8 @@
 
         // Event 2: "User" turns off the automatic rule (sets it to not enabled)
         zenRule.setEnabled(false);
-        mZenModeHelper.updateAutomaticZenRule(id, zenRule, UPDATE_ORIGIN_SYSTEM_OR_SYSTEMUI, "",
-                Process.SYSTEM_UID);
+        mZenModeHelper.updateAutomaticZenRule(id, zenRule,
+                modesApiFlag.mOriginForUserActionInSystemUi, "", Process.SYSTEM_UID);
 
         // Add a new system rule
         AutomaticZenRule systemRule = new AutomaticZenRule("systemRule",
@@ -2442,8 +2476,8 @@
                 UPDATE_ORIGIN_SYSTEM_OR_SYSTEMUI, Process.SYSTEM_UID);
 
         // Event 4: "User" deletes the rule
-        mZenModeHelper.removeAutomaticZenRule(systemId, UPDATE_ORIGIN_SYSTEM_OR_SYSTEMUI, "",
-                Process.SYSTEM_UID);
+        mZenModeHelper.removeAutomaticZenRule(systemId, modesApiFlag.mOriginForUserActionInSystemUi,
+                "", Process.SYSTEM_UID);
 
         // In total, this represents 4 events
         assertEquals(4, mZenModeEventLogger.numLoggedChanges());
@@ -2499,20 +2533,109 @@
     }
 
     @Test
-    public void testZenModeEventLog_policyChanges() throws IllegalArgumentException {
+    @EnableFlags(Flags.FLAG_MODES_API)
+    public void testZenModeEventLog_automaticRuleActivatedFromAppByAppAndUser()
+            throws IllegalArgumentException {
+        mTestFlagResolver.setFlagOverride(LOG_DND_STATE_EVENTS, true);
+        setupZenConfig();
+
+        // Ann app adds an automatic zen rule
+        AutomaticZenRule zenRule = new AutomaticZenRule("name",
+                null,
+                new ComponentName(CUSTOM_PKG_NAME, "ScheduleConditionProvider"),
+                ZenModeConfig.toScheduleConditionId(new ScheduleInfo()),
+                null,
+                NotificationManager.INTERRUPTION_FILTER_PRIORITY, true);
+        String id = mZenModeHelper.addAutomaticZenRule(mContext.getPackageName(), zenRule,
+                UPDATE_ORIGIN_APP, "test", CUSTOM_PKG_UID);
+
+        // Event 1: Mimic the rule coming on manually when the user turns it on in the app
+        // ("Turn on bedtime now" because user goes to bed earlier).
+        mZenModeHelper.setAutomaticZenRuleState(id,
+                new Condition(zenRule.getConditionId(), "", STATE_TRUE, SOURCE_USER_ACTION),
+                UPDATE_ORIGIN_USER, CUSTOM_PKG_UID);
+
+        // Event 2: App deactivates the rule automatically (it's 8 AM, bedtime schedule ends)
+        mZenModeHelper.setAutomaticZenRuleState(id,
+                new Condition(zenRule.getConditionId(), "", STATE_FALSE, SOURCE_SCHEDULE),
+                UPDATE_ORIGIN_APP, CUSTOM_PKG_UID);
+
+        // Event 3: App activates the rule automatically (it's now 11 PM, bedtime schedule starts)
+        mZenModeHelper.setAutomaticZenRuleState(id,
+                new Condition(zenRule.getConditionId(), "", STATE_TRUE, SOURCE_SCHEDULE),
+                UPDATE_ORIGIN_APP, CUSTOM_PKG_UID);
+
+        // Event 4: User deactivates the rule manually (they get up before 8 AM on the next day)
+        mZenModeHelper.setAutomaticZenRuleState(id,
+                new Condition(zenRule.getConditionId(), "", STATE_FALSE, SOURCE_USER_ACTION),
+                UPDATE_ORIGIN_USER, CUSTOM_PKG_UID);
+
+        // In total, this represents 4 events
+        assertEquals(4, mZenModeEventLogger.numLoggedChanges());
+
+        // Automatic rule turning on manually:
+        //   - event ID: DND_TURNED_ON
+        //   - 1 rule (newly) active
+        //   - is a user action
+        //   - package UID is the calling package
+        assertEquals(ZenModeEventLogger.ZenStateChangedEvent.DND_TURNED_ON.getId(),
+                mZenModeEventLogger.getEventId(0));
+        assertEquals(1, mZenModeEventLogger.getNumRulesActive(0));
+        assertTrue(mZenModeEventLogger.getIsUserAction(0));
+        assertEquals(CUSTOM_PKG_UID, mZenModeEventLogger.getPackageUid(0));
+
+        // Automatic rule turned off automatically by app:
+        //   - event ID: DND_TURNED_OFF
+        //   - 0 rules active
+        //   - is not a user action
+        //   - package UID is the calling package
+        assertEquals(ZenModeEventLogger.ZenStateChangedEvent.DND_TURNED_OFF.getId(),
+                mZenModeEventLogger.getEventId(1));
+        assertEquals(0, mZenModeEventLogger.getNumRulesActive(1));
+        assertFalse(mZenModeEventLogger.getIsUserAction(1));
+        assertEquals(CUSTOM_PKG_UID, mZenModeEventLogger.getPackageUid(1));
+
+        // Automatic rule turned on automatically by app:
+        //   - event ID: DND_TURNED_ON
+        //   - 1 rule (newly) active
+        //   - is not a user action
+        //   - package UID is the calling package
+        assertEquals(ZenModeEventLogger.ZenStateChangedEvent.DND_TURNED_ON.getId(),
+                mZenModeEventLogger.getEventId(2));
+        assertEquals(DNDProtoEnums.AUTOMATIC_RULE, mZenModeEventLogger.getChangedRuleType(2));
+        assertEquals(1, mZenModeEventLogger.getNumRulesActive(2));
+        assertFalse(mZenModeEventLogger.getIsUserAction(2));
+        assertEquals(CUSTOM_PKG_UID, mZenModeEventLogger.getPackageUid(2));
+
+        // Automatic rule turned off automatically by the user:
+        //   - event ID: DND_TURNED_ON
+        //   - 0 rules active
+        //   - is a user action
+        //   - package UID is the calling package
+        assertEquals(ZenModeEventLogger.ZenStateChangedEvent.DND_TURNED_OFF.getId(),
+                mZenModeEventLogger.getEventId(3));
+        assertEquals(0, mZenModeEventLogger.getNumRulesActive(3));
+        assertTrue(mZenModeEventLogger.getIsUserAction(3));
+        assertEquals(CUSTOM_PKG_UID, mZenModeEventLogger.getPackageUid(3));
+    }
+
+    @Test
+    public void testZenModeEventLog_policyChanges(@TestParameter ModesApiFlag modesApiFlag)
+            throws IllegalArgumentException {
+        modesApiFlag.applyFlag(mSetFlagsRule);
         mTestFlagResolver.setFlagOverride(LOG_DND_STATE_EVENTS, true);
         setupZenConfig();
 
         // First just turn zen mode on
         mZenModeHelper.setManualZenMode(ZEN_MODE_IMPORTANT_INTERRUPTIONS, null,
-                UPDATE_ORIGIN_SYSTEM_OR_SYSTEMUI, "", null, Process.SYSTEM_UID);
+                modesApiFlag.mOriginForUserActionInSystemUi, "", null, Process.SYSTEM_UID);
 
         // Now change the policy slightly; want to confirm that this'll be reflected in the logs
         ZenModeConfig newConfig = mZenModeHelper.mConfig.copy();
         newConfig.allowAlarms = true;
         newConfig.allowRepeatCallers = false;
         mZenModeHelper.setNotificationPolicy(newConfig.toNotificationPolicy(),
-                UPDATE_ORIGIN_SYSTEM_OR_SYSTEMUI, Process.SYSTEM_UID);
+                modesApiFlag.mOriginForUserActionInSystemUi, Process.SYSTEM_UID);
 
         // Turn zen mode off; we want to make sure policy changes do not get logged when zen mode
         // is off.
@@ -2523,7 +2646,7 @@
         newConfig.allowMessages = false;
         newConfig.allowRepeatCallers = true;
         mZenModeHelper.setNotificationPolicy(newConfig.toNotificationPolicy(),
-                UPDATE_ORIGIN_SYSTEM_OR_SYSTEMUI, Process.SYSTEM_UID);
+                modesApiFlag.mOriginForUserActionInSystemUi, Process.SYSTEM_UID);
 
         // Total events: we only expect ones for turning on, changing policy, and turning off
         assertEquals(3, mZenModeEventLogger.numLoggedChanges());
@@ -2556,7 +2679,9 @@
     }
 
     @Test
-    public void testZenModeEventLog_ruleCounts() throws IllegalArgumentException {
+    public void testZenModeEventLog_ruleCounts(@TestParameter ModesApiFlag modesApiFlag)
+            throws IllegalArgumentException {
+        modesApiFlag.applyFlag(mSetFlagsRule);
         mTestFlagResolver.setFlagOverride(LOG_DND_STATE_EVENTS, true);
         setupZenConfig();
 
@@ -2659,8 +2784,10 @@
     }
 
     @Test
-    public void testZenModeEventLog_noLogWithNoConfigChange() throws IllegalArgumentException {
+    public void testZenModeEventLog_noLogWithNoConfigChange(
+            @TestParameter ModesApiFlag modesApiFlag) throws IllegalArgumentException {
         // If evaluateZenMode is called independently of a config change, don't log.
+        modesApiFlag.applyFlag(mSetFlagsRule);
         mTestFlagResolver.setFlagOverride(LOG_DND_STATE_EVENTS, true);
         setupZenConfig();
 
@@ -2677,9 +2804,11 @@
     }
 
     @Test
-    public void testZenModeEventLog_reassignUid() throws IllegalArgumentException {
+    public void testZenModeEventLog_reassignUid(@TestParameter ModesApiFlag modesApiFlag)
+            throws IllegalArgumentException {
         // Test that, only in specific cases, we reassign the calling UID to one associated with
         // the automatic rule owner.
+        modesApiFlag.applyFlag(mSetFlagsRule);
         mTestFlagResolver.setFlagOverride(LOG_DND_STATE_EVENTS, true);
         setupZenConfig();
 
@@ -2691,7 +2820,7 @@
                 null,
                 NotificationManager.INTERRUPTION_FILTER_PRIORITY, true);
         String id = mZenModeHelper.addAutomaticZenRule(mContext.getPackageName(), zenRule,
-                UPDATE_ORIGIN_SYSTEM_OR_SYSTEMUI, "test", Process.SYSTEM_UID);
+                UPDATE_ORIGIN_APP, "test", Process.SYSTEM_UID);
 
         // Rule 2, same as rule 1 but owned by the system
         AutomaticZenRule zenRule2 = new AutomaticZenRule("name2",
@@ -2701,11 +2830,11 @@
                 null,
                 NotificationManager.INTERRUPTION_FILTER_PRIORITY, true);
         String id2 = mZenModeHelper.addAutomaticZenRule(mContext.getPackageName(), zenRule2,
-                UPDATE_ORIGIN_SYSTEM_OR_SYSTEMUI, "test", Process.SYSTEM_UID);
+                modesApiFlag.mOriginForUserActionInSystemUi, "test", Process.SYSTEM_UID);
 
         // Turn on rule 1; call looks like it's from the system. Because setting a condition is
         // typically an automatic (non-user-initiated) action, expect the calling UID to be
-        // re-evaluated to the one associat.d with CUSTOM_PKG_NAME.
+        // re-evaluated to the one associated with CUSTOM_PKG_NAME.
         mZenModeHelper.setAutomaticZenRuleState(id,
                 new Condition(zenRule.getConditionId(), "", STATE_TRUE),
                 UPDATE_ORIGIN_SYSTEM_OR_SYSTEMUI, Process.SYSTEM_UID);
@@ -2719,8 +2848,8 @@
         // Disable rule 1. Because this looks like a user action, the UID should not be modified
         // from the system-provided one.
         zenRule.setEnabled(false);
-        mZenModeHelper.updateAutomaticZenRule(id, zenRule, UPDATE_ORIGIN_SYSTEM_OR_SYSTEMUI, "",
-                Process.SYSTEM_UID);
+        mZenModeHelper.updateAutomaticZenRule(id, zenRule,
+                modesApiFlag.mOriginForUserActionInSystemUi, "", Process.SYSTEM_UID);
 
         // Add a manual rule. Any manual rule changes should not get calling uids reassigned.
         mZenModeHelper.setManualZenMode(ZEN_MODE_IMPORTANT_INTERRUPTIONS, null, UPDATE_ORIGIN_APP,
@@ -2777,8 +2906,10 @@
     }
 
     @Test
-    public void testZenModeEventLog_channelsBypassingChanges() {
+    public void testZenModeEventLog_channelsBypassingChanges(
+            @TestParameter ModesApiFlag modesApiFlag) {
         // Verify that the right thing happens when the canBypassDnd value changes.
+        modesApiFlag.applyFlag(mSetFlagsRule);
         mTestFlagResolver.setFlagOverride(LOG_DND_STATE_EVENTS, true);
         setupZenConfig();
 
@@ -2874,14 +3005,11 @@
         // Second message where we change the policy:
         //   - DND_POLICY_CHANGED (indicates only the policy changed and nothing else)
         //   - rule type: unknown (it's a policy change, not a rule change)
-        //   - user action (because it comes from a "system" uid)
         //   - change is in allow channels, and final policy
         assertThat(mZenModeEventLogger.getEventId(1))
                 .isEqualTo(ZenModeEventLogger.ZenStateChangedEvent.DND_POLICY_CHANGED.getId());
         assertThat(mZenModeEventLogger.getChangedRuleType(1))
                 .isEqualTo(DNDProtoEnums.UNKNOWN_RULE);
-        assertThat(mZenModeEventLogger.getIsUserAction(1)).isTrue();
-        assertThat(mZenModeEventLogger.getPackageUid(1)).isEqualTo(Process.SYSTEM_UID);
         DNDPolicyProto dndProto = mZenModeEventLogger.getPolicyProto(1);
         assertThat(dndProto.getAllowChannels().getNumber())
                 .isEqualTo(DNDProtoEnums.CHANNEL_TYPE_NONE);
diff --git a/tests/testables/src/android/testing/TestWithLooperRule.java b/tests/testables/src/android/testing/TestWithLooperRule.java
index 99b303e..37b39c3 100644
--- a/tests/testables/src/android/testing/TestWithLooperRule.java
+++ b/tests/testables/src/android/testing/TestWithLooperRule.java
@@ -19,7 +19,6 @@
 import android.testing.TestableLooper.LooperFrameworkMethod;
 import android.testing.TestableLooper.RunWithLooper;
 
-import org.junit.internal.runners.statements.InvokeMethod;
 import org.junit.rules.MethodRule;
 import org.junit.runner.RunWith;
 import org.junit.runners.model.FrameworkMethod;
@@ -79,13 +78,11 @@
             while (next != null) {
                 switch (next.getClass().getSimpleName()) {
                     case "RunAfters":
-                        this.<List<FrameworkMethod>>wrapFieldMethodFor(next,
-                                next.getClass(), "afters", method, target);
+                        this.wrapFieldMethodFor(next, "afters", method, target);
                         next = getNextStatement(next, "next");
                         break;
                     case "RunBefores":
-                        this.<List<FrameworkMethod>>wrapFieldMethodFor(next,
-                                next.getClass(), "befores", method, target);
+                        this.wrapFieldMethodFor(next, "befores", method, target);
                         next = getNextStatement(next, "next");
                         break;
                     case "FailOnTimeout":
@@ -95,8 +92,10 @@
                         next = getNextStatement(next, "originalStatement");
                         break;
                     case "InvokeMethod":
-                        this.<FrameworkMethod>wrapFieldMethodFor(next,
-                                InvokeMethod.class, "testMethod", method, target);
+                        this.wrapFieldMethodFor(next, "testMethod", method, target);
+                        return;
+                    case "InvokeParameterizedMethod":
+                        this.wrapFieldMethodFor(next, "frameworkMethod", method, target);
                         return;
                     default:
                         throw new Exception(
@@ -112,12 +111,11 @@
 
     // Wrapping the befores, afters, and InvokeMethods with LooperFrameworkMethod
     // within the statement.
-    private <T> void wrapFieldMethodFor(Statement base, Class<?> targetClass, String fieldStr,
-            FrameworkMethod method, Object target)
-            throws NoSuchFieldException, IllegalAccessException {
-        Field field = targetClass.getDeclaredField(fieldStr);
+    private void wrapFieldMethodFor(Statement base, String fieldStr, FrameworkMethod method,
+            Object target) throws NoSuchFieldException, IllegalAccessException {
+        Field field = base.getClass().getDeclaredField(fieldStr);
         field.setAccessible(true);
-        T fieldInstance = (T) field.get(base);
+        Object fieldInstance = field.get(base);
         if (fieldInstance instanceof FrameworkMethod) {
             field.set(base, looperWrap(method, target, (FrameworkMethod) fieldInstance));
         } else {