Add {En,Dis}ableBalanceCredits test

token/client/src/token.rs

@@ -966,4 +966,40 @@ where
         )
         .await
     }
+
+    /// Enable confidential transfer `Deposit` and `Transfer` instructions for a token account
+    pub async fn confidential_transfer_enable_balance_credits<S2: Signer>(
+        &self,
+        token_account: &Pubkey,
+        authority: &S2,
+    ) -> TokenResult<T::Output> {
+        self.process_ixs(
+            &[confidential_transfer::instruction::enable_balance_credits(
+                &self.program_id,
+                token_account,
+                &authority.pubkey(),
+                &[],
+            )?],
+            &[authority],
+        )
+        .await
+    }
+
+    /// Disable confidential transfer `Deposit` and `Transfer` instructions for a token account
+    pub async fn confidential_transfer_disable_balance_credits<S2: Signer>(
+        &self,
+        token_account: &Pubkey,
+        authority: &S2,
+    ) -> TokenResult<T::Output> {
+        self.process_ixs(
+            &[confidential_transfer::instruction::disable_balance_credits(
+                &self.program_id,
+                token_account,
+                &authority.pubkey(),
+                &[],
+            )?],
+            &[authority],
+        )
+        .await
+    }
 }

token/program-2022-test/tests/confidential_transfer.rs

@@ -121,7 +121,7 @@ async fn ct_configure_token_account() {
 
     token
         .confidential_transfer_configure_token_account(
-            dbg!(&alice_token_account),
+            &alice_token_account,
             &alice,
             alice_elgamal_keypair.public,
             alice_ae_key.encrypt(0_u64),
@@ -134,6 +134,7 @@ async fn ct_configure_token_account() {
         .get_extension::<ConfidentialTransferAccount>()
         .unwrap();
     assert!(!bool::from(&extension.approved));
+    assert!(bool::from(&extension.allow_balance_credits));
     assert_eq!(
         extension.elgamal_pubkey,
         alice_elgamal_keypair.public.into()
@@ -156,3 +157,59 @@ async fn ct_configure_token_account() {
         .unwrap();
     assert!(bool::from(&extension.approved));
 }
+
+#[tokio::test]
+async fn ct_enable_disable_balance_credits() {
+    let ct_mint = ConfidentialTransferMint::default();
+    let mut context = TestContext::new().await;
+    context
+        .init_token_with_mint(vec![
+            ExtensionInitializationParams::ConfidentialTransferMint { ct_mint },
+        ])
+        .await
+        .unwrap();
+
+    let TokenContext { token, alice, .. } = context.token_context.unwrap();
+
+    let alice_token_account = token
+        .create_auxiliary_token_account_with_extension_space(
+            &alice,
+            &alice.pubkey(),
+            vec![ExtensionType::ConfidentialTransferAccount],
+        )
+        .await
+        .unwrap();
+
+    let alice_elgamal_keypair = ElGamalKeypair::new_rand();
+    let alice_ae_key = AeKey::new(&alice, &alice_token_account).unwrap();
+
+    token
+        .confidential_transfer_configure_token_account(
+            &alice_token_account,
+            &alice,
+            alice_elgamal_keypair.public,
+            alice_ae_key.encrypt(0_u64),
+        )
+        .await
+        .unwrap();
+
+    token
+        .confidential_transfer_disable_balance_credits(&alice_token_account, &alice)
+        .await
+        .unwrap();
+    let state = token.get_account_info(&alice_token_account).await.unwrap();
+    let extension = state
+        .get_extension::<ConfidentialTransferAccount>()
+        .unwrap();
+    assert!(!bool::from(&extension.allow_balance_credits));
+
+    token
+        .confidential_transfer_enable_balance_credits(&alice_token_account, &alice)
+        .await
+        .unwrap();
+    let state = token.get_account_info(&alice_token_account).await.unwrap();
+    let extension = state
+        .get_extension::<ConfidentialTransferAccount>()
+        .unwrap();
+    assert!(bool::from(&extension.allow_balance_credits));
+}

token/program-2022/src/extension/confidential_transfer/instruction.rs

@@ -653,7 +653,6 @@ pub fn transfer(
 /// Create a inner `ApplyPendingBalance` instruction
 ///
 /// This instruction is suitable for use with a cross-program `invoke`
-#[allow(clippy::too_many_arguments)]
 pub fn inner_apply_pending_balance(
     token_program_id: &Pubkey,
     token_account: &Pubkey,
@@ -705,13 +704,13 @@ pub fn apply_pending_balance(
     ])
 }
 
-/// Create a `EnableBalanceCredits` instruction
-pub fn enable_balance_credits(
+fn enable_or_disable_balance_credits(
+    instruction: ConfidentialTransferInstruction,
     token_program_id: &Pubkey,
     token_account: &Pubkey,
     authority: &Pubkey,
     multisig_signers: &[&Pubkey],
-) -> Result<Vec<Instruction>, ProgramError> {
+) -> Result<Instruction, ProgramError> {
     check_program_account(token_program_id)?;
     let mut accounts = vec![
         AccountMeta::new(*token_account, false),
@@ -722,36 +721,42 @@ pub fn enable_balance_credits(
         accounts.push(AccountMeta::new_readonly(**multisig_signer, true));
     }
 
-    Ok(vec![encode_instruction(
+    Ok(encode_instruction(
         token_program_id,
         accounts,
-        ConfidentialTransferInstruction::EnableBalanceCredits,
+        instruction,
         &(),
-    )])
+    ))
+}
+
+/// Create a `EnableBalanceCredits` instruction
+pub fn enable_balance_credits(
+    token_program_id: &Pubkey,
+    token_account: &Pubkey,
+    authority: &Pubkey,
+    multisig_signers: &[&Pubkey],
+) -> Result<Instruction, ProgramError> {
+    enable_or_disable_balance_credits(
+        ConfidentialTransferInstruction::EnableBalanceCredits,
+        token_program_id,
+        token_account,
+        authority,
+        multisig_signers,
+    )
 }
 
 /// Create a `DisableBalanceCredits` instruction
-#[cfg(not(target_arch = "bpf"))]
 pub fn disable_balance_credits(
     token_program_id: &Pubkey,
     token_account: &Pubkey,
     authority: &Pubkey,
     multisig_signers: &[&Pubkey],
-) -> Result<Vec<Instruction>, ProgramError> {
-    check_program_account(token_program_id)?;
-    let mut accounts = vec![
-        AccountMeta::new(*token_account, false),
-        AccountMeta::new_readonly(*authority, multisig_signers.is_empty()),
-    ];
-
-    for multisig_signer in multisig_signers.iter() {
-        accounts.push(AccountMeta::new_readonly(**multisig_signer, true));
-    }
-
-    Ok(vec![encode_instruction(
-        token_program_id,
-        accounts,
+) -> Result<Instruction, ProgramError> {
+    enable_or_disable_balance_credits(
         ConfidentialTransferInstruction::DisableBalanceCredits,
-        &(),
-    )])
+        token_program_id,
+        token_account,
+        authority,
+        multisig_signers,
+    )
 }