github.com/thanos-io/thanos@v0.32.5/pkg/receive/limiter_test.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package receive
     5  
     6  import (
     7  	"context"
     8  	"os"
     9  	"path"
    10  	"testing"
    11  	"time"
    12  
    13  	"github.com/thanos-io/thanos/pkg/extkingpin"
    14  
    15  	"github.com/efficientgo/core/testutil"
    16  	"github.com/go-kit/log"
    17  )
    18  
    19  func TestLimiter_StartConfigReloader(t *testing.T) {
    20  	origLimitsFile, err := os.ReadFile(path.Join("testdata", "limits_config", "good_limits.yaml"))
    21  	testutil.Ok(t, err)
    22  	copyLimitsFile := path.Join(t.TempDir(), "limits.yaml")
    23  	testutil.Ok(t, os.WriteFile(copyLimitsFile, origLimitsFile, 0666))
    24  
    25  	goodLimits, err := extkingpin.NewStaticPathContent(copyLimitsFile)
    26  	if err != nil {
    27  		t.Fatalf("error trying to save static limit config: %s", err)
    28  	}
    29  	invalidLimitsPath := path.Join("./testdata", "limits_config", "invalid_limits.yaml")
    30  	invalidLimits, err := os.ReadFile(invalidLimitsPath)
    31  	if err != nil {
    32  		t.Fatalf("could not load test content at %s: %s", invalidLimitsPath, err)
    33  	}
    34  
    35  	limiter, err := NewLimiter(goodLimits, nil, RouterIngestor, log.NewLogfmtLogger(os.Stdout), 1*time.Second)
    36  	testutil.Ok(t, err)
    37  
    38  	ctx, cancel := context.WithCancel(context.Background())
    39  	defer cancel()
    40  	err = limiter.StartConfigReloader(ctx)
    41  	testutil.Ok(t, err)
    42  
    43  	time.Sleep(1 * time.Second)
    44  	testutil.Ok(t, goodLimits.Rewrite(invalidLimits))
    45  }
    46  
    47  type emptyPathFile struct{}
    48  
    49  func (e emptyPathFile) Content() ([]byte, error) {
    50  	return []byte{}, nil
    51  }
    52  
    53  func (e emptyPathFile) Path() string {
    54  	return ""
    55  }
    56  
    57  func TestLimiter_CanReload(t *testing.T) {
    58  	validLimitsPath, err := extkingpin.NewStaticPathContent(
    59  		path.Join("testdata", "limits_config", "good_limits.yaml"),
    60  	)
    61  	testutil.Ok(t, err)
    62  	emptyLimitsPath := emptyPathFile{}
    63  
    64  	type args struct {
    65  		configFilePath fileContent
    66  	}
    67  	tests := []struct {
    68  		name       string
    69  		args       args
    70  		wantReload bool
    71  	}{
    72  		{
    73  			name:       "Nil config file path cannot be reloaded",
    74  			args:       args{configFilePath: nil},
    75  			wantReload: false,
    76  		},
    77  		{
    78  			name:       "Empty config file path cannot be reloaded",
    79  			args:       args{configFilePath: emptyLimitsPath},
    80  			wantReload: false,
    81  		},
    82  		{
    83  			name:       "Valid config file path can be reloaded",
    84  			args:       args{configFilePath: validLimitsPath},
    85  			wantReload: true,
    86  		},
    87  	}
    88  	for _, tt := range tests {
    89  		t.Run(tt.name, func(t *testing.T) {
    90  			configFile := tt.args.configFilePath
    91  			limiter, err := NewLimiter(configFile, nil, RouterIngestor, log.NewLogfmtLogger(os.Stdout), 1*time.Second)
    92  			testutil.Ok(t, err)
    93  			if tt.wantReload {
    94  				testutil.Assert(t, limiter.CanReload())
    95  			} else {
    96  				testutil.Assert(t, !limiter.CanReload())
    97  			}
    98  		})
    99  	}
   100  }