github.com/prysmaticlabs/prysm@v1.4.4/slasher/beaconclient/receivers_test.go (about)

     1  package beaconclient
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/golang/mock/gomock"
     9  	ethpb "github.com/prysmaticlabs/prysm/proto/eth/v1alpha1"
    10  	"github.com/prysmaticlabs/prysm/shared/event"
    11  	"github.com/prysmaticlabs/prysm/shared/mock"
    12  	"github.com/prysmaticlabs/prysm/shared/slotutil"
    13  	"github.com/prysmaticlabs/prysm/shared/testutil/require"
    14  	testDB "github.com/prysmaticlabs/prysm/slasher/db/testing"
    15  	"google.golang.org/protobuf/types/known/emptypb"
    16  )
    17  
    18  func TestService_ReceiveBlocks(t *testing.T) {
    19  	ctrl := gomock.NewController(t)
    20  	defer ctrl.Finish()
    21  	client := mock.NewMockBeaconChainClient(ctrl)
    22  
    23  	bs := Service{
    24  		cfg:       &Config{BeaconClient: client},
    25  		blockFeed: new(event.Feed),
    26  	}
    27  	stream := mock.NewMockBeaconChain_StreamBlocksClient(ctrl)
    28  	ctx, cancel := context.WithCancel(context.Background())
    29  	client.EXPECT().StreamBlocks(
    30  		gomock.Any(),
    31  		&ethpb.StreamBlocksRequest{},
    32  	).Return(stream, nil)
    33  	stream.EXPECT().Context().Return(ctx).AnyTimes()
    34  	stream.EXPECT().Recv().Return(
    35  		&ethpb.SignedBeaconBlock{},
    36  		nil,
    37  	).Do(func() {
    38  		cancel()
    39  	})
    40  	bs.ReceiveBlocks(ctx)
    41  }
    42  
    43  func TestService_ReceiveAttestations(t *testing.T) {
    44  	ctrl := gomock.NewController(t)
    45  	defer ctrl.Finish()
    46  	client := mock.NewMockBeaconChainClient(ctrl)
    47  
    48  	bs := Service{
    49  		cfg:                         &Config{BeaconClient: client},
    50  		blockFeed:                   new(event.Feed),
    51  		receivedAttestationsBuffer:  make(chan *ethpb.IndexedAttestation, 1),
    52  		collectedAttestationsBuffer: make(chan []*ethpb.IndexedAttestation, 1),
    53  	}
    54  	stream := mock.NewMockBeaconChain_StreamIndexedAttestationsClient(ctrl)
    55  	ctx, cancel := context.WithCancel(context.Background())
    56  	att := &ethpb.IndexedAttestation{
    57  		Data: &ethpb.AttestationData{
    58  			Slot: 5,
    59  		},
    60  	}
    61  	client.EXPECT().StreamIndexedAttestations(
    62  		gomock.Any(),
    63  		&emptypb.Empty{},
    64  	).Return(stream, nil)
    65  	stream.EXPECT().Context().Return(ctx).AnyTimes()
    66  	stream.EXPECT().Recv().Return(
    67  		att,
    68  		nil,
    69  	).Do(func() {
    70  		cancel()
    71  	})
    72  	bs.ReceiveAttestations(ctx)
    73  }
    74  
    75  func TestService_ReceiveAttestations_Batched(t *testing.T) {
    76  	ctrl := gomock.NewController(t)
    77  	defer ctrl.Finish()
    78  	client := mock.NewMockBeaconChainClient(ctrl)
    79  
    80  	bs := Service{
    81  		cfg: &Config{
    82  			BeaconClient: client,
    83  			SlasherDB:    testDB.SetupSlasherDB(t, false),
    84  		},
    85  		blockFeed:                   new(event.Feed),
    86  		attestationFeed:             new(event.Feed),
    87  		receivedAttestationsBuffer:  make(chan *ethpb.IndexedAttestation, 1),
    88  		collectedAttestationsBuffer: make(chan []*ethpb.IndexedAttestation, 1),
    89  	}
    90  	stream := mock.NewMockBeaconChain_StreamIndexedAttestationsClient(ctrl)
    91  	ctx, cancel := context.WithCancel(context.Background())
    92  	att := &ethpb.IndexedAttestation{
    93  		Data: &ethpb.AttestationData{
    94  			Slot: 5,
    95  			Target: &ethpb.Checkpoint{
    96  				Epoch: 5,
    97  				Root:  []byte("test root 1"),
    98  			},
    99  		},
   100  		Signature: []byte{1, 2},
   101  	}
   102  	client.EXPECT().StreamIndexedAttestations(
   103  		gomock.Any(),
   104  		&emptypb.Empty{},
   105  	).Return(stream, nil)
   106  	stream.EXPECT().Context().Return(ctx).AnyTimes()
   107  	stream.EXPECT().Recv().Return(
   108  		att,
   109  		nil,
   110  	).Do(func() {
   111  		// Let a slot pass for the ticker.
   112  		time.Sleep(slotutil.DivideSlotBy(1))
   113  		cancel()
   114  	})
   115  
   116  	go bs.ReceiveAttestations(ctx)
   117  	bs.receivedAttestationsBuffer <- att
   118  	att.Data.Target.Root = []byte("test root 2")
   119  	bs.receivedAttestationsBuffer <- att
   120  	att.Data.Target.Root = []byte("test root 3")
   121  	bs.receivedAttestationsBuffer <- att
   122  	atts := <-bs.collectedAttestationsBuffer
   123  	require.Equal(t, 3, len(atts), "Unexpected number of attestations batched")
   124  }