diff --git a/nssa/program_methods/guest/src/bin/amm.rs b/nssa/program_methods/guest/src/bin/amm.rs
index ef1a505..39f11fa 100644
--- a/nssa/program_methods/guest/src/bin/amm.rs
+++ b/nssa/program_methods/guest/src/bin/amm.rs
@@ -258,7 +258,7 @@ fn new_definition(
             definition_token_b_id,
             vault_a_addr: vault_a.account_id.clone(),
             vault_b_addr: vault_b.account_id.clone(),
-            liquidity_pool_id: pool_lp.account_id.clone(),
+            liquidity_pool_id: TokenHolding::parse(&pool_lp.account.data).unwrap().definition_id,
             liquidity_pool_cap: amount_a,
             reserve_a: amount_a,
             reserve_b: amount_b,
@@ -269,6 +269,7 @@ fn new_definition(
 
     let mut chained_call = Vec::new();
    
+    //Chain call for Token A (User_A -> Vault_A)
     let mut instruction: [u8;32] = [0; 32];
     instruction[0] = 1;      
     instruction[1..17].copy_from_slice(&amount_a.to_le_bytes());
@@ -280,6 +281,7 @@ fn new_definition(
             pre_states: vec![user_a.clone(), vault_a.clone()]
         };
         
+    //Chain call for Token B (User_B -> Vault_B)
     instruction[1..17].copy_from_slice(&amount_b.to_le_bytes());
     let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap();
 
@@ -289,6 +291,7 @@ fn new_definition(
             pre_states: vec![user_b.clone(), vault_b.clone()]
         };
 
+    //Chain call for LP (Pool_LP -> User_LP)
     instruction[1..17].copy_from_slice(&amount_a.to_le_bytes());
     let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap();
 
@@ -335,38 +338,32 @@ fn swap(
     let vault1_data = TokenHolding::parse(&vault1.account.data).unwrap();
     let vault2_data = TokenHolding::parse(&vault2.account.data).unwrap();
 
-    let mut vault_a = AccountWithMetadata::default();
-    let mut vault_b = AccountWithMetadata::default();
-
-    if vault1_data.definition_id == pool_def_data.definition_token_a_id {
-        vault_a = vault1.clone();
+    let vault_a = if vault1_data.definition_id == pool_def_data.definition_token_a_id {
+        vault1.clone()
     } else if vault2_data.definition_id == pool_def_data.definition_token_a_id {
-        vault_a = vault2.clone();
+        vault2.clone()
     } else {
         panic!("Vault A was not provided");
-    }
+    };
         
-    if vault1_data.definition_id == pool_def_data.definition_token_b_id {
-       vault_b = vault1.clone();
+    let vault_b = if vault1_data.definition_id == pool_def_data.definition_token_b_id {
+       vault1.clone()
     } else if vault2_data.definition_id == pool_def_data.definition_token_b_id {
-        vault_b = vault2.clone();
+        vault2.clone()
     } else {
         panic!("Vault B was not provided");
-    }
+    };
 
     // 1. Identify swap direction (a -> b or b -> a)
-    let mut deposit_a = 0;
-    let mut deposit_b = 0;
-    let a_to_b;
-    if token_id == pool_def_data.definition_token_a_id {
-        deposit_a = amount;
-        a_to_b = true;
-    } else if token_id == pool_def_data.definition_token_b_id {
-        deposit_b = amount;
-        a_to_b = false;
-    } else {
+    let a_to_b = if token_id == pool_def_data.definition_token_a_id {
+        true
+    } else if token_id == pool_def_data.definition_token_b_id { false }
+    else {
         panic!("Address is not a token type for the pool");
-    }
+    };
+
+    let deposit_a = if a_to_b { amount } else { 0 };
+    let deposit_b = if a_to_b { 0 } else { amount }; 
 
     // 2. fetch pool reserves
     //validates reserves is at least the vaults' balances
@@ -387,8 +384,10 @@ fn swap(
 
     // 4. Slippage check
     if a_to_b {
+        assert!(withdraw_b != 0);
         assert!(withdraw_a == 0); }
     else {
+        assert!(withdraw_a != 0);
         assert!(withdraw_b == 0); }
 
     // 5. Update pool account
@@ -409,57 +408,54 @@ fn swap(
 
     let mut chained_call = Vec::new();
     
-    let call_token_a = ChainedCall::default();
-    let call_token_b = ChainedCall::default();
-
-    if a_to_b {
+    let call_token_a = if a_to_b { 
         let mut instruction_data = [0; 23];
         instruction_data[0] = 1;
         instruction_data[1..17].copy_from_slice(&deposit_a.to_le_bytes());
         let instruction_data = risc0_zkvm::serde::to_vec(&instruction_data).unwrap();
 
-        let call_token_a = ChainedCall{
+        ChainedCall{
                 program_id: pool_def_data.token_program_id,
                 instruction_data: instruction_data,
                 pre_states: vec![user_a.clone(), vault_a.clone()]
-            };
-
-        let mut instruction_data = [0; 23];
-        instruction_data[0] = 1;
-        instruction_data[1..17].copy_from_slice(&withdraw_b.to_le_bytes());
-        let instruction_data = risc0_zkvm::serde::to_vec(&instruction_data).unwrap();
-
-        let call_token_b = ChainedCall{
-                program_id: pool_def_data.token_program_id,
-                instruction_data: instruction_data,
-                pre_states: vec![user_b.clone(), vault_b.clone()]
-            };
-     
+            }
     } else {
         let mut instruction_data = [0; 23];
         instruction_data[0] = 1;
         instruction_data[1..17].copy_from_slice(&withdraw_a.to_le_bytes());
         let instruction_data = risc0_zkvm::serde::to_vec(&instruction_data).unwrap();
 
-        let call_token_a = ChainedCall{
+        ChainedCall{
                 program_id: pool_def_data.token_program_id,
                 instruction_data: instruction_data,
                 pre_states: vec![vault_a.clone(), user_a.clone()]
-            };
-
+            }
+    };
+    
+    let call_token_b = if a_to_b {
+        let mut instruction_data = [0; 23];
+        instruction_data[0] = 1;
+        instruction_data[1..17].copy_from_slice(&withdraw_b.to_le_bytes());
+        let instruction_data = risc0_zkvm::serde::to_vec(&instruction_data).unwrap();
+        ChainedCall{
+                program_id: pool_def_data.token_program_id,
+                instruction_data: instruction_data,
+                pre_states: vec![vault_b.clone(), user_b.clone()]
+            }
+    } else {
         let mut instruction_data = [0; 23];
         instruction_data[0] = 1;
         instruction_data[1..17].copy_from_slice(&deposit_b.to_le_bytes());
         let instruction_data = risc0_zkvm::serde::to_vec(&instruction_data).unwrap();
 
-        let call_token_b = ChainedCall{
+        ChainedCall{
                 program_id: pool_def_data.token_program_id,
                 instruction_data: instruction_data,
-                pre_states: vec![vault_b.clone(), user_b.clone()]
-            };
-    }
-    
-    
+                pre_states: vec![user_b.clone(), vault_b.clone()]
+            }
+
+    };
+
     chained_call.push(call_token_a);
     chained_call.push(call_token_b);
     
@@ -523,10 +519,7 @@ fn add_liquidity(pre_states: &[AccountWithMetadata],
         panic!("Both max-balances must be nonzero");
     }
     
-    // 2. Determine deposit amounts
-    let mut actual_amount_a = 0;
-    let mut actual_amount_b = 0;
-
+    // 2. Determine deposit amount
     let vault_b_balance = TokenHolding::parse(&vault_b.account.data).unwrap().balance;
     let vault_a_balance = TokenHolding::parse(&vault_a.account.data).unwrap().balance;
     if vault_a_balance == 0 || vault_b_balance == 0 {
@@ -537,15 +530,22 @@ fn add_liquidity(pre_states: &[AccountWithMetadata],
         panic!("Reserves must be nonzero");
     }
 
-    if main_token == pool_def_data.definition_token_a_id {
-        actual_amount_a += max_amount_a;
-        actual_amount_b += (pool_def_data.reserve_b*actual_amount_a)/pool_def_data.reserve_a;
+    //Calculate actual_amounts
+    let actual_amount_a = if main_token == pool_def_data.definition_token_a_id {
+        max_amount_a
     } else if main_token == pool_def_data.definition_token_b_id {
-        actual_amount_b += max_amount_b;
-        actual_amount_a += (pool_def_data.reserve_a*actual_amount_b)/pool_def_data.reserve_b;
+        (pool_def_data.reserve_a*max_amount_b)/pool_def_data.reserve_b
     } else {
         panic!("Mismatch of token types"); //main token does not match with vaults.
-    }
+    };
+
+    let actual_amount_b = if main_token == pool_def_data.definition_token_a_id {
+        (pool_def_data.reserve_b*max_amount_a)/pool_def_data.reserve_a
+    } else if main_token == pool_def_data.definition_token_b_id {
+        max_amount_b
+    } else {
+        panic!("Mismatch of token types"); //main token does not match with vaults.
+    };
 
     // 3. Validate amounts
     let user_a_balance = TokenHolding::parse(&user_a.account.data).unwrap().balance;
@@ -583,6 +583,7 @@ fn add_liquidity(pre_states: &[AccountWithMetadata],
     pool_post.data = pool_post_definition.into_data();
     let mut chained_call = Vec::new();
 
+    // Chain call for Token A (User_A -> Vault_A)
     let mut instruction_data = [0; 23];
     instruction_data[0] = 1;
     instruction_data[1..17].copy_from_slice(&actual_amount_a.to_le_bytes());
@@ -593,6 +594,7 @@ fn add_liquidity(pre_states: &[AccountWithMetadata],
             pre_states: vec![user_a.clone(), vault_a]
         };
 
+    // Chain call for Token B (User_B -> Vault_B)        
     let mut instruction_data = [0; 23];
     instruction_data[0] = 1;
     instruction_data[1..17].copy_from_slice(&actual_amount_b.to_le_bytes());
@@ -602,7 +604,8 @@ fn add_liquidity(pre_states: &[AccountWithMetadata],
             instruction_data: instruction_data,
             pre_states: vec![user_b.clone(), vault_b]
         };
-        
+
+    // Chain call for LP (User_LP -> Pool_LP)    
     let mut instruction_data = [0; 23];
     instruction_data[0] = 1;
     instruction_data[1..17].copy_from_slice(&delta_lp.to_le_bytes());
@@ -702,6 +705,7 @@ fn remove_liquidity(pre_states: &[AccountWithMetadata]) -> (Vec<Account>, Vec<Ch
     let mut instruction_data = [0; 23];
     instruction_data[0] = 1;
 
+    //Chaincall for Token A withdraw
     let mut instruction: [u8;32] = [0; 32];
     instruction[0] = 1;      
     instruction[1..17].copy_from_slice(&withdraw_amount_a.to_le_bytes());
@@ -712,6 +716,7 @@ fn remove_liquidity(pre_states: &[AccountWithMetadata]) -> (Vec<Account>, Vec<Ch
             pre_states: vec![vault_a, user_a.clone()]
         };
 
+    //Chaincall for Token B withdraw
     let mut instruction: [u8;32] = [0; 32];
     instruction[0] = 1;      
     instruction[1..17].copy_from_slice(&withdraw_amount_b.to_le_bytes());
@@ -722,6 +727,7 @@ fn remove_liquidity(pre_states: &[AccountWithMetadata]) -> (Vec<Account>, Vec<Ch
             pre_states: vec![vault_b, user_b.clone()]
         };
 
+    //Chaincall for LP adjustment        
     let mut instruction: [u8;32] = [0; 32];
     instruction[0] = 1;      
     instruction[1..17].copy_from_slice(&delta_lp.to_le_bytes());
@@ -736,7 +742,6 @@ fn remove_liquidity(pre_states: &[AccountWithMetadata]) -> (Vec<Account>, Vec<Ch
     chained_call.push(call_token_b);
     chained_call.push(call_token_a);
 
-
     let post_states = vec!
         [pool_post.clone(), 
         pre_states[1].account.clone(),
@@ -746,13 +751,12 @@ fn remove_liquidity(pre_states: &[AccountWithMetadata]) -> (Vec<Account>, Vec<Ch
         pre_states[5].account.clone(),
         pre_states[6].account.clone()];
 
-    (post_states.clone(), chained_call)
-
+    (post_states, chained_call)
 }
 
 #[cfg(test)]
 mod tests {
-    use nssa_core::{account::{Account, AccountId, AccountWithMetadata}, address::Address};
+    use nssa_core::{account::{Account, AccountId, AccountWithMetadata}, program::ChainedCall};
 
     use crate::{PoolDefinition, TOKEN_HOLDING_TYPE, TokenHolding, add_liquidity, new_definition, remove_liquidity, swap};
 
@@ -1193,20 +1197,21 @@ mod tests {
         let mut user_b = Account::default();
         let mut pool_lp = Account::default();
 
+        let definition_token_a_id = AccountId::new([1;32]);
+        let definition_token_b_id = AccountId::new([2;32]); 
+
         user_a.data = TokenHolding::into_data(
-            TokenHolding{
-                account_type: 0u8,
-                definition_id: AccountId::new([1;32]),
-                balance: 200u128,
-            }
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_a_id.clone(),
+                balance: 15u128 }
         );
+
         user_b.data = TokenHolding::into_data(
-            TokenHolding{
-                account_type: 0u8,
-                definition_id: AccountId::new([1;32]),
-                balance: 100u128,
-            }
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_a_id.clone(),
+                balance: 15u128 }
         );
+
         
         vault_a.data = vec![
                 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
@@ -1261,28 +1266,30 @@ mod tests {
     }
 
     #[test]
-    fn test_call_new_definition_success() {
+    fn test_call_new_definition_chain_call_success() {
         let mut pool = Account::default();
         let mut vault_a = Account::default();
         let mut vault_b = Account::default();
         let mut user_a = Account::default();
         let mut user_b = Account::default();
         let mut pool_lp = Account::default();
+        let mut user_lp = Account::default();
+
+        let definition_token_a_id = AccountId::new([1;32]);
+        let definition_token_b_id = AccountId::new([2;32]); 
 
         user_a.data = TokenHolding::into_data(
-            TokenHolding{
-                account_type: 0u8,
-                definition_id: AccountId::new([1;32]),
-                balance: 200u128,
-            }
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_a_id.clone(),
+                balance: 15u128 }
         );
+
         user_b.data = TokenHolding::into_data(
-            TokenHolding{
-                account_type: 0u8,
-                definition_id: AccountId::new([2;32]),
-                balance: 100u128,
-            }
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_b_id.clone(),
+                balance: 15u128 }
         );
+
         
         vault_a.data = vec![
                 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
@@ -1300,39 +1307,102 @@ mod tests {
             ];
 
 
+        let user_a = AccountWithMetadata {
+            account: user_a.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([5; 32])};
+    
+        let user_b = AccountWithMetadata {
+            account: user_b.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([6; 32])};
+
+        let user_lp = AccountWithMetadata {
+            account: user_lp.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([7; 32])
+        };
+
+        let vault_a = AccountWithMetadata {
+            account: vault_a.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([2; 32])};
+
+        let vault_b = AccountWithMetadata {
+            account: vault_b.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([3; 32])};
+
+        let pool_lp = AccountWithMetadata {
+            account: pool_lp.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([4; 32])};
+
         let pre_states = vec![AccountWithMetadata {
-            account: pool,
+            account: pool.clone(),
             is_authorized: true,
             account_id: AccountId::new([0; 32])},
-            AccountWithMetadata {
-            account: vault_a,
-            is_authorized: true,
-            account_id: AccountId::new([2; 32])},
-            AccountWithMetadata {
-            account: vault_b,
-            is_authorized: true,
-            account_id: AccountId::new([3; 32])},
-            AccountWithMetadata {
-            account: pool_lp,
-            is_authorized: true,
-            account_id: AccountId::new([4; 32])},
-            AccountWithMetadata {
-            account: user_a,
-            is_authorized: true,
-            account_id: AccountId::new([5; 32])},
-            AccountWithMetadata {
-            account: user_b,
-            is_authorized: true,
-            account_id: AccountId::new([6; 32])},
-            AccountWithMetadata {
-            account: Account::default(),
-            is_authorized: true,
-            account_id: AccountId::new([7; 32])},
+            vault_a.clone(),
+            vault_b.clone(),
+            pool_lp.clone(),
+            user_a.clone(),
+            user_b.clone(),
+            user_lp.clone(),
         ];
         let balance_a = 15u128;
         let balance_b = 15u128;
         let token_program_id: [u32;8] = [0; 8];
-        let _post_states = new_definition(&pre_states, &[balance_a, balance_b], token_program_id);
+        let (post_states, chained_calls) = new_definition(&pre_states, &[balance_a, balance_b], token_program_id);
+
+        let chained_call_lp = chained_calls[0].clone();
+        let chained_call_b = chained_calls[1].clone();
+        let chained_call_a = chained_calls[2].clone();
+        
+        //Expected chain_call for Token A
+        let mut instruction: [u8;32] = [0; 32];
+        instruction[0] = 1;      
+        instruction[1..17].copy_from_slice(&balance_a.to_le_bytes());
+        let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap();
+        let expected_chained_call_a = ChainedCall{
+            program_id: token_program_id,
+            instruction_data,
+            pre_states: vec![user_a.clone(), vault_a.clone()],
+        };
+
+        //Expected chain call for Token B
+        let mut instruction: [u8;32] = [0; 32];
+        instruction[0] = 1;      
+        instruction[1..17].copy_from_slice(&balance_b.to_le_bytes());
+        let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap();
+        let expected_chained_call_b = ChainedCall{
+            program_id: token_program_id,
+            instruction_data,
+            pre_states: vec![user_b.clone(), vault_b.clone()],
+        };
+
+        //Expected chain call for LP
+        let mut instruction: [u8;32] = [0; 32];
+        instruction[0] = 1;      
+        instruction[1..17].copy_from_slice(&balance_a.to_le_bytes());
+        let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap();
+        let expected_chained_call_lp = ChainedCall{
+            program_id: token_program_id,
+            instruction_data,
+            pre_states: vec![pool_lp.clone(), user_lp.clone()],
+        };
+        
+        assert!(chained_call_a.program_id == expected_chained_call_a.program_id);
+        assert!(chained_call_a.instruction_data == expected_chained_call_a.instruction_data);
+        assert!(chained_call_a.pre_states[0].account == expected_chained_call_a.pre_states[0].account);
+        assert!(chained_call_a.pre_states[1].account == expected_chained_call_a.pre_states[1].account);
+        assert!(chained_call_b.program_id == expected_chained_call_b.program_id);
+        assert!(chained_call_b.instruction_data == expected_chained_call_b.instruction_data);
+        assert!(chained_call_b.pre_states[0].account == expected_chained_call_b.pre_states[0].account);
+        assert!(chained_call_b.pre_states[1].account == expected_chained_call_b.pre_states[1].account);
+        assert!(chained_call_lp.program_id == expected_chained_call_lp.program_id);
+        assert!(chained_call_lp.instruction_data == expected_chained_call_lp.instruction_data);
+        assert!(chained_call_lp.pre_states[0].account == expected_chained_call_lp.pre_states[0].account);
+        assert!(chained_call_lp.pre_states[1].account == expected_chained_call_lp.pre_states[1].account);
     }
 
     #[should_panic(expected = "Invalid number of input accounts")]
@@ -1551,7 +1621,6 @@ mod tests {
         let mut vault2 = Account::default();
         let mut pool_lp = Account::default();
 
-
         let definition_token_a_id = AccountId::new([1;32]);
         let definition_token_b_id = AccountId::new([2;32]); 
 
@@ -1625,6 +1694,174 @@ mod tests {
         let _post_states = remove_liquidity(&pre_states);
     }
 
+    #[test]
+    fn test_call_remove_liquidity_chain_call_success() {
+        let mut pool = Account::default();
+        let mut vault_a = Account::default();
+        let mut vault_b = Account::default();
+        let mut user_a = Account::default();
+        let mut user_b = Account::default();
+        let mut pool_lp = Account::default();
+        let mut user_lp = Account::default();
+
+        let definition_token_a_id = AccountId::new([1;32]);
+        let definition_token_b_id = AccountId::new([2;32]); 
+        
+        vault_a.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_a_id.clone(),
+                balance: 15u128 }
+        );
+
+        vault_b.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_b_id.clone(),
+                balance: 15u128 }
+        );
+
+        pool_lp.data = vec![
+                1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+                1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+            ];
+
+        let vault_a_addr = AccountId::new([5;32]);
+        let vault_b_addr = AccountId::new([6;32]);
+        let liquidity_pool_id = AccountId::new([7;32]);
+        let liquidity_pool_cap: u128 = 30u128;
+        let reserve_a: u128 = 30;
+        let reserve_b: u128 = 20;
+        let user_lp_amt: u128 = 10;
+        let token_program_id: [u32;8] = [0; 8];
+
+        pool.data = PoolDefinition::into_data( PoolDefinition {
+            definition_token_a_id: definition_token_a_id.clone(),
+            definition_token_b_id: definition_token_b_id.clone(),
+            vault_a_addr: vault_a_addr.clone(),
+            vault_b_addr: vault_b_addr.clone(),
+            liquidity_pool_id,
+            liquidity_pool_cap: reserve_a,
+            reserve_a,
+            reserve_b,
+            token_program_id,
+        });
+        
+        user_a.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_a_id.clone(),
+                balance: 5u128 }
+        );
+
+        user_b.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_b_id.clone(),
+                balance: 5u128 }
+        );
+
+        user_lp.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id: AccountId::new([3;32]),
+                balance: user_lp_amt }
+        );
+
+        let user_a = AccountWithMetadata {
+            account: user_a.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([5; 32])};
+    
+        let user_b = AccountWithMetadata {
+            account: user_b.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([6; 32])};
+
+        let user_lp = AccountWithMetadata {
+            account: user_lp.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([7; 32])
+        };
+
+        let vault_a = AccountWithMetadata {
+            account: vault_a.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([2; 32])};
+
+        let vault_b = AccountWithMetadata {
+            account: vault_b.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([3; 32])};
+
+        let pool_lp = AccountWithMetadata {
+            account: pool_lp.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([4; 32])};
+
+        let pre_states = vec![AccountWithMetadata {
+            account: pool.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([0; 32])},
+            vault_a.clone(),
+            vault_b.clone(),
+            pool_lp.clone(),
+            user_a.clone(),
+            user_b.clone(),
+            user_lp.clone(),
+        ];
+        let (post_states, chained_calls) = remove_liquidity(&pre_states);
+
+        let chained_call_lp = chained_calls[0].clone();
+        let chained_call_b = chained_calls[1].clone();
+        let chained_call_a = chained_calls[2].clone();
+        
+        //Expected withdraw
+        let withdraw_amount_a = reserve_a * (user_lp_amt/reserve_a);
+        let withdraw_amount_b = reserve_b * (user_lp_amt/reserve_a);
+
+        //Expected chain_call for Token A
+        let mut instruction: [u8;32] = [0; 32];
+        instruction[0] = 1;      
+        instruction[1..17].copy_from_slice(&withdraw_amount_a.to_le_bytes());
+        let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap();
+        let expected_chained_call_a = ChainedCall{
+            program_id: token_program_id,
+            instruction_data,
+            pre_states: vec![vault_a.clone(), user_a.clone()],
+        };
+
+        //Expected chain call for Token B
+        let mut instruction: [u8;32] = [0; 32];
+        instruction[0] = 1;      
+        instruction[1..17].copy_from_slice(&withdraw_amount_b.to_le_bytes());
+        let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap();
+        let expected_chained_call_b = ChainedCall{
+            program_id: token_program_id,
+            instruction_data,
+            pre_states: vec![vault_b.clone(), user_b.clone()],
+        };
+
+        //Expected chain call for LP
+        let mut instruction: [u8;32] = [0; 32];
+        instruction[0] = 1;      
+        instruction[1..17].copy_from_slice(&user_lp_amt.to_le_bytes());
+        let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap();
+        let expected_chained_call_lp = ChainedCall{
+            program_id: token_program_id,
+            instruction_data,
+            pre_states: vec![user_lp.clone(), pool_lp.clone()],
+        };
+        
+        assert!(chained_call_a.program_id == expected_chained_call_a.program_id);
+        assert!(chained_call_a.instruction_data == expected_chained_call_a.instruction_data);
+        assert!(chained_call_a.pre_states[0].account == expected_chained_call_a.pre_states[0].account);
+        assert!(chained_call_a.pre_states[1].account == expected_chained_call_a.pre_states[1].account);
+        assert!(chained_call_b.program_id == expected_chained_call_b.program_id);
+        assert!(chained_call_b.instruction_data == expected_chained_call_b.instruction_data);
+        assert!(chained_call_b.pre_states[0].account == expected_chained_call_b.pre_states[0].account);
+        assert!(chained_call_b.pre_states[1].account == expected_chained_call_b.pre_states[1].account);
+        assert!(chained_call_lp.program_id == expected_chained_call_lp.program_id);
+        assert!(chained_call_lp.instruction_data == expected_chained_call_lp.instruction_data);
+        assert!(chained_call_lp.pre_states[0].account == expected_chained_call_lp.pre_states[0].account);
+        assert!(chained_call_lp.pre_states[1].account == expected_chained_call_lp.pre_states[1].account);
+    }
+
     #[should_panic(expected = "Invalid number of input accounts")]
     #[test]    
     fn test_call_add_liquidity_with_invalid_number_of_accounts_1() {
@@ -2834,6 +3071,7 @@ mod tests {
         let main_token = definition_token_b_id.clone();
         let _post_states = add_liquidity(&pre_states, &[balance_a,balance_b], main_token);
     }
+
     #[should_panic(expected = "Reserves must be nonzero")]
     #[test]
     fn test_call_add_liquidity_reserves_zero() {
@@ -2842,7 +3080,6 @@ mod tests {
         let mut vault2 = Account::default();
         let mut pool_lp = Account::default();
 
-
         let definition_token_a_id = AccountId::new([1;32]);
         let definition_token_b_id = AccountId::new([2;32]); 
 
@@ -2921,6 +3158,354 @@ mod tests {
         let _post_states = add_liquidity(&pre_states, &[balance_a,balance_b], main_token);
     }
 
+    #[test]
+    fn test_call_add_liquidity_chain_call_success_1() {
+        let mut pool = Account::default();
+        let mut vault_a = Account::default();
+        let mut vault_b = Account::default();
+        let mut user_a = Account::default();
+        let mut user_b = Account::default();
+        let mut pool_lp = Account::default();
+        let mut user_lp = Account::default();
+
+        let definition_token_a_id = AccountId::new([1;32]);
+        let definition_token_b_id = AccountId::new([2;32]); 
+        
+        vault_a.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_a_id.clone(),
+                balance: 15u128 }
+        );
+
+        vault_b.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_b_id.clone(),
+                balance: 15u128 }
+        );
+
+        pool_lp.data = vec![
+                1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+                1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+            ];
+
+        let vault_a_addr = AccountId::new([5;32]);
+        let vault_b_addr = AccountId::new([6;32]);
+        let liquidity_pool_id = AccountId::new([7;32]);
+        let liquidity_pool_cap: u128 = 30u128;
+        let reserve_a: u128 = 30;
+        let reserve_b: u128 = 20;
+        let user_lp_amt: u128 = 10;
+        let token_program_id: [u32;8] = [0; 8];
+
+        pool.data = PoolDefinition::into_data( PoolDefinition {
+            definition_token_a_id: definition_token_a_id.clone(),
+            definition_token_b_id: definition_token_b_id.clone(),
+            vault_a_addr: vault_a_addr.clone(),
+            vault_b_addr: vault_b_addr.clone(),
+            liquidity_pool_id,
+            liquidity_pool_cap: reserve_a,
+            reserve_a,
+            reserve_b,
+            token_program_id,
+        });
+        
+        user_a.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_a_id.clone(),
+                balance: 50u128 }
+        );
+
+        user_b.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_b_id.clone(),
+                balance: 50u128 }
+        );
+
+        user_lp.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id: AccountId::new([3;32]),
+                balance: user_lp_amt }
+        );
+
+        let user_a = AccountWithMetadata {
+            account: user_a.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([5; 32])};
+    
+        let user_b = AccountWithMetadata {
+            account: user_b.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([6; 32])};
+
+        let user_lp = AccountWithMetadata {
+            account: user_lp.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([7; 32])
+        };
+
+        let vault_a = AccountWithMetadata {
+            account: vault_a.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([2; 32])};
+
+        let vault_b = AccountWithMetadata {
+            account: vault_b.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([3; 32])};
+
+        let pool_lp = AccountWithMetadata {
+            account: pool_lp.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([4; 32])};
+
+        let pre_states = vec![AccountWithMetadata {
+            account: pool.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([0; 32])},
+            vault_a.clone(),
+            vault_b.clone(),
+            pool_lp.clone(),
+            user_a.clone(),
+            user_b.clone(),
+            user_lp.clone(),
+        ];
+
+        let balance_a = 10u128;
+        let balance_b = 30u128;
+        let main_token = definition_token_a_id.clone();
+        let (post_states, chained_calls) = add_liquidity(&pre_states, &[balance_a,balance_b], main_token);
+
+        let chained_call_lp = chained_calls[0].clone();
+        let chained_call_b = chained_calls[1].clone();
+        let chained_call_a = chained_calls[2].clone();
+        
+        //Expected amounts
+        let expected_actual_amount_a = balance_a;
+        //Uses: (pool_def_data.reserve_b*actual_amount_a)/pool_def_data.reserve_a
+        let expected_actual_amount_b = (reserve_b*expected_actual_amount_a)/reserve_a;
+        let expected_delta_lp = (reserve_a * expected_actual_amount_b)/reserve_b;
+
+        //Expected chain_call for Token A
+        let mut instruction: [u8;23] = [0; 23];
+        instruction[0] = 1;      
+        instruction[1..17].copy_from_slice(&expected_actual_amount_a.to_le_bytes());
+        let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap();
+        let expected_chained_call_a = ChainedCall{
+            program_id: token_program_id,
+            instruction_data,
+            pre_states: vec![user_a.clone(), vault_a.clone()],
+        };
+
+        //Expected chain call for Token B
+        let mut instruction: [u8;23] = [0; 23];
+        instruction[0] = 1;      
+        instruction[1..17].copy_from_slice(&expected_actual_amount_b.to_le_bytes());
+        let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap();
+        let expected_chained_call_b = ChainedCall{
+            program_id: token_program_id,
+            instruction_data,
+            pre_states: vec![user_b.clone(), vault_b.clone()],
+        };
+
+        //Expected chain call for LP
+        let mut instruction: [u8;23] = [0; 23];
+        instruction[0] = 1;      
+        instruction[1..17].copy_from_slice(&expected_delta_lp.to_le_bytes());
+        let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap();
+        let expected_chained_call_lp = ChainedCall{
+            program_id: token_program_id,
+            instruction_data,
+            pre_states: vec![pool_lp.clone(), user_lp.clone()],
+        };
+        
+        assert!(chained_call_a.program_id == expected_chained_call_a.program_id);
+        assert!(chained_call_a.instruction_data == expected_chained_call_a.instruction_data);
+        assert!(chained_call_a.pre_states[0].account == expected_chained_call_a.pre_states[0].account);
+        assert!(chained_call_a.pre_states[1].account == expected_chained_call_a.pre_states[1].account);
+        assert!(chained_call_b.program_id == expected_chained_call_b.program_id);
+        assert!(chained_call_b.instruction_data == expected_chained_call_b.instruction_data);
+        assert!(chained_call_b.pre_states[0].account == expected_chained_call_b.pre_states[0].account);
+        assert!(chained_call_b.pre_states[1].account == expected_chained_call_b.pre_states[1].account);
+        assert!(chained_call_lp.program_id == expected_chained_call_lp.program_id);
+        assert!(chained_call_lp.instruction_data == expected_chained_call_lp.instruction_data);
+        assert!(chained_call_lp.pre_states[0].account == expected_chained_call_lp.pre_states[0].account);
+        assert!(chained_call_lp.pre_states[1].account == expected_chained_call_lp.pre_states[1].account);
+    }
+  
+    #[test]
+    fn test_call_add_liquidity_chain_call_success_2() {
+        let mut pool = Account::default();
+        let mut vault_a = Account::default();
+        let mut vault_b = Account::default();
+        let mut user_a = Account::default();
+        let mut user_b = Account::default();
+        let mut pool_lp = Account::default();
+        let mut user_lp = Account::default();
+
+        let definition_token_a_id = AccountId::new([1;32]);
+        let definition_token_b_id = AccountId::new([2;32]); 
+        
+        vault_a.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_a_id.clone(),
+                balance: 15u128 }
+        );
+
+        vault_b.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_b_id.clone(),
+                balance: 15u128 }
+        );
+
+        pool_lp.data = vec![
+                1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+                1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+            ];
+
+        let vault_a_addr = AccountId::new([5;32]);
+        let vault_b_addr = AccountId::new([6;32]);
+        let liquidity_pool_id = AccountId::new([7;32]);
+        let liquidity_pool_cap: u128 = 30u128;
+        let reserve_a: u128 = 30;
+        let reserve_b: u128 = 20;
+        let user_lp_amt: u128 = 10;
+        let token_program_id: [u32;8] = [0; 8];
+
+        pool.data = PoolDefinition::into_data( PoolDefinition {
+            definition_token_a_id: definition_token_a_id.clone(),
+            definition_token_b_id: definition_token_b_id.clone(),
+            vault_a_addr: vault_a_addr.clone(),
+            vault_b_addr: vault_b_addr.clone(),
+            liquidity_pool_id,
+            liquidity_pool_cap: reserve_a,
+            reserve_a,
+            reserve_b,
+            token_program_id,
+        });
+        
+        user_a.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_a_id.clone(),
+                balance: 50u128 }
+        );
+
+        user_b.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_b_id.clone(),
+                balance: 50u128 }
+        );
+
+        user_lp.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id: AccountId::new([3;32]),
+                balance: user_lp_amt }
+        );
+
+        let user_a = AccountWithMetadata {
+            account: user_a.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([5; 32])};
+    
+        let user_b = AccountWithMetadata {
+            account: user_b.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([6; 32])};
+
+        let user_lp = AccountWithMetadata {
+            account: user_lp.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([7; 32])
+        };
+
+        let vault_a = AccountWithMetadata {
+            account: vault_a.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([2; 32])};
+
+        let vault_b = AccountWithMetadata {
+            account: vault_b.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([3; 32])};
+
+        let pool_lp = AccountWithMetadata {
+            account: pool_lp.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([4; 32])};
+
+        let pre_states = vec![AccountWithMetadata {
+            account: pool.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([0; 32])},
+            vault_a.clone(),
+            vault_b.clone(),
+            pool_lp.clone(),
+            user_a.clone(),
+            user_b.clone(),
+            user_lp.clone(),
+        ];
+
+        let balance_a = 40u128;
+        let balance_b = 20u128;
+        let main_token = definition_token_b_id.clone();
+        let (post_states, chained_calls) = add_liquidity(&pre_states, &[balance_a,balance_b], main_token);
+
+        let chained_call_lp = chained_calls[0].clone();
+        let chained_call_b = chained_calls[1].clone();
+        let chained_call_a = chained_calls[2].clone();
+        
+        //Expected amounts
+        let expected_actual_amount_b = balance_b;
+        //Uses: (pool_def_data.reserve_b*actual_amount_a)/pool_def_data.reserve_a
+        let expected_actual_amount_a = (reserve_a*expected_actual_amount_b)/reserve_b;
+        let expected_delta_lp = (reserve_a * expected_actual_amount_b)/reserve_b;
+
+        //Expected chain_call for Token A
+        let mut instruction: [u8;23] = [0; 23];
+        instruction[0] = 1;      
+        instruction[1..17].copy_from_slice(&expected_actual_amount_a.to_le_bytes());
+        let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap();
+        let expected_chained_call_a = ChainedCall{
+            program_id: token_program_id,
+            instruction_data,
+            pre_states: vec![user_a.clone(), vault_a.clone()],
+        };
+
+        //Expected chain call for Token B
+        let mut instruction: [u8;23] = [0; 23];
+        instruction[0] = 1;      
+        instruction[1..17].copy_from_slice(&expected_actual_amount_b.to_le_bytes());
+        let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap();
+        let expected_chained_call_b = ChainedCall{
+            program_id: token_program_id,
+            instruction_data,
+            pre_states: vec![user_b.clone(), vault_b.clone()],
+        };
+
+        //Expected chain call for LP
+        let mut instruction: [u8;23] = [0; 23];
+        instruction[0] = 1;      
+        instruction[1..17].copy_from_slice(&expected_delta_lp.to_le_bytes());
+        let instruction_data = risc0_zkvm::serde::to_vec(&instruction).unwrap();
+        let expected_chained_call_lp = ChainedCall{
+            program_id: token_program_id,
+            instruction_data,
+            pre_states: vec![pool_lp.clone(), user_lp.clone()],
+        };
+        
+        assert!(chained_call_a.program_id == expected_chained_call_a.program_id);
+        assert!(chained_call_a.instruction_data == expected_chained_call_a.instruction_data);
+        assert!(chained_call_a.pre_states[0].account == expected_chained_call_a.pre_states[0].account);
+        assert!(chained_call_a.pre_states[1].account == expected_chained_call_a.pre_states[1].account);
+        assert!(chained_call_b.program_id == expected_chained_call_b.program_id);
+        assert!(chained_call_b.instruction_data == expected_chained_call_b.instruction_data);
+        assert!(chained_call_b.pre_states[0].account == expected_chained_call_b.pre_states[0].account);
+        assert!(chained_call_b.pre_states[1].account == expected_chained_call_b.pre_states[1].account);
+        assert!(chained_call_lp.program_id == expected_chained_call_lp.program_id);
+        assert!(chained_call_lp.instruction_data == expected_chained_call_lp.instruction_data);
+        assert!(chained_call_lp.pre_states[0].account == expected_chained_call_lp.pre_states[0].account);
+        assert!(chained_call_lp.pre_states[1].account == expected_chained_call_lp.pre_states[1].account);
+    }
+
     #[should_panic(expected = "Invalid number of input accounts")]
     #[test]    
     fn test_call_swap_with_invalid_number_of_accounts_1() {
@@ -3227,4 +3812,501 @@ mod tests {
         let vault_addr = AccountId::new([0;32]);
         let _post_states = swap(&pre_states, amount, vault_addr);
     }
-}
+
+    #[test]
+    fn test_call_swap_successful_chain_call_1() {
+        let mut pool = Account::default();
+        let mut vault_a = Account::default();
+        let mut vault_b = Account::default();
+        let mut pool_lp = Account::default();
+        let mut user_a = Account::default();
+        let mut user_b = Account::default();
+
+        let definition_token_a_id = AccountId::new([1;32]);
+        let definition_token_b_id = AccountId::new([2;32]); 
+
+        vault_a.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_a_id.clone(),
+                balance: 15u128 }
+        );
+
+        vault_b.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_b_id.clone(),
+                balance: 20u128 }
+        );
+
+        pool_lp.data = vec![
+                1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+                1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+            ];
+        user_a.data = vec![
+                0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+                1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+            ];
+        user_b.data = vec![
+                1, 1, 1, 1, 1, 1, 1, 12, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+                1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+            ];
+
+        let definition_token_a_id = AccountId::new([1;32]);
+        let definition_token_b_id = AccountId::new([2;32]);
+        let vault_a_addr = AccountId::new([5;32]);
+        let vault_b_addr = AccountId::new([6;32]);
+        let liquidity_pool_id = AccountId::new([7;32]);
+        let liquidity_pool_cap: u128 = 30u128;
+        let reserve_a: u128 = 15;
+        let reserve_b: u128 = 20;
+        let token_program_id: [u32;8] = [5; 8];
+
+        pool.data = PoolDefinition::into_data(
+            PoolDefinition {
+                definition_token_a_id: definition_token_a_id.clone(),
+                definition_token_b_id: definition_token_b_id.clone(),
+                vault_a_addr: vault_a_addr.clone(),
+                vault_b_addr: vault_b_addr.clone(),
+                liquidity_pool_id,
+                liquidity_pool_cap,
+                reserve_a,
+                reserve_b,
+                token_program_id,
+            }
+         );
+
+        let pre_states = vec![AccountWithMetadata {
+            account: pool.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([0; 32])},
+            AccountWithMetadata {
+            account: vault_a.clone(),
+            is_authorized: true,
+            account_id: vault_a_addr.clone()},
+            AccountWithMetadata {
+            account: vault_b.clone(),
+            is_authorized: true,
+            account_id: vault_b_addr.clone()},
+            AccountWithMetadata {
+            account: user_a.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([4; 32])},
+            AccountWithMetadata {
+            account: user_b.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([5; 32])}
+        ];
+        let amount = 15u128;
+        let token_addr = definition_token_a_id;
+        let (post_accounts, chain_calls) = swap(&pre_states, amount, token_addr);
+
+        let pool_post = post_accounts[0].clone();
+        let pool_pre_data = PoolDefinition::parse(&pool.data).unwrap();
+        let pool_post_data = PoolDefinition::parse(&pool_post.data).unwrap();
+        assert!(pool_post_data.reserve_a == pool_pre_data.reserve_a + amount);
+
+        let expected_withdraw = (pool_pre_data.reserve_b * amount)/(pool_pre_data.reserve_a + amount);
+        assert!(pool_post_data.reserve_b  == pool_pre_data.reserve_b - expected_withdraw);
+
+        let chain_call_a = chain_calls[0].clone();
+        let chain_call_b = chain_calls[1].clone();
+
+        assert!(chain_call_b.program_id == token_program_id);
+        assert!(chain_call_a.program_id == token_program_id);
+
+        let mut instruction_data = [0; 23];
+        instruction_data[0] = 1;
+        instruction_data[1..17].copy_from_slice(&amount.to_le_bytes());
+        let expected_instruction_data_0 = risc0_zkvm::serde::to_vec(&instruction_data).unwrap();
+        let mut instruction_data = [0; 23];
+        instruction_data[0] = 1;
+        instruction_data[1..17].copy_from_slice(&expected_withdraw.to_le_bytes());
+        let expected_instruction_data_1 = risc0_zkvm::serde::to_vec(&instruction_data).unwrap();
+
+        let chain_call_a_account0 = chain_call_a.pre_states[0].account.clone();
+        let chain_call_a_account1 = chain_call_a.pre_states[1].account.clone();
+
+        let chain_call_b_account0 = chain_call_b.pre_states[0].account.clone();
+        let chain_call_b_account1 = chain_call_b.pre_states[1].account.clone();
+
+        assert!(chain_call_a.instruction_data == expected_instruction_data_0);
+        assert!(chain_call_a_account0 == user_a);
+        assert!(chain_call_a_account1 == vault_a);
+        assert!(chain_call_b.instruction_data == expected_instruction_data_1);
+        assert!(chain_call_b_account0 == vault_b);
+        assert!(chain_call_b_account1 == user_b);
+
+    }
+
+    #[test]
+    fn test_call_swap_successful_chain_call_2() {
+        let mut pool = Account::default();
+        let mut vault_a = Account::default();
+        let mut vault_b = Account::default();
+        let mut pool_lp = Account::default();
+        let mut user_a = Account::default();
+        let mut user_b = Account::default();
+
+        let definition_token_a_id = AccountId::new([1;32]);
+        let definition_token_b_id = AccountId::new([2;32]); 
+
+        vault_a.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_a_id.clone(),
+                balance: 15u128 }
+        );
+
+        vault_b.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_b_id.clone(),
+                balance: 20u128 }
+        );
+
+        pool_lp.data = vec![
+                1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+                1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+            ];
+        user_a.data = vec![
+                0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+                1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+            ];
+        user_b.data = vec![
+                1, 1, 1, 1, 1, 1, 1, 12, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+                1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+            ];
+
+        let definition_token_a_id = AccountId::new([1;32]);
+        let definition_token_b_id = AccountId::new([2;32]);
+        let vault_a_addr = AccountId::new([5;32]);
+        let vault_b_addr = AccountId::new([6;32]);
+        let liquidity_pool_id = AccountId::new([7;32]);
+        let liquidity_pool_cap: u128 = 30u128;
+        let reserve_a: u128 = 15;
+        let reserve_b: u128 = 20;
+        let token_program_id: [u32;8] = [5; 8];
+
+        pool.data = PoolDefinition::into_data(
+            PoolDefinition {
+                definition_token_a_id: definition_token_a_id.clone(),
+                definition_token_b_id: definition_token_b_id.clone(),
+                vault_a_addr: vault_a_addr.clone(),
+                vault_b_addr: vault_b_addr.clone(),
+                liquidity_pool_id,
+                liquidity_pool_cap,
+                reserve_a,
+                reserve_b,
+                token_program_id,
+            }
+         );
+         //swapped order of vault_a and vault_b
+        let pre_states = vec![AccountWithMetadata {
+            account: pool.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([0; 32])},
+            AccountWithMetadata {
+            account: vault_b.clone(),
+            is_authorized: true,
+            account_id: vault_b_addr.clone()},
+            AccountWithMetadata {
+            account: vault_a.clone(),
+            is_authorized: true,
+            account_id: vault_a_addr.clone()},
+            AccountWithMetadata {
+            account: user_a.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([4; 32])},
+            AccountWithMetadata {
+            account: user_b.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([5; 32])}
+        ];
+        let amount = 15u128;
+        let token_addr = definition_token_a_id;
+        let (post_accounts, chain_calls) = swap(&pre_states, amount, token_addr);
+
+        let pool_post = post_accounts[0].clone();
+        let pool_pre_data = PoolDefinition::parse(&pool.data).unwrap();
+        let pool_post_data = PoolDefinition::parse(&pool_post.data).unwrap();
+        assert!(pool_post_data.reserve_a == pool_pre_data.reserve_a + amount);
+
+        let expected_withdraw = (pool_pre_data.reserve_b * amount)/(pool_pre_data.reserve_a + amount);
+        assert!(pool_post_data.reserve_b  == pool_pre_data.reserve_b - expected_withdraw);
+
+        let chain_call_a = chain_calls[0].clone();
+        let chain_call_b = chain_calls[1].clone();
+
+        assert!(chain_call_b.program_id == token_program_id);
+        assert!(chain_call_a.program_id == token_program_id);
+
+        let mut instruction_data = [0; 23];
+        instruction_data[0] = 1;
+        instruction_data[1..17].copy_from_slice(&amount.to_le_bytes());
+        let expected_instruction_data_0 = risc0_zkvm::serde::to_vec(&instruction_data).unwrap();
+        let mut instruction_data = [0; 23];
+        instruction_data[0] = 1;
+        instruction_data[1..17].copy_from_slice(&expected_withdraw.to_le_bytes());
+        let expected_instruction_data_1 = risc0_zkvm::serde::to_vec(&instruction_data).unwrap();
+
+        let chain_call_a_account0 = chain_call_a.pre_states[0].account.clone();
+        let chain_call_a_account1 = chain_call_a.pre_states[1].account.clone();
+
+        let chain_call_b_account0 = chain_call_b.pre_states[0].account.clone();
+        let chain_call_b_account1 = chain_call_b.pre_states[1].account.clone();
+
+        assert!(chain_call_a.instruction_data == expected_instruction_data_0);
+        assert!(chain_call_a_account0 == user_a);
+        assert!(chain_call_a_account1 == vault_a);
+        assert!(chain_call_b.instruction_data == expected_instruction_data_1);
+        assert!(chain_call_b_account0 == vault_b);
+        assert!(chain_call_b_account1 == user_b);
+
+    }
+
+    #[test]
+    fn test_call_swap_successful_chain_call_3() {
+        let mut pool = Account::default();
+        let mut vault_a = Account::default();
+        let mut vault_b = Account::default();
+        let mut pool_lp = Account::default();
+        let mut user_a = Account::default();
+        let mut user_b = Account::default();
+
+        let definition_token_a_id = AccountId::new([1;32]);
+        let definition_token_b_id = AccountId::new([2;32]); 
+
+        vault_a.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_a_id.clone(),
+                balance: 15u128 }
+        );
+
+        vault_b.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_b_id.clone(),
+                balance: 20u128 }
+        );
+
+        pool_lp.data = vec![
+                1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+                1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+            ];
+        user_a.data = vec![
+                0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+                1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+            ];
+        user_b.data = vec![
+                1, 1, 1, 1, 1, 1, 1, 12, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+                1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+            ];
+
+        let definition_token_a_id = AccountId::new([1;32]);
+        let definition_token_b_id = AccountId::new([2;32]);
+        let vault_a_addr = AccountId::new([5;32]);
+        let vault_b_addr = AccountId::new([6;32]);
+        let liquidity_pool_id = AccountId::new([7;32]);
+        let liquidity_pool_cap: u128 = 30u128;
+        let reserve_a: u128 = 15;
+        let reserve_b: u128 = 20;
+        let token_program_id: [u32;8] = [5; 8];
+
+        pool.data = PoolDefinition::into_data(
+            PoolDefinition {
+                definition_token_a_id: definition_token_a_id.clone(),
+                definition_token_b_id: definition_token_b_id.clone(),
+                vault_a_addr: vault_a_addr.clone(),
+                vault_b_addr: vault_b_addr.clone(),
+                liquidity_pool_id,
+                liquidity_pool_cap,
+                reserve_a,
+                reserve_b,
+                token_program_id,
+            }
+         );
+
+        let pre_states = vec![AccountWithMetadata {
+            account: pool.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([0; 32])},
+            AccountWithMetadata {
+            account: vault_a.clone(),
+            is_authorized: true,
+            account_id: vault_a_addr.clone()},
+            AccountWithMetadata {
+            account: vault_b.clone(),
+            is_authorized: true,
+            account_id: vault_b_addr.clone()},
+            AccountWithMetadata {
+            account: user_a.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([4; 32])},
+            AccountWithMetadata {
+            account: user_b.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([5; 32])}
+        ];
+        let amount = 15u128;
+        let token_addr = definition_token_b_id;
+        let (post_accounts, chain_calls) = swap(&pre_states, amount, token_addr);
+
+        let pool_post = post_accounts[0].clone();
+        let pool_pre_data = PoolDefinition::parse(&pool.data).unwrap();
+        let pool_post_data = PoolDefinition::parse(&pool_post.data).unwrap();
+        assert!(pool_post_data.reserve_b == pool_pre_data.reserve_b + amount);
+
+        let expected_withdraw = (pool_pre_data.reserve_a * amount)/(pool_pre_data.reserve_b + amount);
+        assert!(pool_post_data.reserve_a  == pool_pre_data.reserve_a - expected_withdraw);
+
+        let chain_call_a = chain_calls[0].clone();
+        let chain_call_b = chain_calls[1].clone();
+
+        assert!(chain_call_b.program_id == token_program_id);
+        assert!(chain_call_a.program_id == token_program_id);
+
+        let mut instruction_data = [0; 23];
+        instruction_data[0] = 1;
+        instruction_data[1..17].copy_from_slice(&expected_withdraw.to_le_bytes());
+        let expected_instruction_data_0 = risc0_zkvm::serde::to_vec(&instruction_data).unwrap();
+        let mut instruction_data = [0; 23];
+        instruction_data[0] = 1;
+        instruction_data[1..17].copy_from_slice(&amount.to_le_bytes());
+        let expected_instruction_data_1 = risc0_zkvm::serde::to_vec(&instruction_data).unwrap();
+
+        let chain_call_a_account0 = chain_call_a.pre_states[0].account.clone();
+        let chain_call_a_account1 = chain_call_a.pre_states[1].account.clone();
+
+        let chain_call_b_account0 = chain_call_b.pre_states[0].account.clone();
+        let chain_call_b_account1 = chain_call_b.pre_states[1].account.clone();
+
+        assert!(chain_call_a.instruction_data == expected_instruction_data_0);
+        assert!(chain_call_a_account0 == vault_a);
+        assert!(chain_call_a_account1 == user_a);
+        assert!(chain_call_b.instruction_data == expected_instruction_data_1);
+        assert!(chain_call_b_account0 == user_b);
+        assert!(chain_call_b_account1 == vault_b);
+
+    }
+
+    #[test]
+    fn test_call_swap_successful_chain_call_4() {
+        let mut pool = Account::default();
+        let mut vault_a = Account::default();
+        let mut vault_b = Account::default();
+        let mut pool_lp = Account::default();
+        let mut user_a = Account::default();
+        let mut user_b = Account::default();
+
+        let definition_token_a_id = AccountId::new([1;32]);
+        let definition_token_b_id = AccountId::new([2;32]); 
+
+        vault_a.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_a_id.clone(),
+                balance: 15u128 }
+        );
+
+        vault_b.data = TokenHolding::into_data(
+            TokenHolding { account_type: TOKEN_HOLDING_TYPE,
+                definition_id:definition_token_b_id.clone(),
+                balance: 20u128 }
+        );
+
+        pool_lp.data = vec![
+                1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+                1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+            ];
+        user_a.data = vec![
+                0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+                1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+            ];
+        user_b.data = vec![
+                1, 1, 1, 1, 1, 1, 1, 12, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+                1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
+            ];
+
+        let definition_token_a_id = AccountId::new([1;32]);
+        let definition_token_b_id = AccountId::new([2;32]);
+        let vault_a_addr = AccountId::new([5;32]);
+        let vault_b_addr = AccountId::new([6;32]);
+        let liquidity_pool_id = AccountId::new([7;32]);
+        let liquidity_pool_cap: u128 = 30u128;
+        let reserve_a: u128 = 15;
+        let reserve_b: u128 = 20;
+        let token_program_id: [u32;8] = [5; 8];
+
+        pool.data = PoolDefinition::into_data(
+            PoolDefinition {
+                definition_token_a_id: definition_token_a_id.clone(),
+                definition_token_b_id: definition_token_b_id.clone(),
+                vault_a_addr: vault_a_addr.clone(),
+                vault_b_addr: vault_b_addr.clone(),
+                liquidity_pool_id,
+                liquidity_pool_cap,
+                reserve_a,
+                reserve_b,
+                token_program_id,
+            }
+         );
+
+        //swapped order of vaults
+        let pre_states = vec![AccountWithMetadata {
+            account: pool.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([0; 32])},
+            AccountWithMetadata {
+            account: vault_b.clone(),
+            is_authorized: true,
+            account_id: vault_b_addr.clone()},
+            AccountWithMetadata {
+            account: vault_a.clone(),
+            is_authorized: true,
+            account_id: vault_b_addr.clone()},
+            AccountWithMetadata {
+            account: user_a.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([4; 32])},
+            AccountWithMetadata {
+            account: user_b.clone(),
+            is_authorized: true,
+            account_id: AccountId::new([5; 32])}
+        ];
+        let amount = 15u128;
+        let token_addr = definition_token_b_id;
+        let (post_accounts, chain_calls) = swap(&pre_states, amount, token_addr);
+
+        let pool_post = post_accounts[0].clone();
+        let pool_pre_data = PoolDefinition::parse(&pool.data).unwrap();
+        let pool_post_data = PoolDefinition::parse(&pool_post.data).unwrap();
+        assert!(pool_post_data.reserve_b == pool_pre_data.reserve_b + amount);
+
+        let expected_withdraw = (pool_pre_data.reserve_a * amount)/(pool_pre_data.reserve_b + amount);
+        assert!(pool_post_data.reserve_a  == pool_pre_data.reserve_a - expected_withdraw);
+
+        let chain_call_a = chain_calls[0].clone();
+        let chain_call_b = chain_calls[1].clone();
+
+        assert!(chain_call_b.program_id == token_program_id);
+        assert!(chain_call_a.program_id == token_program_id);
+
+        let mut instruction_data = [0; 23];
+        instruction_data[0] = 1;
+        instruction_data[1..17].copy_from_slice(&expected_withdraw.to_le_bytes());
+        let expected_instruction_data_0 = risc0_zkvm::serde::to_vec(&instruction_data).unwrap();
+        let mut instruction_data = [0; 23];
+        instruction_data[0] = 1;
+        instruction_data[1..17].copy_from_slice(&amount.to_le_bytes());
+        let expected_instruction_data_1 = risc0_zkvm::serde::to_vec(&instruction_data).unwrap();
+
+        let chain_call_a_account0 = chain_call_a.pre_states[0].account.clone();
+        let chain_call_a_account1 = chain_call_a.pre_states[1].account.clone();
+
+        let chain_call_b_account0 = chain_call_b.pre_states[0].account.clone();
+        let chain_call_b_account1 = chain_call_b.pre_states[1].account.clone();
+
+        assert!(chain_call_a.instruction_data == expected_instruction_data_0);
+        assert!(chain_call_a_account0 == vault_a);
+        assert!(chain_call_a_account1 == user_a);
+        assert!(chain_call_b.instruction_data == expected_instruction_data_1);
+        assert!(chain_call_b_account0 == user_b);
+        assert!(chain_call_b_account1 == vault_b);
+
+    }
+}
\ No newline at end of file
diff --git a/nssa/src/state.rs b/nssa/src/state.rs
index 389c3c3..8a1aadf 100644
--- a/nssa/src/state.rs
+++ b/nssa/src/state.rs
@@ -2161,6 +2161,7 @@ pub mod tests {
         ));
     }
 
