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

     1  package discovery
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	mocknet "github.com/libp2p/go-libp2p/p2p/net/mock"
     9  	"github.com/stretchr/testify/require"
    10  )
    11  
    12  func TestSet_TryAdd(t *testing.T) {
    13  	m := mocknet.New()
    14  	h, err := m.GenPeer()
    15  	require.NoError(t, err)
    16  
    17  	set := newLimitedSet(1)
    18  	set.Add(h.ID())
    19  	require.True(t, set.Contains(h.ID()))
    20  }
    21  
    22  func TestSet_Remove(t *testing.T) {
    23  	m := mocknet.New()
    24  	h, err := m.GenPeer()
    25  	require.NoError(t, err)
    26  
    27  	set := newLimitedSet(1)
    28  	set.Add(h.ID())
    29  	set.Remove(h.ID())
    30  	require.False(t, set.Contains(h.ID()))
    31  }
    32  
    33  func TestSet_Peers(t *testing.T) {
    34  	m := mocknet.New()
    35  	h1, err := m.GenPeer()
    36  	require.NoError(t, err)
    37  	h2, err := m.GenPeer()
    38  	require.NoError(t, err)
    39  
    40  	set := newLimitedSet(2)
    41  	set.Add(h1.ID())
    42  	set.Add(h2.ID())
    43  
    44  	ctx, cancel := context.WithTimeout(context.Background(), time.Second*1)
    45  	t.Cleanup(cancel)
    46  
    47  	peers, err := set.Peers(ctx)
    48  	require.NoError(t, err)
    49  	require.True(t, len(peers) == 2)
    50  }
    51  
    52  // TestSet_WaitPeers ensures that `Peers` will be unblocked once
    53  // a new peer was discovered.
    54  func TestSet_WaitPeers(t *testing.T) {
    55  	m := mocknet.New()
    56  	h1, err := m.GenPeer()
    57  	require.NoError(t, err)
    58  
    59  	set := newLimitedSet(2)
    60  	go func() {
    61  		time.Sleep(time.Millisecond * 500)
    62  		set.Add(h1.ID())
    63  	}()
    64  
    65  	ctx, cancel := context.WithTimeout(context.Background(), time.Second*2)
    66  	t.Cleanup(cancel)
    67  
    68  	// call `Peers` on empty set will block until a new peer will be discovered
    69  	peers, err := set.Peers(ctx)
    70  	require.NoError(t, err)
    71  	require.True(t, len(peers) == 1)
    72  }
    73  
    74  func TestSet_Size(t *testing.T) {
    75  	m := mocknet.New()
    76  	h1, err := m.GenPeer()
    77  	require.NoError(t, err)
    78  	h2, err := m.GenPeer()
    79  	require.NoError(t, err)
    80  
    81  	set := newLimitedSet(2)
    82  	set.Add(h1.ID())
    83  	set.Add(h2.ID())
    84  	require.EqualValues(t, 2, set.Size())
    85  	set.Remove(h2.ID())
    86  	require.EqualValues(t, 1, set.Size())
    87  }