From aa656caeee9d29b52452549db877b0651ca0655e Mon Sep 17 00:00:00 2001
From: Kaze <230549489+kaze-cow@users.noreply.github.com>
Date: Mon, 7 Sep 2026 23:27:00 +0900
Subject: [PATCH] Rerun the buffer suite against Token-2022

The Token-2022 paths are worth little if they only hold for the handful of
cases someone thought to write twice, so instead of a parallel suite this
reruns the suite that already exists against the second program.

`common::also_under_token_2022!(some_test)` sits in front of a test and
generates `some_test_token_2022`, which runs the same body with Token-2022 as
the thread-local active program. Nothing in the body changes: the token helpers
build against the program that owns the account they are handed, and
`payer_signed_tx` / `signed_tx` repoint the legacy program id in every
instruction they assemble. 19 tests are covered this way, across buffer
creation and reclamation -- error paths as much as happy ones.

`CreateBuffer` and `ReclaimBuffer` name their token program as an account, so
repointing reaches them. `common::buffer::ensure_buffer_exists` reads the
program off the mint instead of taking it as a parameter, a buffer being a
token account of its mint and so bound to the mint's own program; that also
lets one test build buffers under both programs at once, which is what
`ensure_buffer_exists_for` used to be for.

Naming the test in front of it, rather than wrapping the body, keeps the
indentation and makes a stale name a compile error instead of a test that
quietly stopped being generated. A test that can only hold under one program
goes without and says why -- the one pinned to the legacy native mint.

Three tests are Token-2022-only, covering what has no legacy analogue -- a
buffer for a mint with a `TransferFeeConfig`, which needs a `TransferFeeAmount`
on every account holding it and so must be longer than the base layout:

- `creates_buffer_sized_for_a_mint_with_extensions`
- `recreating_an_extension_mint_buffer_is_idempotent`
- `reclaims_a_buffer_sized_for_an_extension_mint`

`bench-report.json` gains the CU, account, and transaction-byte readings for
every generated test that is benched.

