mock network
Juan Batiz-Benet committed
Dec 16, 2014 at 11:27 UTC
fd3cd399e51e2cd5bf41c82834eddfef2d4cca29
3 files changed
+397
-3
net/mock/mock.go
new
+291
@@ -0,0 +1,291 @@
1
+// Package mocknet provides a mock net.Network to test with.
2
+package mocknet
3
+
4
+import (
5
+ "io"
6
+ "sync"
7
+
8
+ inet "github.com/jbenet/go-ipfs/net"
9
+ peer "github.com/jbenet/go-ipfs/peer"
10
+
11
+ context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
12
+ ctxgroup "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-ctxgroup"
13
+ ma "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
14
+)
15
+
16
+type Stream struct {
17
+ io.Reader
18
+ io.Writer
19
+ conn *Conn
20
+}
21
+
22
+func (s *Stream) Close() error {
23
+ s.conn.removeStream(s)
24
+ if r, ok := (s.Reader).(io.Closer); ok {
25
+ r.Close()
26
+ }
27
+ if w, ok := (s.Writer).(io.Closer); ok {
28
+ return w.Close()
29
+ }
30
+ return nil
31
+}
32
+
33
+func (s *Stream) Conn() inet.Conn {
34
+ return s.conn
35
+}
36
+
37
+// wire pipe between two network conns. yay io.
38
+func newStreamPair(n1 *Network, p2 peer.Peer) (*Stream, *Stream) {
39
+ p1 := n1.local
40
+ r1, w1 := io.Pipe()
41
+ r2, w2 := io.Pipe()
42
+
43
+ s1 := &Stream{Reader: r1, Writer: w2}
44
+ s2 := &Stream{Reader: r2, Writer: w1}
45
+
46
+ n1.Lock()
47
+ n1.conns[p2].addStream(s1)
48
+ n2 := n1.conns[p2].remote
49
+ n1.Unlock()
50
+
51
+ n2.Lock()
52
+ n2.conns[p1].addStream(s2)
53
+ n2.Unlock()
54
+ n2.handle(s2)
55
+
56
+ return s1, s2
57
+}
58
+
59
+type Conn struct {
60
+ connected bool
61
+ local *Network
62
+ remote *Network
63
+ streams []*Stream
64
+ sync.RWMutex
65
+}
66
+
67
+func (c *Conn) Close() error {
68
+ c.Lock()
69
+ defer c.Unlock()
70
+
71
+ c.connected = false
72
+ for _, s := range c.streams {
73
+ go s.Close()
74
+ }
75
+ c.streams = nil
76
+ return nil
77
+}
78
+
79
+func (c *Conn) addStream(s *Stream) {
80
+ c.Lock()
81
+ defer c.Unlock()
82
+
83
+ s.conn = c
84
+ c.streams = append(c.streams, s)
85
+}
86
+
87
+func (c *Conn) removeStream(s *Stream) {
88
+ c.Lock()
89
+ defer c.Unlock()
90
+
91
+ strs := make([]*Stream, 0, len(c.streams))
92
+ for _, s2 := range c.streams {
93
+ if s2 != s {
94
+ strs = append(strs, s2)
95
+ }
96
+ }
97
+}
98
+
99
+func (c *Conn) NewStreamWithProtocol(pr inet.ProtocolID, p peer.Peer) (inet.Stream, error) {
100
+
101
+ ss, _ := newStreamPair(c.local, p)
102
+
103
+ if err := inet.WriteProtocolHeader(pr, ss); err != nil {
104
+ ss.Close()
105
+ return nil, err
106
+ }
107
+
108
+ return ss, nil
109
+}
110
+
111
+// LocalMultiaddr is the Multiaddr on this side
112
+func (c *Conn) LocalMultiaddr() ma.Multiaddr {
113
+ return nil
114
+}
115
+
116
+// LocalPeer is the Peer on our side of the connection
117
+func (c *Conn) LocalPeer() peer.Peer {
118
+ return c.local.local
119
+}
120
+
121
+// RemoteMultiaddr is the Multiaddr on the remote side
122
+func (c *Conn) RemoteMultiaddr() ma.Multiaddr {
123
+ return nil
124
+}
125
+
126
+// RemotePeer is the Peer on the remote side
127
+func (c *Conn) RemotePeer() peer.Peer {
128
+ return c.remote.local
129
+}
130
+
131
+// network implements the Network interface,
132
+type Network struct {
133
+ local peer.Peer // local peer
134
+ mux inet.Mux // protocol multiplexing
135
+
136
+ conns map[peer.Peer]*Conn
137
+ sync.RWMutex
138
+
139
+ cg ctxgroup.ContextGroup // for Context closing
140
+}
141
+
142
+func MakeNetworks(ctx context.Context, peers []peer.Peer) (nets []*Network, err error) {
143
+ nets = make([]*Network, len(peers))
144
+ for i, p := range peers {
145
+ ps := peer.NewPeerstore()
146
+ nets[i], err = newNetwork(ctx, p, ps)
147
+ if err != nil {
148
+ return nil, err
149
+ }
150
+ }
151
+
152
+ for _, n1 := range nets {
153
+ for _, n2 := range nets {
154
+ if n1 == n2 {
155
+ continue
156
+ }
157
+
158
+ n1.conns[n2.local] = &Conn{local: n1, remote: n2}
159
+ }
160
+ }
161
+
162
+ return nets, nil
163
+}
164
+
165
+// NewNetwork constructs a new Mock network
166
+func newNetwork(ctx context.Context, local peer.Peer, peers peer.Peerstore) (*Network, error) {
167
+
168
+ n := &Network{
169
+ local: local,
170
+ mux: inet.Mux{Handlers: inet.StreamHandlerMap{}},
171
+ cg: ctxgroup.WithContext(ctx),
172
+ conns: map[peer.Peer]*Conn{},
173
+ }
174
+
175
+ n.cg.SetTeardown(n.close)
176
+ return n, nil
177
+}
178
+
179
+func (n *Network) handle(s inet.Stream) {
180
+ go n.mux.Handle(s)
181
+}
182
+
183
+// DialPeer attempts to establish a connection to a given peer.
184
+// Respects the context.
185
+func (n *Network) DialPeer(ctx context.Context, p peer.Peer) error {
186
+ n.Lock()
187
+ defer n.Unlock()
188
+
189
+ n.conns[p].connected = true
190
+ return nil
191
+}
192
+
193
+// CtxGroup returns the network's ContextGroup
194
+func (n *Network) CtxGroup() ctxgroup.ContextGroup {
195
+ return n.cg
196
+}
197
+
198
+// LocalPeer the network's LocalPeer
199
+func (n *Network) LocalPeer() peer.Peer {
200
+ return n.local
201
+}
202
+
203
+// Peers returns the connected peers
204
+func (n *Network) Peers() []peer.Peer {
205
+ n.RLock()
206
+ defer n.RUnlock()
207
+
208
+ peers := make([]peer.Peer, 0, len(n.conns))
209
+ for _, c := range n.conns {
210
+ if c.connected {
211
+ peers = append(peers, c.RemotePeer())
212
+ }
213
+ }
214
+ return peers
215
+}
216
+
217
+// Conns returns the connected peers
218
+func (n *Network) Conns() []inet.Conn {
219
+ n.RLock()
220
+ defer n.RUnlock()
221
+
222
+ out := make([]inet.Conn, 0, len(n.conns))
223
+ for _, c := range n.conns {
224
+ if c.connected {
225
+ out = append(out, c)
226
+ }
227
+ }
228
+ return out
229
+}
230
+
231
+// ClosePeer connection to peer
232
+func (n *Network) ClosePeer(p peer.Peer) error {
233
+ return n.conns[p].Close()
234
+}
235
+
236
+// close is the real teardown function
237
+func (n *Network) close() error {
238
+ for _, c := range n.conns {
239
+ c.Close()
240
+ }
241
+ return nil
242
+}
243
+
244
+// Close calls the ContextCloser func
245
+func (n *Network) Close() error {
246
+ return n.cg.Close()
247
+}
248
+
249
+// BandwidthTotals returns the total amount of bandwidth transferred
250
+func (n *Network) BandwidthTotals() (in uint64, out uint64) {
251
+ // need to implement this. probably best to do it in swarm this time.
252
+ // need a "metrics" object
253
+ return 0, 0
254
+}
255
+
256
+// ListenAddresses returns a list of addresses at which this network listens.
257
+func (n *Network) ListenAddresses() []ma.Multiaddr {
258
+ return []ma.Multiaddr{}
259
+}
260
+
261
+// InterfaceListenAddresses returns a list of addresses at which this network
262
+// listens. It expands "any interface" addresses (/ip4/0.0.0.0, /ip6/::) to
263
+// use the known local interfaces.
264
+func (n *Network) InterfaceListenAddresses() ([]ma.Multiaddr, error) {
265
+ return []ma.Multiaddr{}, nil
266
+}
267
+
268
+// Connectedness returns a state signaling connection capabilities
269
+// For now only returns Connecter || NotConnected. Expand into more later.
270
+func (n *Network) Connectedness(p peer.Peer) inet.Connectedness {
271
+ n.Lock()
272
+ defer n.Unlock()
273
+
274
+ if _, found := n.conns[p]; found && n.conns[p].connected {
275
+ return inet.Connected
276
+ }
277
+ return inet.NotConnected
278
+}
279
+
280
+// NewStream returns a new stream to given peer p.
281
+// If there is no connection to p, attempts to create one.
282
+// If ProtocolID is "", writes no header.
283
+func (c *Network) NewStream(pr inet.ProtocolID, p peer.Peer) (inet.Stream, error) {
284
+ return c.conns[p].NewStreamWithProtocol(pr, p)
285
+}
286
+
287
+// SetHandler sets the protocol handler on the Network's Muxer.
288
+// This operation is threadsafe.
289
+func (n *Network) SetHandler(p inet.ProtocolID, h inet.StreamHandler) {
290
+ n.mux.SetHandler(p, h)
291
+}
net/mock/mock_test.go
new
+103
@@ -0,0 +1,103 @@
1
+package mocknet
2
+
3
+import (
4
+ "bytes"
5
+ "io"
6
+ "testing"
7
+
8
+ inet "github.com/jbenet/go-ipfs/net"
9
+ peer "github.com/jbenet/go-ipfs/peer"
10
+
11
+ context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
12
+ testutil "github.com/jbenet/go-ipfs/util/testutil"
13
+)
14
+
15
+func TestNetworkSetup(t *testing.T) {
16
+
17
+ p1 := testutil.RandPeer()
18
+ p2 := testutil.RandPeer()
19
+ p3 := testutil.RandPeer()
20
+ peers := []peer.Peer{p1, p2, p3}
21
+
22
+ nets, err := MakeNetworks(context.Background(), peers)
23
+ if err != nil {
24
+ t.Fatal(err)
25
+ }
26
+
27
+ // check things
28
+
29
+ if len(nets) != 3 {
30
+ t.Error("nets must be 3")
31
+ }
32
+
33
+ for i, n := range nets {
34
+ if n.local != peers[i] {
35
+ t.Error("peer mismatch")
36
+ }
37
+
38
+ if len(n.conns) != (len(nets) - 1) {
39
+ t.Error("conn mismatch")
40
+ }
41
+
42
+ for _, c := range n.conns {
43
+ if c.remote.local == n.local {
44
+ t.Error("conn to self")
45
+ }
46
+
47
+ if c.remote.conns[n.local] == nil {
48
+ t.Error("conn other side fail")
49
+ }
50
+ if c.remote.conns[n.local].remote.local != n.local {
51
+ t.Error("conn other side fail")
52
+ }
53
+ }
54
+
55
+ }
56
+
57
+}
58
+
59
+func TestStreams(t *testing.T) {
60
+
61
+ p1 := testutil.RandPeer()
62
+ p2 := testutil.RandPeer()
63
+ p3 := testutil.RandPeer()
64
+ peers := []peer.Peer{p1, p2, p3}
65
+
66
+ nets, err := MakeNetworks(context.Background(), peers)
67
+ if err != nil {
68
+ t.Fatal(err)
69
+ }
70
+
71
+ nets[1].SetHandler(inet.ProtocolDHT, func(s inet.Stream) {
72
+ go func() {
73
+ b := make([]byte, 4)
74
+ if _, err := io.ReadFull(s, b); err != nil {
75
+ panic(err)
76
+ }
77
+ if !bytes.Equal(b, []byte("beep")) {
78
+ panic("bytes mismatch")
79
+ }
80
+ if _, err := s.Write([]byte("boop")); err != nil {
81
+ panic(err)
82
+ }
83
+ s.Close()
84
+ }()
85
+ })
86
+
87
+ s, err := nets[0].NewStream(inet.ProtocolDHT, nets[1].local)
88
+ if err != nil {
89
+ t.Fatal(err)
90
+ }
91
+
92
+ if _, err := s.Write([]byte("beep")); err != nil {
93
+ panic(err)
94
+ }
95
+ b := make([]byte, 4)
96
+ if _, err := io.ReadFull(s, b); err != nil {
97
+ panic(err)
98
+ }
99
+ if !bytes.Equal(b, []byte("boop")) {
100
+ panic("bytes mismatch 2")
101
+ }
102
+
103
+}
net/net.go
+3
-3
@@ -52,7 +52,7 @@ func (c *conn_) NewStreamWithProtocol(pr ProtocolID, p peer.Peer) (Stream, error
52
53
ss := (*stream)(s)
54
55
- if err := writeProtocolHeader(pr, ss); err != nil {
55
+ if err := WriteProtocolHeader(pr, ss); err != nil {
56
ss.Close()
57
return nil, err
58
}
@@ -209,7 +209,7 @@ func (c *network) NewStream(pr ProtocolID, p peer.Peer) (Stream, error) {
209
210
ss := (*stream)(s)
211
212
- if err := writeProtocolHeader(pr, ss); err != nil {
212
+ if err := WriteProtocolHeader(pr, ss); err != nil {
213
ss.Close()
214
return nil, err
215
}
@@ -223,7 +223,7 @@ func (n *network) SetHandler(p ProtocolID, h StreamHandler) {
223
n.mux.SetHandler(p, h)
224
}
225
226
-func writeProtocolHeader(pr ProtocolID, s Stream) error {
226
+func WriteProtocolHeader(pr ProtocolID, s Stream) error {
227
if pr != "" { // only write proper protocol headers
228
if err := WriteLengthPrefix(s, string(pr)); err != nil {
229
return err