-
Notifications
You must be signed in to change notification settings - Fork 43
/
mock_peer_discoverer.go
62 lines (53 loc) · 1.28 KB
/
mock_peer_discoverer.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
package discv5
import (
"context"
"sync"
"github.com/libp2p/go-libp2p/core/peer"
"github.com/waku-org/go-waku/waku/v2/service"
)
// TestPeerDiscoverer is mock peer discoverer for testing
type TestPeerDiscoverer struct {
sync.RWMutex
peerMap map[peer.ID]struct{}
}
// NewTestPeerDiscoverer is a constructor for TestPeerDiscoverer
func NewTestPeerDiscoverer() *TestPeerDiscoverer {
result := &TestPeerDiscoverer{
peerMap: make(map[peer.ID]struct{}),
}
return result
}
// Subscribe is for subscribing to peer discoverer
func (t *TestPeerDiscoverer) Subscribe(ctx context.Context, ch <-chan service.PeerData) {
go func() {
for {
select {
case <-ctx.Done():
return
case p := <-ch:
t.Lock()
t.peerMap[p.AddrInfo.ID] = struct{}{}
t.Unlock()
}
}
}()
}
// HasPeer is for checking if a peer is present in peer discoverer
func (t *TestPeerDiscoverer) HasPeer(p peer.ID) bool {
t.RLock()
defer t.RUnlock()
_, ok := t.peerMap[p]
return ok
}
// PeerCount is for getting the number of peers in peer discoverer
func (t *TestPeerDiscoverer) PeerCount() int {
t.RLock()
defer t.RUnlock()
return len(t.peerMap)
}
// Clear is for clearing the peer discoverer
func (t *TestPeerDiscoverer) Clear() {
t.Lock()
defer t.Unlock()
t.peerMap = make(map[peer.ID]struct{})
}