Skip to content

Commit 0474aa9

Browse files
committed
fix: Improve teardown and timeout of peer tests
1 parent 077b02e commit 0474aa9

File tree

1 file changed

+42
-11
lines changed

1 file changed

+42
-11
lines changed

peer/conn_test.go

Lines changed: 42 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,8 @@ func TestConn(t *testing.T) {
9191
// Create a channel that closes on disconnect.
9292
channel, err := server.CreateChannel(context.Background(), "wow", nil)
9393
assert.NoError(t, err)
94+
defer channel.Close()
95+
9496
err = wan.Stop()
9597
require.NoError(t, err)
9698
// Once the connection is marked as disconnected, this
@@ -107,10 +109,13 @@ func TestConn(t *testing.T) {
107109
t.Parallel()
108110
client, server, _ := createPair(t)
109111
exchange(t, client, server)
110-
cch, err := client.CreateChannel(context.Background(), "hello", &peer.ChannelOptions{})
112+
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
113+
defer cancel()
114+
cch, err := client.CreateChannel(ctx, "hello", &peer.ChannelOptions{})
111115
require.NoError(t, err)
116+
defer cch.Close()
112117

113-
sch, err := server.Accept(context.Background())
118+
sch, err := server.Accept(ctx)
114119
require.NoError(t, err)
115120
defer sch.Close()
116121

@@ -123,9 +128,12 @@ func TestConn(t *testing.T) {
123128
t.Parallel()
124129
client, server, wan := createPair(t)
125130
exchange(t, client, server)
126-
cch, err := client.CreateChannel(context.Background(), "hello", &peer.ChannelOptions{})
131+
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
132+
defer cancel()
133+
cch, err := client.CreateChannel(ctx, "hello", &peer.ChannelOptions{})
127134
require.NoError(t, err)
128-
sch, err := server.Accept(context.Background())
135+
defer cch.Close()
136+
sch, err := server.Accept(ctx)
129137
require.NoError(t, err)
130138
defer sch.Close()
131139

@@ -170,13 +178,29 @@ func TestConn(t *testing.T) {
170178
srv, err := net.Listen("tcp", "127.0.0.1:0")
171179
require.NoError(t, err)
172180
defer srv.Close()
181+
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
182+
defer cancel()
173183
go func() {
174-
sch, err := server.Accept(context.Background())
175-
assert.NoError(t, err)
184+
sch, err := server.Accept(ctx)
185+
if err != nil {
186+
assert.NoError(t, err)
187+
return
188+
}
189+
defer sch.Close()
190+
176191
nc2 := sch.NetConn()
192+
defer nc2.Close()
193+
177194
nc1, err := net.Dial("tcp", srv.Addr().String())
178-
assert.NoError(t, err)
195+
if err != nil {
196+
assert.NoError(t, err)
197+
return
198+
}
199+
defer nc1.Close()
200+
179201
go func() {
202+
defer nc1.Close()
203+
defer nc2.Close()
180204
_, _ = io.Copy(nc1, nc2)
181205
}()
182206
_, _ = io.Copy(nc2, nc1)
@@ -204,7 +228,7 @@ func TestConn(t *testing.T) {
204228
c := http.Client{
205229
Transport: defaultTransport,
206230
}
207-
req, err := http.NewRequestWithContext(context.Background(), "GET", "http://localhost/", nil)
231+
req, err := http.NewRequestWithContext(ctx, "GET", "http://localhost/", nil)
208232
require.NoError(t, err)
209233
resp, err := c.Do(req)
210234
require.NoError(t, err)
@@ -272,14 +296,21 @@ func TestConn(t *testing.T) {
272296
t.Parallel()
273297
client, server, _ := createPair(t)
274298
exchange(t, client, server)
299+
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
300+
defer cancel()
275301
go func() {
276-
channel, err := client.CreateChannel(context.Background(), "test", nil)
277-
assert.NoError(t, err)
302+
channel, err := client.CreateChannel(ctx, "test", nil)
303+
if err != nil {
304+
assert.NoError(t, err)
305+
return
306+
}
307+
defer channel.Close()
278308
_, err = channel.Write([]byte{1, 2})
279309
assert.NoError(t, err)
280310
}()
281-
channel, err := server.Accept(context.Background())
311+
channel, err := server.Accept(ctx)
282312
require.NoError(t, err)
313+
defer channel.Close()
283314
data := make([]byte, 1)
284315
_, err = channel.Read(data)
285316
require.NoError(t, err)

0 commit comments

Comments
 (0)