diff --git a/tests/src/com/android/launcher3/util/rule/TestStabilityRule.java b/tests/src/com/android/launcher3/util/rule/TestStabilityRule.java index 08fa0980a2..858cb384d7 100644 --- a/tests/src/com/android/launcher3/util/rule/TestStabilityRule.java +++ b/tests/src/com/android/launcher3/util/rule/TestStabilityRule.java @@ -21,6 +21,7 @@ import android.content.pm.PackageManager; import android.os.Build; import android.util.Log; +import androidx.test.InstrumentationRegistry; import androidx.test.uiautomator.UiDevice; import org.junit.rules.TestRule; @@ -86,6 +87,19 @@ public class TestStabilityRule implements TestRule { } private static int getRunFlavor() { + final String flavorOverride = InstrumentationRegistry.getArguments().getString("flavor"); + + if (flavorOverride != null) { + Log.d(TAG, "Flavor override: " + flavorOverride); + try { + return (int) TestStabilityRule.class.getField(flavorOverride).get(null); + } catch (NoSuchFieldException e) { + throw new AssertionError("Unrecognized run flavor override: " + flavorOverride); + } catch (IllegalAccessException e) { + throw new RuntimeException(e); + } + } + final String launcherVersion; try { launcherVersion = getInstrumentation().