diff --git a/tests/robotests/src/com/android/settings/notification/zen/ZenModeBackendTest.java b/tests/robotests/src/com/android/settings/notification/zen/ZenModeBackendTest.java
index 5c3b5b8..32bf9af 100644
--- a/tests/robotests/src/com/android/settings/notification/zen/ZenModeBackendTest.java
+++ b/tests/robotests/src/com/android/settings/notification/zen/ZenModeBackendTest.java
@@ -16,9 +16,9 @@
 
 import static junit.framework.Assert.assertEquals;
 
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -31,7 +31,6 @@
 import android.service.notification.ZenModeConfig;
 
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -189,7 +188,6 @@
         }
     }
 
-    @Ignore
     @Test
     public void saveConversationSenders_importantToNone() {
         when(mNotificationManager.getNotificationPolicy()).thenReturn(
@@ -204,7 +202,11 @@
         mBackend.saveConversationSenders(CONVERSATION_SENDERS_NONE);
 
         ArgumentCaptor<Policy> captor = ArgumentCaptor.forClass(Policy.class);
-        verify(mNotificationManager, times(1)).setNotificationPolicy(captor.capture());
+        if (android.app.Flags.modesApi()) {
+            verify(mNotificationManager).setNotificationPolicy(captor.capture(), eq(true));
+        } else {
+            verify(mNotificationManager).setNotificationPolicy(captor.capture());
+        }
 
         Policy expected = new Policy(
                 PRIORITY_CATEGORY_CALLS | PRIORITY_CATEGORY_MESSAGES | PRIORITY_CATEGORY_ALARMS,
@@ -215,7 +217,6 @@
         assertEquals(expected, captor.getValue());
     }
 
-    @Ignore
     @Test
     public void saveConversationSenders_noneToAll() {
         when(mNotificationManager.getNotificationPolicy()).thenReturn(new Policy(
@@ -229,7 +230,11 @@
         mBackend.saveConversationSenders(CONVERSATION_SENDERS_ANYONE);
 
         ArgumentCaptor<Policy> captor = ArgumentCaptor.forClass(Policy.class);
-        verify(mNotificationManager, times(1)).setNotificationPolicy(captor.capture());
+        if (android.app.Flags.modesApi()) {
+            verify(mNotificationManager).setNotificationPolicy(captor.capture(), eq(true));
+        } else {
+            verify(mNotificationManager).setNotificationPolicy(captor.capture());
+        }
 
         Policy expected = new Policy(PRIORITY_CATEGORY_CONVERSATIONS
                 | PRIORITY_CATEGORY_CALLS | PRIORITY_CATEGORY_MESSAGES | PRIORITY_CATEGORY_ALARMS,
@@ -240,7 +245,6 @@
         assertEquals(expected, captor.getValue());
     }
 
-    @Ignore
     @Test
     public void saveSenders_doesNotChangeConversations() {
         when(mNotificationManager.getNotificationPolicy()).thenReturn(
@@ -255,7 +259,11 @@
         mBackend.saveSenders(PRIORITY_CATEGORY_CALLS, PRIORITY_SENDERS_ANY);
 
         ArgumentCaptor<Policy> captor = ArgumentCaptor.forClass(Policy.class);
-        verify(mNotificationManager, times(1)).setNotificationPolicy(captor.capture());
+        if (android.app.Flags.modesApi()) {
+            verify(mNotificationManager).setNotificationPolicy(captor.capture(), eq(true));
+        } else {
+            verify(mNotificationManager).setNotificationPolicy(captor.capture());
+        }
 
         Policy expected = new Policy(PRIORITY_CATEGORY_CONVERSATIONS
                 | PRIORITY_CATEGORY_CALLS | PRIORITY_CATEGORY_MESSAGES | PRIORITY_CATEGORY_ALARMS,
diff --git a/tests/robotests/src/com/android/settings/notification/zen/ZenModeSliceBuilderTest.java b/tests/robotests/src/com/android/settings/notification/zen/ZenModeSliceBuilderTest.java
index 8b76004..3f9e486 100644
--- a/tests/robotests/src/com/android/settings/notification/zen/ZenModeSliceBuilderTest.java
+++ b/tests/robotests/src/com/android/settings/notification/zen/ZenModeSliceBuilderTest.java
@@ -17,13 +17,18 @@
 package com.android.settings.notification.zen;
 
 import static android.app.slice.Slice.EXTRA_TOGGLE_STATE;
+import static android.provider.Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS;
+import static android.provider.Settings.Global.ZEN_MODE_OFF;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.verify;
+
 import android.app.NotificationManager;
 import android.content.Context;
 import android.content.Intent;
-import android.provider.Settings;
 
 import androidx.slice.Slice;
 import androidx.slice.SliceMetadata;
@@ -37,12 +42,14 @@
 
 import org.junit.After;
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
+import org.robolectric.shadows.ShadowApplication;
 
 import java.util.List;
 
@@ -52,10 +59,17 @@
 
     private Context mContext;
 
+    @Mock
+    private NotificationManager mNm;
+
     @Before
     public void setUp() {
         mContext = RuntimeEnvironment.application;
 
+        MockitoAnnotations.initMocks(this);
+        ShadowApplication shadowApplication = ShadowApplication.getInstance();
+        shadowApplication.setSystemService(Context.NOTIFICATION_SERVICE, mNm);
+
         // Set-up specs for SliceMetadata.
         SliceProvider.setSpecs(SliceLiveData.SUPPORTED_SPECS);
     }
@@ -96,30 +110,31 @@
         assertThat(primaryAction.getIcon()).isNull();
     }
 
-    @Ignore
     @Test
     public void handleUriChange_turnOn_zenModeTurnsOn() {
         final Intent intent = new Intent();
         intent.putExtra(EXTRA_TOGGLE_STATE, true);
-        NotificationManager.from(mContext).setZenMode(Settings.Global.ZEN_MODE_OFF, null, "");
 
         ZenModeSliceBuilder.handleUriChange(mContext, intent);
 
-        final int zenMode = NotificationManager.from(mContext).getZenMode();
-        assertThat(zenMode).isEqualTo(Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS);
+        if (android.app.Flags.modesApi()) {
+            verify(mNm).setZenMode(eq(ZEN_MODE_IMPORTANT_INTERRUPTIONS), any(), any(), eq(true));
+        } else {
+            verify(mNm).setZenMode(eq(ZEN_MODE_IMPORTANT_INTERRUPTIONS), any(), any());
+        }
     }
 
-    @Ignore
     @Test
     public void handleUriChange_turnOff_zenModeTurnsOff() {
         final Intent intent = new Intent();
         intent.putExtra(EXTRA_TOGGLE_STATE, false);
-        NotificationManager.from(mContext).setZenMode(
-                Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS, null, "");
 
         ZenModeSliceBuilder.handleUriChange(mContext, intent);
 
-        final int zenMode = NotificationManager.from(mContext).getZenMode();
-        assertThat(zenMode).isEqualTo(Settings.Global.ZEN_MODE_OFF);
+        if (android.app.Flags.modesApi()) {
+            verify(mNm).setZenMode(eq(ZEN_MODE_OFF), any(), any(), eq(true));
+        } else {
+            verify(mNm).setZenMode(eq(ZEN_MODE_OFF), any(), any());
+        }
     }
 }
diff --git a/tests/robotests/src/com/android/settings/notification/zen/ZenOnboardingActivityTest.java b/tests/robotests/src/com/android/settings/notification/zen/ZenOnboardingActivityTest.java
index 6063efb..01360fd 100644
--- a/tests/robotests/src/com/android/settings/notification/zen/ZenOnboardingActivityTest.java
+++ b/tests/robotests/src/com/android/settings/notification/zen/ZenOnboardingActivityTest.java
@@ -27,10 +27,13 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.app.Flags;
 import android.app.NotificationManager;
 import android.app.NotificationManager.Policy;
 import android.content.Context;
@@ -42,7 +45,6 @@
 import com.android.settings.testutils.FakeFeatureFactory;
 
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -91,7 +93,6 @@
         verify(mMetricsLogger).visible(MetricsEvent.SETTINGS_ZEN_ONBOARDING);
     }
 
-    @Ignore
     @Test
     public void saveNewSetting() {
         Policy policy = new Policy(PRIORITY_CATEGORY_ALARMS, 0, 0, SUPPRESSED_EFFECT_SCREEN_ON);
@@ -103,7 +104,11 @@
         verify(mMetricsLogger).action(MetricsEvent.ACTION_ZEN_ONBOARDING_OK);
 
         ArgumentCaptor<Policy> captor = ArgumentCaptor.forClass(Policy.class);
-        verify(mNm).setNotificationPolicy(captor.capture());
+        if (android.app.Flags.modesApi()) {
+            verify(mNm).setNotificationPolicy(captor.capture(), eq(true));
+        } else {
+            verify(mNm).setNotificationPolicy(captor.capture());
+        }
 
         Policy actual = captor.getValue();
         assertThat(actual.priorityCategories).isEqualTo(PRIORITY_CATEGORY_ALARMS
@@ -123,7 +128,11 @@
         mActivity.save(null);
 
         verify(mMetricsLogger).action(MetricsEvent.ACTION_ZEN_ONBOARDING_KEEP_CURRENT_SETTINGS);
-        verify(mNm, never()).setNotificationPolicy(any());
+        if (Flags.modesApi()) {
+            verify(mNm, never()).setNotificationPolicy(any(), anyBoolean());
+        } else {
+            verify(mNm, never()).setNotificationPolicy(any());
+        }
     }
 
     @Test