github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/datastore/crdb/pool/health_test.go (about)

     1  package pool
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/stretchr/testify/require"
     8  	"golang.org/x/time/rate"
     9  )
    10  
    11  func TestNodeHealthTracker(t *testing.T) {
    12  	tracker := &NodeHealthTracker{
    13  		healthyNodes:  make(map[uint32]struct{}),
    14  		nodesEverSeen: make(map[uint32]*rate.Limiter),
    15  		newLimiter: func() *rate.Limiter {
    16  			return rate.NewLimiter(rate.Every(1*time.Minute), 2)
    17  		},
    18  	}
    19  
    20  	tracker.SetNodeHealth(1, true)
    21  	require.True(t, tracker.IsHealthy(1))
    22  	require.False(t, tracker.IsHealthy(2))
    23  	require.Equal(t, tracker.HealthyNodeCount(), 1)
    24  
    25  	tracker.SetNodeHealth(2, true)
    26  	require.True(t, tracker.IsHealthy(1))
    27  	require.True(t, tracker.IsHealthy(2))
    28  	require.Equal(t, tracker.HealthyNodeCount(), 2)
    29  
    30  	// just 1 mark isn't enough to trigger false
    31  	tracker.SetNodeHealth(1, false)
    32  	require.True(t, tracker.IsHealthy(1))
    33  	require.True(t, tracker.IsHealthy(2))
    34  	require.Equal(t, tracker.HealthyNodeCount(), 2)
    35  
    36  	tracker.SetNodeHealth(1, false)
    37  	tracker.SetNodeHealth(1, false)
    38  	require.False(t, tracker.IsHealthy(1))
    39  	require.True(t, tracker.IsHealthy(2))
    40  	require.Equal(t, tracker.HealthyNodeCount(), 1)
    41  }