feat(pubsub/pstest): support listening on custom address (#11606)

larryanz · hongalex · web-flow · commit 63865a202b95 · 2025-03-11T17:45:16.000-04:00
Co-authored-by: Alex Hong &lt;9397363+hongalex@users.noreply.github.com&gt;
diff --git a/pubsub/pstest/fake.go b/pubsub/pstest/fake.go
@@ -112,18 +112,33 @@ func NewServer(opts ...ServerReactorOption) *Server {
 	return NewServerWithPort(0, opts...)
 }
 
-// NewServerWithPort creates a new fake server running in the current process at the specified port.
+// NewServerWithPort creates a new fake server running in the current process at
+// the specified port.
 func NewServerWithPort(port int, opts ...ServerReactorOption) *Server {
-	return NewServerWithCallback(port, func(*grpc.Server) { /* empty */ }, opts...)
+	return NewServerWithAddress(fmt.Sprintf("localhost:%d", port), opts...)
 }
 
-// NewServerWithCallback creates new fake server running in the current process at the specified port.
-// Before starting the server, the provided callback is called to allow caller to register additional fakes
-// into grpc server.
+// NewServerWithAddress creates a new fake server running in the current process
+// at the specified address (host and port).
+func NewServerWithAddress(address string, opts ...ServerReactorOption) *Server {
+	return initNewServer(address, func(*grpc.Server) { /* empty */ }, opts...)
+}
+
+// NewServerWithCallback creates new fake server running in the current process
+// at the specified port. Before starting the server, the provided callback is
+// called to allow caller to register additional fakes into grpc server.
 func NewServerWithCallback(port int, callback func(*grpc.Server), opts ...ServerReactorOption) *Server {
-	srv, err := testutil.NewServerWithPort(port)
+	return initNewServer(fmt.Sprintf("localhost:%d", port), callback, opts...)
+}
+
+// NewServerByAddressWithCallback creates new fake server running in the current
+// process at with the provided address (host and port).
+// Before starting the server, the provided callback is called to allow caller
+// to register additional fakes into grpc server.
+func initNewServer(address string, callback func(*grpc.Server), opts ...ServerReactorOption) *Server {
+	srv, err := testutil.NewServerWithAddress(address)
 	if err != nil {
-		panic(fmt.Sprintf("pstest.NewServerWithPort: %v", err))
+		panic(fmt.Sprintf("pstest.initNewServer: %v", err))
 	}
 	reactorOptions := ReactorOptions{}
 	for _, opt := range opts {
diff --git a/pubsub/pstest/fake_test.go b/pubsub/pstest/fake_test.go
@@ -41,42 +41,49 @@ import (
 )
 
 func TestNewServerWithPort(t *testing.T) {
-	// Allocate an available port to use with NewServerWithPort and then close it so it's available.
-	// Note: There is no guarantee that the port does not become used between closing
-	// the listener and creating the new server with NewServerWithPort, but the chances are
-	// very small.
-	l, err := net.Listen("tcp", ":0")
-	if err != nil {
-		t.Fatal(err)
-	}
-	port := l.Addr().(*net.TCPAddr).Port
-	l.Close()
-
-	// Pass a non 0 port to demonstrate we can pass a hardcoded port for the server to listen on
+	port := getFreePort(t)
 	srv := NewServerWithPort(port)
+
+	conn, err := grpc.NewClient(srv.Addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	if err != nil {
 		t.Fatal(err)
 	}
-	defer srv.Close()
-	conn, err := grpc.Dial(srv.Addr, grpc.WithInsecure())
-	if err != nil {
-		t.Fatal(err)
-	}
-	defer conn.Close()
+
+	t.Cleanup(func() {
+		conn.Close()
+		srv.Close()
+	})
 }
 
-func TestNewServerWithCallback(t *testing.T) {
-	// Allocate an available port to use with NewServerWithPort and then close it so it's available.
-	// Note: There is no guarantee that the port does not become used between closing
-	// the listener and creating the new server with NewServerWithPort, but the chances are
-	// very small.
-	l, err := net.Listen("tcp", ":0")
-	if err != nil {
-		t.Fatal(err)
+func TestNewServerWithAddress(t *testing.T) {
+	hosts := []string{
+		"",
+		"0.0.0.0",
+		"127.0.0.1",
+		"localhost",
+	}
+	for _, h := range hosts {
+		port := getFreePort(t)
+		address := fmt.Sprintf("%s:%d", h, port)
+		t.Run(fmt.Sprintf("Init new server succeed with address %s", address), func(t *testing.T) {
+			srv := NewServerWithAddress(address)
+
+			conn, err := grpc.NewClient(srv.Addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			t.Cleanup(func() {
+				conn.Close()
+				srv.Close()
+			})
+		})
 	}
-	port := l.Addr().(*net.TCPAddr).Port
-	l.Close()
 
+}
+
+func TestNewServerWithCallback(t *testing.T) {
+	port := getFreePort(t)
 	additionalFake := struct {
 		iampb.UnimplementedIAMPolicyServer
 	}{}
@@ -90,20 +97,35 @@ func TestNewServerWithCallback(t *testing.T) {
 
 	// Pass a non 0 port to demonstrate we can pass a hardcoded port for the server to listen on
 	srv := NewServerWithCallback(port, callback)
-	if err != nil {
-		t.Fatal(err)
-	}
-	defer srv.Close()
 
 	conn, err := grpc.NewClient(srv.Addr, grpc.WithTransportCredentials(insecure.NewCredentials()))
 	if err != nil {
 		t.Fatal(err)
 	}
-	defer conn.Close()
 
 	if !verifyCallback {
 		t.Fatal("callback was not invoked")
 	}
+
+	t.Cleanup(func() {
+		conn.Close()
+		srv.Close()
+	})
+}
+
+// getFreePort allocates an available port then close it so it's available.
+// Note: There is no guarantee that the port does not become used between closing
+// the listener and creating the new server with the invocation function, but
+// the chances are very small.
+func getFreePort(t *testing.T) int {
+	l, err := net.Listen("tcp", ":0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	port := l.Addr().(*net.TCPAddr).Port
+	l.Close()
+
+	return port
 }
 
 func TestTopics(t *testing.T) {
