From 10ee1521be103c6119c5906e452eeb39ff33b0b4 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Thu, 24 Jan 2019 18:22:27 +0700
Subject: [PATCH] fix race conditions in packet handler map unit tests

---
 packet_handler_map_test.go | 28 ++++++++++++++++++----------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/packet_handler_map_test.go b/packet_handler_map_test.go
index 6d9e83a83..22a412f6d 100644
--- a/packet_handler_map_test.go
+++ b/packet_handler_map_test.go
@@ -44,6 +44,19 @@ var _ = Describe("Packet Handler Map", func() {
 		handler = newPacketHandlerMap(conn, 5, utils.DefaultLogger).(*packetHandlerMap)
 	})
 
+	AfterEach(func() {
+		// delete sessions and the server before closing
+		// They might be mock implementations, and we'd have to register the expected calls before otherwise.
+		handler.mutex.Lock()
+		for connID := range handler.handlers {
+			delete(handler.handlers, connID)
+		}
+		handler.server = nil
+		handler.mutex.Unlock()
+		handler.Close()
+		Eventually(handler.listening).Should(BeClosed())
+	})
+
 	It("closes", func() {
 		getMultiplexer() // make the sync.Once execute
 		// replace the clientMuxer. getClientMultiplexer will now return the MockMultiplexer
@@ -89,11 +102,6 @@ var _ = Describe("Packet Handler Map", func() {
 			conn.dataToRead <- getPacket(connID2)
 			Eventually(handledPacket1).Should(BeClosed())
 			Eventually(handledPacket2).Should(BeClosed())
-
-			// makes the listen go routine return
-			packetHandler1.EXPECT().destroy(gomock.Any()).AnyTimes()
-			packetHandler2.EXPECT().destroy(gomock.Any()).AnyTimes()
-			close(conn.dataToRead)
 		})
 
 		It("drops unparseable packets", func() {
@@ -125,10 +133,14 @@ var _ = Describe("Packet Handler Map", func() {
 			handler.deleteRetiredSessionsAfter = time.Hour
 			connID := protocol.ConnectionID{1, 2, 3, 4, 5, 6, 7, 8}
 			packetHandler := NewMockPacketHandler(mockCtrl)
-			packetHandler.EXPECT().handlePacket(gomock.Any())
+			handled := make(chan struct{})
+			packetHandler.EXPECT().handlePacket(gomock.Any()).Do(func(p *receivedPacket) {
+				close(handled)
+			})
 			handler.Add(connID, packetHandler)
 			handler.Retire(connID)
 			handler.handlePacket(nil, nil, getPacket(connID))
+			Eventually(handled).Should(BeClosed())
 		})
 
 		It("drops packets for unknown receivers", func() {
@@ -200,10 +212,6 @@ var _ = Describe("Packet Handler Map", func() {
 					Expect(p.rcvTime).To(BeTemporally("~", now, scaleDuration(20*time.Millisecond)))
 					Expect(p.buffer.refCount).To(Equal(3))
 				}
-
-				// makes the listen go routine return
-				packetHandler.EXPECT().destroy(gomock.Any()).AnyTimes()
-				close(conn.dataToRead)
 			})
 
 			It("ignores coalesced packet parts if the connection IDs don't match", func() {