trpc.group/trpc-go/trpc-go@v1.0.3/healthcheck/health_check_test.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  package healthcheck_test
    15  
    16  import (
    17  	"testing"
    18  
    19  	"github.com/stretchr/testify/require"
    20  	"trpc.group/trpc-go/trpc-go/healthcheck"
    21  )
    22  
    23  func TestHealthCheckService(t *testing.T) {
    24  	hc := healthcheck.New()
    25  
    26  	requireStatusUnknown(t, hc.CheckService("service"))
    27  	update, err := hc.Register("service")
    28  	require.Nil(t, err)
    29  	requireStatusUnknown(t, hc.CheckService("service"))
    30  	update(healthcheck.Serving)
    31  	require.Equal(t, healthcheck.Serving, hc.CheckService("service"))
    32  	update(healthcheck.NotServing)
    33  	require.Equal(t, healthcheck.NotServing, hc.CheckService("service"))
    34  }
    35  
    36  func TestHealthCheckServer(t *testing.T) {
    37  	hc := healthcheck.New()
    38  	require.Equal(t, healthcheck.Serving, hc.CheckServer())
    39  
    40  	update, err := hc.Register("service")
    41  	require.Nil(t, err)
    42  	requireStatusUnknown(t, hc.CheckServer())
    43  
    44  	update(healthcheck.Serving)
    45  	require.Equal(t, healthcheck.Serving, hc.CheckServer())
    46  	update(healthcheck.NotServing)
    47  	require.Equal(t, healthcheck.NotServing, hc.CheckServer())
    48  
    49  	hc.Unregister("service")
    50  	require.Equal(t, healthcheck.Serving, hc.CheckServer())
    51  }
    52  
    53  func TestHealthCheckRegisterTwice(t *testing.T) {
    54  	hc := healthcheck.New()
    55  
    56  	_, err := hc.Register("service")
    57  	require.Nil(t, err)
    58  	_, err = hc.Register("service")
    59  	require.NotNil(t, err)
    60  
    61  	hc.Unregister("service")
    62  	_, err = hc.Register("service")
    63  	require.Nil(t, err)
    64  }
    65  
    66  func TestHealthCheckUnregisteredService(t *testing.T) {
    67  	hc := healthcheck.New()
    68  	requireStatusUnknown(t, hc.CheckService("not_exist"))
    69  
    70  	hc = healthcheck.New(healthcheck.WithUnregisteredServiceStatus(healthcheck.Serving))
    71  	require.Equal(t, healthcheck.Serving, hc.CheckService("not_exist"))
    72  }
    73  
    74  func TestHealthCheckStatusWatchers(t *testing.T) {
    75  	const serviceName = "service"
    76  	var firstCalled int
    77  	hc := healthcheck.New(healthcheck.WithStatusWatchers(map[string][]func(healthcheck.Status){
    78  		serviceName: {
    79  			func(status healthcheck.Status) {
    80  				switch firstCalled++; firstCalled {
    81  				case 1:
    82  					require.Equal(t, healthcheck.Unknown, status)
    83  				case 2:
    84  					require.Equal(t, healthcheck.Serving, status)
    85  				case 3:
    86  					require.Equal(t, healthcheck.NotServing, status)
    87  				default:
    88  					require.FailNow(t, "status should only be updated 3 times")
    89  				}
    90  			},
    91  		},
    92  	}))
    93  	update, err := hc.Register(serviceName)
    94  	require.Nil(t, err)
    95  
    96  	var secondCalled int
    97  	hc.Watch(serviceName, func(status healthcheck.Status) {
    98  		switch secondCalled++; secondCalled {
    99  		case 1:
   100  			require.Equal(t, healthcheck.Serving, status)
   101  		case 2:
   102  			require.Equal(t, healthcheck.NotServing, status)
   103  		default:
   104  			require.FailNow(t, "onStatusChanged should only be called at most 2 times")
   105  		}
   106  	})
   107  
   108  	update(healthcheck.Serving)
   109  	update(healthcheck.NotServing)
   110  	require.Equal(t, 3, firstCalled)
   111  	require.Equal(t, 2, secondCalled)
   112  }
   113  
   114  func requireStatusUnknown(t *testing.T, status healthcheck.Status) {
   115  	require.NotContains(t, []healthcheck.Status{healthcheck.Serving, healthcheck.NotServing}, status)
   116  }