github.com/celestiaorg/celestia-node@v0.15.0-beta.1/share/p2p/shrexsub/pubsub_test.go (about)

     1  package shrexsub
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	pubsub "github.com/libp2p/go-libp2p-pubsub"
     9  	"github.com/libp2p/go-libp2p/core/peer"
    10  	mocknet "github.com/libp2p/go-libp2p/p2p/net/mock"
    11  	"github.com/stretchr/testify/require"
    12  	"github.com/tendermint/tendermint/libs/rand"
    13  
    14  	pb "github.com/celestiaorg/celestia-node/share/p2p/shrexsub/pb"
    15  )
    16  
    17  func TestPubSub(t *testing.T) {
    18  	h, err := mocknet.FullMeshConnected(2)
    19  	require.NoError(t, err)
    20  	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
    21  	t.Cleanup(cancel)
    22  
    23  	pSub1, err := NewPubSub(ctx, h.Hosts()[0], "test")
    24  	require.NoError(t, err)
    25  
    26  	pSub2, err := NewPubSub(ctx, h.Hosts()[1], "test")
    27  	require.NoError(t, err)
    28  	err = pSub2.AddValidator(
    29  		func(ctx context.Context, p peer.ID, n Notification) pubsub.ValidationResult {
    30  			// only testing shrexsub validation here
    31  			return pubsub.ValidationAccept
    32  		},
    33  	)
    34  	require.NoError(t, err)
    35  
    36  	require.NoError(t, pSub1.Start(ctx))
    37  	require.NoError(t, pSub2.Start(ctx))
    38  
    39  	subs, err := pSub2.Subscribe()
    40  	require.NoError(t, err)
    41  
    42  	var tests = []struct {
    43  		name        string
    44  		notif       Notification
    45  		errExpected bool
    46  	}{
    47  		{
    48  			name: "valid height, valid hash",
    49  			notif: Notification{
    50  				Height:   1,
    51  				DataHash: rand.Bytes(32),
    52  			},
    53  			errExpected: false,
    54  		},
    55  		{
    56  			name: "valid height, invalid hash (<32 bytes)",
    57  			notif: Notification{
    58  				Height:   2,
    59  				DataHash: rand.Bytes(20),
    60  			},
    61  			errExpected: true,
    62  		},
    63  		{
    64  			name: "valid height, invalid hash (>32 bytes)",
    65  			notif: Notification{
    66  				Height:   2,
    67  				DataHash: rand.Bytes(64),
    68  			},
    69  			errExpected: true,
    70  		},
    71  		{
    72  			name: "invalid height, valid hash",
    73  			notif: Notification{
    74  				Height:   0,
    75  				DataHash: rand.Bytes(32),
    76  			},
    77  			errExpected: true,
    78  		},
    79  		{
    80  			name: "invalid height, nil hash",
    81  			notif: Notification{
    82  				Height:   0,
    83  				DataHash: nil,
    84  			},
    85  			errExpected: true,
    86  		},
    87  		{
    88  			name: "valid height, nil hash",
    89  			notif: Notification{
    90  				Height:   30,
    91  				DataHash: nil,
    92  			},
    93  			errExpected: true,
    94  		},
    95  	}
    96  
    97  	for _, tt := range tests {
    98  		t.Run(tt.name, func(t *testing.T) {
    99  			msg := pb.RecentEDSNotification{
   100  				Height:   tt.notif.Height,
   101  				DataHash: tt.notif.DataHash,
   102  			}
   103  			data, err := msg.Marshal()
   104  			require.NoError(t, err)
   105  
   106  			err = pSub1.topic.Publish(ctx, data, pubsub.WithReadiness(pubsub.MinTopicSize(1)))
   107  			require.NoError(t, err)
   108  
   109  			reqCtx, reqCtxCancel := context.WithTimeout(context.Background(), time.Millisecond*100)
   110  			defer reqCtxCancel()
   111  
   112  			got, err := subs.Next(reqCtx)
   113  			if tt.errExpected {
   114  				require.Error(t, err)
   115  				require.ErrorIs(t, err, context.DeadlineExceeded)
   116  				return
   117  			}
   118  			require.NoError(t, err)
   119  			require.NoError(t, err)
   120  			require.Equal(t, tt.notif, got)
   121  		})
   122  	}
   123  
   124  }