diff --git a/internal/consensus/common_test.go b/internal/consensus/common_test.go index 0a555ad05..b13d9c9d4 100644 --- a/internal/consensus/common_test.go +++ b/internal/consensus/common_test.go @@ -366,18 +366,29 @@ func validatePrecommit( } } -func subscribeToVoter(ctx context.Context, t *testing.T, cs *State, addr []byte) <-chan tmpubsub.Message { +// subscribeToVoter allows one to subscribe to vote messages from the validator +// associated with the given address. If supplied, only the specified vote +// types will be published via the return channel; if not supplied, all vote +// types will be published via the return channel. +func subscribeToVoter(ctx context.Context, t *testing.T, cs *State, addr []byte, voteTypes ...tmproto.SignedMsgType) <-chan tmpubsub.Message { t.Helper() + vt := make(map[tmproto.SignedMsgType]struct{}) + for _, t := range voteTypes { + vt[t] = struct{}{} + } + 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 if bytes.Equal(addr, vote.Vote.ValidatorAddress) { - select { - case <-ctx.Done(): - return ctx.Err() - case ch <- msg: + if _, wanted := vt[vote.Vote.Type]; len(vt) == 0 || wanted { + select { + case <-ctx.Done(): + return ctx.Err() + case ch <- msg: + } } } return nil diff --git a/internal/consensus/state_test.go b/internal/consensus/state_test.go index 03cd49a36..ac3f865c0 100644 --- a/internal/consensus/state_test.go +++ b/internal/consensus/state_test.go @@ -2211,7 +2211,7 @@ func TestPrepareProposalReceivesVoteExtensions(t *testing.T) { pv1, err := cs1.privValidator.GetPubKey(ctx) require.NoError(t, err) addr := pv1.Address() - voteCh := subscribeToVoter(ctx, t, cs1, addr) + voteCh := subscribeToVoter(ctx, t, cs1, addr, tmproto.PrecommitType) // ensure that the height is committed. ensurePrecommit(t, voteCh, height, round)