From 3e6d8172a40118ad66ad96cea212277455a726b8 Mon Sep 17 00:00:00 2001
From: agureev <artguril@proton.me>
Date: Tue, 23 Jun 2026 13:45:05 +0400
Subject: [PATCH] tests: improve the circuit test

---
 .../privacy_preserving_transaction/circuit.rs | 73 ++++++++++---------
 1 file changed, 37 insertions(+), 36 deletions(-)

diff --git a/lee/state_machine/src/privacy_preserving_transaction/circuit.rs b/lee/state_machine/src/privacy_preserving_transaction/circuit.rs
index d6f98309..c240dbb4 100644
--- a/lee/state_machine/src/privacy_preserving_transaction/circuit.rs
+++ b/lee/state_machine/src/privacy_preserving_transaction/circuit.rs
@@ -356,11 +356,16 @@ mod tests {
             Commitment::new(&recipient_account_id, &expected_private_account_2),
         ];
 
-        let shared_secret_1 =
-            SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0).0;
+        let view_tag_1 =
+            EncryptedAccountData::compute_view_tag(&sender_keys.npk(), &sender_keys.vpk());
+        let view_tag_2 =
+            EncryptedAccountData::compute_view_tag(&recipient_keys.npk(), &recipient_keys.vpk());
 
-        let shared_secret_2 =
-            SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 1).0;
+        let (shared_secret_1, epk_1) =
+            SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0);
+
+        let (shared_secret_2, epk_2) =
+            SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 1);
 
         let (output, proof) = execute_and_prove(
             vec![sender_pre, recipient],
@@ -370,11 +375,8 @@ mod tests {
             .unwrap(),
             vec![
                 InputAccountIdentity::PrivateAuthorizedUpdate {
-                    epk: EphemeralPublicKey(Vec::new()),
-                    view_tag: EncryptedAccountData::compute_view_tag(
-                        &sender_keys.npk(),
-                        &sender_keys.vpk(),
-                    ),
+                    epk: epk_1.clone(),
+                    view_tag: view_tag_1,
                     ssk: shared_secret_1,
                     nsk: sender_keys.nsk,
                     membership_proof: commitment_set
@@ -383,11 +385,8 @@ mod tests {
                     identifier: 0,
                 },
                 InputAccountIdentity::PrivateUnauthorized {
-                    epk: EphemeralPublicKey(Vec::new()),
-                    view_tag: EncryptedAccountData::compute_view_tag(
-                        &recipient_keys.npk(),
-                        &recipient_keys.vpk(),
-                    ),
+                    epk: epk_2.clone(),
+                    view_tag: view_tag_2,
                     npk: recipient_keys.npk(),
                     ssk: shared_secret_2,
                     identifier: 0,
@@ -400,34 +399,36 @@ 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],
-            ),
+        let mut expected_encrypted = vec![
+            EncryptedAccountData {
+                ciphertext: EncryptionScheme::encrypt(
+                    &expected_private_account_1,
+                    &PrivateAccountKind::Regular(0),
+                    &shared_secret_1,
+                    &expected_new_commitments[0],
+                ),
+                epk: epk_1,
+                view_tag: view_tag_1,
+            },
+            EncryptedAccountData {
+                ciphertext: EncryptionScheme::encrypt(
+                    &expected_private_account_2,
+                    &PrivateAccountKind::Regular(0),
+                    &shared_secret_2,
+                    &expected_new_commitments[1],
+                ),
+                epk: epk_2,
+                view_tag: view_tag_2,
+            },
         ];
         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()));
+        expected_encrypted
+            .sort_unstable_by(|a, b| a.ciphertext.as_bytes().cmp(b.ciphertext.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
-                .iter()
-                .map(|e| e.ciphertext.clone())
-                .collect::<Vec<_>>(),
-            expected_ciphertexts,
-        );
+        assert_eq!(output.encrypted_private_post_states, expected_encrypted);
     }
 
     #[test]