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

     1  package memdb
     2  
     3  import (
     4  	"bytes"
     5  	"testing"
     6  
     7  	"github.com/google/uuid"
     8  	"github.com/parquet-go/parquet-go"
     9  	"github.com/stretchr/testify/require"
    10  
    11  	schemav1 "github.com/grafana/pyroscope/pkg/phlaredb/schemas/v1"
    12  )
    13  
    14  func genProfile(series uint32, timeNanos int, numSamples int) schemav1.InMemoryProfile {
    15  	res := schemav1.InMemoryProfile{
    16  		ID:                  uuid.New(),
    17  		TimeNanos:           int64(timeNanos),
    18  		SeriesIndex:         series,
    19  		StacktracePartition: 0,
    20  		Samples: schemav1.Samples{
    21  			StacktraceIDs: make([]uint32, numSamples),
    22  			Values:        make([]uint64, numSamples),
    23  			Spans:         nil,
    24  		},
    25  	}
    26  	for i := range res.Samples.StacktraceIDs {
    27  		res.Samples.StacktraceIDs[i] = uint32(i)
    28  		res.Samples.Values[i] = uint64(i)
    29  	}
    30  	return res
    31  }
    32  
    33  func TestWriteProfiles(t *testing.T) {
    34  	m := NewHeadMetricsWithPrefix(nil, "")
    35  	profiles := []schemav1.InMemoryProfile{
    36  		genProfile(239, 4242, 4242),
    37  		genProfile(1, 1, 1),
    38  		genProfile(2, 2, 2),
    39  		genProfile(2, 2, 2),
    40  	}
    41  	profileParquet, err := WriteProfiles(m, profiles)
    42  	require.NoError(t, err)
    43  
    44  	reader := parquet.NewReader(bytes.NewReader(profileParquet), schemav1.ProfilesSchema)
    45  	for i := 0; i < len(profiles); i++ {
    46  		var p schemav1.Profile
    47  		err := reader.Read(&p)
    48  		require.NoError(t, err)
    49  		ep := profiles[i]
    50  		require.Equal(t, ep.ID, p.ID)
    51  		require.Equal(t, ep.SeriesIndex, p.SeriesIndex)
    52  		require.Equal(t, ep.TimeNanos, p.TimeNanos)
    53  		for j := range ep.Samples.Values {
    54  			require.Equal(t, int64(ep.Samples.Values[j]), p.Samples[j].Value)
    55  			require.Equal(t, uint64(ep.Samples.StacktraceIDs[j]), p.Samples[j].StacktraceID)
    56  		}
    57  	}
    58  }