diff --git a/packages/vault-client/src/tests/test_rebalance_integration.rs b/packages/vault-client/src/tests/test_rebalance_integration.rs new file mode 100644 index 0000000..a44371c --- /dev/null +++ b/packages/vault-client/src/tests/test_rebalance_integration.rs @@ -0,0 +1,88 @@ +/// A partial `withdraw()` call that happens to drain the entire Blend +/// position (i.e. the amount pulled from Blend exactly matches the +/// outstanding Blend balance) must update CurrentProtocol to "none", +/// even though the user only requested part of their own shares. +/// +/// This is distinct from `test_integration_withdraw_all_updates_current_protocol_to_none`: +/// it exercises the `withdraw` entry point (not `withdraw_all`), which goes +/// through a different reconciliation path before delegating to +/// `withdraw_from_blend`. +#[test] +fn test_integration_withdraw_updates_current_protocol_to_none() { + let env = Env::default(); + env.mock_all_auths(); + + let (contract_id, _agent, owner, usdc_token, blend_pool) = + setup_vault_with_token_and_blend(&env); + let client = NeuroWealthVaultClient::new(&env, &contract_id); + let blend_client = MockBlendPoolClient::new(&env, &blend_pool); + + client.set_blend_pool(&owner, &blend_pool); + + let deposit_amount = 25_000_000_i128; + let user = Address::generate(&env); + mint_and_deposit(&env, &client, &usdc_token, &user, deposit_amount); + + // Move all funds into Blend so the vault itself holds zero idle USDC. + client.rebalance(&symbol_short!("blend"), &800_i128, &0_i128); + assert_eq!(client.get_current_protocol(), symbol_short!("blend")); + assert_eq!(blend_client.supplied(&usdc_token), deposit_amount); + assert_eq!(vault_usdc_balance(&env, &usdc_token, &contract_id), 0); + + // User withdraws their entire entitled amount via `withdraw` (not + // `withdraw_all`). Since the vault has zero idle balance, the full + // amount must be pulled from Blend, leaving exactly zero remaining. + client.withdraw(&user, &deposit_amount); + + // Vault should have pulled everything out of Blend. + assert_eq!(blend_client.supplied(&usdc_token), 0); + + // CurrentProtocol must reset to "none" as part of the same withdrawal, + // with no separate rebalance call required. + assert_eq!( + client.get_current_protocol(), + symbol_short!("none"), + "CurrentProtocol should switch to 'none' when a partial withdraw() fully drains Blend" + ); +} + +/// A `withdraw()` that only partially drains the Blend position must leave +/// CurrentProtocol untouched ("blend"), since funds are still deployed. +/// This guards against an overly-eager reset that flips CurrentProtocol to +/// "none" any time `withdraw_from_blend` is invoked, regardless of whether +/// the position was actually fully exited. +#[test] +fn test_integration_withdraw_partial_blend_drain_keeps_current_protocol() { + let env = Env::default(); + env.mock_all_auths(); + + let (contract_id, _agent, owner, usdc_token, blend_pool) = + setup_vault_with_token_and_blend(&env); + let client = NeuroWealthVaultClient::new(&env, &contract_id); + let blend_client = MockBlendPoolClient::new(&env, &blend_pool); + + client.set_blend_pool(&owner, &blend_pool); + + let deposit_amount = 25_000_000_i128; + let user = Address::generate(&env); + mint_and_deposit(&env, &client, &usdc_token, &user, deposit_amount); + + // Move all funds into Blend so the vault itself holds zero idle USDC. + client.rebalance(&symbol_short!("blend"), &800_i128, &0_i128); + assert_eq!(client.get_current_protocol(), symbol_short!("blend")); + + // Withdraw less than the full position; Blend should still hold a + // residual balance afterward. + let partial_amount = 10_000_000_i128; + client.withdraw(&user, &partial_amount); + + assert!( + blend_client.supplied(&usdc_token) > 0, + "Blend should still hold a residual position after a partial withdrawal" + ); + assert_eq!( + client.get_current_protocol(), + symbol_short!("blend"), + "CurrentProtocol must remain 'blend' while funds are still deployed" + ); +} \ No newline at end of file