From 074df56c3d7b47c17b35198894c1aa850e4c0c57 Mon Sep 17 00:00:00 2001
From: Viktor Liu <17948409+lixmal@users.noreply.github.com>
Date: Mon, 26 Jan 2026 16:30:00 +0800
Subject: [PATCH] [client] Fix flaky JWT SSH test (#5181)

---
 client/ssh/server/jwt_test.go | 12 ++++++------
 client/ssh/server/test.go     |  3 +--
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/client/ssh/server/jwt_test.go b/client/ssh/server/jwt_test.go
index dbef011ac..b2f3ac6a0 100644
--- a/client/ssh/server/jwt_test.go
+++ b/client/ssh/server/jwt_test.go
@@ -54,7 +54,7 @@ func TestJWTEnforcement(t *testing.T) {
 		server.SetAllowRootLogin(true)
 
 		serverAddr := StartTestServer(t, server)
-		defer require.NoError(t, server.Stop())
+		defer func() { require.NoError(t, server.Stop()) }()
 
 		host, portStr, err := net.SplitHostPort(serverAddr)
 		require.NoError(t, err)
@@ -88,7 +88,7 @@ func TestJWTEnforcement(t *testing.T) {
 		serverNoJWT.SetAllowRootLogin(true)
 
 		serverAddrNoJWT := StartTestServer(t, serverNoJWT)
-		defer require.NoError(t, serverNoJWT.Stop())
+		defer func() { require.NoError(t, serverNoJWT.Stop()) }()
 
 		hostNoJWT, portStrNoJWT, err := net.SplitHostPort(serverAddrNoJWT)
 		require.NoError(t, err)
@@ -213,7 +213,7 @@ func TestJWTDetection(t *testing.T) {
 	server.SetAllowRootLogin(true)
 
 	serverAddr := StartTestServer(t, server)
-	defer require.NoError(t, server.Stop())
+	defer func() { require.NoError(t, server.Stop()) }()
 
 	host, portStr, err := net.SplitHostPort(serverAddr)
 	require.NoError(t, err)
@@ -341,7 +341,7 @@ func TestJWTFailClose(t *testing.T) {
 			server.SetAllowRootLogin(true)
 
 			serverAddr := StartTestServer(t, server)
-			defer require.NoError(t, server.Stop())
+			defer func() { require.NoError(t, server.Stop()) }()
 
 			host, portStr, err := net.SplitHostPort(serverAddr)
 			require.NoError(t, err)
@@ -596,7 +596,7 @@ func TestJWTAuthentication(t *testing.T) {
 			server.UpdateSSHAuth(authConfig)
 
 			serverAddr := StartTestServer(t, server)
-			defer require.NoError(t, server.Stop())
+			defer func() { require.NoError(t, server.Stop()) }()
 
 			host, portStr, err := net.SplitHostPort(serverAddr)
 			require.NoError(t, err)
@@ -715,7 +715,7 @@ func TestJWTMultipleAudiences(t *testing.T) {
 			server.UpdateSSHAuth(authConfig)
 
 			serverAddr := StartTestServer(t, server)
-			defer require.NoError(t, server.Stop())
+			defer func() { require.NoError(t, server.Stop()) }()
 
 			host, portStr, err := net.SplitHostPort(serverAddr)
 			require.NoError(t, err)
diff --git a/client/ssh/server/test.go b/client/ssh/server/test.go
index f8abd1752..454d3afa3 100644
--- a/client/ssh/server/test.go
+++ b/client/ssh/server/test.go
@@ -8,19 +8,18 @@ import (
 	"time"
 )
 
+// StartTestServer starts the SSH server and returns the address it's listening on.
 func StartTestServer(t *testing.T, server *Server) string {
 	started := make(chan string, 1)
 	errChan := make(chan error, 1)
 
 	go func() {
-		// Use port 0 to let the OS assign a free port
 		addrPort := netip.MustParseAddrPort("127.0.0.1:0")
 		if err := server.Start(context.Background(), addrPort); err != nil {
 			errChan <- err
 			return
 		}
 
-		// Get the actual listening address from the server
 		actualAddr := server.Addr()
 		if actualAddr == nil {
 			errChan <- fmt.Errorf("server started but no listener address available")