github.com/ledgerwatch/erigon-lib@v1.0.0/txpool/test_util.go (about)

     1  /*
     2     Copyright 2021 Erigon contributors
     3  
     4     Licensed under the Apache License, Version 2.0 (the "License");
     5     you may not use this file except in compliance with the License.
     6     You may obtain a copy of the License at
     7  
     8         http://www.apache.org/licenses/LICENSE-2.0
     9  
    10     Unless required by applicable law or agreed to in writing, software
    11     distributed under the License is distributed on an "AS IS" BASIS,
    12     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13     See the License for the specific language governing permissions and
    14     limitations under the License.
    15  */
    16  
    17  package txpool
    18  
    19  import (
    20  	"context"
    21  	"sync"
    22  
    23  	"github.com/ledgerwatch/erigon-lib/gointerfaces"
    24  	"github.com/ledgerwatch/erigon-lib/gointerfaces/sentry"
    25  	"github.com/ledgerwatch/erigon-lib/types"
    26  	"google.golang.org/protobuf/types/known/emptypb"
    27  )
    28  
    29  //go:generate moq -stub -out mocks_test.go . Pool
    30  
    31  type MockSentry struct {
    32  	ctx context.Context
    33  	*sentry.SentryServerMock
    34  	streams      map[sentry.MessageId][]sentry.Sentry_MessagesServer
    35  	peersStreams []sentry.Sentry_PeerEventsServer
    36  	StreamWg     sync.WaitGroup
    37  	lock         sync.RWMutex
    38  }
    39  
    40  func NewMockSentry(ctx context.Context) *MockSentry {
    41  	return &MockSentry{ctx: ctx, SentryServerMock: &sentry.SentryServerMock{}}
    42  }
    43  
    44  var peerID types.PeerID = gointerfaces.ConvertHashToH512([64]byte{0x12, 0x34, 0x50}) // "12345"
    45  
    46  func (ms *MockSentry) Send(req *sentry.InboundMessage) (errs []error) {
    47  	ms.lock.RLock()
    48  	defer ms.lock.RUnlock()
    49  	for _, stream := range ms.streams[req.Id] {
    50  		if err := stream.Send(req); err != nil {
    51  			errs = append(errs, err)
    52  		}
    53  	}
    54  	return errs
    55  }
    56  
    57  func (ms *MockSentry) SetStatus(context.Context, *sentry.StatusData) (*sentry.SetStatusReply, error) {
    58  	return &sentry.SetStatusReply{}, nil
    59  }
    60  func (ms *MockSentry) HandShake(context.Context, *emptypb.Empty) (*sentry.HandShakeReply, error) {
    61  	return &sentry.HandShakeReply{Protocol: sentry.Protocol_ETH68}, nil
    62  }
    63  func (ms *MockSentry) Messages(req *sentry.MessagesRequest, stream sentry.Sentry_MessagesServer) error {
    64  	ms.lock.Lock()
    65  	if ms.streams == nil {
    66  		ms.streams = map[sentry.MessageId][]sentry.Sentry_MessagesServer{}
    67  	}
    68  	for _, id := range req.Ids {
    69  		ms.streams[id] = append(ms.streams[id], stream)
    70  	}
    71  	ms.lock.Unlock()
    72  	ms.StreamWg.Done()
    73  	select {
    74  	case <-ms.ctx.Done():
    75  		return nil
    76  	case <-stream.Context().Done():
    77  		return nil
    78  	}
    79  }
    80  
    81  func (ms *MockSentry) PeerEvents(req *sentry.PeerEventsRequest, stream sentry.Sentry_PeerEventsServer) error {
    82  	ms.lock.Lock()
    83  	ms.peersStreams = append(ms.peersStreams, stream)
    84  	ms.lock.Unlock()
    85  	ms.StreamWg.Done()
    86  	select {
    87  	case <-ms.ctx.Done():
    88  		return nil
    89  	case <-stream.Context().Done():
    90  		return nil
    91  	}
    92  }
    93  
    94  func toHashes(h ...byte) (out types.Hashes) {
    95  	for i := range h {
    96  		hash := [32]byte{h[i]}
    97  		out = append(out, hash[:]...)
    98  	}
    99  	return out
   100  }
   101  
   102  func testRlps(num int) [][]byte {
   103  	rlps := make([][]byte, num)
   104  	for i := 0; i < num; i++ {
   105  		rlps[i] = []byte{1}
   106  	}
   107  	return rlps
   108  }
   109  
   110  func toPeerIDs(h ...byte) (out []types.PeerID) {
   111  	for i := range h {
   112  		hash := [64]byte{h[i]}
   113  		out = append(out, gointerfaces.ConvertHashToH512(hash))
   114  	}
   115  	return out
   116  }