test: prohibit dkg (#489)

* Allow to prohibit DKG changes at governance layer.
* Allow test.App to broadcast RoundEvent

10 files changed, 301 insertions, 21 deletions

diff --git a/core/blockchain_test.go b/core/blockchain_test.go
index c5023ff..1996fbb 100644
--- a/core/blockchain_test.go
+++ b/core/blockchain_test.go
@@ -144,7 +144,7 @@ func (s *BlockChainTestSuite) newBlockChain(initB *types.Block,
 		initConfig.SetRoundBeginHeight(0)
 	}
 	return newBlockChain(s.nID, s.dMoment, initB, initConfig,
-		test.NewApp(0, nil), &testTSigVerifierGetter{}, s.signer,
+		test.NewApp(0, nil, nil), &testTSigVerifierGetter{}, s.signer,
 		&common.NullLogger{})
 }
 
diff --git a/core/consensus_test.go b/core/consensus_test.go
index d721bd8..8ebfb02 100644
--- a/core/consensus_test.go
+++ b/core/consensus_test.go
@@ -192,7 +192,7 @@ func (s *ConsensusTestSuite) prepareConsensus(
 	conn *networkConnection) (
 	*test.App, *Consensus) {
 
-	app := test.NewApp(0, nil)
+	app := test.NewApp(0, nil, nil)
 	dbInst, err := db.NewMemBackedDB()
 	s.Require().NoError(err)
 	nID := types.NewNodeID(prvKey.PublicKey())
@@ -211,7 +211,7 @@ func (s *ConsensusTestSuite) prepareConsensusWithDB(
 	dbInst db.Database) (
 	*test.App, *Consensus) {
 
-	app := test.NewApp(0, nil)
+	app := test.NewApp(0, nil, nil)
 	nID := types.NewNodeID(prvKey.PublicKey())
 	network := conn.newNetwork(nID)
 	con := NewConsensus(
diff --git a/core/test/app.go b/core/test/app.go
index 516974c..0a17f13 100644
--- a/core/test/app.go
+++ b/core/test/app.go
@@ -24,6 +24,7 @@ import (
 
 	"github.com/dexon-foundation/dexon-consensus/common"
 	"github.com/dexon-foundation/dexon-consensus/core/types"
+	"github.com/dexon-foundation/dexon-consensus/core/utils"
 )
 
 var (
@@ -84,23 +85,40 @@ type App struct {
 	deliveredLock         sync.RWMutex
 	state                 *State
 	gov                   *Governance
+	rEvt                  *utils.RoundEvent
+	hEvt                  *common.Event
 	lastPendingHeightLock sync.RWMutex
 	LastPendingHeight     uint64
 	roundToNotify         uint64
 }
 
 // NewApp constructs a TestApp instance.
-func NewApp(initRound uint64, gov *Governance) (app *App) {
+func NewApp(initRound uint64, gov *Governance, rEvt *utils.RoundEvent) (
+	app *App) {
 	app = &App{
 		Confirmed:       make(map[common.Hash]*types.Block),
 		Delivered:       make(map[common.Hash]*AppDeliveredRecord),
 		DeliverSequence: common.Hashes{},
 		gov:             gov,
+		rEvt:            rEvt,
+		hEvt:            common.NewEvent(),
 		roundToNotify:   initRound,
 	}
 	if gov != nil {
 		app.state = gov.State()
 	}
+	if rEvt != nil {
+		app.hEvt.RegisterHeight(
+			utils.GetNextRoundCheckpoint(rEvt.LastPeriod()), func(h uint64) {
+				rEvt.ValidateNextRound(h)
+			})
+		rEvt.Register(func(evts []utils.RoundEventParam) {
+			app.hEvt.RegisterHeight(evts[len(evts)-1].NextRoundCheckpoint(),
+				func(h uint64) {
+					rEvt.ValidateNextRound(h)
+				})
+		})
+	}
 	return app
 }
 
@@ -191,8 +209,8 @@ func (app *App) BlockConfirmed(b types.Block) {
 }
 
 // BlockDelivered implements Application interface.
-func (app *App) BlockDelivered(
-	blockHash common.Hash, pos types.Position, result types.FinalizationResult) {
+func (app *App) BlockDelivered(blockHash common.Hash, pos types.Position,
+	result types.FinalizationResult) {
 	func() {
 		app.deliveredLock.Lock()
 		defer app.deliveredLock.Unlock()
@@ -236,6 +254,7 @@ func (app *App) BlockDelivered(
 			}
 		}
 	}()
+	app.hEvt.NotifyHeight(result.Height)
 }
 
 // GetLatestDeliveredPosition would return the latest position of delivered
diff --git a/core/test/app_test.go b/core/test/app_test.go
index 79518ea..7ac9ef9 100644
--- a/core/test/app_test.go
+++ b/core/test/app_test.go
@@ -19,16 +19,82 @@ package test
 
 import (
 	"bytes"
+	"context"
+	"fmt"
 	"testing"
 	"time"
 
 	"github.com/dexon-foundation/dexon-consensus/common"
+	"github.com/dexon-foundation/dexon-consensus/core"
+	"github.com/dexon-foundation/dexon-consensus/core/crypto"
+	"github.com/dexon-foundation/dexon-consensus/core/crypto/dkg"
 	"github.com/dexon-foundation/dexon-consensus/core/types"
+	typesDKG "github.com/dexon-foundation/dexon-consensus/core/types/dkg"
+	"github.com/dexon-foundation/dexon-consensus/core/utils"
 	"github.com/stretchr/testify/suite"
 )
 
+func getCRS(round, reset uint64) []byte {
+	return []byte(fmt.Sprintf("r#%d,reset#%d", round, reset))
+}
+
+type evtParamToCheck struct {
+	round  uint64
+	reset  uint64
+	height uint64
+	crs    common.Hash
+}
+
 type AppTestSuite struct {
 	suite.Suite
+
+	pubKeys []crypto.PublicKey
+	signers []*utils.Signer
+	logger  common.Logger
+}
+
+func (s *AppTestSuite) SetupSuite() {
+	prvKeys, pubKeys, err := NewKeys(4)
+	s.Require().NoError(err)
+	s.pubKeys = pubKeys
+	for _, k := range prvKeys {
+		s.signers = append(s.signers, utils.NewSigner(k))
+	}
+	s.logger = &common.NullLogger{}
+}
+
+func (s *AppTestSuite) prepareGov() *Governance {
+	gov, err := NewGovernance(
+		NewState(1, s.pubKeys, 100*time.Millisecond, s.logger, true),
+		core.ConfigRoundShift)
+	s.Require().NoError(err)
+	return gov
+}
+
+func (s *AppTestSuite) proposeMPK(gov *Governance, round uint64, count int) {
+	for idx, pubKey := range s.pubKeys[:count] {
+		_, pubShare := dkg.NewPrivateKeyShares(utils.GetDKGThreshold(
+			gov.Configuration(round)))
+		mpk := &typesDKG.MasterPublicKey{
+			Round:           round,
+			DKGID:           typesDKG.NewID(types.NewNodeID(pubKey)),
+			PublicKeyShares: *pubShare,
+		}
+		s.Require().NoError(s.signers[idx].SignDKGMasterPublicKey(mpk))
+		gov.AddDKGMasterPublicKey(round, mpk)
+	}
+}
+
+func (s *AppTestSuite) proposeFinalize(gov *Governance, round uint64,
+	count int) {
+	for idx, pubKey := range s.pubKeys[:count] {
+		final := &typesDKG.Finalize{
+			ProposerID: types.NewNodeID(pubKey),
+			Round:      round,
+		}
+		s.Require().NoError(s.signers[idx].SignDKGFinalize(final))
+		gov.AddDKGFinalize(round, final)
+	}
 }
 
 func (s *AppTestSuite) deliverBlockWithTimeFromSequenceLength(
@@ -58,7 +124,7 @@ func (s *AppTestSuite) TestCompare() {
 		}
 	)
 	// Prepare an OK App instance.
-	app1 := NewApp(0, nil)
+	app1 := NewApp(0, nil, nil)
 	app1.BlockConfirmed(b0)
 	app1.BlockConfirmed(b1)
 	app1.BlockDelivered(b0.Hash, b0.Position, types.FinalizationResult{
@@ -69,7 +135,7 @@ func (s *AppTestSuite) TestCompare() {
 		Height:    2,
 		Timestamp: now.Add(1 * time.Second),
 	})
-	app2 := NewApp(0, nil)
+	app2 := NewApp(0, nil, nil)
 	s.Require().Equal(ErrEmptyDeliverSequence.Error(),
 		app1.Compare(app2).Error())
 	app2.BlockConfirmed(b0)
@@ -88,7 +154,7 @@ func (s *AppTestSuite) TestCompare() {
 	})
 	s.Require().Equal(ErrMismatchBlockHashSequence.Error(),
 		app1.Compare(app2).Error())
-	app2 = NewApp(0, nil)
+	app2 = NewApp(0, nil, nil)
 	app2.BlockConfirmed(b0)
 	app2.BlockDelivered(b0.Hash, b0.Position, types.FinalizationResult{
 		Height:    1,
@@ -107,13 +173,13 @@ func (s *AppTestSuite) TestVerify() {
 			Position: types.Position{Height: 1},
 		}
 	)
-	app := NewApp(0, nil)
+	app := NewApp(0, nil, nil)
 	s.Require().Equal(ErrEmptyDeliverSequence.Error(), app.Verify().Error())
 	app.BlockDelivered(b0.Hash, b0.Position, types.FinalizationResult{})
 	app.BlockDelivered(b1.Hash, b1.Position, types.FinalizationResult{Height: 1})
 	s.Require().Equal(
 		ErrDeliveredBlockNotConfirmed.Error(), app.Verify().Error())
-	app = NewApp(0, nil)
+	app = NewApp(0, nil, nil)
 	app.BlockConfirmed(b0)
 	app.BlockDelivered(b0.Hash, b0.Position, types.FinalizationResult{
 		Height:    1,
@@ -126,7 +192,7 @@ func (s *AppTestSuite) TestVerify() {
 	})
 	s.Require().Equal(ErrConsensusTimestampOutOfOrder.Error(),
 		app.Verify().Error())
-	app = NewApp(0, nil)
+	app = NewApp(0, nil, nil)
 	app.BlockConfirmed(b0)
 	app.BlockConfirmed(b1)
 	app.BlockDelivered(b0.Hash, b0.Position, types.FinalizationResult{
@@ -141,7 +207,7 @@ func (s *AppTestSuite) TestVerify() {
 
 func (s *AppTestSuite) TestWitness() {
 	// Deliver several blocks, there is only one chain only.
-	app := NewApp(0, nil)
+	app := NewApp(0, nil, nil)
 	deliver := func(b *types.Block) {
 		app.BlockConfirmed(*b)
 		app.BlockDelivered(b.Hash, b.Position, b.Finalization)
@@ -205,6 +271,78 @@ func (s *AppTestSuite) TestWitness() {
 	s.Require().Equal(0, bytes.Compare(w.Data, b02.Hash[:]))
 }
 
+func (s *AppTestSuite) TestAttachedWithRoundEvent() {
+	// This test case is copied/modified from
+	// integraion.RoundEventTestSuite.TestFromRoundN, the difference is the
+	// calls to utils.RoundEvent.ValidateNextRound is not explicitly called but
+	// triggered by App.BlockDelivered.
+	gov := s.prepareGov()
+	s.Require().NoError(gov.State().RequestChange(StateChangeRoundLength,
+		uint64(100)))
+	gov.CatchUpWithRound(22)
+	for r := uint64(2); r <= uint64(20); r++ {
+		gov.ProposeCRS(r, getCRS(r, 0))
+	}
+	// Reset round#20 twice, then make it done DKG preparation.
+	gov.ResetDKG(getCRS(20, 1))
+	gov.ResetDKG(getCRS(20, 2))
+	s.proposeMPK(gov, 20, 3)
+	s.proposeFinalize(gov, 20, 3)
+	s.Require().Equal(gov.DKGResetCount(20), uint64(2))
+	// Propose CRS for round#21, and it works without reset.
+	gov.ProposeCRS(21, getCRS(21, 0))
+	s.proposeMPK(gov, 21, 3)
+	s.proposeFinalize(gov, 21, 3)
+	// Propose CRS for round#22, and it works without reset.
+	gov.ProposeCRS(22, getCRS(22, 0))
+	s.proposeMPK(gov, 22, 3)
+	s.proposeFinalize(gov, 22, 3)
+	// Prepare utils.RoundEvent, starts from round#19, reset(for round#20)#1.
+	rEvt, err := utils.NewRoundEvent(context.Background(), gov, s.logger, 19,
+		1900, 2019, core.ConfigRoundShift)
+	s.Require().NoError(err)
+	// Register a handler to collects triggered events.
+	var evts []evtParamToCheck
+	rEvt.Register(func(params []utils.RoundEventParam) {
+		for _, p := range params {
+			evts = append(evts, evtParamToCheck{
+				round:  p.Round,
+				reset:  p.Reset,
+				height: p.BeginHeight,
+				crs:    p.CRS,
+			})
+		}
+	})
+	// Setup App instance.
+	app := NewApp(19, gov, rEvt)
+	deliver := func(round, start, end uint64) {
+		for i := start; i <= end; i++ {
+			b := &types.Block{
+				Hash:         common.NewRandomHash(),
+				Position:     types.Position{Round: round, Height: i},
+				Finalization: types.FinalizationResult{Height: i},
+			}
+			app.BlockConfirmed(*b)
+			app.BlockDelivered(b.Hash, b.Position, b.Finalization)
+		}
+	}
+	// Deliver blocks from height=2020 to height=2081.
+	deliver(0, 0, 2019)
+	deliver(19, 2020, 2091)
+	s.Require().Len(evts, 2)
+	s.Require().Equal(evts[0], evtParamToCheck{19, 2, 2100, gov.CRS(19)})
+	s.Require().Equal(evts[1], evtParamToCheck{20, 0, 2200, gov.CRS(20)})
+	// Deliver blocks from height=2082 to height=2281.
+	deliver(19, 2092, 2199)
+	deliver(20, 2200, 2291)
+	s.Require().Len(evts, 3)
+	s.Require().Equal(evts[2], evtParamToCheck{21, 0, 2300, gov.CRS(21)})
+	// Deliver blocks from height=2282 to height=2381.
+	deliver(20, 2292, 2299)
+	deliver(21, 2300, 2391)
+	s.Require().Equal(evts[3], evtParamToCheck{22, 0, 2400, gov.CRS(22)})
+}
+
 func TestApp(t *testing.T) {
 	suite.Run(t, new(AppTestSuite))
 }
diff --git a/core/test/governance.go b/core/test/governance.go
index 937c953..74eca85 100644
--- a/core/test/governance.go
+++ b/core/test/governance.go
@@ -43,6 +43,7 @@ type Governance struct {
 	stateModule          *State
 	networkModule        *Network
 	pendingConfigChanges map[uint64]map[StateChangeType]interface{}
+	prohibitedTypes      map[StateChangeType]struct{}
 	lock                 sync.RWMutex
 }
 
@@ -53,6 +54,7 @@ func NewGovernance(state *State, roundShift uint64) (g *Governance, err error) {
 		roundShift:           roundShift,
 		pendingConfigChanges: make(map[uint64]map[StateChangeType]interface{}),
 		stateModule:          state,
+		prohibitedTypes:      make(map[StateChangeType]struct{}),
 	}
 	return
 }
@@ -140,6 +142,9 @@ func (g *Governance) ProposeCRS(round uint64, signedCRS []byte) {
 // AddDKGComplaint add a DKGComplaint.
 func (g *Governance) AddDKGComplaint(
 	round uint64, complaint *typesDKG.Complaint) {
+	if g.isProhibited(StateAddDKGComplaint) {
+		return
+	}
 	if round != complaint.Round {
 		return
 	}
@@ -161,6 +166,9 @@ func (g *Governance) DKGComplaints(round uint64) []*typesDKG.Complaint {
 // AddDKGMasterPublicKey adds a DKGMasterPublicKey.
 func (g *Governance) AddDKGMasterPublicKey(
 	round uint64, masterPublicKey *typesDKG.MasterPublicKey) {
+	if g.isProhibited(StateAddDKGMasterPublicKey) {
+		return
+	}
 	if round != masterPublicKey.Round {
 		return
 	}
@@ -208,6 +216,9 @@ func (g *Governance) IsDKGMPKReady(round uint64) bool {
 
 // AddDKGFinalize adds a DKG finalize message.
 func (g *Governance) AddDKGFinalize(round uint64, final *typesDKG.Finalize) {
+	if g.isProhibited(StateAddDKGFinal) {
+		return
+	}
 	if round != final.Round {
 		return
 	}
@@ -436,3 +447,33 @@ func (g *Governance) SwitchToRemoteMode(n *Network) {
 	g.networkModule = n
 	n.addStateModule(g.stateModule)
 }
+
+// Prohibit would prohibit DKG related state change requests.
+//
+// Note this method only prevents local modification, state changes related to
+// DKG from others won't be prohibited.
+func (g *Governance) Prohibit(t StateChangeType) {
+	g.lock.Lock()
+	defer g.lock.Unlock()
+	switch t {
+	case StateAddDKGMasterPublicKey, StateAddDKGFinal, StateAddDKGComplaint:
+		g.prohibitedTypes[t] = struct{}{}
+	default:
+		panic(fmt.Errorf("unsupported state change type to prohibit: %s", t))
+	}
+}
+
+// Unprohibit would unprohibit DKG related state change requests.
+func (g *Governance) Unprohibit(t StateChangeType) {
+	g.lock.Lock()
+	defer g.lock.Unlock()
+	delete(g.prohibitedTypes, t)
+}
+
+// isProhibited checks if a state change request is prohibited or not.
+func (g *Governance) isProhibited(t StateChangeType) (prohibited bool) {
+	g.lock.RLock()
+	defer g.lock.RUnlock()
+	_, prohibited = g.prohibitedTypes[t]
+	return
+}
diff --git a/core/test/governance_test.go b/core/test/governance_test.go
index af2b5c5..cb280a8 100644
--- a/core/test/governance_test.go
+++ b/core/test/governance_test.go
@@ -22,6 +22,11 @@ import (
 	"time"
 
 	"github.com/dexon-foundation/dexon-consensus/common"
+	"github.com/dexon-foundation/dexon-consensus/core/crypto"
+	"github.com/dexon-foundation/dexon-consensus/core/crypto/dkg"
+	"github.com/dexon-foundation/dexon-consensus/core/types"
+	typesDKG "github.com/dexon-foundation/dexon-consensus/core/types/dkg"
+	"github.com/dexon-foundation/dexon-consensus/core/utils"
 	"github.com/stretchr/testify/suite"
 )
 
@@ -98,6 +103,74 @@ func (s *GovernanceTestSuite) TestRegisterChange() {
 	req.Equal(g.Configuration(7).DKGSetSize, uint32(40))
 }
 
+func (s *GovernanceTestSuite) TestProhibit() {
+	round := uint64(1)
+	prvKeys, genesisNodes, err := NewKeys(4)
+	s.Require().NoError(err)
+	gov, err := NewGovernance(NewState(
+		1, genesisNodes, 100*time.Millisecond, &common.NullLogger{}, true), 2)
+	s.Require().NoError(err)
+	// Test MPK.
+	proposeMPK := func(k crypto.PrivateKey) {
+		signer := utils.NewSigner(k)
+		_, pubShare := dkg.NewPrivateKeyShares(utils.GetDKGThreshold(
+			gov.Configuration(round)))
+		mpk := &typesDKG.MasterPublicKey{
+			Round:           round,
+			DKGID:           typesDKG.NewID(types.NewNodeID(k.PublicKey())),
+			PublicKeyShares: *pubShare,
+		}
+		s.Require().NoError(signer.SignDKGMasterPublicKey(mpk))
+		gov.AddDKGMasterPublicKey(round, mpk)
+	}
+	proposeMPK(prvKeys[0])
+	s.Require().Len(gov.DKGMasterPublicKeys(round), 1)
+	gov.Prohibit(StateAddDKGMasterPublicKey)
+	proposeMPK(prvKeys[1])
+	s.Require().Len(gov.DKGMasterPublicKeys(round), 1)
+	gov.Unprohibit(StateAddDKGMasterPublicKey)
+	proposeMPK(prvKeys[1])
+	s.Require().Len(gov.DKGMasterPublicKeys(round), 2)
+	// Test Complaint.
+	proposeComplaint := func(k crypto.PrivateKey) {
+		signer := utils.NewSigner(k)
+		comp := &typesDKG.Complaint{
+			ProposerID: types.NewNodeID(k.PublicKey()),
+			Round:      round,
+		}
+		s.Require().NoError(signer.SignDKGComplaint(comp))
+		gov.AddDKGComplaint(round, comp)
+	}
+	proposeComplaint(prvKeys[0])
+	s.Require().Len(gov.DKGComplaints(round), 1)
+	gov.Prohibit(StateAddDKGComplaint)
+	proposeComplaint(prvKeys[1])
+	s.Require().Len(gov.DKGComplaints(round), 1)
+	gov.Unprohibit(StateAddDKGComplaint)
+	proposeComplaint(prvKeys[1])
+	s.Require().Len(gov.DKGComplaints(round), 2)
+	// Test DKG Final.
+	proposeFinal := func(k crypto.PrivateKey) {
+		signer := utils.NewSigner(k)
+		final := &typesDKG.Finalize{
+			Round:      round,
+			ProposerID: types.NewNodeID(k.PublicKey()),
+		}
+		s.Require().NoError(signer.SignDKGFinalize(final))
+		gov.AddDKGFinalize(round, final)
+	}
+	gov.Prohibit(StateAddDKGFinal)
+	for _, k := range prvKeys {
+		proposeFinal(k)
+	}
+	s.Require().False(gov.IsDKGFinal(round))
+	gov.Unprohibit(StateAddDKGFinal)
+	for _, k := range prvKeys {
+		proposeFinal(k)
+	}
+	s.Require().True(gov.IsDKGFinal(round))
+}
+
 func TestGovernance(t *testing.T) {
 	suite.Run(t, new(GovernanceTestSuite))
 }
diff --git a/core/test/network_test.go b/core/test/network_test.go
index 3e9b557..863fee2 100644
--- a/core/test/network_test.go
+++ b/core/test/network_test.go
@@ -244,6 +244,7 @@ func (s *NetworkTestSuite) TestBroadcastToSet() {
 	var (
 		req       = s.Require()
 		peerCount = 5
+		round     = uint64(1)
 	)
 	_, pubKeys, err := NewKeys(peerCount)
 	req.NoError(err)
@@ -252,13 +253,14 @@ func (s *NetworkTestSuite) TestBroadcastToSet() {
 	req.NoError(err)
 	req.NoError(gov.State().RequestChange(StateChangeDKGSetSize, uint32(1)))
 	req.NoError(gov.State().RequestChange(StateChangeNotarySetSize, uint32(1)))
+	gov.NotifyRound(round)
 	networks := s.setupNetworks(pubKeys)
 	cache := utils.NewNodeSetCache(gov)
 	// Cache required set of nodeIDs.
-	dkgSet, err := cache.GetDKGSet(0)
+	dkgSet, err := cache.GetDKGSet(round)
 	req.NoError(err)
 	req.Len(dkgSet, 1)
-	notarySet, err := cache.GetNotarySet(0)
+	notarySet, err := cache.GetNotarySet(round)
 	req.NoError(err)
 	req.Len(notarySet, 1)
 	var (
@@ -290,13 +292,14 @@ func (s *NetworkTestSuite) TestBroadcastToSet() {
 	nerd.AddNodeSetCache(cache)
 	// Try broadcasting with datum from round 0, and make sure only node belongs
 	// to that set receiving the message.
-	nerd.BroadcastVote(&types.Vote{})
+	nerd.BroadcastVote(&types.Vote{VoteHeader: types.VoteHeader{
+		Position: types.Position{Round: round}}})
 	req.IsType(&types.Vote{}, <-notaryNode.ReceiveChan())
-	nerd.BroadcastDKGPrivateShare(&typesDKG.PrivateShare{})
+	nerd.BroadcastDKGPrivateShare(&typesDKG.PrivateShare{Round: round})
 	req.IsType(&typesDKG.PrivateShare{}, <-dkgNode.ReceiveChan())
-	nerd.BroadcastDKGPartialSignature(&typesDKG.PartialSignature{})
+	nerd.BroadcastDKGPartialSignature(&typesDKG.PartialSignature{Round: round})
 	req.IsType(&typesDKG.PartialSignature{}, <-dkgNode.ReceiveChan())
-	nerd.BroadcastBlock(&types.Block{})
+	nerd.BroadcastBlock(&types.Block{Position: types.Position{Round: round}})
 	req.IsType(&types.Block{}, <-notaryNode.ReceiveChan())
 }
 
diff --git a/core/utils/utils.go b/core/utils/utils.go
index b8bd95e..d714068 100644
--- a/core/utils/utils.go
+++ b/core/utils/utils.go
@@ -138,3 +138,9 @@ func LaunchDummyReceiver(
 func GetDKGThreshold(config *types.Config) int {
 	return int(config.DKGSetSize/3) + 1
 }
+
+// GetNextRoundCheckpoint returns the block height to check if the next round
+// is ready.
+func GetNextRoundCheckpoint(begin, length uint64) uint64 {
+	return begin + length*9/10
+}
diff --git a/integration_test/byzantine_test.go b/integration_test/byzantine_test.go
index c90c9a1..a709870 100644
--- a/integration_test/byzantine_test.go
+++ b/integration_test/byzantine_test.go
@@ -81,7 +81,7 @@ func (s *ByzantineTestSuite) setupNodes(
 		gov.SwitchToRemoteMode(networkModule)
 		gov.NotifyRound(0)
 		networkModule.AddNodeSetCache(utils.NewNodeSetCache(gov))
-		app := test.NewApp(1, gov)
+		app := test.NewApp(1, gov, nil)
 		nodes[nID] = &node{nID, nil, app, gov, dbInst, networkModule}
 		go func() {
 			defer wg.Done()
diff --git a/integration_test/consensus_test.go b/integration_test/consensus_test.go
index 4f916c3..78d0de5 100644
--- a/integration_test/consensus_test.go
+++ b/integration_test/consensus_test.go
@@ -78,7 +78,7 @@ func (s *ConsensusTestSuite) setupNodes(
 		gov.SwitchToRemoteMode(networkModule)
 		gov.NotifyRound(0)
 		networkModule.AddNodeSetCache(utils.NewNodeSetCache(gov))
-		app := test.NewApp(1, gov)
+		app := test.NewApp(1, gov, nil)
 		nID := types.NewNodeID(k.PublicKey())
 		nodes[nID] = &node{nID, nil, app, gov, dbInst, networkModule}
 		go func() {