Files
grafana/pkg/tsdb/grafana-pyroscope-datasource/annotation/annotation_test.go
T
Aleksandar Petrov e3f5a65372 Pyroscope: Process and display sampling annotations (#109707)
* pyroscope: process sampling annotations

* Enable annotations in classic explore

* Run prettier

* Revert unneeded change to plugin.json

* Tweak wording in sampling annotation

* Fix test

* Disable annotations by default
2025-08-29 13:14:22 +02:00

356 lines
10 KiB
Go

package annotation
import (
"testing"
"time"
"github.com/grafana/grafana-plugin-sdk-go/data"
typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
"github.com/stretchr/testify/require"
)
func TestConvertAnnotation(t *testing.T) {
t.Run("converts a valid throttling annotation", func(t *testing.T) {
rawAnnotation := `{"body":{"periodType":"day","periodLimitMb":1024,"limitResetTime":1609459200}}`
timedAnnotation := &TimedAnnotation{
Timestamp: 1609455600000,
Annotation: &typesv1.ProfileAnnotation{
Key: string(ProfileAnnotationKeyThrottled),
Value: rawAnnotation,
},
}
processed, err := convertAnnotation(timedAnnotation)
require.NoError(t, err)
require.NotNil(t, processed)
require.Contains(t, processed.text, "Ingestion limit (1.0 GiB/day) reached")
require.Contains(t, processed.text, "day")
require.Equal(t, int64(1609455600000), processed.time)
require.Equal(t, int64(1609459200000), processed.timeEnd) // LimitResetTime * 1000
})
t.Run("converts a valid sampling annotation", func(t *testing.T) {
rawAnnotation := `{"body":{"source": {"usageGroup":"group-1","probability":0.1}}}`
timedAnnotation := &TimedAnnotation{
Timestamp: 1609455600000,
Annotation: &typesv1.ProfileAnnotation{
Key: string(ProfileAnnotationKeySampled),
Value: rawAnnotation,
},
}
processed, err := convertAnnotation(timedAnnotation)
require.NoError(t, err)
require.NotNil(t, processed)
require.Contains(t, processed.text, "Profile volume reduced by 90.00% for this service.")
require.Equal(t, int64(1609455600000), processed.time)
require.Equal(t, int64(1609455600000), processed.timeEnd)
})
t.Run("ignores non-throttling annotations", func(t *testing.T) {
timedAnnotation := &TimedAnnotation{
Timestamp: 1000,
Annotation: &typesv1.ProfileAnnotation{
Key: "some.other.key",
Value: `{"test":"value"}`,
},
}
processed, err := convertAnnotation(timedAnnotation)
require.NoError(t, err)
require.Nil(t, processed)
})
t.Run("handles invalid annotation data", func(t *testing.T) {
timedAnnotation := &TimedAnnotation{
Timestamp: 1000,
Annotation: &typesv1.ProfileAnnotation{
Key: string(ProfileAnnotationKeyThrottled),
Value: `invalid json`,
},
}
processed, err := convertAnnotation(timedAnnotation)
require.Error(t, err)
require.Nil(t, processed)
require.Contains(t, err.Error(), "error parsing annotation data")
})
}
func TestProcessAnnotations(t *testing.T) {
rawAnnotation := `{"body":{"periodType":"day","periodLimitMb":1024,"limitResetTime":1609459200}}`
t.Run("processes multiple annotations", func(t *testing.T) {
annotations := []*TimedAnnotation{
{
Timestamp: 1609455600000,
Annotation: &typesv1.ProfileAnnotation{
Key: string(ProfileAnnotationKeyThrottled),
Value: rawAnnotation,
},
},
{
Timestamp: 1609459200000,
Annotation: &typesv1.ProfileAnnotation{
Key: string(ProfileAnnotationKeyThrottled),
Value: rawAnnotation,
},
},
}
result, err := processAnnotations(annotations)
require.NoError(t, err)
require.Equal(t, 1, len(result.times))
require.Equal(t, 1, len(result.timeEnds))
require.Equal(t, 1, len(result.texts))
require.Equal(t, 1, len(result.isRegions))
})
t.Run("handles empty annotations list", func(t *testing.T) {
result, err := processAnnotations([]*TimedAnnotation{})
require.NoError(t, err)
require.Equal(t, 0, len(result.times))
require.Equal(t, 0, len(result.timeEnds))
require.Equal(t, 0, len(result.texts))
require.Equal(t, 0, len(result.isRegions))
})
t.Run("handles nil annotations", func(t *testing.T) {
annotations := []*TimedAnnotation{nil}
result, err := processAnnotations(annotations)
require.NoError(t, err)
require.Equal(t, 0, len(result.times))
})
t.Run("handles invalid annotation data", func(t *testing.T) {
annotations := []*TimedAnnotation{
{
Timestamp: 1000,
Annotation: &typesv1.ProfileAnnotation{
Key: string(ProfileAnnotationKeyThrottled),
Value: `invalid json`,
},
},
}
result, err := processAnnotations(annotations)
require.Error(t, err)
require.Nil(t, result)
require.Contains(t, err.Error(), "error parsing annotation data")
})
}
func TestGrafanaAnnotationDataAdd(t *testing.T) {
t.Run("adds first annotation", func(t *testing.T) {
ga := &grafanaAnnotationData{
ids: []string{},
times: []time.Time{},
timeEnds: []time.Time{},
texts: []string{},
isRegions: []bool{},
}
annotation := &processedProfileAnnotation{
id: "test-id-1",
text: "Test annotation 1",
time: 1609455600000,
timeEnd: 1609459200000,
isRegion: true,
}
ga.add(annotation)
require.Equal(t, 1, len(ga.ids))
require.Equal(t, "test-id-1", ga.ids[0])
require.Equal(t, time.UnixMilli(1609455600000), ga.times[0])
require.Equal(t, time.UnixMilli(1609459200000), ga.timeEnds[0])
require.Equal(t, "Test annotation 1", ga.texts[0])
require.Equal(t, true, ga.isRegions[0])
})
t.Run("adds different annotations", func(t *testing.T) {
ga := &grafanaAnnotationData{
ids: []string{},
times: []time.Time{},
timeEnds: []time.Time{},
texts: []string{},
isRegions: []bool{},
}
annotation1 := &processedProfileAnnotation{
id: "test-id-1",
text: "Test annotation 1",
time: 1609455600000,
timeEnd: 1609459200000,
isRegion: true,
}
annotation2 := &processedProfileAnnotation{
id: "test-id-2",
text: "Test annotation 2",
time: 1609463800000,
timeEnd: 1609467400000,
isRegion: false,
}
ga.add(annotation1)
ga.add(annotation2)
require.Equal(t, 2, len(ga.ids))
require.Equal(t, "test-id-1", ga.ids[0])
require.Equal(t, "test-id-2", ga.ids[1])
require.Equal(t, time.UnixMilli(1609455600000), ga.times[0])
require.Equal(t, time.UnixMilli(1609463800000), ga.times[1])
})
t.Run("removes duplicates and extends timeEnd", func(t *testing.T) {
ga := &grafanaAnnotationData{
ids: []string{},
times: []time.Time{},
timeEnds: []time.Time{},
texts: []string{},
isRegions: []bool{},
}
annotation1 := &processedProfileAnnotation{
id: "duplicate-id",
text: "First occurrence",
time: 1609455600000,
timeEnd: 1609459200000,
isRegion: true,
}
annotation2 := &processedProfileAnnotation{
id: "duplicate-id",
text: "Second occurrence (should be ignored)",
time: 1609460000000,
timeEnd: 1609463600000,
isRegion: false,
}
ga.add(annotation1)
ga.add(annotation2)
require.Equal(t, 1, len(ga.ids))
require.Equal(t, 1, len(ga.times))
require.Equal(t, 1, len(ga.timeEnds))
require.Equal(t, 1, len(ga.texts))
require.Equal(t, 1, len(ga.isRegions))
require.Equal(t, "duplicate-id", ga.ids[0])
require.Equal(t, time.UnixMilli(1609455600000), ga.times[0]) // Original time
require.Equal(t, time.UnixMilli(1609463600000), ga.timeEnds[0]) // Extended timeEnd
require.Equal(t, "First occurrence", ga.texts[0]) // Original text
require.Equal(t, true, ga.isRegions[0]) // Original isRegion
})
t.Run("handles multiple duplicates correctly", func(t *testing.T) {
ga := &grafanaAnnotationData{
ids: []string{},
times: []time.Time{},
timeEnds: []time.Time{},
texts: []string{},
isRegions: []bool{},
}
annotation1 := &processedProfileAnnotation{
id: "id-1",
text: "Annotation 1",
time: 1609455600000,
timeEnd: 1609459200000,
isRegion: true,
}
// Add duplicate of first
annotation1Duplicate := &processedProfileAnnotation{
id: "id-1",
text: "Annotation 1 duplicate",
time: 1609460000000,
timeEnd: 1609470000000,
isRegion: false,
}
// Add a second, unique annotation
annotation2 := &processedProfileAnnotation{
id: "id-2",
text: "Annotation 2",
time: 1609480000000,
timeEnd: 1609490000000,
isRegion: false,
}
// Add duplicate of second
annotation2Duplicate := &processedProfileAnnotation{
id: "id-2",
text: "Annotation 2 duplicate",
time: 1609500000000,
timeEnd: 1609510000000,
isRegion: true,
}
ga.add(annotation1)
ga.add(annotation1Duplicate)
ga.add(annotation2)
ga.add(annotation2Duplicate)
require.Equal(t, 2, len(ga.ids))
require.Equal(t, "id-1", ga.ids[0])
require.Equal(t, "id-2", ga.ids[1])
// The first annotation should have an extended timeEnd
require.Equal(t, time.UnixMilli(1609455600000), ga.times[0])
require.Equal(t, time.UnixMilli(1609470000000), ga.timeEnds[0])
require.Equal(t, "Annotation 1", ga.texts[0])
require.Equal(t, true, ga.isRegions[0])
// The second annotation should have an extended timeEnd
require.Equal(t, time.UnixMilli(1609480000000), ga.times[1])
require.Equal(t, time.UnixMilli(1609510000000), ga.timeEnds[1])
require.Equal(t, "Annotation 2", ga.texts[1])
require.Equal(t, false, ga.isRegions[1])
})
}
func TestCreateAnnotationFrame(t *testing.T) {
rawAnnotation := `{"body":{"periodType":"day","periodLimitMb":1024,"limitResetTime":1609459200}}`
t.Run("creates frame with correct fields", func(t *testing.T) {
annotations := []*TimedAnnotation{
{
Timestamp: 1609455600000,
Annotation: &typesv1.ProfileAnnotation{
Key: string(ProfileAnnotationKeyThrottled),
Value: rawAnnotation,
},
},
}
frame, err := CreateAnnotationFrame(annotations)
require.NoError(t, err)
require.NotNil(t, frame)
require.Equal(t, "annotations", frame.Name)
require.Equal(t, data.DataTopicAnnotations, frame.Meta.DataTopic)
require.Equal(t, 5, len(frame.Fields))
require.Equal(t, "time", frame.Fields[0].Name)
require.Equal(t, "timeEnd", frame.Fields[1].Name)
require.Equal(t, "text", frame.Fields[2].Name)
require.Equal(t, "isRegion", frame.Fields[3].Name)
require.Equal(t, "color", frame.Fields[4].Name)
require.Equal(t, 1, frame.Fields[0].Len())
require.Equal(t, time.UnixMilli(1609455600000), frame.Fields[0].At(0))
require.Equal(t, time.UnixMilli(1609459200000), frame.Fields[1].At(0))
require.Contains(t, frame.Fields[2].At(0).(string), "Ingestion limit")
})
t.Run("handles empty annotations list", func(t *testing.T) {
frame, err := CreateAnnotationFrame([]*TimedAnnotation{})
require.NoError(t, err)
require.NotNil(t, frame)
require.Equal(t, 5, len(frame.Fields))
require.Equal(t, 0, frame.Fields[0].Len())
})
}