swarm_notif_test.go 4.49 KB
Newer Older
1
2
3
4
5
6
package swarm

import (
	"testing"
	"time"

Jeromy's avatar
Jeromy committed
7
8
	ma "gx/ipfs/QmR3JkmZBKYXgNMNsNZawm914455Qof3PEopwuVSeXG7aV/go-multiaddr"
	context "gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"
Juan Batiz-Benet's avatar
Juan Batiz-Benet committed
9

Juan Batiz-Benet's avatar
Juan Batiz-Benet committed
10
11
	inet "github.com/ipfs/go-libp2p/p2p/net"
	peer "github.com/ipfs/go-libp2p/p2p/peer"
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
)

func TestNotifications(t *testing.T) {
	ctx := context.Background()
	swarms := makeSwarms(ctx, t, 5)
	defer func() {
		for _, s := range swarms {
			s.Close()
		}
	}()

	timeout := 5 * time.Second

	// signup notifs
	notifiees := make([]*netNotifiee, len(swarms))
	for i, swarm := range swarms {
		n := newNetNotifiee()
		swarm.Notify(n)
		notifiees[i] = n
	}

	connectSwarms(t, ctx, swarms)

	<-time.After(time.Millisecond)
	// should've gotten 5 by now.

	// test everyone got the correct connection opened calls
	for i, s := range swarms {
		n := notifiees[i]
Jeromy's avatar
Jeromy committed
41
42
43
		notifs := make(map[peer.ID][]inet.Conn)
		for j, s2 := range swarms {
			if i == j {
44
45
46
				continue
			}

Jeromy's avatar
Jeromy committed
47
48
			// this feels a little sketchy, but its probably okay
			for len(s.ConnectionsToPeer(s2.LocalPeer())) != len(notifs[s2.LocalPeer()]) {
Juan Batiz-Benet's avatar
Juan Batiz-Benet committed
49
50
				select {
				case c := <-n.connected:
Jeromy's avatar
Jeromy committed
51
52
					nfp := notifs[c.RemotePeer()]
					notifs[c.RemotePeer()] = append(nfp, c)
Juan Batiz-Benet's avatar
Juan Batiz-Benet committed
53
54
55
56
				case <-time.After(timeout):
					t.Fatal("timeout")
				}
			}
Jeromy's avatar
Jeromy committed
57
58
59
60
61
62
63
		}

		for p, cons := range notifs {
			expect := s.ConnectionsToPeer(p)
			if len(expect) != len(cons) {
				t.Fatal("got different number of connections")
			}
Juan Batiz-Benet's avatar
Juan Batiz-Benet committed
64

Jeromy's avatar
Jeromy committed
65
66
			for _, c := range cons {
				var found bool
Juan Batiz-Benet's avatar
Juan Batiz-Benet committed
67
				for _, c2 := range expect {
Jeromy's avatar
Jeromy committed
68
					if c == c2 {
Juan Batiz-Benet's avatar
Juan Batiz-Benet committed
69
70
						found = true
						break
71
72
					}
				}
Jeromy's avatar
Jeromy committed
73

Juan Batiz-Benet's avatar
Juan Batiz-Benet committed
74
				if !found {
Jeromy's avatar
Jeromy committed
75
					t.Fatal("connection not found!")
Juan Batiz-Benet's avatar
Juan Batiz-Benet committed
76
77
				}
			}
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
		}
	}

	complement := func(c inet.Conn) (*Swarm, *netNotifiee, *Conn) {
		for i, s := range swarms {
			for _, c2 := range s.Connections() {
				if c.LocalMultiaddr().Equal(c2.RemoteMultiaddr()) &&
					c2.LocalMultiaddr().Equal(c.RemoteMultiaddr()) {
					return s, notifiees[i], c2
				}
			}
		}
		t.Fatal("complementary conn not found", c)
		return nil, nil, nil
	}

	testOCStream := func(n *netNotifiee, s inet.Stream) {
		var s2 inet.Stream
		select {
		case s2 = <-n.openedStream:
			t.Log("got notif for opened stream")
		case <-time.After(timeout):
			t.Fatal("timeout")
		}
		if s != s2 {
			t.Fatal("got incorrect stream", s.Conn(), s2.Conn())
		}

		select {
		case s2 = <-n.closedStream:
			t.Log("got notif for closed stream")
		case <-time.After(timeout):
			t.Fatal("timeout")
		}
		if s != s2 {
			t.Fatal("got incorrect stream", s.Conn(), s2.Conn())
		}
	}

	streams := make(chan inet.Stream)
	for _, s := range swarms {
		s.SetStreamHandler(func(s inet.Stream) {
			streams <- s
			s.Close()
		})
	}

	// open a streams in each conn
	for i, s := range swarms {
		for _, c := range s.Connections() {
			_, n2, _ := complement(c)

			st1, err := c.NewStream()
			if err != nil {
				t.Error(err)
			} else {
				st1.Write([]byte("hello"))
				st1.Close()
				testOCStream(notifiees[i], st1)
				st2 := <-streams
				testOCStream(n2, st2)
			}
		}
	}

	// close conns
	for i, s := range swarms {
		n := notifiees[i]
		for _, c := range s.Connections() {
			_, n2, c2 := complement(c)
			c.Close()
			c2.Close()

			var c3, c4 inet.Conn
			select {
			case c3 = <-n.disconnected:
			case <-time.After(timeout):
				t.Fatal("timeout")
			}
			if c != c3 {
				t.Fatal("got incorrect conn", c, c3)
			}

			select {
			case c4 = <-n2.disconnected:
			case <-time.After(timeout):
				t.Fatal("timeout")
			}
			if c2 != c4 {
				t.Fatal("got incorrect conn", c, c2)
			}
		}
	}
}

type netNotifiee struct {
Juan Batiz-Benet's avatar
Juan Batiz-Benet committed
174
175
	listen       chan ma.Multiaddr
	listenClose  chan ma.Multiaddr
176
177
178
179
180
181
182
183
	connected    chan inet.Conn
	disconnected chan inet.Conn
	openedStream chan inet.Stream
	closedStream chan inet.Stream
}

func newNetNotifiee() *netNotifiee {
	return &netNotifiee{
Juan Batiz-Benet's avatar
Juan Batiz-Benet committed
184
185
		listen:       make(chan ma.Multiaddr),
		listenClose:  make(chan ma.Multiaddr),
186
187
188
189
190
191
192
		connected:    make(chan inet.Conn),
		disconnected: make(chan inet.Conn),
		openedStream: make(chan inet.Stream),
		closedStream: make(chan inet.Stream),
	}
}

Juan Batiz-Benet's avatar
Juan Batiz-Benet committed
193
194
195
196
197
198
func (nn *netNotifiee) Listen(n inet.Network, a ma.Multiaddr) {
	nn.listen <- a
}
func (nn *netNotifiee) ListenClose(n inet.Network, a ma.Multiaddr) {
	nn.listenClose <- a
}
199
200
201
202
203
204
205
206
207
208
209
210
func (nn *netNotifiee) Connected(n inet.Network, v inet.Conn) {
	nn.connected <- v
}
func (nn *netNotifiee) Disconnected(n inet.Network, v inet.Conn) {
	nn.disconnected <- v
}
func (nn *netNotifiee) OpenedStream(n inet.Network, v inet.Stream) {
	nn.openedStream <- v
}
func (nn *netNotifiee) ClosedStream(n inet.Network, v inet.Stream) {
	nn.closedStream <- v
}