feat: Add a vault pause mechanism, refactor total

Author: oyinade247

contracts/vault/src/lib.rs

@@ -175,7 +175,6 @@ impl YieldVault {
 
     /// Read the total underlying assets represented by the vault.
     pub fn total_assets(env: Env) -> i128 {
-        Self::get_state(&env).total_assets
         let idle_assets = env.storage().instance().get::<_, i128>(&DataKey::TotalAssets).unwrap_or(0);
 
         let strategy_assets = if let Some(strategy_addr) = Self::strategy(env.clone()) {
@@ -630,8 +629,6 @@ impl YieldVault {
         // Transfer assets from user to vault
         token_client.transfer(&user, &env.current_contract_address(), &amount);
 
-        token_client.transfer(&user, &env.current_contract_address(), &amount);
-
         // Update state
         let ta = Self::total_assets(env.clone());
         env.storage()
@@ -711,14 +708,10 @@ impl YieldVault {
             .set(&DataKey::TotalShares, &Self::checked_sub(ts, shares)?);
 
         env.storage().instance().set(
-            &DataKey::ShareBalance(user.clone()),
+            &user_key,
             &Self::checked_sub(user_shares, shares)?,
         );
 
-        env.storage()
-            .instance()
-            .set(&user_key, &(user_shares - shares));
-
         env.events().publish(
             (symbol_short!("withdraw"), user),
             (assets_to_return, shares),
@@ -748,6 +741,7 @@ impl YieldVault {
 
         // Update idle assets
         env.storage().instance().set(&DataKey::TotalAssets, &(idle_ta - amount));
+        Ok(())
     }
 
     /// Recall funds from the strategy.
@@ -784,6 +778,8 @@ impl YieldVault {
             .instance()
             .get::<_, i128>(&DataKey::TotalAssets)
             .unwrap_or(0);
+        // Update total assets state
+        let ta = env.storage().instance().get::<_, i128>(&DataKey::TotalAssets).unwrap_or(0);
         env.storage().instance().set(&DataKey::TotalAssets, &(ta + amount));
 
         let mut state = Self::get_state(&env);
@@ -796,12 +792,7 @@ impl YieldVault {
         );
     }
 
-    /// Legacy admin function retained for compatibility.
-    pub fn accrue_yield(env: Env, amount: i128) {
-        Self::distribute_yield(env, amount);
-    }
-
-    pub fn report_benji_yield(env: Env, strategy: Address, amount: i128) {
+    pub fn report_benji_yield(env: Env, strategy: Address, amount: i128) -> Result<(), VaultError> {
         strategy.require_auth();
         if amount <= 0 {
             return Err(VaultError::InvalidAmount);
@@ -810,7 +801,7 @@ impl YieldVault {
         let token_addr = Self::token(env.clone());
         let token_client = token::Client::new(&env, &token_addr);
 
-        token_client.transfer(&admin, &env.current_contract_address(), &amount);
+        token_client.transfer(&strategy, &env.current_contract_address(), &amount);
 
         let ta = Self::total_assets(env.clone());
         env.storage()

contracts/vault/src/test.rs

@@ -767,3 +767,49 @@ fn test_multiple_deposits_atomic_state_updates() {
     assert_eq!(vault.total_shares(), 200);
     assert_eq!(vault.total_assets(), 200);
 }
+
+// ─── Pause Mechanism Tests ───────────────────────────────────────────────
+
+#[test]
+fn test_pause_mechanism() {
+    let env = Env::default();
+    env.mock_all_auths();
+
+    let admin = Address::generate(&env);
+    let user = Address::generate(&env);
+    let token_admin = Address::generate(&env);
+    let usdc = create_token_contract(&env, &token_admin);
+    let usdc_admin_client = token::StellarAssetClient::new(&env, &usdc.address);
+    usdc_admin_client.mint(&user, &1000);
+
+    let vault_id = env.register(YieldVault, ());
+    let vault = YieldVaultClient::new(&env, &vault_id);
+    vault.initialize(&admin, &usdc.address);
+
+    // 1. Initially NOT paused
+    assert_eq!(vault.is_paused(), false);
+    vault.deposit(&user, &100);
+    assert_eq!(vault.balance(&user), 100);
+
+    // 2. Pause the vault
+    vault.set_pause(&true);
+    assert_eq!(vault.is_paused(), true);
+
+    // 3. Verify deposit reverts when paused
+    let deposit_result = vault.try_deposit(&user, &100);
+    assert_eq!(deposit_result, Err(Ok(VaultError::ContractPaused)));
+
+    // 4. Verify withdraw reverts when paused
+    let withdraw_result = vault.try_withdraw(&user, &50);
+    assert_eq!(withdraw_result, Err(Ok(VaultError::ContractPaused)));
+
+    // 5. Unpause the vault
+    vault.set_pause(&false);
+    assert_eq!(vault.is_paused(), false);
+
+    // 6. Verify operations resume
+    vault.deposit(&user, &100);
+    assert_eq!(vault.balance(&user), 200);
+    vault.withdraw(&user, &50);
+    assert_eq!(vault.balance(&user), 150);
+}