From a70189d3180a9607dfd5715be3edfd08cb8b8a26 Mon Sep 17 00:00:00 2001
From: Doug <6060466+pixlwave@users.noreply.github.com>
Date: Thu, 23 Jan 2025 12:39:53 +0000
Subject: [PATCH] Make sure the Recovery Key option is shown on the
 IdentityConfirmationScreen when available. (#3699)

There was a race condition where 2 different states could be processed at the same time.

* Fix unit test
The view model wasn't stored, so relying the current value publisher didn't have time to fire before deinit.
---
 .../IdentityConfirmationScreenViewModel.swift     | 15 +++++----------
 .../View/IdentityConfirmationScreen.swift         |  4 +++-
 2 files changed, 8 insertions(+), 11 deletions(-)

diff --git a/ElementX/Sources/Screens/Onboarding/IdentityConfirmationScreen/IdentityConfirmationScreenViewModel.swift b/ElementX/Sources/Screens/Onboarding/IdentityConfirmationScreen/IdentityConfirmationScreenViewModel.swift
index 1ab5c5761..a903f6708 100644
--- a/ElementX/Sources/Screens/Onboarding/IdentityConfirmationScreen/IdentityConfirmationScreenViewModel.swift
+++ b/ElementX/Sources/Screens/Onboarding/IdentityConfirmationScreen/IdentityConfirmationScreenViewModel.swift
@@ -25,18 +25,13 @@ class IdentityConfirmationScreenViewModel: IdentityConfirmationScreenViewModelTy
         
         super.init(initialViewState: IdentityConfirmationScreenViewState(learnMoreURL: appSettings.encryptionURL))
         
-        userSession.sessionSecurityStatePublisher
-            .receive(on: DispatchQueue.main)
-            .sink { [weak self] state in
-                Task {
-                    await self?.updateWithSessionSecurityState(state)
-                }
+        Task { [weak self] in
+            for await state in userSession.sessionSecurityStatePublisher.values {
+                // We need to call this inside an AsyncSequence otherwise there's a race condition when the method suspends.
+                await self?.updateWithSessionSecurityState(state)
             }
-            .store(in: &cancellables)
-        
-        Task {
-            await updateWithSessionSecurityState(userSession.sessionSecurityStatePublisher.value)
         }
+        .store(in: &cancellables)
     }
     
     // MARK: - Public
diff --git a/ElementX/Sources/Screens/Onboarding/IdentityConfirmationScreen/View/IdentityConfirmationScreen.swift b/ElementX/Sources/Screens/Onboarding/IdentityConfirmationScreen/View/IdentityConfirmationScreen.swift
index a183adcbb..002dfef31 100644
--- a/ElementX/Sources/Screens/Onboarding/IdentityConfirmationScreen/View/IdentityConfirmationScreen.swift
+++ b/ElementX/Sources/Screens/Onboarding/IdentityConfirmationScreen/View/IdentityConfirmationScreen.swift
@@ -104,6 +104,8 @@ struct IdentityConfirmationScreen: View {
 // MARK: - Previews
 
 struct IdentityConfirmationScreen_Previews: PreviewProvider, TestablePreview {
+    static var viewModel = makeViewModel()
+    
     static var previews: some View {
         NavigationStack {
             IdentityConfirmationScreen(context: viewModel.context)
@@ -113,7 +115,7 @@ struct IdentityConfirmationScreen_Previews: PreviewProvider, TestablePreview {
         })
     }
     
-    private static var viewModel: IdentityConfirmationScreenViewModel {
+    static func makeViewModel() -> IdentityConfirmationScreenViewModel {
         let clientProxy = ClientProxyMock(.init())
         let userSession = UserSessionMock(.init(clientProxy: clientProxy))
         userSession.sessionSecurityStatePublisher = CurrentValuePublisher<SessionSecurityState, Never>(.init(verificationState: .unverified, recoveryState: .enabled))