prysm-pulse/beacon-chain/state/stateutil/trie_helpers_test.go
Nishant Das 1bb12c3568
Add Field Trie to State (#5118)
* add new helpers

* make zerohash public

* remove unused method

* add more tests

* cleanup

* add in new tests

* fix all tests

* Apply suggestions from code review

Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>
Co-authored-by: Raul Jordan <raul@prysmaticlabs.com>
2020-03-18 10:09:31 +08:00

140 lines
4.1 KiB
Go

package stateutil_test
import (
"testing"
ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
"github.com/prysmaticlabs/prysm/beacon-chain/state/stateutil"
"github.com/prysmaticlabs/prysm/shared/bytesutil"
"github.com/prysmaticlabs/prysm/shared/params"
"github.com/prysmaticlabs/prysm/shared/testutil"
)
func TestReturnTrieLayer_OK(t *testing.T) {
newState, _ := testutil.DeterministicGenesisState(t, 32)
root, err := stateutil.RootsArrayHashTreeRoot(newState.BlockRoots(), params.BeaconConfig().SlotsPerHistoricalRoot, "BlockRoots")
if err != nil {
t.Fatal(err)
}
blockRts := newState.BlockRoots()
roots := make([][32]byte, 0, len(blockRts))
for _, rt := range blockRts {
roots = append(roots, bytesutil.ToBytes32(rt))
}
layers := stateutil.ReturnTrieLayer(roots, uint64(len(roots)))
newRoot := *layers[len(layers)-1][0]
if newRoot != root {
t.Errorf("Wanted root of %#x but got %#x", root, newRoot)
}
}
func TestReturnTrieLayerVariable_OK(t *testing.T) {
newState, _ := testutil.DeterministicGenesisState(t, 32)
root, err := stateutil.ValidatorRegistryRoot(newState.Validators())
if err != nil {
t.Fatal(err)
}
validators := newState.Validators()
roots := make([][32]byte, 0, len(validators))
for _, val := range validators {
rt, err := stateutil.ValidatorRoot(val)
if err != nil {
t.Fatal(err)
}
roots = append(roots, rt)
}
layers := stateutil.ReturnTrieLayerVariable(roots, params.BeaconConfig().ValidatorRegistryLimit)
newRoot := *layers[len(layers)-1][0]
newRoot, err = stateutil.AddInMixin(newRoot, uint64(len(validators)))
if err != nil {
t.Fatal(err)
}
if newRoot != root {
t.Errorf("Wanted root of %#x but got %#x", root, newRoot)
}
}
func TestRecomputeFromLayer_FixedSizedArray(t *testing.T) {
newState, _ := testutil.DeterministicGenesisState(t, 32)
blockRts := newState.BlockRoots()
roots := make([][32]byte, 0, len(blockRts))
for _, rt := range blockRts {
roots = append(roots, bytesutil.ToBytes32(rt))
}
layers := stateutil.ReturnTrieLayer(roots, uint64(len(roots)))
changedIdx := []uint64{24, 41}
changedRoots := [][32]byte{{'A', 'B', 'C'}, {'D', 'E', 'F'}}
newState.UpdateBlockRootAtIndex(changedIdx[0], changedRoots[0])
newState.UpdateBlockRootAtIndex(changedIdx[1], changedRoots[1])
expectedRoot, err := stateutil.RootsArrayHashTreeRoot(newState.BlockRoots(), params.BeaconConfig().SlotsPerHistoricalRoot, "BlockRoots")
if err != nil {
t.Fatal(err)
}
root, _, err := stateutil.RecomputeFromLayer(changedRoots, changedIdx, layers)
if err != nil {
t.Fatal(err)
}
if root != expectedRoot {
t.Errorf("Wanted root of %#x but got %#x", expectedRoot, root)
}
}
func TestRecomputeFromLayer_VariableSizedArray(t *testing.T) {
newState, _ := testutil.DeterministicGenesisState(t, 32)
validators := newState.Validators()
roots := make([][32]byte, 0, len(validators))
for _, val := range validators {
rt, err := stateutil.ValidatorRoot(val)
if err != nil {
t.Fatal(err)
}
roots = append(roots, rt)
}
layers := stateutil.ReturnTrieLayerVariable(roots, params.BeaconConfig().ValidatorRegistryLimit)
changedIdx := []uint64{2, 29}
val1, err := newState.ValidatorAtIndex(10)
if err != nil {
t.Fatal(err)
}
val2, err := newState.ValidatorAtIndex(11)
if err != nil {
t.Fatal(err)
}
val1.Slashed = true
val1.ExitEpoch = 20
val2.Slashed = true
val2.ExitEpoch = 40
changedVals := []*ethpb.Validator{val1, val2}
newState.UpdateValidatorAtIndex(changedIdx[0], changedVals[0])
newState.UpdateValidatorAtIndex(changedIdx[1], changedVals[1])
expectedRoot, err := stateutil.ValidatorRegistryRoot(newState.Validators())
if err != nil {
t.Fatal(err)
}
roots = make([][32]byte, 0, len(changedVals))
for _, val := range changedVals {
rt, err := stateutil.ValidatorRoot(val)
if err != nil {
t.Fatal(err)
}
roots = append(roots, rt)
}
root, _, err := stateutil.RecomputeFromLayerVariable(roots, changedIdx, layers)
if err != nil {
t.Fatal(err)
}
root, err = stateutil.AddInMixin(root, uint64(len(validators)))
if err != nil {
t.Fatal(err)
}
if root != expectedRoot {
t.Errorf("Wanted root of %#x but got %#x", expectedRoot, root)
}
}