From aaae1d0c24a84abb99762bebea0444b43c5034e0 Mon Sep 17 00:00:00 2001
From: Graeme Arthur <2030310+graeme@users.noreply.github.com>
Date: Fri, 20 Dec 2024 11:33:28 +0100
Subject: [PATCH] Make sure flow isnt ended before switch finishes

---
 .../Preferences/Model/SyncPreferences.swift   | 40 ++++++++++---------
 .../ViewModels/ManagementDialogModel.swift    |  1 -
 2 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/DuckDuckGo/Preferences/Model/SyncPreferences.swift b/DuckDuckGo/Preferences/Model/SyncPreferences.swift
index 2193c7da59..0ebdd3a67b 100644
--- a/DuckDuckGo/Preferences/Model/SyncPreferences.swift
+++ b/DuckDuckGo/Preferences/Model/SyncPreferences.swift
@@ -764,8 +764,10 @@ extension SyncPreferences: ManagementDialogModelDelegate {
             managementDialogModel.shouldShowSwitchAccountsMessage = true
             PixelKit.fire(SyncSwitchAccountPixelKitEvent.syncAskUserToSwitchAccount.withoutMacPrefix)
         } else {
-            switchAccounts(recoveryKey: recoveryKey)
-            managementDialogModel.endFlow()
+            Task { @MainActor in
+                await switchAccounts(recoveryKey: recoveryKey)
+                managementDialogModel.endFlow()
+            }
         }
         PixelKit.fire(DebugEvent(GeneralPixel.syncLoginExistingAccountError(error: SyncError.accountAlreadyExists)))
     }
@@ -775,27 +777,27 @@ extension SyncPreferences: ManagementDialogModelDelegate {
         guard let recoveryKey = try? SyncCode.decodeBase64String(recoveryCode).recovery else {
             return
         }
-        switchAccounts(recoveryKey: recoveryKey)
+        Task {
+            await switchAccounts(recoveryKey: recoveryKey)
+            managementDialogModel.endFlow()
+        }
     }
 
-    private func switchAccounts(recoveryKey: SyncCode.RecoveryKey) {
-        Task { [weak self] in
-            guard let self else { return }
-            do {
-                try await syncService.disconnect()
-            } catch {
-                PixelKit.fire(SyncSwitchAccountPixelKitEvent.syncUserSwitchedLogoutError.withoutMacPrefix)
-            }
+    private func switchAccounts(recoveryKey: SyncCode.RecoveryKey) async {
+        do {
+            try await syncService.disconnect()
+        } catch {
+            PixelKit.fire(SyncSwitchAccountPixelKitEvent.syncUserSwitchedLogoutError.withoutMacPrefix)
+        }
 
-            do {
-                let device = deviceInfo()
-                let registeredDevices = try await syncService.login(recoveryKey, deviceName: device.name, deviceType: device.type)
-                await mapDevices(registeredDevices)
-            } catch {
-                PixelKit.fire(SyncSwitchAccountPixelKitEvent.syncUserSwitchedLoginError.withoutMacPrefix)
-            }
-            PixelKit.fire(SyncSwitchAccountPixelKitEvent.syncUserSwitchedAccount.withoutMacPrefix)
+        do {
+            let device = deviceInfo()
+            let registeredDevices = try await syncService.login(recoveryKey, deviceName: device.name, deviceType: device.type)
+            await mapDevices(registeredDevices)
+        } catch {
+            PixelKit.fire(SyncSwitchAccountPixelKitEvent.syncUserSwitchedLoginError.withoutMacPrefix)
         }
+        PixelKit.fire(SyncSwitchAccountPixelKitEvent.syncUserSwitchedAccount.withoutMacPrefix)
     }
 
     func switchAccountsCancelled() {
diff --git a/LocalPackages/SyncUI/Sources/SyncUI/ViewModels/ManagementDialogModel.swift b/LocalPackages/SyncUI/Sources/SyncUI/ViewModels/ManagementDialogModel.swift
index cc3d602a79..92cb4e5574 100644
--- a/LocalPackages/SyncUI/Sources/SyncUI/ViewModels/ManagementDialogModel.swift
+++ b/LocalPackages/SyncUI/Sources/SyncUI/ViewModels/ManagementDialogModel.swift
@@ -68,7 +68,6 @@ public final class ManagementDialogModel: ObservableObject {
 
     public func userConfirmedSwitchAccounts(recoveryCode: String) {
         delegate?.userConfirmedSwitchAccounts(recoveryCode: recoveryCode)
-        doEndFlow()
     }
 
     private func doEndFlow() {