diff --git a/apps/contracts/vault/src/test.rs b/apps/contracts/vault/src/test.rs
index d5af5203..f696f7ee 100755
--- a/apps/contracts/vault/src/test.rs
+++ b/apps/contracts/vault/src/test.rs
@@ -172,5 +172,5 @@ mod deposit;
 mod admin;
 mod invest;
 mod withdraw;
-// mod emergency_withdraw;
-// mod rebalance;
+mod emergency_withdraw;
+mod rebalance;
diff --git a/apps/contracts/vault/src/test/emergency_withdraw.rs b/apps/contracts/vault/src/test/emergency_withdraw.rs
index 2ea7ac00..1c014738 100644
--- a/apps/contracts/vault/src/test/emergency_withdraw.rs
+++ b/apps/contracts/vault/src/test/emergency_withdraw.rs
@@ -2,7 +2,10 @@ use soroban_sdk::{vec as sorobanvec, String, Vec};
 
 use crate::test::{
     create_strategy_params_token0,
-    defindex_vault::{AssetStrategySet, Investment},
+    defindex_vault::{AssetStrategySet,
+    AssetInvestmentAllocation,  
+    StrategyInvestment
+    },
     DeFindexVaultTest,
 };
 
@@ -30,7 +33,7 @@ fn test_emergency_withdraw_success() {
         &String::from_str(&test.env, "dfToken"),
         &String::from_str(&test.env, "DFT"),
     );
-    let amount = 1000i128;
+    let amount = 987654321i128;
 
     let users = DeFindexVaultTest::generate_random_users(&test.env, 1);
 
@@ -49,27 +52,34 @@ fn test_emergency_withdraw_success() {
     );
 
     let df_balance = test.defindex_contract.balance(&users[0]);
-    assert_eq!(df_balance, amount);
+    assert_eq!(df_balance, amount - 1000);   
 
     // Balance of the token0 on the vault should be `amount` since it is deposited into the vault first
     let vault_balance_of_token = test.token0.balance(&test.defindex_contract.address);
     assert_eq!(vault_balance_of_token, amount);
-
-    // Should invest the funds
+    
     let investments = sorobanvec![
         &test.env,
-        Investment {
-            amount: amount.clone(),
-            strategy: strategy_params_token0.first().unwrap().address.clone()
-        }
+        Some(AssetInvestmentAllocation {
+            asset: test.token0.address.clone(),
+            strategy_investments: sorobanvec![
+                &test.env,
+                Some(StrategyInvestment {
+                    strategy: test.strategy_client_token0.address.clone(),
+                    amount: amount,
+                }),
+            ],
+        }),
     ];
+
+
     test.defindex_contract.invest(&investments);
 
     // Balance of the token0 on the vault should be 0
     let vault_balance_of_token = test.token0.balance(&test.defindex_contract.address);
     assert_eq!(vault_balance_of_token, 0);
 
-    // Balance of the strategy should be `amount`
+    // Balance of the vault on the strategy contract should be `amount`
     let strategy_balance = test
         .strategy_client_token0
         .balance(&test.defindex_contract.address);
@@ -80,7 +90,7 @@ fn test_emergency_withdraw_success() {
         &test.emergency_manager,
     );
 
-    // Balance of the strategy should be 0
+    // Balance of the vault on the strategy should be 0
     let strategy_balance = test
         .strategy_client_token0
         .balance(&test.defindex_contract.address);
diff --git a/apps/contracts/vault/src/test/rebalance.rs b/apps/contracts/vault/src/test/rebalance.rs
index 7fd734ca..dffb1e7a 100644
--- a/apps/contracts/vault/src/test/rebalance.rs
+++ b/apps/contracts/vault/src/test/rebalance.rs
@@ -3,7 +3,12 @@ use soroban_sdk::{vec as sorobanvec, String, Vec};
 use crate::test::{
     create_strategy_params_token0,
     defindex_vault::{
-        ActionType, AssetStrategySet, Instruction, Investment, OptionalSwapDetailsExactIn,
+        ActionType, 
+        AssetStrategySet, 
+        Instruction, 
+        AssetInvestmentAllocation,  
+        StrategyInvestment,
+        OptionalSwapDetailsExactIn,
         OptionalSwapDetailsExactOut,
     },
     DeFindexVaultTest,
@@ -33,7 +38,7 @@ fn rebalance() {
         &String::from_str(&test.env, "dfToken"),
         &String::from_str(&test.env, "DFT"),
     );
-    let amount = 1000i128;
+    let amount = 987654321i128;
 
     let users = DeFindexVaultTest::generate_random_users(&test.env, 1);
 
@@ -51,14 +56,20 @@ fn rebalance() {
     );
 
     let df_balance = test.defindex_contract.balance(&users[0]);
-    assert_eq!(df_balance, amount);
+    assert_eq!(df_balance, amount - 1000);
 
     let investments = sorobanvec![
         &test.env,
-        Investment {
-            amount: amount,
-            strategy: test.strategy_client_token0.address.clone()
-        }
+        Some(AssetInvestmentAllocation {
+            asset: test.token0.address.clone(),
+            strategy_investments: sorobanvec![
+                &test.env,
+                Some(StrategyInvestment {
+                    strategy: test.strategy_client_token0.address.clone(),
+                    amount: amount,
+                }),
+            ],
+        }),
     ];
 
     test.defindex_contract.invest(&investments);