From 3e737717bd9beea0ee65ba487766e6ff8ff19542 Mon Sep 17 00:00:00 2001
From: Andy Wang <cbeuw.andy@gmail.com>
Date: Wed, 23 Dec 2020 12:25:58 +0000
Subject: [PATCH] Use assert.Eventually to correctly handle more timing
 sensitive tests

---
 internal/multiplex/stream_test.go | 37 +++++++++++++------------------
 1 file changed, 16 insertions(+), 21 deletions(-)

diff --git a/internal/multiplex/stream_test.go b/internal/multiplex/stream_test.go
index c0b86fb..35fae9d 100644
--- a/internal/multiplex/stream_test.go
+++ b/internal/multiplex/stream_test.go
@@ -152,15 +152,6 @@ func TestStream_Close(t *testing.T) {
 			return
 		}
 
-		// we read something to wait for the test frame to reach our recvBuffer.
-		// if it's empty by the point we call stream.Close(), the incoming
-		// frame will be dropped
-		readBuf := make([]byte, len(testPayload))
-		_, err = io.ReadFull(stream, readBuf[:1])
-		if err != nil {
-			t.Errorf("can't read any data before active closing")
-		}
-
 		err = stream.Close()
 		if err != nil {
 			t.Error("failed to actively close stream", err)
@@ -175,10 +166,11 @@ func TestStream_Close(t *testing.T) {
 		}
 		sesh.streamsM.Unlock()
 
-		_, err = io.ReadFull(stream, readBuf[1:])
-		if err != nil {
-			t.Errorf("can't read residual data %v", err)
-		}
+		readBuf := make([]byte, len(testPayload))
+		assert.Eventually(t, func() bool {
+			_, err = io.ReadFull(stream, readBuf)
+			return err == nil
+		}, time.Second, 10*time.Millisecond, "can't read residual data", err)
 		if !bytes.Equal(readBuf, testPayload) {
 			t.Errorf("read wrong data")
 		}
@@ -324,10 +316,12 @@ func TestStream_Read(t *testing.T) {
 				writingEnd.Write(obfsBuf[:i])
 				stream, _ := sesh.Accept()
 				stream.Close()
-				i, err := stream.Read(buf)
-				if err != nil {
-					t.Error("failed to read", err)
-				}
+
+				var err error
+				assert.Eventually(t, func() bool {
+					i, err = stream.Read(buf)
+					return err == nil
+				}, time.Second, 10*time.Millisecond, "failed to read", err)
 				if i != smallPayloadLen {
 					t.Errorf("expected read %v, got %v", smallPayloadLen, i)
 				}
@@ -348,10 +342,11 @@ func TestStream_Read(t *testing.T) {
 				writingEnd.Write(obfsBuf[:i])
 				stream, _ := sesh.Accept()
 				sesh.Close()
-				i, err := stream.Read(buf)
-				if err != nil {
-					t.Error("failed to read", err)
-				}
+				var err error
+				assert.Eventually(t, func() bool {
+					i, err = stream.Read(buf)
+					return err == nil
+				}, time.Second, 10*time.Millisecond, "failed to read", err)
 				if i != smallPayloadLen {
 					t.Errorf("expected read %v, got %v", smallPayloadLen, i)
 				}