Make token publisher to return optional #4392

AlphaWallet/ActiveWalletCoordinator.swift

@@ -236,7 +236,8 @@ class ActiveWalletCoordinator: NSObject, Coordinator, DappRequestHandlerDelegate
     private func setupWalletSessions() {
         var walletSessions: ServerDictionary<WalletSession> = .init()
         for each in config.enabledServers {
-            let tokenBalanceService = SingleChainTokenBalanceService(wallet: wallet, server: each, tokenBalanceProvider: walletBalanceService)
+            let etherToken = MultipleChainsTokensDataStore.functional.etherToken(forServer: each)
+            let tokenBalanceService = SingleChainTokenBalanceService(wallet: wallet, server: each, etherToken: etherToken, tokenBalanceProvider: walletBalanceService)
             let session = WalletSession(account: wallet, server: each, config: config, tokenBalanceService: tokenBalanceService)
 
             walletSessions[each] = session

AlphaWallet/Alerts/ViewControllers/EditPriceAlertViewController.swift

@@ -101,7 +101,7 @@ class EditPriceAlertViewController: UIViewController {
                 .sink { [weak self] viewModel in
                     guard let strongSelf = self else { return }
 
-                    strongSelf.viewModel.set(marketPrice: viewModel.ticker?.price_usd)
+                    strongSelf.viewModel.set(marketPrice: viewModel?.ticker?.price_usd)
                     strongSelf.configure(viewModel: strongSelf.viewModel)
                 }.store(in: &cancelable)
         case .erc875, .erc721, .erc721ForTickets, .erc1155:

AlphaWallet/Core/Coordinators/WalletBalance/WalletBalanceFetcher.swift

@@ -15,17 +15,23 @@ protocol WalletBalanceFetcherDelegate: AnyObject {
     func didAddToken(in fetcher: WalletBalanceFetcherType)
     func didUpdate(in fetcher: WalletBalanceFetcherType)
 }
+protocol WalletBalanceFetcherTypeTests {
+    func setBalanceTestsOnly(_ value: BigInt, forToken token: TokenObject)
+    func deleteTokenTestsOnly(token: TokenObject)
+    func addOrUpdateTokenTestsOnly(token: TokenObject)
+}
 
-protocol WalletBalanceFetcherType: AnyObject {
+protocol WalletBalanceFetcherType: AnyObject, WalletBalanceFetcherTypeTests {
     var tokenObjects: [Activity.AssignedToken] { get }
     var balance: WalletBalance { get }
     var walletBalance: AnyPublisher<WalletBalance, Never> { get }
 
-    func tokenBalancePublisher(_ addressAndRPCServer: AddressAndRPCServer) -> AnyPublisher<BalanceBaseViewModel, Never>
-    func tokenBalance(_ key: AddressAndRPCServer) -> BalanceBaseViewModel
+    func tokenBalancePublisher(_ addressAndRPCServer: AddressAndRPCServer) -> AnyPublisher<BalanceBaseViewModel?, Never>
+    func tokenBalance(_ key: AddressAndRPCServer) -> BalanceBaseViewModel?
     func start()
     func stop() 
     func refreshBalance(updatePolicy: PrivateBalanceFetcher.RefreshBalancePolicy, force: Bool) -> Promise<Void>
+    func update(servers: [RPCServer])
 }
 
 class WalletBalanceFetcher: NSObject, WalletBalanceFetcherType {
@@ -36,16 +42,14 @@ class WalletBalanceFetcher: NSObject, WalletBalanceFetcherType {
     private var balanceFetchers: ServerDictionary<PrivateBalanceFetcherType> = .init()
     private let queue: DispatchQueue
     private let coinTickersFetcher: CoinTickersFetcherType
-    private lazy var tokensDataStore: TokensDataStore = {
-        return MultipleChainsTokensDataStore(realm: realm, servers: config.enabledServers)
-    }()
+    private let tokensDataStore: TokensDataStore
     private let nftProvider: NFTProvider
-    private lazy var transactionsStorage = TransactionDataStore(realm: realm)
-    private lazy var realm = Wallet.functional.realm(forAccount: wallet)
+    private let transactionsStorage: TransactionDataStore
     private var cancelable = Set<AnyCancellable>()
     private let config: Config
     private let balanceUpdateSubject = PassthroughSubject<Void, Never>()
     private lazy var walletBalanceSubject: CurrentValueSubject<WalletBalance, Never> = .init(balance)
+    private var servers: [RPCServer]
 
     weak var delegate: WalletBalanceFetcherDelegate?
     var tokenObjects: [Activity.AssignedToken] {
@@ -60,25 +64,22 @@ class WalletBalanceFetcher: NSObject, WalletBalanceFetcherType {
             .eraseToAnyPublisher()
     }
 
-    required init(wallet: Wallet, nftProvider: NFTProvider, config: Config, assetDefinitionStore: AssetDefinitionStore, queue: DispatchQueue, coinTickersFetcher: CoinTickersFetcherType) {
+    required init(wallet: Wallet, servers: [RPCServer], tokensDataStore: TokensDataStore, transactionsStorage: TransactionDataStore, nftProvider: NFTProvider, config: Config, assetDefinitionStore: AssetDefinitionStore, queue: DispatchQueue, coinTickersFetcher: CoinTickersFetcherType) {
         self.wallet = wallet
         self.nftProvider = nftProvider
         self.assetDefinitionStore = assetDefinitionStore
         self.queue = queue
+        self.tokensDataStore = tokensDataStore
         self.coinTickersFetcher = coinTickersFetcher
         self.config = config
+        self.transactionsStorage = transactionsStorage
+        self.servers = servers
         super.init()
 
-        for each in config.enabledServers {
+        for each in servers {
             balanceFetchers[each] = createBalanceFetcher(wallet: wallet, server: each)
         }
 
-        config.enabledServersPublisher
-            .receive(on: RunLoop.main)
-            .sink { servers in
-                self.update(servers: servers)
-            }.store(in: &cancelable)
-
         coinTickersFetcher
             .tickersUpdatedPublisher
             .receive(on: RunLoop.main)
@@ -106,7 +107,9 @@ class WalletBalanceFetcher: NSObject, WalletBalanceFetcherType {
         }
     }
 
-    private func update(servers: [RPCServer]) {
+    func update(servers: [RPCServer]) {
+        self.servers = servers
+
         for each in servers {
             //NOTE: when we change servers it might happen the case when native
             tokensDataStore.addEthToken(forServer: each)
@@ -127,19 +130,19 @@ class WalletBalanceFetcher: NSObject, WalletBalanceFetcherType {
         delegate?.didUpdate(in: self)
     }
 
-    private func balanceViewModel(key tokenObject: TokenObject) -> BalanceBaseViewModel {
-        let ticker = coinTickersFetcher.ticker(for: tokenObject.addressAndRPCServer)
+    private func balanceViewModel(forToken token: TokenObject) -> BalanceBaseViewModel {
+        let ticker = coinTickersFetcher.ticker(for: token.addressAndRPCServer)
 
-        switch tokenObject.type {
+        switch token.type {
         case .nativeCryptocurrency:
-            let balance = Balance(value: BigInt(tokenObject.value, radix: 10) ?? BigInt(0))
-            return NativecryptoBalanceViewModel(server: tokenObject.server, balance: balance, ticker: ticker)
+            let balance = Balance(value: BigInt(token.value, radix: 10) ?? BigInt(0))
+            return NativecryptoBalanceViewModel(server: token.server, balance: balance, ticker: ticker)
         case .erc20:
-            let balance = ERC20Balance(tokenObject: tokenObject)
-            return ERC20BalanceViewModel(server: tokenObject.server, balance: balance, ticker: ticker)
+            let balance = ERC20Balance(tokenObject: token)
+            return ERC20BalanceViewModel(server: token.server, balance: balance, ticker: ticker)
         case .erc875, .erc721, .erc721ForTickets, .erc1155:
-            let balance = NFTBalance(tokenObject: tokenObject)
-            return NFTBalanceViewModel(server: tokenObject.server, balance: balance, ticker: ticker)
+            let balance = NFTBalance(tokenObject: token)
+            return NFTBalanceViewModel(server: token.server, balance: balance, ticker: ticker)
         }
     }
 
@@ -147,33 +150,26 @@ class WalletBalanceFetcher: NSObject, WalletBalanceFetcherType {
         balanceUpdateSubject.send(())
     }
 
-    func tokenBalance(_ key: AddressAndRPCServer) -> BalanceBaseViewModel {
-        guard let tokenObject = tokensDataStore.token(forContract: key.address, server: key.server) else {
-            let ticker = coinTickersFetcher.ticker(for: key)
-            return NativecryptoBalanceViewModel(server: key.server, balance: Balance(value: .zero), ticker: ticker)
+    func tokenBalance(_ key: AddressAndRPCServer) -> BalanceBaseViewModel? {
+        guard let token = tokensDataStore.token(forContract: key.address, server: key.server) else {
+            return nil
         }
 
-        return balanceViewModel(key: tokenObject)
+        return balanceViewModel(forToken: token)
     }
 
-    func tokenBalancePublisher(_ key: AddressAndRPCServer) -> AnyPublisher<BalanceBaseViewModel, Never> {
-        guard let tokenObject = tokensDataStore.token(forContract: key.address, server: key.server) else {
-            let ticker = coinTickersFetcher.ticker(for: key)
-            let viewModel: BalanceBaseViewModel = NativecryptoBalanceViewModel(server: key.server, balance: Balance(value: .zero), ticker: ticker)
-            let publisher = Just(viewModel)
-                .eraseToAnyPublisher()
+    func tokenBalancePublisher(_ key: AddressAndRPCServer) -> AnyPublisher<BalanceBaseViewModel?, Never> {
+        let tokenPublisher = tokensDataStore
+            .tokenValuePublisher(forContract: key.address, server: key.server)
+            .replaceError(with: nil)
 
-            return publisher
-        }
+        let forceReloadBalanceWhenServersChange = balanceUpdateSubject
+            .prepend(())
 
-        let publisher = tokenObject
-            .publisher(for: \.value, options: [.new, .initial])
-            .combineLatest(balanceUpdateSubject) { _, _ -> Void in return () }
-            .map { _ in self.balanceViewModel(key: tokenObject) }
-            .prepend(self.balanceViewModel(key: tokenObject))
+        return Publishers.CombineLatest(forceReloadBalanceWhenServersChange, tokenPublisher)
+            .map { $1 }
+            .map { [weak self] in $0.flatMap { self?.balanceViewModel(forToken: $0) } }
             .eraseToAnyPublisher()
-
-        return publisher
     }
 
     var balance: WalletBalance {
@@ -243,6 +239,36 @@ extension WalletBalanceFetcher: PrivateTokensDataStoreDelegate {
     }
 }
 
+extension WalletBalanceFetcher: WalletBalanceFetcherTypeTests {
+
+    func setBalanceTestsOnly(_ value: BigInt, forToken token: TokenObject) {
+        tokensDataStore.updateToken(primaryKey: token.primaryKey, action: .value(value))
+    }
+
+    func deleteTokenTestsOnly(token: TokenObject) {
+        tokensDataStore.deleteTestsOnly(tokens: [token])
+    }
+
+    func addOrUpdateTokenTestsOnly(token: TokenObject) {
+        tokensDataStore.addTokenObjects(values: [
+            .tokenObject(token)
+        ])
+    }
+
+}
+
+extension Optional where Wrapped == TokenChange {
+    var initialValueOrBalanceChanged: Bool {
+        guard let strongSelf = self else { return true }
+        switch strongSelf.change {
+        case .initial:
+            return true
+        case .changed(let properties):
+            return properties.isBalanceUpdate
+        }
+    }
+}
+
 fileprivate extension Array where Element == PropertyChange {
     var isBalanceUpdate: Bool {
         contains(where: { $0.name == "value" || $0.name == "balance" })

AlphaWallet/Core/Coordinators/WalletBalance/WalletBalanceService.swift

@@ -9,14 +9,21 @@ import UIKit
 import BigInt
 import PromiseKit
 import Combine
+import RealmSwift
 
 protocol CoinTickerProvider: AnyObject {
     func coinTicker(_ addressAndRPCServer: AddressAndRPCServer) -> CoinTicker?
 }
 
-protocol TokenBalanceProvider: AnyObject {
-    func tokenBalance(_ key: AddressAndRPCServer, wallet: Wallet) -> BalanceBaseViewModel
-    func tokenBalancePublisher(_ addressAndRPCServer: AddressAndRPCServer, wallet: Wallet) -> AnyPublisher<BalanceBaseViewModel, Never>
+protocol TokenBalanceProviderTests {
+    func setBalanceTestsOnly(_ value: BigInt, forToken token: TokenObject, wallet: Wallet)
+    func deleteTokenTestsOnly(token: TokenObject, wallet: Wallet)
+    func addOrUpdateTokenTestsOnly(token: TokenObject, wallet: Wallet)
+}
+
+protocol TokenBalanceProvider: AnyObject, TokenBalanceProviderTests {
+    func tokenBalance(_ key: AddressAndRPCServer, wallet: Wallet) -> BalanceBaseViewModel?
+    func tokenBalancePublisher(_ addressAndRPCServer: AddressAndRPCServer, wallet: Wallet) -> AnyPublisher<BalanceBaseViewModel?, Never>
     func refreshBalance(updatePolicy: PrivateBalanceFetcher.RefreshBalancePolicy, wallets: [Wallet], force: Bool) -> Promise<Void>
 }
 
@@ -30,8 +37,8 @@ protocol WalletBalanceService: TokenBalanceProvider, CoinTickerProvider {
 class MultiWalletBalanceService: NSObject, WalletBalanceService {
     private let keystore: Keystore
     private let config: Config
-    private let assetDefinitionStore: AssetDefinitionStore
-    private var coinTickersFetcher: CoinTickersFetcherType
+    let assetDefinitionStore: AssetDefinitionStore
+    var coinTickersFetcher: CoinTickersFetcherType
     private var balanceFetchers: [Wallet: WalletBalanceFetcherType] = [:]
     private lazy var walletsSummarySubject: CurrentValueSubject<WalletSummary, Never> = {
         let balances = balanceFetchers.map { $0.value.balance }
@@ -74,14 +81,25 @@ class MultiWalletBalanceService: NSObject, WalletBalanceService {
 
                 strongSelf.notifyWalletsSummary()
             }.store(in: &cancelable)
+
+        config.enabledServersPublisher
+            .receive(on: RunLoop.main)
+            .sink { [weak self] servers in
+                guard let strongSelf = self else { return }
+
+                for wallet in strongSelf.walletAddressesStore.wallets {
+                    let fetcher = strongSelf.getOrCreateBalanceFetcher(for: wallet)
+                    fetcher.update(servers: servers)
+                }
+            }.store(in: &cancelable)
     }
 
-    func tokenBalance(_ key: AddressAndRPCServer, wallet: Wallet) -> BalanceBaseViewModel {
+    func tokenBalance(_ key: AddressAndRPCServer, wallet: Wallet) -> BalanceBaseViewModel? {
         return getOrCreateBalanceFetcher(for: wallet)
             .tokenBalance(key)
     }
 
-    func tokenBalancePublisher(_ addressAndRPCServer: AddressAndRPCServer, wallet: Wallet) -> AnyPublisher<BalanceBaseViewModel, Never> {
+    func tokenBalancePublisher(_ addressAndRPCServer: AddressAndRPCServer, wallet: Wallet) -> AnyPublisher<BalanceBaseViewModel?, Never> {
         return getOrCreateBalanceFetcher(for: wallet)
             .tokenBalancePublisher(addressAndRPCServer)
     }
@@ -153,8 +171,12 @@ class MultiWalletBalanceService: NSObject, WalletBalanceService {
         })
     }
 
-    private func createWalletBalanceFetcher(wallet: Wallet) -> WalletBalanceFetcherType {
-        let fetcher = WalletBalanceFetcher(wallet: wallet, nftProvider: nftProvider, config: config, assetDefinitionStore: assetDefinitionStore, queue: queue, coinTickersFetcher: coinTickersFetcher)
+    /// NOTE: internal for test ourposes
+    func createWalletBalanceFetcher(wallet: Wallet) -> WalletBalanceFetcherType {
+        let realm: Realm = Wallet.functional.realm(forAccount: wallet)
+        let tokensDataStore = MultipleChainsTokensDataStore(realm: realm, servers: config.enabledServers)
+        let transactionsStorage = TransactionDataStore(realm: realm)
+        let fetcher = WalletBalanceFetcher(wallet: wallet, servers: config.enabledServers, tokensDataStore: tokensDataStore, transactionsStorage: transactionsStorage, nftProvider: nftProvider, config: config, assetDefinitionStore: assetDefinitionStore, queue: queue, coinTickersFetcher: coinTickersFetcher)
         fetcher.delegate = self
 
         return fetcher
@@ -183,6 +205,23 @@ class MultiWalletBalanceService: NSObject, WalletBalanceService {
     }
 }
 
+extension MultiWalletBalanceService {
+    func setBalanceTestsOnly(_ value: BigInt, forToken token: TokenObject, wallet: Wallet) {
+        getOrCreateBalanceFetcher(for: wallet)
+            .setBalanceTestsOnly(value, forToken: token)
+    }
+
+    func deleteTokenTestsOnly(token: TokenObject, wallet: Wallet) {
+        getOrCreateBalanceFetcher(for: wallet)
+            .deleteTokenTestsOnly(token: token)
+    }
+
+    func addOrUpdateTokenTestsOnly(token: TokenObject, wallet: Wallet) {
+        getOrCreateBalanceFetcher(for: wallet)
+            .addOrUpdateTokenTestsOnly(token: token)
+    }
+}
+
 extension MultiWalletBalanceService: WalletBalanceFetcherDelegate {
 
     func didAddToken(in fetcher: WalletBalanceFetcherType) {

AlphaWallet/Core/TokenBalanceService.swift

@@ -5,39 +5,36 @@ import Combine
 
 protocol TokenBalanceService {
     var etherToken: TokenObject { get }
-    var ethBalanceViewModel: BalanceBaseViewModel { get }
-    var etherBalance: AnyPublisher<BalanceBaseViewModel, Never> { get }
+    var ethBalanceViewModel: BalanceBaseViewModel? { get }
+    var etherBalance: AnyPublisher<BalanceBaseViewModel?, Never> { get }
     var etherToFiatRatePublisher: AnyPublisher<Double?, Never> { get }
     var etherToFiatRate: Double? { get }
 
     func start()
     func refresh(refreshBalancePolicy: PrivateBalanceFetcher.RefreshBalancePolicy)
     func coinTicker(_ addressAndRPCServer: AddressAndRPCServer) -> CoinTicker?
-    func tokenBalancePublisher(_ addressAndRPCServer: AddressAndRPCServer) -> AnyPublisher<BalanceBaseViewModel, Never>
+    func tokenBalance(_ addressAndRPCServer: AddressAndRPCServer) -> BalanceBaseViewModel?
+    func tokenBalancePublisher(_ addressAndRPCServer: AddressAndRPCServer) -> AnyPublisher<BalanceBaseViewModel?, Never>
 }
 
 class SingleChainTokenBalanceService: NSObject, TokenBalanceService {
     private let wallet: Wallet
     private let server: RPCServer
-    private var emptyEtherBalance: BalanceBaseViewModel {
+    private var emptyEtherBalance: BalanceBaseViewModel? {
         return balanceProvider.tokenBalance(etherToken.addressAndRPCServer, wallet: wallet)
     }
     private let balanceProvider: TokenBalanceProvider & CoinTickerProvider
-    private var cancelable = Set<AnyCancellable>()
-    private lazy var etherBalanceSubject: CurrentValueSubject<BalanceBaseViewModel, Never> = .init(emptyEtherBalance)
 
-    private (set) lazy var etherToken: TokenObject = {
-        return MultipleChainsTokensDataStore.functional.etherToken(forServer: server)
-    }()
+    let etherToken: TokenObject
 
     private (set) lazy var etherToFiatRatePublisher: AnyPublisher<Double?, Never> = {
         return etherBalance
-            .map { $0.ticker?.price_usd }
+            .map { $0?.ticker?.price_usd }
             .eraseToAnyPublisher()
     }()
 
-    private (set) lazy var etherBalance: AnyPublisher<BalanceBaseViewModel, Never> = {
-        return etherBalanceSubject
+    private (set) lazy var etherBalance: AnyPublisher<BalanceBaseViewModel?, Never> = {
+        return tokenBalancePublisher(etherToken.addressAndRPCServer)
             .eraseToAnyPublisher()
     }()
 
@@ -46,28 +43,31 @@ class SingleChainTokenBalanceService: NSObject, TokenBalanceService {
             .flatMap { $0.price_usd }
     }
 
-    var ethBalanceViewModel: BalanceBaseViewModel {
-        return etherBalanceSubject.value
+    var ethBalanceViewModel: BalanceBaseViewModel? {
+        return tokenBalance(etherToken.addressAndRPCServer)
     }
 
-    init(wallet: Wallet, server: RPCServer, tokenBalanceProvider: TokenBalanceProvider & CoinTickerProvider) {
+    init(wallet: Wallet, server: RPCServer, etherToken: TokenObject, tokenBalanceProvider: TokenBalanceProvider & CoinTickerProvider) {
         self.wallet = wallet
+        self.etherToken = etherToken
         self.server = server
         self.balanceProvider = tokenBalanceProvider
         super.init()
     }
 
     func start() {
-        tokenBalancePublisher(etherToken.addressAndRPCServer)
-            .assign(to: \.value, on: etherBalanceSubject)
-            .store(in: &cancelable)
+        //no-op
     }
 
     func coinTicker(_ addressAndRPCServer: AddressAndRPCServer) -> CoinTicker? {
         balanceProvider.coinTicker(addressAndRPCServer)
     }
 
-    func tokenBalancePublisher(_ addressAndRPCServer: AddressAndRPCServer) -> AnyPublisher<BalanceBaseViewModel, Never> {
+    func tokenBalance(_ addressAndRPCServer: AddressAndRPCServer) -> BalanceBaseViewModel? {
+        balanceProvider.tokenBalance(addressAndRPCServer, wallet: wallet)
+    }
+
+    func tokenBalancePublisher(_ addressAndRPCServer: AddressAndRPCServer) -> AnyPublisher<BalanceBaseViewModel?, Never> {
         balanceProvider.tokenBalancePublisher(addressAndRPCServer, wallet: wallet)
     }
 

AlphaWallet/Core/Types/RealmStore.swift

@@ -35,10 +35,6 @@ final class RealmStore {
 }
 
 extension Realm {
-    var threadSafe: Realm {
-         try! Realm(configuration: configuration)
-    }
-
     public func safeWrite(_ block: (() throws -> Void)) throws {
         if isInWriteTransaction {
             try block()

AlphaWallet/Export/Coordinators/PromptBackupCoordinator.swift

@@ -115,7 +115,7 @@ class PromptBackupCoordinator: Coordinator {
             .receive(on: RunLoop.main)
             .sink { [weak self] viewModel in
                 guard let strongSelf = self else { return }
-                let dollarValue = viewModel.currencyAmountWithoutSymbol ?? 0
+                let dollarValue = viewModel?.currencyAmountWithoutSymbol ?? 0
                 if !dollarValue.isZero {
                     strongSelf.showCreateBackupAfterExceedThresholdPrompt(valueInUsd: dollarValue)
                 }

AlphaWallet/Market/Coordinators/ImportMagicLinkCoordinator.swift

@@ -403,10 +403,10 @@ class ImportMagicLinkCoordinator: Coordinator {
             .etherBalance
             .first()
             .receive(on: RunLoop.main)
-            .sink { [weak self] balance in
-                guard let celf = self else { return }
+            .sink { [weak self] viewModel in
+                guard let celf = self, let viewModel = viewModel else { return }
 
-                if balance.value > signedOrder.order.price {
+                if viewModel.value > signedOrder.order.price {
                     celf.handlePaidImportsImpl(signedOrder: signedOrder)
                 } else {
                     celf.notEnoughEthForPaidImport(signedOrder: signedOrder)

AlphaWallet/Tokens/Types/TokensDataStore.swift

@@ -5,11 +5,28 @@ import BigInt
 import RealmSwift
 import Combine
 
+enum DataStoreError: Error {
+    case objectTypeMismatch
+    case objectNotFound
+    case objectDeleted
+    case general(error: Error)
+}
+
+struct TokenChange {
+    let token: TokenObject
+    let change: TokenPropertiesChange
+}
+
+enum TokenPropertiesChange {
+    case initial
+    case changed(properties: [PropertyChange])
+}
+
 /// Multiple-chains tokens data store
 protocol TokensDataStore: NSObjectProtocol {
     func enabledTokenObjectsChangesetPublisher(forServers servers: [RPCServer]) -> AnyPublisher<ChangeSet<[TokenObject]>, Never>
     func enabledTokenObjects(forServers servers: [RPCServer]) -> [TokenObject]
-
+    func tokenValuePublisher(forContract contract: AlphaWallet.Address, server: RPCServer) -> AnyPublisher<TokenObject?, DataStoreError>
     func deletedContracts(forServer server: RPCServer) -> [DeletedContract]
     func delegateContracts(forServer server: RPCServer) -> [DelegateContract]
     func hiddenContracts(forServer server: RPCServer) -> [HiddenContract]
@@ -18,10 +35,8 @@ protocol TokensDataStore: NSObjectProtocol {
     func token(forContract contract: AlphaWallet.Address, server: RPCServer) -> TokenObject?
     @discardableResult func addCustom(tokens: [ERCToken], shouldUpdateBalance: Bool) -> [TokenObject]
     func add(hiddenContracts: [HiddenContract])
-
     func deleteTestsOnly(tokens: [TokenObject])
     func updateOrderedTokens(with orderedTokens: [TokenObject])
-
     func add(tokenUpdates updates: [TokenUpdate])
     @discardableResult func updateToken(primaryKey: String, action: TokenUpdateAction) -> Bool?
     @discardableResult func addTokenObjects(values: [SingleChainTokensAutodetector.AddTokenObjectOperation]) -> [TokenObject]
@@ -72,6 +87,39 @@ enum TokenUpdateAction {
         return publisher
     }
 
+    func tokenValuePublisher(forContract contract: AlphaWallet.Address, server: RPCServer) -> AnyPublisher<TokenObject?, DataStoreError> {
+            let predicate = MultipleChainsTokensDataStore
+                .functional
+                .tokenPredicate(server: server, contract: contract)
+
+            let publisher: CurrentValueSubject<TokenObject?, DataStoreError> = .init(nil)
+            var cancelable: AnyCancellable?
+
+            store.performSync { realm in
+                guard let token = realm.objects(TokenObject.self).filter(predicate).first else {
+                    publisher.send(completion: .failure(DataStoreError.objectNotFound))
+                    return
+                }
+                let valuePublisher = token
+                    .publisher(for: \.value, options: [.initial, .new])
+                    .map { _ -> TokenObject in return token }
+                    .setFailureType(to: DataStoreError.self)
+
+                cancelable = valuePublisher
+                    .sink(receiveCompletion: { compl in
+                        publisher.send(completion: compl)
+                    }, receiveValue: { token in
+                        publisher.send(token)
+                    })
+            }
+
+            return publisher
+                .handleEvents(receiveCancel: {
+                    cancelable?.cancel()
+                })
+                .eraseToAnyPublisher()
+        }
+
     func enabledTokenObjects(forServers servers: [RPCServer]) -> [TokenObject] {
         var tokens: [TokenObject] = []
         store.performSync { realm in

AlphaWallet/Tokens/ViewControllers/SelectTokenViewController.swift

@@ -144,7 +144,7 @@ extension SelectTokenViewController: UITableViewDataSource {
             cell.configure(viewModel: .init(
                 token: token,
                 ticker: session.tokenBalanceService.coinTicker(token.addressAndRPCServer),
-                currencyAmount: session.tokenBalanceService.ethBalanceViewModel.currencyAmountWithoutSymbol,
+                currencyAmount: session.tokenBalanceService.ethBalanceViewModel?.currencyAmountWithoutSymbol,
                 assetDefinitionStore: assetDefinitionStore
             ))
             cell.accessoryType = viewModel.accessoryType(selectedToken, indexPath: indexPath)

AlphaWallet/Tokens/ViewControllers/TokenViewController.swift

@@ -220,7 +220,7 @@ class TokenViewController: UIViewController {
                 .etherBalance
                 .receive(on: RunLoop.main)
                 .sink { [weak self] viewModel in
-                    guard let celf = self else { return }
+                    guard let celf = self, let viewModel = viewModel else { return }
 
                     celf.tokenInfoPageView.viewModel.title = "\(viewModel.amountShort) \(viewModel.symbol)"
                     celf.tokenInfoPageView.viewModel.ticker = viewModel.ticker
@@ -237,7 +237,7 @@ class TokenViewController: UIViewController {
                 .tokenBalancePublisher(token.addressAndRPCServer)
                 .receive(on: RunLoop.main)
                 .sink { [weak self] viewModel in
-                    guard let strongSelf = self else { return }
+                    guard let strongSelf = self, let viewModel = viewModel else { return }
 
                     strongSelf.tokenInfoPageView.viewModel.currencyAmount = viewModel.currencyAmount
                     strongSelf.configure(viewModel: strongSelf.viewModel)

AlphaWallet/Tokens/ViewControllers/TokensViewController.swift

@@ -438,7 +438,7 @@ extension TokensViewController: UITableViewDataSource {
                     cell.configure(viewModel: .init(
                         token: token,
                         ticker: session.tokenBalanceService.coinTicker(token.addressAndRPCServer),
-                        currencyAmount: session.tokenBalanceService.ethBalanceViewModel.currencyAmountWithoutSymbol,
+                        currencyAmount: session.tokenBalanceService.ethBalanceViewModel?.currencyAmountWithoutSymbol,
                         assetDefinitionStore: assetDefinitionStore
                     ))
 

AlphaWallet/Tokens/ViewModels/TokenViewControllerViewModel.swift

@@ -95,7 +95,7 @@ struct TokenViewControllerViewModel {
     var fungibleBalance: BigInt? {
         switch transactionType {
         case .nativeCryptocurrency:
-            return session.tokenBalanceService.ethBalanceViewModel.value
+            return session.tokenBalanceService.ethBalanceViewModel?.value
         case .erc20Token(let tokenObject, _, _):
             return tokenObject.valueBigInt
         case .erc875Token, .erc875TokenOrder, .erc721Token, .erc721ForTicketToken, .erc1155Token, .dapp, .tokenScript, .claimPaidErc875MagicLink, .prebuilt:

AlphaWallet/Transactions/ViewControllers/TransactionViewController.swift

@@ -15,7 +15,7 @@ class TransactionViewController: UIViewController {
             transactionRow: transactionRow,
             chainState: session.chainState,
             wallet: session.account,
-            currencyRate: session.tokenBalanceService.ethBalanceViewModel.currencyRate
+            currencyRate: session.tokenBalanceService.ethBalanceViewModel?.currencyRate
         )
     }()
     private let roundedBackground = RoundedBackground()

AlphaWallet/Transfer/ViewModels/ConfigureTransactionViewModel.swift

@@ -19,7 +19,7 @@ struct ConfigureTransactionViewModel {
         configurator.session.server
     }
     private var currencyRate: CurrencyRate? {
-        configurator.session.tokenBalanceService.ethBalanceViewModel.currencyRate
+        configurator.session.tokenBalanceService.ethBalanceViewModel?.currencyRate
     }
 
     var recoveryMode: RecoveryMode
@@ -170,7 +170,7 @@ struct ConfigureTransactionViewModel {
         let isSelected = selectedConfigurationType == configurationType
         let configuration = configurations[configurationType]!
         //TODO if subscribable price are resolved or changes, will be good to refresh, but not essential
-        let ethPrice = configurator.session.tokenBalanceService.ethBalanceViewModel.ticker?.price_usd
+        let ethPrice = configurator.session.tokenBalanceService.ethBalanceViewModel?.ticker?.price_usd
         return .init(configuration: configuration, configurationType: configurationType, cryptoToDollarRate: ethPrice, symbol: server.symbol, title: configurationType.title, isSelected: isSelected)
     }
 
@@ -178,7 +178,7 @@ struct ConfigureTransactionViewModel {
         let isSelected = selectedConfigurationType == configurationType
         let configuration = configurations[configurationType]!
         //TODO if subscribable price are resolved or changes, will be good to refresh, but not essential
-        let ethPrice = configurator.session.tokenBalanceService.ethBalanceViewModel.ticker?.price_usd
+        let ethPrice = configurator.session.tokenBalanceService.ethBalanceViewModel?.ticker?.price_usd
         return .init(configuration: configuration, configurationType: configurationType, cryptoToDollarRate: ethPrice, symbol: server.symbol, title: configurationType.title, isSelected: isSelected)
     }
 

AlphaWallet/Transfer/ViewModels/SendViewModel.swift

@@ -32,27 +32,6 @@ struct SendViewModel {
         return Colors.appBackground
     }
 
-    var token: TokenObject? {
-        switch transactionType {
-        case .nativeCryptocurrency:
-            return nil
-        case .erc20Token(let token, _, _):
-            return token
-        case .erc875Token(let token, _):
-            return token
-        case .erc875TokenOrder(let token, _):
-            return token
-        case .erc721Token(let token, _):
-            return token
-        case .erc721ForTicketToken(let token, _):
-            return token
-        case .erc1155Token(let token, _, _):
-            return token
-        case .dapp, .tokenScript, .claimPaidErc875MagicLink, .prebuilt:
-            return nil
-        }
-    }
-
     var textFieldsLabelTextColor: UIColor {
         return Colors.appGrayLabel
     }
@@ -96,10 +75,11 @@ struct SendViewModel {
     var availableLabelText: String? {
         switch transactionType {
         case .nativeCryptocurrency:
-            return R.string.localizable.sendAvailable(session.tokenBalanceService.ethBalanceViewModel.amountShort)
+            return session.tokenBalanceService.ethBalanceViewModel
+                .flatMap { return R.string.localizable.sendAvailable($0.amountShort) }
         case .erc20Token(let token, _, _):
-            let value = EtherNumberFormatter.short.string(from: token.valueBigInt, decimals: token.decimals)
-            return R.string.localizable.sendAvailable("\(value) \(transactionType.symbol)")
+            return session.tokenBalanceService.tokenBalance(token.addressAndRPCServer)
+                .flatMap { R.string.localizable.sendAvailable("\($0.amountShort) \(transactionType.symbol)") }
         case .dapp, .erc721ForTicketToken, .erc721Token, .erc875Token, .erc875TokenOrder, .erc1155Token, .tokenScript, .claimPaidErc875MagicLink, .prebuilt:
             break
         }
@@ -112,8 +92,7 @@ struct SendViewModel {
         case .nativeCryptocurrency:
             return false
         case .erc20Token(let token, _, _):
-            let tokenBalance = tokensDataStore.token(forContract: token.contractAddress, server: session.server)?.valueBigInt
-            return tokenBalance == nil
+            return session.tokenBalanceService.tokenBalance(token.addressAndRPCServer) == nil
         case .dapp, .erc721ForTicketToken, .erc721Token, .erc875Token, .erc1155Token, .erc875TokenOrder, .tokenScript, .claimPaidErc875MagicLink, .prebuilt:
             break
         }
@@ -132,14 +111,15 @@ struct SendViewModel {
     var allFundsFormattedValues: (allFundsFullValue: NSDecimalNumber?, allFundsShortValue: String)? {
         switch transactionType {
         case .nativeCryptocurrency:
-            let balance = session.tokenBalanceService.ethBalanceViewModel
+            guard let balance = session.tokenBalanceService.ethBalanceViewModel else { return nil }
             let fullValue = EtherNumberFormatter.plain.string(from: balance.value, units: .ether).droppedTrailingZeros
             let shortValue = EtherNumberFormatter.shortPlain.string(from: balance.value, units: .ether).droppedTrailingZeros
 
             return (fullValue.optionalDecimalValue, shortValue)
         case .erc20Token(let token, _, _):
-            let fullValue = EtherNumberFormatter.plain.string(from: token.valueBigInt, decimals: token.decimals).droppedTrailingZeros
-            let shortValue = EtherNumberFormatter.shortPlain.string(from: token.valueBigInt, decimals: token.decimals).droppedTrailingZeros
+            guard let balance = session.tokenBalanceService.tokenBalance(token.addressAndRPCServer) else { return nil }
+            let fullValue = EtherNumberFormatter.plain.string(from: balance.value, decimals: token.decimals).droppedTrailingZeros
+            let shortValue = EtherNumberFormatter.shortPlain.string(from: balance.value, decimals: token.decimals).droppedTrailingZeros
 
             return (fullValue.optionalDecimalValue, shortValue)
         case .dapp, .erc721ForTicketToken, .erc721Token, .erc875Token, .erc1155Token, .erc875TokenOrder, .tokenScript, .claimPaidErc875MagicLink, .prebuilt:
@@ -173,11 +153,11 @@ struct SendViewModel {
 
         switch transactionType {
         case .nativeCryptocurrency:
-            if session.tokenBalanceService.ethBalanceViewModel.value < value {
+            if let viewModel = session.tokenBalanceService.ethBalanceViewModel, viewModel.value < value {
                 return nil
             }
         case .erc20Token(let token, _, _):
-            if let tokenBalance = tokensDataStore.token(forContract: token.contractAddress, server: session.server)?.valueBigInt, tokenBalance < value {
+            if let viewModel = session.tokenBalanceService.tokenBalance(token.addressAndRPCServer), viewModel.value < value {
                 return nil
             }
         case .dapp, .erc721ForTicketToken, .erc721Token, .erc875Token, .erc1155Token, .erc875TokenOrder, .tokenScript, .claimPaidErc875MagicLink, .prebuilt:

AlphaWalletTests/Factories/WalletSession.swift

@@ -1,6 +1,9 @@
 // Copyright SIX DAY LLC. All rights reserved.
 
 import Foundation
+import PromiseKit
+import Combine
+import RealmSwift
 @testable import AlphaWallet
 
 extension WalletSession {
@@ -23,11 +26,12 @@ extension WalletSession {
         server: RPCServer = .main,
         config: Config = .make()
     ) -> WalletSession {
+        let tokenBalanceService = FakeSingleChainTokenBalanceService(wallet: account, server: server, etherToken: TokenObject(contract: AlphaWallet.Address.make(), server: server, value: "0", type: .nativeCryptocurrency))
         return WalletSession(
             account: account,
             server: server,
             config: config,
-            tokenBalanceService: FakeSingleChainTokenBalanceService(wallet: account, server: server)
+            tokenBalanceService: tokenBalanceService
         )
     }
 
@@ -37,65 +41,81 @@ extension WalletSession {
         config: Config = .make(),
         tokenBalanceService: TokenBalanceService
     ) -> WalletSession {
+        let tokenBalanceService = FakeSingleChainTokenBalanceService(wallet: account, server: server, etherToken: TokenObject(contract: AlphaWallet.Address.make(), server: server, value: "0", type: .nativeCryptocurrency))
         return WalletSession(
             account: account,
             server: server,
             config: config,
-            tokenBalanceService: FakeSingleChainTokenBalanceService(wallet: account, server: server)
+            tokenBalanceService: tokenBalanceService
         )
     }
 }
 
-import PromiseKit
-import Combine
-
-private final class FakeTokenBalanceProvider: TokenBalanceProvider, CoinTickerProvider {
-    private var balanceSubject = CurrentValueSubject<Balance?, Never>(nil)
-
-    var balance: Balance? {
-        didSet {
-            balanceSubject.value = balance
-        }
+private class FakeNftProvider: NFTProvider {
+    func nonFungible(wallet: Wallet, server: RPCServer) -> Promise<OpenSeaNonFungiblesToAddress> {
+        return .value([:])
     }
+}
 
-    func tokenBalance(_ key: AddressAndRPCServer, wallet: Wallet) -> BalanceBaseViewModel {
-        let b: Balance = balance ?? .init(value: .zero)
-        return NativecryptoBalanceViewModel(server: key.server, balance: b, ticker: nil)
+extension Realm {
+    static func fake(forWallet wallet: Wallet) -> Realm {
+        return try! Realm(configuration: Realm.Configuration(inMemoryIdentifier: "MyInMemoryRealm-\(wallet.address.eip55String)"))
     }
+}
 
-    func coinTicker(_ addressAndRPCServer: AddressAndRPCServer) -> CoinTicker? {
-        return nil
+class FakeMultiWalletBalanceService: MultiWalletBalanceService {
+    private var servers: [RPCServer] = []
+    private let wallet: Wallet
+    lazy var tokensDataStore = FakeTokensDataStore(account: wallet, servers: servers)
+
+    init(wallet: Wallet = .make(), servers: [RPCServer] = [.main]) {
+        self.servers = servers
+        self.wallet = wallet
+
+        let tickersFetcher = FakeCoinTickersFetcher()
+        var walletAddressesStore = EtherKeystore.migratedWalletAddressesStore(userDefaults: .test)
+        walletAddressesStore.addToListOfWatchEthereumAddresses(wallet.address)
+
+        let keystore = FakeKeystore(wallets: [wallet], recentlyUsedWallet: wallet)
+        super.init(keystore: keystore, config: .make(), assetDefinitionStore: .init(), coinTickersFetcher: tickersFetcher, walletAddressesStore: walletAddressesStore)
     }
 
-    func tokenBalancePublisher(_ addressAndRPCServer: AddressAndRPCServer, wallet: Wallet) -> AnyPublisher<BalanceBaseViewModel, Never> {
-        return balanceSubject
-            .map { $0 ?? Balance(value: .zero) }
-            .map { NativecryptoBalanceViewModel(server: addressAndRPCServer.server, balance: $0, ticker: nil) }
-            .print("XXX.tokenBalancePublisher")
-            .eraseToAnyPublisher()
+    override func createWalletBalanceFetcher(wallet: Wallet) -> WalletBalanceFetcherType {
+        let nftProvider = FakeNftProvider()
+        let fetcher = WalletBalanceFetcher(wallet: wallet, servers: servers, tokensDataStore: tokensDataStore, transactionsStorage: FakeTransactionsStorage(), nftProvider: nftProvider, config: .make(), assetDefinitionStore: assetDefinitionStore, queue: .main, coinTickersFetcher: coinTickersFetcher)
+        fetcher.delegate = self
+
+        return fetcher
     }
+}
 
-    func refreshBalance(for wallet: Wallet) -> Promise<Void> {
-        return .init()
+class FakeSingleChainTokenBalanceService: SingleChainTokenBalanceService {
+    private let balanceService: FakeMultiWalletBalanceService
+    private let wallet: Wallet
+
+    var tokensDataStore: TokensDataStore {
+        balanceService.tokensDataStore
     }
 
-    func refreshEthBalance(for wallet: Wallet) -> Promise<Void> {
-        return .init()
+    init(wallet: Wallet, server: RPCServer, etherToken: TokenObject) {
+        self.wallet = wallet
+        balanceService = FakeMultiWalletBalanceService(wallet: wallet, servers: [server])
+        super.init(wallet: wallet, server: server, etherToken: etherToken, tokenBalanceProvider: balanceService)
     }
 
-    func refreshBalance(updatePolicy: PrivateBalanceFetcher.RefreshBalancePolicy, wallets: [Wallet], force: Bool) -> Promise<Void> {
-        return .init()
+    func setBalanceTestsOnly(balance: Balance, forToken token: TokenObject) {
+        balanceService.setBalanceTestsOnly(balance.value, forToken: token, wallet: wallet)
     }
-}
 
-class FakeSingleChainTokenBalanceService: SingleChainTokenBalanceService {
-    private let balanceProvider = FakeTokenBalanceProvider()
+    func addOrUpdateTokenTestsOnly(token: TokenObject) {
+        balanceService.addOrUpdateTokenTestsOnly(token: token, wallet: wallet)
+    }
 
-    var balance: Balance? {
-        didSet { balanceProvider.balance = balance }
+    func deleteTokenTestsOnly(token: TokenObject) {
+        balanceService.deleteTokenTestsOnly(token: token, wallet: wallet)
     }
 
-    init(wallet: Wallet, server: RPCServer) {
-        super.init(wallet: wallet, server: server, tokenBalanceProvider: balanceProvider)
+    override func refresh(refreshBalancePolicy: PrivateBalanceFetcher.RefreshBalancePolicy) {
+        //no-op
     }
 }

AlphaWalletTests/Settings/Coordinators/SettingsCoordinatorTests.swift

@@ -95,62 +95,3 @@ class SettingsCoordinatorTests: XCTestCase {
         XCTAssertEqual(1, walletAddressesStore.wallets.count)
     }
 }
-
-final class FakeMultiWalletBalanceService: WalletBalanceService {
-    func walletBalance(wallet: Wallet) -> AnyPublisher<WalletBalance, Never> {
-        return Just(WalletBalance(wallet: wallet, values: .init()))
-            .eraseToAnyPublisher()
-    }
-
-    var walletsSummary: AnyPublisher<WalletSummary, Never> {
-        Just(WalletSummary(balances: []))
-            .eraseToAnyPublisher()
-    }
-    func tokenBalance(_ key: AddressAndRPCServer, wallet: Wallet) -> BalanceBaseViewModel {
-        return NativecryptoBalanceViewModel(server: key.server, balance: Balance(value: .zero), ticker: nil)
-    }
-
-    func coinTicker(_ addressAndRPCServer: AddressAndRPCServer) -> CoinTicker? {
-        return nil
-    }
-
-    func subscribableWalletBalance(wallet: Wallet) -> Subscribable<WalletBalance> {
-        return .init(nil)
-    }
-
-    func tokenBalancePublisher(_ addressAndRPCServer: AddressAndRPCServer, wallet: Wallet) -> AnyPublisher<BalanceBaseViewModel, Never> {
-        let viewModel = NativecryptoBalanceViewModel(server: addressAndRPCServer.server, balance: Balance(value: .zero), ticker: nil)
-        return Just(viewModel)
-            .eraseToAnyPublisher()
-    }
-
-    func refreshBalance(for wallet: Wallet) -> Promise<Void> {
-        return .value(())
-    }
-
-    func refreshEthBalance(for wallet: Wallet) -> Promise<Void> {
-        return .value(())
-    }
-
-    var subscribableWalletsSummary: Subscribable<WalletSummary> = .init(nil)
-
-    init(config: Config = .make(), account: Wallet = .make()) {
-
-    }
-
-    func start() {
-
-    }
-
-    func refreshBalance() -> Promise<Void> {
-        return .value(())
-    }
-
-    func refreshEthBalance() -> Promise<Void> {
-        return .value(())
-    }
-
-    func refreshBalance(updatePolicy: PrivateBalanceFetcher.RefreshBalancePolicy, wallets: [Wallet], force: Bool) -> Promise<Void> {
-        return .value(())
-    }
-}

AlphaWalletTests/Transfer/ViewControllers/SendViewControllerTests.swift

@@ -11,26 +11,77 @@ import XCTest
 import BigInt
 
 class SendViewControllerTests: XCTestCase {
-    private let tokensDataStore = FakeTokensDataStore()
-
-    private let tokenBalanceService = FakeSingleChainTokenBalanceService(wallet: .make(), server: .main)
-    private let nativeCryptocurrencyTransactionType: TransactionType = {
-        let token = TokenObject(contract: AlphaWallet.Address.make(), server: .main, value: "0", type: .nativeCryptocurrency)
+    private lazy var tokenBalanceService = FakeSingleChainTokenBalanceService(wallet: .make(), server: .main, etherToken: token)
+    private let token = TokenObject(contract: AlphaWallet.Address.make(), server: .main, value: "0", type: .nativeCryptocurrency)
+    private lazy var nativeCryptocurrencyTransactionType: TransactionType = {
         return .nativeCryptocurrency(token, destination: nil, amount: nil)
     }()
 
     private lazy var session: WalletSession = {
+        tokenBalanceService.addOrUpdateTokenTestsOnly(token: token)
         tokenBalanceService.start()
         return .make(tokenBalanceService: tokenBalanceService)
     }()
 
+    func testBalanceUpdates() {
+        let token = TokenObject(contract: AlphaWallet.Address.make(), server: .main, decimals: 18, value: "2000000020224719101120", type: .erc20)
+        var balance = tokenBalanceService.tokenBalance(token.addressAndRPCServer)
+        XCTAssertNil(balance)
+
+        let isNilExpectation = self.expectation(description: "Non nil value")
+        _ = tokenBalanceService.tokenBalancePublisher(token.addressAndRPCServer).sink { value in
+            guard value == nil else { return }
+            isNilExpectation.fulfill()
+        }
+        waitForExpectations(timeout: 10)
+
+        tokenBalanceService.addOrUpdateTokenTestsOnly(token: token)
+
+        balance = tokenBalanceService.tokenBalance(token.addressAndRPCServer)
+        XCTAssertNotNil(balance)
+
+        let initialExpectation = self.expectation(description: "Initial value")
+
+        _ = tokenBalanceService.tokenBalancePublisher(token.addressAndRPCServer).sink { value in
+            guard value != nil else { return }
+            initialExpectation.fulfill()
+        }
+
+        waitForExpectations(timeout: 10)
+    }
+
+    func testUpdateNativeCryptoBalance() {
+        let token = TokenObject(contract: AlphaWallet.Address.make(), server: .main, value: "0", type: .nativeCryptocurrency)
+        let tokenBalanceService = FakeSingleChainTokenBalanceService(wallet: .make(), server: .main, etherToken: token)
+        let session: WalletSession = {
+            tokenBalanceService.addOrUpdateTokenTestsOnly(token: token)
+            tokenBalanceService.start()
+            return .make(tokenBalanceService: tokenBalanceService)
+        }()
+
+        XCTAssertEqual(session.tokenBalanceService.etherToken, token)
+        XCTAssertNotNil(session.tokenBalanceService.ethBalanceViewModel)
+        XCTAssertEqual(session.tokenBalanceService.ethBalanceViewModel!.value, .zero)
+
+        let testValue1 = BigInt("10000000000000000000000")
+        tokenBalanceService.setBalanceTestsOnly(balance: .init(value: testValue1), forToken: token)
+
+        XCTAssertEqual(session.tokenBalanceService.ethBalanceViewModel!.value, testValue1)
+
+        let testValue2 = BigInt("20000000000000000000000")
+        tokenBalanceService.setBalanceTestsOnly(balance: .init(value: testValue2), forToken: token)
+
+        XCTAssertNotNil(session.tokenBalanceService.ethBalanceViewModel)
+        XCTAssertEqual(session.tokenBalanceService.ethBalanceViewModel!.value, testValue2)
+    }
+
     func testNativeCryptocurrencyAllFundsValueSpanish() {
         let vc = createSendViewControllerAndSetLocale(locale: .spanish, transactionType: nativeCryptocurrencyTransactionType)
 
         XCTAssertEqual(vc.amountTextField.value, "")
 
         let testValue = BigInt("10000000000000000000000")
-        tokenBalanceService.balance = .some(.init(value: testValue))
+        tokenBalanceService.setBalanceTestsOnly(balance: .init(value: testValue), forToken: token)
 
         vc.allFundsSelected()
         XCTAssertEqual(vc.amountTextField.value, "10000")
@@ -45,7 +96,7 @@ class SendViewControllerTests: XCTestCase {
         XCTAssertEqual(vc.amountTextField.value, "")
 
         let testValue = BigInt("10000000000000000000000")
-        tokenBalanceService.balance = .some(.init(value: testValue))
+        tokenBalanceService.setBalanceTestsOnly(balance: .init(value: testValue), forToken: token)
 
         vc.allFundsSelected()
 
@@ -62,7 +113,7 @@ class SendViewControllerTests: XCTestCase {
         XCTAssertEqual(vc.amountTextField.value, "")
 
         let testValue = BigInt("10000000000000")
-        tokenBalanceService.balance = .some(.init(value: testValue))
+        tokenBalanceService.setBalanceTestsOnly(balance: .init(value: testValue), forToken: token)
 
         vc.allFundsSelected()
 
@@ -75,25 +126,28 @@ class SendViewControllerTests: XCTestCase {
 
     func testERC20AllFunds() {
         let token = TokenObject(contract: AlphaWallet.Address.make(), server: .main, decimals: 18, value: "2000000020224719101120", type: .erc20)
-
+        tokenBalanceService.addOrUpdateTokenTestsOnly(token: token)
         let vc = createSendViewControllerAndSetLocale(locale: .spanish, transactionType: .erc20Token(token, destination: .none, amount: nil))
-
+        tokenBalanceService.setBalanceTestsOnly(balance: .init(value: BigInt("2000000020224719101120")), forToken: token)
         XCTAssertEqual(vc.amountTextField.value, "")
-
         vc.allFundsSelected()
 
         XCTAssertEqual(vc.amountTextField.value, "2000")
         XCTAssertNotNil(vc.shortValueForAllFunds)
         XCTAssertTrue((vc.shortValueForAllFunds ?? "").nonEmpty)
 
+        tokenBalanceService.setBalanceTestsOnly(balance: .init(value: .zero), forToken: token)
+        vc.allFundsSelected()
+        XCTAssertEqual(vc.amountTextField.value, "0")
+
         Config.setLocale(AppLocale.system)
     }
 
     func testERC20AllFundsSpanish() {
-        let token = TokenObject(contract: AlphaWallet.Address.make(), server: .main, decimals: 18, value: "2020224719101120", type: .erc20)
-
+        let token = TokenObject(contract: AlphaWallet.Address.make(), server: .main, decimals: 18, value: "0", type: .erc20)
+        tokenBalanceService.addOrUpdateTokenTestsOnly(token: token)
         let vc = createSendViewControllerAndSetLocale(locale: .spanish, transactionType: .erc20Token(token, destination: .none, amount: nil))
-
+        tokenBalanceService.setBalanceTestsOnly(balance: .init(value: BigInt("2020224719101120")), forToken: token)
         XCTAssertEqual(vc.amountTextField.value, "")
 
         vc.allFundsSelected()
@@ -105,9 +159,30 @@ class SendViewControllerTests: XCTestCase {
         Config.setLocale(AppLocale.system)
     }
 
-    func testERC20AllFundsEnglish() {
+    func testTokenBalance() {
         let token = TokenObject(contract: AlphaWallet.Address.make(), server: .main, decimals: 18, value: "2020224719101120", type: .erc20)
+        tokenBalanceService.addOrUpdateTokenTestsOnly(token: token)
+
+        let tokens = tokenBalanceService.tokensDataStore.enabledTokenObjects(forServers: [.main])
+
+        XCTAssertTrue(tokens.contains(token))
+
+        let viewModel = tokenBalanceService.tokenBalance(token.addressAndRPCServer)
+        XCTAssertNotNil(viewModel)
+        XCTAssertEqual(viewModel!.value, BigInt("2020224719101120")!)
+
+        tokenBalanceService.setBalanceTestsOnly(balance: .init(value: BigInt("10000000000000")), forToken: token)
+
+        let viewModel_2 = tokenBalanceService.tokenBalance(token.addressAndRPCServer)
+        XCTAssertNotNil(viewModel_2)
+        XCTAssertEqual(viewModel_2!.value, BigInt("10000000000000")!)
+    }
+
+    func testERC20AllFundsEnglish() {
+        let token = TokenObject(contract: AlphaWallet.Address.make(), server: .main, decimals: 18, value: "0", type: .erc20)
+        tokenBalanceService.addOrUpdateTokenTestsOnly(token: token)
         let vc = createSendViewControllerAndSetLocale(locale: .english, transactionType: .erc20Token(token, destination: .none, amount: nil))
+        tokenBalanceService.setBalanceTestsOnly(balance: .init(value: BigInt("2020224719101120")), forToken: token)
 
         XCTAssertEqual(vc.amountTextField.value, "")
 
@@ -121,9 +196,10 @@ class SendViewControllerTests: XCTestCase {
     }
 
     func testERC20English() {
-        let token = TokenObject(contract: AlphaWallet.Address.make(), server: .main, decimals: 18, value: "2020224719101120", type: .erc20)
+        let token = TokenObject(contract: AlphaWallet.Address.make(), server: .main, decimals: 18, value: "0", type: .erc20)
+        tokenBalanceService.addOrUpdateTokenTestsOnly(token: token)
         let vc = createSendViewControllerAndSetLocale(locale: .english, transactionType: .erc20Token(token, destination: .none, amount: nil))
-
+        tokenBalanceService.setBalanceTestsOnly(balance: .init(value: BigInt("2020224719101120")), forToken: token)
         XCTAssertEqual(vc.amountTextField.value, "")
 
         XCTAssertNil(vc.shortValueForAllFunds)
@@ -137,10 +213,10 @@ class SendViewControllerTests: XCTestCase {
         Config.setLocale(locale)
 
         let vc = SendViewController(session: session,
-                                    tokensDataStore: tokensDataStore,
+                                    tokensDataStore: tokenBalanceService.tokensDataStore,
                                     transactionType: nativeCryptocurrencyTransactionType)
 
-        vc.configure(viewModel: .init(transactionType: transactionType, session: session, tokensDataStore: tokensDataStore))
+        vc.configure(viewModel: .init(transactionType: transactionType, session: session, tokensDataStore: tokenBalanceService.tokensDataStore))
 
         return vc
     }