github.com/grpc-ecosystem/grpc-gateway/v2@v2.19.1/examples/internal/integration/main_test.go (about)

     1  package integration_test
     2  
     3  import (
     4  	"context"
     5  	"flag"
     6  	"fmt"
     7  	"net/http"
     8  	"os"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/grpc-ecosystem/grpc-gateway/v2/examples/internal/gateway"
    13  	"github.com/grpc-ecosystem/grpc-gateway/v2/examples/internal/server"
    14  	gwruntime "github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
    15  	"google.golang.org/grpc/grpclog"
    16  )
    17  
    18  var (
    19  	endpoint   = flag.String("endpoint", "localhost:9090", "endpoint of the gRPC service")
    20  	network    = flag.String("network", "tcp", `one of "tcp" or "unix". Must be consistent to -endpoint`)
    21  	openAPIDir = flag.String("openapi_dir", "examples/internal/proto/examplepb", "path to the directory which contains OpenAPI definitions")
    22  )
    23  
    24  func runGateway(ctx context.Context, addr string, opts ...gwruntime.ServeMuxOption) error {
    25  	return gateway.Run(ctx, gateway.Options{
    26  		Addr: addr,
    27  		GRPCServer: gateway.Endpoint{
    28  			Network: *network,
    29  			Addr:    *endpoint,
    30  		},
    31  		OpenAPIDir: *openAPIDir,
    32  		Mux:        opts,
    33  	})
    34  }
    35  
    36  func waitForGateway(ctx context.Context, port uint16) error {
    37  	ch := time.After(10 * time.Second)
    38  
    39  	for {
    40  		r, err := http.Get(fmt.Sprintf("http://localhost:%d/healthz", port))
    41  		if err == nil && r.StatusCode == http.StatusOK {
    42  			return nil
    43  		}
    44  
    45  		grpclog.Infof("Waiting for localhost:%d to get ready", port)
    46  		select {
    47  		case <-ctx.Done():
    48  			return err
    49  		case <-ch:
    50  			return err
    51  		case <-time.After(10 * time.Millisecond):
    52  		}
    53  	}
    54  }
    55  
    56  func runServers(ctx context.Context) <-chan error {
    57  	ch := make(chan error, 3)
    58  	go func() {
    59  		if err := server.Run(ctx, *network, *endpoint); err != nil {
    60  			ch <- fmt.Errorf("cannot run grpc service: %v", err)
    61  		}
    62  	}()
    63  	go func() {
    64  		if err := runGateway(ctx, ":8088"); err != nil {
    65  			ch <- fmt.Errorf("cannot run gateway service: %v", err)
    66  		}
    67  	}()
    68  	go func() {
    69  		if err := server.RunInProcessGateway(ctx, ":8089"); err != nil {
    70  			ch <- fmt.Errorf("cannot run in process gateway service: %v", err)
    71  		}
    72  	}()
    73  	return ch
    74  }
    75  
    76  func TestMain(m *testing.M) {
    77  	flag.Parse()
    78  
    79  	ctx, cancel := context.WithCancel(context.Background())
    80  	defer cancel()
    81  	errCh := runServers(ctx)
    82  
    83  	ch := make(chan int, 1)
    84  	go func() {
    85  		if err := waitForGateway(ctx, 8088); err != nil {
    86  			grpclog.Errorf("waitForGateway(ctx, 8088) failed with %v; want success", err)
    87  		}
    88  		ch <- m.Run()
    89  	}()
    90  
    91  	select {
    92  	case err := <-errCh:
    93  		fmt.Fprintln(os.Stderr, err)
    94  		os.Exit(1)
    95  	case status := <-ch:
    96  		cancel()
    97  		os.Exit(status)
    98  	}
    99  }