Add port open timeout handling

leonardinius · Apr 3, 2024 · a527dfa · a527dfa
1 parent 92d6879
commit a527dfa
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 38 deletions.
diff --git a/app/main_test.go b/app/main_test.go
@@ -7,6 +7,7 @@ import (
 	"log"
 	"net"
 	"os"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
@@ -86,31 +87,32 @@ smtpd-proxy:
 }
 
 func waitForPortListenStart(ctx context.Context, t *testing.T, port int) (conn net.Conn) {
-	var d net.Dialer
-	var err error
-	addr := fmt.Sprintf("%s:%d", bindHost, port)
+	t.Helper()
+
+	addr := net.JoinHostPort(bindHost, strconv.Itoa(port))
+
 	poll := time.NewTicker(50 * time.Millisecond)
 	defer poll.Stop()
+
 	timeout := time.NewTimer(5 * time.Second)
 	defer timeout.Stop()
 
-	select {
-	case <-poll.C:
-		conn, _ = checkAddr(ctx, &d, addr)
-		if conn != nil {
-			break
+	for {
+		select {
+		case <-timeout.C:
+			t.Fatalf("%s port open timeout", addr)
+
+		case <-ctx.Done():
+			t.Fatalf("%s port open error, parent context is done: %v", addr, ctx.Err())
+
+		case <-poll.C:
+			var d net.Dialer
+			conn, _ = checkAddr(ctx, &d, addr)
+			if conn != nil {
+				return conn
+			}
 		}
-	case <-timeout.C:
-		t.Fatal("SMTP open timeout")
-		break
-	}
-
-	require.NotNil(t, conn)
-	err = conn.SetDeadline(time.Now().Add(100 * time.Millisecond))
-	if err != nil {
-		t.Fatal("SMTP set connection deadline error", err)
 	}
-	return conn
 }
 
 func checkAddr(ctx context.Context, d *net.Dialer, addr string) (net.Conn, error) {

diff --git a/tests/infra.go b/tests/infra.go
@@ -2,10 +2,10 @@ package systemtest
 
 import (
 	"context"
-	"fmt"
 	"log"
 	"net"
 	"os"
+	"strconv"
 	"sync"
 	"testing"
 	"time"
@@ -39,35 +39,38 @@ func RunMainWithConfig(ctx context.Context, t *testing.T, yamlConfig string, por
 	}()
 
 	conn := waitForPortListenStart(ctx, t, port)
+	err = conn.SetDeadline(time.Now().Add(500 * time.Millisecond))
+	require.NoError(t, err, "SMTP set connection deadline error")
 	test(t, conn)
 }
 
 func waitForPortListenStart(ctx context.Context, t *testing.T, port int) (conn net.Conn) {
-	var d net.Dialer
-	var err error
-	addr := fmt.Sprintf("%s:%d", BindHost, port)
-	poll := time.NewTicker(20 * time.Millisecond)
+	t.Helper()
+
+	addr := net.JoinHostPort(BindHost, strconv.Itoa(port))
+
+	poll := time.NewTicker(50 * time.Millisecond)
 	defer poll.Stop()
+
 	timeout := time.NewTimer(5 * time.Second)
 	defer timeout.Stop()
 
-	select {
-	case <-poll.C:
-		conn, _ = checkAddr(ctx, &d, addr)
-		if conn != nil {
-			break
+	for {
+		select {
+		case <-timeout.C:
+			t.Fatalf("%s port open timeout", addr)
+
+		case <-ctx.Done():
+			t.Fatalf("%s port open error, parent context is done: %v", addr, ctx.Err())
+
+		case <-poll.C:
+			var d net.Dialer
+			conn, _ = checkAddr(ctx, &d, addr)
+			if conn != nil {
+				return conn
+			}
 		}
-	case <-timeout.C:
-		t.Fatal("SMTP open timeout")
-		break
-	}
-
-	require.NotNil(t, conn)
-	err = conn.SetDeadline(time.Now().Add(500 * time.Millisecond))
-	if err != nil {
-		t.Fatal("SMTP set connection deadline error", err)
 	}
-	return conn
 }
 
 func checkAddr(ctx context.Context, d *net.Dialer, addr string) (net.Conn, error) {