consensus: update tests

2026-02-10 14:00:33 +00:00 · 2021-03-16 14:05:32 -04:00
parent 2b60852e23
commit c4520b19c5
9 changed files with 105 additions and 63 deletions
--- a/consensus/byzantine_test.go
+++ b/consensus/byzantine_test.go
@@ -203,7 +203,7 @@ func TestByzantinePrevoteEquivocation(t *testing.T) {
 		propBlockID := types.BlockID{Hash: block.Hash(), PartSetHeader: blockParts.Header()}
 		proposal := types.NewProposal(height, round, lazyProposer.ValidRound, propBlockID)
 		p := proposal.ToProto()
-		if err := lazyProposer.privValidator.SignProposal(lazyProposer.state.ChainID, p); err == nil {
+		if err := lazyProposer.privValidator.SignProposal(context.Background(), lazyProposer.state.ChainID, p); err == nil {
 			proposal.Signature = p.Signature

 			// send proposal and block parts on internal msg queue
@@ -251,7 +251,7 @@ func TestByzantinePrevoteEquivocation(t *testing.T) {
 		close(done)
 	}()

-	pubkey, err := bcs.privValidator.GetPubKey()
+	pubkey, err := bcs.privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)

 	select {
@@ -426,7 +426,7 @@ func byzantineDecideProposalFunc(t *testing.T, height int64, round int32, cs *St
 	polRound, propBlockID := cs.ValidRound, types.BlockID{Hash: block1.Hash(), PartSetHeader: blockParts1.Header()}
 	proposal1 := types.NewProposal(height, round, polRound, propBlockID)
 	p1 := proposal1.ToProto()
-	if err := cs.privValidator.SignProposal(cs.state.ChainID, p1); err != nil {
+	if err := cs.privValidator.SignProposal(context.Background(), cs.state.ChainID, p1); err != nil {
 		t.Error(err)
 	}

@@ -440,7 +440,7 @@ func byzantineDecideProposalFunc(t *testing.T, height int64, round int32, cs *St
 	polRound, propBlockID = cs.ValidRound, types.BlockID{Hash: block2.Hash(), PartSetHeader: blockParts2.Header()}
 	proposal2 := types.NewProposal(height, round, polRound, propBlockID)
 	p2 := proposal2.ToProto()
-	if err := cs.privValidator.SignProposal(cs.state.ChainID, p2); err != nil {
+	if err := cs.privValidator.SignProposal(context.Background(), cs.state.ChainID, p2); err != nil {
 		t.Error(err)
 	}

--- a/consensus/common_test.go
+++ b/consensus/common_test.go
@@ -91,7 +91,7 @@ func (vs *validatorStub) signVote(
 	hash []byte,
 	header types.PartSetHeader) (*types.Vote, error) {

-	pubKey, err := vs.PrivValidator.GetPubKey()
+	pubKey, err := vs.PrivValidator.GetPubKey(context.Background())
 	if err != nil {
 		return nil, fmt.Errorf("can't get pubkey: %w", err)
 	}
@@ -106,7 +106,7 @@ func (vs *validatorStub) signVote(
 		BlockID:          types.BlockID{Hash: hash, PartSetHeader: header},
 	}
 	v := vote.ToProto()
-	err = vs.PrivValidator.SignVote(config.ChainID(), v)
+	err = vs.PrivValidator.SignVote(context.Background(), config.ChainID(), v)
 	vote.Signature = v.Signature

 	return vote, err
@@ -152,11 +152,11 @@ func (vss ValidatorStubsByPower) Len() int {
 }

 func (vss ValidatorStubsByPower) Less(i, j int) bool {
-	vssi, err := vss[i].GetPubKey()
+	vssi, err := vss[i].GetPubKey(context.Background())
 	if err != nil {
 		panic(err)
 	}
-	vssj, err := vss[j].GetPubKey()
+	vssj, err := vss[j].GetPubKey(context.Background())
 	if err != nil {
 		panic(err)
 	}
@@ -203,7 +203,7 @@ func decideProposal(
 	polRound, propBlockID := validRound, types.BlockID{Hash: block.Hash(), PartSetHeader: blockParts.Header()}
 	proposal = types.NewProposal(height, round, polRound, propBlockID)
 	p := proposal.ToProto()
-	if err := vs.SignProposal(chainID, p); err != nil {
+	if err := vs.SignProposal(context.Background(), chainID, p); err != nil {
 		panic(err)
 	}

@@ -231,7 +231,7 @@ func signAddVotes(

 func validatePrevote(t *testing.T, cs *State, round int32, privVal *validatorStub, blockHash []byte) {
 	prevotes := cs.Votes.Prevotes(round)
-	pubKey, err := privVal.GetPubKey()
+	pubKey, err := privVal.GetPubKey(context.Background())
 	require.NoError(t, err)
 	address := pubKey.Address()
 	var vote *types.Vote
@@ -251,7 +251,7 @@ func validatePrevote(t *testing.T, cs *State, round int32, privVal *validatorStu

 func validateLastPrecommit(t *testing.T, cs *State, privVal *validatorStub, blockHash []byte) {
 	votes := cs.LastCommit
-	pv, err := privVal.GetPubKey()
+	pv, err := privVal.GetPubKey(context.Background())
 	require.NoError(t, err)
 	address := pv.Address()
 	var vote *types.Vote
@@ -273,7 +273,7 @@ func validatePrecommit(
 	lockedBlockHash []byte,
 ) {
 	precommits := cs.Votes.Precommits(thisRound)
-	pv, err := privVal.GetPubKey()
+	pv, err := privVal.GetPubKey(context.Background())
 	require.NoError(t, err)
 	address := pv.Address()
 	var vote *types.Vote
--- a/consensus/invalid_test.go
+++ b/consensus/invalid_test.go
@@ -1,6 +1,7 @@
 package consensus

 import (
+	"context"
 	"testing"

 	"github.com/tendermint/tendermint/libs/bytes"
@@ -62,7 +63,7 @@ func invalidDoPrevoteFunc(t *testing.T, height int64, round int32, cs *State, sw
 	go func() {
 		cs.mtx.Lock()
 		cs.privValidator = pv
-		pubKey, err := cs.privValidator.GetPubKey()
+		pubKey, err := cs.privValidator.GetPubKey(context.Background())
 		if err != nil {
 			panic(err)
 		}
@@ -83,7 +84,7 @@ func invalidDoPrevoteFunc(t *testing.T, height int64, round int32, cs *State, sw
 				PartSetHeader: types.PartSetHeader{Total: 1, Hash: tmrand.Bytes(32)}},
 		}
 		p := precommit.ToProto()
-		err = cs.privValidator.SignVote(cs.state.ChainID, p)
+		err = cs.privValidator.SignVote(context.Background(), cs.state.ChainID, p)
 		if err != nil {
 			t.Error(err)
 		}
--- a/consensus/msgs_test.go
+++ b/consensus/msgs_test.go
@@ -1,6 +1,7 @@
 package consensus

 import (
+	"context"
 	"encoding/hex"
 	"math"
 	"testing"
@@ -58,7 +59,7 @@ func TestMsgToProto(t *testing.T) {
 	pbProposal := proposal.ToProto()

 	pv := types.NewMockPV()
-	pk, err := pv.GetPubKey()
+	pk, err := pv.GetPubKey(context.Background())
 	require.NoError(t, err)
 	val := types.NewValidator(pk, 100)

--- a/consensus/reactor_test.go
+++ b/consensus/reactor_test.go
@@ -316,7 +316,7 @@ func TestReactorVotingPowerChange(t *testing.T) {
 	// map of active validators
 	activeVals := make(map[string]struct{})
 	for i := 0; i < nVals; i++ {
-		pubKey, err := css[i].privValidator.GetPubKey()
+		pubKey, err := css[i].privValidator.GetPubKey(context.Background())
 		require.NoError(t, err)
 		addr := pubKey.Address()
 		activeVals[string(addr)] = struct{}{}
@@ -330,7 +330,7 @@ func TestReactorVotingPowerChange(t *testing.T) {
 	//---------------------------------------------------------------------------
 	logger.Debug("---------------------------- Testing changing the voting power of one validator a few times")

-	val1PubKey, err := css[0].privValidator.GetPubKey()
+	val1PubKey, err := css[0].privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)

 	val1PubKeyABCI, err := cryptoenc.PubKeyToProto(val1PubKey)
@@ -400,7 +400,7 @@ func TestReactorValidatorSetChanges(t *testing.T) {
 	// map of active validators
 	activeVals := make(map[string]struct{})
 	for i := 0; i < nVals; i++ {
-		pubKey, err := css[i].privValidator.GetPubKey()
+		pubKey, err := css[i].privValidator.GetPubKey(context.Background())
 		require.NoError(t, err)
 		activeVals[string(pubKey.Address())] = struct{}{}
 	}
@@ -413,7 +413,7 @@ func TestReactorValidatorSetChanges(t *testing.T) {
 	//---------------------------------------------------------------------------
 	logger.Info("---------------------------- Testing adding one validator")

-	newValidatorPubKey1, err := css[nVals].privValidator.GetPubKey()
+	newValidatorPubKey1, err := css[nVals].privValidator.GetPubKey(context.Background())
 	assert.NoError(t, err)
 	valPubKey1ABCI, err := cryptoenc.PubKeyToProto(newValidatorPubKey1)
 	assert.NoError(t, err)
@@ -442,7 +442,7 @@ func TestReactorValidatorSetChanges(t *testing.T) {
 	//---------------------------------------------------------------------------
 	logger.Info("---------------------------- Testing changing the voting power of one validator")

-	updateValidatorPubKey1, err := css[nVals].privValidator.GetPubKey()
+	updateValidatorPubKey1, err := css[nVals].privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	updatePubKey1ABCI, err := cryptoenc.PubKeyToProto(updateValidatorPubKey1)
 	require.NoError(t, err)
@@ -464,13 +464,13 @@ func TestReactorValidatorSetChanges(t *testing.T) {
 	//---------------------------------------------------------------------------
 	logger.Info("---------------------------- Testing adding two validators at once")

-	newValidatorPubKey2, err := css[nVals+1].privValidator.GetPubKey()
+	newValidatorPubKey2, err := css[nVals+1].privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	newVal2ABCI, err := cryptoenc.PubKeyToProto(newValidatorPubKey2)
 	require.NoError(t, err)
 	newValidatorTx2 := kvstore.MakeValSetChangeTx(newVal2ABCI, testMinPower)

-	newValidatorPubKey3, err := css[nVals+2].privValidator.GetPubKey()
+	newValidatorPubKey3, err := css[nVals+2].privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	newVal3ABCI, err := cryptoenc.PubKeyToProto(newValidatorPubKey3)
 	require.NoError(t, err)
--- a/consensus/replay_test.go
+++ b/consensus/replay_test.go
@@ -356,7 +356,7 @@ func TestSimulateValidatorsChange(t *testing.T) {
 	// HEIGHT 2
 	height++
 	incrementHeight(vss...)
-	newValidatorPubKey1, err := css[nVals].privValidator.GetPubKey()
+	newValidatorPubKey1, err := css[nVals].privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	valPubKey1ABCI, err := cryptoenc.PubKeyToProto(newValidatorPubKey1)
 	require.NoError(t, err)
@@ -369,7 +369,7 @@ func TestSimulateValidatorsChange(t *testing.T) {

 	proposal := types.NewProposal(vss[1].Height, round, -1, blockID)
 	p := proposal.ToProto()
-	if err := vss[1].SignProposal(config.ChainID(), p); err != nil {
+	if err := vss[1].SignProposal(context.Background(), config.ChainID(), p); err != nil {
 		t.Fatal("failed to sign bad proposal", err)
 	}
 	proposal.Signature = p.Signature
@@ -386,7 +386,7 @@ func TestSimulateValidatorsChange(t *testing.T) {
 	// HEIGHT 3
 	height++
 	incrementHeight(vss...)
-	updateValidatorPubKey1, err := css[nVals].privValidator.GetPubKey()
+	updateValidatorPubKey1, err := css[nVals].privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	updatePubKey1ABCI, err := cryptoenc.PubKeyToProto(updateValidatorPubKey1)
 	require.NoError(t, err)
@@ -399,7 +399,7 @@ func TestSimulateValidatorsChange(t *testing.T) {

 	proposal = types.NewProposal(vss[2].Height, round, -1, blockID)
 	p = proposal.ToProto()
-	if err := vss[2].SignProposal(config.ChainID(), p); err != nil {
+	if err := vss[2].SignProposal(context.Background(), config.ChainID(), p); err != nil {
 		t.Fatal("failed to sign bad proposal", err)
 	}
 	proposal.Signature = p.Signature
@@ -416,14 +416,14 @@ func TestSimulateValidatorsChange(t *testing.T) {
 	// HEIGHT 4
 	height++
 	incrementHeight(vss...)
-	newValidatorPubKey2, err := css[nVals+1].privValidator.GetPubKey()
+	newValidatorPubKey2, err := css[nVals+1].privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	newVal2ABCI, err := cryptoenc.PubKeyToProto(newValidatorPubKey2)
 	require.NoError(t, err)
 	newValidatorTx2 := kvstore.MakeValSetChangeTx(newVal2ABCI, testMinPower)
 	err = assertMempool(css[0].txNotifier).CheckTx(newValidatorTx2, nil, mempl.TxInfo{})
 	assert.Nil(t, err)
-	newValidatorPubKey3, err := css[nVals+2].privValidator.GetPubKey()
+	newValidatorPubKey3, err := css[nVals+2].privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	newVal3ABCI, err := cryptoenc.PubKeyToProto(newValidatorPubKey3)
 	require.NoError(t, err)
@@ -439,10 +439,10 @@ func TestSimulateValidatorsChange(t *testing.T) {

 	valIndexFn := func(cssIdx int) int {
 		for i, vs := range newVss {
-			vsPubKey, err := vs.GetPubKey()
+			vsPubKey, err := vs.GetPubKey(context.Background())
 			require.NoError(t, err)

-			cssPubKey, err := css[cssIdx].privValidator.GetPubKey()
+			cssPubKey, err := css[cssIdx].privValidator.GetPubKey(context.Background())
 			require.NoError(t, err)

 			if vsPubKey.Equals(cssPubKey) {
@@ -456,7 +456,7 @@ func TestSimulateValidatorsChange(t *testing.T) {

 	proposal = types.NewProposal(vss[3].Height, round, -1, blockID)
 	p = proposal.ToProto()
-	if err := vss[3].SignProposal(config.ChainID(), p); err != nil {
+	if err := vss[3].SignProposal(context.Background(), config.ChainID(), p); err != nil {
 		t.Fatal("failed to sign bad proposal", err)
 	}
 	proposal.Signature = p.Signature
@@ -515,7 +515,7 @@ func TestSimulateValidatorsChange(t *testing.T) {
 	selfIndex = valIndexFn(0)
 	proposal = types.NewProposal(vss[1].Height, round, -1, blockID)
 	p = proposal.ToProto()
-	if err := vss[1].SignProposal(config.ChainID(), p); err != nil {
+	if err := vss[1].SignProposal(context.Background(), config.ChainID(), p); err != nil {
 		t.Fatal("failed to sign bad proposal", err)
 	}
 	proposal.Signature = p.Signature
@@ -675,7 +675,8 @@ func testHandshakeReplay(t *testing.T, config *cfg.Config, nBlocks int, mode uin
 		walFile := tempWALWithData(walBody)
 		config.Consensus.SetWalFile(walFile)

-		privVal := privval.LoadFilePV(config.PrivValidatorKeyFile(), config.PrivValidatorStateFile())
+		privVal, err := privval.LoadFilePV(config.PrivValidatorKeyFile(), config.PrivValidatorStateFile())
+		require.NoError(t, err)

 		wal, err := NewWAL(walFile)
 		require.NoError(t, err)
@@ -689,7 +690,7 @@ func testHandshakeReplay(t *testing.T, config *cfg.Config, nBlocks int, mode uin
 		})
 		chain, commits, err = makeBlockchainFromWAL(wal)
 		require.NoError(t, err)
-		pubKey, err := privVal.GetPubKey()
+		pubKey, err := privVal.GetPubKey(context.Background())
 		require.NoError(t, err)
 		stateDB, genesisState, store = stateAndStore(config, pubKey, kvstore.ProtocolVersion)

@@ -887,9 +888,12 @@ func TestHandshakePanicsIfAppReturnsWrongAppHash(t *testing.T) {
 	//		- 0x03
 	config := ResetConfig("handshake_test_")
 	t.Cleanup(func() { os.RemoveAll(config.RootDir) })
-	privVal := privval.LoadFilePV(config.PrivValidatorKeyFile(), config.PrivValidatorStateFile())
+
+	privVal, err := privval.LoadFilePV(config.PrivValidatorKeyFile(), config.PrivValidatorStateFile())
+	require.NoError(t, err)
+
 	const appVersion = 0x0
-	pubKey, err := privVal.GetPubKey()
+	pubKey, err := privVal.GetPubKey(context.Background())
 	require.NoError(t, err)
 	stateDB, state, store := stateAndStore(config, pubKey, appVersion)
 	stateStore := sm.NewStore(stateDB)
@@ -1223,8 +1227,10 @@ func TestHandshakeUpdatesValidators(t *testing.T) {
 	config := ResetConfig("handshake_test_")
 	t.Cleanup(func() { _ = os.RemoveAll(config.RootDir) })

-	privVal := privval.LoadFilePV(config.PrivValidatorKeyFile(), config.PrivValidatorStateFile())
-	pubKey, err := privVal.GetPubKey()
+	privVal, err := privval.LoadFilePV(config.PrivValidatorKeyFile(), config.PrivValidatorStateFile())
+	require.NoError(t, err)
+
+	pubKey, err := privVal.GetPubKey(context.Background())
 	require.NoError(t, err)
 	stateDB, state, store := stateAndStore(config, pubKey, 0x0)
 	stateStore := sm.NewStore(stateDB)
--- a/consensus/state.go
+++ b/consensus/state.go
@@ -2,6 +2,7 @@ package consensus

 import (
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"io/ioutil"
@@ -1083,7 +1084,11 @@ func (cs *State) defaultDecideProposal(height int64, round int32) {
 	propBlockID := types.BlockID{Hash: block.Hash(), PartSetHeader: blockParts.Header()}
 	proposal := types.NewProposal(height, round, cs.ValidRound, propBlockID)
 	p := proposal.ToProto()
-	if err := cs.privValidator.SignProposal(cs.state.ChainID, p); err == nil {
+
+	ctx, cancel := context.WithTimeout(context.TODO(), cs.config.TimeoutPropose)
+	defer cancel()
+
+	if err := cs.privValidator.SignProposal(ctx, cs.state.ChainID, p); err == nil {
 		proposal.Signature = p.Signature

 		// send proposal and block parts on internal msg queue
@@ -2086,7 +2091,24 @@ func (cs *State) signVote(
 		BlockID:          types.BlockID{Hash: hash, PartSetHeader: header},
 	}
 	v := vote.ToProto()
-	err := cs.privValidator.SignVote(cs.state.ChainID, v)
+
+	// If the signedMessageType is for precommit, use our local precommit Timeout
+	// as the max wait time for getting a singed commit. The same goes for prevote.
+	var timeout time.Duration
+
+	switch msgType {
+	case tmproto.PrecommitType:
+		timeout = cs.config.TimeoutPrecommit
+	case tmproto.PrevoteType:
+		timeout = cs.config.TimeoutPrevote
+	default:
+		timeout = time.Second
+	}
+
+	ctx, cancel := context.WithTimeout(context.TODO(), timeout)
+	defer cancel()
+
+	err := cs.privValidator.SignVote(ctx, cs.state.ChainID, v)
 	vote.Signature = v.Signature

 	return vote, err
@@ -2154,10 +2176,21 @@ func (cs *State) updatePrivValidatorPubKey() error {
 		return nil
 	}

-	pubKey, err := cs.privValidator.GetPubKey()
+	var timeout time.Duration
+	if cs.config.TimeoutPrecommit > cs.config.TimeoutPrevote {
+		timeout = cs.config.TimeoutPrecommit
+	} else {
+		timeout = cs.config.TimeoutPrevote
+	}
+
+	ctx, cancel := context.WithTimeout(context.TODO(), timeout)
+	defer cancel()
+
+	pubKey, err := cs.privValidator.GetPubKey(ctx)
 	if err != nil {
 		return err
 	}
+
 	cs.privValidatorPubKey = pubKey
 	return nil
 }
--- a/consensus/state_test.go
+++ b/consensus/state_test.go
@@ -69,7 +69,7 @@ func TestStateProposerSelection0(t *testing.T) {

 	// Commit a block and ensure proposer for the next height is correct.
 	prop := cs1.GetRoundState().Validators.GetProposer()
-	pv, err := cs1.privValidator.GetPubKey()
+	pv, err := cs1.privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	address := pv.Address()
 	if !bytes.Equal(prop.Address, address) {
@@ -86,7 +86,7 @@ func TestStateProposerSelection0(t *testing.T) {
 	ensureNewRound(newRoundCh, height+1, 0)

 	prop = cs1.GetRoundState().Validators.GetProposer()
-	pv1, err := vss[1].GetPubKey()
+	pv1, err := vss[1].GetPubKey(context.Background())
 	require.NoError(t, err)
 	addr := pv1.Address()
 	if !bytes.Equal(prop.Address, addr) {
@@ -112,7 +112,7 @@ func TestStateProposerSelection2(t *testing.T) {
 	// everyone just votes nil. we get a new proposer each round
 	for i := int32(0); int(i) < len(vss); i++ {
 		prop := cs1.GetRoundState().Validators.GetProposer()
-		pvk, err := vss[int(i+round)%len(vss)].GetPubKey()
+		pvk, err := vss[int(i+round)%len(vss)].GetPubKey(context.Background())
 		require.NoError(t, err)
 		addr := pvk.Address()
 		correctProposer := addr
@@ -208,7 +208,7 @@ func TestStateBadProposal(t *testing.T) {
 	blockID := types.BlockID{Hash: propBlock.Hash(), PartSetHeader: propBlockParts.Header()}
 	proposal := types.NewProposal(vs2.Height, round, -1, blockID)
 	p := proposal.ToProto()
-	if err := vs2.SignProposal(config.ChainID(), p); err != nil {
+	if err := vs2.SignProposal(context.Background(), config.ChainID(), p); err != nil {
 		t.Fatal("failed to sign bad proposal", err)
 	}

@@ -262,7 +262,7 @@ func TestStateOversizedBlock(t *testing.T) {
 	blockID := types.BlockID{Hash: propBlock.Hash(), PartSetHeader: propBlockParts.Header()}
 	proposal := types.NewProposal(height, round, -1, blockID)
 	p := proposal.ToProto()
-	if err := vs2.SignProposal(config.ChainID(), p); err != nil {
+	if err := vs2.SignProposal(context.Background(), config.ChainID(), p); err != nil {
 		t.Fatal("failed to sign bad proposal", err)
 	}
 	proposal.Signature = p.Signature
@@ -594,7 +594,7 @@ func TestStateLockPOLRelock(t *testing.T) {

 	timeoutWaitCh := subscribe(cs1.eventBus, types.EventQueryTimeoutWait)
 	proposalCh := subscribe(cs1.eventBus, types.EventQueryCompleteProposal)
-	pv1, err := cs1.privValidator.GetPubKey()
+	pv1, err := cs1.privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	addr := pv1.Address()
 	voteCh := subscribeToVoter(cs1, addr)
@@ -693,7 +693,7 @@ func TestStateLockPOLUnlock(t *testing.T) {
 	timeoutWaitCh := subscribe(cs1.eventBus, types.EventQueryTimeoutWait)
 	newRoundCh := subscribe(cs1.eventBus, types.EventQueryNewRound)
 	unlockCh := subscribe(cs1.eventBus, types.EventQueryUnlock)
-	pv1, err := cs1.privValidator.GetPubKey()
+	pv1, err := cs1.privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	addr := pv1.Address()
 	voteCh := subscribeToVoter(cs1, addr)
@@ -783,7 +783,7 @@ func TestStateLockPOLUnlockOnUnknownBlock(t *testing.T) {

 	timeoutWaitCh := subscribe(cs1.eventBus, types.EventQueryTimeoutWait)
 	proposalCh := subscribe(cs1.eventBus, types.EventQueryCompleteProposal)
-	pv1, err := cs1.privValidator.GetPubKey()
+	pv1, err := cs1.privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	addr := pv1.Address()
 	voteCh := subscribeToVoter(cs1, addr)
@@ -911,7 +911,7 @@ func TestStateLockPOLSafety1(t *testing.T) {
 	timeoutProposeCh := subscribe(cs1.eventBus, types.EventQueryTimeoutPropose)
 	timeoutWaitCh := subscribe(cs1.eventBus, types.EventQueryTimeoutWait)
 	newRoundCh := subscribe(cs1.eventBus, types.EventQueryNewRound)
-	pv1, err := cs1.privValidator.GetPubKey()
+	pv1, err := cs1.privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	addr := pv1.Address()
 	voteCh := subscribeToVoter(cs1, addr)
@@ -1030,7 +1030,7 @@ func TestStateLockPOLSafety2(t *testing.T) {
 	timeoutWaitCh := subscribe(cs1.eventBus, types.EventQueryTimeoutWait)
 	newRoundCh := subscribe(cs1.eventBus, types.EventQueryNewRound)
 	unlockCh := subscribe(cs1.eventBus, types.EventQueryUnlock)
-	pv1, err := cs1.privValidator.GetPubKey()
+	pv1, err := cs1.privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	addr := pv1.Address()
 	voteCh := subscribeToVoter(cs1, addr)
@@ -1085,7 +1085,7 @@ func TestStateLockPOLSafety2(t *testing.T) {
 	// in round 2 we see the polkad block from round 0
 	newProp := types.NewProposal(height, round, 0, propBlockID0)
 	p := newProp.ToProto()
-	if err := vs3.SignProposal(config.ChainID(), p); err != nil {
+	if err := vs3.SignProposal(context.Background(), config.ChainID(), p); err != nil {
 		t.Fatal(err)
 	}

@@ -1128,7 +1128,7 @@ func TestProposeValidBlock(t *testing.T) {
 	timeoutProposeCh := subscribe(cs1.eventBus, types.EventQueryTimeoutPropose)
 	newRoundCh := subscribe(cs1.eventBus, types.EventQueryNewRound)
 	unlockCh := subscribe(cs1.eventBus, types.EventQueryUnlock)
-	pv1, err := cs1.privValidator.GetPubKey()
+	pv1, err := cs1.privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	addr := pv1.Address()
 	voteCh := subscribeToVoter(cs1, addr)
@@ -1217,7 +1217,7 @@ func TestSetValidBlockOnDelayedPrevote(t *testing.T) {
 	timeoutWaitCh := subscribe(cs1.eventBus, types.EventQueryTimeoutWait)
 	newRoundCh := subscribe(cs1.eventBus, types.EventQueryNewRound)
 	validBlockCh := subscribe(cs1.eventBus, types.EventQueryValidBlock)
-	pv1, err := cs1.privValidator.GetPubKey()
+	pv1, err := cs1.privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	addr := pv1.Address()
 	voteCh := subscribeToVoter(cs1, addr)
@@ -1279,7 +1279,7 @@ func TestSetValidBlockOnDelayedProposal(t *testing.T) {
 	timeoutProposeCh := subscribe(cs1.eventBus, types.EventQueryTimeoutPropose)
 	newRoundCh := subscribe(cs1.eventBus, types.EventQueryNewRound)
 	validBlockCh := subscribe(cs1.eventBus, types.EventQueryValidBlock)
-	pv1, err := cs1.privValidator.GetPubKey()
+	pv1, err := cs1.privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	addr := pv1.Address()
 	voteCh := subscribeToVoter(cs1, addr)
@@ -1352,7 +1352,7 @@ func TestWaitingTimeoutProposeOnNewRound(t *testing.T) {

 	timeoutWaitCh := subscribe(cs1.eventBus, types.EventQueryTimeoutPropose)
 	newRoundCh := subscribe(cs1.eventBus, types.EventQueryNewRound)
-	pv1, err := cs1.privValidator.GetPubKey()
+	pv1, err := cs1.privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	addr := pv1.Address()
 	voteCh := subscribeToVoter(cs1, addr)
@@ -1388,7 +1388,7 @@ func TestRoundSkipOnNilPolkaFromHigherRound(t *testing.T) {

 	timeoutWaitCh := subscribe(cs1.eventBus, types.EventQueryTimeoutWait)
 	newRoundCh := subscribe(cs1.eventBus, types.EventQueryNewRound)
-	pv1, err := cs1.privValidator.GetPubKey()
+	pv1, err := cs1.privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	addr := pv1.Address()
 	voteCh := subscribeToVoter(cs1, addr)
@@ -1424,7 +1424,7 @@ func TestWaitTimeoutProposeOnNilPolkaForTheCurrentRound(t *testing.T) {

 	timeoutProposeCh := subscribe(cs1.eventBus, types.EventQueryTimeoutPropose)
 	newRoundCh := subscribe(cs1.eventBus, types.EventQueryNewRound)
-	pv1, err := cs1.privValidator.GetPubKey()
+	pv1, err := cs1.privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	addr := pv1.Address()
 	voteCh := subscribeToVoter(cs1, addr)
@@ -1545,7 +1545,7 @@ func TestStartNextHeightCorrectlyAfterTimeout(t *testing.T) {

 	newRoundCh := subscribe(cs1.eventBus, types.EventQueryNewRound)
 	newBlockHeader := subscribe(cs1.eventBus, types.EventQueryNewBlockHeader)
-	pv1, err := cs1.privValidator.GetPubKey()
+	pv1, err := cs1.privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	addr := pv1.Address()
 	voteCh := subscribeToVoter(cs1, addr)
@@ -1605,7 +1605,7 @@ func TestResetTimeoutPrecommitUponNewHeight(t *testing.T) {

 	newRoundCh := subscribe(cs1.eventBus, types.EventQueryNewRound)
 	newBlockHeader := subscribe(cs1.eventBus, types.EventQueryNewBlockHeader)
-	pv1, err := cs1.privValidator.GetPubKey()
+	pv1, err := cs1.privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	addr := pv1.Address()
 	voteCh := subscribeToVoter(cs1, addr)
@@ -1744,7 +1744,7 @@ func TestStateHalt1(t *testing.T) {
 	timeoutWaitCh := subscribe(cs1.eventBus, types.EventQueryTimeoutWait)
 	newRoundCh := subscribe(cs1.eventBus, types.EventQueryNewRound)
 	newBlockCh := subscribe(cs1.eventBus, types.EventQueryNewBlock)
-	pv1, err := cs1.privValidator.GetPubKey()
+	pv1, err := cs1.privValidator.GetPubKey(context.Background())
 	require.NoError(t, err)
 	addr := pv1.Address()
 	voteCh := subscribeToVoter(cs1, addr)
--- a/consensus/types/height_vote_set_test.go
+++ b/consensus/types/height_vote_set_test.go
@@ -1,6 +1,7 @@
 package types

 import (
+	"context"
 	"fmt"
 	"os"
 	"testing"
@@ -57,7 +58,7 @@ func TestPeerCatchupRounds(t *testing.T) {

 func makeVoteHR(t *testing.T, height int64, valIndex, round int32, privVals []types.PrivValidator) *types.Vote {
 	privVal := privVals[valIndex]
-	pubKey, err := privVal.GetPubKey()
+	pubKey, err := privVal.GetPubKey(context.Background())
 	if err != nil {
 		panic(err)
 	}
@@ -76,7 +77,7 @@ func makeVoteHR(t *testing.T, height int64, valIndex, round int32, privVals []ty
 	chainID := config.ChainID()

 	v := vote.ToProto()
-	err = privVal.SignVote(chainID, v)
+	err = privVal.SignVote(context.Background(), chainID, v)
 	if err != nil {
 		panic(fmt.Sprintf("Error signing vote: %v", err))
 	}