cleanup tests code

1 files changed, 51 insertions, 31 deletions

diff --git a/pkg/test/test.go b/pkg/test/test.go
index 44289c9..6d83ad0 100644
--- a/pkg/test/test.go
+++ b/pkg/test/test.go
@@ -14,10 +14,26 @@ import (
 	"tunnel/pkg/server"
 )
 
-const xData = "Hello, World!"
+const dummy = "Hello, World!"
 
 type env struct {
 	*testing.T
+
+	trash []func()
+}
+
+func newEnv(t *testing.T) *env {
+	return &env{T: t}
+}
+
+func (e *env) Free() {
+	for _, f := range e.trash {
+		f()
+	}
+}
+
+func (e *env) add(f func()) {
+	e.trash = append(e.trash, f)
 }
 
 func getSocketPath(id string) string {
@@ -28,19 +44,19 @@ func getSocketPath(id string) string {
 type Client struct {
 	*client.Client
 
-	t *testing.T
+	e *env
 }
 
 type Server struct {
 	*server.Server
 }
 
-func newClientServer(t *testing.T) (*Client, *Server) {
-	socket := getSocketPath(t.Name())
+func (e *env) newInstance() *Client {
+	socket := getSocketPath(e.Name())
 
 	s, err := server.New(socket)
 	if err != nil {
-		t.Fatal(err)
+		e.Fatal(err)
 	}
 
 	go s.Serve()
@@ -48,15 +64,15 @@ func newClientServer(t *testing.T) (*Client, *Server) {
 	c, err := client.New(socket)
 	if err != nil {
 		s.Stop()
-		t.Fatal(err)
+		e.Fatal(err)
 	}
 
-	return &Client{c, t}, &Server{s}
-}
+	e.add(func() {
+		c.Close()
+		s.Stop()
+	})
 
-func closeClientServer(c *Client, s *Server) {
-	c.Close()
-	s.Stop()
+	return &Client{c, e}
 }
 
 func (c *Client) Send(format string, args ...interface{}) string {
@@ -65,7 +81,7 @@ func (c *Client) Send(format string, args ...interface{}) string {
 
 	r, err := c.Client.Send(t)
 	if err != nil {
-		c.t.Fatal(err)
+		c.e.Fatal(err)
 	}
 
 	return r
@@ -82,34 +98,34 @@ func (c *Client) Set(name string, value interface{}) {
 func (c *Client) Exec(format string, args ...interface{}) {
 	s := c.Send(format, args...)
 	if s != "" {
-		c.t.Fatal(s)
+		c.e.Fatal(s)
 	}
 }
 
-func (c *Client) AddListenTunnel(name string, format string, args ...interface{}) string {
-	t := append([]interface{}{name}, args...)
-	c.Exec("add name %s listen,addr=127.0.0.1:0 "+format, t...)
-	return c.Send("get tunnel.%s.listen", name)
-}
-
-func xListen(t *testing.T, network, address string) net.Listener {
+func (e *env) Listen(network, address string) net.Listener {
 	listen, err := net.Listen(network, address)
 	if err != nil {
-		t.Fatal(err)
+		e.Fatal(err)
 	}
+	e.add(func() {
+		listen.Close()
+	})
 	return listen
 }
 
-func xDial(t *testing.T, network, address string) net.Conn {
+func (e *env) Dial(network, address string) net.Conn {
 	d := net.Dialer{Timeout: 100 * time.Millisecond}
 	conn, err := d.Dial(network, address)
 	if err != nil {
-		t.Fatal(err)
+		e.Fatal(err)
 	}
+	e.add(func() {
+		conn.Close()
+	})
 	return conn
 }
 
-func xAccept(t *testing.T, listen net.Listener) net.Conn {
+func (e *env) Accept(listen net.Listener) net.Conn {
 	var conn net.Conn
 
 	c := make(chan error, 1)
@@ -123,16 +139,20 @@ func xAccept(t *testing.T, listen net.Listener) net.Conn {
 	select {
 	case err := <-c:
 		if err != nil {
-			t.Fatal(err)
+			e.Fatal(err)
 		}
 	case <-timer.C:
-		t.Fatal("accept timeout")
+		e.Fatal("accept timeout")
 	}
 
+	e.add(func() {
+		conn.Close()
+	})
+
 	return conn
 }
 
-func xWrite(t *testing.T, conn net.Conn, i interface{}) {
+func (e *env) Write(conn net.Conn, i interface{}) {
 	var buf []byte
 
 	switch v := i.(type) {
@@ -141,20 +161,20 @@ func xWrite(t *testing.T, conn net.Conn, i interface{}) {
 	case []byte:
 		buf = v
 	default:
-		t.Fatalf("unexpected type %T", i)
+		e.Fatalf("unexpected type %T", i)
 	}
 
 	conn.SetDeadline(time.Now().Add(100 * time.Millisecond))
 	if _, err := conn.Write(buf); err != nil {
-		t.Fatal("write to conn:", err)
+		e.Fatal("write to conn:", err)
 	}
 	conn.SetDeadline(time.Time{})
 }
 
-func xReadFull(t *testing.T, conn net.Conn, buf []byte) {
+func (e *env) ReadFull(conn net.Conn, buf []byte) {
 	conn.SetDeadline(time.Now().Add(100 * time.Millisecond))
 	if _, err := io.ReadFull(conn, buf); err != nil {
-		t.Fatal("read from conn:", err)
+		e.Fatal("read from conn:", err)
 	}
 	conn.SetDeadline(time.Time{})
 }