Avoid timing issues in the integration tests

examples/gateway/BUILD.bazel

@@ -15,5 +15,6 @@ go_library(
         "//runtime:go_default_library",
         "@com_github_golang_glog//:go_default_library",
         "@org_golang_google_grpc//:go_default_library",
+        "@org_golang_google_grpc//connectivity:go_default_library",
     ],
 )

examples/gateway/gateway.go

@@ -7,24 +7,13 @@ import (
 	"net/http"
 	"time"
 
-	"github.com/golang/glog"
 	"github.com/grpc-ecosystem/grpc-gateway/examples/proto/examplepb"
 	gwruntime "github.com/grpc-ecosystem/grpc-gateway/runtime"
 	"google.golang.org/grpc"
 )
 
 // newGateway returns a new gateway server which translates HTTP into gRPC.
-func newGateway(ctx context.Context, network, addr string, opts []gwruntime.ServeMuxOption) (http.Handler, error) {
-	conn, err := dial(ctx, network, addr)
-	if err != nil {
-		return nil, err
-	}
-	go func() {
-		<-ctx.Done()
-		if err := conn.Close(); err != nil {
-			glog.Errorf("Failed to close a client connection to the gRPC server: %v", err)
-		}
-	}()
+func newGateway(ctx context.Context, conn *grpc.ClientConn, opts []gwruntime.ServeMuxOption) (http.Handler, error) {
 
 	mux := gwruntime.NewServeMux(opts...)
 

examples/gateway/handlers.go

@@ -1,11 +1,14 @@
 package gateway
 
 import (
+	"fmt"
 	"net/http"
 	"path"
 	"strings"
 
 	"github.com/golang/glog"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/connectivity"
 )
 
 func swaggerServer(dir string) http.HandlerFunc {
@@ -44,5 +47,15 @@ func preflightHandler(w http.ResponseWriter, r *http.Request) {
 	methods := []string{"GET", "HEAD", "POST", "PUT", "DELETE"}
 	w.Header().Set("Access-Control-Allow-Methods", strings.Join(methods, ","))
 	glog.Infof("preflight request for %s", r.URL.Path)
-	return
+}
+
+func healthzServer(conn *grpc.ClientConn) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "text/plain")
+		if s := conn.GetState(); s != connectivity.Ready {
+			http.Error(w, fmt.Sprintf("grpc server is %s", s), http.StatusBadGateway)
+			return
+		}
+		fmt.Fprintln(w, "ok")
+	}
 }

examples/gateway/main.go

@@ -35,10 +35,22 @@ func Run(ctx context.Context, opts Options) error {
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
+	conn, err := dial(ctx, opts.GRPCServer.Network, opts.GRPCServer.Addr)
+	if err != nil {
+		return err
+	}
+	go func() {
+		<-ctx.Done()
+		if err := conn.Close(); err != nil {
+			glog.Errorf("Failed to close a client connection to the gRPC server: %v", err)
+		}
+	}()
+
 	mux := http.NewServeMux()
 	mux.HandleFunc("/swagger/", swaggerServer(opts.SwaggerDir))
+	mux.HandleFunc("/healthz", healthzServer(conn))
 
-	gw, err := newGateway(ctx, opts.GRPCServer.Network, opts.GRPCServer.Addr, opts.Mux)
+	gw, err := newGateway(ctx, conn, opts.Mux)
 	if err != nil {
 		return err
 	}

examples/integration/integration_test.go

@@ -13,7 +13,6 @@ import (
 	"strings"
 	"sync"
 	"testing"
-	"time"
 
 	"github.com/golang/protobuf/jsonpb"
 	"github.com/golang/protobuf/proto"
@@ -60,8 +59,9 @@ func TestForwardResponseOption(t *testing.T) {
 			return
 		}
 	}()
-
-	time.Sleep(100 * time.Millisecond)
+	if err := waitForGateway(ctx, 8081); err != nil {
+		t.Errorf("waitForGateway(ctx, 8081) failed with %v; want success", err)
+	}
 	testEcho(t, 8081, "application/vnd.docker.plugins.v1.1+json")
 }
 

examples/integration/main_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"net/http"
 	"os"
 	"testing"
 	"time"
@@ -32,6 +33,29 @@ func runGateway(ctx context.Context, addr string, opts ...gwruntime.ServeMuxOpti
 	})
 }
 
+func waitForGateway(ctx context.Context, port uint16) error {
+	ch := time.After(10 * time.Second)
+
+	var err error
+	for {
+		if r, err := http.Get(fmt.Sprintf("http://localhost:%d/healthz", port)); err == nil {
+			if r.StatusCode == http.StatusOK {
+				return nil
+			}
+			err = fmt.Errorf("server localhost:%d returned an unexpected status %d", port, r.StatusCode)
+		}
+
+		glog.Infof("Waiting for localhost:%d to get ready", port)
+		select {
+		case <-ctx.Done():
+			return err
+		case <-ch:
+			return err
+		case <-time.After(10 * time.Millisecond):
+		}
+	}
+}
+
 func runServers(ctx context.Context) <-chan error {
 	ch := make(chan error, 2)
 	go func() {
@@ -57,7 +81,9 @@ func TestMain(m *testing.M) {
 
 	ch := make(chan int, 1)
 	go func() {
-		time.Sleep(100 * time.Millisecond)
+		if err := waitForGateway(ctx, 8080); err != nil {
+			glog.Errorf("waitForGateway(ctx, 8080) failed with %v; want success", err)
+		}
 		ch <- m.Run()
 	}()
 

examples/integration/proto_error_test.go

@@ -29,11 +29,9 @@ func TestWithProtoErrorHandler(t *testing.T) {
 
 	const port = 8082
 	go runServer(ctx, t, port)
-
-	// Waiting for the server's getting available.
-	// TODO(yugui) find a better way to wait
-	time.Sleep(100 * time.Millisecond)
-
+	if err := waitForGateway(ctx, 8082); err != nil {
+		t.Errorf("waitForGateway(ctx, 8082) failed with %v; want success", err)
+	}
 	testEcho(t, port, "application/json")
 	testEchoBody(t, port)
 }
@@ -50,9 +48,9 @@ func TestABEWithProtoErrorHandler(t *testing.T) {
 
 	const port = 8083
 	go runServer(ctx, t, port)
-	// Waiting for the server's getting available.
-	// TODO(yugui) find a better way to wait
-	time.Sleep(100 * time.Millisecond)
+	if err := waitForGateway(ctx, 8083); err != nil {
+		t.Errorf("waitForGateway(ctx, 8083) failed with %v; want success", err)
+	}
 
 	testABECreate(t, port)
 	testABECreateBody(t, port)
@@ -122,10 +120,9 @@ func TestUnknownPathWithProtoError(t *testing.T) {
 
 	const port = 8084
 	go runServer(ctx, t, port)
-
-	// Waiting for the server's getting available.
-	// TODO(yugui) find a better way to wait
-	time.Sleep(100 * time.Millisecond)
+	if err := waitForGateway(ctx, 8084); err != nil {
+		t.Errorf("waitForGateway(ctx, 8084) failed with %v; want success", err)
+	}
 
 	url := fmt.Sprintf("http://localhost:%d", port)
 	resp, err := http.Post(url, "application/json", strings.NewReader("{}"))