fix: separate eth and token balances

MishkaRogachev · Oct 1, 2024 · de9fcc0 · de9fcc0
1 parent 8200ff6
commit de9fcc0
Show file tree

Hide file tree

Showing 6 changed files with 78 additions and 71 deletions.
diff --git a/src/core/provider.rs b/src/core/provider.rs
@@ -69,42 +69,42 @@ impl Provider {
         Ok(result.answer as f64 / 10f64.powi(8))
     }
 
-    pub async fn get_balances(&self, account: Address, tokens: &TokenList) -> anyhow::Result<Balances> {
+    pub async fn get_eth_balance(&self, account: Address) -> anyhow::Result<Balance> {
+        let wei = self.web3.eth().balance(account, None).await?;
+        let eth = wei_to_eth(wei);
+        let eth_usd_rate = self.get_eth_usd_rate().await?;
+        Ok(Balance::new(eth, eth_usd_rate * eth, "ETH", self.chain.is_test_network()))
+    }
+
+    pub async fn get_token_balances(&self, account: Address, tokens: &TokenList) -> anyhow::Result<Balances> {
         let mut balances = Vec::new();
         let eth_usd_rate = self.get_eth_usd_rate().await?;
 
         for token in tokens {
-            let balance = if token.symbol == "ETH" {
-                // Handle ETH balance
-                let wei = self.web3.eth().balance(account, None).await?;
-                let eth = wei_to_eth(wei);
-                Balance::new(eth, eth_usd_rate * eth, &token.symbol, self.chain.is_test_network())
-            } else {
-                // Handle ERC-20 token balances
-                let token_chain_data = match token.get_chain_data(&self.chain) {
-                    Some(token_chain_data) => token_chain_data,
-                    None => continue,
-                };
-                let contract = match Contract::from_json(self.web3.eth(), token_chain_data.contract_address, ERC20_BALANCE_ABI) {
-                    Ok(contract) => contract,
-                    Err(_) => {
-                        log::warn!("Failed to create contract for token {} on {}", token.symbol, self.chain);
-                        continue
-                    },
-                };
-
-                let balance: U256 = match contract
-                    .query("balanceOf", (account,), None, Options::default(), None).await {
-                    Ok(balance) => balance,
-                    Err(_) => {
-                        log::warn!("Failed to get balance for token {} on {}", token.symbol, self.chain);
-                        continue
-                    },
-                };
-
-                let balance_f64 = balance.as_u128() as f64 / 10f64.powi(token_chain_data.decimals as i32);
-                Balance::new(balance_f64, balance_f64 * eth_usd_rate, &token.symbol, self.chain.is_test_network())
+            // Handle ERC-20 token balances
+            let token_chain_data = match token.get_chain_data(&self.chain) {
+                Some(token_chain_data) => token_chain_data,
+                None => continue,
+            };
+            let contract = match Contract::from_json(self.web3.eth(), token_chain_data.contract_address, ERC20_BALANCE_ABI) {
+                Ok(contract) => contract,
+                Err(_) => {
+                    log::warn!("Failed to create contract for token {} on {}", token.symbol, self.chain);
+                    continue
+                },
             };
+
+            let balance: U256 = match contract
+                .query("balanceOf", (account,), None, Options::default(), None).await {
+                Ok(balance) => balance,
+                Err(_) => {
+                    log::warn!("Failed to get balance for token {} on {}", token.symbol, self.chain);
+                    continue
+                },
+            };
+
+            let balance_f64 = balance.as_u128() as f64 / 10f64.powi(token_chain_data.decimals as i32);
+            let balance = Balance::new(balance_f64, balance_f64 * eth_usd_rate, &token.symbol, self.chain.is_test_network());
             balances.push(balance);
         }
         Ok(balances)

diff --git a/src/core/provider_test.rs b/src/core/provider_test.rs
@@ -40,20 +40,34 @@ mod tests {
     #[test_case(Chain::EthereumSepolia)]
     #[test_case(Chain::OptimismMainnet)]
     #[tokio::test]
-    async fn test_get_balances(chain: Chain) -> anyhow::Result<()> {
+    async fn test_get_eth_balance(chain: Chain) -> anyhow::Result<()> {
         let endpoint_url = get_endpoint_url()?;
 
-        let tokens = vec![
-            Token::new("Ethereum", "ETH"),
-        ];
-
         let account = web3::types::Address::from_low_u64_be(0);
         let provider = Provider::new(&endpoint_url, chain)?;
 
-        let balances = provider.get_balances(account, &tokens).await?;
-        assert_eq!(balances.len(), tokens.len());
-        assert_eq!(balances[0].currency, "ETH");
-        assert_ne!(balances[0].value, 0.0);
+        let balance = provider.get_eth_balance(account).await?;
+        assert_eq!(balance.currency, "ETH");
+        assert_ne!(balance.value, 0.0);
+
+        Ok(())
+    }
+
+    #[test_case("0x6B175474E89094C44Da98b954EedeAC495271d0F", "DAI", "Dai Stablecoin", 18)]
+    #[test_case("0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606EB48", "USDC", "USD Coin", 6)]
+    #[test_case("0xdAC17F958D2ee523a2206206994597C13D831ec7", "USDT", "Tether USD", 6)]
+    #[tokio::test]
+    async fn test_get_token_balances(contract_address: &str, symbol: &str, name: &str, decimals: u16) -> anyhow::Result<()> {
+        let endpoint_url = get_endpoint_url()?;
+
+        let token = Token::new(name, symbol).with_chain_data(Chain::EthereumMainnet, contract_address.parse()?, decimals);
+
+        let account = web3::types::Address::from_low_u64_be(0);
+        let provider = Provider::new(&endpoint_url, Chain::EthereumMainnet)?;
+
+        let balances = provider.get_token_balances(account, &vec![token]).await?;
+        assert_eq!(balances.len(), 1);
+        assert_eq!(balances[0].currency, symbol);
 
         Ok(())
     }

diff --git a/src/service/crypto.rs b/src/service/crypto.rs
@@ -6,6 +6,8 @@ use crate::{
     persistence::db::Db
 };
 
+const BALANCES_FETCH_PROVIDER_DELAY: std::time::Duration = std::time::Duration::from_millis(100);
+
 #[derive(Clone)]
 pub struct Crypto {
     db: Arc<Db>,
@@ -80,16 +82,26 @@ impl Crypto {
             for account in accounts.iter() {
                 let mut sum_balances: Balances = Vec::new();
                 for (_, provider) in providers.iter() {
-                    let response = provider.get_balances(*account, &token_list).await;
-                    match response {
-                        Ok(balances) => {
-                            sum_balances = Balance::extend_balances(&sum_balances, &balances);
+                    let balance = provider.get_eth_balance(*account).await;
+                    match balance {
+                        Ok(balance) => {
+                            sum_balances = Balance::extend_balances(&sum_balances, &vec![balance]);
+                        }
+                        Err(_err) => {
+                            log::error!("Failed to fetch ETH balance for {}", account);
+                        }
+                    }
+
+                    let token_balances = provider.get_token_balances(*account, &token_list).await;
+                    match token_balances {
+                        Ok(token_balances) => {
+                            sum_balances = Balance::extend_balances(&sum_balances, &token_balances);
                         }
                         Err(_err) => {
-                            // TODO: just log it
-                            // eprintln!("Failed to fetch balance for {}: {:?}", account, err);
+                            log::error!("Failed to fetch token balances for {}", account);
                         }
                     }
+                    tokio::time::sleep(BALANCES_FETCH_PROVIDER_DELAY).await;
                 }
                 account_balances.write().await.insert(*account, sum_balances);
             }

diff --git a/src/tui/screens/porfolio_accounts.rs b/src/tui/screens/porfolio_accounts.rs
@@ -13,7 +13,7 @@ use crate::tui::{widgets::{controls, account}, app::AppScreen};
 const SUMMARY_HEIGHT: u16 = 2;
 const SUMMARY_TEXT: &str = "Summary balance";
 
-const UPDATE_INTERVAL: tokio::time::Duration = tokio::time::Duration::from_millis(2000);
+const UPDATE_INTERVAL: tokio::time::Duration = tokio::time::Duration::from_secs(10);
 
 pub struct Screen {
     crypto: Arc<Mutex<Crypto>>,
@@ -40,7 +40,7 @@ impl Screen {
         let mut usd_summary = None;
         let mut test_network = false;
         for account in &self.accounts {
-            if let Some((_, usd_value, from_test_network)) = &account.get_total_balances() {
+            if let Some((usd_value, from_test_network)) = &account.get_total_usd_balance() {
                 usd_summary = Some(usd_summary.unwrap_or(0.0) + usd_value);
                 test_network = test_network || *from_test_network;
             }

diff --git a/src/tui/widgets/account.rs b/src/tui/widgets/account.rs
@@ -28,11 +28,10 @@ impl Account {
         }
     }
 
-    pub fn get_total_balances(&self) -> Option<(f64, f64, bool)> {
+    pub fn get_total_usd_balance(&self) -> Option<(f64, bool)> {
         self.balances.as_ref().map(|balances| {
-            balances.iter().fold((0.0, 0.0, false), |(total_value, total_usd, from_test), balance| {
+            balances.iter().fold((0.0, false), |(total_usd, from_test), balance| {
                 (
-                    total_value + balance.value,
                     total_usd + balance.usd_value,
                     from_test || balance.from_test_network,
                 )
@@ -45,8 +44,8 @@ impl Account {
     }
 
     fn render_total_balances(&mut self, frame: &mut Frame, area: Rect) {
-        if let Some((total_value, total_usd, from_test_network)) = self.get_total_balances() {
-            let balances_str = format!("{:.6} ETH ({:.2} USD)", total_value, total_usd);
+        if let Some((total_usd, from_test_network)) = self.get_total_usd_balance() {
+            let balances_str = format!("{:.2} USD", total_usd);
             let balances_color = if from_test_network { Color::Red } else { Color::Yellow };
             Paragraph::new(balances_str)
                 .style(Style::default().fg(balances_color).add_modifier(ratatui::style::Modifier::BOLD))

diff --git a/token_list.json b/token_list.json
@@ -1,22 +1,4 @@
 [{
-        "name": "Ethereum",
-        "symbol": "ETH",
-        "chain_data": {
-            "EthereumMainnet": {
-                "contract_address": "0x0000000000000000000000000000000000000000",
-                "decimals": 18
-            },
-            "OptimismMainnet": {
-                "contract_address": "0x4200000000000000000000000000000000000006",
-                "decimals": 18
-            },
-            "ArbitrumMainnet": {
-                "contract_address": "0x82af49447d8a07e3bd95bd0d56f35241523fbab1",
-                "decimals": 18
-            }
-        }
-    },
-    {
         "name": "Polygon",
         "symbol": "MATIC",
         "chain_data": {