From a3b64e402106232261bb78ed1b9fa97f59fd92bf Mon Sep 17 00:00:00 2001
From: Sergio Chouhy <sergio.chouhy@gmail.com>
Date: Sun, 10 Aug 2025 10:09:23 -0300
Subject: [PATCH] add test for missing outputs

---
 nssa/src/program.rs                           | 18 ++++++++++++++----
 nssa/src/state.rs                             |  3 ++-
 nssa/src/tests/valid_execution_tests.rs       | 19 ++++++++++++++++++-
 .../bin/{extra_outputs.rs => extra_output.rs} |  0
 .../guest/src/bin/missing_output.rs           | 16 ++++++++++++++++
 5 files changed, 50 insertions(+), 6 deletions(-)
 rename nssa/test_program_methods/guest/src/bin/{extra_outputs.rs => extra_output.rs} (100%)
 create mode 100644 nssa/test_program_methods/guest/src/bin/missing_output.rs

diff --git a/nssa/src/program.rs b/nssa/src/program.rs
index 2b6c041..45a42b7 100644
--- a/nssa/src/program.rs
+++ b/nssa/src/program.rs
@@ -87,12 +87,22 @@ impl Program {
     }
 
     /// A program that produces more output accounts than the inputs it received
-    pub fn extra_outputs_program() -> Self {
-        use test_program_methods::{EXTRA_OUTPUTS_ELF, EXTRA_OUTPUTS_ID};
+    pub fn extra_output_program() -> Self {
+        use test_program_methods::{EXTRA_OUTPUT_ELF, EXTRA_OUTPUT_ID};
 
         Program {
-            id: EXTRA_OUTPUTS_ID,
-            elf: EXTRA_OUTPUTS_ELF,
+            id: EXTRA_OUTPUT_ID,
+            elf: EXTRA_OUTPUT_ELF,
+        }
+    }
+
+    /// A program that produces less output accounts than the inputs it received
+    pub fn missing_output_program() -> Self {
+        use test_program_methods::{MISSING_OUTPUT_ELF, MISSING_OUTPUT_ID};
+
+        Program {
+            id: MISSING_OUTPUT_ID,
+            elf: MISSING_OUTPUT_ELF,
         }
     }
 }
diff --git a/nssa/src/state.rs b/nssa/src/state.rs
index b699ab0..195727f 100644
--- a/nssa/src/state.rs
+++ b/nssa/src/state.rs
@@ -147,7 +147,8 @@ impl V01State {
     /// Include test programs in the builtin programs map
     pub fn with_test_programs(mut self) -> Self {
         self.insert_program(Program::nonce_changer_program());
-        self.insert_program(Program::extra_outputs_program());
+        self.insert_program(Program::extra_output_program());
+        self.insert_program(Program::missing_output_program());
         self
     }
 }
diff --git a/nssa/src/tests/valid_execution_tests.rs b/nssa/src/tests/valid_execution_tests.rs
index 66ec49d..f010c06 100644
--- a/nssa/src/tests/valid_execution_tests.rs
+++ b/nssa/src/tests/valid_execution_tests.rs
@@ -24,7 +24,24 @@ fn test_program_should_fail_if_output_accounts_exceed_inputs() {
     let mut state = V01State::new_with_genesis_accounts(&initial_data).with_test_programs();
     let addresses = vec![Address::new([1; 32])];
     let nonces = vec![];
-    let program_id = Program::extra_outputs_program().id();
+    let program_id = Program::extra_output_program().id();
+    let message = public_transaction::Message::new(program_id, addresses, nonces, 0);
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx);
+
+    assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+}
+
+
+#[test]
+fn test_program_should_fail_with_missing_output_accounts() {
+    let initial_data = [([1; 32], 100)];
+    let mut state = V01State::new_with_genesis_accounts(&initial_data).with_test_programs();
+    let addresses = vec![Address::new([1; 32]), Address::new([2; 32])];
+    let nonces = vec![];
+    let program_id = Program::missing_output_program().id();
     let message = public_transaction::Message::new(program_id, addresses, nonces, 0);
     let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
     let tx = PublicTransaction::new(message, witness_set);
diff --git a/nssa/test_program_methods/guest/src/bin/extra_outputs.rs b/nssa/test_program_methods/guest/src/bin/extra_output.rs
similarity index 100%
rename from nssa/test_program_methods/guest/src/bin/extra_outputs.rs
rename to nssa/test_program_methods/guest/src/bin/extra_output.rs
diff --git a/nssa/test_program_methods/guest/src/bin/missing_output.rs b/nssa/test_program_methods/guest/src/bin/missing_output.rs
new file mode 100644
index 0000000..7857ea5
--- /dev/null
+++ b/nssa/test_program_methods/guest/src/bin/missing_output.rs
@@ -0,0 +1,16 @@
+use nssa_core::account::AccountWithMetadata;
+use risc0_zkvm::guest::env;
+
+fn main() {
+    let input_accounts: Vec<AccountWithMetadata> = env::read();
+    let _instruction_data: u128 = env::read();
+
+    let [pre1, _] = match input_accounts.try_into() {
+        Ok(array) => array,
+        Err(_) => return,
+    };
+
+    let account_pre1 = pre1.account;
+
+    env::commit(&vec![account_pre1]);
+}