diff --git a/programs/rewards/src/error.rs b/programs/rewards/src/error.rs
index be8c454..d4035aa 100644
--- a/programs/rewards/src/error.rs
+++ b/programs/rewards/src/error.rs
@@ -90,6 +90,25 @@ pub enum MplxRewardsError {
     /// Failed to derive PDA.
     #[error("Failed to derive PDA")]
     DerivationError,
+
+    /// 21
+    #[error("Mining already restricted")]
+    MiningAlreadyRestricted,
+
+    /// 22
+    /// Mining is not restricted
+    #[error("Mining is not restricted")]
+    MiningNotRestricted,
+
+    /// 23
+    /// Claiming is restricted
+    #[error("Claiming is restricted")]
+    ClaimingRestricted,
+
+    /// 24
+    /// Withdrawal is restricted while claiming is restricted
+    #[error("Withdrawal is restricted while claiming is restricted")]
+    WithdrawalRestricted,
 }
 
 impl PrintProgramError for MplxRewardsError {
diff --git a/programs/rewards/src/instruction.rs b/programs/rewards/src/instruction.rs
index a1c6954..bb3678a 100644
--- a/programs/rewards/src/instruction.rs
+++ b/programs/rewards/src/instruction.rs
@@ -144,6 +144,21 @@ pub enum RewardsInstruction {
         staked_amount: u64,
         new_delegate: Pubkey,
     },
+
+    /// Prevents the mining account from rewards withdrawing
+    #[account(0, signer, name = "deposit_authority", desc = "The address of the Staking program's Registrar, which is PDA and is responsible for signing CPIs")]
+    #[account(1, name = "reward_pool", desc = "The address of the reward pool")]
+    #[account(2, writable, name = "mining", desc = "The address of the mining account which belongs to the user and stores info about user's rewards")]
+    RestrictTokenFlow {
+        mining_owner: Pubkey,
+    },
+
+    #[account(0, signer, name = "deposit_authority", desc = "The address of the Staking program's Registrar, which is PDA and is responsible for signing CPIs")]
+    #[account(1, name = "reward_pool", desc = "The address of the reward pool")]
+    #[account(2, writable, name = "mining", desc = "The address of the mining account which belongs to the user and stores info about user's rewards")]
+    AllowTokenFlow {
+        mining_owner: Pubkey,
+    },
 }
 
 /// Creates 'InitializePool' instruction.
@@ -431,3 +446,47 @@ pub fn change_delegate(
         accounts,
     )
 }
+
+pub fn restrict_tokenflow(
+    program_id: &Pubkey,
+    deposit_authority: &Pubkey,
+    reward_pool: &Pubkey,
+    mining: &Pubkey,
+    mining_owner: &Pubkey,
+) -> Instruction {
+    let accounts = vec![
+        AccountMeta::new_readonly(*deposit_authority, true),
+        AccountMeta::new_readonly(*reward_pool, false),
+        AccountMeta::new(*mining, false),
+    ];
+
+    Instruction::new_with_borsh(
+        *program_id,
+        &RewardsInstruction::RestrictTokenFlow {
+            mining_owner: *mining_owner,
+        },
+        accounts,
+    )
+}
+
+pub fn allow_tokenflow(
+    program_id: &Pubkey,
+    deposit_authority: &Pubkey,
+    reward_pool: &Pubkey,
+    mining: &Pubkey,
+    mining_owner: &Pubkey,
+) -> Instruction {
+    let accounts = vec![
+        AccountMeta::new_readonly(*deposit_authority, true),
+        AccountMeta::new_readonly(*reward_pool, false),
+        AccountMeta::new(*mining, false),
+    ];
+
+    Instruction::new_with_borsh(
+        *program_id,
+        &RewardsInstruction::AllowTokenFlow {
+            mining_owner: *mining_owner,
+        },
+        accounts,
+    )
+}
diff --git a/programs/rewards/src/instructions/claim.rs b/programs/rewards/src/instructions/claim.rs
index d409f5a..c84432e 100644
--- a/programs/rewards/src/instructions/claim.rs
+++ b/programs/rewards/src/instructions/claim.rs
@@ -1,5 +1,6 @@
 use crate::{
     asserts::{assert_account_key, assert_account_owner},
+    error::MplxRewardsError,
     state::{WrappedMining, WrappedRewardPool},
     utils::{spl_transfer, AccountLoader},
 };
