github.com/grafana/pyroscope@v1.18.0/pkg/validation/retention_overrides_test.go (about)

     1  package validation
     2  
     3  import (
     4  	"bytes"
     5  	"flag"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/prometheus/common/model"
    10  	"github.com/stretchr/testify/assert"
    11  	"github.com/stretchr/testify/require"
    12  
    13  	"github.com/grafana/pyroscope/pkg/metastore/index/cleaner/retention"
    14  )
    15  
    16  func TestRetentionOverridesYAML(t *testing.T) {
    17  	yamlConfig := `
    18  overrides:
    19    tenant-default: {}
    20    tenant-short:
    21      retention_period: 24h
    22    tenant-long:
    23      retention_period: 30d
    24    tenant-infinite:
    25      retention_period: 0
    26  `
    27  
    28  	runtimeConfig, err := LoadRuntimeConfig(bytes.NewReader([]byte(yamlConfig)))
    29  	require.NoError(t, err)
    30  
    31  	var testLimits Limits
    32  	fs := flag.NewFlagSet("test", flag.PanicOnError)
    33  	testLimits.RegisterFlags(fs)
    34  	testLimits.Retention.RegisterFlags(fs)
    35  	expectedDefault := model.Duration(31 * 24 * time.Hour)
    36  
    37  	overrides, err := NewOverrides(testLimits, &mockTenantLimits{
    38  		limits: runtimeConfig.TenantLimits,
    39  		config: runtimeConfig,
    40  	})
    41  	require.NoError(t, err)
    42  	defaults, overridesIter := overrides.Retention()
    43  	assert.Equal(t, expectedDefault, defaults.RetentionPeriod)
    44  
    45  	tenantOverrides := make(map[string]retention.Config)
    46  	for tenantID, config := range overridesIter {
    47  		tenantOverrides[tenantID] = config
    48  	}
    49  
    50  	expectedOverrides := map[string]model.Duration{
    51  		"tenant-short":    model.Duration(24 * time.Hour),
    52  		"tenant-long":     model.Duration(720 * time.Hour),
    53  		"tenant-infinite": 0, // Infinite retention.
    54  	}
    55  
    56  	for tenantID, expectedPeriod := range expectedOverrides {
    57  		config, exists := tenantOverrides[tenantID]
    58  		require.True(t, exists, "tenant %s should have override config", tenantID)
    59  		assert.Equal(t, expectedPeriod, config.RetentionPeriod,
    60  			"tenant %s should have retention period %v", tenantID, expectedPeriod)
    61  	}
    62  }
    63  
    64  func TestRetentionOverrides(t *testing.T) {
    65  	expectedDefault := model.Duration(31 * 24 * time.Hour)
    66  
    67  	tests := []struct {
    68  		name          string
    69  		tenantLimits  map[string]*Limits
    70  		wantOverrides map[string]model.Duration
    71  	}{
    72  		{
    73  			name:          "no overrides",
    74  			tenantLimits:  nil,
    75  			wantOverrides: map[string]model.Duration{},
    76  		},
    77  		{
    78  			name: "with tenant overrides",
    79  			tenantLimits: map[string]*Limits{
    80  				"tenant-a": {Retention: retention.Config{RetentionPeriod: model.Duration(30 * 24 * time.Hour)}},
    81  				"tenant-b": {Retention: retention.Config{RetentionPeriod: 0}}, // infinite
    82  			},
    83  			wantOverrides: map[string]model.Duration{
    84  				"tenant-a": model.Duration(30 * 24 * time.Hour),
    85  				"tenant-b": 0,
    86  			},
    87  		},
    88  	}
    89  
    90  	for _, tt := range tests {
    91  		t.Run(tt.name, func(t *testing.T) {
    92  			var testLimits Limits
    93  			fs := flag.NewFlagSet("test", flag.PanicOnError)
    94  			testLimits.RegisterFlags(fs)
    95  			testLimits.Retention.RegisterFlags(fs)
    96  
    97  			var tenantLimits TenantLimits
    98  			if tt.tenantLimits != nil {
    99  				tenantLimits = &mockTenantLimits{
   100  					limits: tt.tenantLimits,
   101  					config: &RuntimeConfigValues{TenantLimits: tt.tenantLimits},
   102  				}
   103  			}
   104  
   105  			overrides, err := NewOverrides(testLimits, tenantLimits)
   106  			require.NoError(t, err)
   107  
   108  			defaults, overridesIter := overrides.Retention()
   109  			assert.Equal(t, expectedDefault, defaults.RetentionPeriod)
   110  
   111  			actualOverrides := make(map[string]model.Duration)
   112  			for tenantID, config := range overridesIter {
   113  				actualOverrides[tenantID] = config.RetentionPeriod
   114  			}
   115  			assert.Equal(t, tt.wantOverrides, actualOverrides)
   116  
   117  			secondPass := make(map[string]model.Duration)
   118  			for tenantID, config := range overridesIter {
   119  				secondPass[tenantID] = config.RetentionPeriod
   120  			}
   121  			assert.Equal(t, actualOverrides, secondPass)
   122  		})
   123  	}
   124  }