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

     1  package model
     2  
     3  import (
     4  	"math"
     5  	"testing"
     6  
     7  	"github.com/prometheus/common/model"
     8  	"github.com/stretchr/testify/require"
     9  	"go.uber.org/goleak"
    10  
    11  	"github.com/grafana/pyroscope/pkg/iter"
    12  )
    13  
    14  var (
    15  	aLabels = LabelsFromStrings("foo", "a")
    16  	bLabels = LabelsFromStrings("foo", "b")
    17  	cLabels = LabelsFromStrings("foo", "c")
    18  )
    19  
    20  type profile struct {
    21  	labels    Labels
    22  	timestamp model.Time
    23  }
    24  
    25  func (p profile) Labels() Labels {
    26  	return p.labels
    27  }
    28  
    29  func (p profile) Timestamp() model.Time {
    30  	return p.timestamp
    31  }
    32  
    33  func TestMergeIterator(t *testing.T) {
    34  	for _, tt := range []struct {
    35  		name        string
    36  		deduplicate bool
    37  		input       [][]profile
    38  		expected    []profile
    39  	}{
    40  		{
    41  			name:        "deduplicate exact",
    42  			deduplicate: true,
    43  			input: [][]profile{
    44  				{
    45  					{labels: aLabels, timestamp: 1},
    46  					{labels: aLabels, timestamp: 2},
    47  					{labels: aLabels, timestamp: 3},
    48  				},
    49  				{
    50  					{labels: aLabels, timestamp: 1},
    51  					{labels: aLabels, timestamp: 2},
    52  					{labels: aLabels, timestamp: 3},
    53  				},
    54  				{
    55  					{labels: aLabels, timestamp: 1},
    56  					{labels: aLabels, timestamp: 2},
    57  					{labels: aLabels, timestamp: 3},
    58  				},
    59  			},
    60  			expected: []profile{
    61  				{labels: aLabels, timestamp: 1},
    62  				{labels: aLabels, timestamp: 2},
    63  				{labels: aLabels, timestamp: 3},
    64  			},
    65  		},
    66  		{
    67  			name: "no deduplicate",
    68  			input: [][]profile{
    69  				{
    70  					{labels: aLabels, timestamp: 1},
    71  					{labels: aLabels, timestamp: 2},
    72  					{labels: aLabels, timestamp: 3},
    73  				},
    74  				{
    75  					{labels: aLabels, timestamp: 1},
    76  					{labels: aLabels, timestamp: 3},
    77  				},
    78  				{
    79  					{labels: aLabels, timestamp: 2},
    80  				},
    81  			},
    82  			expected: []profile{
    83  				{labels: aLabels, timestamp: 1},
    84  				{labels: aLabels, timestamp: 1},
    85  				{labels: aLabels, timestamp: 2},
    86  				{labels: aLabels, timestamp: 2},
    87  				{labels: aLabels, timestamp: 3},
    88  				{labels: aLabels, timestamp: 3},
    89  			},
    90  		},
    91  		{
    92  			name:        "deduplicate and sort",
    93  			deduplicate: true,
    94  			input: [][]profile{
    95  				{
    96  					{labels: aLabels, timestamp: 1},
    97  					{labels: aLabels, timestamp: 2},
    98  					{labels: aLabels, timestamp: 3},
    99  					{labels: aLabels, timestamp: 4},
   100  				},
   101  				{
   102  					{labels: aLabels, timestamp: 1},
   103  					{labels: cLabels, timestamp: 2},
   104  					{labels: aLabels, timestamp: 3},
   105  				},
   106  				{
   107  					{labels: aLabels, timestamp: 2},
   108  					{labels: bLabels, timestamp: 4},
   109  				},
   110  			},
   111  			expected: []profile{
   112  				{labels: aLabels, timestamp: 1},
   113  				{labels: aLabels, timestamp: 2},
   114  				{labels: cLabels, timestamp: 2},
   115  				{labels: aLabels, timestamp: 3},
   116  				{labels: aLabels, timestamp: 4},
   117  				{labels: bLabels, timestamp: 4},
   118  			},
   119  		},
   120  	} {
   121  		tt := tt
   122  		t.Run(tt.name, func(t *testing.T) {
   123  			iters := make([]iter.Iterator[profile], len(tt.input))
   124  			for i, input := range tt.input {
   125  				iters[i] = iter.NewSliceIterator(input)
   126  			}
   127  			it := NewMergeIterator(
   128  				profile{timestamp: math.MaxInt64},
   129  				tt.deduplicate,
   130  				iters...)
   131  			actual := []profile{}
   132  			for it.Next() {
   133  				actual = append(actual, it.At())
   134  			}
   135  			require.NoError(t, it.Err())
   136  			require.NoError(t, it.Close())
   137  			require.Equal(t, tt.expected, actual)
   138  		})
   139  	}
   140  }
   141  
   142  func Test_BufferedIterator(t *testing.T) {
   143  	for _, tc := range []struct {
   144  		name string
   145  		size int
   146  		in   []profile
   147  	}{
   148  		{
   149  			name: "empty",
   150  			size: 1,
   151  			in:   nil,
   152  		},
   153  		{
   154  			name: "smaller than buffer",
   155  			size: 1000,
   156  			in:   generatesProfiles(t, 100),
   157  		},
   158  		{
   159  			name: "bigger than buffer",
   160  			size: 10,
   161  			in:   generatesProfiles(t, 100),
   162  		},
   163  	} {
   164  		t.Run(tc.name, func(t *testing.T) {
   165  			actual, err := iter.Slice(
   166  				iter.NewBufferedIterator(
   167  					iter.NewSliceIterator(tc.in), tc.size),
   168  			)
   169  			require.NoError(t, err)
   170  			require.Equal(t, tc.in, actual)
   171  		})
   172  	}
   173  }
   174  
   175  func Test_BufferedIteratorClose(t *testing.T) {
   176  	defer goleak.VerifyNone(t, goleak.IgnoreCurrent())
   177  
   178  	it := iter.NewBufferedIterator(
   179  		iter.NewSliceIterator(generatesProfiles(t, 100)), 10)
   180  	require.NoError(t, it.Close())
   181  }
   182  
   183  func generatesProfiles(t *testing.T, n int) []profile {
   184  	t.Helper()
   185  	profiles := make([]profile, n)
   186  	for i := range profiles {
   187  		profiles[i] = profile{labels: aLabels, timestamp: model.Time(i)}
   188  	}
   189  	return profiles
   190  }
   191  
   192  // todo test timedRangeIterator