From 5f3201faa72867a9467d892cab7b8b01192826dc Mon Sep 17 00:00:00 2001
From: Oleg Koretsky <oko@concordium.com>
Date: Tue, 4 Jun 2024 18:20:36 +0300
Subject: [PATCH] Run account update when unshielding completes

---
 .../more/unshielding/UnshieldingViewModel.kt  | 56 +++++++++----------
 1 file changed, 27 insertions(+), 29 deletions(-)

diff --git a/app/src/main/java/com/concordium/wallet/ui/more/unshielding/UnshieldingViewModel.kt b/app/src/main/java/com/concordium/wallet/ui/more/unshielding/UnshieldingViewModel.kt
index d4f38770..2d038810 100644
--- a/app/src/main/java/com/concordium/wallet/ui/more/unshielding/UnshieldingViewModel.kt
+++ b/app/src/main/java/com/concordium/wallet/ui/more/unshielding/UnshieldingViewModel.kt
@@ -27,6 +27,7 @@ import com.concordium.wallet.data.room.Account
 import com.concordium.wallet.data.room.Transfer
 import com.concordium.wallet.data.room.WalletDatabase
 import com.concordium.wallet.ui.account.common.accountupdater.AccountUpdater
+import com.concordium.wallet.ui.account.common.accountupdater.TotalBalancesData
 import com.concordium.wallet.ui.common.BackendErrorHandler
 import com.concordium.wallet.ui.transaction.sendfunds.SendFundsViewModel
 import com.concordium.wallet.util.DateTimeUtil
@@ -243,9 +244,6 @@ class UnshieldingViewModel(application: Application) : AndroidViewModel(applicat
             amount = amount,
             submissionId = submissionId
         )
-
-        _isUnshieldEnabledLiveData.postValue(true)
-        _waitingLiveData.postValue(false)
     }
 
     private suspend fun submitTransfer(
@@ -266,26 +264,13 @@ class UnshieldingViewModel(application: Application) : AndroidViewModel(applicat
         amount: BigInteger,
         submissionId: String,
     ) {
-        var newSelfEncryptedAmount: String? = null
-
-        if (createTransferOutput.remaining != null) {
-            newSelfEncryptedAmount = createTransferOutput.remaining
-            val remainingAmount =
-                accountUpdater.decryptAndSaveAmount(
-                    credentialsOutput.encryptionSecretKey,
-                    createTransferOutput.remaining
-                )
+        val newSelfEncryptedAmount: String? = createTransferOutput.remaining
 
-            account.finalizedEncryptedBalance?.let { encBalance ->
-                val oldDecryptedAmount =
-                    accountUpdater.lookupMappedAmount(encBalance.selfAmount)
-                oldDecryptedAmount?.let {
-                    accountUpdater.saveDecryptedAmount(
-                        createTransferOutput.remaining,
-                        remainingAmount.toString()
-                    )
-                }
-            }
+        if (newSelfEncryptedAmount != null) {
+            accountUpdater.decryptAndSaveAmount(
+                credentialsOutput.encryptionSecretKey,
+                newSelfEncryptedAmount
+            )
         }
 
         val newTransfer = Transfer(
@@ -308,14 +293,27 @@ class UnshieldingViewModel(application: Application) : AndroidViewModel(applicat
         )
         transferRepository.insert(newTransfer)
 
-        _finishWithResultLiveData.postValue(
-            Event(
-                UnshieldingResult(
-                    unshieldedAmount = amount,
-                    accountAddress = account.address,
+        accountUpdater.setUpdateListener(object : AccountUpdater.UpdateListener {
+            override fun onError(stringRes: Int) {
+                _errorLiveData.postValue(Event(stringRes))
+                _isUnshieldEnabledLiveData.postValue(true)
+                _waitingLiveData.postValue(false)
+            }
+
+            override fun onDone(totalBalances: TotalBalancesData) {
+                _finishWithResultLiveData.postValue(
+                    Event(
+                        UnshieldingResult(
+                            unshieldedAmount = amount,
+                            accountAddress = account.address,
+                        )
+                    )
                 )
-            )
-        )
+            }
+
+            override fun onNewAccountFinalized(accountName: String) {}
+        })
+        accountUpdater.updateForAccount(account)
     }
 
     private fun getBalanceAtDisposal(): BigInteger =