From 9b83ac230b44001f3b6fe38fe96a6ae77ace8287 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Mon, 29 Jan 2024 11:26:03 +0700
Subject: [PATCH] fix flaky handshake limiting test (#4281)

---
 server_test.go | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/server_test.go b/server_test.go
index 802488a7..7679034a 100644
--- a/server_test.go
+++ b/server_test.go
@@ -664,6 +664,8 @@ var _ = Describe("Server", func() {
 
 				handshakeChan := make(chan struct{})
 				connChan := make(chan *MockQUICConn, 1)
+				var wg sync.WaitGroup
+				wg.Add(2 * limit)
 				serv.newConn = func(
 					_ sendConn,
 					runner connRunner,
@@ -687,7 +689,7 @@ var _ = Describe("Server", func() {
 					conn.EXPECT().handlePacket(gomock.Any())
 					conn.EXPECT().run()
 					conn.EXPECT().Context().Return(context.Background())
-					conn.EXPECT().HandshakeComplete().Return(handshakeChan)
+					conn.EXPECT().HandshakeComplete().Return(handshakeChan).Do(func() <-chan struct{} { wg.Done(); return nil })
 					return conn
 				}
 
@@ -727,6 +729,7 @@ var _ = Describe("Server", func() {
 					connChan <- conn
 					serv.handlePacket(getInitialWithRandomDestConnID())
 				}
+				wg.Wait()
 			})
 
 			It("limits the number of total handshakes", func() {