nautilus_model/defi/
wallet.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2025 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16use std::{collections::HashSet, fmt::Display};
17
18use alloy_primitives::{Address, U256};
19
20use crate::{
21    defi::Token,
22    types::{Money, Quantity},
23};
24
25/// Represents the balance of a specific ERC-20 token held in a wallet.
26///
27/// This struct tracks the raw token amount along with optional USD valuation
28/// and the token metadata.
29#[derive(Debug)]
30pub struct TokenBalance {
31    /// The raw token amount as a 256-bit unsigned integer.
32    pub amount: U256,
33    /// The optional USD equivalent value of the token balance.
34    pub amount_usd: Option<Quantity>,
35    /// The token metadata including chain, address, name, symbol, and decimals.
36    pub token: Token,
37}
38
39impl TokenBalance {
40    /// Creates a new [`TokenBalance`] instance.
41    pub const fn new(amount: U256, token: Token) -> Self {
42        Self {
43            amount,
44            token,
45            amount_usd: None,
46        }
47    }
48
49    /// Converts the raw token amount to a human-readable [`Quantity`].
50    ///
51    /// # Errors
52    ///
53    /// Returns an error if the U256 amount cannot be converted to a `Quantity`.
54    pub fn as_quantity(&self) -> anyhow::Result<Quantity> {
55        Quantity::from_u256(self.amount, self.token.decimals)
56    }
57
58    /// Sets the USD equivalent value for this token balance.
59    pub fn set_amount_usd(&mut self, amount_usd: Quantity) {
60        self.amount_usd = Some(amount_usd);
61    }
62}
63
64impl Display for TokenBalance {
65    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
66        let quantity = self.as_quantity().unwrap_or_default();
67        match &self.amount_usd {
68            Some(usd) => write!(
69                f,
70                "TokenBalance(token={}, amount={}, usd=${:.2})",
71                self.token.symbol,
72                quantity.as_decimal(),
73                usd.as_f64()
74            ),
75            None => write!(
76                f,
77                "TokenBalance(token={}, amount={})",
78                self.token.symbol,
79                quantity.as_decimal()
80            ),
81        }
82    }
83}
84
85/// Represents the complete balance state of a blockchain wallet.
86///
87/// Tracks both the native currency balance (e.g., ETH, ARB) and ERC-20 token
88/// balances for a wallet address. The `token_universe` defines which tokens
89/// should be tracked for balance fetching.
90#[derive(Debug)]
91pub struct WalletBalance {
92    /// The balance of the chain's native currency
93    pub native_currency: Option<Money>,
94    /// Collection of ERC-20 token balances held in the wallet.
95    pub token_balances: Vec<TokenBalance>,
96    /// Set of token addresses to track for balance updates.
97    pub token_universe: HashSet<Address>,
98}
99
100impl WalletBalance {
101    /// Creates a new [`WalletBalance`] with the specified token universe.
102    pub const fn new(token_universe: HashSet<Address>) -> Self {
103        Self {
104            native_currency: None,
105            token_balances: vec![],
106            token_universe,
107        }
108    }
109
110    /// Returns `true` if the token universe has been initialized with token addresses.
111    pub fn is_token_universe_initialized(&self) -> bool {
112        !self.token_universe.is_empty()
113    }
114
115    /// Sets the native currency balance for the wallet.
116    pub fn set_native_currency_balance(&mut self, balance: Money) {
117        self.native_currency = Some(balance);
118    }
119
120    /// Adds an ERC-20 token balance to the wallet.
121    pub fn add_token_balance(&mut self, token_balance: TokenBalance) {
122        self.token_balances.push(token_balance);
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use std::sync::Arc;
129
130    use alloy_primitives::{U256, address};
131    use rstest::rstest;
132
133    use super::*;
134    use crate::defi::{
135        SharedChain, Token,
136        chain::chains,
137        stubs::{arbitrum, usdc, weth},
138    };
139
140    // Helper to create a token with specific decimals
141    fn create_token(symbol: &str, decimals: u8) -> Token {
142        Token::new(
143            Arc::new(chains::ETHEREUM.clone()),
144            address!("0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48"),
145            format!("{symbol} Token"),
146            symbol.to_string(),
147            decimals,
148        )
149    }
150
151    #[rstest]
152    fn test_token_balance_as_quantity_18_decimals(#[from(arbitrum)] chain: SharedChain) {
153        // Test case: NU token with 18 decimals
154        // Raw amount: 10342000000000000000000 (10342 * 10^18)
155        // Expected: 10342.000000000000000000
156        let token = Token::new(
157            chain,
158            address!("0x4fE83213D56308330EC302a8BD641f1d0113A4Cc"),
159            "NuCypher".to_string(),
160            "NU".to_string(),
161            18,
162        );
163        let amount = U256::from(10342u64) * U256::from(10u64).pow(U256::from(18u64));
164        let balance = TokenBalance::new(amount, token);
165
166        let quantity = balance.as_quantity().unwrap();
167        assert_eq!(
168            quantity.as_decimal().to_string(),
169            "10342.000000000000000000"
170        );
171    }
172
173    #[rstest]
174    fn test_token_balance_as_quantity_6_decimals() {
175        // Test case: USDC with 6 decimals
176        // Raw amount: 92220728254 (92220.728254 * 10^6)
177        // Expected: 92220.728254
178        let token = create_token("USDC", 6);
179        let amount = U256::from(92220728254u64);
180        let balance = TokenBalance::new(amount, token);
181
182        let quantity = balance.as_quantity().unwrap();
183        assert_eq!(quantity.as_decimal().to_string(), "92220.728254");
184    }
185
186    #[rstest]
187    fn test_token_balance_as_quantity_fractional_18_decimals(#[from(arbitrum)] chain: SharedChain) {
188        // Test case: mETH with 18 decimals and fractional amount
189        // Raw amount: 758325512078001391
190        // Expected: 0.758325512078001391
191        let token = Token::new(
192            chain,
193            address!("0xd5F7838F5C461fefF7FE49ea5ebaF7728bB0ADfa"),
194            "mETH".to_string(),
195            "mETH".to_string(),
196            18,
197        );
198        let amount = U256::from(758325512078001391u64);
199        let balance = TokenBalance::new(amount, token);
200
201        let quantity = balance.as_quantity().unwrap();
202        assert_eq!(quantity.as_decimal().to_string(), "0.758325512078001391");
203    }
204
205    #[rstest]
206    fn test_token_balance_display_18_decimals(#[from(arbitrum)] chain: SharedChain) {
207        // Test Display implementation with 18 decimal token
208        let token = Token::new(
209            chain,
210            address!("0x912CE59144191C1204E64559FE8253a0e49E6548"),
211            "Arbitrum".to_string(),
212            "ARB".to_string(),
213            18,
214        );
215        // 7922.013795343949480329 ARB
216        let amount = U256::from_str_radix("7922013795343949480329", 10).unwrap();
217        let balance = TokenBalance::new(amount, token);
218
219        let display = balance.to_string();
220        assert!(display.contains("ARB"));
221        assert!(display.contains("7922.013795343949480329"));
222    }
223
224    #[rstest]
225    fn test_token_balance_display_6_decimals() {
226        // Test Display implementation with 6 decimal token (USDC)
227        let token = create_token("USDC", 6);
228        let amount = U256::from(92220728254u64); // 92220.728254 USDC
229        let balance = TokenBalance::new(amount, token);
230
231        let display = balance.to_string();
232        assert!(display.contains("USDC"));
233        assert!(display.contains("92220.728254"));
234    }
235
236    #[rstest]
237    fn test_token_balance_set_amount_usd(weth: Token) {
238        let amount = U256::from(1u64) * U256::from(10u64).pow(U256::from(18u64));
239        let mut balance = TokenBalance::new(amount, weth);
240
241        assert!(balance.amount_usd.is_none());
242
243        let usd_value = Quantity::from("3500.00");
244        balance.set_amount_usd(usd_value);
245
246        assert!(balance.amount_usd.is_some());
247        assert_eq!(
248            balance.amount_usd.unwrap().as_decimal().to_string(),
249            "3500.00"
250        );
251    }
252
253    #[rstest]
254    fn test_wallet_balance_new_empty() {
255        let wallet = WalletBalance::new(HashSet::new());
256
257        assert!(wallet.native_currency.is_none());
258        assert!(wallet.token_balances.is_empty());
259        assert!(!wallet.is_token_universe_initialized());
260    }
261
262    #[rstest]
263    fn test_wallet_balance_with_token_universe() {
264        let mut tokens = HashSet::new();
265        tokens.insert(address!("0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48")); // USDC
266        tokens.insert(address!("0x912CE59144191C1204E64559FE8253a0e49E6548")); // ARB
267
268        let wallet = WalletBalance::new(tokens);
269
270        assert!(wallet.is_token_universe_initialized());
271        assert_eq!(wallet.token_universe.len(), 2);
272    }
273
274    #[rstest]
275    fn test_wallet_balance_set_native_currency() {
276        let mut wallet = WalletBalance::new(HashSet::new());
277
278        assert!(wallet.native_currency.is_none());
279
280        let eth_balance = Money::new(50.936054, crate::types::Currency::ETH());
281        wallet.set_native_currency_balance(eth_balance);
282
283        assert!(wallet.native_currency.is_some());
284    }
285
286    #[rstest]
287    fn test_wallet_balance_add_token_balance(usdc: Token, weth: Token) {
288        let mut wallet = WalletBalance::new(HashSet::new());
289
290        let usdc_balance = TokenBalance::new(U256::from(100_000_000u64), usdc); // 100 USDC
291        let weth_balance = TokenBalance::new(U256::from(10u64).pow(U256::from(18u64)), weth); // 1 WETH
292
293        wallet.add_token_balance(usdc_balance);
294        wallet.add_token_balance(weth_balance);
295
296        assert_eq!(wallet.token_balances.len(), 2);
297        assert_eq!(wallet.token_balances[0].token.symbol, "USDC");
298        assert_eq!(wallet.token_balances[1].token.symbol, "WETH");
299    }
300}
nautilus_model/defi/wallet.rs

nautilus_model/defi/
wallet.rs
