diff --git a/tests/src/com/android/launcher3/model/GridSizeMigrationUtilTest.kt b/tests/src/com/android/launcher3/model/GridSizeMigrationUtilTest.kt
index a85fa3a..ca269a9 100644
--- a/tests/src/com/android/launcher3/model/GridSizeMigrationUtilTest.kt
+++ b/tests/src/com/android/launcher3/model/GridSizeMigrationUtilTest.kt
@@ -33,6 +33,7 @@
 import com.android.launcher3.provider.LauncherDbUtils
 import com.android.launcher3.util.LauncherModelHelper
 import com.android.launcher3.util.LauncherModelHelper.*
+import com.android.launcher3.util.TestUtil
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
 import org.junit.Before
@@ -750,20 +751,14 @@
     }
 
     private fun enableNewMigrationLogic(srcGridSize: String) {
-        context
-            .getSharedPreferences(FeatureFlags.FLAGS_PREF_NAME, Context.MODE_PRIVATE)
-            .edit()
-            .putBoolean(FeatureFlags.ENABLE_NEW_MIGRATION_LOGIC.key, true)
-            .commit()
         LauncherPrefs.get(context).putSync(WORKSPACE_SIZE.to(srcGridSize))
+        TestUtil.overrideBooleanFlagValue(context,
+                FeatureFlags.ENABLE_NEW_MIGRATION_LOGIC, true);
         FeatureFlags.initialize(context)
     }
 
     private fun disableNewMigrationLogic() {
-        context
-            .getSharedPreferences(FeatureFlags.FLAGS_PREF_NAME, Context.MODE_PRIVATE)
-            .edit()
-            .putBoolean(FeatureFlags.ENABLE_NEW_MIGRATION_LOGIC.key, false)
-            .commit()
+        TestUtil.overrideBooleanFlagValue(context,
+                FeatureFlags.ENABLE_NEW_MIGRATION_LOGIC, false);
     }
 }
diff --git a/tests/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncherTest.java b/tests/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncherTest.java
index 082e243..33a7f5c 100644
--- a/tests/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncherTest.java
+++ b/tests/src/com/android/launcher3/secondarydisplay/SecondaryDisplayLauncherTest.java
@@ -38,6 +38,7 @@
 import com.android.launcher3.tapl.LauncherInstrumentation;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.util.LauncherModelHelper;
+import com.android.launcher3.util.TestUtil;
 
 import org.junit.After;
 import org.junit.Ignore;
@@ -278,10 +279,8 @@
 
     private void setDragNDropFlag(Boolean status) {
         Context context = new LauncherModelHelper().sandboxContext;
-        context.getSharedPreferences(FeatureFlags.FLAGS_PREF_NAME, Context.MODE_PRIVATE).edit()
-                .putBoolean(FeatureFlags.SECONDARY_DRAG_N_DROP_TO_PIN.key, status)
-                .commit();
-        FeatureFlags.initialize(context);
+        TestUtil.overrideBooleanFlagValue(
+                context, FeatureFlags.SECONDARY_DRAG_N_DROP_TO_PIN, status);
         startSecondaryDisplayActivity();
     }
 }
diff --git a/tests/src/com/android/launcher3/util/TestUtil.java b/tests/src/com/android/launcher3/util/TestUtil.java
index d7c6c4f..38de4c3 100644
--- a/tests/src/com/android/launcher3/util/TestUtil.java
+++ b/tests/src/com/android/launcher3/util/TestUtil.java
@@ -19,14 +19,18 @@
 import static androidx.test.InstrumentationRegistry.getInstrumentation;
 import static androidx.test.InstrumentationRegistry.getTargetContext;
 
+import android.content.Context;
 import android.content.pm.LauncherApps;
 import android.content.res.Resources;
 import android.os.Handler;
 import android.os.Looper;
 import android.os.UserHandle;
 
+import androidx.annotation.VisibleForTesting;
 import androidx.test.uiautomator.UiDevice;
 
+import com.android.launcher3.config.FeatureFlags;
+
 import org.junit.Assert;
 
 import java.io.FileOutputStream;
@@ -68,6 +72,18 @@
         }
     }
 
+    @VisibleForTesting
+    // Override feature flag, mainly to be used ONLY in tests
+    public static void overrideBooleanFlagValue(
+            Context context, FeatureFlags.BooleanFlag flagToOverride,
+            boolean bool) {
+        context.getSharedPreferences(FeatureFlags.FLAGS_PREF_NAME, Context.MODE_PRIVATE)
+                .edit()
+                .putBoolean(flagToOverride.key, bool)
+                .commit();
+        FeatureFlags.initialize(context);
+    }
+
     public static void uninstallDummyApp() throws IOException {
         UiDevice.getInstance(getInstrumentation()).executeShellCommand(
                 "pm uninstall " + DUMMY_PACKAGE);