From c889b609ad4937405dcbf27e8e60c3c7f9c01996 Mon Sep 17 00:00:00 2001
From: Andrey Sokolov <andrey.sokolov@amplitude.com>
Date: Wed, 19 Jul 2023 20:49:09 +0400
Subject: [PATCH] fix: changes after merge

---
 .../Migration/RemnantDataMigration.swift      | 35 ++++++++++---------
 Tests/AmplitudeTests/AmplitudeTests.swift     | 10 ++----
 .../Migration/RemnantDataMigrationTests.swift |  4 +--
 3 files changed, 24 insertions(+), 25 deletions(-)

diff --git a/Sources/Amplitude/Migration/RemnantDataMigration.swift b/Sources/Amplitude/Migration/RemnantDataMigration.swift
index 67038dd9..f2aa0945 100644
--- a/Sources/Amplitude/Migration/RemnantDataMigration.swift
+++ b/Sources/Amplitude/Migration/RemnantDataMigration.swift
@@ -31,40 +31,43 @@ class RemnantDataMigration {
         }
 
         if maxEventId > 0 {
-            let currentLastEventId = amplitude.sessions.lastEventId
-            if currentLastEventId <= 0 {
-                amplitude.sessions.lastEventId = maxEventId
+            let currentLastEventId: Int64? = amplitude.storage.read(key: StorageKey.LAST_EVENT_ID)
+            if currentLastEventId == nil || currentLastEventId! <= 0 {
+                try? amplitude.storage.write(key: StorageKey.LAST_EVENT_ID, value: maxEventId)
             }
         }
     }
 
     private func moveDeviceAndUserId() {
-        let deviceId = storage.getValue(RemnantDataMigration.DEVICE_ID_KEY)
-        let userId = storage.getValue(RemnantDataMigration.USER_ID_KEY)
-
-        if deviceId != nil {
-            amplitude.setDeviceId(deviceId: deviceId)
+        let currentDeviceId: String? = amplitude.storage.read(key: StorageKey.DEVICE_ID)
+        if currentDeviceId == nil || currentDeviceId! == "" {
+            if let deviceId = storage.getValue(RemnantDataMigration.DEVICE_ID_KEY) {
+                try? amplitude.storage.write(key: StorageKey.DEVICE_ID, value: deviceId)
+            }
         }
 
-        if userId != nil {
-            amplitude.setUserId(userId: userId)
+        let currentUserId: String? = amplitude.storage.read(key: StorageKey.USER_ID)
+        if currentUserId == nil || currentUserId == "" {
+            if let userId = storage.getValue(RemnantDataMigration.USER_ID_KEY) {
+                try? amplitude.storage.write(key: StorageKey.USER_ID, value: userId)
+            }
         }
     }
 
     private func moveSessionData() {
-        let currentSessionId = amplitude.sessions.sessionId
-        let currentLastEventTime = amplitude.sessions.lastEventTime
+        let currentSessionId: Int64? = amplitude.storage.read(key: StorageKey.PREVIOUS_SESSION_ID)
+        let currentLastEventTime: Int64? = amplitude.storage.read(key: StorageKey.LAST_EVENT_TIME)
 
         let previousSessionId = storage.getLongValue(RemnantDataMigration.PREVIOUS_SESSION_ID_KEY)
         let lastEventTime = storage.getLongValue(RemnantDataMigration.PREVIOUS_SESSION_TIME_KEY)
 
-        if currentSessionId < 0 && previousSessionId != nil && previousSessionId! >= 0 {
-            amplitude.sessions.sessionId = previousSessionId!
+        if (currentSessionId == nil || currentSessionId! < 0) && previousSessionId != nil && previousSessionId! >= 0 {
+            try? amplitude.storage.write(key: StorageKey.PREVIOUS_SESSION_ID, value: previousSessionId)
             storage.removeLongValue(RemnantDataMigration.PREVIOUS_SESSION_ID_KEY)
         }
 
-        if currentLastEventTime < 0 && lastEventTime != nil && lastEventTime! >= 0 {
-            amplitude.sessions.lastEventTime = lastEventTime!
+        if (currentLastEventTime == nil || currentLastEventTime! < 0) && lastEventTime != nil && lastEventTime! >= 0 {
+            try? amplitude.storage.write(key: StorageKey.LAST_EVENT_TIME, value: lastEventTime)
             storage.removeLongValue(RemnantDataMigration.PREVIOUS_SESSION_TIME_KEY)
         }
     }
diff --git a/Tests/AmplitudeTests/AmplitudeTests.swift b/Tests/AmplitudeTests/AmplitudeTests.swift
index bcd1cd65..d2b85d67 100644
--- a/Tests/AmplitudeTests/AmplitudeTests.swift
+++ b/Tests/AmplitudeTests/AmplitudeTests.swift
@@ -51,11 +51,7 @@ final class AmplitudeTests: XCTestCase {
         let amplitude = Amplitude(configuration: configurationWithFakeStorage)
         XCTAssertEqual(amplitude.getDeviceId() != nil, true)
         let deviceIdUuid = amplitude.getDeviceId()!
-        XCTAssertEqual(storage.haveBeenCalledWith, [
-            "read(key: device_id)",
-            "read(key: user_id)",
-            "write(key: \(StorageKey.DEVICE_ID.rawValue), \(deviceIdUuid))"
-        ])
+        XCTAssertEqual(storage.haveBeenCalledWith.last, "write(key: \(StorageKey.DEVICE_ID.rawValue), \(deviceIdUuid))")
     }
 
     func testContext() {
@@ -119,7 +115,7 @@ final class AmplitudeTests: XCTestCase {
 
         amplitude.setUserId(userId: "test-user")
         XCTAssertEqual(amplitude.getUserId(), "test-user")
-        XCTAssertEqual(storage.haveBeenCalledWith[3], "write(key: \(StorageKey.USER_ID.rawValue), test-user)")
+        XCTAssertEqual(storage.haveBeenCalledWith.last, "write(key: \(StorageKey.USER_ID.rawValue), test-user)")
     }
 
     func testSetDeviceId() {
@@ -129,7 +125,7 @@ final class AmplitudeTests: XCTestCase {
 
         amplitude.setDeviceId(deviceId: "test-device")
         XCTAssertEqual(amplitude.getDeviceId(), "test-device")
-        XCTAssertEqual(storage.haveBeenCalledWith[3], "write(key: \(StorageKey.DEVICE_ID.rawValue), test-device)")
+        XCTAssertEqual(storage.haveBeenCalledWith.last, "write(key: \(StorageKey.DEVICE_ID.rawValue), test-device)")
     }
 
     func testInterceptedIdentifyIsSentOnFlush() {
diff --git a/Tests/AmplitudeTests/Migration/RemnantDataMigrationTests.swift b/Tests/AmplitudeTests/Migration/RemnantDataMigrationTests.swift
index 717e1ad7..eb333e6a 100644
--- a/Tests/AmplitudeTests/Migration/RemnantDataMigrationTests.swift
+++ b/Tests/AmplitudeTests/Migration/RemnantDataMigrationTests.swift
@@ -44,8 +44,8 @@ final class RemnantDataMigrationTests: XCTestCase {
         let configuration = Configuration(
             apiKey: apiKey,
             instanceName: instanceName,
-            storageProvider: PersistentStorage(apiKey: apiKey, storagePrefix: "\(instanceName)-default"),
-            identifyStorageProvider: PersistentStorage(apiKey: apiKey, storagePrefix: "\(instanceName)-identify"),
+            storageProvider: PersistentStorage(storagePrefix: "\(instanceName)-default"),
+            identifyStorageProvider: PersistentStorage(storagePrefix: "\(instanceName)-identify"),
             migrateLegacyData: migrateLegacyData
         )
         let amplitude = Amplitude(configuration: configuration)