From c5b4d2bf1fd002340300e427cedbc75113560fc8 Mon Sep 17 00:00:00 2001
From: Chris Leonavicius <chris.leonavicius@amplitude.com>
Date: Tue, 12 Nov 2024 15:02:33 -0800
Subject: [PATCH] fix: properly decode embedded arrays (#242)

---
 .../Utilities/CodableExtension.swift          | 28 ++++++++-----------
 .../Events/BaseEventTests.swift               | 21 ++++++++++++--
 2 files changed, 30 insertions(+), 19 deletions(-)

diff --git a/Sources/Amplitude/Utilities/CodableExtension.swift b/Sources/Amplitude/Utilities/CodableExtension.swift
index 5df58f4..d131df9 100644
--- a/Sources/Amplitude/Utilities/CodableExtension.swift
+++ b/Sources/Amplitude/Utilities/CodableExtension.swift
@@ -27,18 +27,10 @@ struct JSONCodingKeys: CodingKey {
 }
 
 extension KeyedDecodingContainer {
-    func decode(_ type: [String: Any].Type, forKey key: K) throws -> [String: Any] {
-        let container = try self.nestedContainer(keyedBy: JSONCodingKeys.self, forKey: key)
-        return try container.decode(type)
-    }
 
-    func decode(_ type: [[String: Any]].Type, forKey key: K) throws -> [[String: Any]] {
-        var container = try self.nestedUnkeyedContainer(forKey: key)
-        if let decodedData = try container.decode([Any].self) as? [[String: Any]] {
-            return decodedData
-        } else {
-            return []
-        }
+    func decode(_ type: [String: Any].Type, forKey key: K) throws -> [String: Any] {
+        let container = try nestedContainer(keyedBy: JSONCodingKeys.self, forKey: key)
+        return try container.asTypedDictionary(type)
     }
 
     func decodeIfPresent(_ type: [String: Any].Type, forKey key: K) throws -> [String: Any]? {
@@ -52,8 +44,8 @@ extension KeyedDecodingContainer {
     }
 
     func decode(_ type: [Any].Type, forKey key: K) throws -> [Any] {
-        var container = try self.nestedUnkeyedContainer(forKey: key)
-        return try container.decode(type)
+        var container = try nestedUnkeyedContainer(forKey: key)
+        return try container.asTypedArray(type)
     }
 
     func decodeIfPresent(_ type: [Any].Type, forKey key: K) throws -> [Any]? {
@@ -66,7 +58,7 @@ extension KeyedDecodingContainer {
         return try decode(type, forKey: key)
     }
 
-    func decode(_ type: [String: Any].Type) throws -> [String: Any] {
+    func asTypedDictionary(_ type: [String: Any].Type) throws -> [String: Any] {
         var dictionary = [String: Any]()
         for key in allKeys {
             if let boolValue = try? decode(Bool.self, forKey: key) {
@@ -88,7 +80,8 @@ extension KeyedDecodingContainer {
 }
 
 extension UnkeyedDecodingContainer {
-    mutating func decode(_ type: [Any].Type) throws -> [Any] {
+
+    mutating func asTypedArray(_ type: [Any].Type) throws -> [Any] {
         var array: [Any] = []
         while isAtEnd == false {
             // See if the current value in the JSON array is `null` first
@@ -103,7 +96,8 @@ extension UnkeyedDecodingContainer {
                 array.append(value)
             } else if let nestedDictionary = try? decode(Dictionary<String, Any>.self) {
                 array.append(nestedDictionary)
-            } else if let nestedArray = try? decode(Array<Any>.self) {
+            } else if var nestedUnkeyedContainer = try? nestedUnkeyedContainer(),
+                      let nestedArray = try? nestedUnkeyedContainer.asTypedArray([Any].self) {
                 array.append(nestedArray)
             }
         }
@@ -112,7 +106,7 @@ extension UnkeyedDecodingContainer {
 
     mutating func decode(_ type: [String: Any].Type) throws -> [String: Any] {
         let nestedContainer = try self.nestedContainer(keyedBy: JSONCodingKeys.self)
-        return try nestedContainer.decode(type)
+        return try nestedContainer.asTypedDictionary(type)
     }
 }
 
diff --git a/Tests/AmplitudeTests/Events/BaseEventTests.swift b/Tests/AmplitudeTests/Events/BaseEventTests.swift
index 2d012b0..3fcc289 100644
--- a/Tests/AmplitudeTests/Events/BaseEventTests.swift
+++ b/Tests/AmplitudeTests/Events/BaseEventTests.swift
@@ -27,7 +27,9 @@ final class BaseEventTests: XCTestCase {
             "dict": ["a": 1, "b": 2, "c": "d"],
             "embeddedArray": ["a": [1, 2, 3]],
             "embeddedDict": [1, 2, ["a": 3]],
-            "array2": ["a", 1, 2, "b"]
+            "array2": ["a", 1, 2, "b"],
+            "nestedarray": [[["a": ["b": 1]]]],
+            "nesteddictionary": ["a": ["b": [[1]]]],
         ]
         let baseEvent = BaseEvent(
             plan: Plan(
@@ -111,6 +113,10 @@ final class BaseEventTests: XCTestCase {
             baseEventDict!["ingestion_metadata"]!["source_version" as NSString] as! String,
             "test-source-version"
         )
+        XCTAssertEqual((baseEventDict!["event_properties"]!["nestedarray" as NSString] as! NSArray),
+                       [[["a": ["b": 1]]]])
+        XCTAssertEqual((baseEventDict!["event_properties"]!["nesteddictionary" as NSString] as! NSDictionary),
+                       ["a": ["b": [[1]]]])
     }
 
     func testToString_withNilValues() {
@@ -159,7 +165,9 @@ final class BaseEventTests: XCTestCase {
                 "event_properties": {
                     "integer": 1,
                     "string": "stringValue",
-                    "array": [1, 2, 3]
+                    "array": [1, 2, 3],
+                    "nestedarray": [[{"a": {"b": 1}}]],
+                    "nesteddictionary": {"a": {"b": [[1]]}},
                 },
                 "plan": {
                     "branch": "test-branch",
@@ -191,6 +199,15 @@ final class BaseEventTests: XCTestCase {
             event?.eventProperties!["array"] as! [Double],
             [1, 2, 3]
         )
+        XCTAssertEqual(
+            event?.eventProperties!["nestedarray"] as! NSArray,
+            [[["a": ["b": 1]]]]
+        )
+        XCTAssertEqual(
+            event?.eventProperties!["nesteddictionary"] as! NSDictionary,
+            ["a": ["b": [[1]]]]
+        )
+
         XCTAssertEqual(
             event?.plan?.branch,
             "test-branch"