Fix mocks

Signed-off-by: Andrew Dye <[email protected]>
flyteorg · Mar 3, 2023 · 10a894c · 10a894c
1 parent aee4a77
commit 10a894c
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 20 deletions.
diff --git a/clients/go/admin/mocks/TokenSource.go b/clients/go/admin/mocks/TokenSource.go
diff --git 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
@@ -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