1// Copyright 2018 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//      http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package profiler
16
17import (
18	"bytes"
19	"io"
20	"testing"
21
22	"cloud.google.com/go/profiler/testdata"
23	"github.com/google/pprof/profile"
24)
25
26func TestGoHeapProfile(t *testing.T) {
27	oldStartCPUProfile, oldStopCPUProfile, oldWriteHeapProfile, oldSleep := startCPUProfile, stopCPUProfile, writeHeapProfile, sleep
28	defer func() {
29		startCPUProfile, stopCPUProfile, writeHeapProfile, sleep = oldStartCPUProfile, oldStopCPUProfile, oldWriteHeapProfile, oldSleep
30	}()
31
32	tests := []struct {
33		name    string
34		profile *profile.Profile
35		wantErr bool
36	}{
37		{
38			name:    "valid heap profile",
39			profile: testdata.HeapProfileCollected1,
40		},
41		{
42			name:    "profile with too few sample types",
43			profile: testdata.AllocProfileUploaded,
44			wantErr: true,
45		},
46		{
47			name: "profile with incorrect sample types",
48			profile: &profile.Profile{
49				DurationNanos: 10e9,
50				SampleType: []*profile.ValueType{
51					{Type: "objects", Unit: "count"},
52					{Type: "alloc_space", Unit: "bytes"},
53					{Type: "inuse_objects", Unit: "count"},
54					{Type: "inuse_space", Unit: "bytes"},
55				},
56			},
57			wantErr: true,
58		},
59	}
60
61	for _, tc := range tests {
62		var profileBytes bytes.Buffer
63		tc.profile.Write(&profileBytes)
64		writeHeapProfile = func(w io.Writer) error {
65			w.Write(profileBytes.Bytes())
66			return nil
67		}
68		_, err := goHeapProfile()
69		if tc.wantErr {
70			if err == nil {
71				t.Errorf("%s: goHeapProfile() got no error, want error", tc.name)
72			}
73			continue
74		}
75		if err != nil {
76			t.Errorf("%s: goHeapProfile() got %q, want no error", tc.name, err)
77		}
78	}
79}
80