github.com/koko1123/flow-go-1@v0.29.6/utils/liveness/check_test.go (about)

     1  package liveness
     2  
     3  import (
     4  	"net/http"
     5  	"net/http/httptest"
     6  	"testing"
     7  	"time"
     8  )
     9  
    10  func Test_BasicCheck(t *testing.T) {
    11  	mc := NewCheckCollector(time.Millisecond * 20)
    12  	if !mc.IsLive(0) {
    13  		t.Errorf("Multicheck with no checks should always pass")
    14  	}
    15  
    16  	c1 := mc.NewCheck()
    17  	if !mc.IsLive(0) {
    18  		t.Errorf("Just made check should pass")
    19  	}
    20  
    21  	time.Sleep(time.Millisecond * 30)
    22  	if mc.IsLive(0) {
    23  		t.Errorf("Multi check should have failed")
    24  	}
    25  
    26  	c1.CheckIn()
    27  	if !mc.IsLive(0) {
    28  		t.Errorf("Checker should passed after checkin")
    29  	}
    30  
    31  	c2 := mc.NewCheck()
    32  	if !mc.IsLive(0) {
    33  		t.Errorf("Just made check 2 should pass")
    34  	}
    35  
    36  	time.Sleep(time.Millisecond * 30)
    37  	c1.CheckIn()
    38  	// don't checkIn c2
    39  
    40  	if mc.IsLive(0) {
    41  		t.Errorf("Multi check should have failed by c2")
    42  	}
    43  
    44  	c2.CheckIn()
    45  	if !mc.IsLive(0) {
    46  		t.Errorf("Check 2 should pass after checkin")
    47  	}
    48  }
    49  
    50  func Test_CheckHTTP(t *testing.T) {
    51  	c := NewCheckCollector(time.Millisecond * 20)
    52  	r := httptest.NewRequest(http.MethodGet, "/live", nil)
    53  	wr := httptest.NewRecorder()
    54  
    55  	c1 := c.NewCheck()
    56  	_ = c.NewCheck() // never check-in c2
    57  
    58  	c.ServeHTTP(wr, r)
    59  	if wr.Code != http.StatusOK {
    60  		t.Errorf("Check should have passed")
    61  	}
    62  
    63  	time.Sleep(time.Millisecond * 30)
    64  	c1.CheckIn()
    65  
    66  	wr = httptest.NewRecorder()
    67  	c.ServeHTTP(wr, r)
    68  	if wr.Code != http.StatusServiceUnavailable {
    69  		t.Errorf("Check should not have passed")
    70  	}
    71  }
    72  
    73  func Test_CheckHTTPOverride(t *testing.T) {
    74  	c := NewCheckCollector(time.Millisecond * 20)
    75  	r := httptest.NewRequest(http.MethodGet, "/live", nil)
    76  	r.Header.Add(ToleranceHeader, "30s")
    77  	wr := httptest.NewRecorder()
    78  
    79  	c1 := c.NewCheck()
    80  
    81  	c1.CheckIn()
    82  
    83  	c.ServeHTTP(wr, r)
    84  	if wr.Code != http.StatusOK {
    85  		t.Errorf("Check should have passed")
    86  	}
    87  
    88  	time.Sleep(time.Millisecond * 60)
    89  
    90  	wr = httptest.NewRecorder()
    91  	c.ServeHTTP(wr, r)
    92  	if wr.Code != http.StatusOK {
    93  		t.Errorf("Check should still have passed")
    94  	}
    95  
    96  	r.Header.Del(ToleranceHeader)
    97  	wr = httptest.NewRecorder()
    98  	c.ServeHTTP(wr, r)
    99  	if wr.Code != http.StatusServiceUnavailable {
   100  		t.Errorf("Check should not have passed")
   101  	}
   102  }