Repository: pactus-project/pactus
Pull request #518 (2023-06-18 07:53:46)

refactor: optimization of total power calculation
## Description

Prior to this PR, the total power (the total staked coins) was calculated at each height by iterating over all validators and summing up the power (stakes). This was not an efficient method for calculating total power.

With this PR, a more efficient approach has been implemented. Now, whenever a validator bonds or unbonds a stake, we record this as a power delta (change) and apply it to the total power. This removes the need for iterating over all validators at every height, thereby optimizing the total power calculation process.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pactus-project/pactus/pull/518?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pactus-project) Report
> Merging [#518](https://app.codecov.io/gh/pactus-project/pactus/pull/518?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pactus-project) (c8aefb8) into [main](https://app.codecov.io/gh/pactus-project/pactus/commit/5e2f1357a770a84454f2c5b6f8eb5e6f30518936?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pactus-project) (5e2f135) will **increase** coverage by `0.15%`.
> The diff coverage is `97.29%`.


```diff
@@            Coverage Diff             @@
##             main     #518      +/-   ##
==========================================
+ Coverage   83.15%   83.31%   +0.15%     
==========================================
  Files         154      154              
  Lines        7306     7322      +16     
==========================================
+ Hits         6075     6100      +25     
+ Misses        951      942       -9     
  Partials      280      280              
```


Ja7ad: @b00f All thing is okay for merge

Patch:
diff --git a/execution/executor/bond.go b/execution/executor/bond.go
index 4fc03337..b8163055 100644
--- a/execution/executor/bond.go
+++ b/execution/executor/bond.go
@@ -77,6 +77,7 @@ func (e *BondExecutor) Execute(trx *tx.Tx, sb sandbox.Sandbox) error {
 	receiverVal.AddToStake(pld.Stake)
 	receiverVal.UpdateLastBondingHeight(sb.CurrentHeight())
 
+	sb.UpdatePowerDelta(pld.Stake)
 	sb.UpdateAccount(pld.Sender, senderAcc)
 	sb.UpdateValidator(receiverVal)
 
diff --git a/execution/executor/unbond.go b/execution/executor/unbond.go
index c4512acc..c3a48f99 100644
--- a/execution/executor/unbond.go
+++ b/execution/executor/unbond.go
@@ -53,6 +53,11 @@ func (e *UnbondExecutor) Execute(trx *tx.Tx, sb sandbox.Sandbox) error {
 
 	val.IncSequence()
 	val.UpdateUnbondingHeight(sb.CurrentHeight())
+
+	// At this point, the validator's power is zero.
+	// However, we know the validator's stake.
+	// So, we can update the power delta with the negative of the validator's stake.
+	sb.UpdatePowerDelta(-1 * val.Power())
 	sb.UpdateValidator(val)
 
 	return nil
diff --git a/sandbox/interface.go b/sandbox/interface.go
index 57617415..2567f49c 100644
--- a/sandbox/interface.go
+++ b/sandbox/interface.go
@@ -20,6 +20,8 @@ type Sandbox interface {
 	Validator(crypto.Address) *validator.Validator
 	MakeNewValidator(*bls.PublicKey) *validator.Validator
 	UpdateValidator(*validator.Validator)
+	UpdatePowerDelta(delta int64)
+	PowerDelta() int64
 
 	VerifyProof(hash.Stamp, sortition.Proof, *validator.Validator) bool
 	Committee() committee.Reader
diff --git a/sandbox/mock.go b/sandbox/mock.go
index 861c597a..e81092b7 100644
--- a/sandbox/mock.go
+++ b/sandbox/mock.go
@@ -21,7 +21,8 @@ type MockSandbox struct {
 	TestStore            *store.MockStore
 	TestCommittee        committee.Committee
 	TestCommitteeSigners []crypto.Signer
-	AcceptTestSortition  bool
+	TestAcceptSortition  bool
+	TestPowerDelta       int64
 }
 
 func MockingSandbox() *MockSandbox {
@@ -97,7 +98,12 @@ func (m *MockSandbox) IterateValidators(consumer func(*validator.Validator, bool
 func (m *MockSandbox) Committee() committee.Reader {
 	return m.TestCommittee
 }
-
+func (m *MockSandbox) UpdatePowerDelta(delta int64) {
+	m.TestPowerDelta += delta
+}
+func (m *MockSandbox) PowerDelta() int64 {
+	return m.TestPowerDelta
+}
 func (m *MockSandbox) VerifyProof(hash.Stamp, sortition.Proof, *validator.Validator) bool {
-	return m.AcceptTestSortition
+	return m.TestAcceptSortition
 }
diff --git a/sandbox/sandbox.go b/sandbox/sandbox.go
index d589af93..87c475bd 100644
--- a/sandbox/sandbox.go
+++ b/sandbox/sandbox.go
@@ -28,6 +28,8 @@ type sandbox struct {
 	params          param.Params
 	totalAccounts   int32
 	totalValidators int32
+	totalPower      int64
+	powerDelta      int64
 }
 
 type sandboxValidator struct {
@@ -40,11 +42,13 @@ type sandboxAccount struct {
 	updated bool
 }
 
-func NewSandbox(store store.Reader, params param.Params, committee committee.Reader) Sandbox {
+func NewSandbox(store store.Reader, params param.Params,
+	committee committee.Reader, totalPower int64) Sandbox {
 	sb := &sandbox{
-		store:     store,
-		committee: committee,
-		params:    params,
+		store:      store,
+		committee:  committee,
+		totalPower: totalPower,
+		params:     params,
 	}
 
 	sb.accounts = make(map[crypto.Address]*sandboxAccount)
@@ -223,7 +227,22 @@ func (sb *sandbox) Committee() committee.Reader {
 	return sb.committee
 }
 
-// TODO: write test for me.
+// UpdatePowerDelta updates the change in the total power of the blockchain.
+// The delta is the amount of change in the total power and can be either positive or negative.
+func (sb *sandbox) UpdatePowerDelta(delta int64) {
+	sb.lk.Lock()
+	defer sb.lk.Unlock()
+
+	sb.powerDelta += delta
+}
+
+func (sb *sandbox) PowerDelta() int64 {
+	sb.lk.RLock()
+	defer sb.lk.RUnlock()
+
+	return sb.powerDelta
+}
+
 // VerifyProof verifies proof of a sortition transaction.
 func (sb *sandbox) VerifyProof(stamp hash.Stamp, proof sortition.Proof, val *validator.Validator) bool {
 	_, b := sb.store.RecentBlockByStamp(stamp)
@@ -231,10 +250,5 @@ func (sb *sandbox) VerifyProof(stamp hash.Stamp, proof sortition.Proof, val *val
 		return false
 	}
 	seed := b.Header().SortitionSeed()
-	total := int64(0) // TODO: we can get it from state
-	sb.store.IterateValidators(func(val *validator.Validator) bool {
-		total += val.Power()
-		return false
-	})
-	return sortition.VerifyProof(seed, proof, val.PublicKey(), total, val.Power())
+	return sortition.VerifyProof(seed, proof, val.PublicKey(), sb.totalPower, val.Power())
 }
diff --git a/state/state.go b/state/state.go
index 6248ba54..9336f039 100644
--- a/state/state.go
+++ b/state/state.go
@@ -37,6 +37,7 @@ type state struct {
 	params          param.Params
 	txPool          txpool.TxPool
 	committee       committee.Committee
+	totalPower      int64
 	lastInfo        *lastinfo.LastInfo
 	accountMerkle   *persistentmerkle.Tree
 	validatorMerkle *persistentmerkle.Tree
@@ -78,6 +79,8 @@ func LoadOrNewState(
 		}
 	}
 
+	st.totalPower = st.retrieveTotalPower()
+
 	st.loadMerkels()
 
 	txPool.SetNewSandboxAndRecheck(st.concreteSandbox())
@@ -88,7 +91,7 @@ func LoadOrNewState(
 }
 
 func (st *state) concreteSandbox() sandbox.Sandbox {
-	return sandbox.NewSandbox(st.store, st.params, st.committee)
+	return sandbox.NewSandbox(st.store, st.params, st.committee, st.totalPower)
 }
 
 func (st *state) tryLoadLastInfo() error {
@@ -171,6 +174,15 @@ func (st *state) loadMerkels() {
 	})
 }
 
+func (st *state) retrieveTotalPower() int64 {
+	totalPower := int64(0)
+	st.store.IterateValidators(func(val *validator.Validator) (stop bool) {
+		totalPower += val.Power()
+		return false
+	})
+	return totalPower
+}
+
 func (st *state) stateRoot() hash.Hash {
 	accRoot := st.accountMerkle.Root()
 	valRoot := st.validatorMerkle.Root()
@@ -445,7 +457,6 @@ func (st *state) CommitBlock(height uint32, block *block.Block, cert *block.Cert
 }
 
 func (st *state) evaluateSortition() bool {
-	totalPower := st.totalPower()
 	evaluated := false
 	for _, signer := range st.signers {
 		val, _ := st.store.Validator(signer.Address())
@@ -464,7 +475,7 @@ func (st *state) evaluateSortition() bool {
 			continue
 		}
 
-		ok, proof := sortition.EvaluateSortition(st.lastInfo.SortitionSeed(), signer, totalPower, val.Power())
+		ok, proof := sortition.EvaluateSortition(st.lastInfo.SortitionSeed(), signer, st.totalPower, val.Power())
 		if ok {
 			trx := tx.NewSortitionTx(st.lastInfo.BlockHash().Stamp(), val.Sequence()+1, val.Address(), proof)
 			signer.SignMsg(trx)
@@ -517,6 +528,8 @@ func (st *state) commitSandbox(sb sandbox.Sandbox, round int16) {
 			st.validatorMerkle.SetHash(int(val.Number()), val.Hash())
 		}
 	})
+
+	st.totalPower += sb.PowerDelta()
 }
 
 func (st *state) validateBlockTime(t time.Time) error {
@@ -543,29 +556,13 @@ func (st *state) TotalPower() int64 {
 	st.lk.RLock()
 	defer st.lk.RUnlock()
 
-	return st.totalPower()
+	return st.totalPower
 }
 
 func (st *state) CommitteePower() int64 {
 	st.lk.RLock()
 	defer st.lk.RUnlock()
 
-	return st.committeePower()
-}
-
-// TODO: add test for me when a validator is parked (unbonded)
-// TODO: Improve performance of remember total power
-// TODO: sandbox has the same logic.
-func (st *state) totalPower() int64 {
-	p := int64(0)
-	st.store.IterateValidators(func(val *validator.Validator) bool {
-		p += val.Power()
-		return false
-	})
-	return p
-}
-
-func (st *state) committeePower() int64 {
 	return st.committee.TotalPower()
 }
 


Test patch:
diff --git a/execution/execution_test.go b/execution/execution_test.go
index 4a625f12..d268741c 100644
--- a/execution/execution_test.go
+++ b/execution/execution_test.go
@@ -142,7 +142,7 @@ func TestLockTime(t *testing.T) {
 		val := sb.MakeNewValidator(pub)
 		sb.UpdateValidator(val)
 
-		sb.AcceptTestSortition = true
+		sb.TestAcceptSortition = true
 		pld := &payload.SortitionPayload{
 			Address: pub.Address(),
 			Proof:   sortition.GenerateRandomProof(),
diff --git a/execution/executor/bond_test.go b/execution/executor/bond_test.go
index 17e2ab50..7be8d36b 100644
--- a/execution/executor/bond_test.go
+++ b/execution/executor/bond_test.go
@@ -17,11 +17,12 @@ func TestExecuteBondTx(t *testing.T) {
 	senderAddr, senderAcc := tSandbox.TestStore.RandomTestAcc()
 	senderBalance := senderAcc.Balance()
 	pub, _ := bls.GenerateTestKeyPair()
+	receiverAddr := pub.Address()
 	fee, amt := randomAmountAndFee(senderBalance / 2)
 
 	t.Run("Should fail, invalid sender", func(t *testing.T) {
 		trx := tx.NewBondTx(tStamp500000, 1, crypto.GenerateTestAddress(),
-			pub.Address(), pub, amt, fee, "invalid sender")
+			receiverAddr, pub, amt, fee, "invalid sender")
 
 		err := exe.Execute(trx, tSandbox)
 		assert.Equal(t, errors.Code(err), errors.ErrInvalidAddress)
@@ -37,7 +38,7 @@ func TestExecuteBondTx(t *testing.T) {
 
 	t.Run("Should fail, invalid sequence", func(t *testing.T) {
 		trx := tx.NewBondTx(tStamp500000, senderAcc.Sequence()+2, senderAddr,
-			pub.Address(), pub, amt, fee, "invalid sequence")
+			receiverAddr, pub, amt, fee, "invalid sequence")
 
 		err := exe.Execute(trx, tSandbox)
 		assert.Equal(t, errors.Code(err), errors.ErrInvalidSequence)
@@ -45,7 +46,7 @@ func TestExecuteBondTx(t *testing.T) {
 
 	t.Run("Should fail, insufficient balance", func(t *testing.T) {
 		trx := tx.NewBondTx(tStamp500000, senderAcc.Sequence()+1, senderAddr,
-			pub.Address(), pub, senderBalance+1, 0, "insufficient balance")
+			receiverAddr, pub, senderBalance+1, 0, "insufficient balance")
 
 		err := exe.Execute(trx, tSandbox)
 		assert.Equal(t, errors.Code(err), errors.ErrInsufficientFunds)
@@ -74,7 +75,7 @@ func TestExecuteBondTx(t *testing.T) {
 
 	t.Run("Should fail, public key is not set", func(t *testing.T) {
 		trx := tx.NewBondTx(tStamp500000, senderAcc.Sequence()+1, senderAddr,
-			pub.Address(), nil, amt, fee, "no public key")
+			receiverAddr, nil, amt, fee, "no public key")
 
 		err := exe.Execute(trx, tSandbox)
 		assert.Equal(t, errors.Code(err), errors.ErrInvalidPublicKey)
@@ -82,7 +83,7 @@ func TestExecuteBondTx(t *testing.T) {
 
 	t.Run("Ok", func(t *testing.T) {
 		trx := tx.NewBondTx(tStamp500000, senderAcc.Sequence()+1, senderAddr,
-			pub.Address(), pub, amt, fee, "ok")
+			receiverAddr, pub, amt, fee, "ok")
 
 		assert.NoError(t, exe.Execute(trx, tSandbox), "Ok")
 		assert.Error(t, exe.Execute(trx, tSandbox), "Execute again, should fail")
@@ -90,19 +91,17 @@ func TestExecuteBondTx(t *testing.T) {
 
 	t.Run("Should fail, public key should not set for existing validators", func(t *testing.T) {
 		trx := tx.NewBondTx(tStamp500000, senderAcc.Sequence()+2, senderAddr,
-			pub.Address(), pub, amt, fee, "with public key")
+			receiverAddr, pub, amt, fee, "with public key")
 
 		err := exe.Execute(trx, tSandbox)
 		assert.Equal(t, errors.Code(err), errors.ErrInvalidPublicKey)
 	})
 
-	assert.Equal(t, tSandbox.Account(senderAddr).Balance(),
-		senderBalance-(amt+fee))
-	assert.Equal(t, tSandbox.Validator(pub.Address()).Stake(), amt)
-	assert.Equal(t, tSandbox.Validator(pub.Address()).LastBondingHeight(),
-		tSandbox.CurrentHeight())
+	assert.Equal(t, tSandbox.Account(senderAddr).Balance(), senderBalance-(amt+fee))
+	assert.Equal(t, tSandbox.Validator(receiverAddr).Stake(), amt)
+	assert.Equal(t, tSandbox.Validator(receiverAddr).LastBondingHeight(), tSandbox.CurrentHeight())
+	assert.Equal(t, tSandbox.PowerDelta(), amt)
 	assert.Equal(t, exe.Fee(), fee)
-
 	checkTotalCoin(t, fee)
 }
 
diff --git a/execution/executor/sortition_test.go b/execution/executor/sortition_test.go
index 544c5456..c9d84f45 100644
--- a/execution/executor/sortition_test.go
+++ b/execution/executor/sortition_test.go
@@ -33,7 +33,7 @@ func TestExecuteSortitionTx(t *testing.T) {
 
 	t.Run("Should fail, Invalid address", func(t *testing.T) {
 		trx := tx.NewSortitionTx(tStamp500000, 1, crypto.GenerateTestAddress(), proof)
-		tSandbox.AcceptTestSortition = true
+		tSandbox.TestAcceptSortition = true
 		assert.Equal(t, errors.Code(exe.Execute(trx, tSandbox)), errors.ErrInvalidAddress)
 	})
 
@@ -41,7 +41,7 @@ func TestExecuteSortitionTx(t *testing.T) {
 	tSandbox.UpdateValidator(newVal)
 	t.Run("Should fail, Bonding period", func(t *testing.T) {
 		trx := tx.NewSortitionTx(tStamp500000, newVal.Sequence()+1, newVal.Address(), proof)
-		tSandbox.AcceptTestSortition = true
+		tSandbox.TestAcceptSortition = true
 		assert.Equal(t, errors.Code(exe.Execute(trx, tSandbox)), errors.ErrInvalidHeight)
 	})
 
@@ -49,25 +49,25 @@ func TestExecuteSortitionTx(t *testing.T) {
 
 	t.Run("Should fail, Invalid sequence", func(t *testing.T) {
 		trx := tx.NewSortitionTx(tStamp500000, newVal.Sequence()+2, newVal.Address(), proof)
-		tSandbox.AcceptTestSortition = true
+		tSandbox.TestAcceptSortition = true
 		assert.Equal(t, errors.Code(exe.Execute(trx, tSandbox)), errors.ErrInvalidSequence)
 	})
 
 	t.Run("Should fail, Invalid proof", func(t *testing.T) {
 		trx := tx.NewSortitionTx(tStamp500000, newVal.Sequence()+1, newVal.Address(), proof)
-		tSandbox.AcceptTestSortition = false
+		tSandbox.TestAcceptSortition = false
 		assert.Equal(t, errors.Code(exe.Execute(trx, tSandbox)), errors.ErrInvalidProof)
 	})
 
 	t.Run("Should fail, Committee has free seats and validator is in the committee", func(t *testing.T) {
 		trx := tx.NewSortitionTx(tStamp500000, existingVal.Sequence()+1, existingVal.Address(), proof)
-		tSandbox.AcceptTestSortition = true
+		tSandbox.TestAcceptSortition = true
 		assert.Equal(t, errors.Code(exe.Execute(trx, tSandbox)), errors.ErrInvalidTx)
 	})
 
 	t.Run("Should be ok", func(t *testing.T) {
 		trx := tx.NewSortitionTx(tStamp500000, newVal.Sequence()+1, newVal.Address(), proof)
-		tSandbox.AcceptTestSortition = true
+		tSandbox.TestAcceptSortition = true
 		assert.NoError(t, exe.Execute(trx, tSandbox))
 
 		// Execute again, should fail
@@ -88,7 +88,7 @@ func TestSortitionNonStrictMode(t *testing.T) {
 	val := tSandbox.TestStore.RandomTestVal()
 	proof := sortition.GenerateRandomProof()
 
-	tSandbox.AcceptTestSortition = true
+	tSandbox.TestAcceptSortition = true
 	trx := tx.NewSortitionTx(tStamp500000, val.Sequence(), val.Address(), proof)
 	assert.Error(t, exe1.Execute(trx, tSandbox))
 	assert.NoError(t, exe2.Execute(trx, tSandbox))
@@ -119,7 +119,7 @@ func TestChangePower1(t *testing.T) {
 	proof3 := sortition.GenerateRandomProof()
 
 	tSandbox.TestParams.CommitteeSize = 4
-	tSandbox.AcceptTestSortition = true
+	tSandbox.TestAcceptSortition = true
 	trx1 := tx.NewSortitionTx(tStamp500000, val1.Sequence()+1, val1.Address(), proof1)
 	assert.NoError(t, exe.Execute(trx1, tSandbox))
 
@@ -162,7 +162,7 @@ func TestChangePower2(t *testing.T) {
 	proof4 := sortition.GenerateRandomProof()
 
 	tSandbox.TestParams.CommitteeSize = 7
-	tSandbox.AcceptTestSortition = true
+	tSandbox.TestAcceptSortition = true
 	trx1 := tx.NewSortitionTx(tStamp500000, val1.Sequence()+1, val1.Address(), proof1)
 	assert.NoError(t, exe.Execute(trx1, tSandbox))
 
@@ -196,7 +196,7 @@ func TestOldestDidNotPropose(t *testing.T) {
 	}
 
 	tSandbox.TestParams.CommitteeSize = 7
-	tSandbox.AcceptTestSortition = true
+	tSandbox.TestAcceptSortition = true
 
 	stamp := tStamp500000
 	for i := 0; i < 8; i = i + 2 {
diff --git a/execution/executor/unbond_test.go b/execution/executor/unbond_test.go
index b757531b..2cf0f521 100644
--- a/execution/executor/unbond_test.go
+++ b/execution/executor/unbond_test.go
@@ -15,6 +15,7 @@ func TestExecuteUnbondTx(t *testing.T) {
 	exe := NewUnbondExecutor(true)
 
 	pub, _ := bls.GenerateTestKeyPair()
+	valAddr := pub.Address()
 	val := tSandbox.MakeNewValidator(pub)
 	tSandbox.UpdateValidator(val)
 
@@ -24,7 +25,7 @@ func TestExecuteUnbondTx(t *testing.T) {
 	})
 
 	t.Run("Should fail, Invalid sequence", func(t *testing.T) {
-		trx := tx.NewUnbondTx(tStamp500000, val.Sequence()+2, pub.Address(), "invalid sequence")
+		trx := tx.NewUnbondTx(tStamp500000, val.Sequence()+2, valAddr, "invalid sequence")
 		assert.Equal(t, errors.Code(exe.Execute(trx, tSandbox)), errors.ErrInvalidSequence)
 	})
 
@@ -45,7 +46,7 @@ func TestExecuteUnbondTx(t *testing.T) {
 	})
 
 	t.Run("Ok", func(t *testing.T) {
-		trx := tx.NewUnbondTx(tStamp500000, val.Sequence()+1, pub.Address(), "Ok")
+		trx := tx.NewUnbondTx(tStamp500000, val.Sequence()+1, valAddr, "Ok")
 
 		assert.NoError(t, exe.Execute(trx, tSandbox))
 
@@ -53,9 +54,10 @@ func TestExecuteUnbondTx(t *testing.T) {
 		assert.Error(t, exe.Execute(trx, tSandbox))
 	})
 
-	assert.Zero(t, tSandbox.Validator(pub.Address()).Stake())
-	assert.Zero(t, tSandbox.Validator(pub.Address()).Power())
-	assert.Equal(t, tSandbox.Validator(pub.Address()).UnbondingHeight(), tSandbox.CurrentHeight())
+	assert.Zero(t, tSandbox.Validator(valAddr).Stake())
+	assert.Zero(t, tSandbox.Validator(valAddr).Power())
+	assert.Equal(t, tSandbox.Validator(valAddr).UnbondingHeight(), tSandbox.CurrentHeight())
+	assert.Equal(t, tSandbox.PowerDelta(), -1*val.Stake())
 	assert.Zero(t, exe.Fee())
 
 	checkTotalCoin(t, 0)
diff --git a/sandbox/sandbox_test.go b/sandbox/sandbox_test.go
index 3f72ac1c..a61f745c 100644
--- a/sandbox/sandbox_test.go
+++ b/sandbox/sandbox_test.go
@@ -7,6 +7,7 @@ import (
 	"github.com/pactus-project/pactus/crypto"
 	"github.com/pactus-project/pactus/crypto/bls"
 	"github.com/pactus-project/pactus/crypto/hash"
+	"github.com/pactus-project/pactus/sortition"
 	"github.com/pactus-project/pactus/store"
 	"github.com/pactus-project/pactus/types/account"
 	"github.com/pactus-project/pactus/types/block"
@@ -30,6 +31,7 @@ func setup(t *testing.T) {
 	acc.AddToBalance(21 * 1e14)
 	tStore.UpdateAccount(crypto.TreasuryAddress, acc)
 
+	totalPower := int64(0)
 	for _, val := range committee.Validators() {
 		// For testing purpose we create some test accounts first.
 		// Account number is the validator number plus one,
@@ -37,10 +39,12 @@ func setup(t *testing.T) {
 		acc := account.NewAccount(val.Number() + 1)
 		tStore.UpdateValidator(val)
 		tStore.UpdateAccount(val.Address(), acc)
+
+		totalPower += val.Power()
 	}
 
 	tSigners = signers
-	tSandbox = NewSandbox(tStore, params, committee).(*sandbox)
+	tSandbox = NewSandbox(tStore, params, committee, totalPower).(*sandbox)
 
 	assert.Equal(t, tSandbox.CurrentHeight(), uint32(1))
 	lastHeight := util.RandUint32(144) + 21
@@ -364,3 +368,57 @@ func TestRecentBlockByStamp(t *testing.T) {
 	assert.Equal(t, h, lastHeight)
 	assert.Equal(t, b.Hash(), lastHash)
 }
+
+func TestPowerDelta(t *testing.T) {
+	setup(t)
+
+	assert.Zero(t, tSandbox.PowerDelta())
+	tSandbox.UpdatePowerDelta(1)
+	assert.Equal(t, tSandbox.PowerDelta(), int64(1))
+	tSandbox.UpdatePowerDelta(-1)
+	assert.Zero(t, tSandbox.PowerDelta())
+}
+
+func TestVerifyProof(t *testing.T) {
+	setup(t)
+
+	lastHeight, _ := tStore.LastCertificate()
+	vals := tSandbox.committee.Validators()
+
+	// Try to evaluate a valid sortition
+	var validProof sortition.Proof
+	var validStamp hash.Stamp
+	var validVal *validator.Validator
+	for i := lastHeight; i > 0; i-- {
+		block := tStore.Blocks[i]
+		for i, signer := range tSigners {
+			ok, proof := sortition.EvaluateSortition(
+				block.Header().SortitionSeed(), signer,
+				tSandbox.totalPower, vals[i].Power())
+
+			if ok {
+				validProof = proof
+				validStamp = block.Stamp()
+				validVal = vals[i]
+			}
+		}
+	}
+
+	t.Run("invalid proof", func(t *testing.T) {
+		invalidProof := sortition.GenerateRandomProof()
+		assert.False(t, tSandbox.VerifyProof(validStamp, invalidProof, validVal))
+	})
+	t.Run("invalid stamp", func(t *testing.T) {
+		invalidStamp := hash.GenerateTestStamp()
+		assert.False(t, tSandbox.VerifyProof(invalidStamp, validProof, validVal))
+	})
+
+	t.Run("genesis stamp", func(t *testing.T) {
+		invalidStamp := hash.UndefHash.Stamp()
+		assert.False(t, tSandbox.VerifyProof(invalidStamp, validProof, validVal))
+	})
+
+	t.Run("Ok", func(t *testing.T) {
+		assert.True(t, tSandbox.VerifyProof(validStamp, validProof, validVal))
+	})
+}
diff --git a/state/state_test.go b/state/state_test.go
index 8c747755..cc1e7e59 100644
--- a/state/state_test.go
+++ b/state/state_test.go
@@ -585,7 +585,7 @@ func TestLoadState(t *testing.T) {
 	assert.Equal(t, tState1.store.TotalValidators(), st1Load.(*state).store.TotalValidators())
 	assert.Equal(t, tState1.committee.Committers(), st1Load.(*state).committee.Committers())
 	assert.Equal(t, tState1.committee.TotalPower(), st1Load.(*state).committee.TotalPower())
-	assert.Equal(t, tState1.TotalPower(), st1Load.(*state).totalPower())
+	assert.Equal(t, tState1.TotalPower(), st1Load.(*state).TotalPower())
 	assert.Equal(t, tState1.store.TotalAccounts(), int32(5))
 
 	require.NoError(t, st1Load.CommitBlock(6, b6, c6))
diff --git a/txpool/pool_test.go b/txpool/pool_test.go
index 0f00aaa8..2f9e64cb 100644
--- a/txpool/pool_test.go
+++ b/txpool/pool_test.go
@@ -151,7 +151,7 @@ func TestPrepareBlockTransactions(t *testing.T) {
 		crypto.GenerateTestAddress(), 1000, 1000, "withdraw-tx")
 	val2Signer.SignMsg(withdrawTx)
 
-	tSandbox.AcceptTestSortition = true
+	tSandbox.TestAcceptSortition = true
 	sortitionTx := tx.NewSortitionTx(block1000000.Stamp(), val3.Sequence()+1, val3.Address(),
 		sortition.GenerateRandomProof())
 	val3Signer.SignMsg(sortitionTx)