1 files changed, 21 insertions, 11 deletions
diff --git a/src/runtime/pprof/protomem_test.go b/src/runtime/pprof/protomem_test.go
index 1e30ed93a3..315d5f0b4d 100644
--- a/src/runtime/pprof/protomem_test.go
+++ b/src/runtime/pprof/protomem_test.go
@@ -14,7 +14,6 @@ import (
 func TestConvertMemProfile(t *testing.T) {
 	addr1, addr2, map1, map2 := testPCs(t)
 
-	var buf bytes.Buffer
 	// MemProfileRecord stacks are return PCs, so add one to the
 	// addresses recorded in the "profile". The proto profile
 	// locations are call PCs, so conversion will subtract one
@@ -27,15 +26,6 @@ func TestConvertMemProfile(t *testing.T) {
 		{AllocBytes: 512 * 1024, FreeBytes: 512 * 1024, AllocObjects: 1, FreeObjects: 1, Stack0: [32]uintptr{a1 + 1, a1 + 2, a2 + 3}},
 	}
 
-	if err := writeHeapProto(&buf, rec, rate); err != nil {
-		t.Fatalf("writing profile: %v", err)
-	}
-
-	p, err := profile.Parse(&buf)
-	if err != nil {
-		t.Fatalf("profile.Parse: %v", err)
-	}
-
 	periodType := &profile.ValueType{Type: "space", Unit: "bytes"}
 	sampleType := []*profile.ValueType{
 		{Type: "alloc_objects", Unit: "count"},
@@ -70,5 +60,25 @@ func TestConvertMemProfile(t *testing.T) {
 			NumLabel: map[string][]int64{"bytes": {829411}},
 		},
 	}
-	checkProfile(t, p, rate, periodType, sampleType, samples)
+	for _, tc := range []struct {
+		name              string
+		defaultSampleType string
+	}{
+		{"heap", ""},
+		{"allocs", "alloc_space"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			var buf bytes.Buffer
+			if err := writeHeapProto(&buf, rec, rate, tc.defaultSampleType); err != nil {
+				t.Fatalf("writing profile: %v", err)
+			}
+
+			p, err := profile.Parse(&buf)
+			if err != nil {
+				t.Fatalf("profile.Parse: %v", err)
+			}
+
+			checkProfile(t, p, rate, periodType, sampleType, samples, tc.defaultSampleType)
+		})
+	}
 }