+    
     //TODO: repeated code needs to be cleaned up
     //from token.rs (also repeated in amm.rs)
     const TOKEN_DEFINITION_TYPE: u8 = 0;
@@ -2224,140 +2225,124 @@ pub mod tests {
         }
     }
 
+    const POOL_DEFINITION_DATA_SIZE: usize = 240;
 
-const POOL_DEFINITION_DATA_SIZE: usize = 240;
-
-struct PoolDefinition{
-    definition_token_a_id: AccountId,
-    definition_token_b_id: AccountId,
-    vault_a_addr: AccountId,
-    vault_b_addr: AccountId,
-    liquidity_pool_id: AccountId,
-    liquidity_pool_cap: u128,
-    reserve_a: u128,
-    reserve_b: u128,
-    token_program_id: ProgramId,
-}
-
-
-impl PoolDefinition {
-    fn into_data(self) -> Vec<u8> {
-        /*
-        //TODO: delete?
-        let mut u8_token_program_id: [u8;32] = [0;32];
-        u8_token_program_id[0..4].clone_from_slice(&self.token_program_id[0].to_le_bytes());
-        u8_token_program_id[4..8].clone_from_slice(&self.token_program_id[1].to_le_bytes());
-        u8_token_program_id[8..12].clone_from_slice(&self.token_program_id[2].to_le_bytes());
-        u8_token_program_id[12..16].clone_from_slice(&self.token_program_id[3].to_le_bytes());
-        u8_token_program_id[16..20].clone_from_slice(&self.token_program_id[4].to_le_bytes());
-        u8_token_program_id[20..24].clone_from_slice(&self.token_program_id[5].to_le_bytes());
-        u8_token_program_id[24..28].clone_from_slice(&self.token_program_id[6].to_le_bytes());
-        u8_token_program_id[28..32].clone_from_slice(&self.token_program_id[7].to_le_bytes());
-*/
-
-        let u8_token_program_id : [u8;32] = bytemuck::cast(self.token_program_id);
-
-        let mut bytes = [0; POOL_DEFINITION_DATA_SIZE];
-        bytes[0..32].copy_from_slice(&self.definition_token_a_id.to_bytes());
-        bytes[32..64].copy_from_slice(&self.definition_token_b_id.to_bytes());
-        bytes[64..96].copy_from_slice(&self.vault_a_addr.to_bytes());
-        bytes[96..128].copy_from_slice(&self.vault_b_addr.to_bytes());
-        bytes[128..160].copy_from_slice(&self.liquidity_pool_id.to_bytes());
-        bytes[160..176].copy_from_slice(&self.liquidity_pool_cap.to_le_bytes());
-        bytes[176..192].copy_from_slice(&self.reserve_a.to_le_bytes());
-        bytes[192..208].copy_from_slice(&self.reserve_b.to_le_bytes());
-        bytes[208..].copy_from_slice(&u8_token_program_id);
-        bytes.into()
+    struct PoolDefinition {
+        definition_token_a_id: AccountId,
+        definition_token_b_id: AccountId,
+        vault_a_addr: AccountId,
+        vault_b_addr: AccountId,
+        liquidity_pool_id: AccountId,
+        liquidity_pool_cap: u128,
+        reserve_a: u128,
+        reserve_b: u128,
+        token_program_id: ProgramId,
     }
 
-    fn parse(data: &[u8]) -> Option<Self> {
-        if data.len() != POOL_DEFINITION_DATA_SIZE {
-            None
-        } else {
-            let definition_token_a_id = AccountId::new(data[0..32].try_into().unwrap());
-            let definition_token_b_id = AccountId::new(data[32..64].try_into().unwrap());
-            let vault_a_addr = AccountId::new(data[64..96].try_into().unwrap());
-            let vault_b_addr = AccountId::new(data[96..128].try_into().unwrap());
-            let liquidity_pool_id = AccountId::new(data[128..160].try_into().unwrap());
-            let liquidity_pool_cap = u128::from_le_bytes(data[160..176].try_into().unwrap());
-            let reserve_a = u128::from_le_bytes(data[176..192].try_into().unwrap());
-            let reserve_b = u128::from_le_bytes(data[192..208].try_into().unwrap());
+    impl PoolDefinition {
+        fn into_data(self) -> Vec<u8> {
+            let u8_token_program_id: [u8; 32] = bytemuck::cast(self.token_program_id);
 
-            let token_program_id : &[u32] = bytemuck::cast_slice(&data[208..]);
-            let token_program_id : ProgramId = token_program_id[0..8].try_into().unwrap();
-            Some(Self {
-                definition_token_a_id,
-                definition_token_b_id,
-                vault_a_addr,
-                vault_b_addr,
-                liquidity_pool_id,
-                liquidity_pool_cap,
-                reserve_a,
-                reserve_b,
-                token_program_id,
-            })
+            let mut bytes = [0; POOL_DEFINITION_DATA_SIZE];
+            bytes[0..32].copy_from_slice(&self.definition_token_a_id.to_bytes());
+            bytes[32..64].copy_from_slice(&self.definition_token_b_id.to_bytes());
+            bytes[64..96].copy_from_slice(&self.vault_a_addr.to_bytes());
+            bytes[96..128].copy_from_slice(&self.vault_b_addr.to_bytes());
+            bytes[128..160].copy_from_slice(&self.liquidity_pool_id.to_bytes());
+            bytes[160..176].copy_from_slice(&self.liquidity_pool_cap.to_le_bytes());
+            bytes[176..192].copy_from_slice(&self.reserve_a.to_le_bytes());
+            bytes[192..208].copy_from_slice(&self.reserve_b.to_le_bytes());
+            bytes[208..].copy_from_slice(&u8_token_program_id);
+            bytes.into()
+        }
+
+        fn parse(data: &[u8]) -> Option<Self> {
+            if data.len() != POOL_DEFINITION_DATA_SIZE {
+                None
+            } else {
+                let definition_token_a_id = AccountId::new(data[0..32].try_into().unwrap());
+                let definition_token_b_id = AccountId::new(data[32..64].try_into().unwrap());
+                let vault_a_addr = AccountId::new(data[64..96].try_into().unwrap());
+                let vault_b_addr = AccountId::new(data[96..128].try_into().unwrap());
+                let liquidity_pool_id = AccountId::new(data[128..160].try_into().unwrap());
+                let liquidity_pool_cap = u128::from_le_bytes(data[160..176].try_into().unwrap());
+                let reserve_a = u128::from_le_bytes(data[176..192].try_into().unwrap());
+                let reserve_b = u128::from_le_bytes(data[192..208].try_into().unwrap());
+
+                let token_program_id: &[u32] = bytemuck::cast_slice(&data[208..]);
+                let token_program_id: ProgramId = token_program_id[0..8].try_into().unwrap();
+                Some(Self {
+                    definition_token_a_id,
+                    definition_token_b_id,
+                    vault_a_addr,
+                    vault_b_addr,
+                    liquidity_pool_id,
+                    liquidity_pool_cap,
+                    reserve_a,
+                    reserve_b,
+                    token_program_id,
+                })
+            }
         }
     }
-}
-
-
-
-
-
-#[test]
-    fn test_simple_amm_initialized_attempt_initialize() {
-        let program_token = Program::token().id();
 
+    /// Used for each amm test to initialize
+    /// an AMM pool
+    fn initialize_amm() -> (V02State, Vec<PrivateKey>, Vec<Address>, Vec<u128>) {
         let initial_data = [];
         let mut state =
             V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
 
-
-        let token_a_holding_key = PrivateKey::try_new([1;32]).unwrap();
-        let token_a_definition_key = PrivateKey::try_new([2;32]).unwrap();
-        let token_a_holding_address = Address::from(&PublicKey::new_from_private_key(&token_a_holding_key));
-        let token_a_definition_address = Address::from(&PublicKey::new_from_private_key(&token_a_definition_key));
+        let token_a_holding_key = PrivateKey::try_new([1; 32]).unwrap();
+        let token_a_definition_key = PrivateKey::try_new([2; 32]).unwrap();
+        let token_a_holding_address =
+            Address::from(&PublicKey::new_from_private_key(&token_a_holding_key));
+        let token_a_definition_address =
+            Address::from(&PublicKey::new_from_private_key(&token_a_definition_key));
         let token_a_supply: u128 = 30000;
 
-        let token_b_holding_key = PrivateKey::try_new([3;32]).unwrap();
-        let token_b_definition_key = PrivateKey::try_new([4;32]).unwrap();
-        let token_b_holding_address = Address::from(&PublicKey::new_from_private_key(&token_b_holding_key));
-        let token_b_definition_address = Address::from(&PublicKey::new_from_private_key(&token_b_definition_key));
+        let token_b_holding_key = PrivateKey::try_new([3; 32]).unwrap();
+        let token_b_definition_key = PrivateKey::try_new([4; 32]).unwrap();
+        let token_b_holding_address =
+            Address::from(&PublicKey::new_from_private_key(&token_b_holding_key));
+        let token_b_definition_address =
+            Address::from(&PublicKey::new_from_private_key(&token_b_definition_key));
         let token_b_supply: u128 = 50000;
 
-        let token_lp_holding_key = PrivateKey::try_new([5;32]).unwrap();
-        let token_lp_definition_key = PrivateKey::try_new([6;32]).unwrap();
-        let token_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&token_lp_holding_key));
-        let token_lp_definition_address = Address::from(&PublicKey::new_from_private_key(&token_lp_definition_key));
+        let pool_lp_holding_key = PrivateKey::try_new([5; 32]).unwrap();
+        let pool_lp_definition_key = PrivateKey::try_new([6; 32]).unwrap();
+        let pool_lp_holding_address =
+            Address::from(&PublicKey::new_from_private_key(&pool_lp_holding_key));
+        let pool_lp_definition_address =
+            Address::from(&PublicKey::new_from_private_key(&pool_lp_definition_key));
         let token_lp_supply: u128 = 300000;
 
-        let user_a_holding_key = PrivateKey::try_new([7;32]).unwrap();
-        let user_a_holding_address = Address::from(&PublicKey::new_from_private_key(&user_a_holding_key));
+        let user_a_holding_key = PrivateKey::try_new([7; 32]).unwrap();
+        let user_a_holding_address =
+            Address::from(&PublicKey::new_from_private_key(&user_a_holding_key));
         let user_a_amount: u128 = 10000;
 
-        let user_b_holding_key = PrivateKey::try_new([8;32]).unwrap();
-        let user_b_holding_address = Address::from(&PublicKey::new_from_private_key(&user_b_holding_key));
+        let user_b_holding_key = PrivateKey::try_new([8; 32]).unwrap();
+        let user_b_holding_address =
+            Address::from(&PublicKey::new_from_private_key(&user_b_holding_key));
         let user_b_amount: u128 = 10000;
 
-
-
-        let vault_a_key = PrivateKey::try_new([9;32]).unwrap();
+        let vault_a_key = PrivateKey::try_new([9; 32]).unwrap();
         let vault_a_address = Address::from(&PublicKey::new_from_private_key(&vault_a_key));
 
-        let vault_b_key = PrivateKey::try_new([10;32]).unwrap();
+        let vault_b_key = PrivateKey::try_new([10; 32]).unwrap();
         let vault_b_address = Address::from(&PublicKey::new_from_private_key(&vault_b_key));
 
-        let user_lp_holding_key = PrivateKey::try_new([11;32]).unwrap();
-        let user_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&user_lp_holding_key));
+        let user_lp_holding_key = PrivateKey::try_new([11; 32]).unwrap();
+        let user_lp_holding_address =
+            Address::from(&PublicKey::new_from_private_key(&user_lp_holding_key));
 
-        let pool_lp_holding_key = PrivateKey::try_new([12;32]).unwrap();
-        let pool_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&pool_lp_holding_key));
 