@@ -42,6 +43,10 @@ pub fn process_claim<'a>(program_id: &Pubkey, accounts: &'a [AccountInfo<'a>]) -
             let mining_data = &mut mining.data.borrow_mut();
             let mut wrapped_mining = WrappedMining::from_bytes_mut(mining_data)?;
 
+            if wrapped_mining.mining.is_tokenflow_restricted() {
+                return Err(MplxRewardsError::ClaimingRestricted.into());
+            }
+
             assert_account_owner(reward_pool, program_id)?;
             assert_account_key(mining_owner, &wrapped_mining.mining.owner)?;
             assert_account_key(reward_pool, &wrapped_mining.mining.reward_pool)?;
diff --git a/programs/rewards/src/instructions/mod.rs b/programs/rewards/src/instructions/mod.rs
index c5b8ecd..e313318 100644
--- a/programs/rewards/src/instructions/mod.rs
+++ b/programs/rewards/src/instructions/mod.rs
@@ -12,6 +12,7 @@ mod extend_stake;
 mod fill_vault;
 mod initialize_mining;
 mod initialize_pool;
+mod penalties;
 mod withdraw_mining;
 
 pub(crate) use change_delegate::*;
@@ -23,6 +24,7 @@ pub(crate) use extend_stake::*;
 pub(crate) use fill_vault::*;
 pub(crate) use initialize_mining::*;
 pub(crate) use initialize_pool::*;
+pub(crate) use penalties::*;
 pub(crate) use withdraw_mining::*;
 
 pub fn process_instruction<'a>(
@@ -116,5 +118,13 @@ pub fn process_instruction<'a>(
             msg!("RewardsInstruction: ChangeDelegate");
             process_change_delegate(program_id, accounts, staked_amount, &new_delegate)
         }
+        RewardsInstruction::RestrictTokenFlow { mining_owner } => {
+            msg!("RewardsInstruction: RestrictClaiming");
+            process_restrict_tokenflow(program_id, accounts, &mining_owner)
+        }
+        RewardsInstruction::AllowTokenFlow { mining_owner } => {
+            msg!("RewardsInstruction: AllowClaiming");
+            process_allow_tokenflow(program_id, accounts, &mining_owner)
+        }
     }
 }
