Refactored the test files with helpers to test backend

```
func testRouteToBackend(t *testing.T, front net.Listener, back net.Listener, msg string)
func testNotRouteToBackend(t *testing.T, front net.Listener, back net.Listener, msg string) <-chan bool
```

Author: yangchenyun

Diff for: tcpproxy_test.go

@@ -169,38 +169,90 @@ func testProxy(t *testing.T, front net.Listener) *Proxy {
 	}
 }
 
-func TestProxyAlwaysMatch(t *testing.T) {
-	front := newLocalListener(t)
-	defer front.Close()
-	back := newLocalListener(t)
-	defer back.Close()
+func testRouteToBackendWithExpected(t *testing.T, toFront net.Conn, back net.Listener, msg string, expected string) {
+	io.WriteString(toFront, msg)
+	fromProxy, err := back.Accept()
+	if err != nil {
+		t.Fatal(err)
+	}
 
-	p := testProxy(t, front)
-	p.AddRoute(testFrontAddr, To(back.Addr().String()))
-	if err := p.Start(); err != nil {
+	buf := make([]byte, len(expected))
+	if _, err := io.ReadFull(fromProxy, buf); err != nil {
 		t.Fatal(err)
 	}
+	if string(buf) != expected {
+		t.Fatalf("got %q; want %q", buf, expected)
+	}
+}
 
+func testRouteToBackend(t *testing.T, front net.Listener, back net.Listener, msg string) {
 	toFront, err := net.Dial("tcp", front.Addr().String())
 	if err != nil {
 		t.Fatal(err)
 	}
 	defer toFront.Close()
 
-	fromProxy, err := back.Accept()
+	testRouteToBackendWithExpected(t, toFront, back, msg, msg)
+}
+
+// test the backend is not receiving traffic
+func testNotRouteToBackend(t *testing.T, front net.Listener, back net.Listener, msg string) <-chan bool {
+	done := make(chan bool)
+	toFront, err := net.Dial("tcp", front.Addr().String())
 	if err != nil {
 		t.Fatal(err)
 	}
-	const msg = "message"
-	io.WriteString(toFront, msg)
+	defer toFront.Close()
 
-	buf := make([]byte, len(msg))
-	if _, err := io.ReadFull(fromProxy, buf); err != nil {
+	timeC := time.NewTimer(10 * time.Millisecond).C
+	acceptC := make(chan struct{})
+	go func() {
+		io.WriteString(toFront, msg)
+		fromProxy, err := back.Accept()
+		acceptC <- struct{}{}
+		{
+			if err == nil {
+				buf := make([]byte, len(msg))
+				if _, err := io.ReadFull(fromProxy, buf); err != nil {
+					t.Fatal(err)
+				}
+				t.Fatalf("Expect backend to not receive message, but found %s", string(buf))
+			}
+			err, ok := err.(net.Error)
+			if !ok || !err.Timeout() {
+				t.Fatalf("Expect backend to timeout, but found err: %v", err)
+			}
+		}
+	}()
+	go func() {
+		select {
+		case <-timeC:
+			{
+				done <- true
+			}
+		case <-acceptC:
+			{
+				t.Fatal("Expect backend to not receive message")
+				done <- true
+			}
+		}
+	}()
+	return done
+}
+
+func TestProxyAlwaysMatch(t *testing.T) {
+	front := newLocalListener(t)
+	defer front.Close()
+	back := newLocalListener(t)
+	defer back.Close()
+
+	p := testProxy(t, front)
+	p.AddRoute(testFrontAddr, To(back.Addr().String()))
+	if err := p.Start(); err != nil {
 		t.Fatal(err)
 	}
-	if string(buf) != msg {
-		t.Fatalf("got %q; want %q", buf, msg)
-	}
+
+	testRouteToBackend(t, front, back, "message")
 }
 
 func TestProxyHTTP(t *testing.T) {
@@ -219,27 +271,9 @@ func TestProxyHTTP(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	toFront, err := net.Dial("tcp", front.Addr().String())
-	if err != nil {
-		t.Fatal(err)
-	}
-	defer toFront.Close()
-
-	const msg = "GET / HTTP/1.1\r\nHost: bar.com\r\n\r\n"
-	io.WriteString(toFront, msg)
-
-	fromProxy, err := backBar.Accept()
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	buf := make([]byte, len(msg))
-	if _, err := io.ReadFull(fromProxy, buf); err != nil {
-		t.Fatal(err)
-	}
-	if string(buf) != msg {
-		t.Fatalf("got %q; want %q", buf, msg)
-	}
+	testRouteToBackend(t, front, backBar, "GET / HTTP/1.1\r\nHost: bar.com\r\n\r\n")
+	<-testNotRouteToBackend(t, front, backBar, "GET / HTTP/1.1\r\nHost: boo.com\r\n\r\n")
+	testRouteToBackend(t, front, backFoo, "GET / HTTP/1.1\r\nHost: foo.com\r\n\r\n")
 }
 
 func TestProxySNI(t *testing.T) {
@@ -258,27 +292,32 @@ func TestProxySNI(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	toFront, err := net.Dial("tcp", front.Addr().String())
-	if err != nil {
-		t.Fatal(err)
-	}
-	defer toFront.Close()
+	testRouteToBackend(t, front, backBar, clientHelloRecord(t, "bar.com"))
+	<-testNotRouteToBackend(t, front, backBar, clientHelloRecord(t, "foo.com"))
+	testRouteToBackend(t, front, backFoo, clientHelloRecord(t, "foo.com"))
+}
 
-	msg := clientHelloRecord(t, "bar.com")
-	io.WriteString(toFront, msg)
+func TestProxyRemoveRoute(t *testing.T) {
+	front := newLocalListener(t)
+	defer front.Close()
+	p := testProxy(t, front)
 
-	fromProxy, err := backBar.Accept()
-	if err != nil {
-		t.Fatal(err)
-	}
+	// NOTE: Needs to register testFrontAddr before server starts
+	p.AddSNIRoute(testFrontAddr, "unused.com", noopTarget{})
 
-	buf := make([]byte, len(msg))
-	if _, err := io.ReadFull(fromProxy, buf); err != nil {
+	if err := p.Start(); err != nil {
 		t.Fatal(err)
 	}
-	if string(buf) != msg {
-		t.Fatalf("got %q; want %q", buf, msg)
-	}
+
+	backBar := newLocalListener(t)
+	defer backBar.Close()
+	routeID := p.AddSNIRoute(testFrontAddr, "bar.com", To(backBar.Addr().String()))
+
+	msg := clientHelloRecord(t, "bar.com")
+	testRouteToBackend(t, front, backBar, msg)
+
+	p.RemoveRoute(testFrontAddr, routeID)
+	<-testNotRouteToBackend(t, front, backBar, msg)
 }
 
 func TestProxyPROXYOut(t *testing.T) {
@@ -301,23 +340,8 @@ func TestProxyPROXYOut(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	io.WriteString(toFront, "foo")
-	toFront.Close()
-
-	fromProxy, err := back.Accept()
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	bs, err := ioutil.ReadAll(fromProxy)
-	if err != nil {
-		t.Fatal(err)
-	}
-
 	want := fmt.Sprintf("PROXY TCP4 %s %d %s %d\r\nfoo", toFront.LocalAddr().(*net.TCPAddr).IP, toFront.LocalAddr().(*net.TCPAddr).Port, toFront.RemoteAddr().(*net.TCPAddr).IP, toFront.RemoteAddr().(*net.TCPAddr).Port)
-	if string(bs) != want {
-		t.Fatalf("got %q; want %q", bs, want)
-	}
+	testRouteToBackendWithExpected(t, toFront, back, "foo", want)
 }
 
 type tlsServer struct {