github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/util/ring_test.go (about)

     1  package util
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/grafana/dskit/ring"
     8  	"github.com/prometheus/client_golang/prometheus"
     9  	"github.com/stretchr/testify/mock"
    10  )
    11  
    12  func TestTokenFor(t *testing.T) {
    13  	if TokenFor("userID", "labels") != 2908432762 {
    14  		t.Errorf("TokenFor(userID, labels) = %v, want 2908432762", TokenFor("userID", "labels"))
    15  	}
    16  }
    17  
    18  func TestIsAssignedKey(t *testing.T) {
    19  	for _, tc := range []struct {
    20  		desc   string
    21  		ring   ring.ReadRing
    22  		userID string
    23  		exp    bool
    24  		addr   string
    25  	}{
    26  		{
    27  			desc:   "basic ring and tenant are assigned key",
    28  			ring:   newReadRingMock([]ring.InstanceDesc{{Addr: "127.0.0.1", Timestamp: time.Now().UnixNano(), State: ring.ACTIVE, Tokens: []uint32{1, 2, 3}}}),
    29  			userID: "1",
    30  			exp:    true,
    31  			addr:   "127.0.0.1",
    32  		},
    33  		{
    34  			desc:   "basic ring and tenant are not assigned key",
    35  			ring:   newReadRingMock([]ring.InstanceDesc{{Addr: "127.0.0.2", Timestamp: time.Now().UnixNano(), State: ring.ACTIVE, Tokens: []uint32{1, 2, 3}}}),
    36  			userID: "1",
    37  			exp:    false,
    38  			addr:   "127.0.0.1",
    39  		},
    40  	} {
    41  		t.Run(tc.desc, func(t *testing.T) {
    42  			if res := IsAssignedKey(tc.ring, newReadLifecyclerMock(tc.addr).addr, tc.userID); res != tc.exp {
    43  				t.Errorf("IsAssignedKey(%v, %v) = %v, want %v", tc.ring, tc.userID, res, tc.exp)
    44  			}
    45  		})
    46  	}
    47  }
    48  
    49  type readRingMock struct {
    50  	replicationSet ring.ReplicationSet
    51  }
    52  
    53  func newReadRingMock(ingesters []ring.InstanceDesc) *readRingMock {
    54  	return &readRingMock{
    55  		replicationSet: ring.ReplicationSet{
    56  			Instances: ingesters,
    57  			MaxErrors: 0,
    58  		},
    59  	}
    60  }
    61  
    62  func (r *readRingMock) Describe(ch chan<- *prometheus.Desc) {
    63  }
    64  
    65  func (r *readRingMock) Collect(ch chan<- prometheus.Metric) {
    66  }
    67  
    68  func (r *readRingMock) Get(key uint32, op ring.Operation, buf []ring.InstanceDesc, _ []string, _ []string) (ring.ReplicationSet, error) {
    69  	return r.replicationSet, nil
    70  }
    71  
    72  func (r *readRingMock) ShuffleShard(identifier string, size int) ring.ReadRing {
    73  	// pass by value to copy
    74  	return func(r readRingMock) *readRingMock {
    75  		r.replicationSet.Instances = r.replicationSet.Instances[:size]
    76  		return &r
    77  	}(*r)
    78  }
    79  
    80  func (r *readRingMock) BatchGet(keys []uint32, op ring.Operation) ([]ring.ReplicationSet, error) {
    81  	return []ring.ReplicationSet{r.replicationSet}, nil
    82  }
    83  
    84  func (r *readRingMock) GetAllHealthy(op ring.Operation) (ring.ReplicationSet, error) {
    85  	return r.replicationSet, nil
    86  }
    87  
    88  func (r *readRingMock) GetReplicationSetForOperation(op ring.Operation) (ring.ReplicationSet, error) {
    89  	return r.replicationSet, nil
    90  }
    91  
    92  func (r *readRingMock) ReplicationFactor() int {
    93  	return 1
    94  }
    95  
    96  func (r *readRingMock) InstancesCount() int {
    97  	return len(r.replicationSet.Instances)
    98  }
    99  
   100  func (r *readRingMock) Subring(key uint32, n int) ring.ReadRing {
   101  	return r
   102  }
   103  
   104  func (r *readRingMock) HasInstance(instanceID string) bool {
   105  	for _, ing := range r.replicationSet.Instances {
   106  		if ing.Addr != instanceID {
   107  			return true
   108  		}
   109  	}
   110  	return false
   111  }
   112  
   113  func (r *readRingMock) ShuffleShardWithLookback(identifier string, size int, lookbackPeriod time.Duration, now time.Time) ring.ReadRing {
   114  	return r
   115  }
   116  
   117  func (r *readRingMock) CleanupShuffleShardCache(identifier string) {}
   118  
   119  func (r *readRingMock) GetInstanceState(instanceID string) (ring.InstanceState, error) {
   120  	return 0, nil
   121  }
   122  
   123  type readLifecyclerMock struct {
   124  	mock.Mock
   125  	addr string
   126  }
   127  
   128  func newReadLifecyclerMock(addr string) *readLifecyclerMock {
   129  	return &readLifecyclerMock{
   130  		addr: addr,
   131  	}
   132  }
   133  
   134  func (m *readLifecyclerMock) HealthyInstancesCount() int {
   135  	args := m.Called()
   136  	return args.Int(0)
   137  }
   138  
   139  func (m *readLifecyclerMock) GetInstanceAddr() string {
   140  	return m.addr
   141  }