-        let pool_key = PrivateKey::try_new([13;32]).unwrap();
+        let pool_key = PrivateKey::try_new([13; 32]).unwrap();
         let pool_address = Address::from(&PublicKey::new_from_private_key(&pool_key));
 
-        //initialize Tokens: Token A, Token B and Took LP
-        let mut instruction: [u8;23] = [0;23];
+        //initialize Token A
+        let mut instruction: [u8; 23] = [0; 23];
         instruction[1..17].copy_from_slice(&token_a_supply.to_le_bytes());
         instruction[18] = 0x01; //name is not default.
         instruction[19] = 0x02;
@@ -2370,13 +2355,11 @@ impl PoolDefinition {
         )
         .unwrap();
 
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[],
-        );
+        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
         let tx = PublicTransaction::new(message, witness_set);
         state.transition_from_public_transaction(&tx).unwrap();
 
+        //initialize Token B       
         instruction[1..17].copy_from_slice(&token_b_supply.to_le_bytes());
         instruction[18] = 0x03; //name is not default.
         instruction[19] = 0x02;
@@ -2389,35 +2372,30 @@ impl PoolDefinition {
         )
         .unwrap();
 
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[],
-        );
+        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
         let tx = PublicTransaction::new(message, witness_set);
         state.transition_from_public_transaction(&tx).unwrap();
 
