github.com/MetalBlockchain/subnet-evm@v0.4.9/sync/client/mock_client.go (about)

     1  // (c) 2021-2022, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package statesyncclient
     5  
     6  import (
     7  	"context"
     8  	"fmt"
     9  	"sync/atomic"
    10  
    11  	"github.com/MetalBlockchain/metalgo/codec"
    12  	"github.com/MetalBlockchain/metalgo/ids"
    13  	"github.com/MetalBlockchain/subnet-evm/core/types"
    14  	"github.com/MetalBlockchain/subnet-evm/plugin/evm/message"
    15  	"github.com/MetalBlockchain/subnet-evm/sync/handlers"
    16  	"github.com/ethereum/go-ethereum/common"
    17  	"github.com/ethereum/go-ethereum/rlp"
    18  )
    19  
    20  var (
    21  	_               Client         = &MockClient{}
    22  	mockBlockParser EthBlockParser = &testBlockParser{}
    23  )
    24  
    25  // TODO replace with gomock library
    26  type MockClient struct {
    27  	codec          codec.Manager
    28  	leafsHandler   *handlers.LeafsRequestHandler
    29  	leavesReceived int32
    30  	codesHandler   *handlers.CodeRequestHandler
    31  	codeReceived   int32
    32  	blocksHandler  *handlers.BlockRequestHandler
    33  	blocksReceived int32
    34  	// GetLeafsIntercept is called on every GetLeafs request if set to a non-nil callback.
    35  	// The returned response will be returned by MockClient to the caller.
    36  	GetLeafsIntercept func(req message.LeafsRequest, res message.LeafsResponse) (message.LeafsResponse, error)
    37  	// GetCodesIntercept is called on every GetCode request if set to a non-nil callback.
    38  	// The returned response will be returned by MockClient to the caller.
    39  	GetCodeIntercept func(hashes []common.Hash, codeBytes [][]byte) ([][]byte, error)
    40  	// GetBlocksIntercept is called on every GetBlocks request if set to a non-nil callback.
    41  	// The returned response will be returned by MockClient to the caller.
    42  	GetBlocksIntercept func(blockReq message.BlockRequest, blocks types.Blocks) (types.Blocks, error)
    43  }
    44  
    45  func NewMockClient(
    46  	codec codec.Manager,
    47  	leafHandler *handlers.LeafsRequestHandler,
    48  	codesHandler *handlers.CodeRequestHandler,
    49  	blocksHandler *handlers.BlockRequestHandler,
    50  ) *MockClient {
    51  	return &MockClient{
    52  		codec:         codec,
    53  		leafsHandler:  leafHandler,
    54  		codesHandler:  codesHandler,
    55  		blocksHandler: blocksHandler,
    56  	}
    57  }
    58  
    59  func (ml *MockClient) GetLeafs(ctx context.Context, request message.LeafsRequest) (message.LeafsResponse, error) {
    60  	response, err := ml.leafsHandler.OnLeafsRequest(ctx, ids.GenerateTestNodeID(), 1, request)
    61  	if err != nil {
    62  		return message.LeafsResponse{}, err
    63  	}
    64  
    65  	leafResponseIntf, numLeaves, err := parseLeafsResponse(ml.codec, request, response)
    66  	if err != nil {
    67  		return message.LeafsResponse{}, err
    68  	}
    69  	leafsResponse := leafResponseIntf.(message.LeafsResponse)
    70  	if ml.GetLeafsIntercept != nil {
    71  		leafsResponse, err = ml.GetLeafsIntercept(request, leafsResponse)
    72  	}
    73  	// Increment the number of leaves received by the mock client
    74  	atomic.AddInt32(&ml.leavesReceived, int32(numLeaves))
    75  	return leafsResponse, err
    76  }
    77  
    78  func (ml *MockClient) LeavesReceived() int32 {
    79  	return atomic.LoadInt32(&ml.leavesReceived)
    80  }
    81  
    82  func (ml *MockClient) GetCode(ctx context.Context, hashes []common.Hash) ([][]byte, error) {
    83  	if ml.codesHandler == nil {
    84  		panic("no code handler for mock client")
    85  	}
    86  	request := message.CodeRequest{Hashes: hashes}
    87  	response, err := ml.codesHandler.OnCodeRequest(ctx, ids.GenerateTestNodeID(), 1, request)
    88  	if err != nil {
    89  		return nil, err
    90  	}
    91  
    92  	codeBytesIntf, lenCode, err := parseCode(ml.codec, request, response)
    93  	if err != nil {
    94  		return nil, err
    95  	}
    96  	code := codeBytesIntf.([][]byte)
    97  	if ml.GetCodeIntercept != nil {
    98  		code, err = ml.GetCodeIntercept(hashes, code)
    99  	}
   100  	if err == nil {
   101  		atomic.AddInt32(&ml.codeReceived, int32(lenCode))
   102  	}
   103  	return code, err
   104  }
   105  
   106  func (ml *MockClient) CodeReceived() int32 {
   107  	return atomic.LoadInt32(&ml.codeReceived)
   108  }
   109  
   110  func (ml *MockClient) GetBlocks(ctx context.Context, blockHash common.Hash, height uint64, numParents uint16) ([]*types.Block, error) {
   111  	if ml.blocksHandler == nil {
   112  		panic("no blocks handler for mock client")
   113  	}
   114  	request := message.BlockRequest{
   115  		Hash:    blockHash,
   116  		Height:  height,
   117  		Parents: numParents,
   118  	}
   119  	response, err := ml.blocksHandler.OnBlockRequest(ctx, ids.GenerateTestNodeID(), 1, request)
   120  	if err != nil {
   121  		return nil, err
   122  	}
   123  
   124  	client := &client{blockParser: mockBlockParser} // Hack to avoid duplicate code
   125  	blocksRes, numBlocks, err := client.parseBlocks(ml.codec, request, response)
   126  	if err != nil {
   127  		return nil, err
   128  	}
   129  	blocks := blocksRes.(types.Blocks)
   130  	if ml.GetBlocksIntercept != nil {
   131  		blocks, err = ml.GetBlocksIntercept(request, blocks)
   132  	}
   133  	atomic.AddInt32(&ml.blocksReceived, int32(numBlocks))
   134  	return blocks, err
   135  }
   136  
   137  func (ml *MockClient) BlocksReceived() int32 {
   138  	return atomic.LoadInt32(&ml.blocksReceived)
   139  }
   140  
   141  type testBlockParser struct{}
   142  
   143  func (t *testBlockParser) ParseEthBlock(b []byte) (*types.Block, error) {
   144  	block := new(types.Block)
   145  	if err := rlp.DecodeBytes(b, block); err != nil {
   146  		return nil, fmt.Errorf("%s: %w", errUnmarshalResponse, err)
   147  	}
   148  
   149  	return block, nil
   150  }