hook and socket args check fix, tests

1 files changed, 124 insertions, 8 deletions

diff --git a/pkg/test/test.go b/pkg/test/test.go
index 6a2c776..89722fd 100644
--- a/pkg/test/test.go
+++ b/pkg/test/test.go
@@ -1,25 +1,44 @@
 package test
 
 import (
+	"fmt"
+	"io"
+	"net"
 	"os"
 	"path/filepath"
-	"strconv"
+	"strings"
+	"testing"
+	"time"
 
 	"tunnel/pkg/client"
 	"tunnel/pkg/server"
 )
 
-func getSocketPath() string {
-	s := "tunnel.test." + strconv.Itoa(os.Getpid())
+type env struct {
+	*testing.T
+}
+
+func getSocketPath(id string) string {
+	s := fmt.Sprintf("tunnel.%d.test.%s", os.Getpid(), id)
 	return filepath.Join(os.TempDir(), s)
 }
 
-func newClientServer() (*client.Client, *server.Server, error) {
-	socket := getSocketPath()
+type Client struct {
+	*client.Client
+
+	t *testing.T
+}
+
+type Server struct {
+	*server.Server
+}
+
+func newClientServer(t *testing.T) (*Client, *Server) {
+	socket := getSocketPath(t.Name())
 
 	s, err := server.New(socket)
 	if err != nil {
-		return nil, nil, err
+		t.Fatal(err)
 	}
 
 	go s.Serve()
@@ -27,8 +46,105 @@ func newClientServer() (*client.Client, *server.Server, error) {
 	c, err := client.New(socket)
 	if err != nil {
 		s.Stop()
-		return nil, nil, err
+		t.Fatal(err)
+	}
+
+	return &Client{c, t}, &Server{s}
+}
+
+func closeClientServer(c *Client, s *Server) {
+	c.Close()
+	s.Stop()
+}
+
+func (c *Client) Send(format string, args ...interface{}) string {
+	s := fmt.Sprintf(format, args...)
+	t := strings.Split(s, " ")
+
+	r, err := c.Client.Send(t)
+	if err != nil {
+		c.t.Fatal(err)
 	}
 
-	return c, s, nil
+	return r
+}
+
+func (c *Client) Exec(format string, args ...interface{}) {
+	s := c.Send(format, args...)
+	if s != "" {
+		c.t.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 {
+	listen, err := net.Listen(network, address)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return listen
+}
+
+func xDial(t *testing.T, network, address string) net.Conn {
+	d := net.Dialer{Timeout: 100 * time.Millisecond}
+	conn, err := d.Dial(network, address)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return conn
+}
+
+func xAccept(t *testing.T, listen net.Listener) net.Conn {
+	var conn net.Conn
+
+	c := make(chan error, 1)
+	go func() {
+		var err error
+		conn, err = listen.Accept()
+		c <- err
+	}()
+
+	timer := time.NewTimer(100 * time.Millisecond)
+	select {
+	case err := <-c:
+		if err != nil {
+			t.Fatal(err)
+		}
+	case <-timer.C:
+		t.Fatal("accept timeout")
+	}
+
+	return conn
+}
+
+func xWrite(t *testing.T, conn net.Conn, i interface{}) {
+	var buf []byte
+
+	switch v := i.(type) {
+	case string:
+		buf = []byte(v)
+	case []byte:
+		buf = v
+	default:
+		t.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)
+	}
+	conn.SetDeadline(time.Time{})
+}
+
+func xReadFull(t *testing.T, 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)
+	}
+	conn.SetDeadline(time.Time{})
 }