+
+        //initialize Token LP        
         instruction[1..17].copy_from_slice(&token_lp_supply.to_le_bytes());
         instruction[18] = 0x03; //name is not default.
         instruction[19] = 0x04;
 
         let message = public_transaction::Message::try_new(
             Program::token().id(),
-            vec![token_lp_definition_address, token_lp_holding_address],
+            vec![pool_lp_definition_address, pool_lp_holding_address],
             vec![],
             instruction,
         )
         .unwrap();
 
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[],
-        );
+        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
         let tx = PublicTransaction::new(message, witness_set);
         state.transition_from_public_transaction(&tx).unwrap();
 
-
-        // Initialize User accounts for Token A and B
-        let mut instruction: [u8;23] = [0;23];
+        // Initialize User accounts for Token A
+        let mut instruction: [u8; 23] = [0; 23];
         instruction[0] = 1; //transfer
         instruction[1..17].copy_from_slice(&user_a_amount.to_le_bytes());
 
@@ -2429,13 +2407,13 @@ impl PoolDefinition {
         )
         .unwrap();
 
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_a_holding_key],
-        );
+        let witness_set =
+            public_transaction::WitnessSet::for_message(&message, &[&token_a_holding_key]);
         let tx = PublicTransaction::new(message, witness_set);
         state.transition_from_public_transaction(&tx).unwrap();
 
