From 24a3165634ab2a4599e14c5db8fffcfaf10b0b6d Mon Sep 17 00:00:00 2001
From: Sergio Chouhy <sergio.chouhy@gmail.com>
Date: Wed, 13 Aug 2025 00:43:17 -0300
Subject: [PATCH] add tests

---
 nssa/src/program.rs           |  1 +
 nssa/src/state.rs             |  4 +-
 nssa/src/tests/state_tests.rs | 95 ++++++++++++++++++++++++++++++++++-
 3 files changed, 97 insertions(+), 3 deletions(-)

diff --git a/nssa/src/program.rs b/nssa/src/program.rs
index 5b4a1e8..5baa996 100644
--- a/nssa/src/program.rs
+++ b/nssa/src/program.rs
@@ -8,6 +8,7 @@ use serde::Serialize;
 
 use crate::error::NssaError;
 
+#[derive(Debug, PartialEq, Eq)]
 pub struct Program {
     pub(crate) id: ProgramId,
     pub(crate) elf: &'static [u8],
diff --git a/nssa/src/state.rs b/nssa/src/state.rs
index 690919f..343fd6e 100644
--- a/nssa/src/state.rs
+++ b/nssa/src/state.rs
@@ -5,8 +5,8 @@ use nssa_core::{account::Account, program::ProgramId};
 use std::collections::HashMap;
 
 pub struct V01State {
-    public_state: HashMap<Address, Account>,
-    builtin_programs: HashMap<ProgramId, Program>,
+    pub(crate) public_state: HashMap<Address, Account>,
+    pub(crate) builtin_programs: HashMap<ProgramId, Program>,
 }
 
 impl V01State {
diff --git a/nssa/src/tests/state_tests.rs b/nssa/src/tests/state_tests.rs
index 5acbed8..acf81a2 100644
--- a/nssa/src/tests/state_tests.rs
+++ b/nssa/src/tests/state_tests.rs
@@ -1,8 +1,10 @@
+use std::collections::HashMap;
+
 use crate::{
     Address, PublicKey, PublicTransaction, V01State, error::NssaError, program::Program,
     public_transaction, signature::PrivateKey,
 };
-use nssa_core::account::Account;
+use nssa_core::{account::Account, program::ProgramId};
 
 fn transfer_transaction(
     from: Address,
@@ -20,6 +22,97 @@ fn transfer_transaction(
     PublicTransaction::new(message, witness_set)
 }
 
+#[test]
+fn test_new_with_genesis() {
+    let key1 = PrivateKey::try_new([1; 32]).unwrap();
+    let key2 = PrivateKey::try_new([2; 32]).unwrap();
+    let addr1 = Address::from_public_key(&PublicKey::new(&key1));
+    let addr2 = Address::from_public_key(&PublicKey::new(&key2));
+    let initial_data = [(*addr1.value(), 100u128), (*addr2.value(), 151u128)];
+    let program = Program::authenticated_transfer_program();
+    let expected_public_state = {
+        let mut this = HashMap::new();
+        this.insert(
+            addr1,
+            Account {
+                balance: 100,
+                program_owner: program.id(),
+                ..Account::default()
+            },
+        );
+        this.insert(
+            addr2,
+            Account {
+                balance: 151,
+                program_owner: program.id(),
+                ..Account::default()
+            },
+        );
+        this
+    };
+    let expected_builtin_programs = {
+        let mut this = HashMap::new();
+        this.insert(program.id(), program);
+        this
+    };
+
+    let state = V01State::new_with_genesis_accounts(&initial_data);
+
+    assert_eq!(state.public_state, expected_public_state);
+    assert_eq!(state.builtin_programs, expected_builtin_programs);
+}
+
+#[test]
+fn test_insert_program() {
+    let mut state = V01State::new_with_genesis_accounts(&[]);
+    let program_to_insert = Program::simple_balance_transfer();
+    let program_id = program_to_insert.id();
+    assert!(!state.builtin_programs.contains_key(&program_id));
+
+    state.insert_program(program_to_insert);
+
+    assert!(state.builtin_programs.contains_key(&program_id));
+}
+
+#[test]
+fn test_get_account_by_address_non_default_account() {
+    let key = PrivateKey::try_new([1; 32]).unwrap();
+    let addr = Address::from_public_key(&PublicKey::new(&key));
+    let initial_data = [(*addr.value(), 100u128)];
+    let state = V01State::new_with_genesis_accounts(&initial_data);
+    let expected_account = state.public_state.get(&addr).unwrap();
+
+    let account = state.get_account_by_address(&addr);
+
+    assert_eq!(&account, expected_account);
+}
+
+#[test]
+fn test_get_account_by_address_default_account() {
+    let addr2 = Address::new([0; 32]);
+    let state = V01State::new_with_genesis_accounts(&[]);
+    let expected_account = Account::default();
+
+    let account = state.get_account_by_address(&addr2);
+
+    assert_eq!(account, expected_account);
+}
+
+#[test]
+fn test_builtin_programs_getter() {
+    let state = V01State::new_with_genesis_accounts(&[]);
+    let program = Program::authenticated_transfer_program();
+    let expected_builtin_programs = {
+        let mut this = HashMap::new();
+        this.insert(program.id(), program);
+        this
+    };
+
+    let builtin_programs = state.builtin_programs();
+
+    assert_eq!(builtin_programs, &state.builtin_programs);
+}
+
 #[test]
 fn transition_from_authenticated_transfer_program_invocation_default_account_destination() {
     let key = PrivateKey::try_new([1; 32]).unwrap();