fix flaky client tests

client_test.go

@@ -42,15 +42,21 @@ var _ = Describe("Client", func() {
 		Expect(err).To(MatchError(qerr.PacketTooLarge))
 	})
 
-	It("closes sessions when Close is called", func() {
-		addr, err := net.ResolveUDPAddr("udp", "127.0.0.1:0")
+	It("properly closes the client", func(done Done) {
+		startUDPConn()
+		var stoppedListening bool
+		go func() {
+			err := client.Listen()
+			Expect(err).ToNot(HaveOccurred())
+			stoppedListening = true
+		}()
+
+		err := client.Close()
 		Expect(err).ToNot(HaveOccurred())
-		client.conn, err = net.ListenUDP("udp", addr)
-		Expect(err).ToNot(HaveOccurred())
-		err = client.Close()
-		Expect(err).ToNot(HaveOccurred())
-		Expect(session.closed).To(BeTrue())
+		Eventually(session.closed).Should(BeTrue())
 		Expect(session.closeReason).To(BeNil())
+		Eventually(func() bool { return stoppedListening }).Should(BeTrue())
+		close(done)
 	})
 
 	It("creates new sessions with the right parameters", func() {
@@ -62,6 +68,9 @@ var _ = Describe("Client", func() {
 		Expect(client.session).ToNot(BeNil())
 		Expect(client.session.(*Session).connectionID).To(Equal(client.connectionID))
 		Expect(client.session.(*Session).version).To(Equal(client.version))
+
+		err = client.Close()
+		Expect(err).ToNot(HaveOccurred())
 	})
 
 	Context("handling packets", func() {
@@ -75,11 +84,12 @@ var _ = Describe("Client", func() {
 			serverConn, err := net.DialUDP("udp", nil, client.conn.LocalAddr().(*net.UDPAddr))
 			Expect(err).NotTo(HaveOccurred())
 
+			var stoppedListening bool
 			go func() {
 				defer GinkgoRecover()
-				listenErr := client.Listen()
-				Expect(listenErr).ToNot(HaveOccurred())
-				close(done)
+				_ = client.Listen()
+				// it should continue listening when receiving valid packets
+				stoppedListening = true
 			}()
 
 			Expect(session.packetCount).To(BeZero())
@@ -96,9 +106,11 @@ var _ = Describe("Client", func() {
 
 			Eventually(func() int { return session.packetCount }).Should(Equal(1))
 			Expect(session.closed).To(BeFalse())
+			Eventually(func() bool { return stoppedListening }).Should(BeFalse())
 
 			err = client.Close()
 			Expect(err).ToNot(HaveOccurred())
+			close(done)
 		})
 
 		It("closes the session when encountering an error while handling a packet", func(done Done) {
@@ -109,20 +121,15 @@ var _ = Describe("Client", func() {
 			var listenErr error
 			go func() {
 				defer GinkgoRecover()
-				listenErr = client.Listen()
-				Expect(listenErr).To(HaveOccurred())
-				close(done)
+				_, err = serverConn.Write(bytes.Repeat([]byte{'f'}, 100))
+				Expect(err).ToNot(HaveOccurred())
 			}()
 
-			// cause a PacketTooLarge error
-			_, err = serverConn.Write(bytes.Repeat([]byte{'f'}, 100))
-			Expect(err).ToNot(HaveOccurred())
-
-			Eventually(func() bool { return session.closed }).Should(BeTrue())
+			listenErr = client.Listen()
+			Expect(listenErr).To(HaveOccurred())
+			Eventually(session.closed).Should(BeTrue())
 			Expect(session.closeReason).To(MatchError(listenErr))
-
-			err = client.Close()
-			Expect(err).ToNot(HaveOccurred())
+			close(done)
 		})
 	})
 

server_test.go

@@ -29,8 +29,8 @@ func (s *mockSession) handlePacket(*receivedPacket) {
 
 func (s *mockSession) run() {}
 func (s *mockSession) Close(e error) error {
-	s.closed = true
 	s.closeReason = e
+	s.closed = true
 	return nil
 }