package sessionpeermanager import ( "sync" "testing" "github.com/ipfs/go-bitswap/internal/testutil" peer "github.com/libp2p/go-libp2p-core/peer" ) type fakePeerTagger struct { lk sync.Mutex taggedPeers []peer.ID wait sync.WaitGroup } func (fpt *fakePeerTagger) TagPeer(p peer.ID, tag string, n int) { fpt.wait.Add(1) fpt.lk.Lock() defer fpt.lk.Unlock() fpt.taggedPeers = append(fpt.taggedPeers, p) } func (fpt *fakePeerTagger) UntagPeer(p peer.ID, tag string) { defer fpt.wait.Done() fpt.lk.Lock() defer fpt.lk.Unlock() for i := 0; i < len(fpt.taggedPeers); i++ { if fpt.taggedPeers[i] == p { fpt.taggedPeers[i] = fpt.taggedPeers[len(fpt.taggedPeers)-1] fpt.taggedPeers = fpt.taggedPeers[:len(fpt.taggedPeers)-1] return } } } func TestAddPeers(t *testing.T) { peers := testutil.GeneratePeers(2) spm := New(1, &fakePeerTagger{}) isNew := spm.AddPeer(peers[0]) if !isNew { t.Fatal("Expected peer to be new") } isNew = spm.AddPeer(peers[0]) if isNew { t.Fatal("Expected peer to no longer be new") } isNew = spm.AddPeer(peers[1]) if !isNew { t.Fatal("Expected peer to be new") } } func TestRemovePeers(t *testing.T) { peers := testutil.GeneratePeers(2) spm := New(1, &fakePeerTagger{}) existed := spm.RemovePeer(peers[0]) if existed { t.Fatal("Expected peer not to exist") } spm.AddPeer(peers[0]) spm.AddPeer(peers[1]) existed = spm.RemovePeer(peers[0]) if !existed { t.Fatal("Expected peer to exist") } existed = spm.RemovePeer(peers[1]) if !existed { t.Fatal("Expected peer to exist") } existed = spm.RemovePeer(peers[0]) if existed { t.Fatal("Expected peer not to have existed") } } func TestHasPeers(t *testing.T) { peers := testutil.GeneratePeers(2) spm := New(1, &fakePeerTagger{}) if spm.HasPeers() { t.Fatal("Expected not to have peers yet") } spm.AddPeer(peers[0]) if !spm.HasPeers() { t.Fatal("Expected to have peers") } spm.AddPeer(peers[1]) if !spm.HasPeers() { t.Fatal("Expected to have peers") } spm.RemovePeer(peers[0]) if !spm.HasPeers() { t.Fatal("Expected to have peers") } spm.RemovePeer(peers[1]) if spm.HasPeers() { t.Fatal("Expected to no longer have peers") } } func TestHasPeer(t *testing.T) { peers := testutil.GeneratePeers(2) spm := New(1, &fakePeerTagger{}) if spm.HasPeer(peers[0]) { t.Fatal("Expected not to have peer yet") } spm.AddPeer(peers[0]) if !spm.HasPeer(peers[0]) { t.Fatal("Expected to have peer") } spm.AddPeer(peers[1]) if !spm.HasPeer(peers[1]) { t.Fatal("Expected to have peer") } spm.RemovePeer(peers[0]) if spm.HasPeer(peers[0]) { t.Fatal("Expected not to have peer") } if !spm.HasPeer(peers[1]) { t.Fatal("Expected to have peer") } } func TestPeers(t *testing.T) { peers := testutil.GeneratePeers(2) spm := New(1, &fakePeerTagger{}) if len(spm.Peers()) > 0 { t.Fatal("Expected not to have peers yet") } spm.AddPeer(peers[0]) if len(spm.Peers()) != 1 { t.Fatal("Expected to have one peer") } spm.AddPeer(peers[1]) if len(spm.Peers()) != 2 { t.Fatal("Expected to have two peers") } spm.RemovePeer(peers[0]) if len(spm.Peers()) != 1 { t.Fatal("Expected to have one peer") } } func TestPeersDiscovered(t *testing.T) { peers := testutil.GeneratePeers(2) spm := New(1, &fakePeerTagger{}) if spm.PeersDiscovered() { t.Fatal("Expected not to have discovered peers yet") } spm.AddPeer(peers[0]) if !spm.PeersDiscovered() { t.Fatal("Expected to have discovered peers") } spm.RemovePeer(peers[0]) if !spm.PeersDiscovered() { t.Fatal("Expected to still have discovered peers") } } func TestPeerTagging(t *testing.T) { peers := testutil.GeneratePeers(2) fpt := &fakePeerTagger{} spm := New(1, fpt) spm.AddPeer(peers[0]) if len(fpt.taggedPeers) != 1 { t.Fatal("Expected to have tagged one peer") } spm.AddPeer(peers[0]) if len(fpt.taggedPeers) != 1 { t.Fatal("Expected to have tagged one peer") } spm.AddPeer(peers[1]) if len(fpt.taggedPeers) != 2 { t.Fatal("Expected to have tagged two peers") } spm.RemovePeer(peers[1]) if len(fpt.taggedPeers) != 1 { t.Fatal("Expected to have untagged peer") } } func TestShutdown(t *testing.T) { peers := testutil.GeneratePeers(2) fpt := &fakePeerTagger{} spm := New(1, fpt) spm.AddPeer(peers[0]) spm.AddPeer(peers[1]) if len(fpt.taggedPeers) != 2 { t.Fatal("Expected to have tagged two peers") } spm.Shutdown() if len(fpt.taggedPeers) != 0 { t.Fatal("Expected to have untagged all peers") } }