From cb970564d7bb598df48948b6f2d9793b868e2dde Mon Sep 17 00:00:00 2001
From: Thane Thomson <connect@thanethomson.com>
Date: Thu, 14 Apr 2022 16:41:08 -0400
Subject: [PATCH] consensus: fix flaky
 TestPrepareProposalReceivesVoteExtensions

Signed-off-by: Thane Thomson <connect@thanethomson.com>
---
 internal/consensus/common_test.go | 17 ++++++++---------
 internal/consensus/state_test.go  | 12 ++++++------
 2 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/internal/consensus/common_test.go b/internal/consensus/common_test.go
index acde5aa76..a0603405b 100644
--- a/internal/consensus/common_test.go
+++ b/internal/consensus/common_test.go
@@ -350,20 +350,19 @@ func validatePrecommit(
 		require.True(t, bytes.Equal(vote.BlockID.Hash, votedBlockHash), "Expected precommit to be for proposal block")
 	}
 
-	cs.mtx.RLock()
-	defer cs.mtx.RUnlock()
+	rs := cs.GetRoundState()
 	if lockedBlockHash == nil {
-		require.False(t, cs.LockedRound != lockRound || cs.LockedBlock != nil,
+		require.False(t, rs.LockedRound != lockRound || rs.LockedBlock != nil,
 			"Expected to be locked on nil at round %d. Got locked at round %d with block %v",
 			lockRound,
-			cs.LockedRound,
-			cs.LockedBlock)
+			rs.LockedRound,
+			rs.LockedBlock)
 	} else {
-		require.False(t, cs.LockedRound != lockRound || !bytes.Equal(cs.LockedBlock.Hash(), lockedBlockHash),
+		require.False(t, rs.LockedRound != lockRound || !bytes.Equal(rs.LockedBlock.Hash(), lockedBlockHash),
 			"Expected block to be locked on round %d, got %d. Got locked block %X, expected %X",
 			lockRound,
-			cs.LockedRound,
-			cs.LockedBlock.Hash(),
+			rs.LockedRound,
+			rs.LockedBlock.Hash(),
 			lockedBlockHash)
 	}
 }
@@ -380,7 +379,7 @@ func subscribeToVoter(ctx context.Context, t *testing.T, cs *State, addr []byte,
 		vt[t] = struct{}{}
 	}
 
-	ch := make(chan tmpubsub.Message)
+	ch := make(chan tmpubsub.Message, 1)
 	if err := cs.eventBus.Observe(ctx, func(msg tmpubsub.Message) error {
 		vote := msg.Data().(types.EventDataVote)
 		// we only fire for our own votes
diff --git a/internal/consensus/state_test.go b/internal/consensus/state_test.go
index ac3f865c0..345b65f56 100644
--- a/internal/consensus/state_test.go
+++ b/internal/consensus/state_test.go
@@ -2191,6 +2191,10 @@ func TestPrepareProposalReceivesVoteExtensions(t *testing.T) {
 
 	newRoundCh := subscribe(ctx, t, cs1.eventBus, types.EventQueryNewRound)
 	proposalCh := subscribe(ctx, t, cs1.eventBus, types.EventQueryCompleteProposal)
+	pv1, err := cs1.privValidator.GetPubKey(ctx)
+	require.NoError(t, err)
+	addr := pv1.Address()
+	voteCh := subscribeToVoter(ctx, t, cs1, addr)
 
 	startTestRound(ctx, cs1, height, round)
 	ensureNewRound(t, newRoundCh, height, round)
@@ -2208,14 +2212,10 @@ func TestPrepareProposalReceivesVoteExtensions(t *testing.T) {
 		signAddPrecommitWithExtension(ctx, t, cs1, config.ChainID(), blockID, voteExtensions[i+1], vs)
 	}
 
-	pv1, err := cs1.privValidator.GetPubKey(ctx)
-	require.NoError(t, err)
-	addr := pv1.Address()
-	voteCh := subscribeToVoter(ctx, t, cs1, addr, tmproto.PrecommitType)
+	ensurePrevote(t, voteCh, height, round)
 
 	// ensure that the height is committed.
-	ensurePrecommit(t, voteCh, height, round)
-	validatePrecommit(ctx, t, cs1, round, round, vss[0], blockID.Hash, blockID.Hash)
+	ensurePrecommitMatch(t, voteCh, height, round, blockID.Hash)
 	incrementHeight(vss[1:]...)
 
 	height++