commit b138ec9fc6c5a61142c2dcbd34e2fea52b07bebd
author Sunny Goyal <sunnygoyal@google.com> Thu Mar 02 16:49:59 2023 -0800
committer Sunny Goyal <sunnygoyal@google.com> Fri Mar 03 09:29:42 2023 -0800
tree 8cb3e238d8ef9ae657197d0f560b9e86f3f59dfe
parent c6c33dd115e59618072514ccc055c90dc5f4f604

Adding support for overriding flag values in test

Bug: 270386012
Test: Verified on device
Change-Id: Ic1af9ca369c0cd2ec49d98875f7b97e90a0050ca

src/com/android/launcher3/config/FeatureFlags.java

diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index c98b60f..98b61d1 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -21,9 +21,14 @@
 
 import android.content.Context;
 
+import androidx.annotation.VisibleForTesting;
+
 import com.android.launcher3.BuildConfig;
 import com.android.launcher3.Utilities;
 
+import java.util.function.Predicate;
+import java.util.function.ToIntFunction;
+
 /**
  * Defines a set of flags used to control various launcher behaviors.
  *
@@ -33,6 +38,11 @@
 
     public static final String FLAGS_PREF_NAME = "featureFlags";
 
+    @VisibleForTesting
+    public static Predicate<BooleanFlag> sBooleanReader = f -> f.mCurrentValue;
+    @VisibleForTesting
+    public static ToIntFunction<IntFlag> sIntReader = f -> f.mCurrentValue;
+
     private FeatureFlags() { }
 
     public static boolean showFlagTogglerUi(Context context) {
@@ -385,7 +395,7 @@
         }
 
         public boolean get() {
-            return mCurrentValue;
+            return sBooleanReader.test(this);
         }
     }
 
@@ -401,7 +411,7 @@
         }
 
         public int get() {
-            return mCurrentValue;
+            return sIntReader.applyAsInt(this);
         }
     }
 }

tests/src/com/android/launcher3/util/TestUtil.java

diff --git a/tests/src/com/android/launcher3/util/TestUtil.java b/tests/src/com/android/launcher3/util/TestUtil.java
index d7c6c4f..433fd31 100644
--- a/tests/src/com/android/launcher3/util/TestUtil.java
+++ b/tests/src/com/android/launcher3/util/TestUtil.java
@@ -27,12 +27,18 @@
 
 import androidx.test.uiautomator.UiDevice;
 
+import com.android.launcher3.config.FeatureFlags;
+import com.android.launcher3.config.FeatureFlags.BooleanFlag;
+import com.android.launcher3.config.FeatureFlags.IntFlag;
+
 import org.junit.Assert;
 
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.concurrent.CountDownLatch;
+import java.util.function.Predicate;
+import java.util.function.ToIntFunction;
 
 public class TestUtil {
     public static final String DUMMY_PACKAGE = "com.example.android.aardwolf";
@@ -68,6 +74,36 @@
         }
     }
 
+    /**
+     * Utility class to override a boolean flag during test. Note that the returned SafeCloseable
+     * must be closed to restore the original state
+     */
+    public static SafeCloseable overrideFlag(BooleanFlag flag, boolean value) {
+        Predicate<BooleanFlag> originalProxy = FeatureFlags.sBooleanReader;
+        Predicate<BooleanFlag> testProxy = f -> f == flag ? value : originalProxy.test(f);
+        FeatureFlags.sBooleanReader = testProxy;
+        return () -> {
+            if (FeatureFlags.sBooleanReader == testProxy) {
+                FeatureFlags.sBooleanReader = originalProxy;
+            }
+        };
+    }
+
+    /**
+     * Utility class to override a int flag during test. Note that the returned SafeCloseable
+     * must be closed to restore the original state
+     */
+    public static SafeCloseable overrideFlag(IntFlag flag, int value) {
+        ToIntFunction<IntFlag> originalProxy = FeatureFlags.sIntReader;
+        ToIntFunction<IntFlag> testProxy = f -> f == flag ? value : originalProxy.applyAsInt(f);
+        FeatureFlags.sIntReader = testProxy;
+        return () -> {
+            if (FeatureFlags.sIntReader == testProxy) {
+                FeatureFlags.sIntReader = originalProxy;
+            }
+        };
+    }
+
     public static void uninstallDummyApp() throws IOException {
         UiDevice.getInstance(getInstrumentation()).executeShellCommand(
                 "pm uninstall " + DUMMY_PACKAGE);