From ae9963aec302b723386948883d63281aa4e05481 Mon Sep 17 00:00:00 2001
From: Sergio Chouhy <sergio.chouhy@gmail.com>
Date: Thu, 7 Aug 2025 09:33:54 -0300
Subject: [PATCH] add test

---
 Cargo.lock        |  1 -
 nssa/src/state.rs | 74 +++++++++++++++++++++++++++++++++++------------
 2 files changed, 55 insertions(+), 20 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index fc2f455..07f1410 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2987,7 +2987,6 @@ dependencies = [
  "hex",
  "k256",
  "log",
- "nssa",
  "rand 0.8.5",
  "reqwest 0.11.27",
  "risc0-zkvm 2.3.1 (git+https://github.com/risc0/risc0.git?branch=release-2.3)",
diff --git a/nssa/src/state.rs b/nssa/src/state.rs
index 5310fda..a636859 100644
--- a/nssa/src/state.rs
+++ b/nssa/src/state.rs
@@ -121,49 +121,85 @@ impl V01State {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{
-        public_transaction::{self, WitnessSet},
-        signature::PrivateKey,
-    };
+    use crate::{public_transaction, signature::PrivateKey};
 
-    fn genesis_state_for_tests() -> (V01State, Address) {
-        let account_1 = {
-            let mut this = Account::default();
-            this.program_owner = AuthenticatedTransferProgram::PROGRAM_ID;
-            this.balance = 100;
-            this
-        };
-        let address_1 = Address::new([1; 32]);
-        let public_state = [(address_1.clone(), account_1)].into_iter().collect();
-        (V01State { public_state }, address_1)
+    fn genesis_state_for_tests(balances: &[u128], addresses: &[Address]) -> V01State {
+        assert_eq!(balances.len(), addresses.len());
+        let accounts: Vec<Account> = balances
+            .iter()
+            .map(|balance| {
+                let mut account = Account::default();
+                account.balance = *balance;
+                account.program_owner = AuthenticatedTransferProgram::PROGRAM_ID;
+                account
+            })
+            .collect();
+
+        let public_state = addresses
+            .to_owned()
+            .into_iter()
+            .zip(accounts.into_iter())
+            .collect();
+        V01State { public_state }
     }
 
     fn transfer_transaction_for_tests(
         from: Address,
         from_key: PrivateKey,
+        nonce: u128,
         to: Address,
         balance: u128,
     ) -> PublicTransaction {
         let addresses = vec![from, to];
-        let nonces = vec![0];
+        let nonces = vec![nonce];
         let program_id = AuthenticatedTransferProgram::PROGRAM_ID;
         let message = public_transaction::Message::new(program_id, addresses, nonces, balance);
-        let witness_set = WitnessSet::for_message(&message, &[from_key]);
+        let witness_set = public_transaction::WitnessSet::for_message(&message, &[from_key]);
         PublicTransaction::new(message, witness_set)
     }
 
     #[test]
     fn test_1() {
-        let (mut genesis_state, address) = genesis_state_for_tests();
-        let from = address;
+        let addresses = [Address::new([1; 32])];
+        let balances = [100];
+        let mut genesis_state = genesis_state_for_tests(&balances, &addresses);
+        let from = addresses[0].clone();
         let from_key = PrivateKey(1);
         let to = Address::new([2; 32]);
         let balance_to_move = 5;
-        let tx = transfer_transaction_for_tests(from, from_key, to.clone(), 5);
+        let tx =
+            transfer_transaction_for_tests(from.clone(), from_key, 0, to.clone(), balance_to_move);
         let _ = genesis_state.transition_from_public_transaction(tx);
         assert_eq!(
             genesis_state.get_account_by_address(&to).balance,
             balance_to_move
         );
+        assert_eq!(
+            genesis_state.get_account_by_address(&from).balance,
+            balances[0] - balance_to_move
+        );
+        assert_eq!(genesis_state.get_account_by_address(&from).nonce, 1);
+        assert_eq!(genesis_state.get_account_by_address(&to).nonce, 0);
+    }
+
+    #[test]
+    fn test_2() {
+        let addresses = [Address::new([1; 32]), Address::new([99; 32])];
+        let balances = [100, 200];
+        let mut genesis_state = genesis_state_for_tests(&balances, &addresses);
+        let from = addresses[1].clone();
+        let from_key = PrivateKey(99);
+        let to = addresses[0].clone();
+        let balance_to_move = 8;
+        let to_previous_balance = genesis_state.get_account_by_address(&to).balance;
+        let tx = transfer_transaction_for_tests(from.clone(), from_key, 0, to.clone(), balance_to_move);
+        let _ = genesis_state.transition_from_public_transaction(tx);
+        assert_eq!(genesis_state.get_account_by_address(&to).balance, 108);
+        assert_eq!(
+            genesis_state.get_account_by_address(&from).balance,
+            balances[1] - balance_to_move
+        );
+        assert_eq!(genesis_state.get_account_by_address(&from).nonce, 1);
+        assert_eq!(genesis_state.get_account_by_address(&to).nonce, 0);
     }
 }