feat: add flash swap integration tests

nssa/src/state.rs

@@ -3479,4 +3479,320 @@ pub mod tests {
         let state_from_bytes: V03State = borsh::from_slice(&bytes).unwrap();
         assert_eq!(state, state_from_bytes);
     }
+
+    // ── Flash Swap integration tests ──────────────────────────────────────────
+
+    /// Mirror of the guest `FlashSwapInstruction` enum so we can serialise
+    /// instructions on the host side.
+    #[derive(serde::Serialize, serde::Deserialize)]
+    enum FlashSwapInstruction {
+        Initiate {
+            token_program_id: ProgramId,
+            callback_program_id: ProgramId,
+            amount_out: u128,
+            callback_instruction_data: Vec<u32>,
+            vault_after_transfer: AccountWithMetadata,
+            receiver_after_transfer: AccountWithMetadata,
+            vault_after_callback: AccountWithMetadata,
+        },
+        InvariantCheck {
+            min_vault_balance: u128,
+        },
+    }
+
+    /// Mirror of the guest `CallbackInstruction`.
+    #[derive(serde::Serialize, serde::Deserialize)]
+    struct CallbackInstruction {
+        return_funds: bool,
+        token_program_id: ProgramId,
+        amount: u128,
+        vault_after_return: Option<AccountWithMetadata>,
+        receiver_after_return: Option<AccountWithMetadata>,
+    }
+
+    /// Build a flash-swap `PublicTransaction` ready for execution.
+    ///
+    /// `vault_id` / `receiver_id` are the two accounts passed to the initiator.
+    /// `instruction` is the already-built `FlashSwapInstruction`.
+    fn build_flash_swap_tx(
+        initiator: &Program,
+        vault_id: AccountId,
+        receiver_id: AccountId,
+        instruction: FlashSwapInstruction,
+    ) -> PublicTransaction {
+        let message = public_transaction::Message::try_new(
+            initiator.id(),
+            vec![vault_id, receiver_id],
+            vec![],          // no signers — vault is PDA-authorised
+            instruction,
+        )
+        .unwrap();
+        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+        PublicTransaction::new(message, witness_set)
+    }
+
+    #[test]
+    fn flash_swap_successful() {
+        let initiator = Program::flash_swap_initiator();
+        let callback = Program::flash_swap_callback();
+        let token = Program::authenticated_transfer_program();
+
+        let vault_id = AccountId::from((&initiator.id(), &PdaSeed::new([0u8; 32])));
+        let receiver_id = AccountId::from((&callback.id(), &PdaSeed::new([1u8; 32])));
+
+        let initial_balance: u128 = 1000;
+        let amount_out: u128 = 100;
+
+        let vault_account = Account {
+            program_owner: token.id(),
+            balance: initial_balance,
+            ..Account::default()
+        };
+        let receiver_account = Account {
+            program_owner: token.id(),
+            balance: 0,
+            ..Account::default()
+        };
+
+        let mut state = V03State::new_with_genesis_accounts(&[], &[]).with_test_programs();
+        state.force_insert_account(vault_id, vault_account);
+        state.force_insert_account(receiver_id, receiver_account);
+
+        // Pre-simulated intermediate states:
+        // After transfer (vault→receiver, amount_out):
+        let vault_after_transfer = AccountWithMetadata::new(
+            Account { program_owner: token.id(), balance: initial_balance - amount_out, ..Account::default() },
+            false, vault_id,
+        );
+        let receiver_after_transfer = AccountWithMetadata::new(
+            Account { program_owner: token.id(), balance: amount_out, ..Account::default() },
+            false, receiver_id,
+        );
+
+        // After callback returns funds (receiver→vault, amount_out):
+        let vault_after_callback = AccountWithMetadata::new(
+            Account { program_owner: token.id(), balance: initial_balance, ..Account::default() },
+            false, vault_id,
+        );
+
+        // Callback instruction: return funds
+        let cb_instruction = CallbackInstruction {
+            return_funds: true,
+            token_program_id: token.id(),
+            amount: amount_out,
+            vault_after_return: Some(AccountWithMetadata::new(
+                Account { program_owner: token.id(), balance: initial_balance, ..Account::default() },
+                false, vault_id,
+            )),
+            receiver_after_return: Some(AccountWithMetadata::new(
+                Account { program_owner: token.id(), balance: 0, ..Account::default() },
+                false, receiver_id,
+            )),
+        };
+        let cb_data = Program::serialize_instruction(cb_instruction).unwrap();
+
+        let instruction = FlashSwapInstruction::Initiate {
+            token_program_id: token.id(),
+            callback_program_id: callback.id(),
+            amount_out,
+            callback_instruction_data: cb_data,
+            vault_after_transfer,
+            receiver_after_transfer,
+            vault_after_callback,
+        };
+
+        let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction);
+        let result = state.transition_from_public_transaction(&tx, 1, 0);
+        assert!(result.is_ok(), "flash swap should succeed: {result:?}");
+
+        // Vault balance restored, receiver back to 0
+        assert_eq!(state.get_account_by_id(vault_id).balance, initial_balance);
+        assert_eq!(state.get_account_by_id(receiver_id).balance, 0);
+    }
+
+    #[test]
+    fn flash_swap_callback_keeps_funds_rollback() {
+        let initiator = Program::flash_swap_initiator();
+        let callback = Program::flash_swap_callback();
+        let token = Program::authenticated_transfer_program();
+
+        let vault_id = AccountId::from((&initiator.id(), &PdaSeed::new([0u8; 32])));
+        let receiver_id = AccountId::from((&callback.id(), &PdaSeed::new([1u8; 32])));
+
+        let initial_balance: u128 = 1000;
+        let amount_out: u128 = 100;
+
+        let vault_account = Account {
+            program_owner: token.id(),
+            balance: initial_balance,
+            ..Account::default()
+        };
+        let receiver_account = Account {
+            program_owner: token.id(),
+            balance: 0,
+            ..Account::default()
+        };
+
+        let mut state = V03State::new_with_genesis_accounts(&[], &[]).with_test_programs();
+        state.force_insert_account(vault_id, vault_account);
+        state.force_insert_account(receiver_id, receiver_account);
+
+        // Pre-simulated intermediate states (same as successful case for steps 1-2):
+        let vault_after_transfer = AccountWithMetadata::new(
+            Account { program_owner: token.id(), balance: initial_balance - amount_out, ..Account::default() },
+            false, vault_id,
+        );
+        let receiver_after_transfer = AccountWithMetadata::new(
+            Account { program_owner: token.id(), balance: amount_out, ..Account::default() },
+            false, receiver_id,
+        );
+
+        // After callback that does NOT return funds — vault stays drained:
+        let vault_after_callback = AccountWithMetadata::new(
+            Account { program_owner: token.id(), balance: initial_balance - amount_out, ..Account::default() },
+            false, vault_id,
+        );
+
+        // Callback instruction: do NOT return funds
+        let cb_instruction = CallbackInstruction {
+            return_funds: false,
+            token_program_id: token.id(),
+            amount: amount_out,
+            vault_after_return: None,
+            receiver_after_return: None,
+        };
+        let cb_data = Program::serialize_instruction(cb_instruction).unwrap();
+
+        let instruction = FlashSwapInstruction::Initiate {
+            token_program_id: token.id(),
+            callback_program_id: callback.id(),
+            amount_out,
+            callback_instruction_data: cb_data,
+            vault_after_transfer,
+            receiver_after_transfer,
+            vault_after_callback,
+        };
+
+        let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction);
+        let result = state.transition_from_public_transaction(&tx, 1, 0);
+
+        // Invariant check fails → entire tx rolls back
+        assert!(result.is_err(), "flash swap should fail when callback keeps funds");
+
+        // State unchanged (rollback)
+        assert_eq!(state.get_account_by_id(vault_id).balance, initial_balance);
+        assert_eq!(state.get_account_by_id(receiver_id).balance, 0);
+    }
+
+    #[test]
+    fn flash_swap_self_call_targets_correct_program() {
+        // Zero-amount flash swap: the invariant self-call still runs and succeeds
+        // because vault balance doesn't decrease.
+        let initiator = Program::flash_swap_initiator();
+        let callback = Program::flash_swap_callback();
+        let token = Program::authenticated_transfer_program();
+
+        let vault_id = AccountId::from((&initiator.id(), &PdaSeed::new([0u8; 32])));
+        let receiver_id = AccountId::from((&callback.id(), &PdaSeed::new([1u8; 32])));
+
+        let initial_balance: u128 = 1000;
+
+        let vault_account = Account {
+            program_owner: token.id(),
+            balance: initial_balance,
+            ..Account::default()
+        };
+        let receiver_account = Account {
+            program_owner: token.id(),
+            balance: 0,
+            ..Account::default()
+        };
+
+        let mut state = V03State::new_with_genesis_accounts(&[], &[]).with_test_programs();
+        state.force_insert_account(vault_id, vault_account);
+        state.force_insert_account(receiver_id, receiver_account);
+
+        // Zero-amount transfer: states remain unchanged after transfer
+        let vault_after_transfer = AccountWithMetadata::new(
+            Account { program_owner: token.id(), balance: initial_balance, ..Account::default() },
+            false, vault_id,
+        );
+        let receiver_after_transfer = AccountWithMetadata::new(
+            Account { program_owner: token.id(), balance: 0, ..Account::default() },
+            false, receiver_id,
+        );
+        // Callback with zero amount, return_funds=true (no-op effectively)
+        let vault_after_callback = AccountWithMetadata::new(
+            Account { program_owner: token.id(), balance: initial_balance, ..Account::default() },
+            false, vault_id,
+        );
+
+        let cb_instruction = CallbackInstruction {
+            return_funds: true,
+            token_program_id: token.id(),
+            amount: 0,
+            vault_after_return: Some(AccountWithMetadata::new(
+                Account { program_owner: token.id(), balance: initial_balance, ..Account::default() },
+                false, vault_id,
+            )),
+            receiver_after_return: Some(AccountWithMetadata::new(
+                Account { program_owner: token.id(), balance: 0, ..Account::default() },
+                false, receiver_id,
+            )),
+        };
+        let cb_data = Program::serialize_instruction(cb_instruction).unwrap();
+
+        let instruction = FlashSwapInstruction::Initiate {
+            token_program_id: token.id(),
+            callback_program_id: callback.id(),
+            amount_out: 0,
+            callback_instruction_data: cb_data,
+            vault_after_transfer,
+            receiver_after_transfer,
+            vault_after_callback,
+        };
+
+        let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction);
+        let result = state.transition_from_public_transaction(&tx, 1, 0);
+        assert!(result.is_ok(), "zero-amount flash swap should succeed: {result:?}");
+    }
+
+    #[test]
+    fn flash_swap_standalone_invariant_check_rejected() {
+        // Calling InvariantCheck directly (not as a chained self-call) should fail
+        // because caller_program_id will be None.
+        let initiator = Program::flash_swap_initiator();
+        let token = Program::authenticated_transfer_program();
+
+        let vault_id = AccountId::from((&initiator.id(), &PdaSeed::new([0u8; 32])));
+
+        let vault_account = Account {
+            program_owner: token.id(),
+            balance: 1000,
+            ..Account::default()
+        };
+
+        let mut state = V03State::new_with_genesis_accounts(&[], &[]).with_test_programs();
+        state.force_insert_account(vault_id, vault_account);
+
+        let instruction = FlashSwapInstruction::InvariantCheck {
+            min_vault_balance: 1000,
+        };
+
+        let message = public_transaction::Message::try_new(
+            initiator.id(),
+            vec![vault_id],
+            vec![],
+            instruction,
+        )
+        .unwrap();
+        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+        let tx = PublicTransaction::new(message, witness_set);
+
+        let result = state.transition_from_public_transaction(&tx, 1, 0);
+        assert!(
+            result.is_err(),
+            "standalone InvariantCheck should be rejected (caller_program_id is None)"
+        );
+    }
 }