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

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  // Provenance-includes-location: https://github.com/cortexproject/cortex/blob/master/pkg/util/validation/limits_test.go
     3  // Provenance-includes-license: Apache-2.0
     4  // Provenance-includes-copyright: The Cortex Authors.
     5  
     6  package validation
     7  
     8  import "github.com/grafana/dskit/flagext"
     9  
    10  // mockTenantLimits exposes per-tenant limits based on a provided map
    11  type mockTenantLimits struct {
    12  	limits map[string]*Limits
    13  	config *RuntimeConfigValues
    14  }
    15  
    16  // NewMockTenantLimits creates a new mockTenantLimits that returns per-tenant limits based on
    17  // the given map
    18  func NewMockTenantLimits(limits map[string]*Limits) TenantLimits {
    19  	return &mockTenantLimits{
    20  		limits: limits,
    21  	}
    22  }
    23  
    24  func (l *mockTenantLimits) TenantLimits(userID string) *Limits {
    25  	return l.limits[userID]
    26  }
    27  
    28  func (l *mockTenantLimits) AllByTenantID() map[string]*Limits {
    29  	return l.limits
    30  }
    31  
    32  func (l *mockTenantLimits) RuntimeConfig() *RuntimeConfigValues { return l.config }
    33  
    34  func MockOverrides(customize func(defaults *Limits, tenantLimits map[string]*Limits)) *Overrides {
    35  	defaults := MockDefaultLimits()
    36  	tenantLimits := map[string]*Limits{}
    37  	customize(defaults, tenantLimits)
    38  
    39  	overrides, err := NewOverrides(*defaults, NewMockTenantLimits(tenantLimits))
    40  	if err != nil {
    41  		// This function is expected to be used only in tests, so we're not afraid of panicking.
    42  		panic(err)
    43  	}
    44  
    45  	return overrides
    46  }
    47  
    48  func MockDefaultLimits() *Limits {
    49  	defaults := Limits{}
    50  	flagext.DefaultValues(&defaults)
    51  	defaults.RejectNewerThan = 0
    52  	defaults.RejectOlderThan = 0
    53  	return &defaults
    54  }
    55  
    56  func MockDefaultOverrides() *Overrides {
    57  	defaults := MockDefaultLimits()
    58  	overrides, err := NewOverrides(*defaults, nil)
    59  	if err != nil {
    60  		// This function is expected to be used only in tests, so we're not afraid of panicking.
    61  		panic(err)
    62  	}
    63  
    64  	return overrides
    65  }