Add support for custom headers

Author: arno-stofberg

sdk/src/main/java/com/amplitude/experiment/DefaultExperimentClient.kt

@@ -373,12 +373,18 @@ internal class DefaultExperimentClient internal constructor(
             serverUrl.newBuilder()
                 .addPathSegments("sdk/v2/vardata")
                 .build()
+        val customRequestHeaders = config.customRequestHeaders?.invoke(user)
         val builder =
             Request.Builder()
                 .get()
                 .url(url)
                 .addHeader("Authorization", "Api-Key $apiKey")
                 .addHeader("X-Amp-Exp-User", userBase64)
+                .apply {
+                    customRequestHeaders?.forEach { (name, value) ->
+                        addHeader(name, value)
+                    }
+                }
         if (!options?.flagKeys.isNullOrEmpty()) {
             val flagKeysBase64 =
                 JSONArray(options?.flagKeys)

sdk/src/main/java/com/amplitude/experiment/ExperimentConfig.kt

@@ -58,6 +58,8 @@ class ExperimentConfig internal constructor(
     val analyticsProvider: ExperimentAnalyticsProvider? = Defaults.ANALYTICS_PROVIDER,
     @JvmField
     val exposureTrackingProvider: ExposureTrackingProvider? = Defaults.EXPOSURE_TRACKING_PROVIDER,
+    @JvmField
+    val customRequestHeaders: ((ExperimentUser) -> Map<String, String>)? = Defaults.CUSTOM_REQUEST_HEADERS
 ) {
     /**
      * Construct the default [ExperimentConfig].
@@ -163,6 +165,11 @@ class ExperimentConfig internal constructor(
          * null
          */
         val EXPOSURE_TRACKING_PROVIDER: ExposureTrackingProvider? = null
+
+        /**
+         * null
+         */
+        val CUSTOM_REQUEST_HEADERS: ((ExperimentUser) -> Map<String, String>)? = null
     }
 
     companion object {
@@ -192,6 +199,7 @@ class ExperimentConfig internal constructor(
         private var userProvider = Defaults.USER_PROVIDER
         private var analyticsProvider = Defaults.ANALYTICS_PROVIDER
         private var exposureTrackingProvider = Defaults.EXPOSURE_TRACKING_PROVIDER
+        private var customRequestHeaders = Defaults.CUSTOM_REQUEST_HEADERS
 
         fun debug(debug: Boolean) =
             apply {
@@ -289,6 +297,11 @@ class ExperimentConfig internal constructor(
                 this.exposureTrackingProvider = exposureTrackingProvider
             }
 
+        fun customRequestHeaders(customRequestHeaders: ((ExperimentUser) -> Map<String, String>)?) =
+            apply {
+                this.customRequestHeaders = customRequestHeaders
+            }
+
         fun build(): ExperimentConfig {
             return ExperimentConfig(
                 debug = debug,
@@ -335,5 +348,6 @@ class ExperimentConfig internal constructor(
             .userProvider(userProvider)
             .analyticsProvider(analyticsProvider)
             .exposureTrackingProvider(exposureTrackingProvider)
+            .customRequestHeaders(customRequestHeaders)
     }
 }

sdk/src/test/java/com/amplitude/experiment/ExperimentClientTest.kt

@@ -1355,4 +1355,30 @@ class ExperimentClientTest {
             )
         Assert.assertEquals(900000, client.flagConfigPollingIntervalMillis)
     }
+
+    @Test
+    fun `test config custom request headers added, http call includes headers`() {
+        val mockHttpClient = mockk<OkHttpClient>()
+        val testCountry = "South Africa"
+        val client =
+            DefaultExperimentClient(
+                API_KEY,
+                ExperimentConfig(
+                    customRequestHeaders = { user ->
+                        mapOf("country" to user.country!!)
+                    }
+                ),
+                mockHttpClient,
+                mockStorage,
+                Experiment.executorService,
+            )
+
+        client.fetch(ExperimentUser(country = testCountry))
+
+        verify {
+            mockHttpClient.newCall(match {
+                it.headers["country"] == testCountry
+            })
+        }
+    }
 }