diff --git a/programs/rewards/src/instructions/penalties/allow_tokenflow.rs b/programs/rewards/src/instructions/penalties/allow_tokenflow.rs
new file mode 100644
index 0000000..6f17c4c
--- /dev/null
+++ b/programs/rewards/src/instructions/penalties/allow_tokenflow.rs
@@ -0,0 +1,31 @@
+use crate::{asserts::assert_and_get_pool_and_mining, utils::AccountLoader};
+use solana_program::{account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey};
+
+pub fn process_allow_tokenflow<'a>(
+    program_id: &Pubkey,
+    accounts: &'a [AccountInfo<'a>],
+    mining_owner: &Pubkey,
+) -> ProgramResult {
+    let account_info_iter = &mut accounts.iter().enumerate();
+
+    let deposit_authority = AccountLoader::next_signer(account_info_iter)?;
+    let reward_pool = AccountLoader::next_with_owner(account_info_iter, program_id)?;
+    let mining = AccountLoader::next_with_owner(account_info_iter, program_id)?;
+
+    let reward_pool_data = &mut reward_pool.data.borrow_mut();
+    let mining_data = &mut mining.data.borrow_mut();
+
+    let (_, wrapped_mining) = assert_and_get_pool_and_mining(
+        program_id,
+        mining_owner,
+        mining,
+        reward_pool,
+        deposit_authority,
+        reward_pool_data,
+        mining_data,
+    )?;
+
+    wrapped_mining.mining.allow_tokenflow()?;
+
+    Ok(())
+}
diff --git a/programs/rewards/src/instructions/penalties/mod.rs b/programs/rewards/src/instructions/penalties/mod.rs
new file mode 100644
index 0000000..76095d1
--- /dev/null
+++ b/programs/rewards/src/instructions/penalties/mod.rs
@@ -0,0 +1,5 @@
+pub(crate) use allow_tokenflow::*;
+pub(crate) use restrict_tokenflow::*;
+
+mod allow_tokenflow;
+mod restrict_tokenflow;
diff --git a/programs/rewards/src/instructions/penalties/restrict_tokenflow.rs b/programs/rewards/src/instructions/penalties/restrict_tokenflow.rs
new file mode 100644
index 0000000..70c02ae
--- /dev/null
+++ b/programs/rewards/src/instructions/penalties/restrict_tokenflow.rs
@@ -0,0 +1,31 @@
+use crate::{asserts::assert_and_get_pool_and_mining, utils::AccountLoader};
+use solana_program::{account_info::AccountInfo, entrypoint::ProgramResult, pubkey::Pubkey};
+
+pub fn process_restrict_tokenflow<'a>(
+    program_id: &Pubkey,
+    accounts: &'a [AccountInfo<'a>],
+    mining_owner: &Pubkey,
+) -> ProgramResult {
+    let account_info_iter = &mut accounts.iter().enumerate();
+
+    let deposit_authority = AccountLoader::next_signer(account_info_iter)?;
+    let reward_pool = AccountLoader::next_with_owner(account_info_iter, program_id)?;
+    let mining = AccountLoader::next_with_owner(account_info_iter, program_id)?;
+
+    let reward_pool_data = &mut reward_pool.data.borrow_mut();
+    let mining_data = &mut mining.data.borrow_mut();
+
+    let (_, wrapped_mining) = assert_and_get_pool_and_mining(
+        program_id,
+        mining_owner,
+        mining,
+        reward_pool,
+        deposit_authority,
+        reward_pool_data,
+        mining_data,
+    )?;
+
+    wrapped_mining.mining.restrict_tokenflow()?;
+
+    Ok(())
+}
diff --git a/programs/rewards/src/instructions/withdraw_mining.rs b/programs/rewards/src/instructions/withdraw_mining.rs
index d9f8899..5a0e547 100644
--- a/programs/rewards/src/instructions/withdraw_mining.rs
+++ b/programs/rewards/src/instructions/withdraw_mining.rs
@@ -1,5 +1,6 @@
 use crate::{
     asserts::assert_and_get_pool_and_mining,
+    error::MplxRewardsError,
     utils::{get_delegate_mining, AccountLoader},
 };
 
@@ -33,6 +34,10 @@ pub fn process_withdraw_mining<'a>(
         mining_data,
     )?;
 
+    if wrapped_mining.mining.is_tokenflow_restricted() {
+        return Err(MplxRewardsError::WithdrawalRestricted.into());
+    }
+
     let delegate_mining = get_delegate_mining(delegate_mining, mining)?;
     if let Some(delegate_mining) = delegate_mining {
         verify_delegate_mining_address(program_id, delegate_mining, delegate, reward_pool.key)?
diff --git a/programs/rewards/src/state/mining.rs b/programs/rewards/src/state/mining.rs
index d77c0b4..b91c565 100644
--- a/programs/rewards/src/state/mining.rs
+++ b/programs/rewards/src/state/mining.rs
@@ -32,6 +32,9 @@ pub struct WrappedImmutableMining<'a> {
     pub weighted_stake_diffs: &'a MiningWeightedStakeDiffs,
 }
 
