From 7b569ba0713efea67f72a3313cde016c6ee66884 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Wed, 3 Apr 2019 01:27:31 +0900
Subject: [PATCH] fix flaky server unit test

---
 server_test.go | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/server_test.go b/server_test.go
index 5fb9d291..fb806794 100644
--- a/server_test.go
+++ b/server_test.go
@@ -507,7 +507,8 @@ var _ = Describe("Server", func() {
 		It("never blocks when calling the onHandshakeComplete callback", func() {
 			const num = 50
 
-			done := make(chan struct{}, num)
+			runs := make(chan struct{}, num)
+			contexts := make(chan struct{}, num)
 			serv.newSession = func(
 				_ connection,
 				runner sessionRunner,
@@ -521,10 +522,9 @@ var _ = Describe("Server", func() {
 				_ protocol.VersionNumber,
 			) (quicSession, error) {
 				sess := NewMockQuicSession(mockCtrl)
-				sess.EXPECT().run().Do(func() {})
-				sess.EXPECT().Context().Return(context.Background())
+				sess.EXPECT().run().Do(func() { runs <- struct{}{} })
+				sess.EXPECT().Context().Do(func() { contexts <- struct{}{} }).Return(context.Background())
 				runner.OnHandshakeComplete(sess)
-				done <- struct{}{}
 				return sess, nil
 			}
 
@@ -534,7 +534,8 @@ var _ = Describe("Server", func() {
 					Expect(err).ToNot(HaveOccurred())
 				}
 			}()
-			Eventually(done).Should(HaveLen(num))
+			Eventually(runs).Should(HaveLen(num))
+			Eventually(contexts).Should(HaveLen(num))
 		})
 	})
 })