/
mock.go
60 lines (53 loc) · 1.2 KB
/
mock.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
package routing
import (
"context"
"net/netip"
"sync"
)
type MockRouter struct {
resolver map[string][]netip.AddrPort
self netip.AddrPort
mx sync.RWMutex
}
func NewMockRouter(resolver map[string][]netip.AddrPort, self netip.AddrPort) *MockRouter {
return &MockRouter{
resolver: resolver,
self: self,
}
}
func (m *MockRouter) Ready() (bool, error) {
m.mx.RLock()
defer m.mx.RUnlock()
return len(m.resolver) > 0, nil
}
func (m *MockRouter) Resolve(ctx context.Context, key string, allowSelf bool, count int) (<-chan netip.AddrPort, error) {
peerCh := make(chan netip.AddrPort, count)
peers, ok := m.resolver[key]
// Not found will look forever until timeout.
if !ok {
return peerCh, nil
}
go func() {
m.mx.RLock()
defer m.mx.RUnlock()
for _, peer := range peers {
peerCh <- peer
}
close(peerCh)
}()
return peerCh, nil
}
func (m *MockRouter) Advertise(ctx context.Context, keys []string) error {
m.mx.Lock()
defer m.mx.Unlock()
for _, key := range keys {
m.resolver[key] = []netip.AddrPort{m.self}
}
return nil
}
func (m *MockRouter) LookupKey(key string) ([]netip.AddrPort, bool) {
m.mx.RLock()
defer m.mx.RUnlock()
v, ok := m.resolver[key]
return v, ok
}