From ebf15aa6bbf431e7c8af4ae373dae699ec9b311b Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Sun, 29 Jan 2017 09:02:19 +0100
Subject: [PATCH] fix test cases

---
 proxy/blackhole/blackhole.go            | 12 +++++++-----
 testing/scenarios/common.go             | 15 +++------------
 testing/scenarios/dynamic_vmess_test.go | 23 +++--------------------
 testing/scenarios/shadowsocks_test.go   | 23 +++--------------------
 4 files changed, 16 insertions(+), 57 deletions(-)

diff --git a/proxy/blackhole/blackhole.go b/proxy/blackhole/blackhole.go
index e8819982..43e3ff2e 100644
--- a/proxy/blackhole/blackhole.go
+++ b/proxy/blackhole/blackhole.go
@@ -3,12 +3,17 @@ package blackhole
 
 import (
 	"context"
+	"errors"
 	"time"
 
 	"v2ray.com/core/common"
 	"v2ray.com/core/transport/ray"
 )
 
+var (
+	errConnectionBlocked = errors.New("Blackhole: connection blocked.")
+)
+
 // Handler is an outbound connection that sliently swallow the entire payload.
 type Handler struct {
 	response ResponseConfig
@@ -28,12 +33,9 @@ func New(ctx context.Context, config *Config) (*Handler, error) {
 // Dispatch implements OutboundHandler.Dispatch().
 func (v *Handler) Process(ctx context.Context, outboundRay ray.OutboundRay) error {
 	v.response.WriteTo(outboundRay.OutboundOutput())
-	// CloseError() will immediately close the connection.
 	// Sleep a little here to make sure the response is sent to client.
-	time.Sleep(time.Millisecond * 500)
-	outboundRay.OutboundInput().CloseError()
-	outboundRay.OutboundOutput().CloseError()
-	return nil
+	time.Sleep(time.Second)
+	return errConnectionBlocked
 }
 
 func init() {
diff --git a/testing/scenarios/common.go b/testing/scenarios/common.go
index 3c2b782e..02a1f0f5 100644
--- a/testing/scenarios/common.go
+++ b/testing/scenarios/common.go
@@ -1,6 +1,7 @@
 package scenarios
 
 import (
+	"io"
 	"net"
 	"time"
 
@@ -29,20 +30,10 @@ func xor(b []byte) []byte {
 
 func readFrom(conn net.Conn, timeout time.Duration, length int) []byte {
 	b := make([]byte, 2048)
-	totalBytes := 0
 	deadline := time.Now().Add(timeout)
 	conn.SetReadDeadline(deadline)
-	for totalBytes < length {
-		if time.Now().After(deadline) {
-			break
-		}
-		n, err := conn.Read(b[totalBytes:])
-		if err != nil {
-			break
-		}
-		totalBytes += n
-	}
-	return b[:totalBytes]
+	n, _ := io.ReadFull(conn, b[:length])
+	return b[:n]
 }
 
 func InitializeServerConfig(config *core.Config) error {
diff --git a/testing/scenarios/dynamic_vmess_test.go b/testing/scenarios/dynamic_vmess_test.go
index a6f66d00..40040041 100644
--- a/testing/scenarios/dynamic_vmess_test.go
+++ b/testing/scenarios/dynamic_vmess_test.go
@@ -1,11 +1,10 @@
 package scenarios
 
 import (
-	"fmt"
 	"net"
 	"testing"
+	"time"
 
-	"v2ray.com/core/common/buf"
 	v2net "v2ray.com/core/common/net"
 	"v2ray.com/core/testing/assert"
 	"v2ray.com/core/testing/servers/tcp"
@@ -42,24 +41,8 @@ func TestDynamicVMess(t *testing.T) {
 		assert.Int(nBytes).Equals(len(payload))
 
 		expectedResponse := "Processed: " + payload
-		finished := false
-		response := buf.New()
-		for {
-			err := response.AppendSupplier(buf.ReadFrom(conn))
-			assert.Error(err).IsNil()
-			if err != nil {
-				break
-			}
-			if response.String() == expectedResponse {
-				finished = true
-				break
-			}
-			if response.Len() > len(expectedResponse) {
-				fmt.Printf("Unexpected response: %v\n", response.Bytes())
-				break
-			}
-		}
-		assert.Bool(finished).IsTrue()
+		response := readFrom(conn, time.Second, len(expectedResponse))
+		assert.String(string(response)).Equals(expectedResponse)
 
 		conn.Close()
 	}
diff --git a/testing/scenarios/shadowsocks_test.go b/testing/scenarios/shadowsocks_test.go
index 1517b763..6bd81f31 100644
--- a/testing/scenarios/shadowsocks_test.go
+++ b/testing/scenarios/shadowsocks_test.go
@@ -1,11 +1,10 @@
 package scenarios
 
 import (
-	"fmt"
 	"net"
 	"testing"
+	"time"
 
-	"v2ray.com/core/common/buf"
 	v2net "v2ray.com/core/common/net"
 	"v2ray.com/core/testing/assert"
 	"v2ray.com/core/testing/servers/tcp"
@@ -41,25 +40,9 @@ func TestShadowsocksTCP(t *testing.T) {
 		assert.Error(err).IsNil()
 		assert.Int(nBytes).Equals(len(payload))
 
-		response := buf.New()
-		finished := false
 		expectedResponse := "Processed: " + payload
-		for {
-			err := response.AppendSupplier(buf.ReadFrom(conn))
-			assert.Error(err).IsNil()
-			if err != nil {
-				break
-			}
-			if response.String() == expectedResponse {
-				finished = true
-				break
-			}
-			if response.Len() > len(expectedResponse) {
-				fmt.Printf("Unexpected response: %v\n", response.Bytes())
-				break
-			}
-		}
-		assert.Bool(finished).IsTrue()
+		response := readFrom(conn, time.Second, len(expectedResponse))
+		assert.String(string(response)).Equals(expectedResponse)
 
 		conn.Close()
 	}