gecko/snow/engine/snowman/test_vm.go

100 lines
2.1 KiB
Go

// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.
package snowman
import (
"errors"
"github.com/ava-labs/gecko/ids"
"github.com/ava-labs/gecko/snow/consensus/snowman"
"github.com/ava-labs/gecko/snow/engine/common"
)
var (
errBuildBlock = errors.New("unexpectedly called BuildBlock")
errParseBlock = errors.New("unexpectedly called ParseBlock")
errGetBlock = errors.New("unexpectedly called GetBlock")
)
// VMTest ...
type VMTest struct {
common.VMTest
CantBuildBlock,
CantParseBlock,
CantGetBlock,
CantSetPreference,
CantLastAccepted bool
BuildBlockF func() (snowman.Block, error)
ParseBlockF func([]byte) (snowman.Block, error)
GetBlockF func(ids.ID) (snowman.Block, error)
SetPreferenceF func(ids.ID)
LastAcceptedF func() ids.ID
}
// Default ...
func (vm *VMTest) Default(cant bool) {
vm.VMTest.Default(cant)
vm.CantBuildBlock = cant
vm.CantParseBlock = cant
vm.CantGetBlock = cant
vm.CantSetPreference = cant
vm.CantLastAccepted = cant
}
// BuildBlock ...
func (vm *VMTest) BuildBlock() (snowman.Block, error) {
if vm.BuildBlockF != nil {
return vm.BuildBlockF()
}
if vm.CantBuildBlock && vm.T != nil {
vm.T.Fatal(errBuildBlock)
}
return nil, errBuildBlock
}
// ParseBlock ...
func (vm *VMTest) ParseBlock(b []byte) (snowman.Block, error) {
if vm.ParseBlockF != nil {
return vm.ParseBlockF(b)
}
if vm.CantParseBlock && vm.T != nil {
vm.T.Fatal(errParseBlock)
}
return nil, errParseBlock
}
// GetBlock ...
func (vm *VMTest) GetBlock(id ids.ID) (snowman.Block, error) {
if vm.GetBlockF != nil {
return vm.GetBlockF(id)
}
if vm.CantGetBlock && vm.T != nil {
vm.T.Fatal(errGetBlock)
}
return nil, errGetBlock
}
// SetPreference ...
func (vm *VMTest) SetPreference(id ids.ID) {
if vm.SetPreferenceF != nil {
vm.SetPreferenceF(id)
} else if vm.CantSetPreference && vm.T != nil {
vm.T.Fatalf("Unexpectedly called SetPreference")
}
}
// LastAccepted ...
func (vm *VMTest) LastAccepted() ids.ID {
if vm.LastAcceptedF != nil {
return vm.LastAcceptedF()
}
if vm.CantLastAccepted && vm.T != nil {
vm.T.Fatalf("Unexpectedly called LastAccepted")
}
return ids.ID{}
}