github.com/masterhung0112/hk_server/v5@v5.0.0-20220302090640-ec71aef15e1c/services/remotecluster/ping_test.go (about)

     1  // Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
     2  // See LICENSE.txt for license information.
     3  
     4  package remotecluster
     5  
     6  import (
     7  	"fmt"
     8  	"net/http"
     9  	"net/http/httptest"
    10  	"sync"
    11  	"sync/atomic"
    12  	"testing"
    13  
    14  	"github.com/stretchr/testify/assert"
    15  	"github.com/stretchr/testify/require"
    16  	"github.com/wiggin77/merror"
    17  
    18  	"github.com/masterhung0112/hk_server/v5/model"
    19  )
    20  
    21  const (
    22  	Recent = 60000
    23  )
    24  
    25  func TestPing(t *testing.T) {
    26  	disablePing = false
    27  
    28  	t.Run("No error", func(t *testing.T) {
    29  		var countWebReq int32
    30  		merr := merror.New()
    31  
    32  		wg := &sync.WaitGroup{}
    33  		wg.Add(NumRemotes)
    34  
    35  		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    36  			defer wg.Done()
    37  			defer w.WriteHeader(200)
    38  			atomic.AddInt32(&countWebReq, 1)
    39  
    40  			frame, err := model.RemoteClusterFrameFromJSON(r.Body)
    41  			if err != nil {
    42  				merr.Append(err)
    43  				return
    44  			}
    45  			if len(frame.Msg.Payload) == 0 {
    46  				merr.Append(fmt.Errorf("Payload should not be empty; remote_id=%s", frame.RemoteId))
    47  				return
    48  			}
    49  
    50  			ping, err := model.RemoteClusterPingFromRawJSON(frame.Msg.Payload)
    51  			if err != nil {
    52  				merr.Append(err)
    53  				return
    54  			}
    55  			if !checkRecent(ping.SentAt, Recent) {
    56  				merr.Append(fmt.Errorf("timestamp out of range, got %d", ping.SentAt))
    57  				return
    58  			}
    59  			if ping.RecvAt != 0 {
    60  				merr.Append(fmt.Errorf("timestamp should be 0, got %d", ping.RecvAt))
    61  				return
    62  			}
    63  		}))
    64  		defer ts.Close()
    65  
    66  		mockServer := newMockServer(t, makeRemoteClusters(NumRemotes, ts.URL))
    67  		defer mockServer.Shutdown()
    68  
    69  		service, err := NewRemoteClusterService(mockServer)
    70  		require.NoError(t, err)
    71  
    72  		err = service.Start()
    73  		require.NoError(t, err)
    74  		defer service.Shutdown()
    75  
    76  		wg.Wait()
    77  
    78  		assert.NoError(t, merr.ErrorOrNil())
    79  
    80  		assert.Equal(t, int32(NumRemotes), atomic.LoadInt32(&countWebReq))
    81  		t.Log(fmt.Sprintf("%d web requests counted;  %d expected",
    82  			atomic.LoadInt32(&countWebReq), NumRemotes))
    83  	})
    84  
    85  	t.Run("HTTP errors", func(t *testing.T) {
    86  		var countWebReq int32
    87  		merr := merror.New()
    88  
    89  		wg := &sync.WaitGroup{}
    90  		wg.Add(NumRemotes)
    91  
    92  		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    93  			defer wg.Done()
    94  			atomic.AddInt32(&countWebReq, 1)
    95  
    96  			frame, err := model.RemoteClusterFrameFromJSON(r.Body)
    97  			if err != nil {
    98  				merr.Append(err)
    99  			}
   100  			ping, err := model.RemoteClusterPingFromRawJSON(frame.Msg.Payload)
   101  			if err != nil {
   102  				merr.Append(err)
   103  			}
   104  			if !checkRecent(ping.SentAt, Recent) {
   105  				merr.Append(fmt.Errorf("timestamp out of range, got %d", ping.SentAt))
   106  			}
   107  			if ping.RecvAt != 0 {
   108  				merr.Append(fmt.Errorf("timestamp should be 0, got %d", ping.RecvAt))
   109  			}
   110  			w.WriteHeader(500)
   111  		}))
   112  		defer ts.Close()
   113  
   114  		mockServer := newMockServer(t, makeRemoteClusters(NumRemotes, ts.URL))
   115  		defer mockServer.Shutdown()
   116  
   117  		service, err := NewRemoteClusterService(mockServer)
   118  		require.NoError(t, err)
   119  
   120  		err = service.Start()
   121  		require.NoError(t, err)
   122  		defer service.Shutdown()
   123  
   124  		wg.Wait()
   125  
   126  		assert.NoError(t, merr.ErrorOrNil())
   127  
   128  		assert.Equal(t, int32(NumRemotes), atomic.LoadInt32(&countWebReq))
   129  		t.Log(fmt.Sprintf("%d web requests counted;  %d expected",
   130  			atomic.LoadInt32(&countWebReq), NumRemotes))
   131  	})
   132  }
   133  
   134  func checkRecent(millis int64, within int64) bool {
   135  	now := model.GetMillis()
   136  	return millis > now-within && millis < now+within
   137  }