// +build continuous_profiler_stackdriver

package monitoring

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

	"gitlab.com/gitlab-org/labkit/log"

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

func TestInitProfiler(t *testing.T) {
	tests := []struct {
		name        string
		envParams   string
		writeEnvVar bool
		opts        profilerOpts
		expectedLog string
		profErr     error
	}{
		{
			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`,
		},
		{
			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`,
		},
		{
			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`,
		},
		{
			name:        "unknown driver",
			envParams:   "true?service=gitaly",
			writeEnvVar: true,
			opts:        profilerOpts{},
			expectedLog: `^time=.*level=warning msg="unkown 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"),
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			clearGcpEnvVars(t)

			preStubStart := profStart
			defer func() { profStart = preStubStart }()
			stubStackdriverProfiler(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)
			} 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 stubStackdriverProfiler(err error) {
	profStart = func(cfg profiler.Config, options ...option.ClientOption) error {
		return err
	}
}

func captureLogOutput(t *testing.T, init func(opts profilerOpts), opts profilerOpts) string {
	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) {
	for _, env := range []string{profilerEnvKey, "GAE_SERVICE", "GAE_VERSION", "GOOGLE_CLOUD_PROJECT"} {
		err := os.Unsetenv(env)

		require.NoError(t, err)
	}
}
