gitlab.com/gitlab-org/labkit@v1.21.0/monitoring/profiler_test.go

// +build continuous_profiler_stackdriver

package monitoring

import (
	"bytes"
	"fmt"
	"os"
	"testing"

	"cloud.google.com/go/profiler"
	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/labkit/log"
	"google.golang.org/api/option"
)

func TestInitProfiler(t *testing.T) {
	tests := []struct {
		name            string
		envParams       string
		writeEnvVar     bool
		opts            profilerOpts
		expectedLog     string
		profErr         error
		expectedConfig  profiler.Config
		expectedOptions []option.ClientOption
	}{
		{
			name:           "complete params",
			opts:           profilerOpts{},
			envParams:      "stackdriver?service=gitaly&service_version=1.0.1&project_id=test-123",
			writeEnvVar:    true,
			expectedLog:    `^time=.*level=info msg="profiler enabled" driver=stackdriver projectId=test-123 service=gitaly serviceVersion=1.0.1`,
			expectedConfig: profiler.Config{Service: "gitaly", ServiceVersion: "1.0.1", ProjectID: "test-123", MutexProfiling: true},
		},
		{
			name:           "complete env params and service version given",
			opts:           profilerOpts{ServiceVersion: "9.0.0"},
			envParams:      "stackdriver?service=gitaly&service_version=1.0.1&project_id=test-123",
			writeEnvVar:    true,
			expectedLog:    `^time=.*level=info msg="profiler enabled" driver=stackdriver projectId=test-123 service=gitaly serviceVersion=9.0.0`,
			expectedConfig: profiler.Config{Service: "gitaly", ServiceVersion: "9.0.0", ProjectID: "test-123", MutexProfiling: true},
		},
		{
			name:           "incomplete env params and service version given through args",
			opts:           profilerOpts{ServiceVersion: "9.0.0"},
			envParams:      "stackdriver?service=gitaly",
			expectedLog:    `^time=.*level=info msg="profiler enabled" driver=stackdriver projectId= service=gitaly serviceVersion=9.0.0`,
			expectedConfig: profiler.Config{MutexProfiling: true},
		},
		{
			name:           "debug logging",
			opts:           profilerOpts{ServiceVersion: "9.0.0"},
			envParams:      "stackdriver?service=gitaly&debug_logging=true",
			writeEnvVar:    true,
			expectedLog:    `^time=.*level=info msg="profiler enabled" driver=stackdriver projectId= service=gitaly serviceVersion=9.0.0`,
			expectedConfig: profiler.Config{Service: "gitaly", ServiceVersion: "9.0.0", MutexProfiling: true, DebugLogging: true},
		},
		{
			name:        "unknown driver",
			envParams:   "true?service=gitaly",
			writeEnvVar: true,
			opts:        profilerOpts{},
			expectedLog: `^time=.*level=warning msg="unknown driver" driver=true`,
		},
		{
			name:        "wrong env content format",
			envParams:   ":foo?service=gitaly",
			writeEnvVar: true,
			opts:        profilerOpts{},
			expectedLog: `^time=.*level=warning msg="unable to parse env var content"`,
		},
		{
			name:        "empty logs when no env var is set",
			writeEnvVar: false,
			opts:        profilerOpts{},
		},
		{
			name:           "fails to start profiler",
			envParams:      "stackdriver?service=gitaly&service_version=1.0.1&project_id=test-123",
			writeEnvVar:    true,
			opts:           profilerOpts{},
			expectedLog:    `^time=.*level=warning msg="unable to initialize stackdriver profiler" error="fail to start"`,
			profErr:        fmt.Errorf("fail to start"),
			expectedConfig: profiler.Config{Service: "gitaly", ServiceVersion: "1.0.1", ProjectID: "test-123", MutexProfiling: true},
		},
		{
			name:            "with credentials file",
			envParams:       "stackdriver",
			writeEnvVar:     true,
			opts:            profilerOpts{CredentialsFile: "/path/to/credentials.json"},
			expectedOptions: []option.ClientOption{option.WithCredentialsFile("/path/to/credentials.json")},
			expectedConfig:  profiler.Config{MutexProfiling: true},
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			clearGcpEnvVars(t)

			preStubStart := profStart
			defer func() { profStart = preStubStart }()

			var actualConfig profiler.Config
			var actualOptions []option.ClientOption
			profStart = func(cfg profiler.Config, options ...option.ClientOption) error {
				actualConfig = cfg
				actualOptions = options
				return test.profErr
			}

			if test.writeEnvVar {
				err := os.Setenv(profilerEnvKey, test.envParams)
				require.NoError(t, err)
			}

			logOutput := captureLogOutput(t, initProfiler, test.opts)

			if test.writeEnvVar {
				require.Regexp(t, test.expectedLog, logOutput)
				require.Equal(t, test.expectedConfig, actualConfig)
				require.Equal(t, test.expectedOptions, actualOptions)
			} else {
				require.Empty(t, logOutput)
			}
		})
	}
}

func TestInitProfilerWithNoProfilerEnvKeySet(t *testing.T) {
	t.Run("no GITLAB_CONTINUOUS_PROFILING env var set", func(t *testing.T) {
		clearGcpEnvVars(t)

		logOutput := captureLogOutput(t, initProfiler, profilerOpts{})

		require.Empty(t, logOutput)
	})
}

func captureLogOutput(t *testing.T, init func(opts profilerOpts), opts profilerOpts) string {
	t.Helper()

	buf := &bytes.Buffer{}
	closer, err := log.Initialize(log.WithWriter(buf))
	require.NoError(t, err)
	defer closer.Close()

	init(opts)

	return buf.String()
}

func clearGcpEnvVars(t *testing.T) {
	t.Helper()

	for _, env := range []string{profilerEnvKey, "GAE_SERVICE", "GAE_VERSION", "GOOGLE_CLOUD_PROJECT"} {
		err := os.Unsetenv(env)

		require.NoError(t, err)
	}
}