+
+        // Initialize User accounts for Token A
         instruction[0] = 1; //transfer
         instruction[1..17].copy_from_slice(&user_b_amount.to_le_bytes());
 
@@ -2447,17 +2425,15 @@ impl PoolDefinition {
         )
         .unwrap();
 
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_b_holding_key],
-        );
+        let witness_set =
+            public_transaction::WitnessSet::for_message(&message, &[&token_b_holding_key]);
         let tx = PublicTransaction::new(message, witness_set);
         state.transition_from_public_transaction(&tx).unwrap();
 
-
         //TODO: initialize vaults - ideally, we won't need to do this.
+        // Initialize Vault A
         let temp_amt = 1u128;
-        let mut instruction: [u8;23] = [0;23];
+        let mut instruction: [u8; 23] = [0; 23];
         instruction[0] = 1; //transfer
         instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes());
 
@@ -2469,13 +2445,12 @@ impl PoolDefinition {
         )
         .unwrap();
 
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_a_holding_key],
-        );
+        let witness_set =
+            public_transaction::WitnessSet::for_message(&message, &[&token_a_holding_key]);
         let tx = PublicTransaction::new(message, witness_set);
         state.transition_from_public_transaction(&tx).unwrap();
 
+        // Initialize Vault B
         instruction[0] = 1; //transfer
         instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes());
 
@@ -2487,63 +2462,39 @@ impl PoolDefinition {
         )
         .unwrap();
 
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_b_holding_key],
-        );
+        let witness_set =
+            public_transaction::WitnessSet::for_message(&message, &[&token_b_holding_key]);
         let tx = PublicTransaction::new(message, witness_set);
         state.transition_from_public_transaction(&tx).unwrap();
 
-        //need LP accounts for both AMM and user
+        // Initialize User LP
         instruction[0] = 1; //transfer
         instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes());
 
         let message = public_transaction::Message::try_new(
             Program::token().id(),
-            vec![token_lp_holding_address, user_lp_holding_address],
+            vec![pool_lp_holding_address, user_lp_holding_address],
             vec![0],
             instruction,
         )
         .unwrap();
 
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_lp_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        let amt: u128 = token_lp_supply - 1;
-        //Pool LP should be initialized at full amount (token_lp_supply)
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&amt.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_lp_holding_address, pool_lp_holding_address],
-            vec![1],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_lp_holding_key],
-        );
+        let witness_set =
+            public_transaction::WitnessSet::for_message(&message, &[&pool_lp_holding_key]);
         let tx = PublicTransaction::new(message, witness_set);
         state.transition_from_public_transaction(&tx).unwrap();
 
         //Set up instruction to initialize AMM for (Token-A, Token-B)
         let init_balance_a: u128 = 1000;
         let init_balance_b: u128 = 1000;
-        let token_program_u8: [u8;32] = bytemuck::cast(program_token);
-        
+        let token_program_u8: [u8; 32] = bytemuck::cast(Program::token().id());
+
         let mut instruction: Vec<u8> = Vec::new();
         instruction.push(0);
         instruction.extend_from_slice(&init_balance_a.to_le_bytes());
         instruction.extend_from_slice(&init_balance_b.to_le_bytes());
         instruction.extend_from_slice(&token_program_u8);
-        
+
         let message = public_transaction::Message::try_new(
             Program::amm().id(),
             vec![
@@ -2553,9 +2504,9 @@ impl PoolDefinition {
                 pool_lp_holding_address,
                 user_a_holding_address,
                 user_b_holding_address,
-                user_lp_holding_address
+                user_lp_holding_address,
             ],
-            vec![0,0,0,0],
+            vec![0, 1, 0, 0],
             instruction,
         )
         .unwrap();
@@ -2569,301 +2520,204 @@ impl PoolDefinition {
                 &user_b_holding_key,
             ],
         );
-        
+
         let tx = PublicTransaction::new(message, witness_set);
         state.transition_from_public_transaction(&tx).unwrap();
 
+        let mut vec_address = Vec::new();
+        vec_address.push(token_a_holding_address);
+        vec_address.push(token_a_definition_address);
+        vec_address.push(token_b_holding_address);
+        vec_address.push(token_b_definition_address);
+        vec_address.push(pool_lp_definition_address);
+        vec_address.push(user_a_holding_address);
+        vec_address.push(user_b_holding_address);
+        vec_address.push(vault_a_address);
+        vec_address.push(vault_b_address);
+        vec_address.push(user_lp_holding_address);
+        vec_address.push(pool_address);
+        vec_address.push(pool_lp_holding_address);
+
+        let mut vec_private_keys = Vec::new();
+        vec_private_keys.push(token_a_holding_key);
+        vec_private_keys.push(token_a_definition_key);
+        vec_private_keys.push(token_b_holding_key);
+        vec_private_keys.push(token_b_definition_key);
+        vec_private_keys.push(pool_lp_definition_key);
+        vec_private_keys.push(user_a_holding_key);
+        vec_private_keys.push(user_b_holding_key);
+        vec_private_keys.push(vault_a_key);
+        vec_private_keys.push(vault_b_key);
+        vec_private_keys.push(user_lp_holding_key);
+        vec_private_keys.push(pool_key);
+        vec_private_keys.push(pool_lp_holding_key);
+
+        let mut vec_amounts = Vec::new();
+        vec_amounts.push(temp_amt);
+        vec_amounts.push(init_balance_a);
+        vec_amounts.push(init_balance_b);
+        vec_amounts.push(user_a_amount);
+        vec_amounts.push(user_b_amount);
+
+        (state, vec_private_keys, vec_address, vec_amounts)
+    }
+
+    #[test]
+    fn test_simple_amm_initialize() {
+        let (state, _vec_private_keys, vec_address, vec_amounts) = initialize_amm();
+
+        let temp_amt = vec_amounts[0];
+        let init_balance_a = vec_amounts[1];
+        let init_balance_b = vec_amounts[2];
+        let user_a_amount = vec_amounts[3];
+        let user_b_amount = vec_amounts[4];
+
+        let token_a_holding_address = vec_address[0];
+        let token_a_definition_address = vec_address[1];
+        let token_b_holding_address = vec_address[2];
+        let token_b_definition_address = vec_address[3];
+        let token_lp_definition_address = vec_address[4];
+        let user_a_holding_address = vec_address[5];
+        let user_b_holding_address = vec_address[6];
+        let vault_a_address = vec_address[7];
+        let vault_b_address = vec_address[8];
+        let user_lp_holding_address = vec_address[9];
+        let pool_address = vec_address[10];
+        let pool_lp_holding_address = vec_address[11];
 
         let pool_post = state.get_account_by_address(&pool_address);
         let vault_a_post = state.get_account_by_address(&vault_a_address);
         let vault_b_post = state.get_account_by_address(&vault_b_address);
-        let pool_lp_post = state.get_account_by_address(&pool_lp_holding_address);
         let user_a_post = state.get_account_by_address(&user_a_holding_address);
         let user_b_post = state.get_account_by_address(&user_b_holding_address);
         let user_lp_post = state.get_account_by_address(&user_lp_holding_address);
 
-        //TODO: token accounts initialized with temp_amt balance.
-        // this ensures that amm program does not own any of the token accounts
-        assert!(TokenHolding::parse(&vault_a_post.data).unwrap().balance == init_balance_a+temp_amt);        
-        assert!(TokenHolding::parse(&vault_b_post.data).unwrap().balance == init_balance_b+temp_amt);        
-        assert!(TokenHolding::parse(&user_lp_post.data).unwrap().balance == init_balance_a+temp_amt);     
-        assert!(PoolDefinition::parse(&pool_post.data).unwrap().reserve_a == init_balance_a);
-        assert!(PoolDefinition::parse(&pool_post.data).unwrap().reserve_b == init_balance_b); 
+        let expected_pool = Account {
+            program_owner: Program::amm().id(),
+            balance: 0u128,
+            data: PoolDefinition::into_data(
+                PoolDefinition {
+                    definition_token_a_id: token_a_definition_address,
+                    definition_token_b_id: token_b_definition_address,
+                    vault_a_addr: vault_a_address,
+                    vault_b_addr: vault_b_address,
+                    liquidity_pool_id: token_lp_definition_address,
+                    liquidity_pool_cap: init_balance_a,
+                    reserve_a: init_balance_a,
+                    reserve_b: init_balance_b,
+                    token_program_id: Program::token().id(),
+                }),
+            nonce: 1,
+        };
+
+        let expected_vault_a = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_a_definition_address,
+                    balance: init_balance_a + temp_amt,
+            }),
+            nonce: 0
+        };
+
+        let expected_vault_b = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_b_definition_address,
+                    balance: init_balance_b + temp_amt,
+            }),
+            nonce: 0
+        };
+
+        let expected_user_a = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_a_definition_address,
+                    balance: user_a_amount - init_balance_a,
+            }),
+            nonce: 1
+        };
+
+        let expected_user_b = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_b_definition_address,
+                    balance: user_b_amount - init_balance_b,
+            }),
+            nonce: 1
+        };
+
+        let expected_user_lp = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_lp_definition_address,
+                    balance: init_balance_a + temp_amt,
+            }),
+            nonce: 0
+        };
+
+        assert!(vault_a_post == expected_vault_a);
+        assert!(vault_b_post == expected_vault_b);
+        assert!(user_a_post == expected_user_a);
+        assert!(user_b_post == expected_user_b);
+        assert!(user_lp_post == expected_user_lp);
+        assert!(pool_post == expected_pool);
     }
 
