chore: improve go pprof fix tests

Author: kolesnikovae

pkg/pprof/fix_go_profile_test.go

@@ -7,12 +7,15 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	profilev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
 )
 
 func Test_FixGoProfile(t *testing.T) {
 	p, err := OpenFile("testdata/gotruncatefix/heap_go_truncated_4.pb.gz")
 	require.NoError(t, err)
 
+	total := samplesTotal(p.Profile)
 	f := FixGoProfile(p.Profile)
 	s := make(map[string]struct{})
 	for _, x := range f.StringTable {
@@ -23,6 +26,11 @@ func Test_FixGoProfile(t *testing.T) {
 		}
 	}
 
+	// Assert that the total sum of samples does not
+	// change compared to the source profile.
+	assert.Equal(t, total, samplesTotal(p.Profile))
+	assert.Equal(t, total, samplesTotal(f))
+
 	t.Logf(" * Sample:   %6d -> %-6d", len(p.Sample), len(f.Sample))
 	t.Logf(" * Location: %6d -> %-6d", len(p.Location), len(f.Location))
 	t.Logf(" * Function: %6d -> %-6d", len(p.Function), len(f.Function))
@@ -35,6 +43,20 @@ func Test_FixGoProfile(t *testing.T) {
 	assert.Equal(t, 168, len(p.Location)-len(f.Location))
 	assert.Equal(t, 77, len(p.Function)-len(f.Function))
 	assert.Equal(t, 78, len(p.StringTable)-len(f.StringTable))
+
+	// Now we can normalize the profile aggregate samples.
+	// This must not affect the total sum of values.
+	x := RawFromProto(f)
+	x.Normalize()
+	assert.Equal(t, total, samplesTotal(x.Profile))
+}
+
+func samplesTotal(p *profilev1.Profile) int {
+	var total int
+	for _, s := range p.Sample {
+		total += int(s.Value[0])
+	}
+	return total
 }
 
 func Test_DropGoTypeParameters(t *testing.T) {

pkg/pprof/fix_go_truncated_test.go

@@ -61,29 +61,50 @@ func Test_RepairGoTruncatedStacktraces(t *testing.T) {
 	}
 }
 
+var goTruncatedStacktracesFixtures = []string{
+	"testdata/gotruncatefix/heap_go_truncated_1.pb.gz", // Cortex.
+	"testdata/gotruncatefix/heap_go_truncated_2.pb.gz", // Cortex.
+	"testdata/gotruncatefix/heap_go_truncated_3.pb.gz", // Loki. Pathological.
+	"testdata/gotruncatefix/heap_go_truncated_4.pb.gz", // Pyroscope.
+	"testdata/gotruncatefix/cpu_go_truncated_1.pb.gz",  // Cloudwatch Exporter
+}
+
+func Test_RepairGoTruncatedStacktraces_Fixtures(t *testing.T) {
+	for _, path := range goTruncatedStacktracesFixtures {
+		p, err := OpenFile(path)
+		require.NoError(t, err, path)
+		total := samplesTotal(p.Profile)
+
+		p.Profile = FixGoProfile(p.Profile)
+		assert.Equal(t, total, samplesTotal(p.Profile))
+
+		p.Normalize()
+		assert.Equal(t, total, samplesTotal(p.Profile))
+
+		fixed, err := OpenFile(path + ".fixed")
+		require.NoError(t, err)
+		assert.Equal(t, total, samplesTotal(fixed.Profile))
+	}
+}
+
 func Test_UpdateFixtures_RepairGoTruncatedStacktraces(t *testing.T) {
 	if os.Getenv("UPDATE_FIXTURES") != "true" {
 		t.Skip()
 	}
-	t.Helper()
-	paths := []string{
-		"testdata/gotruncatefix/heap_go_truncated_1.pb.gz", // Cortex.
-		"testdata/gotruncatefix/heap_go_truncated_2.pb.gz", // Cortex.
-		"testdata/gotruncatefix/heap_go_truncated_3.pb.gz", // Loki. Pathological.
-		"testdata/gotruncatefix/heap_go_truncated_4.pb.gz", // Pyroscope.
-		"testdata/gotruncatefix/cpu_go_truncated_1.pb.gz",  // Cloudwatch Exporter
-	}
-	for _, path := range paths {
-		func() {
-			p, err := OpenFile(path)
-			require.NoError(t, err, path)
-			f, err := os.Create(path + ".fixed")
-			require.NoError(t, err, path)
-			defer f.Close()
-			p.Profile = FixGoProfile(p.Profile)
-			RepairGoTruncatedStacktraces(p.Profile)
-			_, err = p.WriteTo(f)
-			require.NoError(t, err, path)
-		}()
+	for _, path := range goTruncatedStacktracesFixtures {
+		p, err := OpenFile(path)
+		require.NoError(t, err, path)
+		total := samplesTotal(p.Profile)
+
+		p.Profile = FixGoProfile(p.Profile)
+		p.Normalize()
+		assert.Equal(t, total, samplesTotal(p.Profile))
+
+		path += ".fixed"
+		fixed, err := os.Create(path)
+		require.NoError(t, err, path)
+		_, err = p.WriteTo(fixed)
+		require.NoError(t, fixed.Close(), path)
+		require.NoError(t, err, path)
 	}
 }