From deb40c2bbbd121517f96ee0eb5adcb69a569cd5d Mon Sep 17 00:00:00 2001
From: gfanton <8671905+gfanton@users.noreply.github.com>
Date: Wed, 10 Jan 2024 13:47:44 +0100
Subject: [PATCH 1/8] fix: add state_test ensure drained channel

Signed-off-by: gfanton <8671905+gfanton@users.noreply.github.com>
---
 tm2/pkg/bft/consensus/common_test.go | 60 +++++++++++++++++++
 tm2/pkg/bft/consensus/state_test.go  | 88 ++++++++++++++++++----------
 2 files changed, 118 insertions(+), 30 deletions(-)

diff --git a/tm2/pkg/bft/consensus/common_test.go b/tm2/pkg/bft/consensus/common_test.go
index ba19881aace..4f9dd0db7c6 100644
--- a/tm2/pkg/bft/consensus/common_test.go
+++ b/tm2/pkg/bft/consensus/common_test.go
@@ -6,6 +6,7 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"reflect"
 	"sort"
 	"sync"
 	"testing"
@@ -563,6 +564,22 @@ func ensureNewEventOnChannel(ch <-chan events.Event) {
 	}
 }
 
+func ensureGetRoundState(to *ConsensusState) (rs *cstypes.RoundState) {
+	time.Sleep(time.Second * 5)
+	crs := make(chan *cstypes.RoundState)
+	go func() {
+		crs <- to.GetRoundState()
+	}()
+
+	select {
+	case <-time.After(ensureTimeout):
+		panic("Timeout expired while waiting for GetRoundState")
+	case rs = <-crs:
+	}
+
+	return
+}
+
 // -------------------------------------------------------------------------------
 // consensus nets
 
