github.com/grafana/pyroscope@v1.18.0/pkg/ingester/pyroscope/ingest_adapter_test.go (about)

     1  package pyroscope
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	pushv1 "github.com/grafana/pyroscope/api/gen/proto/go/push/v1"
     9  	"github.com/grafana/pyroscope/api/model/labelset"
    10  	"github.com/grafana/pyroscope/pkg/og/storage"
    11  	"github.com/grafana/pyroscope/pkg/og/storage/tree"
    12  	"github.com/grafana/pyroscope/pkg/test/mocks/mockpyroscope"
    13  
    14  	"connectrpc.com/connect"
    15  	"github.com/go-kit/log"
    16  	"github.com/stretchr/testify/assert"
    17  	"github.com/stretchr/testify/mock"
    18  	"github.com/stretchr/testify/require"
    19  )
    20  
    21  func TestPutLabelHandling(t *testing.T) {
    22  	tests := []struct {
    23  		name           string
    24  		labels         map[string]string
    25  		expectedLabels map[string]string
    26  	}{
    27  		{
    28  			name: "No service_name adds service_name",
    29  			labels: map[string]string{
    30  				"__name__": "testapp",
    31  			},
    32  			expectedLabels: map[string]string{
    33  				"service_name": "testapp",
    34  			},
    35  		},
    36  		{
    37  			name: "With service_name adds app_name",
    38  			labels: map[string]string{
    39  				"__name__":     "testapp",
    40  				"service_name": "custom-service",
    41  			},
    42  			expectedLabels: map[string]string{
    43  				"service_name": "custom-service",
    44  				"app_name":     "testapp",
    45  			},
    46  		},
    47  		{
    48  			name: "With service_name and app_name doesn't duplicate app_name",
    49  			labels: map[string]string{
    50  				"__name__":     "testapp",
    51  				"service_name": "custom-service",
    52  				"app_name":     "existing-app",
    53  			},
    54  			expectedLabels: map[string]string{
    55  				"service_name": "custom-service",
    56  				"app_name":     "existing-app",
    57  			},
    58  		},
    59  	}
    60  
    61  	for _, tt := range tests {
    62  		t.Run(tt.name, func(t *testing.T) {
    63  			mockService := mockpyroscope.NewMockPushService(t)
    64  			adapter := &pyroscopeIngesterAdapter{
    65  				svc: mockService,
    66  				log: log.NewNopLogger(),
    67  			}
    68  
    69  			mockService.On("Push", mock.Anything, mock.MatchedBy(func(req *connect.Request[pushv1.PushRequest]) bool {
    70  				if len(req.Msg.Series) > 0 {
    71  					labels := req.Msg.Series[0].Labels
    72  
    73  					// Check expected labels
    74  					labelMap := make(map[string]string)
    75  					labelCounts := make(map[string]int)
    76  					for _, label := range labels {
    77  						labelMap[label.Name] = label.Value
    78  						labelCounts[label.Name]++
    79  					}
    80  
    81  					// Verify expected values
    82  					for key, val := range tt.expectedLabels {
    83  						assert.Equal(t, val, labelMap[key], "Label %s should have value %s", key, val)
    84  					}
    85  
    86  					// Verify no duplicates
    87  					for name, count := range labelCounts {
    88  						assert.Equal(t, 1, count, "Label %s appears %d times, should appear exactly once", name, count)
    89  					}
    90  				}
    91  				return true
    92  			})).Return(&connect.Response[pushv1.PushResponse]{}, nil)
    93  
    94  			ls := labelset.New(tt.labels)
    95  			putInput := &storage.PutInput{
    96  				LabelSet:   ls,
    97  				StartTime:  time.Now(),
    98  				Val:        tree.New(),
    99  				SampleRate: 100,
   100  				SpyName:    "testapp",
   101  			}
   102  
   103  			err := adapter.Put(context.Background(), putInput)
   104  			require.NoError(t, err)
   105  		})
   106  	}
   107  }