From 5098298ea98749abfd169fd1c9e423dc1ec6d8b0 Mon Sep 17 00:00:00 2001
From: Marvin Jones <jonesmc8@gmail.com>
Date: Thu, 21 May 2026 19:12:34 -0400
Subject: [PATCH] tps fix

---
 integration_tests/tests/tps.rs | 67 +++++++++++++++++++++++++++++++++-
 test_fixtures/src/lib.rs       | 15 ++++++--
 2 files changed, 77 insertions(+), 5 deletions(-)

diff --git a/integration_tests/tests/tps.rs b/integration_tests/tests/tps.rs
index 22550bb0..0a6a9038 100644
--- a/integration_tests/tests/tps.rs
+++ b/integration_tests/tests/tps.rs
@@ -11,7 +11,7 @@
 
 use std::time::{Duration, Instant};
 
-use anyhow::Result;
+use anyhow::{Context as _, Result};
 use bytesize::ByteSize;
 use common::transaction::NSSATransaction;
 use integration_tests::{TestContext, config::SequencerPartialConfig};
@@ -66,7 +66,64 @@ impl TpsTestManager {
         Duration::from_secs_f64(number_transactions as f64 / self.target_tps as f64)
     }
 
+    /// Claim funds from each account's vault PDA into the account itself.
+    ///
+    /// `GenesisAction::SupplyAccount` funds vault PDAs (not accounts directly), so this step is
+    /// required before sending `authenticated_transfer` transactions from these accounts.
+    /// All claim transactions are submitted at once and then confirmed sequentially.
+    /// After this call every account has nonce 1, so `build_public_txs` must be called after it.
+    pub async fn claim_vault_funds(
+        &self,
+        sequencer_client: &sequencer_service_rpc::SequencerClient,
+    ) -> Result<()> {
+        let vault_program_id = Program::vault().id();
+
+        let mut tx_hashes = Vec::with_capacity(self.public_keypairs.len());
+        for (private_key, account_id) in &self.public_keypairs {
+            let owner_vault_id =
+                vault_core::compute_vault_account_id(vault_program_id, *account_id);
+            let message = putx::Message::try_new(
+                vault_program_id,
+                vec![*account_id, owner_vault_id],
+                vec![Nonce(0_u128)],
+                vault_core::Instruction::Claim { amount: 10 },
+            )
+            .context("Failed to build vault claim message")?;
+            let witness_set =
+                nssa::public_transaction::WitnessSet::for_message(&message, &[private_key]);
+            let tx = PublicTransaction::new(message, witness_set);
+            let hash = sequencer_client
+                .send_transaction(NSSATransaction::Public(tx))
+                .await
+                .context("Failed to submit vault claim")?;
+            tx_hashes.push(hash);
+        }
+
+        let deadline = Instant::now() + Duration::from_secs(300);
+        for (i, tx_hash) in tx_hashes.iter().enumerate() {
+            loop {
+                anyhow::ensure!(
+                    Instant::now() < deadline,
+                    "Vault claims timed out after 5 minutes ({i}/{} confirmed)",
+                    tx_hashes.len()
+                );
+                let found = sequencer_client
+                    .get_transaction(*tx_hash)
+                    .await
+                    .ok()
+                    .flatten()
+                    .is_some();
+                if found {
+                    break;
+                }
+            }
+        }
+        Ok(())
+    }
+
     /// Build a batch of public transactions to submit to the node.
+    ///
+    /// Must be called after `claim_vault_funds`, which sets each account's nonce to 1.
     pub fn build_public_txs(&self) -> Vec<PublicTransaction> {
         // Create valid public transactions
         let program = Program::authenticated_transfer_program();
@@ -78,7 +135,7 @@ impl TpsTestManager {
                 let message = putx::Message::try_new(
                     program.id(),
                     [pair[0].1, pair[1].1].to_vec(),
-                    [Nonce(0_u128)].to_vec(),
+                    [Nonce(1_u128)].to_vec(),
                     authenticated_transfer_core::Instruction::Transfer { amount },
                 )
                 .unwrap();
@@ -127,6 +184,12 @@ pub async fn tps_test() -> Result<()> {
         .build()
         .await?;
 
+    // Genesis funds vault PDAs, not accounts directly. Claim into accounts before measuring.
+    tps_test
+        .claim_vault_funds(ctx.sequencer_client())
+        .await
+        .context("Failed to claim vault funds for TPS accounts")?;
+
     let target_time = tps_test.target_time();
     info!(
         "TPS test begin. Target time is {target_time:?} for {num_transactions} transactions ({target_tps} TPS)"
diff --git a/test_fixtures/src/lib.rs b/test_fixtures/src/lib.rs
index 2c9dfb3a..e9000909 100644
--- a/test_fixtures/src/lib.rs
+++ b/test_fixtures/src/lib.rs
@@ -326,12 +326,21 @@ impl TestContextBuilder {
 
         let initial_public_accounts = config::default_public_accounts_for_wallet();
         let initial_private_accounts = config::default_private_accounts_for_wallet();
+        // Wallet genesis must always be present so that setup_public/private_accounts_with_initial_supply
+        // can claim from the vault PDAs. When a test supplies custom genesis, merge rather than replace.
+        let wallet_genesis =
+            config::genesis_from_accounts(&initial_public_accounts, &initial_private_accounts);
+        let genesis = match genesis_transactions {
+            Some(mut custom) => {
+                custom.extend(wallet_genesis);
+                custom
+            }
+            None => wallet_genesis,
+        };
         let (sequencer_handle, temp_sequencer_dir) = setup_sequencer(
             sequencer_partial_config.unwrap_or_default(),
             bedrock_addr,
-            genesis_transactions.unwrap_or_else(|| {
-                config::genesis_from_accounts(&initial_public_accounts, &initial_private_accounts)
-            }),
+            genesis,
         )
         .await
         .context("Failed to setup Sequencer")?;