From a97f5759c0c07d7d1244f400c200e0f6273e6633 Mon Sep 17 00:00:00 2001
From: Andy Wang <cbeuw.andy@gmail.com>
Date: Wed, 23 Dec 2020 22:34:12 +0000
Subject: [PATCH] Fix race condition in tests

---
 internal/multiplex/mux_test.go    | 12 +++++++-----
 internal/test/integration_test.go | 10 ++++------
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/internal/multiplex/mux_test.go b/internal/multiplex/mux_test.go
index c8c60f4..02b7721 100644
--- a/internal/multiplex/mux_test.go
+++ b/internal/multiplex/mux_test.go
@@ -10,6 +10,7 @@ import (
 	"net"
 	"sync"
 	"testing"
+	"time"
 )
 
 func serveEcho(l net.Listener) {
@@ -19,13 +20,13 @@ func serveEcho(l net.Listener) {
 			// TODO: pass the error back
 			return
 		}
-		go func() {
+		go func(conn net.Conn) {
 			_, err := io.Copy(conn, conn)
 			if err != nil {
 				// TODO: pass the error back
 				return
 			}
-		}()
+		}(conn)
 	}
 }
 
@@ -65,18 +66,19 @@ func makeSessionPair(numConn int) (*Session, *Session, []*connPair) {
 
 func runEchoTest(t *testing.T, conns []net.Conn, msgLen int) {
 	var wg sync.WaitGroup
+	testData := make([]byte, msgLen)
+	rand.Read(testData)
+
 	for _, conn := range conns {
 		wg.Add(1)
 		go func(conn net.Conn) {
-			testData := make([]byte, msgLen)
-			rand.Read(testData)
-
 			n, err := conn.Write(testData)
 			if n != msgLen {
 				t.Fatalf("written only %v, err %v", n, err)
 			}
 
 			recvBuf := make([]byte, msgLen)
+			conn.SetReadDeadline(time.Now().Add(time.Second))
 			_, err = io.ReadFull(conn, recvBuf)
 			if err != nil {
 				t.Fatalf("failed to read back: %v", err)
diff --git a/internal/test/integration_test.go b/internal/test/integration_test.go
index 5812ba2..9f6c252 100644
--- a/internal/test/integration_test.go
+++ b/internal/test/integration_test.go
@@ -30,15 +30,14 @@ func serveTCPEcho(l net.Listener) {
 			log.Error(err)
 			return
 		}
-		go func() {
-			conn := conn
+		go func(conn net.Conn) {
 			_, err := io.Copy(conn, conn)
 			if err != nil {
 				conn.Close()
 				log.Error(err)
 				return
 			}
-		}()
+		}(conn)
 	}
 }
 
@@ -50,8 +49,7 @@ func serveUDPEcho(listener *connutil.PipeListener) {
 			return
 		}
 		const bufSize = 32 * 1024
-		go func() {
-			conn := conn
+		go func(conn net.PacketConn) {
 			defer conn.Close()
 			buf := make([]byte, bufSize)
 			for {
@@ -70,7 +68,7 @@ func serveUDPEcho(listener *connutil.PipeListener) {
 					return
 				}
 			}
-		}()
+		}(conn)
 	}
 }