+pub const ACCOUNT_TYPE_BYTE: usize = 0;
+pub const CLAIMING_RESTRICTION_BYTE: usize = 1;
+
 impl<'a> WrappedMining<'a> {
     pub const LEN: usize = 1776;
 
@@ -99,8 +102,9 @@ pub struct Mining {
     pub bump: u8,
     /// Account type - Mining. This discriminator should exist in order to prevent
     /// shenanigans with customly modified accounts and their fields.
-    /// 1: account type
-    /// 2-7: unused
+    /// 0: account type
+    /// 1: claim is restricted
+    /// 2-6: unused
     pub data: [u8; 7],
 }
 
@@ -114,7 +118,7 @@ impl Mining {
     pub fn initialize(reward_pool: Pubkey, owner: Pubkey, bump: u8) -> Mining {
         let account_type = AccountType::Mining.into();
         let mut data = [0; 7];
-        data[0] = account_type;
+        data[ACCOUNT_TYPE_BYTE] = account_type;
         Mining {
             data,
             reward_pool,
@@ -125,7 +129,7 @@ impl Mining {
     }
 
     pub fn account_type(&self) -> AccountType {
-        AccountType::from(self.data[0])
+        AccountType::from(self.data[ACCOUNT_TYPE_BYTE])
     }
 
     /// Claim reward
@@ -193,11 +197,33 @@ impl Mining {
 
         Ok(())
     }
+
+    pub fn restrict_tokenflow(&mut self) -> ProgramResult {
+        if self.data[CLAIMING_RESTRICTION_BYTE] == 1 {
+            Err(MplxRewardsError::MiningAlreadyRestricted.into())
+        } else {
+            self.data[CLAIMING_RESTRICTION_BYTE] = 1;
+            Ok(())
+        }
+    }
+
+    pub fn allow_tokenflow(&mut self) -> ProgramResult {
+        if self.data[CLAIMING_RESTRICTION_BYTE] == 0 {
+            Err(MplxRewardsError::MiningNotRestricted.into())
+        } else {
+            self.data[CLAIMING_RESTRICTION_BYTE] = 0;
+            Ok(())
+        }
+    }
+
+    pub fn is_tokenflow_restricted(&self) -> bool {
+        self.data[CLAIMING_RESTRICTION_BYTE] == 1
+    }
 }
 
 impl IsInitialized for Mining {
     fn is_initialized(&self) -> bool {
-        self.data[0] == <u8>::from(AccountType::Mining)
+        self.data[ACCOUNT_TYPE_BYTE] == <u8>::from(AccountType::Mining)
     }
 }
 
diff --git a/programs/rewards/tests/rewards/fixtures/mplx_rewards.so b/programs/rewards/tests/rewards/fixtures/mplx_rewards.so
index d33fc33..b3fae7f 100755
Binary files a/programs/rewards/tests/rewards/fixtures/mplx_rewards.so and b/programs/rewards/tests/rewards/fixtures/mplx_rewards.so differ
diff --git a/programs/rewards/tests/rewards/tests.rs b/programs/rewards/tests/rewards/tests.rs
index b51efac..725a8a4 100644
--- a/programs/rewards/tests/rewards/tests.rs
+++ b/programs/rewards/tests/rewards/tests.rs
@@ -11,4 +11,6 @@ mod precision;
 mod utils;
 mod withdraw_mining;
 
+mod tokenflow_restrictions;
+
 mod extend_stake;
diff --git a/programs/rewards/tests/rewards/tokenflow_restrictions.rs b/programs/rewards/tests/rewards/tokenflow_restrictions.rs
new file mode 100644
index 0000000..380ed9a
--- /dev/null
+++ b/programs/rewards/tests/rewards/tokenflow_restrictions.rs
@@ -0,0 +1,312 @@
+use crate::utils::*;
+use assert_custom_on_chain_error::AssertCustomOnChainErr;
+use mplx_rewards::{
+    state::{WrappedMining, WrappedRewardPool},
+    utils::LockupPeriod,
+};
+use solana_program::{program_pack::Pack, pubkey::Pubkey};
+use solana_program_test::*;
+use solana_sdk::{clock::SECONDS_PER_DAY, signature::Keypair, signer::Signer};
+use spl_token::state::Account;
+use std::borrow::{Borrow, BorrowMut};
+
+async fn setup() -> (ProgramTestContext, TestRewards, Pubkey) {
+    let test = ProgramTest::new("mplx_rewards", mplx_rewards::ID, None);
+    let mut context = test.start_with_context().await;
+
+    let owner = &context.payer.pubkey();
+
+    let mint = Keypair::new();
+    create_mint(&mut context, &mint, owner).await.unwrap();
+
+    let test_rewards = TestRewards::new(mint.pubkey());
+    test_rewards.initialize_pool(&mut context).await.unwrap();
+
+    // mint token for fill_authority aka wallet who will fill the vault with tokens
+    let rewarder = Keypair::new();
+    create_token_account(
+        &mut context,
+        &rewarder,
+        &test_rewards.token_mint_pubkey,
+        &test_rewards.fill_authority.pubkey(),
+        0,
+    )
+    .await
+    .unwrap();
+    mint_tokens(
+        &mut context,
+        &test_rewards.token_mint_pubkey,
+        &rewarder.pubkey(),
+        1_000_000,
+    )
+    .await
+    .unwrap();
+
+    (context, test_rewards, rewarder.pubkey())
+}
+
+#[tokio::test]
+async fn claim_restricted() {
+    let (mut context, test_rewards, rewarder) = setup().await;
+
+    let (user_a, user_rewards_a, user_mining_a) =
+        create_end_user(&mut context, &test_rewards).await;
+    test_rewards
+        .deposit_mining(
+            &mut context,
+            &user_mining_a,
+            100,
+            LockupPeriod::ThreeMonths,
+            &user_a.pubkey(),
+            &user_mining_a,
+            &user_a.pubkey(),
+        )
+        .await
+        .unwrap();
+
+    // fill vault with tokens
+    let distribution_ends_at = context
+        .banks_client
+        .get_sysvar::<solana_program::clock::Clock>()
+        .await
+        .unwrap()
+        .unix_timestamp as u64
+        + SECONDS_PER_DAY;
+
+    test_rewards
+        .fill_vault(&mut context, &rewarder, 100, distribution_ends_at)
+        .await
+        .unwrap();
+
+    test_rewards.distribute_rewards(&mut context).await.unwrap();
+
+    // restrict claiming
+    test_rewards
+        .restrict_tokenflow(&mut context, &user_mining_a, &user_a.pubkey())
+        .await
+        .unwrap();
+
+    test_rewards
+        .claim(
+            &mut context,
+            &user_a,
+            &user_mining_a,
+            &user_rewards_a.pubkey(),
+        )
+        .await
+        .assert_on_chain_err(mplx_rewards::error::MplxRewardsError::ClaimingRestricted);
+}
+
+#[tokio::test]
+async fn claim_allowed() {
+    let (mut context, test_rewards, rewarder) = setup().await;
+
+    let (user_a, user_rewards_a, user_mining_a) =
+        create_end_user(&mut context, &test_rewards).await;
+    test_rewards
+        .deposit_mining(
+            &mut context,
+            &user_mining_a,
+            100,
+            LockupPeriod::ThreeMonths,
+            &user_a.pubkey(),
+            &user_mining_a,
+            &user_a.pubkey(),
+        )
+        .await
+        .unwrap();
+
+    // fill vault with tokens
+    let distribution_ends_at = context
+        .banks_client
+        .get_sysvar::<solana_program::clock::Clock>()
+        .await
+        .unwrap()
+        .unix_timestamp as u64
+        + SECONDS_PER_DAY;
+
+    test_rewards
+        .fill_vault(&mut context, &rewarder, 100, distribution_ends_at)
+        .await
+        .unwrap();
+
+    test_rewards.distribute_rewards(&mut context).await.unwrap();
+
+    // restrict claiming
+    test_rewards
+        .restrict_tokenflow(&mut context, &user_mining_a, &user_a.pubkey())
+        .await
+        .unwrap();
+
+    test_rewards
+        .claim(
+            &mut context,
+            &user_a,
+            &user_mining_a,
+            &user_rewards_a.pubkey(),
+        )
+        .await
+        .assert_on_chain_err(mplx_rewards::error::MplxRewardsError::ClaimingRestricted);
+
+    // allow claiming
+    test_rewards
+        .allow_tokenflow(&mut context, &user_mining_a, &user_a.pubkey())
+        .await
+        .unwrap();
+
+    // MUST TO ADVANCE TO AVOID CACHING
+    advance_clock_by_ts(&mut context, 2).await;
+
+    test_rewards
+        .claim(
+            &mut context,
+            &user_a,
+            &user_mining_a,
+            &user_rewards_a.pubkey(),
+        )
+        .await
+        .unwrap();
+
+    let user_reward_account_a = get_account(&mut context, &user_rewards_a.pubkey()).await;
+    let user_rewards_a = Account::unpack(user_reward_account_a.data.borrow()).unwrap();
+
+    assert_eq!(user_rewards_a.amount, 100);
+}
+
+#[tokio::test]
+async fn withdraw_restricted() {
+    let (mut context, test_rewards, rewarder) = setup().await;
+
+    let (user_a, _, user_mining_a) = create_end_user(&mut context, &test_rewards).await;
+    test_rewards
+        .deposit_mining(
+            &mut context,
+            &user_mining_a,
+            100,
+            LockupPeriod::ThreeMonths,
+            &user_a.pubkey(),
+            &user_mining_a,
+            &user_a.pubkey(),
+        )
+        .await
+        .unwrap();
+
+    // fill vault with tokens
+    let distribution_ends_at = context
+        .banks_client
+        .get_sysvar::<solana_program::clock::Clock>()
+        .await
+        .unwrap()
+        .unix_timestamp as u64
+        + SECONDS_PER_DAY;
+
+    test_rewards
+        .fill_vault(&mut context, &rewarder, 100, distribution_ends_at)
+        .await
+        .unwrap();
+
+    test_rewards.distribute_rewards(&mut context).await.unwrap();
+
+    // restrict claiming
+    test_rewards
+        .restrict_tokenflow(&mut context, &user_mining_a, &user_a.pubkey())
+        .await
+        .unwrap();
+
+    test_rewards
+        .withdraw_mining(
+            &mut context,
+            &user_mining_a,
+            &user_mining_a,
+            100,
+            &user_a.pubkey(),
+            &user_a.pubkey(),
+        )
+        .await
+        .assert_on_chain_err(mplx_rewards::error::MplxRewardsError::WithdrawalRestricted);
+}
+
+#[tokio::test]
+async fn withdraw_allowed() {
+    let (mut context, test_rewards, rewarder) = setup().await;
+
+    let (user_a, _, user_mining_a) = create_end_user(&mut context, &test_rewards).await;
+    test_rewards
+        .deposit_mining(
+            &mut context,
+            &user_mining_a,
+            100,
+            LockupPeriod::ThreeMonths,
+            &user_a.pubkey(),
+            &user_mining_a,
+            &user_a.pubkey(),
+        )
+        .await
+        .unwrap();
+
+    // fill vault with tokens
+    let distribution_ends_at = context
+        .banks_client
+        .get_sysvar::<solana_program::clock::Clock>()
+        .await
+        .unwrap()
+        .unix_timestamp as u64
+        + SECONDS_PER_DAY;
+
+    test_rewards
+        .fill_vault(&mut context, &rewarder, 100, distribution_ends_at)
+        .await
+        .unwrap();
+
+    test_rewards.distribute_rewards(&mut context).await.unwrap();
+
+    // restrict claiming
+    test_rewards
+        .restrict_tokenflow(&mut context, &user_mining_a, &user_a.pubkey())
+        .await
+        .unwrap();
+
+    test_rewards
+        .withdraw_mining(
+            &mut context,
+            &user_mining_a,
+            &user_mining_a,
+            100,
+            &user_a.pubkey(),
+            &user_a.pubkey(),
+        )
+        .await
+        .assert_on_chain_err(mplx_rewards::error::MplxRewardsError::WithdrawalRestricted);
+
+    // prevent caching
+    advance_clock_by_ts(&mut context, (distribution_ends_at + 1).try_into().unwrap()).await;
+    test_rewards
+        .allow_tokenflow(&mut context, &user_mining_a, &user_a.pubkey())
+        .await
+        .unwrap();
+
+    test_rewards
+        .withdraw_mining(
+            &mut context,
+            &user_mining_a,
+            &user_mining_a,
+            100,
+            &user_a.pubkey(),
+            &user_a.pubkey(),
+        )
+        .await
+        .unwrap();
+
+    let mut reward_pool_account =
+        get_account(&mut context, &test_rewards.reward_pool.pubkey()).await;
+    let reward_pool_data = &mut reward_pool_account.data.borrow_mut();
+    let wrapped_reward_pool = WrappedRewardPool::from_bytes_mut(reward_pool_data).unwrap();
+    let reward_pool = wrapped_reward_pool.pool;
+
+    assert_eq!(reward_pool.total_share, 0);
+
+    let mut mining_account = get_account(&mut context, &user_mining_a).await;
+    let mining_data = &mut mining_account.data.borrow_mut();
+    let wrapped_mining = WrappedMining::from_bytes_mut(mining_data).unwrap();
+    assert_eq!(wrapped_mining.mining.share, 0);
+}
diff --git a/programs/rewards/tests/rewards/utils.rs b/programs/rewards/tests/rewards/utils.rs
index 6fa5e09..4c8e8fd 100644
--- a/programs/rewards/tests/rewards/utils.rs
+++ b/programs/rewards/tests/rewards/utils.rs
@@ -339,6 +339,50 @@ impl TestRewards {
 
         context.banks_client.process_transaction(tx).await
     }
+
+    pub async fn restrict_tokenflow(
+        &self,
+        context: &mut ProgramTestContext,
+        mining_account: &Pubkey,
+        mining_owner: &Pubkey,
+    ) -> BanksClientResult<()> {
+        let tx = Transaction::new_signed_with_payer(
+            &[mplx_rewards::instruction::restrict_tokenflow(
+                &mplx_rewards::id(),
+                &self.deposit_authority.pubkey(),
+                &self.reward_pool.pubkey(),
+                mining_account,
+                mining_owner,
+            )],
+            Some(&context.payer.pubkey()),
+            &[&context.payer, &self.deposit_authority],
+            context.last_blockhash,
+        );
+
+        context.banks_client.process_transaction(tx).await
+    }
+
+    pub async fn allow_tokenflow(
+        &self,
+        context: &mut ProgramTestContext,
+        mining_account: &Pubkey,
+        mining_owner: &Pubkey,
+    ) -> BanksClientResult<()> {
+        let tx = Transaction::new_signed_with_payer(
+            &[mplx_rewards::instruction::allow_tokenflow(
+                &mplx_rewards::id(),
+                &self.deposit_authority.pubkey(),
+                &self.reward_pool.pubkey(),
+                mining_account,
+                mining_owner,
+            )],
+            Some(&context.payer.pubkey()),
+            &[&context.payer, &self.deposit_authority],
+            context.last_blockhash,
+        );
+
+        context.banks_client.process_transaction(tx).await
+    }
 }
 
 pub async fn create_token_account(