+    #[test]
+    fn test_simple_amm_remove() {
+        let (state, vec_private_keys, vec_address, vec_amounts) = initialize_amm();
+        let mut state: V02State = state;
 
-#[test]
-    fn test_simple_amm_initialized_attempt_remove() {
-        let program_token = Program::token().id();
+        let temp_amt = vec_amounts[0];
+        let init_balance_a = vec_amounts[1];
+        let init_balance_b = vec_amounts[2];
+        let user_a_amount = vec_amounts[3];
+        let user_b_amount = vec_amounts[4];
 
-        let initial_data = [];
-        let mut state =
-            V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
+        let token_a_holding_key = &vec_private_keys[0];
+        let token_a_definition_key = &vec_private_keys[1];
+        let token_b_holding_key = &vec_private_keys[2];
+        let token_b_definition_key = &vec_private_keys[3];
+        let pool_lp_definition_key = &vec_private_keys[4];
+        let user_a_holding_key = &vec_private_keys[5];
+        let user_b_holding_key = &vec_private_keys[6];
+        let vault_a_key = &vec_private_keys[7];
+        let vault_b_key = &vec_private_keys[8];
+        let user_lp_holding_key = &vec_private_keys[9];
+        let pool_key = &vec_private_keys[10];
+        let pool_lp_holding_key = &vec_private_keys[11];
 
-
-        let token_a_holding_key = PrivateKey::try_new([1;32]).unwrap();
-        let token_a_definition_key = PrivateKey::try_new([2;32]).unwrap();
-        let token_a_holding_address = Address::from(&PublicKey::new_from_private_key(&token_a_holding_key));
-        let token_a_definition_address = Address::from(&PublicKey::new_from_private_key(&token_a_definition_key));
-        let token_a_supply: u128 = 30000;
-
-        let token_b_holding_key = PrivateKey::try_new([3;32]).unwrap();
-        let token_b_definition_key = PrivateKey::try_new([4;32]).unwrap();
-        let token_b_holding_address = Address::from(&PublicKey::new_from_private_key(&token_b_holding_key));
-        let token_b_definition_address = Address::from(&PublicKey::new_from_private_key(&token_b_definition_key));
-        let token_b_supply: u128 = 50000;
-
-        let token_lp_holding_key = PrivateKey::try_new([5;32]).unwrap();
-        let token_lp_definition_key = PrivateKey::try_new([6;32]).unwrap();
-        let token_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&token_lp_holding_key));
-        let token_lp_definition_address = Address::from(&PublicKey::new_from_private_key(&token_lp_definition_key));
-        let token_lp_supply: u128 = 300000;
-
-        let user_a_holding_key = PrivateKey::try_new([7;32]).unwrap();
-        let user_a_holding_address = Address::from(&PublicKey::new_from_private_key(&user_a_holding_key));
-        let user_a_amount: u128 = 10000;
-
-        let user_b_holding_key = PrivateKey::try_new([8;32]).unwrap();
-        let user_b_holding_address = Address::from(&PublicKey::new_from_private_key(&user_b_holding_key));
-        let user_b_amount: u128 = 10000;
-
-
-
-        let vault_a_key = PrivateKey::try_new([9;32]).unwrap();
-        let vault_a_address = Address::from(&PublicKey::new_from_private_key(&vault_a_key));
-
-        let vault_b_key = PrivateKey::try_new([10;32]).unwrap();
-        let vault_b_address = Address::from(&PublicKey::new_from_private_key(&vault_b_key));
-
-        let user_lp_holding_key = PrivateKey::try_new([11;32]).unwrap();
-        let user_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&user_lp_holding_key));
-
-        let pool_lp_holding_key = PrivateKey::try_new([12;32]).unwrap();
-        let pool_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&pool_lp_holding_key));
-
-        let pool_key = PrivateKey::try_new([13;32]).unwrap();
-        let pool_address = Address::from(&PublicKey::new_from_private_key(&pool_key));
-
-        //initialize Tokens: Token A, Token B and Took LP
-        let mut instruction: [u8;23] = [0;23];
-        instruction[1..17].copy_from_slice(&token_a_supply.to_le_bytes());
-        instruction[18] = 0x01; //name is not default.
-        instruction[19] = 0x02;
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_a_definition_address, token_a_holding_address],
-            vec![],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        instruction[1..17].copy_from_slice(&token_b_supply.to_le_bytes());
-        instruction[18] = 0x03; //name is not default.
-        instruction[19] = 0x02;
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_b_definition_address, token_b_holding_address],
-            vec![],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        instruction[1..17].copy_from_slice(&token_lp_supply.to_le_bytes());
-        instruction[18] = 0x03; //name is not default.
-        instruction[19] = 0x04;
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_lp_definition_address, token_lp_holding_address],
-            vec![],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-
-        // Initialize User accounts for Token A and B
-        let mut instruction: [u8;23] = [0;23];
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&user_a_amount.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_a_holding_address, user_a_holding_address],
-            vec![0],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_a_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&user_b_amount.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_b_holding_address, user_b_holding_address],
-            vec![0],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_b_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-
-        //TODO: initialize vaults - ideally, we won't need to do this.
-        let temp_amt = 1u128;
-        let mut instruction: [u8;23] = [0;23];
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_a_holding_address, vault_a_address],
-            vec![1],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_a_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_b_holding_address, vault_b_address],
-            vec![1],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_b_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        //need LP accounts for both AMM and user
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_lp_holding_address, user_lp_holding_address],
-            vec![0],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_lp_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        let amt: u128 = token_lp_supply - 1;
-        //Pool LP should be initialized at full amount (token_lp_supply)
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&amt.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_lp_holding_address, pool_lp_holding_address],
-            vec![1],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_lp_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        //Set up instruction to initialize AMM for (Token-A, Token-B)
-        let init_balance_a: u128 = 1000;
-        let init_balance_b: u128 = 1000;
-        let token_program_u8: [u8;32] = bytemuck::cast(program_token);
-        
-        let mut instruction: Vec<u8> = Vec::new();
-        instruction.push(0);
-        instruction.extend_from_slice(&init_balance_a.to_le_bytes());
-        instruction.extend_from_slice(&init_balance_b.to_le_bytes());
-        instruction.extend_from_slice(&token_program_u8);
-
-        let message = public_transaction::Message::try_new(
-            Program::amm().id(),
-            vec![
-                pool_address,
-                vault_a_address,
-                vault_b_address,
-                pool_lp_holding_address,
-                user_a_holding_address,
-                user_b_holding_address,
-                user_lp_holding_address
-            ],
-            vec![0,0,0,0],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[
-                &pool_key,
-                &pool_lp_holding_key,
-                &user_a_holding_key,
-                &user_b_holding_key,
-            ],
-        );
-        
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
+        let token_a_holding_address = vec_address[0];
+        let token_a_definition_address = vec_address[1];
+        let token_b_holding_address = vec_address[2];
+        let token_b_definition_address = vec_address[3];
+        let token_lp_definition_address = vec_address[4];
+        let user_a_holding_address = vec_address[5];
+        let user_b_holding_address = vec_address[6];
+        let vault_a_address = vec_address[7];
+        let vault_b_address = vec_address[8];
+        let user_lp_holding_address = vec_address[9];
+        let pool_address = vec_address[10];
+        let pool_lp_holding_address = vec_address[11];
 
         let mut instruction: Vec<u8> = Vec::new();
         instruction.push(3);
-        
+
         let message = public_transaction::Message::try_new(
             Program::amm().id(),
             vec![
@@ -2873,324 +2727,160 @@ impl PoolDefinition {
                 pool_lp_holding_address,
                 user_a_holding_address,
                 user_b_holding_address,
-                user_lp_holding_address
+                user_lp_holding_address,
             ],
             vec![
                 state.get_account_by_address(&pool_address).nonce,
                 state.get_account_by_address(&user_lp_holding_address).nonce,
                 state.get_account_by_address(&vault_a_address).nonce,
                 state.get_account_by_address(&vault_b_address).nonce,
-                ],
+            ],
             instruction,
         )
         .unwrap();
 
         let witness_set = public_transaction::WitnessSet::for_message(
             &message,
-            &[
-                &pool_key,
-                &user_lp_holding_key,
-                &vault_a_key,
-                &vault_b_key,
-            ],
+            &[&pool_key, &user_lp_holding_key, &vault_a_key, &vault_b_key],
         );
-        
+
         let tx = PublicTransaction::new(message, witness_set);
         state.transition_from_public_transaction(&tx).unwrap();
 
-
-
-        //TODO:
-        //user holdings
         let pool_post = state.get_account_by_address(&pool_address);
         let vault_a_post = state.get_account_by_address(&vault_a_address);
         let vault_b_post = state.get_account_by_address(&vault_b_address);
-        let pool_lp_post = state.get_account_by_address(&pool_lp_holding_address);
         let user_a_post = state.get_account_by_address(&user_a_holding_address);
         let user_b_post = state.get_account_by_address(&user_b_holding_address);
         let user_lp_post = state.get_account_by_address(&user_lp_holding_address);
 
-       //TODO: token accounts initialized with temp_amt balance.
-       // this ensures that amm program does not own any of the token accounts
-        assert!(TokenHolding::parse(&vault_a_post.data).unwrap().balance == temp_amt);        
-        assert!(TokenHolding::parse(&vault_b_post.data).unwrap().balance == temp_amt);  
+        //TODO: temp
+        //    let delta_lp : u128 = (pool_def_data.liquidity_pool_cap*user_lp_amt)/pool_def_data.liquidity_pool_cap;
+        let delta_lp : u128 = (init_balance_a*init_balance_a)/init_balance_a;
 
-        assert!(PoolDefinition::parse(&pool_post.data).unwrap().reserve_a == 0);
-        assert!(PoolDefinition::parse(&pool_post.data).unwrap().reserve_b == 0); 
-         assert!(TokenHolding::parse(&user_lp_post.data).unwrap().balance == 0);         
+        let expected_pool = Account {
+            program_owner: Program::amm().id(),
+            balance: 0u128,
+            data: PoolDefinition::into_data(
+                PoolDefinition {
+                    definition_token_a_id: token_a_definition_address,
+                    definition_token_b_id: token_b_definition_address,
+                    vault_a_addr: vault_a_address,
+                    vault_b_addr: vault_b_address,
+                    liquidity_pool_id: token_lp_definition_address,
+                    liquidity_pool_cap: PoolDefinition::parse(&pool_post.data).unwrap().liquidity_pool_cap, //TODOinit_balance_a - delta_lp,
+                    reserve_a: 0,
+                    reserve_b: 0,
+                    token_program_id: Program::token().id(),
+                }),
+            nonce: 2,
+        };
 
+        let expected_vault_a = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_a_definition_address,
+                    balance: temp_amt,
+            }),
+            nonce: 1
+        };
 
+        let expected_vault_b = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_b_definition_address,
+                    balance: temp_amt,
+            }),
+            nonce: 1
+        };
 
+        let expected_user_a = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_a_definition_address,
+                    balance: user_a_amount,
+            }),
+            nonce: 1
+        };
+
+        let expected_user_b = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_b_definition_address,
+                    balance: user_b_amount,
+            }),
+            nonce: 1
+        };
+
+        let expected_user_lp = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_lp_definition_address,
+                    balance: 0,
+            }),
+            nonce: 1
+        };
+
+        assert!(vault_a_post == expected_vault_a);
+        assert!(vault_b_post == expected_vault_b);
+        assert!(user_a_post == expected_user_a);
+        assert!(user_b_post == expected_user_b);
+        assert!(user_lp_post == expected_user_lp);
+        assert!(pool_post == expected_pool);
     }
 
-#[test]
-    fn test_simple_amm_initialized_attempt_add() {
-        let program_token = Program::token().id();
+    #[test]
+    fn test_simple_amm_add() {
+        let (state, vec_private_keys, vec_address, vec_amounts) = initialize_amm();
+        let mut state: V02State = state;
 
-        let initial_data = [];
-        let mut state =
-            V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
+        let temp_amt = vec_amounts[0];
+        let init_balance_a = vec_amounts[1];
+        let init_balance_b = vec_amounts[2];
+        let user_a_amount = vec_amounts[3];
+        let user_b_amount = vec_amounts[4];
 
+        let _token_a_holding_key = &vec_private_keys[0];
+        let _token_a_definition_key = &vec_private_keys[1];
+        let _token_b_holding_key = &vec_private_keys[2];
+        let _token_b_definition_key = &vec_private_keys[3];
+        let pool_lp_definition_key = &vec_private_keys[4];
+        let user_a_holding_key = &vec_private_keys[5];
+        let user_b_holding_key = &vec_private_keys[6];
+        let vault_a_key = &vec_private_keys[7];
+        let vault_b_key = &vec_private_keys[8];
+        let user_lp_holding_key = &vec_private_keys[9];
+        let pool_key = &vec_private_keys[10];
+        let pool_lp_holding_key = &vec_private_keys[11];
 
-        let token_a_holding_key = PrivateKey::try_new([1;32]).unwrap();
-        let token_a_definition_key = PrivateKey::try_new([2;32]).unwrap();
-        let token_a_holding_address = Address::from(&PublicKey::new_from_private_key(&token_a_holding_key));
-        let token_a_definition_address = Address::from(&PublicKey::new_from_private_key(&token_a_definition_key));
-        let token_a_supply: u128 = 30000;
+        let token_a_holding_address = vec_address[0];
+        let token_a_definition_address = vec_address[1];
+        let token_b_holding_address = vec_address[2];
+        let token_b_definition_address = vec_address[3];
+        let token_lp_definition_address = vec_address[4];
+        let user_a_holding_address = vec_address[5];
+        let user_b_holding_address = vec_address[6];
+        let vault_a_address = vec_address[7];
+        let vault_b_address = vec_address[8];
+        let user_lp_holding_address = vec_address[9];
+        let pool_address = vec_address[10];
+        let pool_lp_holding_address = vec_address[11];
 
-        let token_b_holding_key = PrivateKey::try_new([3;32]).unwrap();
-        let token_b_definition_key = PrivateKey::try_new([4;32]).unwrap();
-        let token_b_holding_address = Address::from(&PublicKey::new_from_private_key(&token_b_holding_key));
-        let token_b_definition_address = Address::from(&PublicKey::new_from_private_key(&token_b_definition_key));
-        let token_b_supply: u128 = 50000;
-
-        let token_lp_holding_key = PrivateKey::try_new([5;32]).unwrap();
-        let token_lp_definition_key = PrivateKey::try_new([6;32]).unwrap();
-        let token_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&token_lp_holding_key));
-        let token_lp_definition_address = Address::from(&PublicKey::new_from_private_key(&token_lp_definition_key));
-        let token_lp_supply: u128 = 300000;
-
-        let user_a_holding_key = PrivateKey::try_new([7;32]).unwrap();
-        let user_a_holding_address = Address::from(&PublicKey::new_from_private_key(&user_a_holding_key));
-        let user_a_amount: u128 = 10000;
-
-        let user_b_holding_key = PrivateKey::try_new([8;32]).unwrap();
-        let user_b_holding_address = Address::from(&PublicKey::new_from_private_key(&user_b_holding_key));
-        let user_b_amount: u128 = 10000;
-
-
-
-        let vault_a_key = PrivateKey::try_new([9;32]).unwrap();
-        let vault_a_address = Address::from(&PublicKey::new_from_private_key(&vault_a_key));
-
-        let vault_b_key = PrivateKey::try_new([10;32]).unwrap();
-        let vault_b_address = Address::from(&PublicKey::new_from_private_key(&vault_b_key));
-
-        let user_lp_holding_key = PrivateKey::try_new([11;32]).unwrap();
-        let user_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&user_lp_holding_key));
-
-        let pool_lp_holding_key = PrivateKey::try_new([12;32]).unwrap();
-        let pool_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&pool_lp_holding_key));
-
-        let pool_key = PrivateKey::try_new([13;32]).unwrap();
-        let pool_address = Address::from(&PublicKey::new_from_private_key(&pool_key));
-
-        //initialize Tokens: Token A, Token B and Took LP
-        let mut instruction: [u8;23] = [0;23];
-        instruction[1..17].copy_from_slice(&token_a_supply.to_le_bytes());
-        instruction[18] = 0x01; //name is not default.
-        instruction[19] = 0x02;
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_a_definition_address, token_a_holding_address],
-            vec![],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        instruction[1..17].copy_from_slice(&token_b_supply.to_le_bytes());
-        instruction[18] = 0x03; //name is not default.
-        instruction[19] = 0x02;
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_b_definition_address, token_b_holding_address],
-            vec![],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        instruction[1..17].copy_from_slice(&token_lp_supply.to_le_bytes());
-        instruction[18] = 0x03; //name is not default.
-        instruction[19] = 0x04;
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_lp_definition_address, token_lp_holding_address],
-            vec![],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-
-        // Initialize User accounts for Token A and B
-        let mut instruction: [u8;23] = [0;23];
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&user_a_amount.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_a_holding_address, user_a_holding_address],
-            vec![0],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_a_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&user_b_amount.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_b_holding_address, user_b_holding_address],
-            vec![0],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_b_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-
-        //TODO: initialize vaults - ideally, we won't need to do this.
-        let temp_amt = 1u128;
-        let mut instruction: [u8;23] = [0;23];
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_a_holding_address, vault_a_address],
-            vec![1],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_a_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_b_holding_address, vault_b_address],
-            vec![1],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_b_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        //need LP accounts for both AMM and user
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_lp_holding_address, user_lp_holding_address],
-            vec![0],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_lp_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        let amt: u128 = token_lp_supply - 1;
-        //Pool LP should be initialized at full amount (token_lp_supply)
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&amt.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_lp_holding_address, pool_lp_holding_address],
-            vec![1],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_lp_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        //Set up instruction to initialize AMM for (Token-A, Token-B)
-        let init_balance_a: u128 = 1000;
-        let init_balance_b: u128 = 1000;
-        let token_program_u8: [u8;32] = bytemuck::cast(program_token);
-        
-        let mut instruction: Vec<u8> = Vec::new();
-        instruction.push(0);
-        instruction.extend_from_slice(&init_balance_a.to_le_bytes());
-        instruction.extend_from_slice(&init_balance_b.to_le_bytes());
-        instruction.extend_from_slice(&token_program_u8);
-        
-        let message = public_transaction::Message::try_new(
-            Program::amm().id(),
-            vec![
-                pool_address,
-                vault_a_address,
-                vault_b_address,
-                pool_lp_holding_address,
-                user_a_holding_address,
-                user_b_holding_address,
-                user_lp_holding_address
-            ],
-            vec![0,0,0,0],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[
-                &pool_key,
-                &pool_lp_holding_key,
-                &user_a_holding_key,
-                &user_b_holding_key,
-            ],
-        );
-        
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
 
         let add_a: u128 = 500;
         let add_b: u128 = 500;
