From c2aeb5b3d15d3ece46579aee317bf20540e9e9f6 Mon Sep 17 00:00:00 2001
From: xtaci <daniel820313@gmail.com>
Date: Thu, 8 Sep 2016 15:54:36 +0800
Subject: [PATCH] update test

---
 session_test.go | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/session_test.go b/session_test.go
index be81c9d..3ec429a 100644
--- a/session_test.go
+++ b/session_test.go
@@ -488,7 +488,6 @@ func BenchmarkConnSmux(b *testing.B) {
 	bench(b, cs, ss)
 }
 
-/*
 func BenchmarkConnTCP(b *testing.B) {
 	cs, ss, err := getTCPConnectionPair()
 	if err != nil {
@@ -498,7 +497,6 @@ func BenchmarkConnTCP(b *testing.B) {
 	defer ss.Close()
 	bench(b, cs, ss)
 }
-*/
 
 func getSmuxStreamPair() (*Stream, *Stream, error) {
 	c1, c2, err := getTCPConnectionPair()
@@ -562,10 +560,25 @@ func getTCPConnectionPair() (net.Conn, net.Conn, error) {
 
 func bench(b *testing.B, rd io.Reader, wr io.Writer) {
 	buf := make([]byte, 128*1024)
+	buf2 := make([]byte, 128*1024)
 	b.SetBytes(128 * 1024)
 	b.ResetTimer()
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		count := 0
+		for {
+			n, _ := rd.Read(buf2)
+			count += n
+			if count == 128*1024*b.N {
+				return
+			}
+		}
+	}()
 	for i := 0; i < b.N; i++ {
 		wr.Write(buf)
-		rd.Read(buf)
 	}
+	wg.Wait()
 }