Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 6916d34

Browse files
authored
fix: Fix cleanup in test helpers, prefer defer in tests (#3113)
* fix: Change uses of t.Cleanup -> defer in test bodies Mixing t.Cleanup and defer can lead to unexpected order of execution. * fix: Ensure t.Cleanup is not aborted by require * chore: Add helper annotations
1 parent c2cd51d commit 6916d34

14 files changed

+60
-53
lines changed

cli/configssh_test.go

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -109,19 +109,19 @@ func TestConfigSSH(t *testing.T) {
109109
agentCloser := agent.New(agentClient.ListenWorkspaceAgent, &agent.Options{
110110
Logger: slogtest.Make(t, nil),
111111
})
112-
t.Cleanup(func() {
112+
defer func() {
113113
_ = agentCloser.Close()
114-
})
114+
}()
115115
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.LatestBuild.ID)
116116
agentConn, err := client.DialWorkspaceAgent(context.Background(), resources[0].Agents[0].ID, nil)
117117
require.NoError(t, err)
118118
defer agentConn.Close()
119119

120120
listener, err := net.Listen("tcp", "127.0.0.1:0")
121121
require.NoError(t, err)
122-
t.Cleanup(func() {
122+
defer func() {
123123
_ = listener.Close()
124-
})
124+
}()
125125
go func() {
126126
for {
127127
conn, err := listener.Accept()
@@ -134,9 +134,6 @@ func TestConfigSSH(t *testing.T) {
134134
go io.Copy(ssh, conn)
135135
}
136136
}()
137-
t.Cleanup(func() {
138-
_ = listener.Close()
139-
})
140137

141138
sshConfigFile, _ := sshConfigFileNames(t)
142139

cli/logout_test.go

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -152,19 +152,19 @@ func TestLogout(t *testing.T) {
152152
err = os.Chmod(string(config), 0500)
153153
require.NoError(t, err)
154154
}
155-
t.Cleanup(func() {
155+
defer func() {
156156
if runtime.GOOS == "windows" {
157157
// Closing the opened files for cleanup.
158158
err = urlFile.Close()
159-
require.NoError(t, err)
159+
assert.NoError(t, err)
160160
err = sessionFile.Close()
161-
require.NoError(t, err)
161+
assert.NoError(t, err)
162162
} else {
163163
// Setting the permissions back for cleanup.
164-
err = os.Chmod(string(config), 0700)
165-
require.NoError(t, err)
164+
err = os.Chmod(string(config), 0o700)
165+
assert.NoError(t, err)
166166
}
167-
})
167+
}()
168168

169169
logoutChan := make(chan struct{})
170170
logout, _ := clitest.New(t, "logout", "--global-config", string(config))

cli/server_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -249,7 +249,7 @@ func TestServer(t *testing.T) {
249249
snapshot <- ss
250250
})
251251
server := httptest.NewServer(r)
252-
t.Cleanup(server.Close)
252+
defer server.Close()
253253

254254
root, _ := clitest.New(t, "server", "--in-memory", "--address", ":0", "--telemetry", "--telemetry-url", server.URL)
255255
errC := make(chan error)

coderd/coderd_test.go

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,8 @@ func TestAuthorizeAllEndpoints(t *testing.T) {
9494
t.Cleanup(func() { close(tickerCh) })
9595

9696
ctx, cancelFunc := context.WithCancel(context.Background())
97+
defer t.Cleanup(cancelFunc) // Defer to ensure cancelFunc is executed first.
98+
9799
lifecycleExecutor := executor.New(
98100
ctx,
99101
db,
@@ -107,11 +109,15 @@ func TestAuthorizeAllEndpoints(t *testing.T) {
107109
return ctx
108110
}
109111
srv.Start()
112+
t.Cleanup(srv.Close)
110113
serverURL, err := url.Parse(srv.URL)
111114
require.NoError(t, err)
112115

113116
turnServer, err := turnconn.New(nil)
114117
require.NoError(t, err)
118+
t.Cleanup(func() {
119+
_ = turnServer.Close()
120+
})
115121

116122
validator, err := idtoken.NewValidator(ctx, option.WithoutAuthentication())
117123
require.NoError(t, err)
@@ -138,9 +144,6 @@ func TestAuthorizeAllEndpoints(t *testing.T) {
138144

139145
_ = coderdtest.NewProvisionerDaemon(t, coderAPI)
140146
t.Cleanup(func() {
141-
cancelFunc()
142-
_ = turnServer.Close()
143-
srv.Close()
144147
_ = coderAPI.Close()
145148
})
146149

coderd/coderdtest/coderdtest.go

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,8 @@ func newWithCloser(t *testing.T, options *Options) (*codersdk.Client, io.Closer)
143143
}
144144

145145
ctx, cancelFunc := context.WithCancel(context.Background())
146+
defer t.Cleanup(cancelFunc) // Defer to ensure cancelFunc is executed first.
147+
146148
lifecycleExecutor := executor.New(
147149
ctx,
148150
db,
@@ -156,6 +158,8 @@ func newWithCloser(t *testing.T, options *Options) (*codersdk.Client, io.Closer)
156158
return ctx
157159
}
158160
srv.Start()
161+
t.Cleanup(srv.Close)
162+
159163
serverURL, err := url.Parse(srv.URL)
160164
require.NoError(t, err)
161165

@@ -166,6 +170,9 @@ func newWithCloser(t *testing.T, options *Options) (*codersdk.Client, io.Closer)
166170

167171
turnServer, err := turnconn.New(nil)
168172
require.NoError(t, err)
173+
t.Cleanup(func() {
174+
_ = turnServer.Close()
175+
})
169176

170177
// We set the handler after server creation for the access URL.
171178
coderAPI := coderd.New(&coderd.Options{
@@ -188,18 +195,16 @@ func newWithCloser(t *testing.T, options *Options) (*codersdk.Client, io.Closer)
188195
Authorizer: options.Authorizer,
189196
Telemetry: telemetry.NewNoop(),
190197
})
198+
t.Cleanup(func() {
199+
_ = coderAPI.Close()
200+
})
191201
srv.Config.Handler = coderAPI.Handler
192202

193203
var provisionerCloser io.Closer = nopcloser{}
194204
if options.IncludeProvisionerD {
195205
provisionerCloser = NewProvisionerDaemon(t, coderAPI)
196206
}
197-
198207
t.Cleanup(func() {
199-
cancelFunc()
200-
_ = turnServer.Close()
201-
srv.Close()
202-
_ = coderAPI.Close()
203208
_ = provisionerCloser.Close()
204209
})
205210

coderd/provisionerjobs_internal_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ func TestProvisionerJobLogs_Unit(t *testing.T) {
3838
}
3939
api := New(&opts)
4040
server := httptest.NewServer(api.Handler)
41-
t.Cleanup(server.Close)
41+
defer server.Close()
4242
userID := uuid.New()
4343
keyID, keySecret, err := generateAPIKeyIDSecret()
4444
require.NoError(t, err)

coderd/provisionerjobs_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ func TestProvisionerJobLogs(t *testing.T) {
4141
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
4242

4343
ctx, cancelFunc := context.WithCancel(context.Background())
44-
t.Cleanup(cancelFunc)
44+
defer cancelFunc()
4545
logs, err := client.WorkspaceBuildLogsAfter(ctx, workspace.LatestBuild.ID, before)
4646
require.NoError(t, err)
4747
for {
@@ -77,7 +77,7 @@ func TestProvisionerJobLogs(t *testing.T) {
7777
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
7878
before := database.Now()
7979
ctx, cancelFunc := context.WithCancel(context.Background())
80-
t.Cleanup(cancelFunc)
80+
defer cancelFunc()
8181
logs, err := client.WorkspaceBuildLogsAfter(ctx, workspace.LatestBuild.ID, before)
8282
require.NoError(t, err)
8383
for {

coderd/templateversions_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -380,7 +380,7 @@ func TestTemplateVersionLogs(t *testing.T) {
380380
}},
381381
})
382382
ctx, cancelFunc := context.WithCancel(context.Background())
383-
t.Cleanup(cancelFunc)
383+
defer cancelFunc()
384384
logs, err := client.TemplateVersionLogsAfter(ctx, version.ID, before)
385385
require.NoError(t, err)
386386
for {

coderd/workspaceagents_test.go

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -108,15 +108,15 @@ func TestWorkspaceAgentListen(t *testing.T) {
108108
agentCloser := agent.New(agentClient.ListenWorkspaceAgent, &agent.Options{
109109
Logger: slogtest.Make(t, nil).Named("agent").Leveled(slog.LevelDebug),
110110
})
111-
t.Cleanup(func() {
111+
defer func() {
112112
_ = agentCloser.Close()
113-
})
113+
}()
114114
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.LatestBuild.ID)
115115
conn, err := client.DialWorkspaceAgent(context.Background(), resources[0].Agents[0].ID, nil)
116116
require.NoError(t, err)
117-
t.Cleanup(func() {
117+
defer func() {
118118
_ = conn.Close()
119-
})
119+
}()
120120
_, err = conn.Ping()
121121
require.NoError(t, err)
122122
})
@@ -233,9 +233,9 @@ func TestWorkspaceAgentTURN(t *testing.T) {
233233
agentCloser := agent.New(agentClient.ListenWorkspaceAgent, &agent.Options{
234234
Logger: slogtest.Make(t, nil),
235235
})
236-
t.Cleanup(func() {
236+
defer func() {
237237
_ = agentCloser.Close()
238-
})
238+
}()
239239
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.LatestBuild.ID)
240240
opts := &peer.ConnOptions{
241241
Logger: slogtest.Make(t, nil).Named("client"),
@@ -244,9 +244,9 @@ func TestWorkspaceAgentTURN(t *testing.T) {
244244
opts.SettingEngine.SetNetworkTypes([]webrtc.NetworkType{webrtc.NetworkTypeTCP4})
245245
conn, err := client.DialWorkspaceAgent(context.Background(), resources[0].Agents[0].ID, opts)
246246
require.NoError(t, err)
247-
t.Cleanup(func() {
247+
defer func() {
248248
_ = conn.Close()
249-
})
249+
}()
250250
_, err = conn.Ping()
251251
require.NoError(t, err)
252252
}
@@ -294,9 +294,9 @@ func TestWorkspaceAgentPTY(t *testing.T) {
294294
agentCloser := agent.New(agentClient.ListenWorkspaceAgent, &agent.Options{
295295
Logger: slogtest.Make(t, nil),
296296
})
297-
t.Cleanup(func() {
297+
defer func() {
298298
_ = agentCloser.Close()
299-
})
299+
}()
300300
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.LatestBuild.ID)
301301

302302
conn, err := client.WorkspaceAgentReconnectingPTY(context.Background(), resources[0].Agents[0].ID, uuid.New(), 80, 80, "/bin/bash")

coderd/workspacebuilds_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -327,7 +327,7 @@ func TestWorkspaceBuildLogs(t *testing.T) {
327327
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
328328
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
329329
ctx, cancelFunc := context.WithCancel(context.Background())
330-
t.Cleanup(cancelFunc)
330+
defer cancelFunc()
331331
logs, err := client.WorkspaceBuildLogsAfter(ctx, workspace.LatestBuild.ID, before.Add(-time.Hour))
332332
require.NoError(t, err)
333333
for {

coderd/wsconncache/wsconncache_test.go

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -40,9 +40,9 @@ func TestCache(t *testing.T) {
4040
cache := wsconncache.New(func(r *http.Request, id uuid.UUID) (*agent.Conn, error) {
4141
return setupAgent(t, agent.Metadata{}, 0), nil
4242
}, 0)
43-
t.Cleanup(func() {
43+
defer func() {
4444
_ = cache.Close()
45-
})
45+
}()
4646
conn1, _, err := cache.Acquire(httptest.NewRequest(http.MethodGet, "/", nil), uuid.Nil)
4747
require.NoError(t, err)
4848
conn2, _, err := cache.Acquire(httptest.NewRequest(http.MethodGet, "/", nil), uuid.Nil)
@@ -56,9 +56,9 @@ func TestCache(t *testing.T) {
5656
called.Add(1)
5757
return setupAgent(t, agent.Metadata{}, 0), nil
5858
}, time.Microsecond)
59-
t.Cleanup(func() {
59+
defer func() {
6060
_ = cache.Close()
61-
})
61+
}()
6262
conn, release, err := cache.Acquire(httptest.NewRequest(http.MethodGet, "/", nil), uuid.Nil)
6363
require.NoError(t, err)
6464
release()
@@ -74,9 +74,9 @@ func TestCache(t *testing.T) {
7474
cache := wsconncache.New(func(r *http.Request, id uuid.UUID) (*agent.Conn, error) {
7575
return setupAgent(t, agent.Metadata{}, 0), nil
7676
}, time.Microsecond)
77-
t.Cleanup(func() {
77+
defer func() {
7878
_ = cache.Close()
79-
})
79+
}()
8080
conn, release, err := cache.Acquire(httptest.NewRequest(http.MethodGet, "/", nil), uuid.Nil)
8181
require.NoError(t, err)
8282
time.Sleep(time.Millisecond)
@@ -87,9 +87,9 @@ func TestCache(t *testing.T) {
8787
t.Parallel()
8888
random, err := net.Listen("tcp", "127.0.0.1:0")
8989
require.NoError(t, err)
90-
t.Cleanup(func() {
90+
defer func() {
9191
_ = random.Close()
92-
})
92+
}()
9393
tcpAddr, valid := random.Addr().(*net.TCPAddr)
9494
require.True(t, valid)
9595

@@ -98,17 +98,17 @@ func TestCache(t *testing.T) {
9898
w.WriteHeader(http.StatusOK)
9999
}),
100100
}
101-
t.Cleanup(func() {
101+
defer func() {
102102
_ = server.Close()
103-
})
103+
}()
104104
go server.Serve(random)
105105

106106
cache := wsconncache.New(func(r *http.Request, id uuid.UUID) (*agent.Conn, error) {
107107
return setupAgent(t, agent.Metadata{}, 0), nil
108108
}, time.Microsecond)
109-
t.Cleanup(func() {
109+
defer func() {
110110
_ = cache.Close()
111-
})
111+
}()
112112

113113
var wg sync.WaitGroup
114114
// Perform many requests in parallel to simulate
@@ -132,7 +132,7 @@ func TestCache(t *testing.T) {
132132
res := httptest.NewRecorder()
133133
proxy.ServeHTTP(res, req)
134134
res.Result().Body.Close()
135-
require.Equal(t, http.StatusOK, res.Result().StatusCode)
135+
assert.Equal(t, http.StatusOK, res.Result().StatusCode)
136136
}()
137137
}
138138
wg.Wait()

peerbroker/proxy_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,9 @@ func TestProxy(t *testing.T) {
4343
Pubsub: pubsub,
4444
})
4545
require.NoError(t, err)
46-
t.Cleanup(func() {
46+
defer func() {
4747
_ = proxyCloser.Close()
48-
})
48+
}()
4949

5050
var wg sync.WaitGroup
5151
wg.Add(1)

provisionerd/provisionerd_test.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -962,6 +962,7 @@ func createProvisionerd(t *testing.T, dialer provisionerd.Dialer, provisioners p
962962
// Creates a provisionerd protobuf client that's connected
963963
// to the server implementation provided.
964964
func createProvisionerDaemonClient(t *testing.T, server provisionerDaemonTestServer) proto.DRPCProvisionerDaemonClient {
965+
t.Helper()
965966
if server.failJob == nil {
966967
// Default to asserting the error from the failure, otherwise
967968
// it can be lost in tests!
@@ -990,6 +991,7 @@ func createProvisionerDaemonClient(t *testing.T, server provisionerDaemonTestSer
990991
// Creates a provisioner protobuf client that's connected
991992
// to the server implementation provided.
992993
func createProvisionerClient(t *testing.T, server provisionerTestServer) sdkproto.DRPCProvisionerClient {
994+
t.Helper()
993995
clientPipe, serverPipe := provisionersdk.TransportPipe()
994996
t.Cleanup(func() {
995997
_ = clientPipe.Close()

provisionersdk/agent_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ func TestAgentScript(t *testing.T) {
3535
render.Status(r, http.StatusOK)
3636
render.Data(rw, r, content)
3737
}))
38-
t.Cleanup(srv.Close)
38+
defer srv.Close()
3939
srvURL, err := url.Parse(srv.URL)
4040
require.NoError(t, err)
4141

0 commit comments

Comments
 (0)