Skip to content

Commit

Permalink
test: add test for uploadProfile
Browse files Browse the repository at this point in the history
  • Loading branch information
marcsanmi committed Sep 9, 2024
1 parent 2ead492 commit dfb81e9
Show file tree
Hide file tree
Showing 4 changed files with 149 additions and 27 deletions.
29 changes: 7 additions & 22 deletions collector_test.go
Original file line number Diff line number Diff line change
@@ -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",
Expand Down Expand Up @@ -45,22 +45,22 @@ 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")
}
}

func Test_StartCPUProfile_blocks_Stop(t *testing.T) {
logger := new(testLogger)
logger := testutil.NewTestLogger()
collector := new(mockCollector)
c := newCPUProfileCollector(
"test",
Expand All @@ -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")
Expand Down Expand Up @@ -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()
}
32 changes: 32 additions & 0 deletions internal/testutil/logging.go
Original file line number Diff line number Diff line change
@@ -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...)
}
13 changes: 8 additions & 5 deletions upstream/remote/remote.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
102 changes: 102 additions & 0 deletions upstream/remote/remote_test.go
Original file line number Diff line number Diff line change
@@ -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)
}

0 comments on commit dfb81e9

Please sign in to comment.