1 files changed, 23 insertions, 0 deletions
diff --git a/qa/qa/page/base.rb b/qa/qa/page/base.rb
index 4708063b2eb..69f58dcb8a5 100644
--- a/qa/qa/page/base.rb
+++ b/qa/qa/page/base.rb
@@ -443,6 +443,29 @@ module QA
         # Some checkboxes and radio buttons are hidden by their labels and cannot be clicked directly
         click_by_js ? page.execute_script("arguments[0].click();", box) : box.click
       end
+
+      def feature_flag_controlled_element(feature_flag, element_when_flag_enabled, element_when_flag_disabled)
+        # Feature flags can change the UI elements shown, but we need admin access to get feature flag values, which
+        # prevents us running the tests on production. Instead we detect the UI element that should be shown when the
+        # feature flag is enabled and otherwise use the element that should be displayed when the feature flag is
+        # disabled.
+
+        # Check both options once quickly so that the test doesn't wait unnecessarily if the UI has loaded
+        # We wait for requests first and wait one second for the element because it can take a moment for a Vue app to
+        # load and render the UI
+        wait_for_requests
+
+        return element_when_flag_enabled if has_element?(element_when_flag_enabled, wait: 1)
+        return element_when_flag_disabled if has_element?(element_when_flag_disabled, wait: 1)
+
+        # Check both options again, this time waiting for the default duration
+        return element_when_flag_enabled if has_element?(element_when_flag_enabled)
+        return element_when_flag_disabled if has_element?(element_when_flag_disabled)
+
+        raise ElementNotFound,
+          "Could not find the expected element as #{element_when_flag_enabled} or #{element_when_flag_disabled}." \
+          "The relevant feature flag is #{feature_flag}"
+      end
     end
   end
 end