feat: add the ability to specifiy the txsim master account (backport #2292)

app/test/square_size_test.go

@@ -136,6 +136,7 @@ func (s *SquareSizeIntegrationTest) fillBlocks(blobSize, blobsPerPFB, pfbsPerBlo
 		[]string{s.rpcAddr},
 		[]string{s.grpcAddr},
 		s.cctx.Keyring,
+		"",
 		rand.Int63(),
 		time.Second,
 		seqs...,

test/cmd/txsim/cli.go

@@ -22,22 +22,23 @@ import (
 
 // A set of environment variables that can be used instead of flags
 const (
-	TxsimGRPC     = "TXSIM_GRPC"
-	TxsimRPC      = "TXSIM_RPC"
-	TxsimSeed     = "TXSIM_SEED"
-	TxsimPoll     = "TXSIM_POLL"
-	TxsimKeypath  = "TXSIM_KEYPATH"
-	TxsimMnemonic = "TXSIM_MNEMONIC"
+	TxsimGRPC          = "TXSIM_GRPC"
+	TxsimRPC           = "TXSIM_RPC"
+	TxsimSeed          = "TXSIM_SEED"
+	TxsimPoll          = "TXSIM_POLL"
+	TxsimKeypath       = "TXSIM_KEYPATH"
+	TxsimMasterAccName = "TXSIM_MASTER_ACC_NAME"
+	TxsimMnemonic      = "TXSIM_MNEMONIC"
 )
 
 // Values for all flags
 var (
-	keyPath, keyMnemonic, rpcEndpoints, grpcEndpoints string
-	blobSizes, blobAmounts                            string
-	seed                                              int64
-	pollTime                                          time.Duration
-	send, sendIterations, sendAmount                  int
-	stake, stakeValue, blob                           int
+	keyPath, masterAccName, keyMnemonic, rpcEndpoints, grpcEndpoints string
+	blobSizes, blobAmounts                                           string
+	seed                                                             int64
+	pollTime                                                         time.Duration
+	send, sendIterations, sendAmount                                 int
+	stake, stakeValue, blob                                          int
 )
 
 func main() {
@@ -101,6 +102,10 @@ well funded account that can act as the master account. The command runs until a
 				}
 			}
 
+			if masterAccName == "" {
+				masterAccName = os.Getenv(TxsimMasterAccName)
+			}
+
 			if stake == 0 && send == 0 && blob == 0 {
 				return errors.New("no sequences specified. Use --stake, --send or --blob")
 			}
@@ -154,6 +159,7 @@ well funded account that can act as the master account. The command runs until a
 				strings.Split(rpcEndpoints, ","),
 				strings.Split(grpcEndpoints, ","),
 				keys,
+				masterAccName,
 				seed,
 				pollTime,
 				sequences...,
@@ -172,6 +178,7 @@ well funded account that can act as the master account. The command runs until a
 func flags() *flag.FlagSet {
 	flags := &flag.FlagSet{}
 	flags.StringVar(&keyPath, "key-path", "", "path to the keyring")
+	flags.StringVar(&masterAccName, "master", "", "the account name of the master account. Leaving empty will result in using the account with the most funds.")
 	flags.StringVar(&keyMnemonic, "key-mnemonic", "", "space separated mnemonic for the keyring. The hdpath used is an empty string")
 	flags.StringVar(&rpcEndpoints, "rpc-endpoints", "", "comma separated list of rpc endpoints")
 	flags.StringVar(&grpcEndpoints, "grpc-endpoints", "", "comma separated list of grpc endpoints")

test/txsim/account.go

@@ -43,7 +43,7 @@ type Account struct {
 	Balance       int64
 }
 
-func NewAccountManager(ctx context.Context, keys keyring.Keyring, txClient *TxClient, queryClient *QueryClient) (*AccountManager, error) {
+func NewAccountManager(ctx context.Context, keys keyring.Keyring, masterAccName string, txClient *TxClient, queryClient *QueryClient) (*AccountManager, error) {
 	records, err := keys.List()
 	if err != nil {
 		return nil, err
@@ -61,8 +61,14 @@ func NewAccountManager(ctx context.Context, keys keyring.Keyring, txClient *TxCl
 		query:    queryClient,
 	}
 
-	if err := am.setupMasterAccount(ctx); err != nil {
-		return nil, err
+	if masterAccName == "" {
+		if err := am.setupDefaultMasterAccount(ctx); err != nil {
+			return nil, err
+		}
+	} else {
+		if err := am.setupSpecifiedMasterAccount(ctx, masterAccName); err != nil {
+			return nil, err
+		}
 	}
 
 	log.Info().
@@ -74,10 +80,10 @@ func NewAccountManager(ctx context.Context, keys keyring.Keyring, txClient *TxCl
 	return am, nil
 }
 
-// setupMasterAccount loops through all accounts in the keyring and picks out the one with
+// setupDefaultMasterAccount loops through all accounts in the keyring and picks out the one with
 // the highest balance as the master account. Accounts that don't yet exist on chain are
 // ignored.
-func (am *AccountManager) setupMasterAccount(ctx context.Context) error {
+func (am *AccountManager) setupDefaultMasterAccount(ctx context.Context) error {
 	am.mtx.Lock()
 	defer am.mtx.Unlock()
 
@@ -127,6 +133,44 @@ func (am *AccountManager) setupMasterAccount(ctx context.Context) error {
 	return nil
 }
 
+func (am *AccountManager) setupSpecifiedMasterAccount(ctx context.Context, masterAccName string) error {
+	masterRecord, err := am.keys.Key(masterAccName)
+	if err != nil {
+		return fmt.Errorf("error getting specified master account %s: %w", masterAccName, err)
+	}
+
+	masterAddress, err := masterRecord.GetAddress()
+	if err != nil {
+		return fmt.Errorf("error getting address for account %s: %w", masterAccName, err)
+	}
+
+	// search for the account on chain
+	masterBalance, err := am.getBalance(ctx, masterAddress)
+	if err != nil {
+		return fmt.Errorf("error getting specified master account %s balance: %w", masterAccName, err)
+	}
+
+	accountNumber, sequence, err := am.getAccountDetails(ctx, masterAddress)
+	if err != nil {
+		return fmt.Errorf("error getting account details for account %s: %w", masterRecord.Name, err)
+	}
+
+	pk, err := masterRecord.GetPubKey()
+	if err != nil {
+		return fmt.Errorf("error getting public key for account %s: %w", masterRecord.Name, err)
+	}
+
+	am.masterAccount = &Account{
+		Address:       masterAddress,
+		PubKey:        pk,
+		Sequence:      sequence,
+		AccountNumber: accountNumber,
+		Balance:       masterBalance,
+	}
+
+	return nil
+}
+
 // AllocateAccounts is used by sequences to specify the number of accounts
 // and the balance of each of those accounts. Not concurrently safe.
 func (am *AccountManager) AllocateAccounts(n, balance int) []types.AccAddress {
@@ -221,8 +265,8 @@ func (am *AccountManager) Submit(ctx context.Context, op Operation) error {
 	return nil
 }
 
-// Generate the pending accounts by sending the adequate funds and setting up the feegrant permissions.
-// This operation is not concurrently safe.
+// Generate the pending accounts by sending the adequate funds. This operation
+// is not concurrently safe.
 func (am *AccountManager) GenerateAccounts(ctx context.Context) error {
 	if len(am.pending) == 0 {
 		return nil
@@ -232,7 +276,7 @@ func (am *AccountManager) GenerateAccounts(ctx context.Context) error {
 	// batch together all the messages needed to create all the accounts
 	for _, acc := range am.pending {
 		if am.masterAccount.Balance < acc.Balance {
-			return fmt.Errorf("master account has insufficient funds")
+			return fmt.Errorf("master account has insufficient funds needed: %v", acc.Balance)
 		}
 
 		bankMsg := bank.NewMsgSend(am.masterAccount.Address, acc.Address, types.NewCoins(types.NewInt64Coin(app.BondDenom, acc.Balance)))

test/txsim/run.go

@@ -27,6 +27,7 @@ func Run(
 	ctx context.Context,
 	rpcEndpoints, grpcEndpoints []string,
 	keys keyring.Keyring,
+	masterAccName string,
 	seed int64,
 	pollTime time.Duration,
 	sequences ...Sequence,
@@ -45,7 +46,7 @@ func Run(
 	defer queryClient.Close()
 
 	// Create the account manager to handle account transactions.
-	manager, err := NewAccountManager(ctx, keys, txClient, queryClient)
+	manager, err := NewAccountManager(ctx, keys, masterAccName, txClient, queryClient)
 	if err != nil {
 		return err
 	}

test/txsim/run_test.go

@@ -87,6 +87,7 @@ func TestTxSimulator(t *testing.T) {
 				[]string{rpcAddr},
 				[]string{grpcAddr},
 				keyring,
+				"",
 				9001,
 				time.Second,
 				tc.sequences...,