Skip to content
This repository was archived by the owner on Mar 11, 2025. It is now read-only.

Commit b6eb64b

Browse files
token-2022: check pending and available balance on confidential token tests (#3145)
* token-2022: check pending and available balance on confidential token tests * token-2022: cargo fmt
1 parent 219b6a1 commit b6eb64b

File tree

1 file changed

+183
-77
lines changed

1 file changed

+183
-77
lines changed

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

Lines changed: 183 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,46 @@ impl ConfidentialTokenAccountMeta {
150150
.unwrap();
151151
meta
152152
}
153+
154+
async fn check_balances<T>(
155+
&self,
156+
token: &Token<T, Keypair>,
157+
expected: ConfidentialTokenAccountBalances,
158+
) where
159+
T: SendTransaction,
160+
{
161+
let state = token.get_account_info(&self.token_account).await.unwrap();
162+
let extension = state
163+
.get_extension::<ConfidentialTransferAccount>()
164+
.unwrap();
165+
166+
assert_eq!(
167+
extension
168+
.pending_balance
169+
.decrypt(&self.elgamal_keypair.secret)
170+
.unwrap(),
171+
expected.pending_balance,
172+
);
173+
assert_eq!(
174+
extension
175+
.available_balance
176+
.decrypt(&self.elgamal_keypair.secret)
177+
.unwrap(),
178+
expected.available_balance,
179+
);
180+
assert_eq!(
181+
self.ae_key
182+
.decrypt(&extension.decryptable_available_balance.try_into().unwrap())
183+
.unwrap(),
184+
expected.decryptable_available_balance,
185+
);
186+
}
187+
}
188+
189+
struct ConfidentialTokenAccountBalances {
190+
pending_balance: u64,
191+
available_balance: u64,
192+
decryptable_available_balance: u64,
153193
}
154194

