From 58be549d127f44cbb24f9f6dcd5a85cadbce9014 Mon Sep 17 00:00:00 2001
From: Tim Yiu <tim.yiu@Tims-MacBook-Pro.local>
Date: Fri, 22 Sep 2023 10:44:04 -0700
Subject: [PATCH] update String.toVariant payload handling

---
 .../com/amplitude/experiment/util/Variant.kt  | 17 +++++++----
 .../com/amplitude/experiment/StorageTest.kt   | 30 +++++++------------
 2 files changed, 21 insertions(+), 26 deletions(-)

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 98fe81d..491ed08 100644
--- a/sdk/src/main/java/com/amplitude/experiment/util/Variant.kt
+++ b/sdk/src/main/java/com/amplitude/experiment/util/Variant.kt
@@ -30,11 +30,7 @@ internal fun String?.toVariant(): Variant? {
     return if (this == null) {
         null
     } else {
-        try {
-            JSONObject(this).toVariant()
-        } catch (e: JSONException) {
-            Variant(key = this, value = this)
-        }
+        JSONObject(this).toVariant()
     }
 }
 
@@ -59,9 +55,18 @@ internal fun JSONObject?.toVariant(): Variant? {
         }
 
         val payload = when {
-            has("payload") -> get("payload")
+            has("payload") -> {
+                val payload = get("payload")
+                if (payload is JSONObject) {
+                    payload.toMap()
+                } else {
+                    payload
+                }
+            }
+
             else -> null
         }
+
         var expKey = when {
             has("expKey") -> getString("expKey")
             else -> null
diff --git a/sdk/src/test/java/com/amplitude/experiment/StorageTest.kt b/sdk/src/test/java/com/amplitude/experiment/StorageTest.kt
index aa845af..b5627ff 100644
--- a/sdk/src/test/java/com/amplitude/experiment/StorageTest.kt
+++ b/sdk/src/test/java/com/amplitude/experiment/StorageTest.kt
@@ -7,16 +7,6 @@ import org.junit.Test
 
 class TransformVariantFromStorageTest {
 
-    private val payload = "payload"
-
-    @Test
-    fun `v0 variant transformation`() {
-        Assert.assertEquals(
-            Variant(key = "on", value = "on"),
-            transformVariantFromStorage("on")
-        )
-    }
-
     @Test
     fun `v1 variant transformation`() {
         val storedVariant = JSONObject(mapOf("value" to "on")).toString()
@@ -31,11 +21,11 @@ class TransformVariantFromStorageTest {
         val storedVariant = JSONObject(
             mapOf(
                 "value" to "on",
-                "payload" to payload
+                "payload" to mapOf("k" to "v")
             )
         ).toString()
         Assert.assertEquals(
-            Variant(key = "on", value = "on", payload = payload),
+            Variant(key = "on", value = "on", payload = mapOf("k" to "v")),
             transformVariantFromStorage(storedVariant)
         )
     }
@@ -45,7 +35,7 @@ class TransformVariantFromStorageTest {
         val storedVariant = JSONObject(
             mapOf(
                 "value" to "on",
-                "payload" to payload,
+                "payload" to mapOf("k" to "v"),
                 "expKey" to "exp-1"
             )
         ).toString()
@@ -53,7 +43,7 @@ class TransformVariantFromStorageTest {
             Variant(
                 key = "on",
                 value = "on",
-                payload = payload,
+                payload = mapOf("k" to "v"),
                 expKey = "exp-1",
                 metadata = mapOf("experimentKey" to "exp-1")
             ),
@@ -81,14 +71,14 @@ class TransformVariantFromStorageTest {
             mapOf(
                 "key" to "treatment",
                 "value" to "on",
-                "payload" to payload
+                "payload" to mapOf("k" to "v")
             )
         ).toString()
         Assert.assertEquals(
             Variant(
                 key = "treatment",
                 value = "on",
-                payload = payload
+                payload = mapOf("k" to "v")
             ),
             transformVariantFromStorage(storedVariant)
         )
@@ -100,7 +90,7 @@ class TransformVariantFromStorageTest {
             mapOf(
                 "key" to "treatment",
                 "value" to "on",
-                "payload" to payload,
+                "payload" to mapOf("k" to "v"),
                 "expKey" to "exp-1"
             )
         ).toString()
@@ -108,7 +98,7 @@ class TransformVariantFromStorageTest {
             Variant(
                 key = "treatment",
                 value = "on",
-                payload = payload,
+                payload = mapOf("k" to "v"),
                 expKey = "exp-1",
                 metadata = mapOf("experimentKey" to "exp-1")
             ),
@@ -122,7 +112,7 @@ class TransformVariantFromStorageTest {
             mapOf(
                 "key" to "treatment",
                 "value" to "on",
-                "payload" to payload,
+                "payload" to mapOf("k" to "v"),
                 "metadata" to mapOf("experimentKey" to "exp-1")
             )
         ).toString()
@@ -130,7 +120,7 @@ class TransformVariantFromStorageTest {
             Variant(
                 key = "treatment",
                 value = "on",
-                payload = payload,
+                payload = mapOf("k" to "v"),
                 expKey = "exp-1",
                 metadata = mapOf("experimentKey" to "exp-1")
             ),