From b7eedb45795452bfe7f05cebcee78ffb65ce71ec Mon Sep 17 00:00:00 2001
From: vadimt <vadimt@google.com>
Date: Fri, 25 Oct 2019 12:32:38 -0700
Subject: [PATCH] Support for overriding test run flavor

Change-Id: I054a1eaa888cd97a7c7fa5fcc65513d6d89c3a1c
---
 .../launcher3/util/rule/TestStabilityRule.java     | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

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().