@@ -3200,7 +2890,7 @@ impl PoolDefinition {
         instruction.extend_from_slice(&add_a.to_le_bytes());
         instruction.extend_from_slice(&add_b.to_le_bytes());
         instruction.extend_from_slice(main_addr.value());
-        
+
         let message = public_transaction::Message::try_new(
             Program::amm().id(),
             vec![
@@ -3210,14 +2900,14 @@ impl PoolDefinition {
                 pool_lp_holding_address,
                 user_a_holding_address,
                 user_b_holding_address,
-                user_lp_holding_address
+                user_lp_holding_address,
             ],
             vec![
                 state.get_account_by_address(&pool_address).nonce,
                 state.get_account_by_address(&pool_lp_holding_address).nonce,
                 state.get_account_by_address(&user_a_holding_address).nonce,
                 state.get_account_by_address(&user_b_holding_address).nonce,
-                ],
+            ],
             instruction,
         )
         .unwrap();
@@ -3231,332 +2921,172 @@ impl PoolDefinition {
                 &user_b_holding_key,
             ],
         );
-        
+
         let tx = PublicTransaction::new(message, witness_set);
         state.transition_from_public_transaction(&tx).unwrap();
 
-
         let pool_post = state.get_account_by_address(&pool_address);
         let vault_a_post = state.get_account_by_address(&vault_a_address);
         let vault_b_post = state.get_account_by_address(&vault_b_address);
-        let pool_lp_post = state.get_account_by_address(&pool_lp_holding_address);
         let user_a_post = state.get_account_by_address(&user_a_holding_address);
         let user_b_post = state.get_account_by_address(&user_b_holding_address);
         let user_lp_post = state.get_account_by_address(&user_lp_holding_address);
 
-       //TODO: token accounts initialized with temp_amt balance.
-       // this ensures that amm program does not own any of the token accounts
-        assert!(TokenHolding::parse(&vault_a_post.data).unwrap().balance == init_balance_a + add_a + temp_amt);        
-        assert!(TokenHolding::parse(&vault_b_post.data).unwrap().balance == init_balance_b + add_b + temp_amt);  
+        let expected_pool = Account {
+            program_owner: Program::amm().id(),
+            balance: 0u128,
+            data: PoolDefinition::into_data(
+                PoolDefinition {
+                    definition_token_a_id: token_a_definition_address,
+                    definition_token_b_id: token_b_definition_address,
+                    vault_a_addr: vault_a_address,
+                    vault_b_addr: vault_b_address,
+                    liquidity_pool_id: token_lp_definition_address,
+                    liquidity_pool_cap: init_balance_a + add_a,
+                    reserve_a: init_balance_a + add_a,
+                    reserve_b: init_balance_b + add_b,
+                    token_program_id: Program::token().id(),
+                }),
+            nonce: 2,
+        };
 
-        assert!(PoolDefinition::parse(&pool_post.data).unwrap().reserve_a == init_balance_a + add_a);
-        assert!(PoolDefinition::parse(&pool_post.data).unwrap().reserve_b == init_balance_b + add_b); 
-        assert!(TokenHolding::parse(&user_lp_post.data).unwrap().balance == init_balance_a + add_a + temp_amt);
+        let expected_vault_a = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_a_definition_address,
+                    balance: init_balance_a + temp_amt + add_a,
+            }),
+            nonce: 0
+        };
+
+        let expected_vault_b = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_b_definition_address,
+                    balance: init_balance_b + temp_amt + add_b,
+            }),
+            nonce: 0
+        };
+
+        let expected_user_a = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_a_definition_address,
+                    balance: user_a_amount - init_balance_a - add_a,
+            }),
+            nonce: 2
+        };
+
+        let expected_user_b = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_b_definition_address,
+                    balance: user_b_amount - init_balance_b - add_b,
+            }),
+            nonce: 2
+        };
+
+        let expected_user_lp = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_lp_definition_address,
+                    balance: temp_amt + init_balance_a + add_a,
+            }),
+            nonce: 0
+        };
+
+        assert!(vault_a_post == expected_vault_a);
+        assert!(vault_b_post == expected_vault_b);
+        assert!(user_a_post == expected_user_a);
+        assert!(user_b_post == expected_user_b);
+        assert!(user_lp_post == expected_user_lp);
+        assert!(pool_post == expected_pool);
 
     }
+    
+    #[test]
+    fn test_simple_amm_swap_1() {
+        let (state, vec_private_keys, vec_address, vec_amounts) = initialize_amm();
+        let mut state: V02State = state;
 
+        let temp_amt = vec_amounts[0];
+        let init_balance_a = vec_amounts[1];
+        let init_balance_b = vec_amounts[2];
+        let user_a_amount = vec_amounts[3];
+        let user_b_amount = vec_amounts[4];
 
-#[test]
-    fn test_simple_amm_initialized_attempt_swap() {
-        let program_token = Program::token().id();
-
-        let initial_data = [];
-        let mut state =
-            V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
-
-
-        let token_a_holding_key = PrivateKey::try_new([1;32]).unwrap();
-        let token_a_definition_key = PrivateKey::try_new([2;32]).unwrap();
-        let token_a_holding_address = Address::from(&PublicKey::new_from_private_key(&token_a_holding_key));
-        let token_a_definition_address = Address::from(&PublicKey::new_from_private_key(&token_a_definition_key));
-        let token_a_supply: u128 = 30000;
-
-        let token_b_holding_key = PrivateKey::try_new([3;32]).unwrap();
-        let token_b_definition_key = PrivateKey::try_new([4;32]).unwrap();
-        let token_b_holding_address = Address::from(&PublicKey::new_from_private_key(&token_b_holding_key));
-        let token_b_definition_address = Address::from(&PublicKey::new_from_private_key(&token_b_definition_key));
-        let token_b_supply: u128 = 50000;
-
-        let token_lp_holding_key = PrivateKey::try_new([5;32]).unwrap();
-        let token_lp_definition_key = PrivateKey::try_new([6;32]).unwrap();
-        let token_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&token_lp_holding_key));
-        let token_lp_definition_address = Address::from(&PublicKey::new_from_private_key(&token_lp_definition_key));
-        let token_lp_supply: u128 = 300000;
-
-        let user_a_holding_key = PrivateKey::try_new([7;32]).unwrap();
-        let user_a_holding_address = Address::from(&PublicKey::new_from_private_key(&user_a_holding_key));
-        let user_a_amount: u128 = 10000;
-
-        let user_b_holding_key = PrivateKey::try_new([8;32]).unwrap();
-        let user_b_holding_address = Address::from(&PublicKey::new_from_private_key(&user_b_holding_key));
-        let user_b_amount: u128 = 10000;
-
-
-
-        let vault_a_key = PrivateKey::try_new([9;32]).unwrap();
-        let vault_a_address = Address::from(&PublicKey::new_from_private_key(&vault_a_key));
-
-        let vault_b_key = PrivateKey::try_new([10;32]).unwrap();
-        let vault_b_address = Address::from(&PublicKey::new_from_private_key(&vault_b_key));
-
-        let user_lp_holding_key = PrivateKey::try_new([11;32]).unwrap();
-        let user_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&user_lp_holding_key));
-
-        let pool_lp_holding_key = PrivateKey::try_new([12;32]).unwrap();
-        let pool_lp_holding_address = Address::from(&PublicKey::new_from_private_key(&pool_lp_holding_key));
-
-        let pool_key = PrivateKey::try_new([13;32]).unwrap();
-        let pool_address = Address::from(&PublicKey::new_from_private_key(&pool_key));
-
-        //initialize Tokens: Token A, Token B and Took LP
-        let mut instruction: [u8;23] = [0;23];
-        instruction[1..17].copy_from_slice(&token_a_supply.to_le_bytes());
-        instruction[18] = 0x01; //name is not default.
-        instruction[19] = 0x02;
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_a_definition_address, token_a_holding_address],
-            vec![],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        instruction[1..17].copy_from_slice(&token_b_supply.to_le_bytes());
-        instruction[18] = 0x03; //name is not default.
-        instruction[19] = 0x02;
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_b_definition_address, token_b_holding_address],
-            vec![],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        instruction[1..17].copy_from_slice(&token_lp_supply.to_le_bytes());
-        instruction[18] = 0x03; //name is not default.
-        instruction[19] = 0x04;
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_lp_definition_address, token_lp_holding_address],
-            vec![],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-
-        // Initialize User accounts for Token A and B
-        let mut instruction: [u8;23] = [0;23];
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&user_a_amount.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_a_holding_address, user_a_holding_address],
-            vec![0],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_a_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&user_b_amount.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_b_holding_address, user_b_holding_address],
-            vec![0],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_b_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-
-        //TODO: initialize vaults - ideally, we won't need to do this.
-        let temp_amt = 1u128;
-        let mut instruction: [u8;23] = [0;23];
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_a_holding_address, vault_a_address],
-            vec![1],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_a_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_b_holding_address, vault_b_address],
-            vec![1],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_b_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        //need LP accounts for both AMM and user
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&temp_amt.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_lp_holding_address, user_lp_holding_address],
-            vec![0],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_lp_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        let amt: u128 = token_lp_supply - 1;
-        //Pool LP should be initialized at full amount (token_lp_supply)
-        instruction[0] = 1; //transfer
-        instruction[1..17].copy_from_slice(&amt.to_le_bytes());
-
-        let message = public_transaction::Message::try_new(
-            Program::token().id(),
-            vec![token_lp_holding_address, pool_lp_holding_address],
-            vec![1],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_lp_holding_key],
-        );
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
-
-        //Set up instruction to initialize AMM for (Token-A, Token-B)
-        let init_balance_a: u128 = 1000;
-        let init_balance_b: u128 = 1000;
-        let token_program_u8: [u8;32] = bytemuck::cast(program_token);
-        
-        let mut instruction: Vec<u8> = Vec::new();
-        instruction.push(0);
-        instruction.extend_from_slice(&init_balance_a.to_le_bytes());
-        instruction.extend_from_slice(&init_balance_b.to_le_bytes());
-        instruction.extend_from_slice(&token_program_u8);
-        
-        let message = public_transaction::Message::try_new(
-            Program::amm().id(),
-            vec![
-                pool_address,
-                vault_a_address,
-                vault_b_address,
-                pool_lp_holding_address,
-                user_a_holding_address,
-                user_b_holding_address,
-                user_lp_holding_address
-            ],
-            vec![0,0,0,0],
-            instruction,
-        )
-        .unwrap();
-
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[
-                &pool_key,
-                &pool_lp_holding_key,
-                &user_a_holding_key,
-                &user_b_holding_key,
-            ],
-        );
-        
-        let tx = PublicTransaction::new(message, witness_set);
-        state.transition_from_public_transaction(&tx).unwrap();
+        let token_a_holding_key = &vec_private_keys[0];
+        let token_a_definition_key = &vec_private_keys[1];
+        let token_b_holding_key = &vec_private_keys[2];
+        let token_b_definition_key = &vec_private_keys[3];
+        let pool_lp_definition_key = &vec_private_keys[4];
+        let user_a_holding_key = &vec_private_keys[5];
+        let user_b_holding_key = &vec_private_keys[6];
+        let vault_a_key = &vec_private_keys[7];
+        let vault_b_key = &vec_private_keys[8];
+        let user_lp_holding_key = &vec_private_keys[9];
+        let pool_key = &vec_private_keys[10];
+        let pool_lp_holding_key = &vec_private_keys[11];
 
+        let token_a_holding_address = vec_address[0];
+        let token_a_definition_address = vec_address[1];
+        let token_b_holding_address = vec_address[2];
+        let token_b_definition_address = vec_address[3];
+        let token_lp_definition_address = vec_address[4];
+        let user_a_holding_address = vec_address[5];
+        let user_b_holding_address = vec_address[6];
+        let vault_a_address = vec_address[7];
+        let vault_b_address = vec_address[8];
+        let user_lp_holding_address = vec_address[9];
+        let pool_address = vec_address[10];
+        let pool_lp_holding_address = vec_address[11];
 
         //Initialize swap user accounts
-        let swap_user_a_holding_key = PrivateKey::try_new([7;32]).unwrap();
-        let swap_user_a_holding_address = Address::from(&PublicKey::new_from_private_key(&user_a_holding_key));
-        let swap_user_a_amount: u128 = 10000;
+        let swap_user_a_holding_key = PrivateKey::try_new([21; 32]).unwrap();
+        let swap_user_a_holding_address =
+            Address::from(&PublicKey::new_from_private_key(&swap_user_a_holding_key));
+        let swap_user_a_amount: u128 = 5000;
 
-        let swap_user_b_holding_key = PrivateKey::try_new([8;32]).unwrap();
-        let swap_user_b_holding_address = Address::from(&PublicKey::new_from_private_key(&user_b_holding_key));
-        let swap_user_b_amount: u128 = 10000;
+        let swap_user_b_holding_key = PrivateKey::try_new([22; 32]).unwrap();
+        let swap_user_b_holding_address =
+            Address::from(&PublicKey::new_from_private_key(&swap_user_b_holding_key));
+        let swap_user_b_amount: u128 = 5000;
 
-        // Initialize Swap User accounts for Token A and B
-        let mut instruction: [u8;23] = [0;23];
+        // Initialize Swap User account for Token A
+        let mut instruction: [u8; 23] = [0; 23];
         instruction[0] = 1; //transfer
         instruction[1..17].copy_from_slice(&swap_user_a_amount.to_le_bytes());
 
         let message = public_transaction::Message::try_new(
             Program::token().id(),
             vec![token_a_holding_address, swap_user_a_holding_address],
-            vec![state.get_account_by_address(&token_a_holding_address).nonce],
+            vec![2],
             instruction,
         )
         .unwrap();
 
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_a_holding_key],
-        );
+        let witness_set =
+            public_transaction::WitnessSet::for_message(&message, &[&token_a_holding_key]);
         let tx = PublicTransaction::new(message, witness_set);
         state.transition_from_public_transaction(&tx).unwrap();
 
-
-        let mut instruction: [u8;23] = [0;23];
+        // Initialize Swap User account for Token B
+        let mut instruction: [u8; 23] = [0; 23];
         instruction[0] = 1; //transfer
         instruction[1..17].copy_from_slice(&swap_user_b_amount.to_le_bytes());
 
@@ -3568,14 +3098,12 @@ impl PoolDefinition {
         )
         .unwrap();
 
-        let witness_set = public_transaction::WitnessSet::for_message(
-            &message,
-            &[&token_b_holding_key],
-        );
+        let witness_set =
+            public_transaction::WitnessSet::for_message(&message, &[&token_b_holding_key]);
         let tx = PublicTransaction::new(message, witness_set);
         state.transition_from_public_transaction(&tx).unwrap();
 
-
+        // Initialize Swap
         let main_addr = token_a_definition_address;
         let swap_a: u128 = 500;
 
@@ -3584,7 +3112,6 @@ impl PoolDefinition {
         instruction.extend_from_slice(&swap_a.to_le_bytes());
         instruction.extend_from_slice(main_addr.value());
 
-   
         let message = public_transaction::Message::try_new(
             Program::amm().id(),
             vec![
@@ -3592,15 +3119,19 @@ impl PoolDefinition {
                 vault_a_address,
                 vault_b_address,
                 swap_user_a_holding_address,
-                swap_user_b_holding_address
+                swap_user_b_holding_address,
             ],
             vec![
                 state.get_account_by_address(&pool_address).nonce,
                 state.get_account_by_address(&vault_a_address).nonce,
                 state.get_account_by_address(&vault_b_address).nonce,
-                state.get_account_by_address(&swap_user_a_holding_address).nonce,
-                state.get_account_by_address(&swap_user_b_holding_address).nonce,
-                ],
+                state
+                    .get_account_by_address(&swap_user_a_holding_address)
+                    .nonce,
+                state
+                    .get_account_by_address(&swap_user_b_holding_address)
+                    .nonce,
+            ],
             instruction,
         )
         .unwrap();
@@ -3615,28 +3146,518 @@ impl PoolDefinition {
                 &swap_user_b_holding_key,
             ],
         );
-        
+
         let tx = PublicTransaction::new(message, witness_set);
         state.transition_from_public_transaction(&tx).unwrap();
 
-        //TODO:
-        //Check AMM pairing:
-        // reserve_a = vault_a
-        // reserve_b = vault_b   
+
         let pool_post = state.get_account_by_address(&pool_address);
         let vault_a_post = state.get_account_by_address(&vault_a_address);
         let vault_b_post = state.get_account_by_address(&vault_b_address);
-        let pool_lp_post = state.get_account_by_address(&pool_lp_holding_address);
-        let swap_a_post = state.get_account_by_address(&swap_user_a_holding_address);
-        let swap_b_post = state.get_account_by_address(&swap_user_b_holding_address);
-        let user_lp_post = state.get_account_by_address(&user_lp_holding_address);
+        let swap_user_a_post = state.get_account_by_address(&swap_user_a_holding_address);
+        let swap_user_b_post = state.get_account_by_address(&swap_user_b_holding_address);
 
