From a52d299a838514e2c80e26f456be56720bcefe47 Mon Sep 17 00:00:00 2001 From: Andrew Dye Date: Thu, 2 Mar 2023 22:30:10 -0800 Subject: [PATCH 1/6] Use TokenCache for ClientCredentialsTokenSourceProvider Signed-off-by: Andrew Dye --- clients/go/admin/mocks/TokenSource.go | 81 ++++++++++ clients/go/admin/token_source_provider.go | 78 ++++++--- clients/go/admin/token_source_test.go | 185 ++++++++++++++++++++++ 3 files changed, 319 insertions(+), 25 deletions(-) create mode 100644 clients/go/admin/mocks/TokenSource.go diff --git a/clients/go/admin/mocks/TokenSource.go b/clients/go/admin/mocks/TokenSource.go new file mode 100644 index 000000000..35338d89b --- /dev/null +++ b/clients/go/admin/mocks/TokenSource.go @@ -0,0 +1,81 @@ +// Code generated by mockery v2.18.0. DO NOT EDIT. + +package mocks + +import ( + mock "github.com/stretchr/testify/mock" + oauth2 "golang.org/x/oauth2" +) + +// TokenSource is an autogenerated mock type for the TokenSource type +type TokenSource struct { + mock.Mock +} + +type TokenSource_Expecter struct { + mock *mock.Mock +} + +func (_m *TokenSource) EXPECT() *TokenSource_Expecter { + return &TokenSource_Expecter{mock: &_m.Mock} +} + +// Token provides a mock function with given fields: +func (_m *TokenSource) Token() (*oauth2.Token, error) { + ret := _m.Called() + + var r0 *oauth2.Token + if rf, ok := ret.Get(0).(func() *oauth2.Token); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*oauth2.Token) + } + } + + var r1 error + if rf, ok := ret.Get(1).(func() error); ok { + r1 = rf() + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// TokenSource_Token_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Token' +type TokenSource_Token_Call struct { + *mock.Call +} + +// Token is a helper method to define mock.On call +func (_e *TokenSource_Expecter) Token() *TokenSource_Token_Call { + return &TokenSource_Token_Call{Call: _e.mock.On("Token")} +} + +func (_c *TokenSource_Token_Call) Run(run func()) *TokenSource_Token_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *TokenSource_Token_Call) Return(_a0 *oauth2.Token, _a1 error) *TokenSource_Token_Call { + _c.Call.Return(_a0, _a1) + return _c +} + +type mockConstructorTestingTNewTokenSource interface { + mock.TestingT + Cleanup(func()) +} + +// NewTokenSource creates a new instance of TokenSource. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +func NewTokenSource(t mockConstructorTestingTNewTokenSource) *TokenSource { + mock := &TokenSource{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/clients/go/admin/token_source_provider.go b/clients/go/admin/token_source_provider.go index c2a520d70..66f63e9b0 100644 --- a/clients/go/admin/token_source_provider.go +++ b/clients/go/admin/token_source_provider.go @@ -68,7 +68,7 @@ func NewTokenSourceProvider(ctx context.Context, cfg *Config, tokenCache cache.T } } - tokenProvider, err = NewClientCredentialsTokenSourceProvider(ctx, cfg, scopes, tokenURL, audienceValue) + tokenProvider, err = NewClientCredentialsTokenSourceProvider(ctx, cfg, scopes, tokenURL, tokenCache, audienceValue) if err != nil { return nil, err } @@ -163,10 +163,12 @@ func GetPKCEAuthTokenSource(ctx context.Context, pkceTokenOrchestrator pkce.Toke type ClientCredentialsTokenSourceProvider struct { ccConfig clientcredentials.Config - TokenRefreshWindow time.Duration + tokenRefreshWindow time.Duration + tokenCache cache.TokenCache } -func NewClientCredentialsTokenSourceProvider(ctx context.Context, cfg *Config, scopes []string, tokenURL string, audience string) (TokenSourceProvider, error) { +func NewClientCredentialsTokenSourceProvider(ctx context.Context, cfg *Config, scopes []string, tokenURL string, + tokenCache cache.TokenCache, audience string) (TokenSourceProvider, error) { var secret string if len(cfg.ClientSecretEnvVar) > 0 { secret = os.Getenv(cfg.ClientSecretEnvVar) @@ -191,56 +193,82 @@ func NewClientCredentialsTokenSourceProvider(ctx context.Context, cfg *Config, s Scopes: scopes, EndpointParams: endpointParams, }, - TokenRefreshWindow: cfg.TokenRefreshWindow.Duration}, nil + tokenRefreshWindow: cfg.TokenRefreshWindow.Duration, + tokenCache: tokenCache}, nil } func (p ClientCredentialsTokenSourceProvider) GetTokenSource(ctx context.Context) (oauth2.TokenSource, error) { - if p.TokenRefreshWindow > 0 { + if p.tokenRefreshWindow > 0 { source := p.ccConfig.TokenSource(ctx) return &customTokenSource{ + ctx: ctx, new: source, mu: sync.Mutex{}, - t: nil, - tokenRefreshWindow: p.TokenRefreshWindow, + tokenRefreshWindow: p.tokenRefreshWindow, + tokenCache: p.tokenCache, }, nil } return p.ccConfig.TokenSource(ctx), nil } type customTokenSource struct { + ctx context.Context new oauth2.TokenSource mu sync.Mutex // guards everything else - t *oauth2.Token refreshTime time.Time failedToRefresh bool tokenRefreshWindow time.Duration + tokenCache cache.TokenCache +} + +// fetchTokenFromCache returns the cached token if available, and a bool indicating if we should try to refresh it. +// This function is not thread safe and should be called with the lock held. +func (s *customTokenSource) fetchTokenFromCache() (*oauth2.Token, bool) { + token, err := s.tokenCache.GetToken() + if err != nil { + logger.Infof(s.ctx, "no token found in cache") + return nil, false + } + if !token.Valid() { + logger.Infof(s.ctx, "cached token invalid") + return nil, false + } + if time.Now().After(s.refreshTime) && !s.failedToRefresh { + logger.Infof(s.ctx, "cached token refresh window exceeded") + return token, true + } + logger.Infof(s.ctx, "using cached token") + return token, false } func (s *customTokenSource) Token() (*oauth2.Token, error) { s.mu.Lock() defer s.mu.Unlock() - if s.t.Valid() { - if time.Now().After(s.refreshTime) && !s.failedToRefresh { - t, err := s.new.Token() - if err != nil { - s.failedToRefresh = true // don't try to refresh again before expiry - return s.t, nil - } - s.t = t - s.refreshTime = s.t.Expiry.Add(-getRandomDuration(s.tokenRefreshWindow)) - s.failedToRefresh = false - return s.t, nil + + cachedToken, needsRefresh := s.fetchTokenFromCache() + if cachedToken != nil && !needsRefresh { + return cachedToken, nil + } + + token, err := s.new.Token() + if err != nil { + if needsRefresh { + logger.Warnf(s.ctx, "failed to refresh token, using last cached token until expired") + s.failedToRefresh = true + return cachedToken, nil + } else { + logger.Errorf(s.ctx, "failed to refresh token") + return nil, err } - return s.t, nil } - t, err := s.new.Token() + logger.Infof(s.ctx, "refreshed token") + err = s.tokenCache.SaveToken(token) if err != nil { - return nil, err + logger.Warnf(s.ctx, "failed to cache token, using anyway") } - s.t = t s.failedToRefresh = false - s.refreshTime = s.t.Expiry.Add(-getRandomDuration(s.tokenRefreshWindow)) - return t, nil + s.refreshTime = token.Expiry.Add(-getRandomDuration(s.tokenRefreshWindow)) + return token, nil } // Get random duration between 0 and maxDuration diff --git a/clients/go/admin/token_source_test.go b/clients/go/admin/token_source_test.go index 9256e5e88..0f4c0d3dc 100644 --- a/clients/go/admin/token_source_test.go +++ b/clients/go/admin/token_source_test.go @@ -2,18 +2,24 @@ package admin import ( "context" + "fmt" "net/url" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" "golang.org/x/oauth2" tokenCacheMocks "github.com/flyteorg/flyteidl/clients/go/admin/cache/mocks" + "github.com/flyteorg/flyteidl/clients/go/admin/mocks" adminMocks "github.com/flyteorg/flyteidl/clients/go/admin/mocks" "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/service" + "github.com/flyteorg/flytestdlib/config" ) +//go:generate mockery --srcpkg golang.org/x/oauth2 --name TokenSource --with-expecter + type DummyTestTokenSource struct { oauth2.TokenSource } @@ -95,3 +101,182 @@ func TestNewTokenSourceProvider(t *testing.T) { assert.Equal(t, url.Values{audienceKey: {test.expectedAudience}}, clientCredSourceProvider.ccConfig.EndpointParams) } } + +func TestCustomTokenSource_fetchTokenFromCache(t *testing.T) { + ctx := context.Background() + cfg := GetConfig(ctx) + cfg.TokenRefreshWindow = config.Duration{Duration: time.Minute} + cfg.ClientSecretLocation = "" + + minuteAgo := time.Now().Add(-time.Minute) + hourAhead := time.Now().Add(time.Hour) + invalidToken := oauth2.Token{AccessToken: "foo", Expiry: minuteAgo} + validToken := oauth2.Token{AccessToken: "foo", Expiry: hourAhead} + + tests := []struct { + name string + refreshTime time.Time + failedToRefresh bool + token *oauth2.Token + expectToken bool + expectNeedsRefresh bool + }{ + { + name: "no token", + refreshTime: hourAhead, + failedToRefresh: false, + token: nil, + expectToken: false, + expectNeedsRefresh: false, + }, + { + name: "invalid token", + refreshTime: hourAhead, + failedToRefresh: false, + token: &invalidToken, + expectToken: false, + expectNeedsRefresh: false, + }, + { + name: "refresh exceeded", + refreshTime: minuteAgo, + failedToRefresh: false, + token: &validToken, + expectToken: false, + expectNeedsRefresh: false, + }, + { + name: "refresh exceeded failed", + refreshTime: minuteAgo, + failedToRefresh: true, + token: &validToken, + expectToken: false, + expectNeedsRefresh: false, + }, + { + name: "valid token", + refreshTime: hourAhead, + failedToRefresh: false, + token: &validToken, + expectToken: false, + expectNeedsRefresh: false, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + tokenCache := &tokenCacheMocks.TokenCache{} + provider, err := NewClientCredentialsTokenSourceProvider(ctx, cfg, []string{}, "", tokenCache, "") + assert.NoError(t, err) + source, err := provider.GetTokenSource(ctx) + assert.NoError(t, err) + customSource, ok := source.(*customTokenSource) + assert.True(t, ok) + + customSource.refreshTime = test.refreshTime + customSource.failedToRefresh = test.failedToRefresh + var tokenErr error = nil + if test.token == nil { + tokenErr = fmt.Errorf("no token!") + } + tokenCache.OnGetToken().Return(nil, tokenErr).Once() + token, needsRefresh := customSource.fetchTokenFromCache() + if test.expectToken { + assert.NotNil(t, token) + } else { + assert.Nil(t, token) + } + assert.Equal(t, test.expectNeedsRefresh, needsRefresh) + }) + } +} + +func TestCustomTokenSource_Token(t *testing.T) { + ctx := context.Background() + cfg := GetConfig(ctx) + cfg.TokenRefreshWindow = config.Duration{Duration: time.Minute} + cfg.ClientSecretLocation = "" + + minuteAgo := time.Now().Add(-time.Minute) + hourAhead := time.Now().Add(time.Hour) + twoHourAhead := time.Now().Add(2 * time.Hour) + invalidToken := oauth2.Token{AccessToken: "foo", Expiry: minuteAgo} + validToken := oauth2.Token{AccessToken: "foo", Expiry: hourAhead} + newToken := oauth2.Token{AccessToken: "foo", Expiry: twoHourAhead} + + tests := []struct { + name string + refreshTime time.Time + failedToRefresh bool + token *oauth2.Token + newToken *oauth2.Token + expectedToken *oauth2.Token + }{ + { + name: "cached token", + refreshTime: hourAhead, + failedToRefresh: false, + token: &validToken, + newToken: nil, + expectedToken: &validToken, + }, + { + name: "failed refresh still valid", + refreshTime: minuteAgo, + failedToRefresh: false, + token: &validToken, + newToken: nil, + expectedToken: &validToken, + }, + { + name: "failed refresh invalid", + refreshTime: minuteAgo, + failedToRefresh: false, + token: &invalidToken, + newToken: nil, + expectedToken: nil, + }, + { + name: "refresh", + refreshTime: minuteAgo, + failedToRefresh: false, + token: &invalidToken, + newToken: &newToken, + expectedToken: &newToken, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + tokenCache := &tokenCacheMocks.TokenCache{} + provider, err := NewClientCredentialsTokenSourceProvider(ctx, cfg, []string{}, "", tokenCache, "") + assert.NoError(t, err) + source, err := provider.GetTokenSource(ctx) + assert.NoError(t, err) + customSource, ok := source.(*customTokenSource) + assert.True(t, ok) + + mockSource := &mocks.TokenSource{} + if test.newToken != nil { + mockSource.On("Token").Return(test.newToken, nil) + } else { + mockSource.On("Token").Return(nil, fmt.Errorf("refresh token failed")) + } + customSource.new = mockSource + customSource.refreshTime = test.refreshTime + customSource.failedToRefresh = test.failedToRefresh + tokenCache.OnGetToken().Return(test.token, nil).Once() + if test.newToken != nil { + tokenCache.OnSaveToken(test.newToken).Return(nil).Once() + } + token, err := source.Token() + if test.expectedToken != nil { + assert.Equal(t, test.expectedToken, token) + assert.NoError(t, err) + } else { + assert.Nil(t, token) + assert.Error(t, err) + } + }) + } +} From 63a4243a40ee752e4418a4a59f9f42a8bf5cc5e0 Mon Sep 17 00:00:00 2001 From: Andrew Dye Date: Thu, 2 Mar 2023 22:41:31 -0800 Subject: [PATCH 2/6] Remove expecter mocks Signed-off-by: Andrew Dye --- clients/go/admin/mocks/TokenSource.go | 30 --------------------------- clients/go/admin/token_source_test.go | 2 +- 2 files changed, 1 insertion(+), 31 deletions(-) diff --git a/clients/go/admin/mocks/TokenSource.go b/clients/go/admin/mocks/TokenSource.go index 35338d89b..d92c16a46 100644 --- a/clients/go/admin/mocks/TokenSource.go +++ b/clients/go/admin/mocks/TokenSource.go @@ -12,14 +12,6 @@ type TokenSource struct { mock.Mock } -type TokenSource_Expecter struct { - mock *mock.Mock -} - -func (_m *TokenSource) EXPECT() *TokenSource_Expecter { - return &TokenSource_Expecter{mock: &_m.Mock} -} - // Token provides a mock function with given fields: func (_m *TokenSource) Token() (*oauth2.Token, error) { ret := _m.Called() @@ -43,28 +35,6 @@ func (_m *TokenSource) Token() (*oauth2.Token, error) { return r0, r1 } -// TokenSource_Token_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Token' -type TokenSource_Token_Call struct { - *mock.Call -} - -// Token is a helper method to define mock.On call -func (_e *TokenSource_Expecter) Token() *TokenSource_Token_Call { - return &TokenSource_Token_Call{Call: _e.mock.On("Token")} -} - -func (_c *TokenSource_Token_Call) Run(run func()) *TokenSource_Token_Call { - _c.Call.Run(func(args mock.Arguments) { - run() - }) - return _c -} - -func (_c *TokenSource_Token_Call) Return(_a0 *oauth2.Token, _a1 error) *TokenSource_Token_Call { - _c.Call.Return(_a0, _a1) - return _c -} - type mockConstructorTestingTNewTokenSource interface { mock.TestingT Cleanup(func()) diff --git a/clients/go/admin/token_source_test.go b/clients/go/admin/token_source_test.go index 0f4c0d3dc..711cf4538 100644 --- a/clients/go/admin/token_source_test.go +++ b/clients/go/admin/token_source_test.go @@ -18,7 +18,7 @@ import ( "github.com/flyteorg/flytestdlib/config" ) -//go:generate mockery --srcpkg golang.org/x/oauth2 --name TokenSource --with-expecter +//go:generate mockery --srcpkg golang.org/x/oauth2 --name TokenSource type DummyTestTokenSource struct { oauth2.TokenSource From 7c687bb878c9130b9237a473f2efef28914e787f Mon Sep 17 00:00:00 2001 From: Andrew Dye Date: Thu, 2 Mar 2023 22:43:01 -0800 Subject: [PATCH 3/6] Remove duplicate import Signed-off-by: Andrew Dye --- clients/go/admin/token_source_test.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/clients/go/admin/token_source_test.go b/clients/go/admin/token_source_test.go index 711cf4538..f7b7c9381 100644 --- a/clients/go/admin/token_source_test.go +++ b/clients/go/admin/token_source_test.go @@ -12,7 +12,6 @@ import ( "golang.org/x/oauth2" tokenCacheMocks "github.com/flyteorg/flyteidl/clients/go/admin/cache/mocks" - "github.com/flyteorg/flyteidl/clients/go/admin/mocks" adminMocks "github.com/flyteorg/flyteidl/clients/go/admin/mocks" "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/service" "github.com/flyteorg/flytestdlib/config" @@ -256,7 +255,7 @@ func TestCustomTokenSource_Token(t *testing.T) { customSource, ok := source.(*customTokenSource) assert.True(t, ok) - mockSource := &mocks.TokenSource{} + mockSource := &adminMocks.TokenSource{} if test.newToken != nil { mockSource.On("Token").Return(test.newToken, nil) } else { From 740e3a5dcf8f8e55c876bd63b3987a5d6cda73c9 Mon Sep 17 00:00:00 2001 From: Andrew Dye Date: Thu, 2 Mar 2023 22:55:27 -0800 Subject: [PATCH 4/6] Create default in-memory token cache Signed-off-by: Andrew Dye --- clients/go/admin/token_source_provider.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/clients/go/admin/token_source_provider.go b/clients/go/admin/token_source_provider.go index 66f63e9b0..e46b72256 100644 --- a/clients/go/admin/token_source_provider.go +++ b/clients/go/admin/token_source_provider.go @@ -185,6 +185,9 @@ func NewClientCredentialsTokenSourceProvider(ctx context.Context, cfg *Config, s endpointParams = url.Values{audienceKey: {audience}} } secret = strings.TrimSpace(secret) + if tokenCache == nil { + tokenCache = &cache.TokenCacheInMemoryProvider{} + } return ClientCredentialsTokenSourceProvider{ ccConfig: clientcredentials.Config{ ClientID: cfg.ClientID, From aee4a77cd42c9dc7039be623a3714bfbf7174224 Mon Sep 17 00:00:00 2001 From: Andrew Dye Date: Fri, 3 Mar 2023 00:20:20 -0800 Subject: [PATCH 5/6] Lint Signed-off-by: Andrew Dye --- clients/go/admin/token_source_provider.go | 5 ++--- clients/go/admin/token_source_test.go | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/clients/go/admin/token_source_provider.go b/clients/go/admin/token_source_provider.go index e46b72256..b62f5245c 100644 --- a/clients/go/admin/token_source_provider.go +++ b/clients/go/admin/token_source_provider.go @@ -259,10 +259,9 @@ func (s *customTokenSource) Token() (*oauth2.Token, error) { logger.Warnf(s.ctx, "failed to refresh token, using last cached token until expired") s.failedToRefresh = true return cachedToken, nil - } else { - logger.Errorf(s.ctx, "failed to refresh token") - return nil, err } + logger.Errorf(s.ctx, "failed to refresh token") + return nil, err } logger.Infof(s.ctx, "refreshed token") err = s.tokenCache.SaveToken(token) diff --git a/clients/go/admin/token_source_test.go b/clients/go/admin/token_source_test.go index f7b7c9381..ef0588e6f 100644 --- a/clients/go/admin/token_source_test.go +++ b/clients/go/admin/token_source_test.go @@ -176,7 +176,7 @@ func TestCustomTokenSource_fetchTokenFromCache(t *testing.T) { customSource.failedToRefresh = test.failedToRefresh var tokenErr error = nil if test.token == nil { - tokenErr = fmt.Errorf("no token!") + tokenErr = fmt.Errorf("no token") } tokenCache.OnGetToken().Return(nil, tokenErr).Once() token, needsRefresh := customSource.fetchTokenFromCache() From 10a894ccd108dcf9fa0c571af3e5929d28800718 Mon Sep 17 00:00:00 2001 From: Andrew Dye Date: Fri, 3 Mar 2023 00:45:48 -0800 Subject: [PATCH 6/6] Fix mocks Signed-off-by: Andrew Dye --- clients/go/admin/mocks/TokenSource.go | 35 ++++++++++++----------- clients/go/admin/token_source_provider.go | 5 ++++ clients/go/admin/token_source_test.go | 6 ++-- 3 files changed, 26 insertions(+), 20 deletions(-) diff --git a/clients/go/admin/mocks/TokenSource.go b/clients/go/admin/mocks/TokenSource.go index d92c16a46..60cc87236 100644 --- a/clients/go/admin/mocks/TokenSource.go +++ b/clients/go/admin/mocks/TokenSource.go @@ -1,4 +1,4 @@ -// Code generated by mockery v2.18.0. DO NOT EDIT. +// Code generated by mockery v1.0.1. DO NOT EDIT. package mocks @@ -12,6 +12,24 @@ type TokenSource struct { mock.Mock } +type TokenSource_Token struct { + *mock.Call +} + +func (_m TokenSource_Token) Return(_a0 *oauth2.Token, _a1 error) *TokenSource_Token { + return &TokenSource_Token{Call: _m.Call.Return(_a0, _a1)} +} + +func (_m *TokenSource) OnToken() *TokenSource_Token { + c_call := _m.On("Token") + return &TokenSource_Token{Call: c_call} +} + +func (_m *TokenSource) OnTokenMatch(matchers ...interface{}) *TokenSource_Token { + c_call := _m.On("Token", matchers...) + return &TokenSource_Token{Call: c_call} +} + // Token provides a mock function with given fields: func (_m *TokenSource) Token() (*oauth2.Token, error) { ret := _m.Called() @@ -34,18 +52,3 @@ func (_m *TokenSource) Token() (*oauth2.Token, error) { return r0, r1 } - -type mockConstructorTestingTNewTokenSource interface { - mock.TestingT - Cleanup(func()) -} - -// NewTokenSource creates a new instance of TokenSource. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. -func NewTokenSource(t mockConstructorTestingTNewTokenSource) *TokenSource { - mock := &TokenSource{} - mock.Mock.Test(t) - - t.Cleanup(func() { mock.AssertExpectations(t) }) - - return mock -} diff --git a/clients/go/admin/token_source_provider.go b/clients/go/admin/token_source_provider.go index b62f5245c..937b6d901 100644 --- a/clients/go/admin/token_source_provider.go +++ b/clients/go/admin/token_source_provider.go @@ -23,6 +23,11 @@ import ( "github.com/flyteorg/flytestdlib/logger" ) +//go:generate mockery -name TokenSource +type TokenSource interface { + Token() (*oauth2.Token, error) +} + const ( audienceKey = "audience" ) diff --git a/clients/go/admin/token_source_test.go b/clients/go/admin/token_source_test.go index ef0588e6f..745237759 100644 --- a/clients/go/admin/token_source_test.go +++ b/clients/go/admin/token_source_test.go @@ -17,8 +17,6 @@ import ( "github.com/flyteorg/flytestdlib/config" ) -//go:generate mockery --srcpkg golang.org/x/oauth2 --name TokenSource - type DummyTestTokenSource struct { oauth2.TokenSource } @@ -257,9 +255,9 @@ func TestCustomTokenSource_Token(t *testing.T) { mockSource := &adminMocks.TokenSource{} if test.newToken != nil { - mockSource.On("Token").Return(test.newToken, nil) + mockSource.OnToken().Return(test.newToken, nil) } else { - mockSource.On("Token").Return(nil, fmt.Errorf("refresh token failed")) + mockSource.OnToken().Return(nil, fmt.Errorf("refresh token failed")) } customSource.new = mockSource customSource.refreshTime = test.refreshTime