diff --git a/bedrock/tests/test_smart_account_usd_vault.rs b/bedrock/tests/test_smart_account_usd_vault.rs index 850bb2c7..7c8f57fc 100644 --- a/bedrock/tests/test_smart_account_usd_vault.rs +++ b/bedrock/tests/test_smart_account_usd_vault.rs @@ -3,13 +3,15 @@ use std::sync::Arc; mod common; use alloy::{ + network::Ethereum, primitives::{ utils::{parse_ether, parse_units}, Address, U256, }, - providers::{ext::AnvilApi, ProviderBuilder}, + providers::{ext::AnvilApi, Provider, ProviderBuilder}, signers::local::PrivateKeySigner, }; +use anyhow::Context; use common::{deploy_safe, set_erc20_balance_for_safe, setup_anvil, IERC20}; use std::str::FromStr; @@ -25,17 +27,6 @@ use crate::common::{ #[tokio::test] async fn test_usd_vault_migration() -> anyhow::Result<()> { - let usdc_address = - Address::from_str("0x79A02482A880bCE3F13e09Da970dC34db4CD24d1").unwrap(); - let sdai_address = - Address::from_str("0x859DBE24b90C9f2f7742083d3cf59cA41f55Be5d").unwrap(); - let usd_vault_address = - Address::from_str("0x6F1D98034D3055684F989f3Ac9832eC37B3F22EC").unwrap(); - let morpho_vault_address = - Address::from_str("0xb1E80387EbE53Ff75a89736097D34dC8D9E9045B").unwrap(); - let bad_morpho_vault_address = - Address::from_str("0x348831b46876d3dF2Db98BdEc5E3B4083329Ab9f").unwrap(); - let owner_signer = PrivateKeySigner::random(); let owner_key_hex = hex::encode(owner_signer.to_bytes()); let owner = owner_signer.address(); @@ -47,10 +38,6 @@ async fn test_usd_vault_migration() -> anyhow::Result<()> { let client = AnvilBackedHttpClient::new(provider.clone()); set_http_client(Arc::new(client)); - let usdc = IERC20::new(usdc_address, &provider); - let sdai = IERC20::new(sdai_address, &provider); - let morpho_vault = IERC20::new(morpho_vault_address, &provider); - let safe_address = deploy_safe(&provider, owner, U256::ZERO).await?; println!("✓ Deployed Safe at: {safe_address}"); @@ -69,8 +56,55 @@ async fn test_usd_vault_migration() -> anyhow::Result<()> { .await?; println!("✓ Set Safe as verified until far future"); - set_erc20_balance_with_slot( + run_usd_vault_migration( &provider, + &safe_account, + safe_address, + "no-limit USDVault", + Address::from_str("0x6F1D98034D3055684F989f3Ac9832eC37B3F22EC").unwrap(), + false, + ) + .await?; + + run_usd_vault_migration( + &provider, + &safe_account, + safe_address, + "main USDVault", + Address::from_str("0xB0e31149c03F1300BD9fF8C165B1fa38fDA2F0bB").unwrap(), + true, + ) + .await?; + + Ok(()) +} + +async fn run_usd_vault_migration

( + provider: &P, + safe_account: &SafeSmartAccount, + safe_address: Address, + vault_name: &str, + usd_vault_address: Address, + limit_withdrawals_to_deposits: bool, +) -> anyhow::Result<()> +where + P: Provider + AnvilApi, +{ + let usdc_address = + Address::from_str("0x79A02482A880bCE3F13e09Da970dC34db4CD24d1").unwrap(); + let sdai_address = + Address::from_str("0x859DBE24b90C9f2f7742083d3cf59cA41f55Be5d").unwrap(); + let morpho_vault_address = + Address::from_str("0xb1E80387EbE53Ff75a89736097D34dC8D9E9045B").unwrap(); + let bad_morpho_vault_address = + Address::from_str("0x348831b46876d3dF2Db98BdEc5E3B4083329Ab9f").unwrap(); + + let usdc = IERC20::new(usdc_address, provider); + let sdai = IERC20::new(sdai_address, provider); + let morpho_vault = IERC20::new(morpho_vault_address, provider); + + set_erc20_balance_with_slot( + provider, usdc_address, usd_vault_address, parse_units("10000", 6).unwrap().into(), @@ -78,7 +112,7 @@ async fn test_usd_vault_migration() -> anyhow::Result<()> { ) .await?; set_erc20_balance_with_slot( - &provider, + provider, sdai_address, usd_vault_address, parse_units("10000", 18).unwrap().into(), @@ -89,8 +123,8 @@ async fn test_usd_vault_migration() -> anyhow::Result<()> { let vault_usdc_balance = usdc.balanceOf(usd_vault_address).call().await?; let vault_sdai_balance = sdai.balanceOf(usd_vault_address).call().await?; - println!("USDVault USDC balance: {vault_usdc_balance}"); - println!("USDVault sDAI balance: {vault_sdai_balance}"); + println!("{vault_name} USDC balance: {vault_usdc_balance}"); + println!("{vault_name} sDAI balance: {vault_sdai_balance}"); let sdai_amount: U256 = parse_units("10", 18).unwrap().into(); @@ -112,11 +146,21 @@ async fn test_usd_vault_migration() -> anyhow::Result<()> { "Expected error message to contain 'Cannot migrate with zero sDAI balance', got: {}", error_message ); - println!("✓ Migration correctly failed with zero sDAI balance error"); + println!("✓ {vault_name} migration correctly failed with zero sDAI balance error"); // Now set up sDAI balance for actual migration test - set_erc20_balance_for_safe(&provider, sdai_address, safe_address, sdai_amount) + set_erc20_balance_for_safe(provider, sdai_address, safe_address, sdai_amount) + .await?; + if limit_withdrawals_to_deposits { + set_erc20_balance_with_slot( + provider, + usd_vault_address, + safe_address, + sdai_amount, + U256::from(1), // USDVault.sDAIBalances is at slot 1 + ) .await?; + } let sdai_balance_before = sdai.balanceOf(safe_address).call().await?; println!("sDAI balance before migration: {sdai_balance_before}"); @@ -142,7 +186,9 @@ async fn test_usd_vault_migration() -> anyhow::Result<()> { "Expected error message to contain 'Asset address mismatch between USDVault and ERC-4626 Vault', got: {}", error_message ); - println!("✓ Migration correctly failed with asset address mismatch error"); + println!( + "✓ {vault_name} migration correctly failed with asset address mismatch error" + ); // Now perform successful migration safe_account @@ -151,8 +197,8 @@ async fn test_usd_vault_migration() -> anyhow::Result<()> { &morpho_vault_address.to_string(), ) .await - .expect("USDVault migration failed"); - println!("✓ Migrated USDVault to MorphoVault"); + .with_context(|| format!("{vault_name} migration failed"))?; + println!("✓ Migrated {vault_name} to MorphoVault"); let sdai_balance_after = sdai.balanceOf(safe_address).call().await?; println!("sDAI balance after migration: {sdai_balance_after}");