-       assert!(TokenHolding::parse(&swap_a_post.data).unwrap().balance == swap_user_a_amount - swap_a);        
-       assert!(TokenHolding::parse(&swap_b_post.data).unwrap().balance == swap_user_b_amount + swap_a);  
-        //todo: reserves reduced
 
+        let withdraw_b = (init_balance_b * swap_a)/(init_balance_a + swap_a);
+
+        let expected_pool = Account {
+            program_owner: Program::amm().id(),
+            balance: 0u128,
+            data: PoolDefinition::into_data(
+                PoolDefinition {
+                    definition_token_a_id: token_a_definition_address,
+                    definition_token_b_id: token_b_definition_address,
+                    vault_a_addr: vault_a_address,
+                    vault_b_addr: vault_b_address,
+                    liquidity_pool_id: token_lp_definition_address,
+                    liquidity_pool_cap: init_balance_a,
+                    reserve_a: init_balance_a + swap_a,
+                    reserve_b: init_balance_b - withdraw_b,
+                    token_program_id: Program::token().id(),
+                }),
+            nonce: 2,
+        };
+
+        let expected_vault_a = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_a_definition_address,
+                    balance: init_balance_a + temp_amt + swap_a,
+            }),
+            nonce: 1
+        };
+
+        let expected_vault_b = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_b_definition_address,
+                    balance: init_balance_b + temp_amt - withdraw_b,
+            }),
+            nonce: 1
+        };
+
+        let expected_swap_user_a = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_a_definition_address,
+                    balance: swap_user_a_amount - swap_a,
+            }),
+            nonce: 1
+        };
+
+        let expected_swap_user_b = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_b_definition_address,
+                    balance: swap_user_b_amount + withdraw_b,
+            }),
+            nonce: 1
+        };
+
+        assert!(vault_a_post == expected_vault_a);
+        assert!(vault_b_post == expected_vault_b);
+        assert!(swap_user_a_post == expected_swap_user_a);
+        assert!(swap_user_b_post == expected_swap_user_b);
+        assert!(pool_post == expected_pool);
+
+    }
+    
+    #[test]
+    fn test_simple_amm_swap_2() {
+        let (state, vec_private_keys, vec_address, vec_amounts) = initialize_amm();
+        let mut state: V02State = state;
+
+        let temp_amt = vec_amounts[0];
+        let init_balance_a = vec_amounts[1];
+        let init_balance_b = vec_amounts[2];
+        let user_a_amount = vec_amounts[3];
+        let user_b_amount = vec_amounts[4];
+
+        let token_a_holding_key = &vec_private_keys[0];
+        let token_a_definition_key = &vec_private_keys[1];
+        let token_b_holding_key = &vec_private_keys[2];
+        let token_b_definition_key = &vec_private_keys[3];
+        let pool_lp_definition_key = &vec_private_keys[4];
+        let user_a_holding_key = &vec_private_keys[5];
+        let user_b_holding_key = &vec_private_keys[6];
+        let vault_a_key = &vec_private_keys[7];
+        let vault_b_key = &vec_private_keys[8];
+        let user_lp_holding_key = &vec_private_keys[9];
+        let pool_key = &vec_private_keys[10];
+        let pool_lp_holding_key = &vec_private_keys[11];
+
+        let token_a_holding_address = vec_address[0];
+        let token_a_definition_address = vec_address[1];
+        let token_b_holding_address = vec_address[2];
+        let token_b_definition_address = vec_address[3];
+        let token_lp_definition_address = vec_address[4];
+        let user_a_holding_address = vec_address[5];
+        let user_b_holding_address = vec_address[6];
+        let vault_a_address = vec_address[7];
+        let vault_b_address = vec_address[8];
+        let user_lp_holding_address = vec_address[9];
+        let pool_address = vec_address[10];
+        let pool_lp_holding_address = vec_address[11];
+
+        //Initialize swap user accounts
+        let swap_user_a_holding_key = PrivateKey::try_new([21; 32]).unwrap();
+        let swap_user_a_holding_address =
+            Address::from(&PublicKey::new_from_private_key(&swap_user_a_holding_key));
+        let swap_user_a_amount: u128 = 5000;
+
+        let swap_user_b_holding_key = PrivateKey::try_new([22; 32]).unwrap();
+        let swap_user_b_holding_address =
+            Address::from(&PublicKey::new_from_private_key(&swap_user_b_holding_key));
+        let swap_user_b_amount: u128 = 5000;
+
+        // Initialize Swap User account for Token A
+        let mut instruction: [u8; 23] = [0; 23];
+        instruction[0] = 1; //transfer
+        instruction[1..17].copy_from_slice(&swap_user_a_amount.to_le_bytes());
+
+        let message = public_transaction::Message::try_new(
+            Program::token().id(),
+            vec![token_a_holding_address, swap_user_a_holding_address],
+            vec![2],
+            instruction,
+        )
+        .unwrap();
+
+        let witness_set =
+            public_transaction::WitnessSet::for_message(&message, &[&token_a_holding_key]);
+        let tx = PublicTransaction::new(message, witness_set);
+        state.transition_from_public_transaction(&tx).unwrap();
+
+        // Initialize Swap User account for Token B
+        let mut instruction: [u8; 23] = [0; 23];
+        instruction[0] = 1; //transfer
+        instruction[1..17].copy_from_slice(&swap_user_b_amount.to_le_bytes());
+
+        let message = public_transaction::Message::try_new(
+            Program::token().id(),
+            vec![token_b_holding_address, swap_user_b_holding_address],
+            vec![state.get_account_by_address(&token_b_holding_address).nonce],
+            instruction,
+        )
+        .unwrap();
+
+        let witness_set =
+            public_transaction::WitnessSet::for_message(&message, &[&token_b_holding_key]);
+        let tx = PublicTransaction::new(message, witness_set);
+        state.transition_from_public_transaction(&tx).unwrap();
+
+        // Swap
+        let main_addr = token_b_definition_address;
+        let swap_b: u128 = 500;
+
+        let mut instruction: Vec<u8> = Vec::new();
+        instruction.push(1);
+        instruction.extend_from_slice(&swap_b.to_le_bytes());
+        instruction.extend_from_slice(main_addr.value());
+
+        let message = public_transaction::Message::try_new(
+            Program::amm().id(),
+            vec![
+                pool_address,
+                vault_a_address,
+                vault_b_address,
+                swap_user_a_holding_address,
+                swap_user_b_holding_address,
+            ],
+            vec![
+                state.get_account_by_address(&pool_address).nonce,
+                state.get_account_by_address(&vault_a_address).nonce,
+                state.get_account_by_address(&vault_b_address).nonce,
+                state
+                    .get_account_by_address(&swap_user_a_holding_address)
+                    .nonce,
+                state
+                    .get_account_by_address(&swap_user_b_holding_address)
+                    .nonce,
+            ],
+            instruction,
+        )
+        .unwrap();
+
+        let witness_set = public_transaction::WitnessSet::for_message(
+            &message,
+            &[
+                &pool_key,
+                &vault_a_key,
+                &vault_b_key,
+                &swap_user_a_holding_key,
+                &swap_user_b_holding_key,
+            ],
+        );
+
+        let tx = PublicTransaction::new(message, witness_set);
+        state.transition_from_public_transaction(&tx).unwrap();
+
+        let pool_post = state.get_account_by_address(&pool_address);
+        let vault_a_post = state.get_account_by_address(&vault_a_address);
+        let vault_b_post = state.get_account_by_address(&vault_b_address);
+        let swap_user_a_post = state.get_account_by_address(&swap_user_a_holding_address);
+        let swap_user_b_post = state.get_account_by_address(&swap_user_b_holding_address);
+
+        let withdraw_a = (init_balance_a * swap_b)/(init_balance_b + swap_b);
+
+        let expected_pool = Account {
+            program_owner: Program::amm().id(),
+            balance: 0u128,
+            data: PoolDefinition::into_data(
+                PoolDefinition {
+                    definition_token_a_id: token_a_definition_address,
+                    definition_token_b_id: token_b_definition_address,
+                    vault_a_addr: vault_a_address,
+                    vault_b_addr: vault_b_address,
+                    liquidity_pool_id: token_lp_definition_address,
+                    liquidity_pool_cap: init_balance_a,
+                    reserve_a: init_balance_a - withdraw_a,
+                    reserve_b: init_balance_b + swap_b,
+                    token_program_id: Program::token().id(),
+                }),
+            nonce: 2,
+        };
+
+        let expected_vault_a = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_a_definition_address,
+                    balance: init_balance_a + temp_amt - withdraw_a,
+            }),
+            nonce: 1
+        };
+
+        let expected_vault_b = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_b_definition_address,
+                    balance: init_balance_b + temp_amt + swap_b,
+            }),
+            nonce: 1
+        };
+
+        let expected_swap_user_a = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_a_definition_address,
+                    balance: swap_user_a_amount + withdraw_a,
+            }),
+            nonce: 1
+        };
+
+        let expected_swap_user_b = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_b_definition_address,
+                    balance: swap_user_b_amount - swap_b,
+            }),
+            nonce: 1
+        };
+
+        assert!(vault_a_post == expected_vault_a);
+        assert!(vault_b_post == expected_vault_b);
+        assert!(swap_user_a_post == expected_swap_user_a);
+        assert!(swap_user_b_post == expected_swap_user_b);
+        assert!(pool_post == expected_pool);
 
     }
 
+        #[test]
+    fn test_simple_amm_swap_3() {
+        let (state, vec_private_keys, vec_address, vec_amounts) = initialize_amm();
+        let mut state: V02State = state;
 
+        let temp_amt = vec_amounts[0];
+        let init_balance_a = vec_amounts[1];
+        let init_balance_b = vec_amounts[2];
+        let user_a_amount = vec_amounts[3];
+        let user_b_amount = vec_amounts[4];
+
+        let token_a_holding_key = &vec_private_keys[0];
+        let token_a_definition_key = &vec_private_keys[1];
+        let token_b_holding_key = &vec_private_keys[2];
+        let token_b_definition_key = &vec_private_keys[3];
+        let pool_lp_definition_key = &vec_private_keys[4];
+        let user_a_holding_key = &vec_private_keys[5];
+        let user_b_holding_key = &vec_private_keys[6];
+        let vault_a_key = &vec_private_keys[7];
+        let vault_b_key = &vec_private_keys[8];
+        let user_lp_holding_key = &vec_private_keys[9];
+        let pool_key = &vec_private_keys[10];
+        let pool_lp_holding_key = &vec_private_keys[11];
+
+        let token_a_holding_address = vec_address[0];
+        let token_a_definition_address = vec_address[1];
+        let token_b_holding_address = vec_address[2];
+        let token_b_definition_address = vec_address[3];
+        let token_lp_definition_address = vec_address[4];
+        let user_a_holding_address = vec_address[5];
+        let user_b_holding_address = vec_address[6];
+        let vault_a_address = vec_address[7];
+        let vault_b_address = vec_address[8];
+        let user_lp_holding_address = vec_address[9];
+        let pool_address = vec_address[10];
+        let pool_lp_holding_address = vec_address[11];
+
+        //Initialize swap user accounts
+        let swap_user_a_holding_key = PrivateKey::try_new([21; 32]).unwrap();
+        let swap_user_a_holding_address =
+            Address::from(&PublicKey::new_from_private_key(&swap_user_a_holding_key));
+        let swap_user_a_amount: u128 = 5000;
+
+        let swap_user_b_holding_key = PrivateKey::try_new([22; 32]).unwrap();
+        let swap_user_b_holding_address =
+            Address::from(&PublicKey::new_from_private_key(&swap_user_b_holding_key));
+        let swap_user_b_amount: u128 = 5000;
+
+        // Initialize Swap User account for Token A
+        let mut instruction: [u8; 23] = [0; 23];
+        instruction[0] = 1; //transfer
+        instruction[1..17].copy_from_slice(&swap_user_a_amount.to_le_bytes());
+
+        let message = public_transaction::Message::try_new(
+            Program::token().id(),
+            vec![token_a_holding_address, swap_user_a_holding_address],
+            vec![2],
+            instruction,
+        )
+        .unwrap();
+
+        let witness_set =
+            public_transaction::WitnessSet::for_message(&message, &[&token_a_holding_key]);
+        let tx = PublicTransaction::new(message, witness_set);
+        state.transition_from_public_transaction(&tx).unwrap();
+
+        // Initialize Swap User account for Token B
+        let mut instruction: [u8; 23] = [0; 23];
+        instruction[0] = 1; //transfer
+        instruction[1..17].copy_from_slice(&swap_user_b_amount.to_le_bytes());
+
+        let message = public_transaction::Message::try_new(
+            Program::token().id(),
+            vec![token_b_holding_address, swap_user_b_holding_address],
+            vec![state.get_account_by_address(&token_b_holding_address).nonce],
+            instruction,
+        )
+        .unwrap();
+
+        let witness_set =
+            public_transaction::WitnessSet::for_message(&message, &[&token_b_holding_key]);
+        let tx = PublicTransaction::new(message, witness_set);
+        state.transition_from_public_transaction(&tx).unwrap();
+
+        // Swap
+        let main_addr = token_b_definition_address;
+        let swap_b: u128 = 500;
+
+        let mut instruction: Vec<u8> = Vec::new();
+        instruction.push(1);
+        instruction.extend_from_slice(&swap_b.to_le_bytes());
+        instruction.extend_from_slice(main_addr.value());
+
+        //order of vault_a and vault_b are swapped.
+        let message = public_transaction::Message::try_new(
+            Program::amm().id(),
+            vec![
+                pool_address,
+                vault_b_address,
+                vault_a_address,
+                swap_user_a_holding_address,
+                swap_user_b_holding_address,
+            ],
+            vec![
+                state.get_account_by_address(&pool_address).nonce,
+                state.get_account_by_address(&vault_b_address).nonce,
+                state.get_account_by_address(&vault_a_address).nonce,
+                state
+                    .get_account_by_address(&swap_user_a_holding_address)
+                    .nonce,
+                state
+                    .get_account_by_address(&swap_user_b_holding_address)
+                    .nonce,
+            ],
+            instruction,
+        )
+        .unwrap();
+
+        let witness_set = public_transaction::WitnessSet::for_message(
+            &message,
+            &[
+                &pool_key,
+                &vault_a_key,
+                &vault_b_key,
+                &swap_user_a_holding_key,
+                &swap_user_b_holding_key,
+            ],
+        );
+
+        let tx = PublicTransaction::new(message, witness_set);
+        state.transition_from_public_transaction(&tx).unwrap();
+
+        let pool_post = state.get_account_by_address(&pool_address);
+        let vault_a_post = state.get_account_by_address(&vault_a_address);
+        let vault_b_post = state.get_account_by_address(&vault_b_address);
+        let swap_user_a_post = state.get_account_by_address(&swap_user_a_holding_address);
+        let swap_user_b_post = state.get_account_by_address(&swap_user_b_holding_address);
+
+        let withdraw_a = (init_balance_a * swap_b)/(init_balance_b + swap_b);
+
+        let expected_pool = Account {
+            program_owner: Program::amm().id(),
+            balance: 0u128,
+            data: PoolDefinition::into_data(
+                PoolDefinition {
+                    definition_token_a_id: token_a_definition_address,
+                    definition_token_b_id: token_b_definition_address,
+                    vault_a_addr: vault_a_address,
+                    vault_b_addr: vault_b_address,
+                    liquidity_pool_id: token_lp_definition_address,
+                    liquidity_pool_cap: init_balance_a,
+                    reserve_a: init_balance_a - withdraw_a,
+                    reserve_b: init_balance_b + swap_b,
+                    token_program_id: Program::token().id(),
+                }),
+            nonce: 2,
+        };
+
+        let expected_vault_a = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_a_definition_address,
+                    balance: init_balance_a + temp_amt - withdraw_a,
+            }),
+            nonce: 1
+        };
+
+        let expected_vault_b = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_b_definition_address,
+                    balance: init_balance_b + temp_amt + swap_b,
+            }),
+            nonce: 1
+        };
+
+        let expected_swap_user_a = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_a_definition_address,
+                    balance: swap_user_a_amount + withdraw_a,
+            }),
+            nonce: 1
+        };
+
+        let expected_swap_user_b = Account {
+            program_owner: Program::token().id(),
+            balance: 0u128,
+            data: TokenHolding::into_data(
+                TokenHolding{
+                    account_type: TOKEN_HOLDING_TYPE,
+                    definition_id: token_b_definition_address,
+                    balance: swap_user_b_amount - swap_b,
+            }),
+            nonce: 1
+        };
+
+        assert!(vault_a_post == expected_vault_a);
+        assert!(vault_b_post == expected_vault_b);
+        assert!(swap_user_a_post == expected_swap_user_a);
+        assert!(swap_user_b_post == expected_swap_user_b);
+        assert!(pool_post == expected_pool);
+
+    }
 }