work in progress refactor of the ensure functions consensus tests

internal/consensus/common_test.go

@@ -544,242 +544,247 @@ func makeState(ctx context.Context, cfg *config.Config, logger log.Logger, nVali
 
 //-------------------------------------------------------------------------------
 
-func ensureNoNewEvent(t *testing.T, ch <-chan tmpubsub.Message, timeout time.Duration,
-	errorMessage string) {
-	t.Helper()
+type eventChecker struct {
+	ctx     context.Context
+	t       *testing.T
+	ch      <-chan tmpubsub.Message
+	timeout time.Duration
+}
+
+func (ec eventChecker) ensureMatchingProposal(height int64, round int32, propID types.BlockID) {
+	ec.t.Helper()
 	select {
-	case <-time.After(timeout):
-		break
-	case <-ch:
-		t.Fatalf("unexpected event: %s", errorMessage)
+	case <-time.After(ensureTimeout):
+		ec.t.Fatalf("Timeout expired while waiting for NewProposal event")
+	case msg := <-ec.ch:
+		proposalEvent, ok := msg.Data().(types.EventDataCompleteProposal)
+		if !ok {
+			ec.t.Fatalf("expected a EventDataCompleteProposal, got %T. Wrong subscription channel?",
+				msg.Data())
+		}
+		if proposalEvent.Height != height {
+			ec.t.Fatalf("expected height %v, got %v", height, proposalEvent.Height)
+		}
+		if proposalEvent.Round != round {
+			ec.t.Fatalf("expected round %v, got %v", round, proposalEvent.Round)
+		}
+		if !proposalEvent.BlockID.Equals(propID) {
+			ec.t.Fatalf("Proposed block does not match expected block (%v != %v)", proposalEvent.BlockID, propID)
+		}
 	}
 }
 
-func ensureNoNewEventOnChannel(t *testing.T, ch <-chan tmpubsub.Message) {
-	t.Helper()
-	ensureNoNewEvent(
-		t,
-		ch,
-		ensureTimeout,
-		"We should be stuck waiting, not receiving new event on the channel")
-}
-
-func ensureNoNewRoundStep(t *testing.T, stepCh <-chan tmpubsub.Message) {
-	t.Helper()
-	ensureNoNewEvent(
-		t,
-		stepCh,
-		ensureTimeout,
-		"We should be stuck waiting, not receiving NewRoundStep event")
-}
-
-func ensureNoNewTimeout(t *testing.T, stepCh <-chan tmpubsub.Message, timeout int64) {
-	t.Helper()
-	timeoutDuration := time.Duration(timeout*10) * time.Nanosecond
-	ensureNoNewEvent(
-		t,
-		stepCh,
-		timeoutDuration,
-		"We should be stuck waiting, not receiving NewTimeout event")
-}
-
-func ensureNewEvent(t *testing.T, ch <-chan tmpubsub.Message, height int64, round int32, timeout time.Duration, errorMessage string) { // nolint: lll
-	t.Helper()
+func (ec eventChecker) ensureNewProposal(height int64, round int32) {
+	ec.t.Helper()
+	to := ec.timeout
+	if to == 0 {
+		to = ensureTimeout
+	}
 	select {
-	case <-time.After(timeout):
-		t.Fatalf("timed out waiting for new event: %s", errorMessage)
-	case msg := <-ch:
+	case <-time.After(to):
+		ec.t.Fatalf("Timeout expired while waiting for NewProposal event")
+	case msg := <-ec.ch:
+		proposalEvent, ok := msg.Data().(types.EventDataCompleteProposal)
+		if !ok {
+			ec.t.Fatalf("expected a EventDataCompleteProposal, got %T. Wrong subscription channel?",
+				msg.Data())
+		}
+		if proposalEvent.Height != height {
+			ec.t.Fatalf("expected height %v, got %v", height, proposalEvent.Height)
+		}
+		if proposalEvent.Round != round {
+			ec.t.Fatalf("expected round %v, got %v", round, proposalEvent.Round)
+		}
+	}
+
+}
+func (ec eventChecker) ensureNewRound(height int64, round int32) {
+	ec.t.Helper()
+	to := ec.timeout
+	if to == 0 {
+		to = ensureTimeout
+	}
+	select {
+	case <-time.After(to):
+		ec.t.Fatal("Timeout expired while waiting for NewRound event")
+	case msg := <-ec.ch:
+		newRoundEvent, ok := msg.Data().(types.EventDataNewRound)
+		if !ok {
+			ec.t.Fatalf("expected a EventDataNewRound, got %T. Wrong subscription channel?", msg.Data())
+		}
+		if newRoundEvent.Height != height {
+			ec.t.Fatalf("expected height %v, got %v", height, newRoundEvent.Height)
+		}
+		if newRoundEvent.Round != round {
+			ec.t.Fatalf("expected round %v, got %v", round, newRoundEvent.Round)
+		}
+	}
+}
+
+func (ec eventChecker) ensureNewEvent() {
+	ec.t.Helper()
+	to := ec.timeout
+	if to == 0 {
+		to = ensureTimeout
+	}
+	select {
+	case <-time.After(to):
+		ec.t.Fatal("timeout occured while waiting for event")
+	case <-ec.ch:
+	}
+}
+func (ec eventChecker) ensureNoNewEvent(errorMessage string) {
+	ec.t.Helper()
+	to := ec.timeout
+	if to == 0 {
+		to = ensureTimeout
+	}
+	select {
+	case <-time.After(to):
+		break
+	case <-ec.ch:
+		ec.t.Fatalf("unexpected event: %s", errorMessage)
+	}
+}
+
+func (ec eventChecker) ensureNoNewRoundStep() {
+	ec.t.Helper()
+	ec.ensureNoNewEvent("We should be stuck waiting, not receiving NewRoundStep event")
+}
+
+func (ec eventChecker) ensureNoNewTimeout() {
+	ec.t.Helper()
+	ec.ensureNoNewEvent("We should be stuck waiting, not receiving NewTimeout event")
+}
+
+func (ec eventChecker) ensureNewRoundState(height int64, round int32, errorMessage string) { // nolint: lll
+	ec.t.Helper()
+	to := ec.timeout
+	if to == 0 {
+		to = ensureTimeout
+	}
+	select {
+	case <-time.After(to):
+		ec.t.Fatalf("timed out waiting for new event: %s", errorMessage)
+	case msg := <-ec.ch:
 		roundStateEvent, ok := msg.Data().(types.EventDataRoundState)
 		if !ok {
-			t.Fatalf("expected a EventDataRoundState, got %T. Wrong subscription channel?", msg.Data())
+			ec.t.Fatalf("expected a EventDataRoundState, got %T. Wrong subscription channel?", msg.Data())
 		}
 		if roundStateEvent.Height != height {
-			t.Fatalf("expected height %v, got %v", height, roundStateEvent.Height)
+			ec.t.Fatalf("expected height %v, got %v", height, roundStateEvent.Height)
 		}
 		if roundStateEvent.Round != round {
-			t.Fatalf("expected round %v, got %v", round, roundStateEvent.Round)
+			ec.t.Fatalf("expected round %v, got %v", round, roundStateEvent.Round)
 		}
 		// TODO: We could check also for a step at this point!
 	}
 }
 
-func ensureNewRound(t *testing.T, roundCh <-chan tmpubsub.Message, height int64, round int32) {
-	t.Helper()
-	select {
-	case <-time.After(ensureTimeout):
-		t.Fatal("Timeout expired while waiting for NewRound event")
-	case msg := <-roundCh:
-		newRoundEvent, ok := msg.Data().(types.EventDataNewRound)
-		if !ok {
-			t.Fatalf("expected a EventDataNewRound, got %T. Wrong subscription channel?", msg.Data())
-		}
-		if newRoundEvent.Height != height {
-			t.Fatalf("expected height %v, got %v", height, newRoundEvent.Height)
-		}
-		if newRoundEvent.Round != round {
-			t.Fatalf("expected round %v, got %v", round, newRoundEvent.Round)
-		}
+func (ec eventChecker) ensureNewTimeout(height int64, round int32) {
+	ec.t.Helper()
+	ec.ensureNewRoundState(height, round, "Timeout expired while waiting for NewTimeout event")
+}
+
+func (ec eventChecker) ensureNewValidBlock(height int64, round int32) {
+	ec.t.Helper()
+	ec.ensureNewRoundState(height, round, "Timeout expired while waiting for NewValidBlock event")
+}
+
+func (ec eventChecker) ensureNewBlock(height int64) {
+	ec.t.Helper()
+	to := ec.timeout
+	if to == 0 {
+		to = ensureTimeout
 	}
-}
-
-func ensureNewTimeout(t *testing.T, timeoutCh <-chan tmpubsub.Message, height int64, round int32, timeout int64) {
-	t.Helper()
-	timeoutDuration := time.Duration(timeout*10) * time.Nanosecond
-	ensureNewEvent(t, timeoutCh, height, round, timeoutDuration,
-		"Timeout expired while waiting for NewTimeout event")
-}
-
-func ensureNewProposal(t *testing.T, proposalCh <-chan tmpubsub.Message, height int64, round int32) {
-	t.Helper()
 	select {
-	case <-time.After(ensureTimeout):
-		t.Fatalf("Timeout expired while waiting for NewProposal event")
-	case msg := <-proposalCh:
-		proposalEvent, ok := msg.Data().(types.EventDataCompleteProposal)
-		if !ok {
-			t.Fatalf("expected a EventDataCompleteProposal, got %T. Wrong subscription channel?",
-				msg.Data())
-		}
-		if proposalEvent.Height != height {
-			t.Fatalf("expected height %v, got %v", height, proposalEvent.Height)
-		}
-		if proposalEvent.Round != round {
-			t.Fatalf("expected round %v, got %v", round, proposalEvent.Round)
-		}
-	}
-}
-
-func ensureNewValidBlock(t *testing.T, validBlockCh <-chan tmpubsub.Message, height int64, round int32) {
-	t.Helper()
-	ensureNewEvent(t, validBlockCh, height, round, ensureTimeout,
-		"Timeout expired while waiting for NewValidBlock event")
-}
-
-func ensureNewBlock(t *testing.T, blockCh <-chan tmpubsub.Message, height int64) {
-	t.Helper()
-	select {
-	case <-time.After(ensureTimeout):
-		t.Fatalf("Timeout expired while waiting for NewBlock event")
-	case msg := <-blockCh:
+	case <-time.After(to):
+		ec.t.Fatalf("Timeout expired while waiting for NewBlock event")
+	case msg := <-ec.ch:
 		blockEvent, ok := msg.Data().(types.EventDataNewBlock)
 		if !ok {
-			t.Fatalf("expected a EventDataNewBlock, got %T. Wrong subscription channel?",
+			ec.t.Fatalf("expected a EventDataNewBlock, got %T. Wrong subscription channel?",
 				msg.Data())
 		}
 		if blockEvent.Block.Height != height {
-			t.Fatalf("expected height %v, got %v", height, blockEvent.Block.Height)
+			ec.t.Fatalf("expected height %v, got %v", height, blockEvent.Block.Height)
 		}
 	}
 }
 
-func ensureNewBlockHeader(t *testing.T, blockCh <-chan tmpubsub.Message, height int64, blockHash tmbytes.HexBytes) {
-	t.Helper()
+func (ec eventChecker) ensureNewBlockHeader(height int64, blockHash tmbytes.HexBytes) {
+	ec.t.Helper()
+	to := ec.timeout
+	if to == 0 {
+		to = ensureTimeout
+
+	}
 	select {
-	case <-time.After(ensureTimeout):
-		t.Fatalf("Timeout expired while waiting for NewBlockHeader event")
-	case msg := <-blockCh:
+	case <-time.After(to):
+		ec.t.Fatalf("Timeout expired while waiting for NewBlockHeader event")
+	case msg := <-ec.ch:
 		blockHeaderEvent, ok := msg.Data().(types.EventDataNewBlockHeader)
 		if !ok {
-			t.Fatalf("expected a EventDataNewBlockHeader, got %T. Wrong subscription channel?",
+			ec.t.Fatalf("expected a EventDataNewBlockHeader, got %T. Wrong subscription channel?",
 				msg.Data())
 		}
 		if blockHeaderEvent.Header.Height != height {
-			t.Fatalf("expected height %v, got %v", height, blockHeaderEvent.Header.Height)
+			ec.t.Fatalf("expected height %v, got %v", height, blockHeaderEvent.Header.Height)
 		}
 		if !bytes.Equal(blockHeaderEvent.Header.Hash(), blockHash) {
-			t.Fatalf("expected header %X, got %X", blockHash, blockHeaderEvent.Header.Hash())
+			ec.t.Fatalf("expected header %X, got %X", blockHash, blockHeaderEvent.Header.Hash())
 		}
 	}
 }
 
-func ensureLock(t *testing.T, lockCh <-chan tmpubsub.Message, height int64, round int32) {
-	t.Helper()
-	ensureNewEvent(t, lockCh, height, round, ensureTimeout,
-		"Timeout expired while waiting for LockValue event")
+func (ec eventChecker) ensureLock(height int64, round int32) {
+	ec.t.Helper()
+	ec.ensureNewRoundState(height, round, "Timeout expired while waiting for LockValue event")
 }
 
-func ensureRelock(t *testing.T, relockCh <-chan tmpubsub.Message, height int64, round int32) {
-	t.Helper()
-	ensureNewEvent(t, relockCh, height, round, ensureTimeout,
-		"Timeout expired while waiting for RelockValue event")
+func (ec eventChecker) ensureRelock(height int64, round int32) {
+	ec.t.Helper()
+	ec.ensureNewRoundState(height, round, "Timeout expired while waiting for RelockValue event")
 }
 
-func ensureProposal(t *testing.T, proposalCh <-chan tmpubsub.Message, height int64, round int32, propID types.BlockID) {
-	t.Helper()
-	select {
-	case <-time.After(ensureTimeout):
-		t.Fatalf("Timeout expired while waiting for NewProposal event")
-	case msg := <-proposalCh:
-		proposalEvent, ok := msg.Data().(types.EventDataCompleteProposal)
-		if !ok {
-			t.Fatalf("expected a EventDataCompleteProposal, got %T. Wrong subscription channel?",
-				msg.Data())
-		}
-		if proposalEvent.Height != height {
-			t.Fatalf("expected height %v, got %v", height, proposalEvent.Height)
-		}
-		if proposalEvent.Round != round {
-			t.Fatalf("expected round %v, got %v", round, proposalEvent.Round)
-		}
-		if !proposalEvent.BlockID.Equals(propID) {
-			t.Fatalf("Proposed block does not match expected block (%v != %v)", proposalEvent.BlockID, propID)
-		}
+func (ec eventChecker) ensurePrecommit(height int64, round int32) {
+	ec.t.Helper()
+	ec.ensureVote(height, round, tmproto.PrecommitType)
+}
+
+func (ec eventChecker) ensurePrevote(height int64, round int32) {
+	ec.t.Helper()
+	ec.ensureVote(height, round, tmproto.PrevoteType)
+}
+
+func (ec eventChecker) ensureVote(height int64, round int32, voteType tmproto.SignedMsgType) {
+	ec.t.Helper()
+	to := ec.timeout
+	if to == 0 {
+		to = ensureTimeout
 	}
-}
-
-func ensurePrecommit(t *testing.T, voteCh <-chan tmpubsub.Message, height int64, round int32) {
-	t.Helper()
-	ensureVote(t, voteCh, height, round, tmproto.PrecommitType)
-}
-
-func ensurePrevote(t *testing.T, voteCh <-chan tmpubsub.Message, height int64, round int32) {
-	t.Helper()
-	ensureVote(t, voteCh, height, round, tmproto.PrevoteType)
-}
-
-func ensureVote(t *testing.T, voteCh <-chan tmpubsub.Message, height int64, round int32,
-	voteType tmproto.SignedMsgType) {
-	t.Helper()
 	select {
-	case <-time.After(ensureTimeout):
-		t.Fatalf("Timeout expired while waiting for NewVote event")
-	case msg := <-voteCh:
+	case <-time.After(to):
+		ec.t.Fatalf("Timeout expired while waiting for NewVote event")
+	case msg := <-ec.ch:
 		voteEvent, ok := msg.Data().(types.EventDataVote)
 		if !ok {
-			t.Fatalf("expected a EventDataVote, got %T. Wrong subscription channel?",
+			ec.t.Fatalf("expected a EventDataVote, got %T. Wrong subscription channel?",
 				msg.Data())
 		}
 		vote := voteEvent.Vote
 		if vote.Height != height {
-			t.Fatalf("expected height %v, got %v", height, vote.Height)
+			ec.t.Fatalf("expected height %v, got %v", height, vote.Height)
 		}
 		if vote.Round != round {
-			t.Fatalf("expected round %v, got %v", round, vote.Round)
+			ec.t.Fatalf("expected round %v, got %v", round, vote.Round)
 		}
 		if vote.Type != voteType {
-			t.Fatalf("expected type %v, got %v", voteType, vote.Type)
+			ec.t.Fatalf("expected type %v, got %v", voteType, vote.Type)
 		}
 	}
 }
 
-func ensurePrecommitTimeout(t *testing.T, ch <-chan tmpubsub.Message) {
-	t.Helper()
-	select {
-	case <-time.After(ensureTimeout):
-		t.Fatalf("Timeout expired while waiting for the Precommit to Timeout")
-	case <-ch:
-	}
-}
-
-func ensureNewEventOnChannel(t *testing.T, ch <-chan tmpubsub.Message) {
-	t.Helper()
-	select {
-	case <-time.After(ensureTimeout):
-		t.Fatalf("Timeout expired while waiting for new activity on the channel")
-	case <-ch:
-	}
-}
-
 //-------------------------------------------------------------------------------
 // consensus nets
 

internal/consensus/mempool_test.go

@@ -10,7 +10,6 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	dbm "github.com/tendermint/tm-db"
 
 	"github.com/tendermint/tendermint/abci/example/code"
 	abci "github.com/tendermint/tendermint/abci/types"
@@ -19,6 +18,7 @@ import (
 	"github.com/tendermint/tendermint/internal/store"
 	"github.com/tendermint/tendermint/libs/log"
 	"github.com/tendermint/tendermint/types"
+	dbm "github.com/tendermint/tm-db"
 )
 
 // for testing
@@ -43,15 +43,19 @@ func TestMempoolNoProgressUntilTxsAvailable(t *testing.T) {
 	cs := newStateWithConfig(ctx, log.TestingLogger(), config, state, privVals[0], NewCounterApplication())
 	assertMempool(cs.txNotifier).EnableTxsAvailable()
 	height, round := cs.Height, cs.Round
-	newBlockCh := subscribe(ctx, t, cs.eventBus, types.EventQueryNewBlock)
+	blockChecker := eventChecker{
+		ctx: ctx,
+		t:   t,
+		ch:  subscribe(ctx, t, cs.eventBus, types.EventQueryNewBlock),
+	}
 	startTestRound(ctx, cs, height, round)
 
-	ensureNewEventOnChannel(t, newBlockCh) // first block gets committed
-	ensureNoNewEventOnChannel(t, newBlockCh)
+	blockChecker.ensureNewBlock(height) // first block gets committed
+	blockChecker.ensureNoNewEvent("unexpected block event")
 	deliverTxsRange(ctx, cs, 0, 1)
-	ensureNewEventOnChannel(t, newBlockCh) // commit txs
-	ensureNewEventOnChannel(t, newBlockCh) // commit updated app hash
-	ensureNoNewEventOnChannel(t, newBlockCh)
+	blockChecker.ensureNewEvent() // commit txs
+	blockChecker.ensureNewEvent() // commit updated app hash
+	blockChecker.ensureNoNewEvent("unexpected block event")
 }
 
 func TestMempoolProgressAfterCreateEmptyBlocksInterval(t *testing.T) {
@@ -71,12 +75,16 @@ func TestMempoolProgressAfterCreateEmptyBlocksInterval(t *testing.T) {
 
 	assertMempool(cs.txNotifier).EnableTxsAvailable()
 
-	newBlockCh := subscribe(ctx, t, cs.eventBus, types.EventQueryNewBlock)
+	blockChecker := eventChecker{
+		ctx: ctx,
+		t:   t,
+		ch:  subscribe(ctx, t, cs.eventBus, types.EventQueryNewBlock),
+	}
 	startTestRound(ctx, cs, cs.Height, cs.Round)
 
-	ensureNewEventOnChannel(t, newBlockCh)   // first block gets committed
-	ensureNoNewEventOnChannel(t, newBlockCh) // then we dont make a block ...
-	ensureNewEventOnChannel(t, newBlockCh)   // until the CreateEmptyBlocksInterval has passed
+	blockChecker.ensureNewEvent()                     // first block gets committed
+	blockChecker.ensureNoNewEvent("unexpected block") // then we dont make a block ...
+	blockChecker.ensureNewEvent()                     // until the CreateEmptyBlocksInterval has passed
 }
 
 func TestMempoolProgressInHigherRound(t *testing.T) {
@@ -95,9 +103,21 @@ func TestMempoolProgressInHigherRound(t *testing.T) {
 	cs := newStateWithConfig(ctx, log.TestingLogger(), config, state, privVals[0], NewCounterApplication())
 	assertMempool(cs.txNotifier).EnableTxsAvailable()
 	height, round := cs.Height, cs.Round
-	newBlockCh := subscribe(ctx, t, cs.eventBus, types.EventQueryNewBlock)
-	newRoundCh := subscribe(ctx, t, cs.eventBus, types.EventQueryNewRound)
-	timeoutCh := subscribe(ctx, t, cs.eventBus, types.EventQueryTimeoutPropose)
+	blockChecker := eventChecker{
+		ctx: ctx,
+		t:   t,
+		ch:  subscribe(ctx, t, cs.eventBus, types.EventQueryNewBlock),
+	}
+	roundChecker := eventChecker{
+		ctx: ctx,
+		t:   t,
+		ch:  subscribe(ctx, t, cs.eventBus, types.EventQueryNewRound),
+	}
+	proposalTimeoutChecker := eventChecker{
+		ctx: ctx,
+		t:   t,
+		ch:  subscribe(ctx, t, cs.eventBus, types.EventQueryTimeoutPropose),
+	}
 	cs.setProposal = func(proposal *types.Proposal) error {
 		if cs.Height == 2 && cs.Round == 0 {
 			// dont set the proposal in round 0 so we timeout and
@@ -109,19 +129,19 @@ func TestMempoolProgressInHigherRound(t *testing.T) {
 	}
 	startTestRound(ctx, cs, height, round)
 
-	ensureNewRound(t, newRoundCh, height, round) // first round at first height
-	ensureNewEventOnChannel(t, newBlockCh)       // first block gets committed
+	roundChecker.ensureNewRound(height, round) // first round at first height
+	blockChecker.ensureNewEvent()              // first block gets committed
 
 	height++ // moving to the next height
 	round = 0
 
-	ensureNewRound(t, newRoundCh, height, round) // first round at next height
-	deliverTxsRange(ctx, cs, 0, 1)               // we deliver txs, but dont set a proposal so we get the next round
-	ensureNewTimeout(t, timeoutCh, height, round, cs.config.TimeoutPropose.Nanoseconds())
+	roundChecker.ensureNewRound(height, round) // first round at next height
+	deliverTxsRange(ctx, cs, 0, 1)             // we deliver txs, but dont set a proposal so we get the next round
+	proposalTimeoutChecker.ensureNewTimeout(height, round)
 
-	round++                                      // moving to the next round
-	ensureNewRound(t, newRoundCh, height, round) // wait for the next round
-	ensureNewEventOnChannel(t, newBlockCh)       // now we can commit the block
+	round++                                    // moving to the next round
+	roundChecker.ensureNewRound(height, round) // wait for the next round
+	blockChecker.ensureNewEvent()              // now we can commit the block
 }
 
 func deliverTxsRange(ctx context.Context, cs *State, start, end int) {

internal/consensus/pbts_test.go

@@ -45,7 +45,7 @@ type pbtsTestHarness struct {
 	chainID string
 
 	// channels for verifying that the observed validator completes certain actions.
-	ensureProposalCh, roundCh, blockCh, ensureVoteCh <-chan tmpubsub.Message
+	proposalChecker, roundChecker, blockChecker, voteChecker eventChecker
 
 	resultCh <-chan heightResult
 
@@ -111,10 +111,10 @@ func newPBTSTestHarness(ctx context.Context, t *testing.T, tc pbtsTestConfigurat
 		validatorClock:        clock,
 		currentHeight:         1,
 		chainID:               cfg.ChainID(),
-		roundCh:               subscribe(ctx, t, cs.eventBus, types.EventQueryNewRound),
-		ensureProposalCh:      subscribe(ctx, t, cs.eventBus, types.EventQueryCompleteProposal),
-		blockCh:               subscribe(ctx, t, cs.eventBus, types.EventQueryNewBlock),
-		ensureVoteCh:          subscribeToVoterBuffered(ctx, t, cs, pubKey.Address()),
+		roundChecker:          eventChecker{ctx: ctx, t: t, ch: subscribe(ctx, t, cs.eventBus, types.EventQueryNewRound)},
+		proposalChecker:       eventChecker{ctx: ctx, t: t, ch: subscribe(ctx, t, cs.eventBus, types.EventQueryCompleteProposal)},
+		blockChecker:          eventChecker{ctx: ctx, t: t, ch: subscribe(ctx, t, cs.eventBus, types.EventQueryNewBlock)},
+		voteChecker:           eventChecker{ctx: ctx, t: t, ch: subscribeToVoterBuffered(ctx, t, cs, pubKey.Address())},
 		resultCh:              resultCh,
 		t:                     t,
 		ctx:                   ctx,
@@ -125,17 +125,17 @@ func (p *pbtsTestHarness) genesisHeight() heightResult {
 	p.validatorClock.On("Now").Return(p.height2ProposedBlockTime).Times(8)
 
 	startTestRound(p.ctx, p.observedState, p.currentHeight, p.currentRound)
-	ensureNewRound(p.t, p.roundCh, p.currentHeight, p.currentRound)
+	p.roundChecker.ensureNewRound(p.currentHeight, p.currentRound)
 	propBlock, partSet := p.observedState.createProposalBlock()
 	bid := types.BlockID{Hash: propBlock.Hash(), PartSetHeader: partSet.Header()}
-	ensureProposal(p.t, p.ensureProposalCh, p.currentHeight, p.currentRound, bid)
-	ensurePrevote(p.t, p.ensureVoteCh, p.currentHeight, p.currentRound)
+	p.proposalChecker.ensureMatchingProposal(p.currentHeight, p.currentRound, bid)
+	p.voteChecker.ensurePrevote(p.currentHeight, p.currentRound)
 	signAddVotes(p.ctx, p.observedState, tmproto.PrevoteType, p.chainID, bid, p.otherValidators...)
 
 	signAddVotes(p.ctx, p.observedState, tmproto.PrecommitType, p.chainID, bid, p.otherValidators...)
-	ensurePrecommit(p.t, p.ensureVoteCh, p.currentHeight, p.currentRound)
+	p.voteChecker.ensurePrecommit(p.currentHeight, p.currentRound)
 
-	ensureNewBlock(p.t, p.blockCh, p.currentHeight)
+	p.blockChecker.ensureNewBlock(p.currentHeight)
 	p.currentHeight++
 	incrementHeight(p.otherValidators...)
 	return <-p.resultCh
@@ -150,7 +150,7 @@ func (p *pbtsTestHarness) height2() heightResult {
 func (p *pbtsTestHarness) nextHeight(proposer types.PrivValidator, deliverTime, proposedTime, nextProposedTime time.Time) heightResult {
 	p.validatorClock.On("Now").Return(nextProposedTime).Times(8)
 
-	ensureNewRound(p.t, p.roundCh, p.currentHeight, p.currentRound)
+	p.roundChecker.ensureNewRound(p.currentHeight, p.currentRound)
 
 	b, _ := p.observedState.createProposalBlock()
 	b.Height = p.currentHeight
@@ -174,13 +174,13 @@ func (p *pbtsTestHarness) nextHeight(proposer types.PrivValidator, deliverTime,
 	if err := p.observedState.SetProposalAndBlock(prop, b, ps, "peerID"); err != nil {
 		p.t.Fatal(err)
 	}
-	ensureProposal(p.t, p.ensureProposalCh, p.currentHeight, 0, bid)
+	p.proposalChecker.ensureMatchingProposal(p.currentHeight, 0, bid)
 
-	ensurePrevote(p.t, p.ensureVoteCh, p.currentHeight, p.currentRound)
+	p.voteChecker.ensurePrevote(p.currentHeight, p.currentRound)
 	signAddVotes(p.ctx, p.observedState, tmproto.PrevoteType, p.chainID, bid, p.otherValidators...)
 
 	signAddVotes(p.ctx, p.observedState, tmproto.PrecommitType, p.chainID, bid, p.otherValidators...)
-	ensurePrecommit(p.t, p.ensureVoteCh, p.currentHeight, p.currentRound)
+	p.voteChecker.ensurePrecommit(p.currentHeight, p.currentRound)
 
 	p.currentHeight++
 	incrementHeight(p.otherValidators...)

internal/consensus/replay_test.go

@@ -348,8 +348,16 @@ func setupSimulator(ctx context.Context, t *testing.T) *simulatorTestSuite {
 
 	partSize := types.BlockPartSizeBytes
 
-	newRoundCh := subscribe(ctx, t, css[0].eventBus, types.EventQueryNewRound)
-	proposalCh := subscribe(ctx, t, css[0].eventBus, types.EventQueryCompleteProposal)
+	roundChecker := eventChecker{
+		ctx: ctx,
+		t:   t,
+		ch:  subscribe(ctx, t, css[0].eventBus, types.EventQueryNewRound),
+	}
+	proposalChecker := eventChecker{
+		ctx: ctx,
+		t:   t,
+		ch:  subscribe(ctx, t, css[0].eventBus, types.EventQueryCompleteProposal),
+	}
 
 	vss := make([]*validatorStub, nPeers)
 	for i := 0; i < nPeers; i++ {
@@ -360,15 +368,15 @@ func setupSimulator(ctx context.Context, t *testing.T) *simulatorTestSuite {
 	// start the machine
 	startTestRound(ctx, css[0], height, round)
 	incrementHeight(vss...)
-	ensureNewRound(t, newRoundCh, height, 0)
-	ensureNewProposal(t, proposalCh, height, round)
+	roundChecker.ensureNewRound(height, 0)
+	proposalChecker.ensureNewProposal(height, round)
 	rs := css[0].GetRoundState()
 
 	signAddVotes(ctx, css[0], tmproto.PrecommitType, sim.Config.ChainID(),
 		types.BlockID{Hash: rs.ProposalBlock.Hash(), PartSetHeader: rs.ProposalBlockParts.Header()},
 		vss[1:nVals]...)
 
-	ensureNewRound(t, newRoundCh, height+1, 0)
+	roundChecker.ensureNewRound(height+1, 0)
 
 	// HEIGHT 2
 	height++
@@ -395,12 +403,12 @@ func setupSimulator(ctx context.Context, t *testing.T) *simulatorTestSuite {
 	if err := css[0].SetProposalAndBlock(proposal, propBlock, propBlockParts, "some peer"); err != nil {
 		t.Fatal(err)
 	}
-	ensureNewProposal(t, proposalCh, height, round)
+	proposalChecker.ensureNewProposal(height, round)
 	rs = css[0].GetRoundState()
 	signAddVotes(ctx, css[0], tmproto.PrecommitType, sim.Config.ChainID(),
 		types.BlockID{Hash: rs.ProposalBlock.Hash(), PartSetHeader: rs.ProposalBlockParts.Header()},
 		vss[1:nVals]...)
-	ensureNewRound(t, newRoundCh, height+1, 0)
+	roundChecker.ensureNewRound(height+1, 0)
 
 	// HEIGHT 3
 	height++
@@ -427,12 +435,12 @@ func setupSimulator(ctx context.Context, t *testing.T) *simulatorTestSuite {
 	if err := css[0].SetProposalAndBlock(proposal, propBlock, propBlockParts, "some peer"); err != nil {
 		t.Fatal(err)
 	}
-	ensureNewProposal(t, proposalCh, height, round)
+	proposalChecker.ensureNewProposal(height, round)
 	rs = css[0].GetRoundState()
 	signAddVotes(ctx, css[0], tmproto.PrecommitType, sim.Config.ChainID(),
 		types.BlockID{Hash: rs.ProposalBlock.Hash(), PartSetHeader: rs.ProposalBlockParts.Header()},
 		vss[1:nVals]...)
-	ensureNewRound(t, newRoundCh, height+1, 0)
+	roundChecker.ensureNewRound(height+1, 0)
 
 	// HEIGHT 4
 	height++
@@ -486,7 +494,7 @@ func setupSimulator(ctx context.Context, t *testing.T) *simulatorTestSuite {
 	if err := css[0].SetProposalAndBlock(proposal, propBlock, propBlockParts, "some peer"); err != nil {
 		t.Fatal(err)
 	}
-	ensureNewProposal(t, proposalCh, height, round)
+	proposalChecker.ensureNewProposal(height, round)
 
 	removeValidatorTx2 := kvstore.MakeValSetChangeTx(newVal2ABCI, 0)
 	err = assertMempool(css[0].txNotifier).CheckTx(ctx, removeValidatorTx2, nil, mempool.TxInfo{})
@@ -503,7 +511,7 @@ func setupSimulator(ctx context.Context, t *testing.T) *simulatorTestSuite {
 			newVss[i])
 	}
 
-	ensureNewRound(t, newRoundCh, height+1, 0)
+	roundChecker.ensureNewRound(height+1, 0)
 
 	// HEIGHT 5
 	height++
@@ -513,7 +521,7 @@ func setupSimulator(ctx context.Context, t *testing.T) *simulatorTestSuite {
 	newVss[newVssIdx].VotingPower = 25
 	sort.Sort(ValidatorStubsByPower(newVss))
 	selfIndex = valIndexFn(0)
-	ensureNewProposal(t, proposalCh, height, round)
+	proposalChecker.ensureNewProposal(height, round)
 	rs = css[0].GetRoundState()
 	for i := 0; i < nVals+1; i++ {
 		if i == selfIndex {
@@ -524,7 +532,7 @@ func setupSimulator(ctx context.Context, t *testing.T) *simulatorTestSuite {
 			types.BlockID{Hash: rs.ProposalBlock.Hash(), PartSetHeader: rs.ProposalBlockParts.Header()},
 			newVss[i])
 	}
-	ensureNewRound(t, newRoundCh, height+1, 0)
+	roundChecker.ensureNewRound(height+1, 0)
 
 	// HEIGHT 6
 	height++
@@ -551,7 +559,7 @@ func setupSimulator(ctx context.Context, t *testing.T) *simulatorTestSuite {
 	if err := css[0].SetProposalAndBlock(proposal, propBlock, propBlockParts, "some peer"); err != nil {
 		t.Fatal(err)
 	}
-	ensureNewProposal(t, proposalCh, height, round)
+	proposalChecker.ensureNewProposal(height, round)
 	rs = css[0].GetRoundState()
 	for i := 0; i < nVals+3; i++ {
 		if i == selfIndex {
@@ -562,7 +570,7 @@ func setupSimulator(ctx context.Context, t *testing.T) *simulatorTestSuite {
 			types.BlockID{Hash: rs.ProposalBlock.Hash(), PartSetHeader: rs.ProposalBlockParts.Header()},
 			newVss[i])
 	}
-	ensureNewRound(t, newRoundCh, height+1, 0)
+	roundChecker.ensureNewRound(height+1, 0)
 
 	sim.Chain = make([]*types.Block, 0)
 	sim.Commits = make([]*types.Commit, 0)