diff --git a/server_test.go b/server_test.go index 2ce8dc19d..5ed4a498d 100644 --- a/server_test.go +++ b/server_test.go @@ -672,10 +672,17 @@ func testServerCreateConnection(t *testing.T, useRetry bool) { tokenGeneratorKey: tokenGeneratorKey, }) + done := make(chan struct{}, 3) c := NewMockQUICConn(mockCtrl) - c.EXPECT().run() - c.EXPECT().Context().Return(context.Background()) - c.EXPECT().HandshakeComplete().Return(make(chan struct{})) + c.EXPECT().run().Do(func() error { done <- struct{}{}; return nil }) + c.EXPECT().Context().DoAndReturn(func() context.Context { + done <- struct{}{} + return context.Background() + }) + c.EXPECT().HandshakeComplete().DoAndReturn(func() <-chan struct{} { + done <- struct{}{} + return make(chan struct{}) + }) recorder := newConnConstructorRecorder(c) server.newConn = recorder.NewConn @@ -725,6 +732,14 @@ func testServerCreateConnection(t *testing.T, useRetry bool) { assert.Zero(t, args.retrySrcConnID) } + for range 3 { + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("timeout") + } + } + // shutdown c.EXPECT().closeWithTransportError(ConnectionRefused) c.EXPECT().destroy(gomock.Any()).AnyTimes()