github.com/grafana/pyroscope@v1.18.0/pkg/distributor/model/push_test.go (about)

     1  package model
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/stretchr/testify/assert"
     8  	"github.com/stretchr/testify/require"
     9  
    10  	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
    11  	"github.com/grafana/pyroscope/pkg/distributor/annotation"
    12  	"github.com/grafana/pyroscope/pkg/distributor/ingestlimits"
    13  )
    14  
    15  func TestProfileSeries_GetLanguage(t *testing.T) {
    16  	tests := []struct {
    17  		labels []*typesv1.LabelPair
    18  		want   string
    19  	}{
    20  		{labels: []*typesv1.LabelPair{{Name: "pyroscope_spy", Value: "gospy"}}, want: "go"},
    21  		{labels: []*typesv1.LabelPair{{Name: "pyroscope_spy", Value: "javaspy"}}, want: "java"},
    22  		{labels: []*typesv1.LabelPair{{Name: "pyroscope_spy", Value: "dotnetspy"}}, want: "dotnet"},
    23  		{labels: []*typesv1.LabelPair{{Name: "pyroscope_spy", Value: "grafana-agent.java"}}, want: "java"},
    24  		{labels: []*typesv1.LabelPair{{Name: "pyroscope_spy", Value: ""}}, want: ""},
    25  	}
    26  	for _, tt := range tests {
    27  		t.Run("", func(t *testing.T) {
    28  			p := &ProfileSeries{
    29  				Labels: tt.labels,
    30  			}
    31  			if got := p.GetLanguage(); got != tt.want {
    32  				t.Errorf("GetLanguage() = %v, want %v", got, tt.want)
    33  			}
    34  		})
    35  	}
    36  }
    37  
    38  func TestMarkThrottledTenant(t *testing.T) {
    39  	tests := []struct {
    40  		name        string
    41  		req         *ProfileSeries
    42  		limit       *ingestlimits.Config
    43  		expectError bool
    44  		verify      func(t *testing.T, req *ProfileSeries)
    45  	}{
    46  		{
    47  			name: "single series",
    48  			req: &ProfileSeries{
    49  				Labels: []*typesv1.LabelPair{
    50  					{Name: "__name__", Value: "cpu"},
    51  				},
    52  			},
    53  			limit: &ingestlimits.Config{
    54  				PeriodType:     "hour",
    55  				PeriodLimitMb:  128,
    56  				LimitResetTime: time.Now().Unix(),
    57  				LimitReached:   true,
    58  			},
    59  			verify: func(t *testing.T, req *ProfileSeries) {
    60  				require.Len(t, req.Annotations, 1)
    61  				assert.Equal(t, annotation.ProfileAnnotationKeyThrottled, req.Annotations[0].Key)
    62  				assert.Contains(t, req.Annotations[0].Value, "\"periodLimitMb\":128")
    63  			},
    64  		},
    65  	}
    66  
    67  	for _, tt := range tests {
    68  		t.Run(tt.name, func(t *testing.T) {
    69  			err := tt.req.MarkThrottledTenant(tt.limit)
    70  			if tt.expectError {
    71  				assert.Error(t, err)
    72  				return
    73  			}
    74  			assert.NoError(t, err)
    75  			if tt.verify != nil {
    76  				tt.verify(t, tt.req)
    77  			}
    78  		})
    79  	}
    80  }
    81  
    82  func TestMarkThrottledUsageGroup(t *testing.T) {
    83  	tests := []struct {
    84  		name        string
    85  		req         *ProfileSeries
    86  		limit       *ingestlimits.Config
    87  		usageGroup  string
    88  		expectError bool
    89  		verify      func(t *testing.T, req *ProfileSeries)
    90  	}{
    91  		{
    92  			name: "single series with usage group",
    93  			req: &ProfileSeries{
    94  				Labels: []*typesv1.LabelPair{
    95  					{Name: "__name__", Value: "cpu"},
    96  				},
    97  			},
    98  			limit: &ingestlimits.Config{
    99  				PeriodType:     "hour",
   100  				PeriodLimitMb:  128,
   101  				LimitResetTime: time.Now().Unix(),
   102  				LimitReached:   true,
   103  				UsageGroups: map[string]ingestlimits.UsageGroup{
   104  					"group-1": {
   105  						PeriodLimitMb: 64,
   106  						LimitReached:  true,
   107  					},
   108  				},
   109  			},
   110  			usageGroup: "group-1",
   111  			verify: func(t *testing.T, req *ProfileSeries) {
   112  				require.Len(t, req.Annotations, 1)
   113  				assert.Equal(t, annotation.ProfileAnnotationKeyThrottled, req.Annotations[0].Key)
   114  				assert.Contains(t, req.Annotations[0].Value, "\"periodLimitMb\":64")
   115  				assert.Contains(t, req.Annotations[0].Value, "\"usageGroup\":\"group-1\"")
   116  			},
   117  		},
   118  		{
   119  			name: "invalid usage group",
   120  			req: &ProfileSeries{
   121  				Labels: []*typesv1.LabelPair{
   122  					{Name: "__name__", Value: "cpu"},
   123  				},
   124  			},
   125  			limit: &ingestlimits.Config{
   126  				PeriodType:     "hour",
   127  				PeriodLimitMb:  128,
   128  				LimitResetTime: time.Now().Unix(),
   129  				LimitReached:   true,
   130  				UsageGroups: map[string]ingestlimits.UsageGroup{
   131  					"group-1": {
   132  						PeriodLimitMb: 64,
   133  						LimitReached:  true,
   134  					},
   135  				},
   136  			},
   137  			usageGroup:  "nonexistent-group",
   138  			expectError: true,
   139  		},
   140  	}
   141  
   142  	for _, tt := range tests {
   143  		t.Run(tt.name, func(t *testing.T) {
   144  			err := tt.req.MarkThrottledUsageGroup(tt.limit, tt.usageGroup)
   145  			if tt.expectError {
   146  				assert.Error(t, err)
   147  				return
   148  			}
   149  			assert.NoError(t, err)
   150  			if tt.verify != nil {
   151  				tt.verify(t, tt.req)
   152  			}
   153  		})
   154  	}
   155  }