github.com/iotexproject/iotex-core@v1.14.1-rc1/pkg/probe/probe_test.go (about)

     1  // Copyright (c) 2019 IoTeX Foundation
     2  // This source code is provided 'as is' and no warranties are given as to title or non-infringement, merchantability
     3  // or fitness for purpose and, to the extent permitted by law, all liability for your use of the code is disclaimed.
     4  // This source code is governed by Apache License 2.0 that can be found in the LICENSE file.
     5  
     6  package probe
     7  
     8  import (
     9  	"context"
    10  	"net/http"
    11  	"testing"
    12  	"time"
    13  
    14  	"github.com/stretchr/testify/assert"
    15  	"github.com/stretchr/testify/require"
    16  
    17  	"github.com/iotexproject/iotex-core/testutil"
    18  )
    19  
    20  type testCase struct {
    21  	endpoint string
    22  	code     int
    23  }
    24  
    25  func testFunc(t *testing.T, ts []testCase) {
    26  	for _, tt := range ts {
    27  		resp, err := http.Get("http://localhost:7788" + tt.endpoint)
    28  		require.NoError(t, err)
    29  		assert.Equal(t, tt.code, resp.StatusCode)
    30  	}
    31  }
    32  
    33  func TestBasicProbe(t *testing.T) {
    34  	s := New(7788)
    35  	ctx := context.Background()
    36  	require.NoError(t, s.Start(ctx))
    37  	require.NoError(t, testutil.WaitUntil(100*time.Millisecond, 2*time.Second, func() (b bool, e error) {
    38  		_, err := http.Get("http://localhost:7788/liveness")
    39  		return err == nil, nil
    40  	}))
    41  	test1 := []testCase{
    42  		{
    43  			endpoint: "/liveness",
    44  			code:     http.StatusOK,
    45  		},
    46  		{
    47  			endpoint: "/readiness",
    48  			code:     http.StatusServiceUnavailable,
    49  		},
    50  		{
    51  			endpoint: "/health",
    52  			code:     http.StatusServiceUnavailable,
    53  		},
    54  	}
    55  	testFunc(t, test1)
    56  
    57  	test2 := []testCase{
    58  		{
    59  			endpoint: "/liveness",
    60  			code:     http.StatusOK,
    61  		},
    62  		{
    63  			endpoint: "/readiness",
    64  			code:     http.StatusOK,
    65  		},
    66  		{
    67  			endpoint: "/health",
    68  			code:     http.StatusOK,
    69  		},
    70  	}
    71  	require.NoError(t, s.TurnOn())
    72  	testFunc(t, test2)
    73  	require.NoError(t, s.TurnOff())
    74  	testFunc(t, test1)
    75  
    76  	require.NoError(t, s.Stop(ctx))
    77  	_, err := http.Get("http://localhost:7788/liveness")
    78  	require.Error(t, err)
    79  }
    80  
    81  func TestReadniessHandler(t *testing.T) {
    82  	ctx := context.Background()
    83  	s := New(7788, WithReadinessHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    84  		w.WriteHeader(http.StatusAccepted)
    85  	})))
    86  	defer s.Stop(ctx)
    87  
    88  	require.NoError(t, s.Start(ctx))
    89  	require.NoError(t, testutil.WaitUntil(100*time.Millisecond, 2*time.Second, func() (b bool, e error) {
    90  		_, err := http.Get("http://localhost:7788/liveness")
    91  		return err == nil, nil
    92  	}))
    93  	test := []testCase{
    94  		{
    95  			endpoint: "/liveness",
    96  			code:     http.StatusOK,
    97  		},
    98  		{
    99  			endpoint: "/readiness",
   100  			code:     http.StatusAccepted,
   101  		},
   102  		{
   103  			endpoint: "/health",
   104  			code:     http.StatusAccepted,
   105  		},
   106  	}
   107  	require.NoError(t, s.TurnOn())
   108  	testFunc(t, test)
   109  }