Skip to content

Commit 6224422

Browse files
sreyasreya
authored
chore: fix test flake in TestProvisionerd (#9709)
1 parent 45eadfc commit 6224422

File tree

1 file changed

+37
-19
lines changed

1 file changed

+37
-19
lines changed

‎provisionerd/provisionerd_test.go

+37-19
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ func TestProvisionerd(t *testing.T) {
5959
close(done)
6060
})
6161
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
62-
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{}), nil
62+
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{}), nil
6363
}, provisionerd.LocalProvisioners{})
6464
require.NoError(t, closer.Close())
6565
})
@@ -91,7 +91,7 @@ func TestProvisionerd(t *testing.T) {
9191
completeChan := make(chan struct{})
9292
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
9393
acquireJobAttempt := 0
94-
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
94+
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
9595
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
9696
if acquireJobAttempt == 1 {
9797
close(completeChan)
@@ -118,7 +118,7 @@ func TestProvisionerd(t *testing.T) {
118118
var closerMutex sync.Mutex
119119
closerMutex.Lock()
120120
closer = createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
121-
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
121+
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
122122
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
123123
return &proto.AcquiredJob{
124124
JobId: "test",
@@ -174,7 +174,7 @@ func TestProvisionerd(t *testing.T) {
174174
)
175175

176176
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
177-
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
177+
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
178178
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
179179
return &proto.AcquiredJob{
180180
JobId: "test",
@@ -214,7 +214,7 @@ func TestProvisionerd(t *testing.T) {
214214
)
215215

216216
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
217-
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
217+
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
218218
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
219219
return &proto.AcquiredJob{
220220
JobId: "test",
@@ -269,7 +269,7 @@ func TestProvisionerd(t *testing.T) {
269269
)
270270

271271
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
272-
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
272+
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
273273
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
274274
if !didAcquireJob.CAS(false, true) {
275275
completeOnce.Do(func() { close(completeChan) })
@@ -361,7 +361,7 @@ func TestProvisionerd(t *testing.T) {
361361
)
362362

363363
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
364-
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
364+
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
365365
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
366366
if !didAcquireJob.CAS(false, true) {
367367
completeOnce.Do(func() { close(completeChan) })
@@ -441,7 +441,7 @@ func TestProvisionerd(t *testing.T) {
441441
)
442442

443443
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
444-
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
444+
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
445445
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
446446
if !didAcquireJob.CAS(false, true) {
447447
completeOnce.Do(func() { close(completeChan) })
@@ -513,7 +513,7 @@ func TestProvisionerd(t *testing.T) {
513513
)
514514

515515
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
516-
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
516+
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
517517
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
518518
if !didAcquireJob.CAS(false, true) {
519519
completeOnce.Do(func() { close(completeChan) })
@@ -612,7 +612,7 @@ func TestProvisionerd(t *testing.T) {
612612
)
613613

614614
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
615-
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
615+
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
616616
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
617617
if !didAcquireJob.CAS(false, true) {
618618
completeOnce.Do(func() { close(completeChan) })
@@ -677,7 +677,7 @@ func TestProvisionerd(t *testing.T) {
677677
updateChan := make(chan struct{})
678678
completeChan := make(chan struct{})
679679
server := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
680-
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
680+
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
681681
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
682682
return &proto.AcquiredJob{
683683
JobId: "test",
@@ -755,7 +755,7 @@ func TestProvisionerd(t *testing.T) {
755755
updateChan := make(chan struct{})
756756
completeChan := make(chan struct{})
757757
server := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
758-
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
758+
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
759759
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
760760
return &proto.AcquiredJob{
761761
JobId: "test",
@@ -844,7 +844,7 @@ func TestProvisionerd(t *testing.T) {
844844
completeOnce sync.Once
845845
)
846846
server := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
847-
client := createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
847+
client := createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
848848
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
849849
if second.Load() {
850850
return &proto.AcquiredJob{}, nil
@@ -927,7 +927,7 @@ func TestProvisionerd(t *testing.T) {
927927
completeOnce sync.Once
928928
)
929929
server := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
930-
client := createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
930+
client := createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
931931
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
932932
if second.Load() {
933933
completeOnce.Do(func() { close(completeChan) })
@@ -1006,7 +1006,7 @@ func TestProvisionerd(t *testing.T) {
10061006
completeOnce := sync.Once{}
10071007

10081008
server := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
1009-
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
1009+
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
10101010
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
10111011
m.Lock()
10121012
defer m.Unlock()
@@ -1118,7 +1118,7 @@ func createProvisionerd(t *testing.T, dialer provisionerd.Dialer, connector prov
11181118

11191119
// Creates a provisionerd protobuf client that's connected
11201120
// to the server implementation provided.
1121-
func createProvisionerDaemonClient(t *testing.T, done <-chan struct{}, server provisionerDaemonTestServer) proto.DRPCProvisionerDaemonClient {
1121+
func createProvisionerDaemonClient(ctx context.Context, t *testing.T, done <-chan struct{}, server provisionerDaemonTestServer) proto.DRPCProvisionerDaemonClient {
11221122
t.Helper()
11231123
if server.failJob == nil {
11241124
// Default to asserting the error from the failure, otherwise
@@ -1137,7 +1137,7 @@ func createProvisionerDaemonClient(t *testing.T, done <-chan struct{}, server pr
11371137
err := proto.DRPCRegisterProvisionerDaemon(mux, &server)
11381138
require.NoError(t, err)
11391139
srv := drpcserver.New(mux)
1140-
ctx, cancelFunc := context.WithCancel(context.Background())
1140+
ctx, cancelFunc := context.WithCancel(ctx)
11411141
closed := make(chan struct{})
11421142
go func() {
11431143
defer close(closed)
@@ -1148,13 +1148,31 @@ func createProvisionerDaemonClient(t *testing.T, done <-chan struct{}, server pr
11481148
<-closed
11491149
select {
11501150
case <-done:
1151-
t.Error("createProvisionerDaemonClient cleanup after test was done!")
1151+
// It's possible to get unlucky since the dialer is run in a retry in a goroutine.
1152+
// The following can occur:
1153+
// 1. The provisionerd.connect goroutine checks if we're closed prior to attempting to establish a connection
1154+
// with coderd, sees that we're not.
1155+
// 2. A test closes the server.
1156+
// 3. The provisionerd.connect goroutine calls the dialer to establish a connection. This
1157+
// function detects that someone has tried to create a client after the test finishes.
1158+
if ctx.Err() == nil {
1159+
t.Error("createProvisionerDaemonClient cleanup after test was done!")
1160+
}
11521161
default:
11531162
}
11541163
})
11551164
select {
11561165
case <-done:
1157-
t.Error("called createProvisionerDaemonClient after test was done!")
1166+
// It's possible to get unlucky since the dialer is run in a retry in a goroutine.
1167+
// The following can occur:
1168+
// 1. The provisionerd.connect goroutine checks if we're closed prior to attempting to establish a connection
1169+
// with coderd, sees that we're not.
1170+
// 2. A test closes the server.
1171+
// 3. The provisionerd.connect goroutine calls the dialer to establish a connection. This
1172+
// function detects that someone has tried to create a client after the test finishes.
1173+
if ctx.Err() == nil {
1174+
t.Error("createProvisionerDaemonClient cleanup after test was done!")
1175+
}
11581176
default:
11591177
}
11601178
return proto.NewDRPCProvisionerDaemonClient(clientPipe)

0 commit comments

Comments
 (0)