github.com/MetalBlockchain/metalgo@v1.11.9/snow/engine/avalanche/vertex/test_storage.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package vertex
     5  
     6  import (
     7  	"context"
     8  	"errors"
     9  	"testing"
    10  
    11  	"github.com/stretchr/testify/require"
    12  
    13  	"github.com/MetalBlockchain/metalgo/ids"
    14  	"github.com/MetalBlockchain/metalgo/snow/consensus/avalanche"
    15  )
    16  
    17  var (
    18  	errGet                = errors.New("unexpectedly called Get")
    19  	errEdge               = errors.New("unexpectedly called Edge")
    20  	errStopVertexAccepted = errors.New("unexpectedly called StopVertexAccepted")
    21  
    22  	_ Storage = (*TestStorage)(nil)
    23  )
    24  
    25  type TestStorage struct {
    26  	T                                            *testing.T
    27  	CantGetVtx, CantEdge, CantStopVertexAccepted bool
    28  	GetVtxF                                      func(context.Context, ids.ID) (avalanche.Vertex, error)
    29  	EdgeF                                        func(context.Context) []ids.ID
    30  	StopVertexAcceptedF                          func(context.Context) (bool, error)
    31  }
    32  
    33  func (s *TestStorage) Default(cant bool) {
    34  	s.CantGetVtx = cant
    35  	s.CantEdge = cant
    36  }
    37  
    38  func (s *TestStorage) GetVtx(ctx context.Context, vtxID ids.ID) (avalanche.Vertex, error) {
    39  	if s.GetVtxF != nil {
    40  		return s.GetVtxF(ctx, vtxID)
    41  	}
    42  	if s.CantGetVtx && s.T != nil {
    43  		require.FailNow(s.T, errGet.Error())
    44  	}
    45  	return nil, errGet
    46  }
    47  
    48  func (s *TestStorage) Edge(ctx context.Context) []ids.ID {
    49  	if s.EdgeF != nil {
    50  		return s.EdgeF(ctx)
    51  	}
    52  	if s.CantEdge && s.T != nil {
    53  		require.FailNow(s.T, errEdge.Error())
    54  	}
    55  	return nil
    56  }
    57  
    58  func (s *TestStorage) StopVertexAccepted(ctx context.Context) (bool, error) {
    59  	if s.StopVertexAcceptedF != nil {
    60  		return s.StopVertexAcceptedF(ctx)
    61  	}
    62  	if s.CantStopVertexAccepted && s.T != nil {
    63  		require.FailNow(s.T, errStopVertexAccepted.Error())
    64  	}
    65  	return false, nil
    66  }