Add TargetListener

1 files changed, 40 insertions, 0 deletions

diff --git a/listener_test.go b/listener_test.go
new file mode 100644
index 0000000..6b713f1
--- /dev/null
+++ b/listener_test.go
@@ -0,0 +1,40 @@
+// Copyright 2017 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package tcpproxy
+
+import (
+	"io"
+	"testing"
+)
+
+func TestListenerAccept(t *testing.T) {
+	tl := new(TargetListener)
+	ch := make(chan interface{}, 1)
+	go func() {
+		for {
+			conn, err := tl.Accept()
+			if err != nil {
+				ch <- err
+				return
+			} else {
+				ch <- conn
+			}
+		}
+	}()
+
+	for i := 0; i < 3; i++ {
+		conn := new(Conn)
+		tl.HandleConn(conn)
+		got := <-ch
+		if got != conn {
+			t.Errorf("Accept conn = %v; want %v", got, conn)
+		}
+	}
+	tl.Close()
+	got := <-ch
+	if got != io.EOF {
+		t.Errorf("Accept error post-Close = %v; want io.EOF", got)
+	}
+}