From dfb81e9b3222c3d4b4fbd5f245d9d82ed055c495 Mon Sep 17 00:00:00 2001 From: Marc Sanmiquel Date: Mon, 9 Sep 2024 13:10:16 +0200 Subject: [PATCH] test: add test for uploadProfile --- collector_test.go | 29 +++------- internal/testutil/logging.go | 32 +++++++++++ upstream/remote/remote.go | 13 +++-- upstream/remote/remote_test.go | 102 +++++++++++++++++++++++++++++++++ 4 files changed, 149 insertions(+), 27 deletions(-) create mode 100644 internal/testutil/logging.go create mode 100644 upstream/remote/remote_test.go diff --git a/collector_test.go b/collector_test.go index 5b13d418..aa2c8017 100644 --- a/collector_test.go +++ b/collector_test.go @@ -1,18 +1,18 @@ package pyroscope import ( - "fmt" "io" "reflect" "sync" "testing" "time" + "github.com/grafana/pyroscope-go/internal/testutil" "github.com/grafana/pyroscope-go/upstream" ) func Test_StartCPUProfile_interrupts_background_profiling(t *testing.T) { - logger := new(testLogger) + logger := testutil.NewTestLogger() collector := new(mockCollector) c := newCPUProfileCollector( "test", @@ -45,14 +45,14 @@ func Test_StartCPUProfile_interrupts_background_profiling(t *testing.T) { c.Stop() - if !reflect.DeepEqual(logger.lines, []string{ + if !reflect.DeepEqual(logger.Lines(), []string{ "starting cpu profile collector", "cpu profile collector interrupted with StartCPUProfile", "cpu profile collector restored", "stopping cpu profile collector", "stopping cpu profile collector stopped", }) { - for _, line := range logger.lines { + for _, line := range logger.Lines() { t.Log(line) } t.Fatal("^ unexpected even sequence") @@ -60,7 +60,7 @@ func Test_StartCPUProfile_interrupts_background_profiling(t *testing.T) { } func Test_StartCPUProfile_blocks_Stop(t *testing.T) { - logger := new(testLogger) + logger := testutil.NewTestLogger() collector := new(mockCollector) c := newCPUProfileCollector( "test", @@ -86,14 +86,14 @@ func Test_StartCPUProfile_blocks_Stop(t *testing.T) { go c.StopCPUProfile() c.Stop() - if !reflect.DeepEqual(logger.lines, []string{ + if !reflect.DeepEqual(logger.Lines(), []string{ "starting cpu profile collector", "cpu profile collector interrupted with StartCPUProfile", "stopping cpu profile collector", "cpu profile collector restored", "stopping cpu profile collector stopped", }) { - for _, line := range logger.lines { + for _, line := range logger.Lines() { t.Log(line) } t.Fatal("^ unexpected even sequence") @@ -146,18 +146,3 @@ type mockUpstream struct{ uploaded []*upstream.UploadJob } func (m *mockUpstream) Upload(j *upstream.UploadJob) { m.uploaded = append(m.uploaded, j) } func (*mockUpstream) Flush() {} - -type testLogger struct { - sync.Mutex - lines []string -} - -func (t *testLogger) Debugf(format string, args ...interface{}) { t.put(format, args...) } -func (t *testLogger) Infof(format string, args ...interface{}) { t.put(format, args...) } -func (t *testLogger) Errorf(format string, args ...interface{}) { t.put(format, args...) } - -func (t *testLogger) put(format string, args ...interface{}) { - t.Lock() - t.lines = append(t.lines, fmt.Sprintf(format, args...)) - t.Unlock() -} diff --git a/internal/testutil/logging.go b/internal/testutil/logging.go new file mode 100644 index 00000000..911707c6 --- /dev/null +++ b/internal/testutil/logging.go @@ -0,0 +1,32 @@ +package testutil + +import ( + "fmt" + "sync" +) + +type TestLogger struct { + sync.Mutex + lines []string +} + +func NewTestLogger() *TestLogger { + return &TestLogger{lines: make([]string, 0)} +} + +func (t *TestLogger) Debugf(format string, args ...interface{}) { t.put(format, args...) } +func (t *TestLogger) Infof(format string, args ...interface{}) { t.put(format, args...) } +func (t *TestLogger) Errorf(format string, args ...interface{}) { t.put(format, args...) } + +func (t *TestLogger) put(format string, args ...interface{}) { + t.Lock() + t.lines = append(t.lines, fmt.Sprintf(format, args...)) + t.Unlock() +} + +// Lines returns a copy of the logged lines +func (t *TestLogger) Lines() []string { + t.Lock() + defer t.Unlock() + return append([]string(nil), t.lines...) +} diff --git a/upstream/remote/remote.go b/upstream/remote/remote.go index 33659847..dd126315 100644 --- a/upstream/remote/remote.go +++ b/upstream/remote/remote.go @@ -21,7 +21,13 @@ import ( var errCloudTokenRequired = errors.New("please provide an authentication token. You can find it here: https://pyroscope.io/cloud") -const cloudHostnameSuffix = "pyroscope.cloud" +const ( + authTokenDeprecationWarning = "Authtoken is specified, but deprecated and ignored. " + + "Please switch to BasicAuthUser and BasicAuthPassword. " + + "If you need to use Bearer token authentication for a custom setup, " + + "you can use the HTTPHeaders option to set the Authorization header manually." + cloudHostnameSuffix = "pyroscope.cloud" +) type Remote struct { cfg Config @@ -195,10 +201,7 @@ func (r *Remote) uploadProfile(j *upstream.UploadJob) error { } else if r.cfg.BasicAuthUser != "" && r.cfg.BasicAuthPassword != "" { request.SetBasicAuth(r.cfg.BasicAuthUser, r.cfg.BasicAuthPassword) } else if r.cfg.AuthToken != "" { - r.logger.Infof("Authtoken is specified, but deprecated and ignored." + - "Please switch to BasicAuthUser and BasicAuthPassword. " + - "If you need to use Bearer token authentication for a custom setup, " + - "you can use the HTTPHeaders option to set the Authorization header manually.") + r.logger.Infof(authTokenDeprecationWarning) } if r.cfg.TenantID != "" { request.Header.Set("X-Scope-OrgID", r.cfg.TenantID) diff --git a/upstream/remote/remote_test.go b/upstream/remote/remote_test.go new file mode 100644 index 00000000..f1e74fce --- /dev/null +++ b/upstream/remote/remote_test.go @@ -0,0 +1,102 @@ +package remote + +import ( + "bytes" + "io" + "net/http" + "testing" + "time" + + "github.com/grafana/pyroscope-go/internal/testutil" + "github.com/grafana/pyroscope-go/upstream" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" +) + +func TestUploadProfile(t *testing.T) { + tests := []struct { + name string + cfg Config + serverAddress string + expectedAuthHeader string + expectWarning bool + }{ + { + name: "OG Pyroscope Cloud with AuthToken", + cfg: Config{ + AuthToken: "test-token", + Address: "https://example.pyroscope.cloud", + }, + expectedAuthHeader: "Bearer test-token", + expectWarning: false, + }, + { + name: "Non-OG Server with BasicAuth", + cfg: Config{ + BasicAuthUser: "user", + BasicAuthPassword: "pass", + Address: "https://example.com", + }, + expectedAuthHeader: "Basic dXNlcjpwYXNz", // Base64 encoded "user:pass" + expectWarning: false, + }, + { + name: "Non-OG Server with AuthToken (Deprecated)", + cfg: Config{ + AuthToken: "deprecated-token", + Address: "https://example.com", + }, + expectedAuthHeader: "", + expectWarning: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + logger := testutil.NewTestLogger() + mockClient := new(MockHTTPClient) + + mockClient.On("Do", mock.Anything).Return(&http.Response{ + StatusCode: 200, + Body: io.NopCloser(bytes.NewBufferString("OK")), + }, nil) + + r := &Remote{ + cfg: tt.cfg, + client: mockClient, + logger: logger, + } + + err := r.uploadProfile(&upstream.UploadJob{ + Name: "test-profile", + StartTime: time.Now(), + EndTime: time.Now().Add(time.Minute), + SpyName: "test-spy", + SampleRate: 100, + Units: "samples", + }) + assert.NoError(t, err) + + if tt.expectWarning { + assert.Contains(t, logger.Lines(), authTokenDeprecationWarning) + } else { + assert.NotContains(t, logger.Lines(), authTokenDeprecationWarning) + } + + mockClient.AssertCalled(t, "Do", mock.MatchedBy(func(req *http.Request) bool { + return req.Header.Get("Authorization") == tt.expectedAuthHeader + })) + + mockClient.AssertExpectations(t) + }) + } +} + +type MockHTTPClient struct { + mock.Mock +} + +func (m *MockHTTPClient) Do(req *http.Request) (*http.Response, error) { + args := m.Called(req) + return args.Get(0).(*http.Response), args.Error(1) +}