-
-
Notifications
You must be signed in to change notification settings - Fork 1.3k
/
conn_test.go
131 lines (114 loc) · 3.06 KB
/
conn_test.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
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
package quic
import (
"errors"
"net"
"time"
. "github.com/onsi/ginkgo"
. "github.com/onsi/gomega"
)
type mockPacketConnWrite struct {
data []byte
to net.Addr
}
type mockPacketConn struct {
addr net.Addr
dataToRead chan []byte
dataReadFrom net.Addr
readErr error
dataWritten chan mockPacketConnWrite
closed bool
}
func newMockPacketConn() *mockPacketConn {
return &mockPacketConn{
addr: &net.UDPAddr{IP: net.IPv6zero, Port: 0x42},
dataToRead: make(chan []byte, 1000),
dataWritten: make(chan mockPacketConnWrite, 1000),
}
}
func (c *mockPacketConn) ReadFrom(b []byte) (int, net.Addr, error) {
if c.readErr != nil {
return 0, nil, c.readErr
}
data, ok := <-c.dataToRead
if !ok {
return 0, nil, errors.New("connection closed")
}
n := copy(b, data)
return n, c.dataReadFrom, nil
}
func (c *mockPacketConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
select {
case c.dataWritten <- mockPacketConnWrite{to: addr, data: b}:
return len(b), nil
default:
panic("channel full")
}
}
func (c *mockPacketConn) Close() error {
if !c.closed {
close(c.dataToRead)
}
c.closed = true
return nil
}
func (c *mockPacketConn) LocalAddr() net.Addr { return c.addr }
func (c *mockPacketConn) SetDeadline(t time.Time) error { panic("not implemented") }
func (c *mockPacketConn) SetReadDeadline(t time.Time) error { panic("not implemented") }
func (c *mockPacketConn) SetWriteDeadline(t time.Time) error { panic("not implemented") }
var _ net.PacketConn = &mockPacketConn{}
var _ = Describe("Connection", func() {
var c *conn
var packetConn *mockPacketConn
BeforeEach(func() {
addr := &net.UDPAddr{
IP: net.IPv4(192, 168, 100, 200),
Port: 1337,
}
packetConn = newMockPacketConn()
c = &conn{
currentAddr: addr,
pconn: packetConn,
}
})
It("writes", func() {
Expect(c.Write([]byte("foobar"))).To(Succeed())
var write mockPacketConnWrite
Expect(packetConn.dataWritten).To(Receive(&write))
Expect(write.to.String()).To(Equal("192.168.100.200:1337"))
Expect(write.data).To(Equal([]byte("foobar")))
})
It("reads", func() {
packetConn.dataToRead <- []byte("foo")
packetConn.dataReadFrom = &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 1336}
p := make([]byte, 10)
n, raddr, err := c.Read(p)
Expect(err).ToNot(HaveOccurred())
Expect(raddr.String()).To(Equal("127.0.0.1:1336"))
Expect(n).To(Equal(3))
Expect(p[0:3]).To(Equal([]byte("foo")))
})
It("gets the remote address", func() {
Expect(c.RemoteAddr().String()).To(Equal("192.168.100.200:1337"))
})
It("gets the local address", func() {
addr := &net.UDPAddr{
IP: net.IPv4(192, 168, 0, 1),
Port: 1234,
}
packetConn.addr = addr
Expect(c.LocalAddr()).To(Equal(addr))
})
It("changes the remote address", func() {
addr := &net.UDPAddr{
IP: net.IPv4(127, 0, 0, 1),
Port: 7331,
}
c.SetCurrentRemoteAddr(addr)
Expect(c.RemoteAddr().String()).To(Equal(addr.String()))
})
It("closes", func() {
err := c.Close()
Expect(err).ToNot(HaveOccurred())
Expect(packetConn.closed).To(BeTrue())
})
})