Fix transptest testing

libbeat/outputs/transport/transptest/testing_test.go

@@ -22,6 +22,7 @@ package transptest
 import (
 	"fmt"
 	"net"
+	"sync"
 	"testing"
 	"time"
 
@@ -34,9 +35,9 @@ import (
 // netSOCKS5Proxy starts a new SOCKS5 proxy server that listens on localhost.
 //
 // Usage:
-//  l, tcpAddr := newSOCKS5Proxy(t)
-//  defer l.Close()
-func newSOCKS5Proxy(t *testing.T) (net.Listener, transport.ProxyConfig) {
+//  l, teardown := newSOCKS5Proxy(t)
+//  defer teardown()
+func newSOCKS5Proxy(t *testing.T) (net.Listener, func()) {
 	// Create a SOCKS5 server
 	conf := &socks5.Config{}
 	server, err := socks5.New(conf)
@@ -50,31 +51,42 @@ func newSOCKS5Proxy(t *testing.T) (net.Listener, transport.ProxyConfig) {
 		t.Fatalf("err: %v", err)
 	}
 
-	// Listen
+	// Listen and serve
+	var wg sync.WaitGroup
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		err := server.Serve(l)
 		if err != nil {
-			t.Log(err)
+			t.Logf("Server error (%T): %+v", err, err)
 		}
 	}()
 
+	cleanup := func() {
+		defer wg.Wait()
+		l.Close()
+	}
+
+	return l, cleanup
+}
+
+func listenerProxyConfig(l net.Listener) *transport.ProxyConfig {
+	if l == nil {
+		return nil
+	}
+
 	tcpAddr := l.Addr().(*net.TCPAddr)
-	config := transport.ProxyConfig{URL: fmt.Sprintf("socks5://%s", tcpAddr.String())}
-	return l, config
+	return &transport.ProxyConfig{
+		URL: fmt.Sprintf("socks5://%s", tcpAddr.String()),
+	}
 }
 
 func TestTransportReconnectsOnConnect(t *testing.T) {
-	l, config := newSOCKS5Proxy(t)
-	defer l.Close()
-
 	certName := "ca_test"
 	timeout := 2 * time.Second
 	GenCertForTestingPurpose(t, "127.0.0.1", certName, "")
 
-	testServer(t, &config, func(t *testing.T, makeServer MockServerFactory, proxy *transport.ProxyConfig) {
-		server := makeServer(t, timeout, certName, proxy)
-		defer server.Close()
-
+	testServer(t, timeout, certName, func(t *testing.T, server *MockServer) {
 		transp, err := server.Transp()
 		if err != nil {
 			t.Fatalf("Failed to create client: %v", err)
@@ -102,50 +114,54 @@ func TestTransportReconnectsOnConnect(t *testing.T) {
 }
 
 func TestTransportFailConnectUnknownAddress(t *testing.T) {
-	l, config := newSOCKS5Proxy(t)
-	defer l.Close()
-
+	timeout := 100 * time.Millisecond
 	certName := "ca_test"
 	GenCertForTestingPurpose(t, "127.0.0.1", certName, "")
 
-	invalidAddr := "invalid.dns.fqdn-unknown.invalid:100"
-
-	run := func(makeTransp TransportFactory, proxy *transport.ProxyConfig) func(*testing.T) {
-		return func(t *testing.T) {
-			transp, err := makeTransp(invalidAddr, proxy)
-			if err != nil {
-				t.Fatalf("failed to generate transport client: %v", err)
-			}
-
-			err = transp.Connect()
-			assert.NotNil(t, err)
-		}
+	transports := map[string]TransportFactory{
+		"tcp": connectTCP(timeout),
+		"tls": connectTLS(timeout, certName),
 	}
 
-	factoryTests := func(f TransportFactory) func(*testing.T) {
-		return func(t *testing.T) {
-			t.Run("connect", run(f, nil))
-			t.Run("socks5", run(f, &config))
-		}
+	modes := map[string]struct {
+		withProxy bool
+	}{
+		"connect": {withProxy: false},
+		"socks5":  {withProxy: true},
 	}
 
-	timeout := 100 * time.Millisecond
-	t.Run("tcp", factoryTests(connectTCP(timeout)))
-	t.Run("tls", factoryTests(connectTLS(timeout, certName)))
+	const invalidAddr = "invalid.dns.fqdn-unknown.invalid:100"
+
+	for name, factory := range transports {
+		t.Run(name, func(t *testing.T) {
+			for mode, test := range modes {
+				t.Run(mode, func(t *testing.T) {
+					var listener net.Listener
+					if test.withProxy {
+						var teardown func()
+						listener, teardown = newSOCKS5Proxy(t)
+						defer teardown()
+					}
+
+					transp, err := factory(invalidAddr, listenerProxyConfig(listener))
+					if err != nil {
+						t.Fatalf("failed to generate transport client: %v", err)
+					}
+
+					err = transp.Connect()
+					assert.NotNil(t, err)
+				})
+			}
+		})
+	}
 }
 
 func TestTransportClosedOnWriteReadError(t *testing.T) {
-	l, config := newSOCKS5Proxy(t)
-	defer l.Close()
-
 	certName := "ca_test"
 	timeout := 2 * time.Second
 	GenCertForTestingPurpose(t, "127.0.0.1", certName, "")
 
-	testServer(t, &config, func(t *testing.T, makeServer MockServerFactory, proxy *transport.ProxyConfig) {
-		server := makeServer(t, timeout, certName, proxy)
-		defer server.Close()
-
+	testServer(t, timeout, certName, func(t *testing.T, server *MockServer) {
 		client, transp, err := server.ConnectPair()
 		if err != nil {
 			t.Fatalf("Failed to create client: %v", err)
@@ -159,20 +175,35 @@ func TestTransportClosedOnWriteReadError(t *testing.T) {
 	})
 }
 
-func testServer(t *testing.T, config *transport.ProxyConfig, run func(*testing.T, MockServerFactory, *transport.ProxyConfig)) {
-	runner := func(f MockServerFactory, c *transport.ProxyConfig) func(t *testing.T) {
-		return func(t *testing.T) {
-			run(t, f, config)
-		}
+func testServer(t *testing.T, timeout time.Duration, cert string, fn func(t *testing.T, server *MockServer)) {
+	transports := map[string]MockServerFactory{
+		"tcp": NewMockServerTCP,
+		"tls": NewMockServerTLS,
 	}
 
-	factoryTests := func(f MockServerFactory) func(t *testing.T) {
-		return func(t *testing.T) {
-			t.Run("connect", runner(f, nil))
-			t.Run("socks5", runner(f, config))
-		}
+	modes := map[string]struct {
+		withProxy bool
+	}{
+		"connect": {withProxy: false},
+		"socks5":  {withProxy: true},
 	}
 
-	t.Run("tcp", factoryTests(NewMockServerTCP))
-	t.Run("tls", factoryTests(NewMockServerTLS))
+	for name, factory := range transports {
+		t.Run(name, func(t *testing.T) {
+			for mode, test := range modes {
+				t.Run(mode, func(t *testing.T) {
+					var listener net.Listener
+					if test.withProxy {
+						var teardown func()
+						listener, teardown = newSOCKS5Proxy(t)
+						defer teardown()
+					}
+
+					server := factory(t, timeout, cert, listenerProxyConfig(listener))
+					defer server.Close()
+					fn(t, server)
+				})
+			}
+		})
+	}
 }