diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 669aaab..ed846ed 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.ui;
 
+import static android.platform.test.flag.junit.SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT;
+
 import static androidx.test.InstrumentationRegistry.getInstrumentation;
 
 import static com.android.launcher3.testing.shared.TestProtocol.ICON_MISSING;
@@ -204,13 +206,7 @@
     public ScreenRecordRule mScreenRecordRule = new ScreenRecordRule();
 
     @Rule
-    public SetFlagsRule mSetFlagsRule = getFlagsRule();
-
-    private SetFlagsRule getFlagsRule() {
-        SetFlagsRule flagsRule = new SetFlagsRule();
-        flagsRule.initAllFlagsToReleaseConfigDefault();
-        return flagsRule;
-    }
+    public SetFlagsRule mSetFlagsRule = new SetFlagsRule(DEVICE_DEFAULT);
 
     protected void clearPackageData(String pkg) throws IOException, InterruptedException {
         final CountDownLatch count = new CountDownLatch(2);
diff --git a/tests/src/com/android/launcher3/ui/BubbleTextViewTest.java b/tests/src/com/android/launcher3/ui/BubbleTextViewTest.java
index 0d15300..58af51c 100644
--- a/tests/src/com/android/launcher3/ui/BubbleTextViewTest.java
+++ b/tests/src/com/android/launcher3/ui/BubbleTextViewTest.java
@@ -63,7 +63,8 @@
  */
 public class BubbleTextViewTest {
 
-    @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+    @Rule public final SetFlagsRule mSetFlagsRule =
+            new SetFlagsRule(SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT);
     private static final StringMatcherUtility.StringMatcher
             MATCHER = StringMatcherUtility.StringMatcher.getInstance();
     private static final UserHandle WORK_HANDLE = new UserHandle(13);
@@ -99,7 +100,6 @@
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
-        mSetFlagsRule.initAllFlagsToReleaseConfigDefault();
         Utilities.enableRunningInTestHarnessForTests();
         mContext = new ActivityContextWrapper(getApplicationContext());
         mBubbleTextView = new BubbleTextView(mContext);