From ef872fc933ea9579c2c6c71fa6340d25a4d84526 Mon Sep 17 00:00:00 2001
From: Oleksandr Redko <oleksandr.red+github@gmail.com>
Date: Tue, 21 Feb 2023 13:38:41 +0200
Subject: [PATCH] Refactor tests for GRPC check (#88)

- Fix "defer is not called in TestMain because of os.Exit".
- Use table-driven approach to reduce repetitive code.
- Listen only localhost instead of all interfaces.
- Choose random listen port.
- Use t.Log instead of log.Fatalf
---
 checks/grpc/check_test.go | 128 +++++++++++++++++---------------------
 1 file changed, 56 insertions(+), 72 deletions(-)

diff --git a/checks/grpc/check_test.go b/checks/grpc/check_test.go
index 096da5b..96d05b0 100644
--- a/checks/grpc/check_test.go
+++ b/checks/grpc/check_test.go
@@ -2,9 +2,7 @@ package grpc
 
 import (
 	"context"
-	"log"
 	"net"
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -13,77 +11,63 @@ import (
 	"google.golang.org/grpc/health/grpc_health_v1"
 )
 
-const (
-	addr    = ":8080"
-	service = "HealthTest"
-)
-
-var healthServer *health.Server
-
-func TestMain(m *testing.M) {
-	healthServer = health.NewServer()
-	healthServer.SetServingStatus(service, grpc_health_v1.HealthCheckResponse_SERVING)
-
-	lis, err := net.Listen("tcp", addr)
-	if err != nil {
-		log.Fatalf("error setting up tcp listener: %v", err)
-	}
-
-	server := grpc.NewServer()
-	grpc_health_v1.RegisterHealthServer(server, healthServer)
-
-	go func() {
-		if err := server.Serve(lis); err != nil {
-			log.Fatalf("failed to serve: %v", err)
-		}
-	}()
-
-	defer server.Stop()
-
-	os.Exit(m.Run())
-}
-
-func TestNew_WithServingStatusServing(t *testing.T) {
-	healthServer.SetServingStatus(service, grpc_health_v1.HealthCheckResponse_SERVING)
-
-	check := New(Config{
-		Target:  addr,
-		Service: service,
-		DialOptions: []grpc.DialOption{
-			grpc.WithInsecure(),
+func TestNew(t *testing.T) {
+	for name, tc := range map[string]struct {
+		servingStatus grpc_health_v1.HealthCheckResponse_ServingStatus
+		requireError  bool
+	}{
+		"serving": {
+			servingStatus: grpc_health_v1.HealthCheckResponse_SERVING,
+			requireError:  false,
 		},
-	})
-
-	err := check(context.Background())
-	require.NoError(t, err)
-}
-
-func TestNew_WithServingStatusUnknown(t *testing.T) {
-	healthServer.SetServingStatus(service, grpc_health_v1.HealthCheckResponse_UNKNOWN)
-
-	check := New(Config{
-		Target:  addr,
-		Service: service,
-		DialOptions: []grpc.DialOption{
-			grpc.WithInsecure(),
+		"unknown": {
+			servingStatus: grpc_health_v1.HealthCheckResponse_UNKNOWN,
+			requireError:  true,
 		},
-	})
-
-	err := check(context.Background())
-	require.Error(t, err)
-}
-
-func TestNew_WithServingStatusNotServing(t *testing.T) {
-	healthServer.SetServingStatus(service, grpc_health_v1.HealthCheckResponse_NOT_SERVING)
-
-	check := New(Config{
-		Target:  addr,
-		Service: service,
-		DialOptions: []grpc.DialOption{
-			grpc.WithInsecure(),
+		"not serving": {
+			servingStatus: grpc_health_v1.HealthCheckResponse_NOT_SERVING,
+			requireError:  true,
 		},
-	})
-
-	err := check(context.Background())
-	require.Error(t, err)
+	} {
+		servingStatus := tc.servingStatus
+		requireError := tc.requireError
+
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			const service = "HealthTest"
+
+			healthServer := health.NewServer()
+			healthServer.SetServingStatus(service, servingStatus)
+
+			lis, err := net.Listen("tcp", "localhost:0")
+			require.NoError(t, err)
+
+			server := grpc.NewServer()
+			grpc_health_v1.RegisterHealthServer(server, healthServer)
+
+			go func() {
+				if err := server.Serve(lis); err != nil {
+					t.Log("Failed to serve GRPC", err)
+				}
+			}()
+			defer server.Stop()
+
+			check := New(Config{
+				Target:  lis.Addr().String(),
+				Service: service,
+				DialOptions: []grpc.DialOption{
+					grpc.WithInsecure(),
+				},
+			})
+
+			err = check(context.Background())
+
+			if requireError {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+		})
+	}
 }