The settlement pair is deliberately left out. `BeginSettle` and `FinalizeSettle`
name their token programs by value rather than by account, so rerunning their
tests needs the slot handling from #128 and belongs on that branch.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 Cargo.lock                                    |   2 +
 Cargo.toml                                    |   1 +
 bench-report.json                             |  21 ++
 programs/settlement/Cargo.toml                |   2 +
 programs/settlement/src/create_buffer.rs      |   2 +-
 programs/settlement/src/reclaim_buffer.rs     |   2 +-
 programs/settlement/tests/common/buffer.rs    |  18 +-
 programs/settlement/tests/common/mod.rs       |  70 ++++-
 programs/settlement/tests/common/token.rs     | 275 +++++++++++++++---
 .../settlement/tests/common/token_2022.rs     |  12 +-
 programs/settlement/tests/create_buffer.rs    |  66 +++--
 programs/settlement/tests/reclaim_buffer.rs   |  31 +-
 12 files changed, 402 insertions(+), 100 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 7428009..4536432 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3327,6 +3327,7 @@ dependencies = [
  "jsonschema",
  "litesvm",
  "litesvm-token",
+ "pastey",
  "pinocchio",
  "pinocchio-system",
  "pinocchio-token",
@@ -3341,6 +3342,7 @@ dependencies = [
  "solana-sdk",
  "solana-sha256-hasher",
  "solana-system-interface 3.2.0",
+ "spl-associated-token-account-interface",
  "spl-token-2022-interface 3.1.1",
  "syn 2.0.119",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index e691330..610866c 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -31,6 +31,7 @@ jsonschema = "0.30"
 litesvm = "0.15"
 litesvm-token = "0.15"
 num_enum = "0.7"
+pastey = "0.2"
 pinocchio = "0.11"
 pinocchio-system = "0.6"
 pinocchio-token = "0.7"
diff --git a/bench-report.json b/bench-report.json
index 8048980..e0a0e2a 100644
--- a/bench-report.json
+++ b/bench-report.json
@@ -3,14 +3,21 @@
     "add_solver/add_with_many_existing_solvers": 5,
     "add_solver/adds_a_solver": 5,
     "create_buffers/happy_path_creates_initialized_buffer_token_account": 6,
+    "create_buffers/happy_path_creates_initialized_buffer_token_account_token_2022": 6,
     "create_buffers/happy_path_creates_multiple_buffers_in_one_instruction": 10,
+    "create_buffers/happy_path_creates_multiple_buffers_in_one_instruction_token_2022": 10,
     "create_buffers/max_buffers_in_one_instruction": 64,
+    "create_buffers/max_buffers_in_one_instruction_token_2022": 46,
     "create_order/happy_path_creates_order_pda_with_expected_body": 4,
     "initialize/happy_path_initializes_state_pda_with_expected_data": 4,
     "reclaim_buffer/funded_buffer_is_skipped": 7,
+    "reclaim_buffer/funded_buffer_is_skipped_token_2022": 7,
     "reclaim_buffer/happy_path_reclaims_empty_buffer_to_the_authority_itself": 7,
+    "reclaim_buffer/happy_path_reclaims_empty_buffer_to_the_authority_itself_token_2022": 7,
     "reclaim_buffer/max_buffers_in_one_instruction": 64,
+    "reclaim_buffer/max_buffers_in_one_instruction_token_2022": 64,
     "reclaim_buffer/reclaims_multiple_buffers_skipping_funded": 9,
+    "reclaim_buffer/reclaims_multiple_buffers_skipping_funded_token_2022": 9,
     "reclaim_order/happy_path_expired_returns_lamports_and_closes_pda": 4,
     "reclaim_order/happy_path_on_chain_order_cancelled_is_reclaimable_before_expiry": 3,
     "reclaim_order/happy_path_on_chain_order_fully_filled_is_reclaimable_before_expiry": 3,
@@ -35,14 +42,21 @@
     "add_solver/add_with_many_existing_solvers": 5074,
     "add_solver/adds_a_solver": 4622,
     "create_buffers/happy_path_creates_initialized_buffer_token_account": 7373,
+    "create_buffers/happy_path_creates_initialized_buffer_token_account_token_2022": 12188,
     "create_buffers/happy_path_creates_multiple_buffers_in_one_instruction": 17283,
+    "create_buffers/happy_path_creates_multiple_buffers_in_one_instruction_token_2022": 31720,
     "create_buffers/max_buffers_in_one_instruction": 169742,
+    "create_buffers/max_buffers_in_one_instruction_token_2022": 214956,
     "create_order/happy_path_creates_order_pda_with_expected_body": 4986,
     "initialize/happy_path_initializes_state_pda_with_expected_data": 4530,
     "reclaim_buffer/funded_buffer_is_skipped": 4860,
+    "reclaim_buffer/funded_buffer_is_skipped_token_2022": 4868,
     "reclaim_buffer/happy_path_reclaims_empty_buffer_to_the_authority_itself": 6010,
+    "reclaim_buffer/happy_path_reclaims_empty_buffer_to_the_authority_itself_token_2022": 7499,
     "reclaim_buffer/max_buffers_in_one_instruction": 124824,
+    "reclaim_buffer/max_buffers_in_one_instruction_token_2022": 169378,
     "reclaim_buffer/reclaims_multiple_buffers_skipping_funded": 7612,
+    "reclaim_buffer/reclaims_multiple_buffers_skipping_funded_token_2022": 9105,
     "reclaim_order/happy_path_expired_returns_lamports_and_closes_pda": 2203,
     "reclaim_order/happy_path_on_chain_order_cancelled_is_reclaimable_before_expiry": 2072,
     "reclaim_order/happy_path_on_chain_order_fully_filled_is_reclaimable_before_expiry": 2080,
@@ -67,14 +81,21 @@
     "add_solver/add_with_many_existing_solvers": 366,
     "add_solver/adds_a_solver": 366,
     "create_buffers/happy_path_creates_initialized_buffer_token_account": 303,
+    "create_buffers/happy_path_creates_initialized_buffer_token_account_token_2022": 303,
     "create_buffers/happy_path_creates_multiple_buffers_in_one_instruction": 435,
+    "create_buffers/happy_path_creates_multiple_buffers_in_one_instruction_token_2022": 435,
     "create_buffers/max_buffers_in_one_instruction": 331,
+    "create_buffers/max_buffers_in_one_instruction_token_2022": 295,
     "create_order/happy_path_creates_order_pda_with_expected_body": 452,
     "initialize/happy_path_initializes_state_pda_with_expected_data": 301,
     "reclaim_buffer/funded_buffer_is_skipped": 400,
+    "reclaim_buffer/funded_buffer_is_skipped_token_2022": 400,
     "reclaim_buffer/happy_path_reclaims_empty_buffer_to_the_authority_itself": 400,
+    "reclaim_buffer/happy_path_reclaims_empty_buffer_to_the_authority_itself_token_2022": 400,
     "reclaim_buffer/max_buffers_in_one_instruction": 332,
+    "reclaim_buffer/max_buffers_in_one_instruction_token_2022": 332,
     "reclaim_buffer/reclaims_multiple_buffers_skipping_funded": 466,
+    "reclaim_buffer/reclaims_multiple_buffers_skipping_funded_token_2022": 466,
     "reclaim_order/happy_path_expired_returns_lamports_and_closes_pda": 236,
     "reclaim_order/happy_path_on_chain_order_cancelled_is_reclaimable_before_expiry": 204,
     "reclaim_order/happy_path_on_chain_order_fully_filled_is_reclaimable_before_expiry": 204,
diff --git a/programs/settlement/Cargo.toml b/programs/settlement/Cargo.toml
index 9f9e01d..437ab5f 100644
--- a/programs/settlement/Cargo.toml
+++ b/programs/settlement/Cargo.toml
@@ -31,6 +31,7 @@ jsonschema.workspace = true
 bincode.workspace = true
 litesvm.workspace = true
 litesvm-token.workspace = true
+pastey.workspace = true
 proptest.workspace = true
 serde_json.workspace = true
 cow-settlement-client.workspace = true
@@ -42,6 +43,7 @@ solana-program-pack.workspace = true
 solana-sdk.workspace = true
 solana-sha256-hasher.workspace = true
 solana-system-interface.workspace = true
+spl-associated-token-account-interface.workspace = true
 spl-token-2022-interface.workspace = true
 syn.workspace = true
 
diff --git a/programs/settlement/src/create_buffer.rs b/programs/settlement/src/create_buffer.rs
index fd9318f..554bb62 100644
--- a/programs/settlement/src/create_buffer.rs
+++ b/programs/settlement/src/create_buffer.rs
@@ -85,7 +85,7 @@ mod tests {
     }
 
     #[test]
-    fn process_create_buffer_rejects_wrong_token_program() {
+    fn process_create_buffer_rejects_unsupported_token_program() {
         let data = create_buffer_data();
         // The three shared accounts plus one (buffer_pda, mint) pair so parsing
         // succeeds and reaches the token-program check. The third account (token
diff --git a/programs/settlement/src/reclaim_buffer.rs b/programs/settlement/src/reclaim_buffer.rs
index 151e6d5..b071a52 100644
--- a/programs/settlement/src/reclaim_buffer.rs
+++ b/programs/settlement/src/reclaim_buffer.rs
@@ -180,7 +180,7 @@ mod tests {
     }
 
     #[test]
-    fn process_reclaim_buffer_rejects_wrong_token_program() {
+    fn process_reclaim_buffer_rejects_unsupported_token_program() {
         let mut accounts = base_accounts();
         accounts[TOKEN_PROGRAM] = fake_account(UNRELATED);
         assert_rejects(accounts, ProgramError::IncorrectProgramId);
diff --git a/programs/settlement/tests/common/buffer.rs b/programs/settlement/tests/common/buffer.rs
index b8f9070..4725043 100644
--- a/programs/settlement/tests/common/buffer.rs
+++ b/programs/settlement/tests/common/buffer.rs
@@ -26,23 +26,17 @@ pub fn ensure_buffer_exists(
     program_id: &Pubkey,
     payer: &Keypair,
     mint: &Pubkey,
-) -> Pubkey {
-    ensure_buffer_exists_for(svm, program_id, payer, mint, TokenProgram::SplToken)
-}
-
-/// [`ensure_buffer_exists`] under a token program of the caller's choosing, for
-/// the tests that need a buffer belonging to Token-2022.
-pub fn ensure_buffer_exists_for(
-    svm: &mut LiteSVM,
-    program_id: &Pubkey,
-    payer: &Keypair,
-    mint: &Pubkey,
-    token_program: TokenProgram,
 ) -> Pubkey {
     let pda = buffer_pda(program_id, mint);
     if svm.get_account(&pda).is_some() {
         return pda;
     }
+    // A buffer is a token account of its mint, so it has to be created under the
+    // mint's own program. Read off the mint rather than left to
+    // `super::aim_at_active_token_program`, because a test may build buffers
+    // under both programs at once.
+    let token_program = TokenProgram::try_from(&token::program_of(svm, mint))
+        .expect("the mint lives under a supported token program");
     let ix = Instruction::from(CreateBuffers {
         program_id: *program_id,
         payer: payer.pubkey(),
diff --git a/programs/settlement/tests/common/mod.rs b/programs/settlement/tests/common/mod.rs
index 995cc7b..b7eee5a 100644
--- a/programs/settlement/tests/common/mod.rs
+++ b/programs/settlement/tests/common/mod.rs
@@ -36,10 +36,6 @@ pub const PROGRAM_SO: &str = concat!(
     "/../../target/deploy/cow_settlement.so"
 );
 
-/// The legacy SPL Token program, which the tests create their buffers and
-/// token accounts under unless they exercise Token-2022 specifically.
-pub const SPL_TOKEN_PROGRAM_ID: Pubkey = TokenProgram::SplToken.address();
-
 pub const CPI_CALLER_SO: &str = concat!(
     env!("CARGO_MANIFEST_DIR"),
     "/../../target/deploy/test_cpi_caller.so"
@@ -265,8 +261,10 @@ pub fn signed_tx(
     owner: &Keypair,
     ix: impl Into<Instruction>,
 ) -> Transaction {
+    let mut instructions = [ix.into()];
+    aim_at_active_token_program(&mut instructions);
     Transaction::new_signed_with_payer(
-        &[ix.into()],
+        &instructions,
         Some(&fee_payer.pubkey()),
         &[fee_payer, owner],
         svm.latest_blockhash(),
@@ -293,8 +291,9 @@ pub fn replace_first_matching_account(instruction: &mut Instruction, from: &Pubk
 pub fn payer_signed_tx(
     svm: &LiteSVM,
     payer: &Keypair,
-    instructions: Vec<Instruction>,
+    mut instructions: Vec<Instruction>,
 ) -> Transaction {
+    aim_at_active_token_program(&mut instructions);
     Transaction::new_signed_with_payer(
         &instructions,
         Some(&payer.pubkey()),
@@ -303,6 +302,65 @@ pub fn payer_signed_tx(
     )
 }
 
+/// Repoint every legacy-SPL-Token account of `instructions` at
+/// [`token::active`], so a test written against the legacy program submits the
+/// same transaction aimed at whichever program it is being run under.
+fn aim_at_active_token_program(instructions: &mut [Instruction]) {
+    let active = token::active();
+    if active == TokenProgram::SplToken {
+        return;
+    }
+    for account in instructions
+        .iter_mut()
+        .flat_map(|instruction| &mut instruction.accounts)
+    {
+        if account.pubkey == TokenProgram::SplToken.address() {
+            account.pubkey = active.address();
+        }
+    }
+}
+
+/// Also run `$test` against Token-2022, as `<test>_token_2022`.
+///
+/// Written in front of the test it applies to:
+///
+/// ```ignore
+/// common::also_under_token_2022!(settles_a_single_order);
+/// #[test]
+/// fn settles_a_single_order() { .. }
+/// ```
+///
+/// The test keeps its own `#[test]`, so it runs twice: once under the legacy SPL
+/// Token program, which is what [`token::active`] reports by default, and once
+/// under Token-2022. Nothing in the body changes — the token helpers and
+/// [`payer_signed_tx`] follow the active program on their own. Naming the test
+/// rather than wrapping it keeps the body's indentation, and a stale name is a
+/// compile error rather than a test that quietly stopped being generated.
+#[allow(
+    unused_macros,
+    reason = "only the suites whose instructions name a token program generate the pair"
+)]
+macro_rules! also_under_token_2022 {
+    ($($test:ident),+ $(,)?) => {
+        $(
+            pastey::paste! {
+                #[test]
+                fn [<$test _token_2022>]() {
+                    $crate::common::token::under_token_program(
+                        cow_settlement_interface::token_program::TokenProgram::Token2022,
+                        $test,
+                    );
+                }
+            }
+        )+
+    };
+}
+#[allow(
+    unused_imports,
+    reason = "re-exported for the suites that use the macro; the others never name it"
+)]
+pub(crate) use also_under_token_2022;
+
 /// Assemble `instructions` into a transaction signed by `payer` and submit it,
 /// surfacing only the transaction-level error on failure (dropping the success
 /// metadata's error wrapper).
diff --git a/programs/settlement/tests/common/token.rs b/programs/settlement/tests/common/token.rs
index a71dd3e..727eb28 100644
--- a/programs/settlement/tests/common/token.rs
+++ b/programs/settlement/tests/common/token.rs
@@ -1,10 +1,22 @@
-//! SPL Token helpers for the settlement integration tests.
+//! Token helpers for the settlement integration tests.
+//!
+//! Every helper that acts on an existing token works under whichever token
+//! program owns it, read back with [`program_of`], so a test settling
+//! Token-2022 accounts uses the same calls as one settling legacy ones.
+//!
+//! Creating a mint is the one thing with nothing to read the program from.
+//! [`create_mint`] takes it from [`active`], the program the running test is
+//! exercising — the legacy one unless [`super::also_under_token_2022`]
+//! generated the test — and [`create_mint_under`] names it outright, for the
+//! tests that build mints under both at once.
 
-use cow_settlement_client::cow_settlement_interface::pda::state::find_state_pda;
+use cow_settlement_client::cow_settlement_interface::{
+    pda::state::find_state_pda, token_program::TokenProgram, Instruction,
+};
 use litesvm::{types::TransactionMetadata, LiteSVM};
 use litesvm_token::{
     spl_token::{instruction::initialize_mint2, state::Mint},
-    Approve, CreateAccount, CreateAssociatedTokenAccount, MintTo, Transfer, TOKEN_ID,
+    CreateAssociatedTokenAccount, TOKEN_ID,
 };
 use solana_program_pack::Pack;
 use solana_sdk::{
@@ -13,16 +25,88 @@ use solana_sdk::{
     transaction::Transaction,
 };
 use solana_system_interface::instruction::create_account as system_create_account;
+use spl_associated_token_account_interface::address::get_associated_token_address_with_program_id;
+use spl_token_2022_interface::{
+    extension::StateWithExtensions,
+    instruction::{
+        approve, initialize_account3, mint_to as mint_to_ix,
+        transfer_checked as transfer_checked_ix,
+    },
+    state::{Account, Mint as Mint2022},
+};
+use std::cell::Cell;
+
+use super::{token_2022::Extensions, unique_keypair};
+
+thread_local! {
+    /// The token program [`active`] reports, scoped to one test by
+    /// [`under_token_program`]. Thread local because the test harness runs each
+    /// test on its own thread, so a per-thread value is a per-test value.
+    static ACTIVE: Cell<TokenProgram> = const { Cell::new(TokenProgram::SplToken) };
+}
+
+/// The token program the running test exercises, which is what [`create_mint`]
+/// creates under and what [`super::token_programs`] tells a settlement to
+/// carry.
+pub fn active() -> TokenProgram {
+    ACTIVE.with(Cell::get)
+}
+
+/// Run `test` with `token_program` as the [`active`] one.
+///
+/// [`super::also_under_token_2022`] is the way tests reach this; call it
+/// directly only to nest a differently-programmed section inside a test.
+pub fn under_token_program(token_program: TokenProgram, test: impl FnOnce()) {
+    ACTIVE.replace(token_program);
+    test();
+}
 
-use super::unique_keypair;
+/// The token program that owns `account`.
+///
+/// A token account always lives under its mint's program, so this answers for a
+/// mint and for the accounts holding it alike — which is what lets the helpers
+/// below take the program from the tokens a test already built, rather than
+/// from [`active`].
+pub fn program_of(svm: &LiteSVM, account: &Pubkey) -> Pubkey {
+    svm.get_account(account)
+        .unwrap_or_else(|| panic!("{account} should exist on-chain"))
+        .owner
+}
 
-/// Create a fresh mint owned by `payer` and return its address.
+/// Re-target a token instruction at `token_program`.
 ///
-/// This open-codes what [`litesvm_token::CreateMint`] does rather than calling
-/// it, because that builder generates the mint keypair with `Keypair::new()`
-/// internally and offers no way to supply one. A mint address is a seed of its
-/// buffer PDA, so a random one makes buffer bumps — and the compute cost of
-/// deriving them — vary between runs. See [`super::unique_pubkey`].
+/// The SPL Token builders refuse to emit an instruction for any program but
+/// their own, so the helpers below build against the legacy program and re-point
+/// the result. Token-2022 encodes each of these instructions exactly as the
+/// legacy program does — the same fact that lets the settlement program issue
+/// one transfer against either — so only the program id needs replacing.
+fn under(mut instruction: Instruction, token_program: &Pubkey) -> Instruction {
+    instruction.program_id = *token_program;
+    instruction
+}
+
+/// Submit `instructions` as one transaction signed by `payer` and `extra`.
+fn send_token_tx(
+    svm: &mut LiteSVM,
+    payer: &Keypair,
+    extra: &[&Keypair],
+    instructions: &[Instruction],
+    what: &str,
+) {
+    let mut signers = vec![payer];
+    signers.extend_from_slice(extra);
+    let tx = Transaction::new_signed_with_payer(
+        instructions,
+        Some(&payer.pubkey()),
+        &signers,
+        svm.latest_blockhash(),
+    );
+    svm.send_transaction(tx)
+        .unwrap_or_else(|error| panic!("{what} should succeed: {error:?}"));
+}
+
+/// Create a fresh mint under [`active`], whose mint authority is `payer`, and
+/// return its address.
 pub fn create_mint(svm: &mut LiteSVM, payer: &Keypair) -> Pubkey {
     create_mint_at(svm, payer, &unique_keypair())
 }
@@ -30,7 +114,46 @@ pub fn create_mint(svm: &mut LiteSVM, payer: &Keypair) -> Pubkey {
 /// [`create_mint`] at `mint`'s address rather than a fresh one. Lets a test
 /// reclaim an address a Token-2022 mint was just closed at, which is the only
 /// way a legacy mint can end up where a Token-2022 one used to be.
+///
+/// Under Token-2022 the mint carries [`Extensions::DEFAULT`] rather than being
+/// bare, so every generated test exercises the longer accounts its extensions
+/// force. [`create_mint_under`] is the way to a bare one.
 pub fn create_mint_at(svm: &mut LiteSVM, payer: &Keypair, mint: &Keypair) -> Pubkey {
+    match active() {
+        TokenProgram::SplToken => {
+            create_mint_at_under(svm, payer, mint, &TokenProgram::SplToken.address())
+        }
+        TokenProgram::Token2022 => {
+            super::token_2022::create_mint(svm, payer, mint, Extensions::DEFAULT)
+        }
+    }
+}
+
+/// The length a buffer for a [`create_mint`] mint is allocated at under
+/// [`active`]: the base layout under the legacy program, and whatever
+/// [`Extensions::DEFAULT`] forces under Token-2022.
+pub fn buffer_len() -> usize {
+    match active() {
+        TokenProgram::SplToken => Account::LEN,
+        TokenProgram::Token2022 => Extensions::DEFAULT.token_account_len(),
+    }
+}
+
+/// [`create_mint`] under `token_program` rather than under [`active`], for the
+/// tests that build mints under both programs at once.
+pub fn create_mint_under(svm: &mut LiteSVM, payer: &Keypair, token_program: &Pubkey) -> Pubkey {
+    create_mint_at_under(svm, payer, &unique_keypair(), token_program)
+}
+
+/// Create a mint at `mint`'s address under `token_program`, whose mint authority
+/// is `payer`, and return its address. Every later helper reads the program back
+/// off the mint, so the wrappers above are the only place a test names it.
+fn create_mint_at_under(
+    svm: &mut LiteSVM,
+    payer: &Keypair,
+    mint: &Keypair,
+    token_program: &Pubkey,
+) -> Pubkey {
     /// `litesvm_token::CreateMint`'s default, kept so the two agree.
     const DECIMALS: u8 = 8;
 
@@ -39,18 +162,17 @@ pub fn create_mint_at(svm: &mut LiteSVM, payer: &Keypair, mint: &Keypair) -> Pub
         &mint.pubkey(),
         svm.minimum_balance_for_rent_exemption(Mint::LEN),
         Mint::LEN as u64,
-        &TOKEN_ID,
+        token_program,
     );
-    let initialize = initialize_mint2(&TOKEN_ID, &mint.pubkey(), &payer.pubkey(), None, DECIMALS)
-        .expect("initialize_mint2 should build");
-    let tx = Transaction::new_signed_with_payer(
-        &[create, initialize],
-        Some(&payer.pubkey()),
-        &[payer, mint],
-        svm.latest_blockhash(),
+    // A mint with no extension data, which is every legacy mint and the shape a
+    // Token-2022 mint takes when nothing asks for more. That is what keeps a
+    // buffer for it at the base layout under either program.
+    let initialize = under(
+        initialize_mint2(&TOKEN_ID, &mint.pubkey(), &payer.pubkey(), None, DECIMALS)
+            .expect("initialize_mint2 should build"),
+        token_program,
     );
-    svm.send_transaction(tx)
-        .expect("mint creation should succeed");
+    send_token_tx(svm, payer, &[mint], &[create, initialize], "mint creation");
     mint.pubkey()
 }
 
@@ -63,26 +185,44 @@ pub fn create_token_account(
     mint: &Pubkey,
     owner: &Pubkey,
 ) -> Pubkey {
-    CreateAccount::new(svm, payer, mint)
-        .owner(owner)
-        // Without this the builder generates the address with `Keypair::new()`;
-        // see [`create_mint`].
-        .account_kp(unique_keypair())
-        .send()
-        .expect("token account creation should succeed")
+    let token_program = program_of(svm, mint);
+    let account = unique_keypair();
+    let create = system_create_account(
+        &payer.pubkey(),
+        &account.pubkey(),
+        svm.minimum_balance_for_rent_exemption(Account::LEN),
+        Account::LEN as u64,
+        &token_program,
+    );
+    let initialize = under(
+        initialize_account3(&TOKEN_ID, &account.pubkey(), mint, owner)
+            .expect("initialize_account3 should build"),
+        &token_program,
+    );
+    send_token_tx(
+        svm,
+        payer,
+        &[&account],
+        &[create, initialize],
+        "token account creation",
+    );
+    account.pubkey()
 }
 
 /// Create `owner`'s associated token account for `mint`, funded by `payer`, and
-/// return its address. The address is the canonical ATA, so `transfer` can
-/// source from it without being told where the tokens live.
+/// return its address. The address is the canonical ATA under `mint`'s own
+/// program, so `transfer` can source from it without being told where the
+/// tokens live.
 pub fn create_associated_token_account(
     svm: &mut LiteSVM,
     payer: &Keypair,
     mint: &Pubkey,
     owner: &Pubkey,
 ) -> Pubkey {
+    let token_program = program_of(svm, mint);
     CreateAssociatedTokenAccount::new(svm, payer, mint)
         .owner(owner)
+        .token_program_id(&token_program)
         .send()
         .expect("associated token account creation should succeed")
 }
@@ -96,9 +236,27 @@ pub fn mint_to(
     destination: &Pubkey,
     amount: u64,
 ) {
-    MintTo::new(svm, payer, mint, destination, amount)
-        .send()
-        .expect("mint_to should succeed");
+    let token_program = program_of(svm, mint);
+    let instruction = under(
+        mint_to_ix(&TOKEN_ID, mint, destination, &payer.pubkey(), &[], amount)
+            .expect("mint_to should build"),
+        &token_program,
+    );
+    send_token_tx(svm, payer, &[], &[instruction], "mint_to");
+}
+
+/// The decimals `mint` was created with.
+///
+/// Read through the extension layout, which covers a mint that has extensions
+/// and one that doesn't alike, so this answers under either program.
+fn decimals_of(svm: &LiteSVM, mint: &Pubkey) -> u8 {
+    let account = svm
+        .get_account(mint)
+        .unwrap_or_else(|| panic!("{mint} should exist on-chain"));
+    StateWithExtensions::<Mint2022>::unpack(&account.data)
+        .expect("the mint should be a valid mint account")
+        .base
+        .decimals
 }
 
 /// Transfer `amount` of `mint` from `owner`'s associated token account into
@@ -110,13 +268,31 @@ pub fn transfer(
     destination: &Pubkey,
     amount: u64,
 ) {
-    Transfer::new(svm, owner, mint, destination, amount)
-        .send()
-        .expect("transfer should succeed");
+    let token_program = program_of(svm, mint);
+    let source =
+        get_associated_token_address_with_program_id(&owner.pubkey(), mint, &token_program);
+    // Checked rather than plain `Transfer`: Token-2022 refuses the unchecked one
+    // for a mint carrying a transfer fee, which [`Extensions::DEFAULT`] does, and
+    // the legacy program accepts it just the same.
+    let instruction = under(
+        transfer_checked_ix(
+            &TOKEN_ID,
+            &source,
+            mint,
+            destination,
+            &owner.pubkey(),
+            &[],
+            amount,
+            decimals_of(svm, mint),
+        )
+        .expect("transfer should build"),
+        &token_program,
+    );
+    send_token_tx(svm, owner, &[], &[instruction], "transfer");
 }
 
 /// Approve `delegate` to spend up to `amount` from `source`. `owner` must be the
-/// SPL owner of `source`. The transaction will be paid by the owner.
+/// token owner of `source`. The transaction will be paid by the owner.
 pub fn delegate(
     svm: &mut LiteSVM,
     owner: &Keypair,
@@ -124,9 +300,13 @@ pub fn delegate(
     delegate: &Pubkey,
     amount: u64,
 ) {
-    Approve::new(svm, owner, delegate, source, amount)
-        .send()
-        .expect("approving a delegate should succeed");
+    let token_program = program_of(svm, source);
+    let instruction = under(
+        approve(&TOKEN_ID, source, delegate, &owner.pubkey(), &[], amount)
+            .expect("approve should build"),
+        &token_program,
+    );
+    send_token_tx(svm, owner, &[], &[instruction], "approving a delegate");
 }
 
 /// Fund `sell_token` with `amount` of its mint and approve the settlement state
@@ -149,7 +329,8 @@ pub fn fund_and_delegate(
     );
 }
 
-/// Read the SPL token balance of `account`.
+/// Read the token balance of `account`. The two programs share the base layout
+/// this reads, so it answers for an account under either.
 pub fn balance(svm: &LiteSVM, account: &Pubkey) -> u64 {
     litesvm_token::get_spl_account::<litesvm_token::spl_token::state::Account>(svm, account)
         .expect("account should exist and be a valid SPL token account")
@@ -163,25 +344,29 @@ pub fn delegated_amount(svm: &LiteSVM, account: &Pubkey) -> u64 {
         .delegated_amount
 }
 
-/// Assert that no SPL Token instruction issued by the transaction references
+/// Assert that no token instruction issued by the transaction references
 /// `account`. Each token transfer the program performs is a CPI recorded in
 /// `transaction.inner_instructions`. We can use that to check the token-program
 /// instructions, so a settlement that must leave one side untouched can prove
 /// no token instruction so much as named it.
+///
+/// Every supported program counts, not just the one the settlement was expected
+/// to use: a transfer issued against the wrong one is exactly the kind of touch
+/// this is meant to catch.
 #[track_caller]
 pub fn assert_no_token_instruction_touching(
     transaction: &TransactionMetadata,
     account_keys: &[Pubkey],
     account: &Pubkey,
 ) {
-    let token_program = Pubkey::new_from_array(litesvm_token::spl_token::ID.to_bytes());
     for instruction in transaction
         .inner_instructions
         .iter()
         .flatten()
         .map(|inner| &inner.instruction)
     {
-        if account_keys[usize::from(instruction.program_id_index)] != token_program {
+        let program = account_keys[usize::from(instruction.program_id_index)];
+        if TokenProgram::try_from(&program).is_err() {
             continue;
         }
         let touches_account = instruction
@@ -190,7 +375,7 @@ pub fn assert_no_token_instruction_touching(
             .any(|&index| account_keys[usize::from(index)] == *account);
         assert!(
             !touches_account,
-            "expected no SPL Token instruction touching {account}, but one did",
+            "expected no token instruction touching {account}, but one did",
         );
     }
 }
@@ -202,12 +387,12 @@ pub fn overwrite_token_account(
     address: &Pubkey,
     mint: &Pubkey,
 ) {
+    let token_program = program_of(svm, mint);
     let template = create_token_account(svm, payer, mint, &payer.pubkey());
     let data = svm
         .get_account(&template)
         .expect("the freshly created template exists")
         .data;
-    let token_program = Pubkey::new_from_array(TOKEN_ID.to_bytes());
     super::create_account_at(svm, *address, &token_program, &data);
 }
 
diff --git a/programs/settlement/tests/common/token_2022.rs b/programs/settlement/tests/common/token_2022.rs
index 284b5e0..678bbc5 100644
--- a/programs/settlement/tests/common/token_2022.rs
+++ b/programs/settlement/tests/common/token_2022.rs
@@ -24,16 +24,16 @@ use spl_token_2022_interface::{
     state::{Account, Mint},
 };
 
-/// The Token-2022 program, the counterpart of [`super::SPL_TOKEN_PROGRAM_ID`].
+/// The Token-2022 program, spelled once so the builders below can take it.
 const TOKEN_2022_PROGRAM_ID: Pubkey = TokenProgram::Token2022.address();
 
 /// Decimals every test mint carries, matching [`super::token::create_mint`] so
 /// a legacy and a Token-2022 mint differ only in their program.
 const DECIMALS: u8 = 8;
 
-/// Transfer-fee parameters for [`Extensions::CloseAuthorityAndTransferFee`]. Arbitrary;
-/// nothing reads them back, but `InitializeTransferFeeConfig` demands values.
-const FEE_BASIS_POINTS: u16 = 50;
+/// Transfer-fee parameters for [`Extensions::CloseAuthorityAndTransferFee`].
+/// Default the fee to 0 so that most tests perform the same as if its a regular SPL transfer.
+const FEE_BASIS_POINTS: u16 = 0;
 const MAXIMUM_FEE: u64 = 1_000;
 
 /// Defines a set of token account/mint configurations we are interested in testing
@@ -65,6 +65,10 @@ impl RequiredInitAccountExtensionType {
 }
 
 impl Extensions {
+    /// Include CloseAuthorityAndTransferFee since it makes it mandatory for its token accounts to include an extension
+    /// for better test coverage.
+    pub const DEFAULT: Self = Self::CloseAuthorityAndTransferFee;
+
     /// The extensions which should be configured on the mint
     fn mint(self) -> &'static [ExtensionType] {
         match self {
diff --git a/programs/settlement/tests/create_buffer.rs b/programs/settlement/tests/create_buffer.rs
index 919a6d8..fdb96da 100644
--- a/programs/settlement/tests/create_buffer.rs
+++ b/programs/settlement/tests/create_buffer.rs
@@ -21,7 +21,6 @@ use solana_compute_budget::{
 use solana_sdk::{
     instruction::{Instruction, InstructionError},
     program_error::ProgramError,
-    program_pack::Pack,
     pubkey::Pubkey,
     signature::{Keypair, Signer},
     transaction::TransactionError,
@@ -35,6 +34,7 @@ use crate::common::{
 
 mod common;
 
+common::also_under_token_2022!(happy_path_creates_initialized_buffer_token_account);
 #[test]
 fn happy_path_creates_initialized_buffer_token_account() {
     let (mut svm, program_id, payer) = common::setup();
@@ -57,13 +57,13 @@ fn happy_path_creates_initialized_buffer_token_account() {
         .expect("buffer PDA should exist after create_buffer");
     assert_eq!(
         account.owner,
-        TokenProgram::SplToken.address(),
-        "buffer must be owned by the SPL Token program"
+        common::token::active().address(),
+        "buffer must be owned by the token program it was created under"
     );
     assert_eq!(
         account.data.len(),
-        TokenAccount::LEN,
-        "buffer must be sized to a token account",
+        common::token::buffer_len(),
+        "buffer must be sized to a token account for its mint",
     );
 
     common::assert_rent_exempt(&svm, &account);
@@ -105,6 +105,7 @@ fn happy_path_creates_initialized_buffer_token_account() {
     );
 }
 
+common::also_under_token_2022!(buffer_can_receive_tokens);
 #[test]
 fn buffer_can_receive_tokens() {
     let (mut svm, program_id, payer) = common::setup();
@@ -141,6 +142,8 @@ fn buffer_can_receive_tokens() {
     );
 }
 
+// Legacy-only: the native mint is the legacy program's, and Token-2022
+// recognizes only its own.
 #[test]
 fn happy_path_creates_native_token_buffer() {
     // `InitializeAccount` special-cases the native mint: it's recognized by key
@@ -176,6 +179,7 @@ fn happy_path_creates_native_token_buffer() {
     );
 }
 
+common::also_under_token_2022!(happy_path_creates_multiple_buffers_in_one_instruction);
 #[test]
 fn happy_path_creates_multiple_buffers_in_one_instruction() {
     let (mut svm, program_id, payer) = common::setup();
@@ -202,13 +206,13 @@ fn happy_path_creates_multiple_buffers_in_one_instruction() {
             .expect("each buffer PDA should exist after create_buffers");
         assert_eq!(
             account.owner,
-            TokenProgram::SplToken.address(),
-            "each buffer must be owned by the SPL Token program"
+            common::token::active().address(),
+            "each buffer must be owned by the token program it was created under"
         );
         assert_eq!(
             account.data.len(),
-            TokenAccount::LEN,
-            "each buffer must be sized to a token account",
+            common::token::buffer_len(),
+            "each buffer must be sized to a token account for its mint",
         );
         common::assert_rent_exempt(&svm, &account);
 
@@ -227,6 +231,8 @@ fn happy_path_creates_multiple_buffers_in_one_instruction() {
     }
 }
 
+// Legacy-only: this fails while parsing, before the instruction so much as
+// looks at its token program.
 #[test]
 fn rejects_no_buffers() {
     let (mut svm, program_id, payer) = common::setup();
@@ -254,6 +260,7 @@ fn rejects_no_buffers() {
     );
 }
 
+common::also_under_token_2022!(rejects_arbitrary_wrong_buffer_pda);
 #[test]
 fn rejects_arbitrary_wrong_buffer_pda() {
     let (mut svm, program_id, payer) = common::setup();
@@ -271,6 +278,7 @@ fn rejects_arbitrary_wrong_buffer_pda() {
     common::pda::assert_rejected_as_noncanonical(&mut svm, tx, &wrong_pda);
 }
 
+common::also_under_token_2022!(rejects_non_canonical_bump_pda);
 #[test]
 fn rejects_non_canonical_bump_pda() {
     let (mut svm, program_id, payer) = common::setup();
@@ -292,7 +300,7 @@ fn rejects_non_canonical_bump_pda() {
 }
 
 #[test]
-fn rejects_non_spl_token_program() {
+fn rejects_unsupported_token_program() {
     let (mut svm, program_id, payer) = common::setup();
     let mint = common::token::create_mint(&mut svm, &payer);
     let (buffer_pda, _bump) = find_buffer_pda(&program_id, &mint);
@@ -331,6 +339,7 @@ fn rejects_non_spl_token_program() {
     );
 }
 
+common::also_under_token_2022!(rejects_invalid_mint);
 #[test]
 fn rejects_invalid_mint() {
     let (mut svm, program_id, payer) = common::setup();
@@ -339,7 +348,9 @@ fn rejects_invalid_mint() {
     // buffer PDA from it and delegates mint validation to InitializeAccount3,
     // which rejects it: a non-mint account isn't owned by the token program, so
     // the CPI fails with IncorrectProgramId after the buffer was allocated,
-    // reverting the whole instruction.
+    // reverting the whole instruction. Sizing doesn't get in the way first: an
+    // account this short can't be a mint with extensions, so it's sized at the
+    // base layout without the token program being consulted.
     let not_a_mint = unique_pubkey();
     let (buffer_pda, _bump) = find_buffer_pda(&program_id, &not_a_mint);
 
@@ -370,6 +381,7 @@ fn rejects_invalid_mint() {
     );
 }
 
+common::also_under_token_2022!(creates_buffer_when_address_is_prefunded);
 #[test]
 fn creates_buffer_when_address_is_prefunded() {
     let (mut svm, program_id, payer) = common::setup();
@@ -387,6 +399,7 @@ fn creates_buffer_when_address_is_prefunded() {
     });
 }
 
+common::also_under_token_2022!(recreating_same_buffer_is_idempotent);
 #[test]
 fn recreating_same_buffer_is_idempotent() {
     let (mut svm, program_id, payer) = common::setup();
@@ -404,6 +417,7 @@ fn recreating_same_buffer_is_idempotent() {
     });
 }
 
+common::also_under_token_2022!(batch_with_existing_buffer_passes_with_no_changes);
 #[test]
 fn batch_with_existing_buffer_passes_with_no_changes() {
     let (mut svm, program_id, payer) = common::setup();
@@ -456,6 +470,7 @@ fn batch_with_existing_buffer_passes_with_no_changes() {
     );
 }
 
+common::also_under_token_2022!(one_failing_buffer_reverts_the_whole_batch);
 #[test]
 fn one_failing_buffer_reverts_the_whole_batch() {
     let (mut svm, program_id, payer) = common::setup();
@@ -489,6 +504,7 @@ fn one_failing_buffer_reverts_the_whole_batch() {
     );
 }
 
+common::also_under_token_2022!(same_mint_twice_in_one_instruction_is_idempotent);
 #[test]
 fn same_mint_twice_in_one_instruction_is_idempotent() {
     let (mut svm, program_id, payer) = common::setup();
@@ -557,6 +573,13 @@ fn sizes_a_token_2022_buffer_to_the_extensions_its_mint_forces() {
     }
 }
 
+fn known_max_buffer_count() -> usize {
+    match common::token::active() {
+        TokenProgram::SplToken => 30,  // Limited by account limit
+        TokenProgram::Token2022 => 21, // Limited by CPI call limit
+    }
+}
+
 /// Largest number of buffers a single ALT-backed `create_buffers` transaction
 /// can carry, bounded by the transaction account-lock limit (litesvm and current
 /// mainnet both cap this at 64).
@@ -570,13 +593,14 @@ fn max_buffers_via_lookup_table(svm: &mut LiteSVM, program_id: &Pubkey, payer: &
         let ix = CreateBuffersRaw {
             program_id: *program_id,
             payer: payer.pubkey(),
-            token_program: TokenProgram::SplToken.address(),
+            token_program: common::token::active().address(),
             buffers: &buffers,
         };
         common::lookup_table::lookup_table_tx(svm, payer, ix)
     })
 }
 
+common::also_under_token_2022!(bench_assert_known_max_buffer_count);
 /// This isn't really a test, it's a way to make it visible that a code change
 /// has changed the amount of buffer accounts that can be created in the same
 /// transaction. If the number increases, great, bump it up! If it decreases and
@@ -584,13 +608,16 @@ fn max_buffers_via_lookup_table(svm: &mut LiteSVM, program_id: &Pubkey, payer: &
 #[test]
 fn bench_assert_known_max_buffer_count() {
     let (mut svm, program_id, payer) = common::setup();
-    let max_buffers = max_buffers_via_lookup_table(&mut svm, &program_id, &payer);
+    let probe = max_buffers_via_lookup_table(&mut svm, &program_id, &payer);
     assert_eq!(
-        max_buffers, 30,
-        "Max buffers that can be created has changed"
+        probe,
+        30,
+        "the account-lock ceiling has changed under {:?}",
+        common::token::active(),
     );
 }
 
+common::also_under_token_2022!(max_buffers_in_one_instruction);
 /// Pack a single `create_buffers` instruction with as many buffers as a
 /// transaction can have. Use Address Lookup Table to reach the real
 /// account-lock ceiling. This is a ceiling on how many buffers one transaction
@@ -622,14 +649,15 @@ fn max_buffers_in_one_instruction() {
     });
     let (state_pda, _) = find_state_pda(&program_id);
 
-    let max_buffers = max_buffers_via_lookup_table(&mut svm, &program_id, &payer);
+    let probe = max_buffers_via_lookup_table(&mut svm, &program_id, &payer);
     // A legacy transaction tops out around 15 buffers (32-byte keys inlined into
     // a 1232-byte packet). The whole point of the lookup table is to beat that;
     // guard against a counterproductive use of lookup tables.
     assert!(
-        max_buffers > 15,
-        "a lookup-table transaction must exceed the legacy packet limit, got {max_buffers}"
+        probe > 15,
+        "a lookup-table transaction must exceed the legacy packet limit, got {probe}"
     );
+    let max_buffers = known_max_buffer_count();
 
     let mints: Vec<Pubkey> = (0..max_buffers)
         .map(|_| common::token::create_mint(&mut svm, &payer))
@@ -638,7 +666,7 @@ fn max_buffers_in_one_instruction() {
     let ix = CreateBuffers {
         program_id,
         payer: payer.pubkey(),
-        token_program: TokenProgram::SplToken,
+        token_program: common::token::active(),
         mints: &mints,
     };
     let tx = common::lookup_table::lookup_table_tx(&mut svm, &payer, ix);
diff --git a/programs/settlement/tests/reclaim_buffer.rs b/programs/settlement/tests/reclaim_buffer.rs
index fd4a51b..b1f26c2 100644
--- a/programs/settlement/tests/reclaim_buffer.rs
+++ b/programs/settlement/tests/reclaim_buffer.rs
@@ -21,6 +21,7 @@ use crate::common::{
 
 mod common;
 
+common::also_under_token_2022!(happy_path_reclaims_to_a_recipient_chosen_by_the_authority);
 #[test]
 fn happy_path_reclaims_to_a_recipient_chosen_by_the_authority() {
     let (
@@ -71,6 +72,7 @@ fn happy_path_reclaims_to_a_recipient_chosen_by_the_authority() {
     );
 }
 
+common::also_under_token_2022!(happy_path_reclaims_empty_buffer_to_the_authority_itself);
 #[test]
 fn happy_path_reclaims_empty_buffer_to_the_authority_itself() {
     let (
@@ -114,6 +116,7 @@ fn happy_path_reclaims_empty_buffer_to_the_authority_itself() {
     );
 }
 
+common::also_under_token_2022!(funded_buffer_is_skipped);
 #[test]
 fn funded_buffer_is_skipped() {
     let (
@@ -149,6 +152,7 @@ fn funded_buffer_is_skipped() {
     );
 }
 
+common::also_under_token_2022!(reclaims_to_the_settlements_own_state_pda);
 /// The recipient isn't required to be a system account: closing only moves
 /// lamports, so a program-owned data account is credited just the same. The
 /// settlement's own state PDA is the sharpest case, since it also occupies the
@@ -207,6 +211,7 @@ fn reclaims_to_the_settlements_own_state_pda() {
     );
 }
 
+common::also_under_token_2022!(reclaims_multiple_buffers_skipping_funded);
 #[test]
 fn reclaims_multiple_buffers_skipping_funded() {
     let (
@@ -248,6 +253,7 @@ fn reclaims_multiple_buffers_skipping_funded() {
     );
 }
 
+common::also_under_token_2022!(rejects_the_same_buffer_twice_in_one_instruction);
 #[test]
 fn rejects_the_same_buffer_twice_in_one_instruction() {
     let (
@@ -278,6 +284,7 @@ fn rejects_the_same_buffer_twice_in_one_instruction() {
     );
 }
 
+common::also_under_token_2022!(rejects_when_signer_is_not_the_configured_reclaim_authority);
 #[test]
 fn rejects_when_signer_is_not_the_configured_reclaim_authority() {
     let (
@@ -308,6 +315,7 @@ fn rejects_when_signer_is_not_the_configured_reclaim_authority() {
     );
 }
 
+common::also_under_token_2022!(rejects_when_the_reclaim_authority_does_not_sign);
 /// Naming the configured authority isn't enough: it has to sign. The builder
 /// always marks it as a signer, so this test strips the flag by hand.
 #[test]
@@ -367,13 +375,7 @@ fn buffer_whose_mint_was_reopened(
     let mint_keypair = common::unique_keypair();
     let mint =
         common::token_2022::create_mint(svm, payer, &mint_keypair, Extensions::CloseAuthorityOnly);
-    let buffer_pda = common::buffer::ensure_buffer_exists_for(
-        svm,
-        program_id,
-        payer,
-        &mint,
-        TokenProgram::Token2022,
-    );
+    let buffer_pda = common::buffer::ensure_buffer_exists(svm, program_id, payer, &mint);
 
     common::token_2022::close_mint(svm, payer, &mint);
     reopen(svm, payer, &mint_keypair);
@@ -459,7 +461,7 @@ fn reclaims_a_buffer_whose_mint_was_reopened_as_a_legacy_mint() {
         svm.get_account(&mint)
             .expect("the reopened mint should exist")
             .owner,
-        common::SPL_TOKEN_PROGRAM_ID,
+        TokenProgram::SplToken.address(),
         "sanity: the mint must now belong to the legacy program"
     );
 
@@ -505,13 +507,14 @@ fn max_buffers_reclaim_via_lookup_table(
             state_pda,
             reclaim_authority: reclaim_authority.pubkey(),
             reclaim_recipient: reclaim_authority.pubkey(),
-            token_program: TokenProgram::SplToken.address(),
+            token_program: common::token::active().address(),
             buffers: &buffers,
         };
         common::lookup_table::lookup_table_tx(svm, reclaim_authority, ix)
     })
 }
 
+common::also_under_token_2022!(bench_assert_known_max_buffer_count);
 /// This isn't really a test, it's a way to make it visible that a code change
 /// has changed the amount of buffer accounts that can be reclaimed in the same
 /// transaction. If the number increases, great, bump it up! If it decreases and
@@ -530,12 +533,16 @@ fn bench_assert_known_max_buffer_count() {
         .expect("airdrop should succeed");
     let max_buffers =
         max_buffers_reclaim_via_lookup_table(&mut svm, &program_id, &reclaim_authority);
+
     assert_eq!(
-        max_buffers, 30,
-        "Max buffers that can be reclaimed has changed"
+        max_buffers,
+        30,
+        "Max buffers that can be reclaimed has changed under {:?}",
+        common::token::active(),
     );
 }
 
+common::also_under_token_2022!(max_buffers_in_one_instruction);
 /// Pack a single `reclaim_buffer` instruction with as many buffers as a
 /// transaction can have, all of them empty and therefore closable. Use an
 /// Address Lookup Table to reach the real account-lock ceiling. This is a
@@ -584,7 +591,7 @@ fn max_buffers_in_one_instruction() {
         program_id,
         reclaim_authority: reclaim_authority.pubkey(),
         reclaim_recipient: reclaim_authority.pubkey(),
-        token_program: TokenProgram::SplToken,
+        token_program: common::token::active(),
         mints: &mints,
     };
     let tx = common::lookup_table::lookup_table_tx(&mut svm, &reclaim_authority, ix);