155195
#[tokio::test]
@@ -411,14 +451,17 @@ async fn ct_deposit() {
411451
assert_eq!(extension.pending_balance_credit_counter, 1.into());
412452
assert_eq!(extension.expected_pending_balance_credit_counter, 0.into());
413453
assert_eq!(extension.actual_pending_balance_credit_counter, 0.into());
414-
assert_ne!(
415-
extension.pending_balance,
416-
zk_token_elgamal::pod::ElGamalCiphertext::zeroed()
417-
);
418-
assert_eq!(
419-
extension.available_balance,
420-
zk_token_elgamal::pod::ElGamalCiphertext::zeroed()
421-
);
454+
455+
alice_meta
456+
.check_balances(
457+
&token,
458+
ConfidentialTokenAccountBalances {
459+
pending_balance: 42,
460+
available_balance: 0,
461+
decryptable_available_balance: 0,
462+
},
463+
)
464+
.await;
422465

423466
let new_decryptable_available_balance = alice_meta.ae_key.encrypt(42_u64);
424467
token
@@ -445,14 +488,17 @@ async fn ct_deposit() {
445488
assert_eq!(extension.pending_balance_credit_counter, 1.into());
446489
assert_eq!(extension.expected_pending_balance_credit_counter, 1.into());
447490
assert_eq!(extension.actual_pending_balance_credit_counter, 1.into());
448-
assert_eq!(
449-
extension.pending_balance,
450-
zk_token_elgamal::pod::ElGamalCiphertext::zeroed()
451-
);
452-
assert_ne!(
453-
extension.available_balance,
454-
zk_token_elgamal::pod::ElGamalCiphertext::zeroed()
455-
);
491+
492+
alice_meta
493+
.check_balances(
494+
&token,
495+
ConfidentialTokenAccountBalances {
496+
pending_balance: 0,
497+
available_balance: 42,
498+
decryptable_available_balance: 42,
499+
},
500+
)
501+
.await;
456502
}
457503

458504
#[tokio::test]
@@ -474,6 +520,7 @@ async fn ct_withdraw() {
474520
decimals,
475521
..
476522
} = context.token_context.unwrap();
523+
477524
let alice_meta =
478525
ConfidentialTokenAccountMeta::with_tokens(&token, &alice, &mint_authority, 42, decimals)
479526
.await;
@@ -503,15 +550,17 @@ async fn ct_withdraw() {
503550
.await
504551
.unwrap();
505552
assert_eq!(state.base.amount, 21);
506-
let extension = state
507-
.get_extension::<ConfidentialTransferAccount>()
508-
.unwrap();
509-
assert_eq!(
510-
alice_meta
511-
.ae_key
512-
.decrypt(&extension.decryptable_available_balance.try_into().unwrap()),
513-
Some(21),
514-
);
553+
554+
alice_meta
555+
.check_balances(
556+
&token,
557+
ConfidentialTokenAccountBalances {
558+
pending_balance: 0,
559+
available_balance: 21,
560+
decryptable_available_balance: 21,
561+
},
562+
)
563+
.await;
515564

516565
token
517566
.confidential_transfer_withdraw(
@@ -532,15 +581,17 @@ async fn ct_withdraw() {
532581
.await
533582
.unwrap();
534583
assert_eq!(state.base.amount, 42);
535-
let extension = state
536-
.get_extension::<ConfidentialTransferAccount>()
537-
.unwrap();
538-
assert_eq!(
539-
alice_meta
540-
.ae_key
541-
.decrypt(&extension.decryptable_available_balance.try_into().unwrap()),
542-
Some(0),
543-
);
584+
585+
alice_meta
586+
.check_balances(
587+
&token,
588+
ConfidentialTokenAccountBalances {
589+
pending_balance: 0,
590+
available_balance: 0,
591+
decryptable_available_balance: 0,
592+
},
593+
)
594+
.await;
544595

545596
token
546597
.confidential_transfer_empty_account(
@@ -591,6 +642,17 @@ async fn ct_transfer() {
591642
.await
592643
.unwrap();
593644

645+
alice_meta
646+
.check_balances(
647+
&token,
648+
ConfidentialTokenAccountBalances {
649+
pending_balance: 0,
650+
available_balance: 42,
651+
decryptable_available_balance: 42,
652+
},
653+
)
654+
.await;
655+
594656
// Self-transfer of N tokens
595657
token
596658
.confidential_transfer_transfer(
@@ -600,11 +662,22 @@ async fn ct_transfer() {
600662
42, // amount
601663
42, // available balance
602664
&alice_meta.elgamal_keypair,
603-
alice_meta.ae_key.encrypt(42_u64),
665+
alice_meta.ae_key.encrypt(0_u64),
604666
)
605667
.await
606668
.unwrap();
607669

670+
alice_meta
671+
.check_balances(
672+
&token,
673+
ConfidentialTokenAccountBalances {
674+
pending_balance: 42,
675+
available_balance: 0,
676+
decryptable_available_balance: 0,
677+
},
678+
)
679+
.await;
680+
608681
token
609682
.confidential_transfer_apply_pending_balance(
610683
&alice_meta.token_account,
@@ -615,19 +688,16 @@ async fn ct_transfer() {
615688
.await
616689
.unwrap();
617690

618-
let state = token
619-
.get_account_info(&alice_meta.token_account)
620-
.await
621-
.unwrap();
622-
let extension = state
623-
.get_extension::<ConfidentialTransferAccount>()
624-
.unwrap();
625-
assert_eq!(
626-
alice_meta
627-
.ae_key
628-
.decrypt(&extension.decryptable_available_balance.try_into().unwrap()),
629-
Some(42),
630-
);
691+
alice_meta
692+
.check_balances(
693+
&token,
694+
ConfidentialTokenAccountBalances {
695+
pending_balance: 0,
696+
available_balance: 42,
697+
decryptable_available_balance: 42,
698+
},
699+
)
700+
.await;
631701

632702
token
633703
.confidential_transfer_transfer(
@@ -642,19 +712,16 @@ async fn ct_transfer() {
642712
.await
643713
.unwrap();
644714

645-
let state = token
646-
.get_account_info(&alice_meta.token_account)
647-
.await
648-
.unwrap();
649-
let extension = state
650-
.get_extension::<ConfidentialTransferAccount>()
651-
.unwrap();
652-
assert_eq!(
653-
alice_meta
654-
.ae_key
655-
.decrypt(&extension.decryptable_available_balance.try_into().unwrap()),
656-
Some(0),
657-
);
715+
alice_meta
716+
.check_balances(
717+
&token,
718+
ConfidentialTokenAccountBalances {
719+
pending_balance: 0,
720+
available_balance: 0,
721+
decryptable_available_balance: 0,
722+
},
723+
)
724+
.await;
658725

659726
token
660727
.confidential_transfer_empty_account(
@@ -687,6 +754,8 @@ async fn ct_transfer() {
687754
let extension = state
688755
.get_extension::<ConfidentialTransferAccount>()
689756
.unwrap();
757+
758+
// TODO: verify bob_meta pending and available balance once syscall lands
690759
assert_eq!(
691760
bob_meta
692761
.ae_key
@@ -711,6 +780,8 @@ async fn ct_transfer() {
711780
let extension = state
712781
.get_extension::<ConfidentialTransferAccount>()
713782
.unwrap();
783+
784+
// TODO: verify bob_meta pending and available balance once syscall lands
714785
assert_eq!(
715786
bob_meta
716787
.ae_key
@@ -769,7 +840,20 @@ async fn ct_transfer_with_fee() {
769840
.await
770841
.unwrap();
771842

772-
// Self-transfer of N tokens
843+
alice_meta
844+
.check_balances(
845+
&token,
846+
ConfidentialTokenAccountBalances {
847+
pending_balance: 0,
848+
available_balance: 100,
849+
decryptable_available_balance: 100,
850+
},
851+
)
852+
.await;
853+
854+
// Fee is 2.5%, so what is left is 97 in Alice account
855+
//
856+
// TODO: make self transfers not not take fees
773857
token
774858
.confidential_transfer_transfer_with_fee(
775859
&alice_meta.token_account,
@@ -778,13 +862,23 @@ async fn ct_transfer_with_fee() {
778862
100, // amount
779863
100, // available balance
780864
&alice_meta.elgamal_keypair,
781-
alice_meta.ae_key.encrypt(100_u64),
865+
alice_meta.ae_key.encrypt(0_u64),
782866
&epoch_info,
783867
)
784868
.await
785869
.unwrap();
786870

787-
// Fee is 2.5%, so what is left in 97 in Alice account
871+
alice_meta
872+
.check_balances(
873+
&token,
874+
ConfidentialTokenAccountBalances {
875+
pending_balance: 97,
876+
available_balance: 0,
877+
decryptable_available_balance: 0,
878+
},
879+
)
880+
.await;
881+
788882
token
789883
.confidential_transfer_apply_pending_balance(
790884
&alice_meta.token_account,
@@ -795,19 +889,16 @@ async fn ct_transfer_with_fee() {
795889
.await
796890
.unwrap();
797891

798-
let state = token
799-
.get_account_info(&alice_meta.token_account)
800-
.await
801-
.unwrap();
802-
let extension = state
803-
.get_extension::<ConfidentialTransferAccount>()
804-
.unwrap();
805-
assert_eq!(
806-
alice_meta
807-
.ae_key
808-
.decrypt(&extension.decryptable_available_balance.try_into().unwrap()),
809-
Some(97),
810-
);
892+
alice_meta
893+
.check_balances(
894+
&token,
895+
ConfidentialTokenAccountBalances {
896+
pending_balance: 0,
897+
available_balance: 97,
898+
decryptable_available_balance: 97,
899+
},
900+
)
901+
.await;
811902

812903
token
813904
.confidential_transfer_transfer_with_fee(
@@ -823,6 +914,17 @@ async fn ct_transfer_with_fee() {
823914
.await
824915
.unwrap();
825916

917+
alice_meta
918+
.check_balances(
919+
&token,
920+
ConfidentialTokenAccountBalances {
921+
pending_balance: 0,
922+
available_balance: 0,
923+
decryptable_available_balance: 0,
924+
},
925+
)
926+
.await;
927+
826928
let state = token
827929
.get_account_info(&alice_meta.token_account)
828930
.await
@@ -878,6 +980,8 @@ async fn ct_transfer_with_fee() {
878980
let extension = state
879981
.get_extension::<ConfidentialTransferAccount>()
880982
.unwrap();
983+
984+
// TODO: check pending and available balance once curve syscall lands
881985
assert_eq!(
882986
bob_meta
883987
.ae_key
@@ -902,6 +1006,8 @@ async fn ct_transfer_with_fee() {
9021006
let extension = state
9031007
.get_extension::<ConfidentialTransferAccount>()
9041008
.unwrap();
1009+
1010+
// TODO: check pending and available balance once curve syscall lands
9051011
assert_eq!(
9061012
bob_meta
9071013
.ae_key

0 commit comments

Comments
 (0)