Extract common test helpers to own files

CHANGELOG.md

@@ -1,6 +1,7 @@
 # [Unreleased]
 
 * Store all the executable `main` packages in `cmd` folder. (#335, @miry)
+* Extract common test helpers to own files. (#336, @miry)
 
 # [2.2.0]
 

go.mod

@@ -6,13 +6,13 @@ require (
 	github.com/gorilla/mux v1.8.0
 	github.com/sirupsen/logrus v1.8.1
 	github.com/urfave/cli/v2 v2.3.0
-	golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1
+	golang.org/x/term v0.0.0-20210927222741-03fcf44c2211
 	gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7
 )
 
 require (
 	github.com/cpuguy83/go-md2man/v2 v2.0.0-20190314233015-f79a8a8ca69d // indirect
 	github.com/russross/blackfriday/v2 v2.0.1 // indirect
 	github.com/shurcooL/sanitized_anchor_name v1.0.0 // indirect
-	golang.org/x/sys v0.0.0-20201119102817-f84b799fce68 // indirect
+	golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1 // indirect
 )

go.sum

@@ -18,10 +18,10 @@ github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXf
 github.com/urfave/cli/v2 v2.3.0 h1:qph92Y649prgesehzOrQjdWyxFOp/QVM+6imKHad91M=
 github.com/urfave/cli/v2 v2.3.0/go.mod h1:LJmUH05zAU44vOAcrfzZQKsZbVcdbOG8rtL3/XcUArI=
 golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20201119102817-f84b799fce68 h1:nxC68pudNYkKU6jWhgrqdreuFiOQWj1Fs7T3VrH4Pjw=
-golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1 h1:v+OssWQX+hTHEmOBgwxdZxK4zHq3yOs8F9J7mk0PY8E=
-golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
+golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1 h1:SrN+KX8Art/Sf4HNj6Zcz06G7VEz+7w9tdXTPOZ7+l4=
+golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/term v0.0.0-20210927222741-03fcf44c2211 h1:JGgROgKl9N8DuW20oFS5gxc+lE67/N3FcwmBPMe7ArY=
+golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ=
 gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw=

proxy_test.go

@@ -4,116 +4,20 @@ import (
 	"bytes"
 	"encoding/hex"
 	"io"
-	"io/ioutil"
 	"net"
 	"testing"
 	"time"
 
-	"github.com/Shopify/toxiproxy/v2"
 	"github.com/sirupsen/logrus"
-	tomb "gopkg.in/tomb.v1"
+
+	"github.com/Shopify/toxiproxy/v2"
+	"github.com/Shopify/toxiproxy/v2/testhelper"
 )
 
 func init() {
 	logrus.SetLevel(logrus.FatalLevel)
 }
 
-func NewTestProxy(name, upstream string) *toxiproxy.Proxy {
-	proxy := toxiproxy.NewProxy()
-
-	proxy.Name = name
-	proxy.Listen = "localhost:0"
-	proxy.Upstream = upstream
-
-	return proxy
-}
-
-func WithTCPServer(t *testing.T, f func(string, chan []byte)) {
-	ln, err := net.Listen("tcp", "localhost:0")
-	if err != nil {
-		t.Fatal("Failed to create TCP server", err)
-	}
-
-	defer ln.Close()
-
-	response := make(chan []byte, 1)
-	tomb := tomb.Tomb{}
-
-	go func() {
-		defer tomb.Done()
-		src, err := ln.Accept()
-		if err != nil {
-			select {
-			case <-tomb.Dying():
-			default:
-				t.Error("Failed to accept client")
-				return
-			}
-			return
-		}
-
-		ln.Close()
-
-		val, err := ioutil.ReadAll(src)
-		if err != nil {
-			t.Error("Failed to read from client")
-			return
-		}
-
-		response <- val
-	}()
-
-	f(ln.Addr().String(), response)
-
-	tomb.Killf("Function body finished")
-	ln.Close()
-	tomb.Wait()
-
-	close(response)
-}
-
-func TestSimpleServer(t *testing.T) {
-	WithTCPServer(t, func(addr string, response chan []byte) {
-		conn, err := net.Dial("tcp", addr)
-		if err != nil {
-			t.Error("Unable to dial TCP server", err)
-		}
-
-		msg := []byte("hello world")
-
-		_, err = conn.Write(msg)
-		if err != nil {
-			t.Error("Failed writing to TCP server", err)
-		}
-
-		err = conn.Close()
-		if err != nil {
-			t.Error("Failed to close TCP connection", err)
-		}
-
-		resp := <-response
-		if !bytes.Equal(resp, msg) {
-			t.Error("Server didn't read bytes from client")
-		}
-	})
-}
-
-func WithTCPProxy(
-	t *testing.T,
-	f func(proxy net.Conn, response chan []byte, proxyServer *toxiproxy.Proxy),
-) {
-	WithTCPServer(t, func(upstream string, response chan []byte) {
-		proxy := NewTestProxy("test", upstream)
-		proxy.Start()
-
-		conn := AssertProxyUp(t, proxy.Listen, true)
-
-		f(conn, response, proxy)
-
-		proxy.Stop()
-	})
-}
-
 func TestProxySimpleMessage(t *testing.T) {
 	WithTCPProxy(t, func(conn net.Conn, response chan []byte, proxy *toxiproxy.Proxy) {
 		msg := []byte("hello world")
@@ -221,7 +125,7 @@ func TestStartTwoProxiesOnSameAddress(t *testing.T) {
 }
 
 func TestStopProxyBeforeStarting(t *testing.T) {
-	WithTCPServer(t, func(upstream string, response chan []byte) {
+	testhelper.WithTCPServer(t, func(upstream string, response chan []byte) {
 		proxy := NewTestProxy("test", upstream)
 		AssertProxyUp(t, proxy.Listen, false)
 
@@ -243,7 +147,7 @@ func TestStopProxyBeforeStarting(t *testing.T) {
 }
 
 func TestProxyUpdate(t *testing.T) {
-	WithTCPServer(t, func(upstream string, response chan []byte) {
+	testhelper.WithTCPServer(t, func(upstream string, response chan []byte) {
 		proxy := NewTestProxy("test", upstream)
 		err := proxy.Start()
 		if err != nil {
@@ -280,7 +184,7 @@ func TestProxyUpdate(t *testing.T) {
 }
 
 func TestRestartFailedToStartProxy(t *testing.T) {
-	WithTCPServer(t, func(upstream string, response chan []byte) {
+	testhelper.WithTCPServer(t, func(upstream string, response chan []byte) {
 		proxy := NewTestProxy("test", upstream)
 		conflict := NewTestProxy("test2", upstream)
 
@@ -309,13 +213,3 @@ func TestRestartFailedToStartProxy(t *testing.T) {
 		AssertProxyUp(t, proxy.Listen, false)
 	})
 }
-
-func AssertProxyUp(t *testing.T, addr string, up bool) net.Conn {
-	conn, err := net.Dial("tcp", addr)
-	if err != nil && up {
-		t.Error("Expected proxy to be up:", err)
-	} else if err == nil && !up {
-		t.Error("Expected proxy to be down")
-	}
-	return conn
-}

testhelper/tcp_server.go

@@ -0,0 +1,68 @@
+package testhelper
+
+import (
+	"io/ioutil"
+	"net"
+	"testing"
+)
+
+func NewTCPServer() (*TCPServer, error) {
+	result := &TCPServer{
+		addr:     "localhost:0",
+		response: make(chan []byte, 1),
+	}
+	err := result.Run()
+	if err != nil {
+		return nil, err
+	}
+	return result, nil
+}
+
+type TCPServer struct {
+	addr     string
+	server   net.Listener
+	response chan []byte
+}
+
+func (server *TCPServer) Run() (err error) {
+	server.server, err = net.Listen("tcp", server.addr)
+	if err != nil {
+		return
+	}
+	server.addr = server.server.Addr().String()
+	return
+}
+
+func (server *TCPServer) handle_connection() (err error) {
+	conn, err := server.server.Accept()
+	if err != nil {
+		return
+	}
+	defer conn.Close()
+
+	val, err := ioutil.ReadAll(conn)
+	if err != nil {
+		return
+	}
+
+	server.response <- val
+	return
+}
+
+func (server *TCPServer) Close() (err error) {
+	return server.server.Close()
+}
+
+func WithTCPServer(t *testing.T, block func(string, chan []byte)) {
+	server, err := NewTCPServer()
+	if err != nil {
+		t.Fatal("Failed to create TCP server", err)
+	}
+	go func(t *testing.T, server *TCPServer) {
+		err := server.handle_connection()
+		if err != nil {
+			t.Error("Failed to handle connection", err)
+		}
+	}(t, server)
+	block(server.addr, server.response)
+}

testhelper/tcp_server_test.go

@@ -0,0 +1,35 @@
+package testhelper_test
+
+import (
+	"bytes"
+	"net"
+	"testing"
+
+	"github.com/Shopify/toxiproxy/v2/testhelper"
+)
+
+func TestSimpleServer(t *testing.T) {
+	testhelper.WithTCPServer(t, func(addr string, response chan []byte) {
+		conn, err := net.Dial("tcp", addr)
+		if err != nil {
+			t.Error("Unable to dial TCP server", err)
+		}
+
+		msg := []byte("hello world")
+
+		_, err = conn.Write(msg)
+		if err != nil {
+			t.Error("Failed writing to TCP server", err)
+		}
+
+		err = conn.Close()
+		if err != nil {
+			t.Error("Failed to close TCP connection", err)
+		}
+
+		resp := <-response
+		if !bytes.Equal(resp, msg) {
+			t.Error("Server didn't read bytes from client")
+		}
+	})
+}

testhelper/testhelper_test.go → testhelper/timeout_after_test.go

@@ -1,17 +1,19 @@
-package testhelper
+package testhelper_test
 
 import (
 	"testing"
 	"time"
+
+	"github.com/Shopify/toxiproxy/v2/testhelper"
 )
 
 func TestTimeoutAfter(t *testing.T) {
-	err := TimeoutAfter(5*time.Millisecond, func() {})
+	err := testhelper.TimeoutAfter(5*time.Millisecond, func() {})
 	if err != nil {
 		t.Fatal("Non blocking function should not timeout.")
 	}
 
-	err = TimeoutAfter(5*time.Millisecond, func() {
+	err = testhelper.TimeoutAfter(5*time.Millisecond, func() {
 		time.Sleep(time.Second)
 	})
 	if err == nil {

toxiproxy_test.go

@@ -0,0 +1,45 @@
+package toxiproxy_test
+
+import (
+	"net"
+	"testing"
+
+	"github.com/Shopify/toxiproxy/v2"
+	"github.com/Shopify/toxiproxy/v2/testhelper"
+)
+
+func NewTestProxy(name, upstream string) *toxiproxy.Proxy {
+	proxy := toxiproxy.NewProxy()
+
+	proxy.Name = name
+	proxy.Listen = "localhost:0"
+	proxy.Upstream = upstream
+
+	return proxy
+}
+
+func WithTCPProxy(
+	t *testing.T,
+	f func(proxy net.Conn, response chan []byte, proxyServer *toxiproxy.Proxy),
+) {
+	testhelper.WithTCPServer(t, func(upstream string, response chan []byte) {
+		proxy := NewTestProxy("test", upstream)
+		proxy.Start()
+
+		conn := AssertProxyUp(t, proxy.Listen, true)
+
+		f(conn, response, proxy)
+
+		proxy.Stop()
+	})
+}
+
+func AssertProxyUp(t *testing.T, addr string, up bool) net.Conn {
+	conn, err := net.Dial("tcp", addr)
+	if err != nil && up {
+		t.Error("Expected proxy to be up:", err)
+	} else if err == nil && !up {
+		t.Error("Expected proxy to be down")
+	}
+	return conn
+}