From 2ac76ff4ce88b4edf6d74bfbf8541185b59910a9 Mon Sep 17 00:00:00 2001
From: jhugman <jhugman@users.noreply.github.com>
Date: Thu, 10 Dec 2020 17:34:31 +0000
Subject: [PATCH] [fenix] Add withExperiment extension method to Nimbus
 (https://github.com/mozilla-mobile/fenix/pull/16926) r=christian, sebastian

* Fixes https://github.com/mozilla-mobile/fenix/issues/16925 Add withExperiment extension method to Nimbus
* Don't call Nimbus at all if not enabled by FeatureFlag
---
 .../main/java/org/mozilla/fenix/ext/Nimbus.kt | 42 +++++++++++++++++++
 .../java/org/mozilla/fenix/home/HomeMenu.kt   | 26 ++++++------
 2 files changed, 54 insertions(+), 14 deletions(-)
 create mode 100644 app/src/main/java/org/mozilla/fenix/ext/Nimbus.kt

diff --git a/app/src/main/java/org/mozilla/fenix/ext/Nimbus.kt b/app/src/main/java/org/mozilla/fenix/ext/Nimbus.kt
new file mode 100644
index 0000000000..e1ab7fa39c
--- /dev/null
+++ b/app/src/main/java/org/mozilla/fenix/ext/Nimbus.kt
@@ -0,0 +1,42 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+package org.mozilla.fenix.ext
+
+import mozilla.components.service.nimbus.NimbusApi
+import mozilla.components.support.base.log.logger.Logger
+import org.mozilla.fenix.FeatureFlags
+
+/**
+ * Gets the branch of the given `experimentId` and transforms it with given closure.
+ *
+ * If we're enrolled in the experiment, the transform is passed the branch id/slug as a `String`.
+ *
+ * If we're not enrolled in the experiment, or the experiment is not valid then the transform
+ * is passed a `null`.
+ */
+@Suppress("TooGenericExceptionCaught")
+fun <T> NimbusApi.withExperiment(experimentId: String, transform: (String?) -> T): T {
+    val branch = if (FeatureFlags.nimbusExperiments) {
+        try {
+            getExperimentBranch(experimentId)
+        } catch (e: Throwable) {
+            Logger.error("Failed to getExperimentBranch($experimentId)", e)
+            null
+        }
+    } else {
+        null
+    }
+    return transform(branch)
+}
+
+/**
+ * The degenerate case of `withExperiment(String, (String?) -> T))`, with an identity transform.
+ *
+ * Short-hand for `mozilla.components.service.nimbus.NimbusApi.getExperimentBranch`.
+ */
+fun NimbusApi.withExperiment(experimentId: String) =
+    this.withExperiment(experimentId, ::identity)
+
+private fun <T> identity(value: T) = value
diff --git a/app/src/main/java/org/mozilla/fenix/home/HomeMenu.kt b/app/src/main/java/org/mozilla/fenix/home/HomeMenu.kt
index b497958691..027e6cb266 100644
--- a/app/src/main/java/org/mozilla/fenix/home/HomeMenu.kt
+++ b/app/src/main/java/org/mozilla/fenix/home/HomeMenu.kt
@@ -26,6 +26,7 @@ import org.mozilla.fenix.experiments.ExperimentBranch
 import org.mozilla.fenix.experiments.Experiments
 import org.mozilla.fenix.ext.components
 import org.mozilla.fenix.ext.settings
+import org.mozilla.fenix.ext.withExperiment
 import org.mozilla.fenix.theme.ThemeManager
 import org.mozilla.fenix.whatsnew.WhatsNew
 
@@ -99,14 +100,12 @@ class HomeMenu(
         }
 
         val experiments = context.components.analytics.experiments
-        val bookmarksIcon = experiments.getExperimentBranch(Experiments.BOOKMARK_ICON)
-            .let {
-                when (it) {
-                    ExperimentBranch.TREATMENT -> R.drawable.ic_bookmark_list
-                    else -> R.drawable.ic_bookmark_filled
-                }
+        val bookmarksIcon = experiments.withExperiment(Experiments.BOOKMARK_ICON) {
+            when (it) {
+                ExperimentBranch.TREATMENT -> R.drawable.ic_bookmark_list
+                else -> R.drawable.ic_bookmark_filled
             }
-
+        }
         val bookmarksItem = BrowserMenuImageText(
             context.getString(R.string.library_bookmarks),
             bookmarksIcon,
@@ -122,14 +121,13 @@ class HomeMenu(
         // user isn't targeted, then we get still get the same treatment.
         // The `let` block is degenerate here, but left here so as to document the form of how experiments
         // are implemented here.
-        val historyIcon = experiments.getExperimentBranch(Experiments.A_A_NIMBUS_VALIDATION)
-            .let {
-                when (it) {
-                    ExperimentBranch.A1 -> R.drawable.ic_history
-                    ExperimentBranch.A2 -> R.drawable.ic_history
-                    else -> R.drawable.ic_history
-                }
+        val historyIcon = experiments.withExperiment(Experiments.A_A_NIMBUS_VALIDATION) {
+            when (it) {
+                ExperimentBranch.A1 -> R.drawable.ic_history
+                ExperimentBranch.A2 -> R.drawable.ic_history
+                else -> R.drawable.ic_history
             }
+        }
         val historyItem = BrowserMenuImageText(
             context.getString(R.string.library_history),
             historyIcon,