diff --git a/token-swap/program/src/processor.rs b/token-swap/program/src/processor.rs
index fae4d9f0..19717ce9 100644
--- a/token-swap/program/src/processor.rs
+++ b/token-swap/program/src/processor.rs
@@ -1145,13 +1145,13 @@ mod tests {
             withdraw_all_token_types, withdraw_single_token_type_exact_amount_out,
         },
     };
-    use solana_program::{instruction::Instruction, program_pack::Pack, program_stubs, rent::Rent};
+    use solana_program::{clock::Clock, entrypoint::SUCCESS, instruction::Instruction, program_pack::Pack, program_stubs, rent::Rent};
     use solana_sdk::account::{
         create_account_for_test, create_is_signer_account_infos, Account as SolanaAccount,
     };
     use spl_token_2022::{
         error::TokenError,
-        extension::ExtensionType,
+        extension::{transfer_fee::instruction::initialize_transfer_fee_config, ExtensionType},
         instruction::{
             approve, close_account, freeze_account, initialize_account, initialize_immutable_owner,
             initialize_mint, initialize_mint_close_authority, mint_to, revoke, set_authority,
@@ -1216,6 +1216,13 @@ mod tests {
                 Err(ProgramError::IncorrectProgramId)
             }
         }
+
+        fn sol_get_clock_sysvar(&self, var_addr: *mut u8) -> u64 {
+            unsafe {
+                *(var_addr as *mut _ as *mut Clock) = Clock::default();
+            }
+            SUCCESS
+        }
     }
 
     fn test_syscall_stubs() {
@@ -1920,7 +1927,10 @@ mod tests {
     ) -> (Pubkey, SolanaAccount) {
         let account_key = Pubkey::new_unique();
         let space = if *program_id == spl_token_2022::id() {
-            ExtensionType::get_account_len::<Account>(&[ExtensionType::ImmutableOwner])
+            ExtensionType::get_account_len::<Account>(&[
+                ExtensionType::ImmutableOwner,
+                ExtensionType::TransferFeeAmount,
+            ])
         } else {
             Account::get_packed_len()
         };
@@ -1977,11 +1987,12 @@ mod tests {
     ) -> (Pubkey, SolanaAccount) {
         let mint_key = Pubkey::new_unique();
         let space = if *program_id == spl_token_2022::id() {
-            ExtensionType::get_account_len::<spl_token_2022::state::Mint>(&[
+            ExtensionType::get_account_len::<Mint>(&[
                 ExtensionType::MintCloseAuthority,
+                ExtensionType::TransferFeeConfig,
             ])
         } else {
-            spl_token_2022::state::Mint::get_packed_len()
+            Mint::get_packed_len()
         };
         let minimum_balance = Rent::default().minimum_balance(space);
         let mut mint_account = SolanaAccount::new(minimum_balance, space, program_id);
@@ -1993,6 +2004,19 @@ mod tests {
                 vec![&mut mint_account],
             )
             .unwrap();
+            do_process_instruction(
+                initialize_transfer_fee_config(
+                    program_id,
+                    &mint_key,
+                    freeze_authority,
+                    freeze_authority,
+                    0,
+                    0,
+                )
+                .unwrap(),
+                vec![&mut mint_account],
+            )
+            .unwrap();
         }
         do_process_instruction(
             initialize_mint(program_id, &mint_key, authority_key, freeze_authority, 2).unwrap(),