github.com/outbrain/consul@v1.4.5/agent/checks/grpc_test.go (about)

     1  package checks
     2  
     3  import (
     4  	"crypto/tls"
     5  	"flag"
     6  	"fmt"
     7  	"log"
     8  	"net"
     9  	"os"
    10  	"testing"
    11  	"time"
    12  
    13  	"google.golang.org/grpc"
    14  	"google.golang.org/grpc/health"
    15  	hv1 "google.golang.org/grpc/health/grpc_health_v1"
    16  )
    17  
    18  var (
    19  	port         int
    20  	server       string
    21  	svcHealthy   string
    22  	svcUnhealthy string
    23  	svcMissing   string
    24  )
    25  
    26  func startServer() (*health.Server, *grpc.Server) {
    27  	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
    28  	if err != nil {
    29  		log.Fatalf("failed to listen: %v", err)
    30  	}
    31  	grpcServer := grpc.NewServer()
    32  	server := health.NewServer()
    33  	hv1.RegisterHealthServer(grpcServer, server)
    34  	go grpcServer.Serve(listener)
    35  	return server, grpcServer
    36  }
    37  
    38  func init() {
    39  	flag.IntVar(&port, "grpc-stub-port", 54321, "port for the gRPC stub server")
    40  }
    41  
    42  func TestMain(m *testing.M) {
    43  	flag.Parse()
    44  
    45  	healthy := "healthy"
    46  	unhealthy := "unhealthy"
    47  	missing := "missing"
    48  
    49  	srv, grpcStubApp := startServer()
    50  	srv.SetServingStatus(healthy, hv1.HealthCheckResponse_SERVING)
    51  	srv.SetServingStatus(unhealthy, hv1.HealthCheckResponse_NOT_SERVING)
    52  
    53  	server = fmt.Sprintf("%s:%d", "localhost", port)
    54  	svcHealthy = fmt.Sprintf("%s/%s", server, healthy)
    55  	svcUnhealthy = fmt.Sprintf("%s/%s", server, unhealthy)
    56  	svcMissing = fmt.Sprintf("%s/%s", server, missing)
    57  
    58  	result := 1
    59  	defer func() {
    60  		grpcStubApp.Stop()
    61  		os.Exit(result)
    62  	}()
    63  
    64  	result = m.Run()
    65  }
    66  
    67  func TestCheck(t *testing.T) {
    68  	type args struct {
    69  		target    string
    70  		timeout   time.Duration
    71  		tlsConfig *tls.Config
    72  	}
    73  	tests := []struct {
    74  		name    string
    75  		args    args
    76  		healthy bool
    77  	}{
    78  		// successes
    79  		{"should pass for healthy server", args{server, time.Second, nil}, true},
    80  		{"should pass for healthy service", args{svcHealthy, time.Second, nil}, true},
    81  
    82  		// failures
    83  		{"should fail for unhealthy service", args{svcUnhealthy, time.Second, nil}, false},
    84  		{"should fail for missing service", args{svcMissing, time.Second, nil}, false},
    85  		{"should timeout for healthy service", args{server, time.Nanosecond, nil}, false},
    86  		{"should fail if probe is secure, but server is not", args{server, time.Second, &tls.Config{}}, false},
    87  	}
    88  	for _, tt := range tests {
    89  		t.Run(tt.name, func(t *testing.T) {
    90  			probe := NewGrpcHealthProbe(tt.args.target, tt.args.timeout, tt.args.tlsConfig)
    91  			actualError := probe.Check()
    92  			actuallyHealthy := actualError == nil
    93  			if tt.healthy != actuallyHealthy {
    94  				t.Errorf("FAIL: %s. Expected healthy %t, but err == %v", tt.name, tt.healthy, actualError)
    95  			}
    96  		})
    97  	}
    98  }