From 10a894ccd108dcf9fa0c571af3e5929d28800718 Mon Sep 17 00:00:00 2001
From: Andrew Dye <andrewwdye@gmail.com>
Date: Fri, 3 Mar 2023 00:45:48 -0800
Subject: [PATCH] Fix mocks

Signed-off-by: Andrew Dye <andrewwdye@gmail.com>
---
 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