From a5cb2e484b66d5edc7acea74006efe19df16d776 Mon Sep 17 00:00:00 2001
From: agureev <artguril@proton.me>
Date: Mon, 22 Jun 2026 23:06:21 +0400
Subject: [PATCH] test: change circuit encryption/decryption test

---
 .../privacy_preserving_transaction/circuit.rs | 51 ++++++++++---------
 1 file changed, 28 insertions(+), 23 deletions(-)

diff --git a/lee/state_machine/src/privacy_preserving_transaction/circuit.rs b/lee/state_machine/src/privacy_preserving_transaction/circuit.rs
index 87c7c5bb..d6f98309 100644
--- a/lee/state_machine/src/privacy_preserving_transaction/circuit.rs
+++ b/lee/state_machine/src/privacy_preserving_transaction/circuit.rs
@@ -204,7 +204,6 @@ mod tests {
             &output.encrypted_private_post_states[idx].ciphertext,
             ssk,
             &output.new_commitments[idx],
-            u32::try_from(idx).expect("idx fits in u32"),
         )
         .unwrap();
         kind
@@ -296,7 +295,6 @@ mod tests {
             &output.encrypted_private_post_states[0].ciphertext,
             &shared_secret,
             &output.new_commitments[0],
-            0,
         )
         .unwrap();
         assert_eq!(recipient_post, expected_recipient_post);
@@ -328,7 +326,7 @@ mod tests {
 
         let mut commitment_set = CommitmentSet::with_capacity(2);
         commitment_set.extend(std::slice::from_ref(&commitment_sender));
-        let expected_new_nullifiers = vec![
+        let mut expected_new_nullifiers = vec![
             (
                 Nullifier::for_account_update(&commitment_sender, &sender_keys.nsk),
                 commitment_set.digest(),
@@ -353,7 +351,7 @@ mod tests {
             nonce: Nonce::private_account_nonce_init(&recipient_account_id),
             ..Default::default()
         };
-        let expected_new_commitments = vec![
+        let mut expected_new_commitments = vec![
             Commitment::new(&sender_account_id, &expected_private_account_1),
             Commitment::new(&recipient_account_id, &expected_private_account_2),
         ];
@@ -402,27 +400,34 @@ mod tests {
         assert!(proof.is_valid_for(&output));
         assert!(output.public_pre_states.is_empty());
         assert!(output.public_post_states.is_empty());
+        let mut expected_ciphertexts = vec![
+            EncryptionScheme::encrypt(
+                &expected_private_account_1,
+                &PrivateAccountKind::Regular(0),
+                &shared_secret_1,
+                &expected_new_commitments[0],
+            ),
+            EncryptionScheme::encrypt(
+                &expected_private_account_2,
+                &PrivateAccountKind::Regular(0),
+                &shared_secret_2,
+                &expected_new_commitments[1],
+            ),
+        ];
+        expected_new_commitments.sort_unstable_by_key(Commitment::to_byte_array);
+        expected_new_nullifiers.sort_unstable_by_key(|entry| entry.0.to_byte_array());
+        expected_ciphertexts.sort_unstable_by(|a, b| a.as_bytes().cmp(b.as_bytes()));
+
         assert_eq!(output.new_commitments, expected_new_commitments);
         assert_eq!(output.new_nullifiers, expected_new_nullifiers);
-        assert_eq!(output.encrypted_private_post_states.len(), 2);
-
-        let (_identifier, sender_post) = EncryptionScheme::decrypt(
-            &output.encrypted_private_post_states[0].ciphertext,
-            &shared_secret_1,
-            &expected_new_commitments[0],
-            0,
-        )
-        .unwrap();
-        assert_eq!(sender_post, expected_private_account_1);
-
-        let (_identifier, recipient_post) = EncryptionScheme::decrypt(
-            &output.encrypted_private_post_states[1].ciphertext,
-            &shared_secret_2,
-            &expected_new_commitments[1],
-            1,
-        )
-        .unwrap();
-        assert_eq!(recipient_post, expected_private_account_2);
+        assert_eq!(
+            output
+                .encrypted_private_post_states
+                .iter()
+                .map(|e| e.ciphertext.clone())
+                .collect::<Vec<_>>(),
+            expected_ciphertexts,
+        );
     }
 
     #[test]