From c553c0cf20a279768ac5714b1864884f52136613 Mon Sep 17 00:00:00 2001
From: Sam Kleinman <garen@tychoish.com>
Date: Wed, 17 Mar 2021 10:39:58 -0400
Subject: [PATCH] cleanup consensus test (#6246)

---
 consensus/byzantine_test.go | 118 +++++++++++++++++++-----------------
 1 file changed, 63 insertions(+), 55 deletions(-)

diff --git a/consensus/byzantine_test.go b/consensus/byzantine_test.go
index 7fed81395..f0b65ee57 100644
--- a/consensus/byzantine_test.go
+++ b/consensus/byzantine_test.go
@@ -40,57 +40,59 @@ func TestByzantinePrevoteEquivocation(t *testing.T) {
 	states := make([]*State, nValidators)
 
 	for i := 0; i < nValidators; i++ {
-		logger := consensusLogger().With("test", "byzantine", "validator", i)
-		stateDB := dbm.NewMemDB() // each state needs its own db
-		stateStore := sm.NewStore(stateDB)
-		state, _ := stateStore.LoadFromDBOrGenesisDoc(genDoc)
-
-		thisConfig := ResetConfig(fmt.Sprintf("%s_%d", testName, i))
-		defer os.RemoveAll(thisConfig.RootDir)
-
-		ensureDir(path.Dir(thisConfig.Consensus.WalFile()), 0700) // dir for wal
-		app := appFunc()
-		vals := types.TM2PB.ValidatorUpdates(state.Validators)
-		app.InitChain(abci.RequestInitChain{Validators: vals})
-
-		blockDB := dbm.NewMemDB()
-		blockStore := store.NewBlockStore(blockDB)
-
-		// one for mempool, one for consensus
-		mtx := new(tmsync.Mutex)
-		proxyAppConnMem := abcicli.NewLocalClient(mtx, app)
-		proxyAppConnCon := abcicli.NewLocalClient(mtx, app)
-
-		// Make Mempool
-		mempool := mempl.NewCListMempool(thisConfig.Mempool, proxyAppConnMem, 0)
-		mempool.SetLogger(log.TestingLogger().With("module", "mempool"))
-		if thisConfig.Consensus.WaitForTxs() {
-			mempool.EnableTxsAvailable()
-		}
+		func() {
+			logger := consensusLogger().With("test", "byzantine", "validator", i)
+			stateDB := dbm.NewMemDB() // each state needs its own db
+			stateStore := sm.NewStore(stateDB)
+			state, _ := stateStore.LoadFromDBOrGenesisDoc(genDoc)
+
+			thisConfig := ResetConfig(fmt.Sprintf("%s_%d", testName, i))
+			defer os.RemoveAll(thisConfig.RootDir)
+
+			ensureDir(path.Dir(thisConfig.Consensus.WalFile()), 0700) // dir for wal
+			app := appFunc()
+			vals := types.TM2PB.ValidatorUpdates(state.Validators)
+			app.InitChain(abci.RequestInitChain{Validators: vals})
+
+			blockDB := dbm.NewMemDB()
+			blockStore := store.NewBlockStore(blockDB)
+
+			// one for mempool, one for consensus
+			mtx := new(tmsync.Mutex)
+			proxyAppConnMem := abcicli.NewLocalClient(mtx, app)
+			proxyAppConnCon := abcicli.NewLocalClient(mtx, app)
+
+			// Make Mempool
+			mempool := mempl.NewCListMempool(thisConfig.Mempool, proxyAppConnMem, 0)
+			mempool.SetLogger(log.TestingLogger().With("module", "mempool"))
+			if thisConfig.Consensus.WaitForTxs() {
+				mempool.EnableTxsAvailable()
+			}
 
-		// Make a full instance of the evidence pool
-		evidenceDB := dbm.NewMemDB()
-		evpool, err := evidence.NewPool(logger.With("module", "evidence"), evidenceDB, stateStore, blockStore)
-		require.NoError(t, err)
-
-		// Make State
-		blockExec := sm.NewBlockExecutor(stateStore, log.TestingLogger(), proxyAppConnCon, mempool, evpool)
-		cs := NewState(thisConfig.Consensus, state, blockExec, blockStore, mempool, evpool)
-		cs.SetLogger(cs.Logger)
-		// set private validator
-		pv := privVals[i]
-		cs.SetPrivValidator(pv)
-
-		eventBus := types.NewEventBus()
-		eventBus.SetLogger(log.TestingLogger().With("module", "events"))
-		err = eventBus.Start()
-		require.NoError(t, err)
-		cs.SetEventBus(eventBus)
-
-		cs.SetTimeoutTicker(tickerFunc())
-		cs.SetLogger(logger)
-
-		states[i] = cs
+			// Make a full instance of the evidence pool
+			evidenceDB := dbm.NewMemDB()
+			evpool, err := evidence.NewPool(logger.With("module", "evidence"), evidenceDB, stateStore, blockStore)
+			require.NoError(t, err)
+
+			// Make State
+			blockExec := sm.NewBlockExecutor(stateStore, log.TestingLogger(), proxyAppConnCon, mempool, evpool)
+			cs := NewState(thisConfig.Consensus, state, blockExec, blockStore, mempool, evpool)
+			cs.SetLogger(cs.Logger)
+			// set private validator
+			pv := privVals[i]
+			cs.SetPrivValidator(pv)
+
+			eventBus := types.NewEventBus()
+			eventBus.SetLogger(log.TestingLogger().With("module", "events"))
+			err = eventBus.Start()
+			require.NoError(t, err)
+			cs.SetEventBus(eventBus)
+
+			cs.SetTimeoutTicker(tickerFunc())
+			cs.SetLogger(logger)
+
+			states[i] = cs
+		}()
 	}
 
 	rts := setup(t, nValidators, states, 100) // buffer must be large enough to not deadlock
@@ -233,11 +235,17 @@ func TestByzantinePrevoteEquivocation(t *testing.T) {
 
 		go func(j int, s types.Subscription) {
 			defer wg.Done()
-
-			for msg := range s.Out() {
-				block := msg.Data().(types.EventDataNewBlock).Block
-				if len(block.Evidence.Evidence) != 0 {
-					evidenceFromEachValidator[j] = block.Evidence.Evidence[0]
+			for {
+				select {
+				case msg := <-s.Out():
+					require.NotNil(t, msg)
+					block := msg.Data().(types.EventDataNewBlock).Block
+					if len(block.Evidence.Evidence) != 0 {
+						evidenceFromEachValidator[j] = block.Evidence.Evidence[0]
+						return
+					}
+				case <-s.Canceled():
+					require.Fail(t, "subscription failed for %d", j)
 					return
 				}
 			}