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