forked from ipfs/go-graphsync
-
Notifications
You must be signed in to change notification settings - Fork 0
/
testconnmanager.go
84 lines (74 loc) · 2.47 KB
/
testconnmanager.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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
package testutil
import (
"sync"
"testing"
"github.com/libp2p/go-libp2p/core/peer"
"github.com/stretchr/testify/require"
)
// TestConnManager implements network.ConnManager and allows you to assert
// behavior
type TestConnManager struct {
protectedConnsLk sync.RWMutex
protectedConns map[peer.ID][]string
}
// NewTestConnManager returns a new TestConnManager
func NewTestConnManager() *TestConnManager {
return &TestConnManager{protectedConns: make(map[peer.ID][]string)}
}
// Protect simulates protecting a connection (just records occurence)
func (tcm *TestConnManager) Protect(p peer.ID, tag string) {
tcm.protectedConnsLk.Lock()
defer tcm.protectedConnsLk.Unlock()
for _, tagCmp := range tcm.protectedConns[p] {
if tag == tagCmp {
return
}
}
tcm.protectedConns[p] = append(tcm.protectedConns[p], tag)
}
// Unprotect simulates unprotecting a connection (just records occurence)
func (tcm *TestConnManager) Unprotect(p peer.ID, tag string) bool {
tcm.protectedConnsLk.Lock()
defer tcm.protectedConnsLk.Unlock()
for i, tagCmp := range tcm.protectedConns[p] {
if tag == tagCmp {
tcm.protectedConns[p] = append(tcm.protectedConns[p][:i], tcm.protectedConns[p][i+1:]...)
break
}
}
return len(tcm.protectedConns[p]) > 0
}
// AssertProtected asserts that the connection is protected by at least one tag
func (tcm *TestConnManager) AssertProtected(t testing.TB, p peer.ID) {
t.Helper()
tcm.protectedConnsLk.RLock()
defer tcm.protectedConnsLk.RUnlock()
require.True(t, len(tcm.protectedConns[p]) > 0)
}
// RefuteProtected refutes that a connection has been protect
func (tcm *TestConnManager) RefuteProtected(t testing.TB, p peer.ID) {
t.Helper()
tcm.protectedConnsLk.RLock()
defer tcm.protectedConnsLk.RUnlock()
require.False(t, len(tcm.protectedConns[p]) > 0)
}
// AssertProtectedWithTags verifies the connection is protected with the given
// tags at least
func (tcm *TestConnManager) AssertProtectedWithTags(t testing.TB, p peer.ID, tags ...string) {
t.Helper()
tcm.protectedConnsLk.RLock()
defer tcm.protectedConnsLk.RUnlock()
for _, tag := range tags {
require.Contains(t, tcm.protectedConns[p], tag)
}
}
// RefuteProtectedWithTags verifies the connection is not protected with any of the given
// tags
func (tcm *TestConnManager) RefuteProtectedWithTags(t testing.TB, p peer.ID, tags ...string) {
t.Helper()
tcm.protectedConnsLk.RLock()
defer tcm.protectedConnsLk.RUnlock()
for _, tag := range tags {
require.NotContains(t, tcm.protectedConns[p], tag)
}
}