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

     1  package usagestats
     2  
     3  import (
     4  	"context"
     5  	"net/http"
     6  	"net/http/httptest"
     7  	"os"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/go-kit/log"
    12  	"github.com/grafana/dskit/kv"
    13  	jsoniter "github.com/json-iterator/go"
    14  	"github.com/prometheus/client_golang/prometheus"
    15  	"github.com/stretchr/testify/require"
    16  
    17  	"github.com/grafana/loki/pkg/storage/chunk/client/local"
    18  )
    19  
    20  func Test_LeaderElection(t *testing.T) {
    21  	stabilityCheckInterval = 100 * time.Millisecond
    22  
    23  	result := make(chan *ClusterSeed, 10)
    24  	objectClient, err := local.NewFSObjectClient(local.FSConfig{
    25  		Directory: t.TempDir(),
    26  	})
    27  	require.NoError(t, err)
    28  	for i := 0; i < 3; i++ {
    29  		go func() {
    30  			r, err := NewReporter(Config{Leader: true, Enabled: true}, kv.Config{
    31  				Store: "inmemory",
    32  			}, objectClient, log.NewLogfmtLogger(os.Stdout), nil)
    33  			require.NoError(t, err)
    34  			r.init(context.Background())
    35  			result <- r.cluster
    36  		}()
    37  	}
    38  	for i := 0; i < 7; i++ {
    39  		go func() {
    40  			r, err := NewReporter(Config{Leader: false, Enabled: true}, kv.Config{
    41  				Store: "inmemory",
    42  			}, objectClient, log.NewLogfmtLogger(os.Stdout), nil)
    43  			require.NoError(t, err)
    44  			r.init(context.Background())
    45  			result <- r.cluster
    46  		}()
    47  	}
    48  
    49  	var UID []string
    50  	for i := 0; i < 10; i++ {
    51  		cluster := <-result
    52  		require.NotNil(t, cluster)
    53  		UID = append(UID, cluster.UID)
    54  	}
    55  	first := UID[0]
    56  	for _, uid := range UID {
    57  		require.Equal(t, first, uid)
    58  	}
    59  	kvClient, err := kv.NewClient(kv.Config{Store: "inmemory"}, JSONCodec, nil, log.NewLogfmtLogger(os.Stdout))
    60  	require.NoError(t, err)
    61  	// verify that the ID found is also correctly stored in the kv store and not overridden by another leader.
    62  	data, err := kvClient.Get(context.Background(), seedKey)
    63  	require.NoError(t, err)
    64  	require.Equal(t, data.(*ClusterSeed).UID, first)
    65  }
    66  
    67  func Test_ReportLoop(t *testing.T) {
    68  	// stub
    69  	reportCheckInterval = 100 * time.Millisecond
    70  	reportInterval = time.Second
    71  	stabilityCheckInterval = 100 * time.Millisecond
    72  
    73  	totalReport := 0
    74  	clusterIDs := []string{}
    75  	server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
    76  		var received Report
    77  		totalReport++
    78  		require.NoError(t, jsoniter.NewDecoder(r.Body).Decode(&received))
    79  		clusterIDs = append(clusterIDs, received.ClusterID)
    80  		rw.WriteHeader(http.StatusOK)
    81  	}))
    82  	usageStatsURL = server.URL
    83  
    84  	objectClient, err := local.NewFSObjectClient(local.FSConfig{
    85  		Directory: t.TempDir(),
    86  	})
    87  	require.NoError(t, err)
    88  
    89  	r, err := NewReporter(Config{Leader: true, Enabled: true}, kv.Config{
    90  		Store: "inmemory",
    91  	}, objectClient, log.NewLogfmtLogger(os.Stdout), prometheus.NewPedanticRegistry())
    92  	require.NoError(t, err)
    93  	ctx, cancel := context.WithCancel(context.Background())
    94  	r.initLeader(ctx)
    95  
    96  	go func() {
    97  		<-time.After(6*time.Second + (stabilityCheckInterval * time.Duration(stabilityMinimunRequired+1)))
    98  		cancel()
    99  	}()
   100  	require.Equal(t, nil, r.running(ctx))
   101  	require.GreaterOrEqual(t, totalReport, 5)
   102  	first := clusterIDs[0]
   103  	for _, uid := range clusterIDs {
   104  		require.Equal(t, first, uid)
   105  	}
   106  	require.Equal(t, first, r.cluster.UID)
   107  }
   108  
   109  func Test_NextReport(t *testing.T) {
   110  	fixtures := map[string]struct {
   111  		interval  time.Duration
   112  		createdAt time.Time
   113  		now       time.Time
   114  
   115  		next time.Time
   116  	}{
   117  		"createdAt aligned with interval and now": {
   118  			interval:  1 * time.Hour,
   119  			createdAt: time.Unix(0, time.Hour.Nanoseconds()),
   120  			now:       time.Unix(0, 2*time.Hour.Nanoseconds()),
   121  			next:      time.Unix(0, 2*time.Hour.Nanoseconds()),
   122  		},
   123  		"createdAt aligned with interval": {
   124  			interval:  1 * time.Hour,
   125  			createdAt: time.Unix(0, time.Hour.Nanoseconds()),
   126  			now:       time.Unix(0, 2*time.Hour.Nanoseconds()+1),
   127  			next:      time.Unix(0, 3*time.Hour.Nanoseconds()),
   128  		},
   129  		"createdAt not aligned": {
   130  			interval:  1 * time.Hour,
   131  			createdAt: time.Unix(0, time.Hour.Nanoseconds()+18*time.Minute.Nanoseconds()+20*time.Millisecond.Nanoseconds()),
   132  			now:       time.Unix(0, 2*time.Hour.Nanoseconds()+1),
   133  			next:      time.Unix(0, 2*time.Hour.Nanoseconds()+18*time.Minute.Nanoseconds()+20*time.Millisecond.Nanoseconds()),
   134  		},
   135  	}
   136  	for name, f := range fixtures {
   137  		t.Run(name, func(t *testing.T) {
   138  			next := nextReport(f.interval, f.createdAt, f.now)
   139  			require.Equal(t, f.next, next)
   140  		})
   141  	}
   142  }
   143  
   144  func TestWrongKV(t *testing.T) {
   145  	objectClient, err := local.NewFSObjectClient(local.FSConfig{
   146  		Directory: t.TempDir(),
   147  	})
   148  	require.NoError(t, err)
   149  
   150  	r, err := NewReporter(Config{Leader: true, Enabled: true}, kv.Config{
   151  		Store: "",
   152  	}, objectClient, log.NewLogfmtLogger(os.Stdout), prometheus.NewPedanticRegistry())
   153  	require.NoError(t, err)
   154  	ctx, cancel := context.WithCancel(context.Background())
   155  	go func() {
   156  		<-time.After(1 * time.Second)
   157  		cancel()
   158  	}()
   159  	require.Equal(t, nil, r.running(ctx))
   160  }