From c1069eec5c7d0e4c8f0f18f9a4c9ca1aae92de95 Mon Sep 17 00:00:00 2001
From: tyiuhc <tim.yiu@amplitude.com>
Date: Mon, 13 Nov 2023 12:00:06 -0800
Subject: [PATCH] add unit test and JSONArray case

---
 .../com/amplitude/experiment/util/Variant.kt  |  4 ++-
 .../experiment/ExperimentClientTest.kt        | 32 +++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/sdk/src/main/java/com/amplitude/experiment/util/Variant.kt b/sdk/src/main/java/com/amplitude/experiment/util/Variant.kt
index f99eeb8..d1def9b 100644
--- a/sdk/src/main/java/com/amplitude/experiment/util/Variant.kt
+++ b/sdk/src/main/java/com/amplitude/experiment/util/Variant.kt
@@ -98,7 +98,9 @@ internal fun EvaluationVariant.convertToVariant(): Variant {
     val payload = when {
         this.payload != null -> {
             if (this.payload is Map<*, *>) {
-                JSONObject(this.payload)
+                this.payload.toJSONObject()
+            } else if (this.payload is Collection<*>) {
+                this.payload.toJSONArray()
             } else {
                 this.payload
             }
diff --git a/sdk/src/test/java/com/amplitude/experiment/ExperimentClientTest.kt b/sdk/src/test/java/com/amplitude/experiment/ExperimentClientTest.kt
index 8f2648b..64bd316 100644
--- a/sdk/src/test/java/com/amplitude/experiment/ExperimentClientTest.kt
+++ b/sdk/src/test/java/com/amplitude/experiment/ExperimentClientTest.kt
@@ -14,6 +14,8 @@ import io.mockk.spyk
 import io.mockk.verify
 import okhttp3.HttpUrl.Companion.toHttpUrl
 import okhttp3.OkHttpClient
+import org.json.JSONArray
+import org.json.JSONObject
 import org.junit.Assert
 import org.junit.Before
 import org.junit.Test
@@ -1120,4 +1122,34 @@ class ExperimentClientTest {
         spyClient.start(null).get()
         verify(exactly = 0) { spyClient.fetchInternal(any(), any(), any(), any()) }
     }
+
+    @Test
+    fun `LocalEvaluation - test payload format`() {
+        val user = ExperimentUser(userId = "test_user")
+        val exposureTrackingProvider = mockk<TestExposureTrackingProvider>()
+        every { exposureTrackingProvider.track(any()) } just Runs
+        val client = DefaultExperimentClient(
+            API_KEY,
+            ExperimentConfig(
+                debug = true,
+                exposureTrackingProvider = exposureTrackingProvider,
+                fetchOnStart = false,
+                source = Source.LOCAL_STORAGE,
+            ),
+            OkHttpClient(),
+            mockStorage,
+            Experiment.executorService,
+        )
+        client.start(user).get()
+        var variant = client.variant("sdk-payload-ci-test")
+        val obj = JSONObject().put("key1", "val1").put("key2", "val2")
+        val array = JSONArray().put(JSONObject().put("key1", "obj1")).put(JSONObject().put("key2", "obj2"))
+        Assert.assertEquals(obj::class, variant.payload!!::class)
+        Assert.assertEquals(obj.toString(), variant.payload.toString())
+        // set null user to get array variant
+        client.start(ExperimentUser()).get()
+        variant = client.variant("sdk-payload-ci-test")
+        Assert.assertEquals(array::class, variant.payload!!::class)
+        Assert.assertEquals(array.toString(), variant.payload.toString())
+    }
 }