@@ -794,3 +811,46 @@ func newPersistentKVStore() abci.Application {
 func newPersistentKVStoreWithPath(dbDir string) abci.Application {
 	return kvstore.NewPersistentKVStoreApplication(dbDir)
 }
+
+// ------------------------------------
+
+func ensureDrainedChannels(t *testing.T, channels ...any) {
+	r := recover()
+	if r == nil {
+		return
+	}
+
+	t.Helper()
+
+	t.Logf("checking for drained channel")
+	leaks := make(map[string]int)
+	for _, ch := range channels {
+		chVal := reflect.ValueOf(ch)
+		if chVal.Kind() != reflect.Chan {
+			panic(chVal.Type().Name() + " not a channel")
+		}
+
+		// Use a select statement with reflection
+		cases := []reflect.SelectCase{
+			{Dir: reflect.SelectRecv, Chan: chVal},
+			{Dir: reflect.SelectDefault},
+		}
+
+		for {
+			chosen, recv, recvOK := reflect.Select(cases)
+			if chosen != 0 || !recvOK {
+				break
+			}
+
+			leaks[reflect.TypeOf(recv.Interface()).String()]++
+			time.Sleep(time.Millisecond * 500)
+		}
+	}
+
+	for leak, count := range leaks {
+		fmt.Printf("channel %q: %d events left\n", leak, count)
+		// assert.Fail(t, "event leak", "channel %q: %d events left", leak, count)
+	}
+
+	panic(r)
+}
diff --git a/tm2/pkg/bft/consensus/state_test.go b/tm2/pkg/bft/consensus/state_test.go
index 35877837ab3..7d51b6feef8 100644
--- a/tm2/pkg/bft/consensus/state_test.go
+++ b/tm2/pkg/bft/consensus/state_test.go
@@ -76,6 +76,7 @@ func TestStateProposerSelection0(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, proposalCh, newRoundCh)
 
 	// Wait for new round so proposer is set.
 	ensureNewRound(newRoundCh, height, round)
@@ -121,19 +122,20 @@ func TestStateProposerSelection2(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, newRoundCh)
 
 	ensureNewRound(newRoundCh, height, round) // wait for the new round
 
 	// everyone just votes nil. we get a new proposer each round
 	for i := 0; i < len(vss); i++ {
-		prop := cs1.GetRoundState().Validators.GetProposer()
+		prop := ensureGetRoundState(cs1).Validators.GetProposer()
 		addr := vss[(i+round)%len(vss)].GetPubKey().Address()
 		correctProposer := addr
 		if prop.Address != correctProposer {
 			panic(fmt.Sprintf("expected RoundState.Validators.GetProposer() to be validator %d. Got %X", (i+2)%len(vss), prop.Address))
 		}
 
-		rs := cs1.GetRoundState()
+		rs := ensureGetRoundState(cs1)
 		signAddVotes(cs1, types.PrecommitType, nil, rs.ProposalBlockParts.Header(), vss[1:]...)
 		ensureNewRound(newRoundCh, height, i+round+1) // wait for the new round event each round
 		incrementRound(vss[1:]...)
@@ -156,13 +158,14 @@ func TestStateEnterProposeNoPrivValidator(t *testing.T) {
 		cs.Stop()
 		cs.Wait()
 	}()
-
+	defer ensureDrainedChannels(t, timeoutCh)
 	// if we're not a validator, EnterPropose should timeout
 	ensureNewTimeout(timeoutCh, height, round, cs.config.TimeoutPropose.Nanoseconds())
 
-	if cs.GetRoundState().Proposal != nil {
+	if ensureGetRoundState(cs).Proposal != nil {
 		t.Error("Expected to make no proposal, since no privValidator")
 	}
+
 }
 
 // a validator should not timeout of the prevote round (TODO: unless the block is really big!)
@@ -183,13 +186,14 @@ func TestStateEnterProposeYesPrivValidator(t *testing.T) {
 		cs.Stop()
 		cs.Wait()
 	}()
+	defer ensureDrainedChannels(t, proposalCh, newRoundCh, timeoutCh)
 
 	// Wait for new round so proposer is set.
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
 
 	// Check that Proposal, ProposalBlock, ProposalBlockParts are set.
-	rs := cs.GetRoundState()
+	rs := ensureGetRoundState(cs)
 	if rs.Proposal == nil {
 		t.Error("rs.Proposal should be set")
 	}
@@ -247,6 +251,7 @@ func TestStateBadProposal(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, proposalCh, voteCh)
 
 	// wait for proposal
 	ensureProposal(proposalCh, height, round, blockID)
@@ -285,11 +290,12 @@ func TestStateFullRound1(t *testing.T) {
 		cs.Stop()
 		cs.Wait()
 	}()
+	defer ensureDrainedChannels(t, newRoundCh, voteCh, propCh)
 
 	ensureNewRound(newRoundCh, height, round)
 
 	ensureNewProposal(propCh, height, round)
-	propBlockHash := cs.GetRoundState().ProposalBlock.Hash()
+	propBlockHash := ensureGetRoundState(cs).ProposalBlock.Hash()
 
 	ensurePrevote(voteCh, height, round) // wait for prevote
 	validatePrevote(cs, round, vss[0], propBlockHash)
@@ -319,6 +325,7 @@ func TestStateFullRoundNil(t *testing.T) {
 		cs.Stop()
 		cs.Wait()
 	}()
+	defer ensureDrainedChannels(t, voteCh)
 
 	ensurePrevote(voteCh, height, round)   // prevote
 	ensurePrecommit(voteCh, height, round) // precommit
@@ -345,11 +352,12 @@ func TestStateFullRound2(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, voteCh, newBlockCh)
 
 	ensurePrevote(voteCh, height, round) // prevote
 
 	// we should be stuck in limbo waiting for more prevotes
-	rs := cs1.GetRoundState()
+	rs := ensureGetRoundState(cs1)
 	propBlockHash, propPartsHeader := rs.ProposalBlock.Hash(), rs.ProposalBlockParts.Header()
 
 	// prevote arrives from vs2:
@@ -390,6 +398,8 @@ func TestStateLockNoPOL(t *testing.T) {
 	proposalCh := subscribe(cs1.evsw, cstypes.EventCompleteProposal{})
 	newRoundCh := subscribe(cs1.evsw, cstypes.EventNewRound{})
 
+	defer ensureDrainedChannels(t, proposalCh, timeoutWaitCh, timeoutProposeCh, newRoundCh, voteCh)
+
 	/*
 		Round1 (cs1, B) // B B // B B2
 	*/
@@ -400,7 +410,7 @@ func TestStateLockNoPOL(t *testing.T) {
 	ensureNewRound(newRoundCh, height, round)
 
 	ensureNewProposal(proposalCh, height, round)
-	roundState := cs1.GetRoundState()
+	roundState := ensureGetRoundState(cs1)
 	theBlockHash := roundState.ProposalBlock.Hash()
 	thePartSetHeader := roundState.ProposalBlockParts.Header()
 
@@ -441,7 +451,7 @@ func TestStateLockNoPOL(t *testing.T) {
 	// now we're on a new round and not the proposer, so wait for timeout
 	ensureNewTimeout(timeoutProposeCh, height, round, cs1.config.Propose(round).Nanoseconds())
 
-	rs := cs1.GetRoundState()
+	rs := ensureGetRoundState(cs1)
 
 	if rs.ProposalBlock != nil {
 		panic("Expected proposal block to be nil")
@@ -483,7 +493,7 @@ func TestStateLockNoPOL(t *testing.T) {
 	incrementRound(vs2)
 
 	ensureNewProposal(proposalCh, height, round)
-	rs = cs1.GetRoundState()
+	rs = ensureGetRoundState(cs1)
 
 	// now we're on a new round and are the proposer
 	if !bytes.Equal(rs.ProposalBlock.Hash(), rs.LockedBlock.Hash()) {
@@ -579,10 +589,11 @@ func TestStateLockPOLRelock(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, proposalCh, timeoutWaitCh, newRoundCh, newBlockCh, voteCh)
 
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
-	rs := cs1.GetRoundState()
+	rs := ensureGetRoundState(cs1)
 	theBlockHash := rs.ProposalBlock.Hash()
 	theBlockParts := rs.ProposalBlockParts.Header()
 
@@ -676,10 +687,11 @@ func TestStateLockPOLUnlock(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, proposalCh, timeoutWaitCh, newRoundCh, voteCh, unlockCh)
 
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
-	rs := cs1.GetRoundState()
+	rs := ensureGetRoundState(cs1)
 	theBlockHash := rs.ProposalBlock.Hash()
 	theBlockParts := rs.ProposalBlockParts.Header()
 
@@ -770,10 +782,11 @@ func TestStateLockPOLSafety1(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, proposalCh, timeoutWaitCh, timeoutProposeCh, newRoundCh, voteCh)
 
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
-	rs := cs1.GetRoundState()
+	rs := ensureGetRoundState(cs1)
 	propBlock := rs.ProposalBlock
 
 	ensurePrevote(voteCh, height, round)
@@ -813,7 +826,7 @@ func TestStateLockPOLSafety1(t *testing.T) {
 
 	ensureNewProposal(proposalCh, height, round)
 
-	rs = cs1.GetRoundState()
+	rs = ensureGetRoundState(cs1)
 
 	if rs.LockedBlock != nil {
 		panic("we should not be locked!")
@@ -854,6 +867,7 @@ func TestStateLockPOLSafety1(t *testing.T) {
 	validatePrevote(cs1, round, vss[0], propBlockHash)
 
 	newStepCh := subscribe(cs1.evsw, cstypes.EventNewRoundStep{})
+	defer ensureDrainedChannels(t, newStepCh)
 
 	// before prevotes from the previous round are added
 	// add prevotes from the earlier round
@@ -912,6 +926,7 @@ func TestStateLockPOLSafety2(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, proposalCh, timeoutWaitCh, newRoundCh, unlockCh, voteCh)
 
 	ensureNewRound(newRoundCh, height, round)
 	if err := cs1.SetProposalAndBlock(prop1, propBlock1, propBlockParts1, "some peer"); err != nil {
@@ -992,9 +1007,11 @@ func TestProposeValidBlock(t *testing.T) {
 		cs1.Wait()
 	}()
 
+	defer ensureDrainedChannels(t, proposalCh, timeoutWaitCh, timeoutProposeCh, newRoundCh, unlockCh, voteCh)
+
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
-	rs := cs1.GetRoundState()
+	rs := ensureGetRoundState(cs1)
 	propBlock := rs.ProposalBlock
 	propBlockHash := propBlock.Hash()
 
@@ -1056,7 +1073,7 @@ func TestProposeValidBlock(t *testing.T) {
 
 	ensureNewProposal(proposalCh, height, round)
 
-	rs = cs1.GetRoundState()
+	rs = ensureGetRoundState(cs1)
 	assert.True(t, bytes.Equal(rs.ProposalBlock.Hash(), propBlockHash))
 	assert.True(t, bytes.Equal(rs.ProposalBlock.Hash(), rs.ValidBlock.Hash()))
 	assert.True(t, rs.Proposal.POLRound == rs.ValidRound)
@@ -1087,10 +1104,11 @@ func TestSetValidBlockOnDelayedPrevote(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, timeoutWaitCh, newRoundCh, validBlockCh, voteCh, proposalCh)
 
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
-	rs := cs1.GetRoundState()
+	rs := ensureGetRoundState(cs1)
 	propBlock := rs.ProposalBlock
 	propBlockHash := propBlock.Hash()
 	propBlockParts := propBlock.MakePartSet(partSize)
@@ -1110,7 +1128,7 @@ func TestSetValidBlockOnDelayedPrevote(t *testing.T) {
 	// we should have precommitted
 	validatePrecommit(t, cs1, round, -1, vss[0], nil, nil)
 
-	rs = cs1.GetRoundState()
+	rs = ensureGetRoundState(cs1)
 
 	assert.True(t, rs.ValidBlock == nil)
 	assert.True(t, rs.ValidBlockParts == nil)
@@ -1121,7 +1139,7 @@ func TestSetValidBlockOnDelayedPrevote(t *testing.T) {
 
 	ensureNewValidBlock(validBlockCh, height, round)
 
-	rs = cs1.GetRoundState()
+	rs = ensureGetRoundState(cs1)
 
 	assert.True(t, bytes.Equal(rs.ValidBlock.Hash(), propBlockHash))
 	assert.True(t, rs.ValidBlockParts.Header().Equals(propBlockParts.Header()))
@@ -1156,6 +1174,7 @@ func TestSetValidBlockOnDelayedProposal(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, timeoutWaitCh, timeoutProposeCh, newRoundCh, validBlockCh, voteCh, proposalCh)
 
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewTimeout(timeoutProposeCh, height, round, cs1.config.Propose(round).Nanoseconds())
@@ -1181,7 +1200,7 @@ func TestSetValidBlockOnDelayedProposal(t *testing.T) {
 	}
 
 	ensureNewProposal(proposalCh, height, round)
-	rs := cs1.GetRoundState()
+	rs := ensureGetRoundState(cs1)
 
 	assert.True(t, bytes.Equal(rs.ValidBlock.Hash(), propBlockHash))
 	assert.True(t, rs.ValidBlockParts.Header().Equals(propBlockParts.Header()))
@@ -1207,6 +1226,7 @@ func TestWaitingTimeoutOnNilPolka(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, timeoutWaitCh, newRoundCh)
 
 	ensureNewRound(newRoundCh, height, round)
 	signAddVotes(cs1, types.PrecommitType, nil, types.PartSetHeader{}, vs2, vs3, vs4)
@@ -1236,6 +1256,7 @@ func TestWaitingTimeoutProposeOnNewRound(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, timeoutWaitCh, newRoundCh, voteCh)
 
 	ensureNewRound(newRoundCh, height, round)
 	ensurePrevote(voteCh, height, round)
@@ -1246,7 +1267,7 @@ func TestWaitingTimeoutProposeOnNewRound(t *testing.T) {
 	round++ // moving to the next round
 	ensureNewRound(newRoundCh, height, round)
 
-	rs := cs1.GetRoundState()
+	rs := ensureGetRoundState(cs1)
 	assert.True(t, rs.Step == cstypes.RoundStepPropose) // P0 does not prevote before timeoutPropose expires
 
 	ensureNewTimeout(timeoutWaitCh, height, round, cs1.config.Propose(round).Nanoseconds())
@@ -1276,6 +1297,7 @@ func TestRoundSkipOnNilPolkaFromHigherRound(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, timeoutWaitCh, newRoundCh, voteCh)
 
 	ensureNewRound(newRoundCh, height, round)
 	ensurePrevote(voteCh, height, round)
@@ -1316,6 +1338,7 @@ func TestWaitTimeoutProposeOnNilPolkaForTheCurrentRound(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, timeoutProposeCh, newRoundCh, voteCh)
 
 	ensureNewRound(newRoundCh, height, round)
 	incrementRound(vss[1:]...)
@@ -1353,13 +1376,14 @@ func TestEmitNewValidBlockEventOnCommitWithoutBlock(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, newRoundCh, newRoundCh, validBlockCh)
 
 	ensureNewRound(newRoundCh, height, round)
 	// vs2, vs3 and vs4 send precommit for propBlock
 	signAddVotes(cs1, types.PrecommitType, propBlockHash, propBlockParts.Header(), vs2, vs3, vs4)
 	ensureNewValidBlock(validBlockCh, height, round)
 
-	rs := cs1.GetRoundState()
+	rs := ensureGetRoundState(cs1)
 	assert.True(t, rs.Step == cstypes.RoundStepCommit)
 	assert.True(t, rs.ProposalBlock == nil)
 	assert.True(t, rs.ProposalBlockParts.Header().Equals(propBlockParts.Header()))
@@ -1391,6 +1415,7 @@ func TestCommitFromPreviousRound(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, proposalCh, newRoundCh, validBlockCh)
 
 	ensureNewRound(newRoundCh, height, round)
 	// vs2, vs3 and vs4 send precommit for propBlock for the previous round
@@ -1398,7 +1423,7 @@ func TestCommitFromPreviousRound(t *testing.T) {
 
 	ensureNewValidBlock(validBlockCh, height, round)
 
-	rs := cs1.GetRoundState()
+	rs := ensureGetRoundState(cs1)
 	assert.True(t, rs.Step == cstypes.RoundStepCommit)
 	assert.True(t, rs.CommitRound == vs2.Round)
 	assert.True(t, rs.ProposalBlock == nil)
@@ -1446,10 +1471,11 @@ func TestStartNextHeightCorrectly(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, proposalCh, newRoundCh, voteCh, newBlockHeader)
 
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
-	rs := cs1.GetRoundState()
+	rs := ensureGetRoundState(cs1)
 	theBlockHash := rs.ProposalBlock.Hash()
 	theBlockParts := rs.ProposalBlockParts.Header()
 
@@ -1468,7 +1494,7 @@ func TestStartNextHeightCorrectly(t *testing.T) {
 	time.Sleep(5 * time.Millisecond)
 	signAddVotes(cs1, types.PrecommitType, theBlockHash, theBlockParts, vs4)
 
-	rs = cs1.GetRoundState()
+	rs = ensureGetRoundState(cs1)
 	assert.True(t, rs.TriggeredTimeoutPrecommit)
 
 	ensureNewBlockHeader(newBlockHeader, height, theBlockHash)
@@ -1478,7 +1504,7 @@ func TestStartNextHeightCorrectly(t *testing.T) {
 	height, round = height+1, 0
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewTimeout(timeoutProposeCh, height, round, cs1.config.Propose(round).Nanoseconds())
-	rs = cs1.GetRoundState()
+	rs = ensureGetRoundState(cs1)
 	assert.False(t, rs.TriggeredTimeoutPrecommit, "triggeredTimeoutPrecommit should be false at the beginning of each round")
 }
 
@@ -1507,10 +1533,11 @@ func TestFlappyResetTimeoutPrecommitUponNewHeight(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, proposalCh, newRoundCh, voteCh, newBlockHeader)
 
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
-	rs := cs1.GetRoundState()
+	rs := ensureGetRoundState(cs1)
 	theBlockHash := rs.ProposalBlock.Hash()
 	theBlockParts := rs.ProposalBlockParts.Header()
 
@@ -1538,7 +1565,7 @@ func TestFlappyResetTimeoutPrecommitUponNewHeight(t *testing.T) {
 	ensureNewRound(newRoundCh, height+1, 0)
 	ensureNewProposal(proposalCh, height+1, 0)
 
-	rs = cs1.GetRoundState()
+	rs = ensureGetRoundState(cs1)
 	assert.False(t, rs.TriggeredTimeoutPrecommit, "triggeredTimeoutPrecommit should be false at the beginning of each height")
 }
 
@@ -1653,10 +1680,11 @@ func TestFlappyStateHalt1(t *testing.T) {
 		cs1.Stop()
 		cs1.Wait()
 	}()
+	defer ensureDrainedChannels(t, proposalCh, timeoutWaitCh, newRoundCh, voteCh, newBlockCh)
 
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
-	rs := cs1.GetRoundState()
+	rs := ensureGetRoundState(cs1)
 	propBlock := rs.ProposalBlock
 	propBlockParts := propBlock.MakePartSet(partSize)
 
@@ -1682,7 +1710,7 @@ func TestFlappyStateHalt1(t *testing.T) {
 	round++ // moving to the next round
 
 	ensureNewRound(newRoundCh, height, round)
-	rs = cs1.GetRoundState()
+	rs = ensureGetRoundState(cs1)
 
 	t.Log("### ONTO ROUND 1")
 	/*Round2

From 455e24f4522ab576439aae19c324a976b610f6a4 Mon Sep 17 00:00:00 2001
From: gfanton <8671905+gfanton@users.noreply.github.com>
Date: Wed, 10 Jan 2024 15:33:01 +0100
Subject: [PATCH 2/8] fix: some state tests

Signed-off-by: gfanton <8671905+gfanton@users.noreply.github.com>
---
 tm2/pkg/bft/consensus/state_test.go | 48 ++++++++++++++++-------------
 1 file changed, 27 insertions(+), 21 deletions(-)

diff --git a/tm2/pkg/bft/consensus/state_test.go b/tm2/pkg/bft/consensus/state_test.go
index 7d51b6feef8..12af229c443 100644
--- a/tm2/pkg/bft/consensus/state_test.go
+++ b/tm2/pkg/bft/consensus/state_test.go
@@ -191,6 +191,7 @@ func TestStateEnterProposeYesPrivValidator(t *testing.T) {
 	// Wait for new round so proposer is set.
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
+	require.FailNow(t, "XXX: new round loop block get round state")
 
 	// Check that Proposal, ProposalBlock, ProposalBlockParts are set.
 	rs := ensureGetRoundState(cs)
@@ -295,6 +296,8 @@ func TestStateFullRound1(t *testing.T) {
 	ensureNewRound(newRoundCh, height, round)
 
 	ensureNewProposal(propCh, height, round)
+
+	require.FailNow(t, "XXX: unable to get the proper order to avoid failing")
 	propBlockHash := ensureGetRoundState(cs).ProposalBlock.Hash()
 
 	ensurePrevote(voteCh, height, round) // wait for prevote
@@ -306,6 +309,7 @@ func TestStateFullRound1(t *testing.T) {
 	ensureNewRound(newRoundCh, height+1, 0)
 
 	validateLastPrecommit(cs, vss[0], propBlockHash)
+
 }
 
 // nil is proposed, so prevote and precommit nil
@@ -410,18 +414,18 @@ func TestStateLockNoPOL(t *testing.T) {
 	ensureNewRound(newRoundCh, height, round)
 
 	ensureNewProposal(proposalCh, height, round)
+	ensurePrevote(voteCh, height, round) // prevote
 	roundState := ensureGetRoundState(cs1)
 	theBlockHash := roundState.ProposalBlock.Hash()
 	thePartSetHeader := roundState.ProposalBlockParts.Header()
-
-	ensurePrevote(voteCh, height, round) // prevote
+	validatePrevote(cs1, round, vss[0], theBlockHash)
 
 	// we should now be stuck in limbo forever, waiting for more prevotes
 	// prevote arrives from vs2:
 	signAddVotes(cs1, types.PrevoteType, theBlockHash, thePartSetHeader, vs2)
-	ensurePrevote(voteCh, height, round) // prevote
-
+	ensurePrevote(voteCh, height, round)   // prevote
 	ensurePrecommit(voteCh, height, round) // precommit
+
 	// the proposed block should now be locked and our precommit added
 	validatePrecommit(t, cs1, round, round, vss[0], theBlockHash, theBlockHash)
 
@@ -451,14 +455,13 @@ func TestStateLockNoPOL(t *testing.T) {
 	// now we're on a new round and not the proposer, so wait for timeout
 	ensureNewTimeout(timeoutProposeCh, height, round, cs1.config.Propose(round).Nanoseconds())
 
+	// wait to finish prevote
+	ensurePrevote(voteCh, height, round)
 	rs := ensureGetRoundState(cs1)
-
 	if rs.ProposalBlock != nil {
 		panic("Expected proposal block to be nil")
 	}
 
-	// wait to finish prevote
-	ensurePrevote(voteCh, height, round)
 	// we should have prevoted our locked block
 	validatePrevote(cs1, round, vss[0], rs.LockedBlock.Hash())
 
@@ -491,8 +494,8 @@ func TestStateLockNoPOL(t *testing.T) {
 	*/
 
 	incrementRound(vs2)
-
 	ensureNewProposal(proposalCh, height, round)
+	ensurePrevote(voteCh, height, round) // prevote
 	rs = ensureGetRoundState(cs1)
 
 	// now we're on a new round and are the proposer
@@ -500,9 +503,7 @@ func TestStateLockNoPOL(t *testing.T) {
 		panic(fmt.Sprintf("Expected proposal block to be locked block. Got %v, Expected %v", rs.ProposalBlock, rs.LockedBlock))
 	}
 
-	ensurePrevote(voteCh, height, round) // prevote
 	validatePrevote(cs1, round, vss[0], rs.LockedBlock.Hash())
-
 	signAddVotes(cs1, types.PrevoteType, hash, rs.ProposalBlock.MakePartSet(partSize).Header(), vs2)
 	ensurePrevote(voteCh, height, round)
 
@@ -593,11 +594,12 @@ func TestStateLockPOLRelock(t *testing.T) {
 
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
+	ensurePrevote(voteCh, height, round) // prevote
 	rs := ensureGetRoundState(cs1)
 	theBlockHash := rs.ProposalBlock.Hash()
 	theBlockParts := rs.ProposalBlockParts.Header()
 
-	ensurePrevote(voteCh, height, round) // prevote
+	validatePrevote(cs1, round, vss[0], theBlockHash)
 
 	signAddVotes(cs1, types.PrevoteType, theBlockHash, theBlockParts, vs2, vs3, vs4)
 
@@ -691,11 +693,11 @@ func TestStateLockPOLUnlock(t *testing.T) {
 
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
+	ensurePrevote(voteCh, height, round)
 	rs := ensureGetRoundState(cs1)
 	theBlockHash := rs.ProposalBlock.Hash()
 	theBlockParts := rs.ProposalBlockParts.Header()
 
-	ensurePrevote(voteCh, height, round)
 	validatePrevote(cs1, round, vss[0], theBlockHash)
 
 	signAddVotes(cs1, types.PrevoteType, theBlockHash, theBlockParts, vs2, vs3, vs4)
@@ -786,10 +788,10 @@ func TestStateLockPOLSafety1(t *testing.T) {
 
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
+	ensurePrevote(voteCh, height, round)
 	rs := ensureGetRoundState(cs1)
 	propBlock := rs.ProposalBlock
 
-	ensurePrevote(voteCh, height, round)
 	validatePrevote(cs1, round, vss[0], propBlock.Hash())
 
 	// the others sign a polka but we don't see it
@@ -824,17 +826,15 @@ func TestStateLockPOLSafety1(t *testing.T) {
 	// a polka happened but we didn't see it!
 	*/
 
+	// go to prevote, prevote for proposal block
 	ensureNewProposal(proposalCh, height, round)
-
+	ensurePrevote(voteCh, height, round)
 	rs = ensureGetRoundState(cs1)
-
 	if rs.LockedBlock != nil {
 		panic("we should not be locked!")
 	}
 	t.Logf("new prop hash %v", fmt.Sprintf("%X", propBlockHash))
 
-	// go to prevote, prevote for proposal block
-	ensurePrevote(voteCh, height, round)
 	validatePrevote(cs1, round, vss[0], propBlockHash)
 
 	// now we see the others prevote for it, so we should lock on it
@@ -1011,11 +1011,11 @@ func TestProposeValidBlock(t *testing.T) {
 
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
+	ensurePrevote(voteCh, height, round)
 	rs := ensureGetRoundState(cs1)
 	propBlock := rs.ProposalBlock
 	propBlockHash := propBlock.Hash()
 
-	ensurePrevote(voteCh, height, round)
 	validatePrevote(cs1, round, vss[0], propBlockHash)
 
 	// the others sign a polka
@@ -1072,7 +1072,7 @@ func TestProposeValidBlock(t *testing.T) {
 	t.Log("### ONTO ROUND 4")
 
 	ensureNewProposal(proposalCh, height, round)
-
+	ensurePrevote(voteCh, height, round)
 	rs = ensureGetRoundState(cs1)
 	assert.True(t, bytes.Equal(rs.ProposalBlock.Hash(), propBlockHash))
 	assert.True(t, bytes.Equal(rs.ProposalBlock.Hash(), rs.ValidBlock.Hash()))
@@ -1108,12 +1108,12 @@ func TestSetValidBlockOnDelayedPrevote(t *testing.T) {
 
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
+	ensurePrevote(voteCh, height, round)
 	rs := ensureGetRoundState(cs1)
 	propBlock := rs.ProposalBlock
 	propBlockHash := propBlock.Hash()
 	propBlockParts := propBlock.MakePartSet(partSize)
 
-	ensurePrevote(voteCh, height, round)
 	validatePrevote(cs1, round, vss[0], propBlockHash)
 
 	// vs2 send prevote for propBlock
@@ -1267,6 +1267,8 @@ func TestWaitingTimeoutProposeOnNewRound(t *testing.T) {
 	round++ // moving to the next round
 	ensureNewRound(newRoundCh, height, round)
 
+	require.FailNow(t, "XXX: this part is flaky")
+
 	rs := ensureGetRoundState(cs1)
 	assert.True(t, rs.Step == cstypes.RoundStepPropose) // P0 does not prevote before timeoutPropose expires
 
@@ -1475,11 +1477,11 @@ func TestStartNextHeightCorrectly(t *testing.T) {
 
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
+	ensurePrevote(voteCh, height, round)
 	rs := ensureGetRoundState(cs1)
 	theBlockHash := rs.ProposalBlock.Hash()
 	theBlockParts := rs.ProposalBlockParts.Header()
 
-	ensurePrevote(voteCh, height, round)
 	validatePrevote(cs1, round, vss[0], theBlockHash)
 
 	signAddVotes(cs1, types.PrevoteType, theBlockHash, theBlockParts, vs2, vs3, vs4)
@@ -1494,6 +1496,9 @@ func TestStartNextHeightCorrectly(t *testing.T) {
 	time.Sleep(5 * time.Millisecond)
 	signAddVotes(cs1, types.PrecommitType, theBlockHash, theBlockParts, vs4)
 
+	require.FailNow(t, "XXX: this part is flaky")
+
+	// XXX: this make the test hang indefinitely need double check
 	rs = ensureGetRoundState(cs1)
 	assert.True(t, rs.TriggeredTimeoutPrecommit)
 
@@ -1504,6 +1509,7 @@ func TestStartNextHeightCorrectly(t *testing.T) {
 	height, round = height+1, 0
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewTimeout(timeoutProposeCh, height, round, cs1.config.Propose(round).Nanoseconds())
+	ensurePrevote(voteCh, height, round)
 	rs = ensureGetRoundState(cs1)
 	assert.False(t, rs.TriggeredTimeoutPrecommit, "triggeredTimeoutPrecommit should be false at the beginning of each round")
 }

From dc6472820c4c5bfd561f35f1990ca87900f95335 Mon Sep 17 00:00:00 2001
From: gfanton <8671905+gfanton@users.noreply.github.com>
Date: Wed, 10 Jan 2024 15:57:45 +0100
Subject: [PATCH 3/8] chore: lint

Signed-off-by: gfanton <8671905+gfanton@users.noreply.github.com>
---
 tm2/pkg/bft/consensus/common_test.go | 7 +++----
 tm2/pkg/bft/consensus/state_test.go  | 2 --
 2 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/tm2/pkg/bft/consensus/common_test.go b/tm2/pkg/bft/consensus/common_test.go
index 4f9dd0db7c6..68f28795256 100644
--- a/tm2/pkg/bft/consensus/common_test.go
+++ b/tm2/pkg/bft/consensus/common_test.go
@@ -815,13 +815,13 @@ func newPersistentKVStoreWithPath(dbDir string) abci.Application {
 // ------------------------------------
 
 func ensureDrainedChannels(t *testing.T, channels ...any) {
+	t.Helper()
+
 	r := recover()
 	if r == nil {
 		return
 	}
 
-	t.Helper()
-
 	t.Logf("checking for drained channel")
 	leaks := make(map[string]int)
 	for _, ch := range channels {
@@ -848,8 +848,7 @@ func ensureDrainedChannels(t *testing.T, channels ...any) {
 	}
 
 	for leak, count := range leaks {
-		fmt.Printf("channel %q: %d events left\n", leak, count)
-		// assert.Fail(t, "event leak", "channel %q: %d events left", leak, count)
+		t.Logf("channel %q: %d events left\n", leak, count)
 	}
 
 	panic(r)
diff --git a/tm2/pkg/bft/consensus/state_test.go b/tm2/pkg/bft/consensus/state_test.go
index 12af229c443..cbd93196af5 100644
--- a/tm2/pkg/bft/consensus/state_test.go
+++ b/tm2/pkg/bft/consensus/state_test.go
@@ -165,7 +165,6 @@ func TestStateEnterProposeNoPrivValidator(t *testing.T) {
 	if ensureGetRoundState(cs).Proposal != nil {
 		t.Error("Expected to make no proposal, since no privValidator")
 	}
-
 }
 
 // a validator should not timeout of the prevote round (TODO: unless the block is really big!)
@@ -309,7 +308,6 @@ func TestStateFullRound1(t *testing.T) {
 	ensureNewRound(newRoundCh, height+1, 0)
 
 	validateLastPrecommit(cs, vss[0], propBlockHash)
-
 }
 
 // nil is proposed, so prevote and precommit nil

From c3da44e15c8d97aa553c288a5df9abcd7a28fee5 Mon Sep 17 00:00:00 2001
From: gfanton <8671905+gfanton@users.noreply.github.com>
Date: Thu, 11 Jan 2024 17:52:18 +0100
Subject: [PATCH 4/8] chore: add fixme note

Signed-off-by: gfanton <8671905+gfanton@users.noreply.github.com>
---
 tm2/pkg/bft/consensus/state_test.go | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/tm2/pkg/bft/consensus/state_test.go b/tm2/pkg/bft/consensus/state_test.go
index cbd93196af5..d3396f69ff1 100644
--- a/tm2/pkg/bft/consensus/state_test.go
+++ b/tm2/pkg/bft/consensus/state_test.go
@@ -190,7 +190,8 @@ func TestStateEnterProposeYesPrivValidator(t *testing.T) {
 	// Wait for new round so proposer is set.
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
-	require.FailNow(t, "XXX: new round loop block get round state")
+
+	// XXX(FIXME): new round loop block get round state
 
 	// Check that Proposal, ProposalBlock, ProposalBlockParts are set.
 	rs := ensureGetRoundState(cs)
@@ -296,7 +297,7 @@ func TestStateFullRound1(t *testing.T) {
 
 	ensureNewProposal(propCh, height, round)
 
-	require.FailNow(t, "XXX: unable to get the proper order to avoid failing")
+	// XXX(FIXME): unable to get the proper order to avoid failing
 	propBlockHash := ensureGetRoundState(cs).ProposalBlock.Hash()
 
 	ensurePrevote(voteCh, height, round) // wait for prevote
@@ -1265,7 +1266,7 @@ func TestWaitingTimeoutProposeOnNewRound(t *testing.T) {
 	round++ // moving to the next round
 	ensureNewRound(newRoundCh, height, round)
 
-	require.FailNow(t, "XXX: this part is flaky")
+	// XXX(FIXME): this part is flaky
 
 	rs := ensureGetRoundState(cs1)
 	assert.True(t, rs.Step == cstypes.RoundStepPropose) // P0 does not prevote before timeoutPropose expires
@@ -1494,9 +1495,7 @@ func TestStartNextHeightCorrectly(t *testing.T) {
 	time.Sleep(5 * time.Millisecond)
 	signAddVotes(cs1, types.PrecommitType, theBlockHash, theBlockParts, vs4)
 
-	require.FailNow(t, "XXX: this part is flaky")
-
-	// XXX: this make the test hang indefinitely need double check
+	// XXX(FIXME): this make the test hang indefinitely
 	rs = ensureGetRoundState(cs1)
 	assert.True(t, rs.TriggeredTimeoutPrecommit)
 

From 12ca714af46230f168073f394c333f2f06884453 Mon Sep 17 00:00:00 2001
From: gfanton <8671905+gfanton@users.noreply.github.com>
Date: Wed, 1 May 2024 14:56:54 +0200
Subject: [PATCH 5/8] fix: add max exp

Signed-off-by: gfanton <8671905+gfanton@users.noreply.github.com>
---
 tm2/pkg/bft/consensus/common_test.go | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tm2/pkg/bft/consensus/common_test.go b/tm2/pkg/bft/consensus/common_test.go
index a0dce2d6546..a0c60f494d7 100644
--- a/tm2/pkg/bft/consensus/common_test.go
+++ b/tm2/pkg/bft/consensus/common_test.go
@@ -798,9 +798,12 @@ func ensureDrainedChannels(t *testing.T, channels ...any) {
 			panic(chVal.Type().Name() + " not a channel")
 		}
 
+		maxExp := time.After(time.Second * 5)
+
 		// Use a select statement with reflection
 		cases := []reflect.SelectCase{
 			{Dir: reflect.SelectRecv, Chan: chVal},
+			{Dir: reflect.SelectRecv, Chan: reflect.ValueOf(maxExp)},
 			{Dir: reflect.SelectDefault},
 		}
 

From 020d82ee59a51a33ee649726c9937389b4e591a2 Mon Sep 17 00:00:00 2001
From: gfanton <8671905+gfanton@users.noreply.github.com>
Date: Tue, 7 May 2024 00:01:30 +0200
Subject: [PATCH 6/8] fix: remove ensureGetRoundState

Signed-off-by: gfanton <8671905+gfanton@users.noreply.github.com>
---
 tm2/pkg/bft/consensus/common_test.go | 16 --------
 tm2/pkg/bft/consensus/state_test.go  | 58 ++++++++++++++--------------
 2 files changed, 29 insertions(+), 45 deletions(-)

diff --git a/tm2/pkg/bft/consensus/common_test.go b/tm2/pkg/bft/consensus/common_test.go
index a0c60f494d7..f657bf3b6d9 100644
--- a/tm2/pkg/bft/consensus/common_test.go
+++ b/tm2/pkg/bft/consensus/common_test.go
@@ -565,22 +565,6 @@ func ensureNewEventOnChannel(ch <-chan events.Event) {
 	}
 }
 
-func ensureGetRoundState(to *ConsensusState) (rs *cstypes.RoundState) {
-	time.Sleep(time.Second * 5)
-	crs := make(chan *cstypes.RoundState)
-	go func() {
-		crs <- to.GetRoundState()
-	}()
-
-	select {
-	case <-time.After(ensureTimeout):
-		panic("Timeout expired while waiting for GetRoundState")
-	case rs = <-crs:
-	}
-
-	return
-}
-
 // -------------------------------------------------------------------------------
 // consensus nets
 
diff --git a/tm2/pkg/bft/consensus/state_test.go b/tm2/pkg/bft/consensus/state_test.go
index d3396f69ff1..4c84f94aa80 100644
--- a/tm2/pkg/bft/consensus/state_test.go
+++ b/tm2/pkg/bft/consensus/state_test.go
@@ -128,14 +128,14 @@ func TestStateProposerSelection2(t *testing.T) {
 
 	// everyone just votes nil. we get a new proposer each round
 	for i := 0; i < len(vss); i++ {
-		prop := ensureGetRoundState(cs1).Validators.GetProposer()
+		prop := cs1.GetRoundState().Validators.GetProposer()
 		addr := vss[(i+round)%len(vss)].GetPubKey().Address()
 		correctProposer := addr
 		if prop.Address != correctProposer {
 			panic(fmt.Sprintf("expected RoundState.Validators.GetProposer() to be validator %d. Got %X", (i+2)%len(vss), prop.Address))
 		}
 
-		rs := ensureGetRoundState(cs1)
+		rs := cs1.GetRoundState()
 		signAddVotes(cs1, types.PrecommitType, nil, rs.ProposalBlockParts.Header(), vss[1:]...)
 		ensureNewRound(newRoundCh, height, i+round+1) // wait for the new round event each round
 		incrementRound(vss[1:]...)
@@ -162,7 +162,7 @@ func TestStateEnterProposeNoPrivValidator(t *testing.T) {
 	// if we're not a validator, EnterPropose should timeout
 	ensureNewTimeout(timeoutCh, height, round, cs.config.TimeoutPropose.Nanoseconds())
 
-	if ensureGetRoundState(cs).Proposal != nil {
+	if cs.GetRoundState().Proposal != nil {
 		t.Error("Expected to make no proposal, since no privValidator")
 	}
 }
@@ -194,7 +194,7 @@ func TestStateEnterProposeYesPrivValidator(t *testing.T) {
 	// XXX(FIXME): new round loop block get round state
 
 	// Check that Proposal, ProposalBlock, ProposalBlockParts are set.
-	rs := ensureGetRoundState(cs)
+	rs := cs.GetRoundState()
 	if rs.Proposal == nil {
 		t.Error("rs.Proposal should be set")
 	}
@@ -298,7 +298,7 @@ func TestStateFullRound1(t *testing.T) {
 	ensureNewProposal(propCh, height, round)
 
 	// XXX(FIXME): unable to get the proper order to avoid failing
-	propBlockHash := ensureGetRoundState(cs).ProposalBlock.Hash()
+	propBlockHash := cs.GetRoundState().ProposalBlock.Hash()
 
 	ensurePrevote(voteCh, height, round) // wait for prevote
 	validatePrevote(cs, round, vss[0], propBlockHash)
@@ -360,7 +360,7 @@ func TestStateFullRound2(t *testing.T) {
 	ensurePrevote(voteCh, height, round) // prevote
 
 	// we should be stuck in limbo waiting for more prevotes
-	rs := ensureGetRoundState(cs1)
+	rs := cs1.GetRoundState()
 	propBlockHash, propPartsHeader := rs.ProposalBlock.Hash(), rs.ProposalBlockParts.Header()
 
 	// prevote arrives from vs2:
@@ -414,7 +414,7 @@ func TestStateLockNoPOL(t *testing.T) {
 
 	ensureNewProposal(proposalCh, height, round)
 	ensurePrevote(voteCh, height, round) // prevote
-	roundState := ensureGetRoundState(cs1)
+	roundState := cs1.GetRoundState()
 	theBlockHash := roundState.ProposalBlock.Hash()
 	thePartSetHeader := roundState.ProposalBlockParts.Header()
 	validatePrevote(cs1, round, vss[0], theBlockHash)
@@ -456,7 +456,7 @@ func TestStateLockNoPOL(t *testing.T) {
 
 	// wait to finish prevote
 	ensurePrevote(voteCh, height, round)
-	rs := ensureGetRoundState(cs1)
+	rs := cs1.GetRoundState()
 	if rs.ProposalBlock != nil {
 		panic("Expected proposal block to be nil")
 	}
@@ -495,7 +495,7 @@ func TestStateLockNoPOL(t *testing.T) {
 	incrementRound(vs2)
 	ensureNewProposal(proposalCh, height, round)
 	ensurePrevote(voteCh, height, round) // prevote
-	rs = ensureGetRoundState(cs1)
+	rs = cs1.GetRoundState()
 
 	// now we're on a new round and are the proposer
 	if !bytes.Equal(rs.ProposalBlock.Hash(), rs.LockedBlock.Hash()) {
@@ -594,7 +594,7 @@ func TestStateLockPOLRelock(t *testing.T) {
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
 	ensurePrevote(voteCh, height, round) // prevote
-	rs := ensureGetRoundState(cs1)
+	rs := cs1.GetRoundState()
 	theBlockHash := rs.ProposalBlock.Hash()
 	theBlockParts := rs.ProposalBlockParts.Header()
 
@@ -693,7 +693,7 @@ func TestStateLockPOLUnlock(t *testing.T) {
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
 	ensurePrevote(voteCh, height, round)
-	rs := ensureGetRoundState(cs1)
+	rs := cs1.GetRoundState()
 	theBlockHash := rs.ProposalBlock.Hash()
 	theBlockParts := rs.ProposalBlockParts.Header()
 
@@ -788,7 +788,7 @@ func TestStateLockPOLSafety1(t *testing.T) {
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
 	ensurePrevote(voteCh, height, round)
-	rs := ensureGetRoundState(cs1)
+	rs := cs1.GetRoundState()
 	propBlock := rs.ProposalBlock
 
 	validatePrevote(cs1, round, vss[0], propBlock.Hash())
@@ -828,7 +828,7 @@ func TestStateLockPOLSafety1(t *testing.T) {
 	// go to prevote, prevote for proposal block
 	ensureNewProposal(proposalCh, height, round)
 	ensurePrevote(voteCh, height, round)
-	rs = ensureGetRoundState(cs1)
+	rs = cs1.GetRoundState()
 	if rs.LockedBlock != nil {
 		panic("we should not be locked!")
 	}
@@ -1011,7 +1011,7 @@ func TestProposeValidBlock(t *testing.T) {
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
 	ensurePrevote(voteCh, height, round)
-	rs := ensureGetRoundState(cs1)
+	rs := cs1.GetRoundState()
 	propBlock := rs.ProposalBlock
 	propBlockHash := propBlock.Hash()
 
@@ -1072,7 +1072,7 @@ func TestProposeValidBlock(t *testing.T) {
 
 	ensureNewProposal(proposalCh, height, round)
 	ensurePrevote(voteCh, height, round)
-	rs = ensureGetRoundState(cs1)
+	rs = cs1.GetRoundState()
 	assert.True(t, bytes.Equal(rs.ProposalBlock.Hash(), propBlockHash))
 	assert.True(t, bytes.Equal(rs.ProposalBlock.Hash(), rs.ValidBlock.Hash()))
 	assert.True(t, rs.Proposal.POLRound == rs.ValidRound)
@@ -1108,7 +1108,7 @@ func TestSetValidBlockOnDelayedPrevote(t *testing.T) {
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
 	ensurePrevote(voteCh, height, round)
-	rs := ensureGetRoundState(cs1)
+	rs := cs1.GetRoundState()
 	propBlock := rs.ProposalBlock
 	propBlockHash := propBlock.Hash()
 	propBlockParts := propBlock.MakePartSet(partSize)
@@ -1127,7 +1127,7 @@ func TestSetValidBlockOnDelayedPrevote(t *testing.T) {
 	// we should have precommitted
 	validatePrecommit(t, cs1, round, -1, vss[0], nil, nil)
 
-	rs = ensureGetRoundState(cs1)
+	rs = cs1.GetRoundState()
 
 	assert.True(t, rs.ValidBlock == nil)
 	assert.True(t, rs.ValidBlockParts == nil)
@@ -1138,7 +1138,7 @@ func TestSetValidBlockOnDelayedPrevote(t *testing.T) {
 
 	ensureNewValidBlock(validBlockCh, height, round)
 
-	rs = ensureGetRoundState(cs1)
+	rs = cs1.GetRoundState()
 
 	assert.True(t, bytes.Equal(rs.ValidBlock.Hash(), propBlockHash))
 	assert.True(t, rs.ValidBlockParts.Header().Equals(propBlockParts.Header()))
@@ -1199,7 +1199,7 @@ func TestSetValidBlockOnDelayedProposal(t *testing.T) {
 	}
 
 	ensureNewProposal(proposalCh, height, round)
-	rs := ensureGetRoundState(cs1)
+	rs := cs1.GetRoundState()
 
 	assert.True(t, bytes.Equal(rs.ValidBlock.Hash(), propBlockHash))
 	assert.True(t, rs.ValidBlockParts.Header().Equals(propBlockParts.Header()))
@@ -1268,7 +1268,7 @@ func TestWaitingTimeoutProposeOnNewRound(t *testing.T) {
 
 	// XXX(FIXME): this part is flaky
 
-	rs := ensureGetRoundState(cs1)
+	rs := cs1.GetRoundState()
 	assert.True(t, rs.Step == cstypes.RoundStepPropose) // P0 does not prevote before timeoutPropose expires
 
 	ensureNewTimeout(timeoutWaitCh, height, round, cs1.config.Propose(round).Nanoseconds())
@@ -1384,7 +1384,7 @@ func TestEmitNewValidBlockEventOnCommitWithoutBlock(t *testing.T) {
 	signAddVotes(cs1, types.PrecommitType, propBlockHash, propBlockParts.Header(), vs2, vs3, vs4)
 	ensureNewValidBlock(validBlockCh, height, round)
 
-	rs := ensureGetRoundState(cs1)
+	rs := cs1.GetRoundState()
 	assert.True(t, rs.Step == cstypes.RoundStepCommit)
 	assert.True(t, rs.ProposalBlock == nil)
 	assert.True(t, rs.ProposalBlockParts.Header().Equals(propBlockParts.Header()))
@@ -1424,7 +1424,7 @@ func TestCommitFromPreviousRound(t *testing.T) {
 
 	ensureNewValidBlock(validBlockCh, height, round)
 
-	rs := ensureGetRoundState(cs1)
+	rs := cs1.GetRoundState()
 	assert.True(t, rs.Step == cstypes.RoundStepCommit)
 	assert.True(t, rs.CommitRound == vs2.Round)
 	assert.True(t, rs.ProposalBlock == nil)
@@ -1477,7 +1477,7 @@ func TestStartNextHeightCorrectly(t *testing.T) {
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
 	ensurePrevote(voteCh, height, round)
-	rs := ensureGetRoundState(cs1)
+	rs := cs1.GetRoundState()
 	theBlockHash := rs.ProposalBlock.Hash()
 	theBlockParts := rs.ProposalBlockParts.Header()
 
@@ -1496,7 +1496,7 @@ func TestStartNextHeightCorrectly(t *testing.T) {
 	signAddVotes(cs1, types.PrecommitType, theBlockHash, theBlockParts, vs4)
 
 	// XXX(FIXME): this make the test hang indefinitely
-	rs = ensureGetRoundState(cs1)
+	rs = cs1.GetRoundState()
 	assert.True(t, rs.TriggeredTimeoutPrecommit)
 
 	ensureNewBlockHeader(newBlockHeader, height, theBlockHash)
@@ -1507,7 +1507,7 @@ func TestStartNextHeightCorrectly(t *testing.T) {
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewTimeout(timeoutProposeCh, height, round, cs1.config.Propose(round).Nanoseconds())
 	ensurePrevote(voteCh, height, round)
-	rs = ensureGetRoundState(cs1)
+	rs = cs1.GetRoundState()
 	assert.False(t, rs.TriggeredTimeoutPrecommit, "triggeredTimeoutPrecommit should be false at the beginning of each round")
 }
 
@@ -1540,7 +1540,7 @@ func TestFlappyResetTimeoutPrecommitUponNewHeight(t *testing.T) {
 
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
-	rs := ensureGetRoundState(cs1)
+	rs := cs1.GetRoundState()
 	theBlockHash := rs.ProposalBlock.Hash()
 	theBlockParts := rs.ProposalBlockParts.Header()
 
@@ -1568,7 +1568,7 @@ func TestFlappyResetTimeoutPrecommitUponNewHeight(t *testing.T) {
 	ensureNewRound(newRoundCh, height+1, 0)
 	ensureNewProposal(proposalCh, height+1, 0)
 
-	rs = ensureGetRoundState(cs1)
+	rs = cs1.GetRoundState()
 	assert.False(t, rs.TriggeredTimeoutPrecommit, "triggeredTimeoutPrecommit should be false at the beginning of each height")
 }
 
@@ -1687,7 +1687,7 @@ func TestFlappyStateHalt1(t *testing.T) {
 
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
-	rs := ensureGetRoundState(cs1)
+	rs := cs1.GetRoundState()
 	propBlock := rs.ProposalBlock
 	propBlockParts := propBlock.MakePartSet(partSize)
 
@@ -1713,7 +1713,7 @@ func TestFlappyStateHalt1(t *testing.T) {
 	round++ // moving to the next round
 
 	ensureNewRound(newRoundCh, height, round)
-	rs = ensureGetRoundState(cs1)
+	rs = cs1.GetRoundState()
 
 	t.Log("### ONTO ROUND 1")
 	/*Round2

From 51f40d5028af0ea1cfa3d5d8b43fa0f5c4da3d0c Mon Sep 17 00:00:00 2001
From: gfanton <8671905+gfanton@users.noreply.github.com>
Date: Tue, 7 May 2024 13:32:04 +0200
Subject: [PATCH 7/8] chore: cleanup `FIXME` comment

Signed-off-by: gfanton <8671905+gfanton@users.noreply.github.com>
---
 tm2/pkg/bft/consensus/state_test.go | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/tm2/pkg/bft/consensus/state_test.go b/tm2/pkg/bft/consensus/state_test.go
index 4c84f94aa80..edfa741340f 100644
--- a/tm2/pkg/bft/consensus/state_test.go
+++ b/tm2/pkg/bft/consensus/state_test.go
@@ -191,8 +191,6 @@ func TestStateEnterProposeYesPrivValidator(t *testing.T) {
 	ensureNewRound(newRoundCh, height, round)
 	ensureNewProposal(proposalCh, height, round)
 
-	// XXX(FIXME): new round loop block get round state
-
 	// Check that Proposal, ProposalBlock, ProposalBlockParts are set.
 	rs := cs.GetRoundState()
 	if rs.Proposal == nil {
@@ -297,7 +295,6 @@ func TestStateFullRound1(t *testing.T) {
 
 	ensureNewProposal(propCh, height, round)
 
-	// XXX(FIXME): unable to get the proper order to avoid failing
 	propBlockHash := cs.GetRoundState().ProposalBlock.Hash()
 
 	ensurePrevote(voteCh, height, round) // wait for prevote
@@ -1266,8 +1263,6 @@ func TestWaitingTimeoutProposeOnNewRound(t *testing.T) {
 	round++ // moving to the next round
 	ensureNewRound(newRoundCh, height, round)
 
-	// XXX(FIXME): this part is flaky
-
 	rs := cs1.GetRoundState()
 	assert.True(t, rs.Step == cstypes.RoundStepPropose) // P0 does not prevote before timeoutPropose expires
 
@@ -1495,7 +1490,6 @@ func TestStartNextHeightCorrectly(t *testing.T) {
 	time.Sleep(5 * time.Millisecond)
 	signAddVotes(cs1, types.PrecommitType, theBlockHash, theBlockParts, vs4)
 
-	// XXX(FIXME): this make the test hang indefinitely
 	rs = cs1.GetRoundState()
 	assert.True(t, rs.TriggeredTimeoutPrecommit)
 

From f17c9b591218bf2aa8dad703b4057bcd1824c62d Mon Sep 17 00:00:00 2001
From: gfanton <8671905+gfanton@users.noreply.github.com>
Date: Tue, 7 May 2024 16:49:28 +0200
Subject: [PATCH 8/8] fix: `TestStateProposerSelection0`

move `ensureProposal` before `GetRoundState`

Signed-off-by: gfanton <8671905+gfanton@users.noreply.github.com>
---
 tm2/pkg/bft/consensus/state_test.go | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tm2/pkg/bft/consensus/state_test.go b/tm2/pkg/bft/consensus/state_test.go
index edfa741340f..201cf8906b3 100644
--- a/tm2/pkg/bft/consensus/state_test.go
+++ b/tm2/pkg/bft/consensus/state_test.go
@@ -81,6 +81,9 @@ func TestStateProposerSelection0(t *testing.T) {
 	// Wait for new round so proposer is set.
 	ensureNewRound(newRoundCh, height, round)
 
+	// Wait for complete proposal.
+	ensureNewProposal(proposalCh, height, round)
+
 	// Commit a block and ensure proposer for the next height is correct.
 	prop := cs1.GetRoundState().Validators.GetProposer()
 	address := cs1.privValidator.GetPubKey().Address()
@@ -88,9 +91,6 @@ func TestStateProposerSelection0(t *testing.T) {
 		t.Fatalf("expected proposer to be validator %d. Got %X", 0, prop.Address)
 	}
 
-	// Wait for complete proposal.
-	ensureNewProposal(proposalCh, height, round)
-
 	rs := cs1.GetRoundState()
 	signAddVotes(cs1, types.PrecommitType, rs.ProposalBlock.Hash(), rs.ProposalBlockParts.Header(), vss[1:]...)