Use cancellableDelay in Account and WireGuard keys controllers

2 files changed, 9 insertions, 9 deletions

diff --git a/ios/MullvadVPN/AccountViewController.swift b/ios/MullvadVPN/AccountViewController.swift
index 6456140cf0..58a35875cc 100644
--- a/ios/MullvadVPN/AccountViewController.swift
+++ b/ios/MullvadVPN/AccountViewController.swift
@@ -64,16 +64,18 @@ class AccountViewController: UIViewController {
     }
 
     @IBAction func copyAccountToken() {
-        UIPasteboard.general.string = Account.shared.token
+        let accountToken = Account.shared.token
+
+        UIPasteboard.general.string = accountToken
 
         accountTokenButton.setTitle(
             NSLocalizedString("COPIED TO PASTEBOARD!", comment: ""),
             for: .normal)
 
         copyToPasteboardSubscriber =
-            Just(()).delay(for: .seconds(3), scheduler: DispatchQueue.main)
-                .sink(receiveValue: { _ in
-                    self.accountTokenButton.setTitle(Account.shared.token, for: .normal)
+            Just(accountToken).cancellableDelay(for: .seconds(3), scheduler: DispatchQueue.main)
+                .sink(receiveValue: { [weak self] (accountToken) in
+                    self?.accountTokenButton.setTitle(accountToken, for: .normal)
                 })
     }
 
diff --git a/ios/MullvadVPN/WireguardKeysViewController.swift b/ios/MullvadVPN/WireguardKeysViewController.swift
index 881200c963..1bfc367f63 100644
--- a/ios/MullvadVPN/WireguardKeysViewController.swift
+++ b/ios/MullvadVPN/WireguardKeysViewController.swift
@@ -110,14 +110,12 @@ class WireguardKeysViewController: UIViewController {
             animated: true)
 
         copyToPasteboardSubscriber =
-            Just(()).delay(for: .seconds(3), scheduler: DispatchQueue.main)
-                .sink(receiveValue: { [weak self] _ in
-                    guard let self = self, let publicKey = self.publicKey else { return }
-
+            Just(publicKey).cancellableDelay(for: .seconds(3), scheduler: DispatchQueue.main)
+                .sink(receiveValue: { [weak self] (publicKey) in
                     let displayKey = publicKey
                         .stringRepresentation(maxLength: kDisplayPublicKeyMaxLength)
 
-                    self.setPublicKeyTitle(string: displayKey, animated: true)
+                    self?.setPublicKeyTitle(string: displayKey, animated: true)
                 })
     }