From 2b873d81230796ef0b4647b5251e322f2252460d Mon Sep 17 00:00:00 2001
From: William Banfield <wbanfield@gmail.com>
Date: Mon, 22 Nov 2021 15:11:54 -0500
Subject: [PATCH] user arg struct

---
 internal/consensus/common_test.go  | 40 +++++++++++++++++-------------
 internal/consensus/mempool_test.go | 20 +++++++++++----
 2 files changed, 38 insertions(+), 22 deletions(-)

diff --git a/internal/consensus/common_test.go b/internal/consensus/common_test.go
index 285a6cd8f..8b0887b33 100644
--- a/internal/consensus/common_test.go
+++ b/internal/consensus/common_test.go
@@ -466,7 +466,9 @@ func loadPrivValidator(t *testing.T, config *cfg.Config) *privval.FilePV {
 
 func makeState(config *cfg.Config, nValidators int) (*State, []*validatorStub) {
 	// Get State
-	state, privVals := makeGenesisState(config, nValidators, 10, types.DefaultConsensusParams())
+	state, privVals := makeGenesisState(config, genesisStateArgs{
+		Validators: nValidators,
+	})
 
 	vss := make([]*validatorStub, nValidators)
 
@@ -850,24 +852,28 @@ func randConsensusNetWithPeers(
 	}
 }
 
-func makeGenesisState(
-	config *cfg.Config,
-	numValidators int,
-	votingPower int64,
-	consensusParams *types.ConsensusParams) (sm.State, []types.PrivValidator) {
-	return makeGenesisStateWithGenesisTime(config, numValidators, votingPower, consensusParams, time.Now())
+type genesisStateArgs struct {
+	Validators int
+	Power      int64
+	Params     *types.ConsensusParams
+	Time       time.Time
 }
 
-func makeGenesisStateWithGenesisTime(
-	config *cfg.Config,
-	numValidators int,
-	votingPower int64,
-	consensusParams *types.ConsensusParams,
-	genesisTime time.Time,
-) (sm.State, []types.PrivValidator) {
-
-	valSet, privValidators := factory.ValidatorSet(numValidators, votingPower)
-	genDoc := factory.GenesisDoc(config, genesisTime, valSet.Validators, consensusParams)
+func makeGenesisState(config *cfg.Config, args genesisStateArgs) (sm.State, []types.PrivValidator) {
+	if args.Power == 0 {
+		args.Power = 1
+	}
+	if args.Validators == 0 {
+		args.Power = 4
+	}
+	valSet, privValidators := factory.ValidatorSet(args.Validators, args.Power)
+	if args.Params == nil {
+		args.Params = types.DefaultConsensusParams()
+	}
+	if args.Time.IsZero() {
+		args.Time = time.Now()
+	}
+	genDoc := factory.GenesisDoc(config, args.Time, valSet.Validators, args.Params)
 	s0, _ := sm.MakeGenesisState(genDoc)
 	return s0, privValidators
 }
diff --git a/internal/consensus/mempool_test.go b/internal/consensus/mempool_test.go
index f9a119536..7302ad9b6 100644
--- a/internal/consensus/mempool_test.go
+++ b/internal/consensus/mempool_test.go
@@ -33,7 +33,9 @@ func TestMempoolNoProgressUntilTxsAvailable(t *testing.T) {
 	t.Cleanup(func() { _ = os.RemoveAll(config.RootDir) })
 
 	config.Consensus.CreateEmptyBlocks = false
-	state, privVals := makeGenesisState(baseConfig, 1, 10, nil)
+	state, privVals := makeGenesisState(baseConfig, genesisStateArgs{
+		Validators: 1,
+		Power:      10})
 	cs := newStateWithConfig(config, state, privVals[0], NewCounterApplication())
 	assertMempool(cs.txNotifier).EnableTxsAvailable()
 	height, round := cs.Height, cs.Round
@@ -55,7 +57,9 @@ func TestMempoolProgressAfterCreateEmptyBlocksInterval(t *testing.T) {
 	t.Cleanup(func() { _ = os.RemoveAll(config.RootDir) })
 
 	config.Consensus.CreateEmptyBlocksInterval = ensureTimeout
-	state, privVals := makeGenesisState(baseConfig, 1, 10, nil)
+	state, privVals := makeGenesisState(baseConfig, genesisStateArgs{
+		Validators: 1,
+		Power:      10})
 	cs := newStateWithConfig(config, state, privVals[0], NewCounterApplication())
 
 	assertMempool(cs.txNotifier).EnableTxsAvailable()
@@ -75,7 +79,9 @@ func TestMempoolProgressInHigherRound(t *testing.T) {
 	t.Cleanup(func() { _ = os.RemoveAll(config.RootDir) })
 
 	config.Consensus.CreateEmptyBlocks = false
-	state, privVals := makeGenesisState(baseConfig, 1, 10, nil)
+	state, privVals := makeGenesisState(baseConfig, genesisStateArgs{
+		Validators: 1,
+		Power:      10})
 	cs := newStateWithConfig(config, state, privVals[0], NewCounterApplication())
 	assertMempool(cs.txNotifier).EnableTxsAvailable()
 	height, round := cs.Height, cs.Round
@@ -123,7 +129,9 @@ func deliverTxsRange(cs *State, start, end int) {
 func TestMempoolTxConcurrentWithCommit(t *testing.T) {
 	config := configSetup(t)
 
-	state, privVals := makeGenesisState(config, 1, 10, nil)
+	state, privVals := makeGenesisState(config, genesisStateArgs{
+		Validators: 1,
+		Power:      10})
 	stateStore := sm.NewStore(dbm.NewMemDB())
 	blockStore := store.NewBlockStore(dbm.NewMemDB())
 	cs := newStateWithConfigAndBlockStore(config, state, privVals[0], NewCounterApplication(), blockStore)
@@ -149,7 +157,9 @@ func TestMempoolTxConcurrentWithCommit(t *testing.T) {
 func TestMempoolRmBadTx(t *testing.T) {
 	config := configSetup(t)
 
-	state, privVals := makeGenesisState(config, 1, 10, nil)
+	state, privVals := makeGenesisState(config, genesisStateArgs{
+		Validators: 1,
+		Power:      10})
 	app := NewCounterApplication()
 	stateStore := sm.NewStore(dbm.NewMemDB())
 	blockStore := store.NewBlockStore(dbm.NewMemDB())