nautilus_model/defi/
wallet.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 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    #[must_use]
42    pub const fn new(amount: U256, token: Token) -> Self {
43        Self {
44            amount,
45            token,
46            amount_usd: None,
47        }
48    }
49
50    /// Converts the raw token amount to a human-readable [`Quantity`].
51    ///
52    /// # Errors
53    ///
54    /// Returns an error if the U256 amount cannot be converted to a `Quantity`.
55    pub fn as_quantity(&self) -> anyhow::Result<Quantity> {
56        Quantity::from_u256(self.amount, self.token.decimals).map_err(Into::into)
57    }
58
59    /// Sets the USD equivalent value for this token balance.
60    pub fn set_amount_usd(&mut self, amount_usd: Quantity) {
61        self.amount_usd = Some(amount_usd);
62    }
63}
64
65impl Display for TokenBalance {
66    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        let quantity = self.as_quantity().unwrap_or_default();
68
69        match &self.amount_usd {
70            Some(usd) => write!(
71                f,
72                "TokenBalance(token={}, amount={}, usd=${:.2})",
73                self.token.symbol,
74                quantity.as_decimal(),
75                usd.as_f64()
76            ),
77            None => write!(
78                f,
79                "TokenBalance(token={}, amount={})",
80                self.token.symbol,
81                quantity.as_decimal()
82            ),
83        }
84    }
85}
86
87/// Represents the complete balance state of a blockchain wallet.
88///
89/// Tracks both the native currency balance (e.g., ETH, ARB) and ERC-20 token
90/// balances for a wallet address. The `token_universe` defines which tokens
91/// should be tracked for balance fetching.
92#[derive(Debug)]
93pub struct WalletBalance {
94    /// The balance of the chain's native currency
95    pub native_currency: Option<Money>,
96    /// Collection of ERC-20 token balances held in the wallet.
97    pub token_balances: Vec<TokenBalance>,
98    /// Set of token addresses to track for balance updates.
99    pub token_universe: HashSet<Address>,
100}
101
102impl WalletBalance {
103    /// Creates a new [`WalletBalance`] with the specified token universe.
104    #[must_use]
105    pub const fn new(token_universe: HashSet<Address>) -> Self {
106        Self {
107            native_currency: None,
108            token_balances: vec![],
109            token_universe,
110        }
111    }
112
113    /// Returns `true` if the token universe has been initialized with token addresses.
114    #[must_use]
115    pub fn is_token_universe_initialized(&self) -> bool {
116        !self.token_universe.is_empty()
117    }
118
119    /// Sets the native currency balance for the wallet.
120    pub fn set_native_currency_balance(&mut self, balance: Money) {
121        self.native_currency = Some(balance);
122    }
123
124    /// Adds an ERC-20 token balance to the wallet.
125    pub fn add_token_balance(&mut self, token_balance: TokenBalance) {
126        self.token_balances.push(token_balance);
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use std::sync::Arc;
133
134    use alloy_primitives::{U256, address};
135    use rstest::rstest;
136
137    use super::*;
138    use crate::defi::{
139        SharedChain, Token,
140        chain::chains,
141        stubs::{arbitrum, usdc, weth},
142    };
143
144    // Helper to create a token with specific decimals
145    fn create_token(symbol: &str, decimals: u8) -> Token {
146        Token::new(
147            Arc::new(chains::ETHEREUM.clone()),
148            address!("0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48"),
149            format!("{symbol} Token"),
150            symbol.to_string(),
151            decimals,
152        )
153    }
154
155    #[rstest]
156    fn test_token_balance_as_quantity_18_decimals(#[from(arbitrum)] chain: SharedChain) {
157        // Test case: NU token with 18 decimals
158        // Raw amount: 10342000000000000000000 (10342 * 10^18)
159        // Expected: 10342.000000000000000000
160        let token = Token::new(
161            chain,
162            address!("0x4fE83213D56308330EC302a8BD641f1d0113A4Cc"),
163            "NuCypher".to_string(),
164            "NU".to_string(),
165            18,
166        );
167        let amount = U256::from(10342u64) * U256::from(10u64).pow(U256::from(18u64));
168        let balance = TokenBalance::new(amount, token);
169
170        let quantity = balance.as_quantity().unwrap();
171        assert_eq!(
172            quantity.as_decimal().to_string(),
173            "10342.000000000000000000"
174        );
175    }
176
177    #[rstest]
178    fn test_token_balance_as_quantity_6_decimals() {
179        // Test case: USDC with 6 decimals
180        // Raw amount: 92220728254 (92220.728254 * 10^6)
181        // Expected: 92220.728254
182        let token = create_token("USDC", 6);
183        let amount = U256::from(92_220_728_254_u64);
184        let balance = TokenBalance::new(amount, token);
185
186        let quantity = balance.as_quantity().unwrap();
187        assert_eq!(quantity.as_decimal().to_string(), "92220.728254");
188    }
189
190    #[rstest]
191    fn test_token_balance_as_quantity_fractional_18_decimals(#[from(arbitrum)] chain: SharedChain) {
192        // Test case: mETH with 18 decimals and fractional amount
193        // Raw amount: 758325512078001391
194        // Expected: 0.758325512078001391
195        let token = Token::new(
196            chain,
197            address!("0xd5F7838F5C461fefF7FE49ea5ebaF7728bB0ADfa"),
198            "mETH".to_string(),
199            "mETH".to_string(),
200            18,
201        );
202        let amount = U256::from(758_325_512_078_001_391_u64);
203        let balance = TokenBalance::new(amount, token);
204
205        let quantity = balance.as_quantity().unwrap();
206        assert_eq!(quantity.as_decimal().to_string(), "0.758325512078001391");
207    }
208
209    #[rstest]
210    fn test_token_balance_display_18_decimals(#[from(arbitrum)] chain: SharedChain) {
211        // Test Display implementation with 18 decimal token
212        let token = Token::new(
213            chain,
214            address!("0x912CE59144191C1204E64559FE8253a0e49E6548"),
215            "Arbitrum".to_string(),
216            "ARB".to_string(),
217            18,
218        );
219        // 7922.013795343949480329 ARB
220        let amount = U256::from_str_radix("7922013795343949480329", 10).unwrap();
221        let balance = TokenBalance::new(amount, token);
222
223        let display = balance.to_string();
224        assert!(display.contains("ARB"));
225        assert!(display.contains("7922.013795343949480329"));
226    }
227
228    #[rstest]
229    fn test_token_balance_display_6_decimals() {
230        // Test Display implementation with 6 decimal token (USDC)
231        let token = create_token("USDC", 6);
232        let amount = U256::from(92_220_728_254_u64); // 92220.728254 USDC
233        let balance = TokenBalance::new(amount, token);
234
235        let display = balance.to_string();
236        assert!(display.contains("USDC"));
237        assert!(display.contains("92220.728254"));
238    }
239
240    #[rstest]
241    fn test_token_balance_set_amount_usd(weth: Token) {
242        let amount = U256::from(1u64) * U256::from(10u64).pow(U256::from(18u64));
243        let mut balance = TokenBalance::new(amount, weth);
244
245        assert!(balance.amount_usd.is_none());
246
247        let usd_value = Quantity::from("3500.00");
248        balance.set_amount_usd(usd_value);
249
250        assert!(balance.amount_usd.is_some());
251        assert_eq!(
252            balance.amount_usd.unwrap().as_decimal().to_string(),
253            "3500.00"
254        );
255    }
256
257    #[rstest]
258    fn test_wallet_balance_new_empty() {
259        let wallet = WalletBalance::new(HashSet::new());
260
261        assert!(wallet.native_currency.is_none());
262        assert!(wallet.token_balances.is_empty());
263        assert!(!wallet.is_token_universe_initialized());
264    }
265
266    #[rstest]
267    fn test_wallet_balance_with_token_universe() {
268        let mut tokens = HashSet::new();
269        tokens.insert(address!("0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48")); // USDC
270        tokens.insert(address!("0x912CE59144191C1204E64559FE8253a0e49E6548")); // ARB
271
272        let wallet = WalletBalance::new(tokens);
273
274        assert!(wallet.is_token_universe_initialized());
275        assert_eq!(wallet.token_universe.len(), 2);
276    }
277
278    #[rstest]
279    fn test_wallet_balance_set_native_currency() {
280        let mut wallet = WalletBalance::new(HashSet::new());
281
282        assert!(wallet.native_currency.is_none());
283
284        let eth_balance = Money::new(50.936_054, crate::types::Currency::ETH());
285        wallet.set_native_currency_balance(eth_balance);
286
287        assert!(wallet.native_currency.is_some());
288    }
289
290    #[rstest]
291    fn test_wallet_balance_add_token_balance(usdc: Token, weth: Token) {
292        let mut wallet = WalletBalance::new(HashSet::new());
293
294        let usdc_balance = TokenBalance::new(U256::from(100_000_000u64), usdc); // 100 USDC
295        let weth_balance = TokenBalance::new(U256::from(10u64).pow(U256::from(18u64)), weth); // 1 WETH
296
297        wallet.add_token_balance(usdc_balance);
298        wallet.add_token_balance(weth_balance);
299
300        assert_eq!(wallet.token_balances.len(), 2);
301        assert_eq!(wallet.token_balances[0].token.symbol, "USDC");
302        assert_eq!(wallet.token_balances[1].token.symbol, "WETH");
303    }
304}
nautilus_model/defi/wallet.rs

nautilus_model/defi/
wallet.rs
