diff --git a/tests/Android.bp b/tests/Android.bp
index 5effa5c..d40efce 100644
--- a/tests/Android.bp
+++ b/tests/Android.bp
@@ -88,7 +88,7 @@
         "truth-prebuilt",
         "platform-test-rules",
         "testables",
-        "com_android_launcher3_flags_lib_debug",
+        "com_android_launcher3_flags_lib",
     ],
     manifest: "AndroidManifest-common.xml",
     platform_apis: true,
@@ -107,7 +107,7 @@
     ],
     static_libs: [
         "Launcher3TestLib",
-        "com_android_launcher3_flags_lib_debug",
+        "com_android_launcher3_flags_lib",
     ],
     libs: [
         "android.test.base",
diff --git a/tests/src/com/android/launcher3/ui/BubbleTextViewTest.java b/tests/src/com/android/launcher3/ui/BubbleTextViewTest.java
index ba17fdc..bba8c89 100644
--- a/tests/src/com/android/launcher3/ui/BubbleTextViewTest.java
+++ b/tests/src/com/android/launcher3/ui/BubbleTextViewTest.java
@@ -30,16 +30,21 @@
 import android.view.ViewGroup;
 
 import com.android.launcher3.BubbleTextView;
+import com.android.launcher3.Flags;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.model.data.ItemInfoWithIcon;
 import com.android.launcher3.search.StringMatcherUtility;
 import com.android.launcher3.util.ActivityContextWrapper;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.TestUtil;
+import com.android.launcher3.util.rule.StaticMockitoRule;
 import com.android.launcher3.views.BaseDragLayer;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
 
 /**
  * Unit tests for testing modifyTitleToSupportMultiLine() in BubbleTextView.java
@@ -50,6 +55,7 @@
  */
 public class BubbleTextViewTest {
 
+    @Rule public StaticMockitoRule mockitoRule = new StaticMockitoRule(Flags.class);
     private static final StringMatcherUtility.StringMatcher
             MATCHER = StringMatcherUtility.StringMatcher.getInstance();
     private static final int ONE_LINE = 1;
@@ -77,6 +83,8 @@
 
     @Before
     public void setUp() throws Exception {
+        MockitoAnnotations.initMocks(this);
+        Mockito.when(Flags.enableTwolineAllapps()).thenReturn(false);
         Utilities.enableRunningInTestHarnessForTests();
         mContext = new ActivityContextWrapper(getApplicationContext());
         mBubbleTextView = new BubbleTextView(mContext);
diff --git a/tests/src/com/android/launcher3/ui/WorkProfileTest.java b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
index ec62058..5b9adcd 100644
--- a/tests/src/com/android/launcher3/ui/WorkProfileTest.java
+++ b/tests/src/com/android/launcher3/ui/WorkProfileTest.java
@@ -15,7 +15,6 @@
  */
 package com.android.launcher3.ui;
 
-import static com.android.launcher3.Flags.FLAG_ENABLE_EXPANDING_PAUSE_WORK_BUTTON;
 import static com.android.launcher3.LauncherState.ALL_APPS;
 import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.allapps.AllAppsStore.DEFER_UPDATES_TEST;
@@ -27,7 +26,6 @@
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assume.assumeTrue;
 
-import android.platform.test.flag.junit.SetFlagsRule;
 import android.util.Log;
 import android.view.View;
 
@@ -46,7 +44,6 @@
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Ignore;
-import org.junit.Rule;
 import org.junit.Test;
 
 import java.util.Objects;
@@ -55,7 +52,6 @@
 public class WorkProfileTest extends AbstractLauncherUiTest {
 
     private static final int WORK_PAGE = ActivityAllAppsContainerView.AdapterHolder.WORK;
-    @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
 
     private int mProfileUserId;
     private boolean mWorkProfileSetupSuccessful;
@@ -64,7 +60,6 @@
     @Before
     @Override
     public void setUp() throws Exception {
-        mSetFlagsRule.disableFlags(FLAG_ENABLE_EXPANDING_PAUSE_WORK_BUTTON);
         super.setUp();
         String output =
                 mDevice.executeShellCommand(