From d02f71b9f81a041575ad7a9a08058b8756d705f8 Mon Sep 17 00:00:00 2001
From: Iaroslav Ciupin <iaroslav@union.ai>
Date: Thu, 7 Sep 2023 09:10:52 +0300
Subject: [PATCH 1/4] logout hook

Signed-off-by: Iaroslav Ciupin <iaroslav@union.ai>
---
 auth/handlers.go         | 17 ++++++++++++++---
 plugins/registry.go      |  1 +
 plugins/registry_test.go | 20 ++++++++++++++++++++
 3 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/auth/handlers.go b/auth/handlers.go
index 0ee2cc776..f49b8d567 100644
--- a/auth/handlers.go
+++ b/auth/handlers.go
@@ -48,6 +48,7 @@ func (e *PreRedirectHookError) Error() string {
 // PreRedirectHookError is the error interface which allows the user to set correct http status code and Message to be set in case the function returns an error
 // without which the current usage in GetCallbackHandler will set this to InternalServerError
 type PreRedirectHookFunc func(ctx context.Context, authCtx interfaces.AuthenticationContext, request *http.Request, w http.ResponseWriter) *PreRedirectHookError
+type LogoutHookFunc func(ctx context.Context, authCtx interfaces.AuthenticationContext, request *http.Request, w http.ResponseWriter) error
 type HTTPRequestToMetadataAnnotator func(ctx context.Context, request *http.Request) metadata.MD
 type UserInfoForwardResponseHandler func(ctx context.Context, w http.ResponseWriter, m protoiface.MessageV1) error
 
@@ -68,7 +69,7 @@ func RegisterHandlers(ctx context.Context, handler interfaces.HandlerRegisterer,
 	handler.HandleFunc(fmt.Sprintf("/%s", OIdCMetadataEndpoint), GetOIdCMetadataEndpointRedirectHandler(ctx, authCtx))
 
 	// These endpoints require authentication
-	handler.HandleFunc("/logout", GetLogoutEndpointHandler(ctx, authCtx))
+	handler.HandleFunc("/logout", GetLogoutEndpointHandler(ctx, authCtx, pluginRegistry))
 }
 
 // Look for access token and refresh token, if both are present and the access token is expired, then attempt to
@@ -466,9 +467,19 @@ func GetOIdCMetadataEndpointRedirectHandler(ctx context.Context, authCtx interfa
 	}
 }
 
-func GetLogoutEndpointHandler(ctx context.Context, authCtx interfaces.AuthenticationContext) http.HandlerFunc {
+func GetLogoutEndpointHandler(ctx context.Context, authCtx interfaces.AuthenticationContext, pluginRegistry *plugins.Registry) http.HandlerFunc {
 	return func(writer http.ResponseWriter, request *http.Request) {
-		logger.Debugf(ctx, "Deleting auth cookies")
+		hook := plugins.Get[LogoutHookFunc](pluginRegistry, plugins.PluginIDLogoutHook)
+		if hook != nil {
+			if err := hook(ctx, authCtx, request, writer); err != nil {
+				logger.Errorf(ctx, "logout hook failed: %v", err)
+				writer.WriteHeader(http.StatusInternalServerError)
+				return
+			}
+			logger.Debugf(ctx, "logout hook called")
+		}
+
+		logger.Debugf(ctx, "deleting auth cookies")
 		authCtx.CookieManager().DeleteCookies(ctx, writer)
 
 		// Redirect if one was given
diff --git a/plugins/registry.go b/plugins/registry.go
index 14682f7e8..90389008b 100644
--- a/plugins/registry.go
+++ b/plugins/registry.go
@@ -13,6 +13,7 @@ const (
 	PluginIDDataProxy              PluginID = "DataProxy"
 	PluginIDUnaryServiceMiddleware PluginID = "UnaryServiceMiddleware"
 	PluginIDPreRedirectHook        PluginID = "PreRedirectHook"
+	PluginIDLogoutHook             PluginID = "LogoutHook"
 )
 
 type AtomicRegistry struct {
diff --git a/plugins/registry_test.go b/plugins/registry_test.go
index 0737c1281..15b0cb93b 100644
--- a/plugins/registry_test.go
+++ b/plugins/registry_test.go
@@ -41,6 +41,26 @@ func TestRedirectHook(t *testing.T) {
 	assert.Equal(t, fmt.Errorf("redirect hook error"), err)
 }
 
+type LogoutHook func(context.Context) error
+
+func TestLogoutHook(t *testing.T) {
+	ar := NewAtomicRegistry(nil)
+	r := NewRegistry()
+
+	hook := LogoutHook(func(ctx context.Context) error {
+		return fmt.Errorf("redirect hook error")
+	})
+	err := r.Register(PluginIDLogoutHook, hook)
+	assert.NoError(t, err)
+
+	ar.Store(r)
+	r = ar.Load()
+	fn := Get[LogoutHook](r, PluginIDLogoutHook)
+	err = fn(context.Background())
+
+	assert.Equal(t, fmt.Errorf("redirect hook error"), err)
+}
+
 func TestRegistry_RegisterDefault(t *testing.T) {
 	r := NewRegistry()
 	r.RegisterDefault("hello", 5)

From 653a8e46bf38d9ab4777260d1d75e6045ba2d3a8 Mon Sep 17 00:00:00 2001
From: Iaroslav Ciupin <iaroslav@union.ai>
Date: Fri, 8 Sep 2023 15:27:46 +0300
Subject: [PATCH 2/4] export GetAuthFlowRedirectURI

Signed-off-by: Iaroslav Ciupin <iaroslav@union.ai>
---
 auth/cookie.go      | 3 ++-
 auth/cookie_test.go | 9 +++++----
 auth/handlers.go    | 4 ++--
 3 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/auth/cookie.go b/auth/cookie.go
index f93679470..113206a54 100644
--- a/auth/cookie.go
+++ b/auth/cookie.go
@@ -163,10 +163,11 @@ func NewRedirectCookie(ctx context.Context, redirectURL string) *http.Cookie {
 	}
 }
 
+// GetAuthFlowEndRedirect returns the redirect URI according to data in request.
 // At the end of the OAuth flow, the server needs to send the user somewhere. This should have been stored as a cookie
 // during the initial /login call. If that cookie is missing from the request, it will default to the one configured
 // in this package's Config object.
-func getAuthFlowEndRedirect(ctx context.Context, authCtx interfaces.AuthenticationContext, request *http.Request) string {
+func GetAuthFlowEndRedirect(ctx context.Context, authCtx interfaces.AuthenticationContext, request *http.Request) string {
 	queryParams := request.URL.Query()
 	// Use the redirect URL specified in the request if one is available.
 	if redirectURL := queryParams.Get(RedirectURLParameter); len(redirectURL) > 0 {
diff --git a/auth/cookie_test.go b/auth/cookie_test.go
index b525ba8b4..575dffe21 100644
--- a/auth/cookie_test.go
+++ b/auth/cookie_test.go
@@ -9,11 +9,12 @@ import (
 	"net/url"
 	"testing"
 
-	"github.com/flyteorg/flyteadmin/auth/config"
-	"github.com/flyteorg/flyteadmin/auth/interfaces/mocks"
 	stdConfig "github.com/flyteorg/flytestdlib/config"
 	"github.com/gorilla/securecookie"
 	"github.com/stretchr/testify/assert"
+
+	"github.com/flyteorg/flyteadmin/auth/config"
+	"github.com/flyteorg/flyteadmin/auth/interfaces/mocks"
 )
 
 func mustParseURL(t testing.TB, u string) url.URL {
@@ -131,7 +132,7 @@ func TestGetAuthFlowEndRedirect(t *testing.T) {
 		assert.NotNil(t, cookie)
 		request.AddCookie(cookie)
 		mockAuthCtx := &mocks.AuthenticationContext{}
-		redirect := getAuthFlowEndRedirect(ctx, mockAuthCtx, request)
+		redirect := GetAuthFlowEndRedirect(ctx, mockAuthCtx, request)
 		assert.Equal(t, "/console", redirect)
 	})
 
@@ -145,7 +146,7 @@ func TestGetAuthFlowEndRedirect(t *testing.T) {
 				RedirectURL: stdConfig.URL{URL: mustParseURL(t, "/api/v1/projects")},
 			},
 		})
-		redirect := getAuthFlowEndRedirect(ctx, mockAuthCtx, request)
+		redirect := GetAuthFlowEndRedirect(ctx, mockAuthCtx, request)
 		assert.Equal(t, "/api/v1/projects", redirect)
 	})
 }
diff --git a/auth/handlers.go b/auth/handlers.go
index f49b8d567..94a28e341 100644
--- a/auth/handlers.go
+++ b/auth/handlers.go
@@ -124,7 +124,7 @@ func RefreshTokensIfExists(ctx context.Context, authCtx interfaces.Authenticatio
 			return
 		}
 
-		redirectURL := getAuthFlowEndRedirect(ctx, authCtx, request)
+		redirectURL := GetAuthFlowEndRedirect(ctx, authCtx, request)
 		http.Redirect(writer, request, redirectURL, http.StatusTemporaryRedirect)
 	}
 }
@@ -211,7 +211,7 @@ func GetCallbackHandler(ctx context.Context, authCtx interfaces.AuthenticationCo
 			}
 			logger.Info(ctx, "Successfully called the preRedirect hook")
 		}
-		redirectURL := getAuthFlowEndRedirect(ctx, authCtx, request)
+		redirectURL := GetAuthFlowEndRedirect(ctx, authCtx, request)
 		http.Redirect(writer, request, redirectURL, http.StatusTemporaryRedirect)
 	}
 }

From 330a781a367e3cf0c363f95de307ba1b865a3977 Mon Sep 17 00:00:00 2001
From: Iaroslav Ciupin <iaroslav@union.ai>
Date: Mon, 11 Sep 2023 21:34:48 +0300
Subject: [PATCH 3/4] increase coverage

Signed-off-by: Iaroslav Ciupin <iaroslav@union.ai>
---
 auth/handlers_test.go                   | 103 ++++++++++++-
 auth/interfaces/cookie.go               |   2 +-
 auth/interfaces/mocks/cookie_handler.go | 196 +++++++++++++++++-------
 3 files changed, 241 insertions(+), 60 deletions(-)

diff --git a/auth/handlers_test.go b/auth/handlers_test.go
index 449b13c4a..0ba0653bc 100644
--- a/auth/handlers_test.go
+++ b/auth/handlers_test.go
@@ -2,6 +2,7 @@ package auth
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -11,8 +12,11 @@ import (
 	"testing"
 
 	"github.com/coreos/go-oidc"
+	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/service"
+	stdConfig "github.com/flyteorg/flytestdlib/config"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
 	"golang.org/x/oauth2"
 	"google.golang.org/protobuf/types/known/structpb"
 
@@ -21,8 +25,6 @@ import (
 	"github.com/flyteorg/flyteadmin/auth/interfaces/mocks"
 	"github.com/flyteorg/flyteadmin/pkg/common"
 	"github.com/flyteorg/flyteadmin/plugins"
-	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/service"
-	stdConfig "github.com/flyteorg/flytestdlib/config"
 )
 
 const (
@@ -50,8 +52,8 @@ func setupMockedAuthContextAtEndpoint(endpoint string) *mocks.AuthenticationCont
 		Timeout: IdpConnectionTimeout,
 	}
 	mockAuthCtx.OnCookieManagerMatch().Return(mockCookieHandler)
-	mockCookieHandler.OnSetTokenCookiesMatch(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
-	mockCookieHandler.OnSetUserInfoCookieMatch(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	mockCookieHandler.EXPECT().SetTokenCookies(mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
+	mockCookieHandler.EXPECT().SetUserInfoCookie(mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
 	mockAuthCtx.OnOAuth2ClientConfigMatch(mock.Anything).Return(&dummyOAuth2Config)
 	mockAuthCtx.OnGetHTTPClient().Return(dummyHTTPClient)
 	return mockAuthCtx
@@ -255,6 +257,99 @@ func TestGetLoginHandler(t *testing.T) {
 	assert.True(t, strings.Contains(w.Header().Get("Set-Cookie"), "flyte_csrf_state="))
 }
 
+func TestGetLogoutHandler(t *testing.T) {
+	ctx := context.Background()
+
+	t.Run("no_hook_no_redirect", func(t *testing.T) {
+		cookieHandler := new(mocks.CookieHandler)
+		authCtx := mocks.AuthenticationContext{}
+		authCtx.OnCookieManager().Return(cookieHandler).Once()
+		w := httptest.NewRecorder()
+		cookieHandler.EXPECT().DeleteCookies(ctx, w).Once()
+		r := plugins.NewRegistry()
+		req, err := http.NewRequest(http.MethodGet, "/logout", nil)
+		require.NoError(t, err)
+
+		GetLogoutEndpointHandler(ctx, &authCtx, r)(w, req)
+
+		assert.Equal(t, http.StatusOK, w.Code)
+		authCtx.AssertExpectations(t)
+		cookieHandler.AssertExpectations(t)
+	})
+
+	t.Run("no_hook_with_redirect", func(t *testing.T) {
+		ctx := context.Background()
+		cookieHandler := new(mocks.CookieHandler)
+		authCtx := mocks.AuthenticationContext{}
+		authCtx.OnCookieManager().Return(cookieHandler).Once()
+		w := httptest.NewRecorder()
+		cookieHandler.EXPECT().DeleteCookies(ctx, w).Once()
+		r := plugins.NewRegistry()
+		req, err := http.NewRequest(http.MethodGet, "/logout?redirect_url=/foo", nil)
+		require.NoError(t, err)
+
+		GetLogoutEndpointHandler(ctx, &authCtx, r)(w, req)
+
+		assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
+		authCtx.AssertExpectations(t)
+		cookieHandler.AssertExpectations(t)
+	})
+
+	t.Run("with_hook_with_redirect", func(t *testing.T) {
+		ctx := context.Background()
+		cookieHandler := new(mocks.CookieHandler)
+		authCtx := mocks.AuthenticationContext{}
+		authCtx.OnCookieManager().Return(cookieHandler).Once()
+		w := httptest.NewRecorder()
+		cookieHandler.EXPECT().DeleteCookies(ctx, w).Once()
+		r := plugins.NewRegistry()
+		hook := new(mock.Mock)
+		err := r.Register(plugins.PluginIDLogoutHook, LogoutHookFunc(func(
+			ctx context.Context,
+			authCtx interfaces.AuthenticationContext,
+			request *http.Request,
+			w http.ResponseWriter) error {
+			return hook.MethodCalled("hook").Error(0)
+		}))
+		hook.On("hook").Return(nil).Once()
+		require.NoError(t, err)
+		req, err := http.NewRequest(http.MethodGet, "/logout?redirect_url=/foo", nil)
+		require.NoError(t, err)
+
+		GetLogoutEndpointHandler(ctx, &authCtx, r)(w, req)
+
+		assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
+		authCtx.AssertExpectations(t)
+		cookieHandler.AssertExpectations(t)
+		hook.AssertExpectations(t)
+	})
+
+	t.Run("hook_error", func(t *testing.T) {
+		ctx := context.Background()
+		authCtx := mocks.AuthenticationContext{}
+		w := httptest.NewRecorder()
+		r := plugins.NewRegistry()
+		hook := new(mock.Mock)
+		err := r.Register(plugins.PluginIDLogoutHook, LogoutHookFunc(func(
+			ctx context.Context,
+			authCtx interfaces.AuthenticationContext,
+			request *http.Request,
+			w http.ResponseWriter) error {
+			return hook.MethodCalled("hook").Error(0)
+		}))
+		hook.On("hook").Return(errors.New("fail")).Once()
+		require.NoError(t, err)
+		req, err := http.NewRequest(http.MethodGet, "/logout?redirect_url=/foo", nil)
+		require.NoError(t, err)
+
+		GetLogoutEndpointHandler(ctx, &authCtx, r)(w, req)
+
+		assert.Equal(t, http.StatusInternalServerError, w.Code)
+		authCtx.AssertExpectations(t)
+		hook.AssertExpectations(t)
+	})
+}
+
 func TestGetHTTPRequestCookieToMetadataHandler(t *testing.T) {
 	ctx := context.Background()
 	// These were generated for unit testing only.
diff --git a/auth/interfaces/cookie.go b/auth/interfaces/cookie.go
index 988709a8f..743f40293 100644
--- a/auth/interfaces/cookie.go
+++ b/auth/interfaces/cookie.go
@@ -9,7 +9,7 @@ import (
 	"golang.org/x/oauth2"
 )
 
-//go:generate mockery -name=CookieHandler -output=mocks/ -case=underscore
+//go:generate mockery --name=CookieHandler --output=mocks/ --case=underscore --with-expecter
 
 type CookieHandler interface {
 	SetTokenCookies(ctx context.Context, writer http.ResponseWriter, token *oauth2.Token) error
diff --git a/auth/interfaces/mocks/cookie_handler.go b/auth/interfaces/mocks/cookie_handler.go
index 74371b0ce..942979a40 100644
--- a/auth/interfaces/mocks/cookie_handler.go
+++ b/auth/interfaces/mocks/cookie_handler.go
@@ -1,4 +1,4 @@
-// Code generated by mockery v1.0.1. DO NOT EDIT.
+// Code generated by mockery v2.18.0. DO NOT EDIT.
 
 package mocks
 
@@ -18,27 +18,41 @@ type CookieHandler struct {
 	mock.Mock
 }
 
+type CookieHandler_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *CookieHandler) EXPECT() *CookieHandler_Expecter {
+	return &CookieHandler_Expecter{mock: &_m.Mock}
+}
+
 // DeleteCookies provides a mock function with given fields: ctx, writer
 func (_m *CookieHandler) DeleteCookies(ctx context.Context, writer http.ResponseWriter) {
 	_m.Called(ctx, writer)
 }
 
-type CookieHandler_RetrieveAuthCodeRequest struct {
+// CookieHandler_DeleteCookies_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DeleteCookies'
+type CookieHandler_DeleteCookies_Call struct {
 	*mock.Call
 }
 
-func (_m CookieHandler_RetrieveAuthCodeRequest) Return(authRequestURL string, err error) *CookieHandler_RetrieveAuthCodeRequest {
-	return &CookieHandler_RetrieveAuthCodeRequest{Call: _m.Call.Return(authRequestURL, err)}
+// DeleteCookies is a helper method to define mock.On call
+//   - ctx context.Context
+//   - writer http.ResponseWriter
+func (_e *CookieHandler_Expecter) DeleteCookies(ctx interface{}, writer interface{}) *CookieHandler_DeleteCookies_Call {
+	return &CookieHandler_DeleteCookies_Call{Call: _e.mock.On("DeleteCookies", ctx, writer)}
 }
 
-func (_m *CookieHandler) OnRetrieveAuthCodeRequest(ctx context.Context, request *http.Request) *CookieHandler_RetrieveAuthCodeRequest {
-	c_call := _m.On("RetrieveAuthCodeRequest", ctx, request)
-	return &CookieHandler_RetrieveAuthCodeRequest{Call: c_call}
+func (_c *CookieHandler_DeleteCookies_Call) Run(run func(ctx context.Context, writer http.ResponseWriter)) *CookieHandler_DeleteCookies_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(http.ResponseWriter))
+	})
+	return _c
 }
 
-func (_m *CookieHandler) OnRetrieveAuthCodeRequestMatch(matchers ...interface{}) *CookieHandler_RetrieveAuthCodeRequest {
-	c_call := _m.On("RetrieveAuthCodeRequest", matchers...)
-	return &CookieHandler_RetrieveAuthCodeRequest{Call: c_call}
+func (_c *CookieHandler_DeleteCookies_Call) Return() *CookieHandler_DeleteCookies_Call {
+	_c.Call.Return()
+	return _c
 }
 
 // RetrieveAuthCodeRequest provides a mock function with given fields: ctx, request
@@ -62,22 +76,28 @@ func (_m *CookieHandler) RetrieveAuthCodeRequest(ctx context.Context, request *h
 	return r0, r1
 }
 
-type CookieHandler_RetrieveTokenValues struct {
+// CookieHandler_RetrieveAuthCodeRequest_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RetrieveAuthCodeRequest'
+type CookieHandler_RetrieveAuthCodeRequest_Call struct {
 	*mock.Call
 }
 
-func (_m CookieHandler_RetrieveTokenValues) Return(idToken string, accessToken string, refreshToken string, err error) *CookieHandler_RetrieveTokenValues {
-	return &CookieHandler_RetrieveTokenValues{Call: _m.Call.Return(idToken, accessToken, refreshToken, err)}
+// RetrieveAuthCodeRequest is a helper method to define mock.On call
+//   - ctx context.Context
+//   - request *http.Request
+func (_e *CookieHandler_Expecter) RetrieveAuthCodeRequest(ctx interface{}, request interface{}) *CookieHandler_RetrieveAuthCodeRequest_Call {
+	return &CookieHandler_RetrieveAuthCodeRequest_Call{Call: _e.mock.On("RetrieveAuthCodeRequest", ctx, request)}
 }
 
-func (_m *CookieHandler) OnRetrieveTokenValues(ctx context.Context, request *http.Request) *CookieHandler_RetrieveTokenValues {
-	c_call := _m.On("RetrieveTokenValues", ctx, request)
-	return &CookieHandler_RetrieveTokenValues{Call: c_call}
+func (_c *CookieHandler_RetrieveAuthCodeRequest_Call) Run(run func(ctx context.Context, request *http.Request)) *CookieHandler_RetrieveAuthCodeRequest_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(*http.Request))
+	})
+	return _c
 }
 
-func (_m *CookieHandler) OnRetrieveTokenValuesMatch(matchers ...interface{}) *CookieHandler_RetrieveTokenValues {
-	c_call := _m.On("RetrieveTokenValues", matchers...)
-	return &CookieHandler_RetrieveTokenValues{Call: c_call}
+func (_c *CookieHandler_RetrieveAuthCodeRequest_Call) Return(authRequestURL string, err error) *CookieHandler_RetrieveAuthCodeRequest_Call {
+	_c.Call.Return(authRequestURL, err)
+	return _c
 }
 
 // RetrieveTokenValues provides a mock function with given fields: ctx, request
@@ -115,22 +135,28 @@ func (_m *CookieHandler) RetrieveTokenValues(ctx context.Context, request *http.
 	return r0, r1, r2, r3
 }
 
-type CookieHandler_RetrieveUserInfo struct {
+// CookieHandler_RetrieveTokenValues_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RetrieveTokenValues'
+type CookieHandler_RetrieveTokenValues_Call struct {
 	*mock.Call
 }
 
-func (_m CookieHandler_RetrieveUserInfo) Return(_a0 *service.UserInfoResponse, _a1 error) *CookieHandler_RetrieveUserInfo {
-	return &CookieHandler_RetrieveUserInfo{Call: _m.Call.Return(_a0, _a1)}
+// RetrieveTokenValues is a helper method to define mock.On call
+//   - ctx context.Context
+//   - request *http.Request
+func (_e *CookieHandler_Expecter) RetrieveTokenValues(ctx interface{}, request interface{}) *CookieHandler_RetrieveTokenValues_Call {
+	return &CookieHandler_RetrieveTokenValues_Call{Call: _e.mock.On("RetrieveTokenValues", ctx, request)}
 }
 
-func (_m *CookieHandler) OnRetrieveUserInfo(ctx context.Context, request *http.Request) *CookieHandler_RetrieveUserInfo {
-	c_call := _m.On("RetrieveUserInfo", ctx, request)
-	return &CookieHandler_RetrieveUserInfo{Call: c_call}
+func (_c *CookieHandler_RetrieveTokenValues_Call) Run(run func(ctx context.Context, request *http.Request)) *CookieHandler_RetrieveTokenValues_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(*http.Request))
+	})
+	return _c
 }
 
-func (_m *CookieHandler) OnRetrieveUserInfoMatch(matchers ...interface{}) *CookieHandler_RetrieveUserInfo {
-	c_call := _m.On("RetrieveUserInfo", matchers...)
-	return &CookieHandler_RetrieveUserInfo{Call: c_call}
+func (_c *CookieHandler_RetrieveTokenValues_Call) Return(idToken string, accessToken string, refreshToken string, err error) *CookieHandler_RetrieveTokenValues_Call {
+	_c.Call.Return(idToken, accessToken, refreshToken, err)
+	return _c
 }
 
 // RetrieveUserInfo provides a mock function with given fields: ctx, request
@@ -156,22 +182,28 @@ func (_m *CookieHandler) RetrieveUserInfo(ctx context.Context, request *http.Req
 	return r0, r1
 }
 
-type CookieHandler_SetAuthCodeCookie struct {
+// CookieHandler_RetrieveUserInfo_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RetrieveUserInfo'
+type CookieHandler_RetrieveUserInfo_Call struct {
 	*mock.Call
 }
 
-func (_m CookieHandler_SetAuthCodeCookie) Return(_a0 error) *CookieHandler_SetAuthCodeCookie {
-	return &CookieHandler_SetAuthCodeCookie{Call: _m.Call.Return(_a0)}
+// RetrieveUserInfo is a helper method to define mock.On call
+//   - ctx context.Context
+//   - request *http.Request
+func (_e *CookieHandler_Expecter) RetrieveUserInfo(ctx interface{}, request interface{}) *CookieHandler_RetrieveUserInfo_Call {
+	return &CookieHandler_RetrieveUserInfo_Call{Call: _e.mock.On("RetrieveUserInfo", ctx, request)}
 }
 
-func (_m *CookieHandler) OnSetAuthCodeCookie(ctx context.Context, writer http.ResponseWriter, authRequestURL string) *CookieHandler_SetAuthCodeCookie {
-	c_call := _m.On("SetAuthCodeCookie", ctx, writer, authRequestURL)
-	return &CookieHandler_SetAuthCodeCookie{Call: c_call}
+func (_c *CookieHandler_RetrieveUserInfo_Call) Run(run func(ctx context.Context, request *http.Request)) *CookieHandler_RetrieveUserInfo_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(*http.Request))
+	})
+	return _c
 }
 
-func (_m *CookieHandler) OnSetAuthCodeCookieMatch(matchers ...interface{}) *CookieHandler_SetAuthCodeCookie {
-	c_call := _m.On("SetAuthCodeCookie", matchers...)
-	return &CookieHandler_SetAuthCodeCookie{Call: c_call}
+func (_c *CookieHandler_RetrieveUserInfo_Call) Return(_a0 *service.UserInfoResponse, _a1 error) *CookieHandler_RetrieveUserInfo_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
 }
 
 // SetAuthCodeCookie provides a mock function with given fields: ctx, writer, authRequestURL
@@ -188,22 +220,29 @@ func (_m *CookieHandler) SetAuthCodeCookie(ctx context.Context, writer http.Resp
 	return r0
 }
 
-type CookieHandler_SetTokenCookies struct {
+// CookieHandler_SetAuthCodeCookie_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetAuthCodeCookie'
+type CookieHandler_SetAuthCodeCookie_Call struct {
 	*mock.Call
 }
 
-func (_m CookieHandler_SetTokenCookies) Return(_a0 error) *CookieHandler_SetTokenCookies {
-	return &CookieHandler_SetTokenCookies{Call: _m.Call.Return(_a0)}
+// SetAuthCodeCookie is a helper method to define mock.On call
+//   - ctx context.Context
+//   - writer http.ResponseWriter
+//   - authRequestURL string
+func (_e *CookieHandler_Expecter) SetAuthCodeCookie(ctx interface{}, writer interface{}, authRequestURL interface{}) *CookieHandler_SetAuthCodeCookie_Call {
+	return &CookieHandler_SetAuthCodeCookie_Call{Call: _e.mock.On("SetAuthCodeCookie", ctx, writer, authRequestURL)}
 }
 
-func (_m *CookieHandler) OnSetTokenCookies(ctx context.Context, writer http.ResponseWriter, token *oauth2.Token) *CookieHandler_SetTokenCookies {
-	c_call := _m.On("SetTokenCookies", ctx, writer, token)
-	return &CookieHandler_SetTokenCookies{Call: c_call}
+func (_c *CookieHandler_SetAuthCodeCookie_Call) Run(run func(ctx context.Context, writer http.ResponseWriter, authRequestURL string)) *CookieHandler_SetAuthCodeCookie_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(http.ResponseWriter), args[2].(string))
+	})
+	return _c
 }
 
-func (_m *CookieHandler) OnSetTokenCookiesMatch(matchers ...interface{}) *CookieHandler_SetTokenCookies {
-	c_call := _m.On("SetTokenCookies", matchers...)
-	return &CookieHandler_SetTokenCookies{Call: c_call}
+func (_c *CookieHandler_SetAuthCodeCookie_Call) Return(_a0 error) *CookieHandler_SetAuthCodeCookie_Call {
+	_c.Call.Return(_a0)
+	return _c
 }
 
 // SetTokenCookies provides a mock function with given fields: ctx, writer, token
@@ -220,22 +259,29 @@ func (_m *CookieHandler) SetTokenCookies(ctx context.Context, writer http.Respon
 	return r0
 }
 
-type CookieHandler_SetUserInfoCookie struct {
+// CookieHandler_SetTokenCookies_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetTokenCookies'
+type CookieHandler_SetTokenCookies_Call struct {
 	*mock.Call
 }
 
-func (_m CookieHandler_SetUserInfoCookie) Return(_a0 error) *CookieHandler_SetUserInfoCookie {
-	return &CookieHandler_SetUserInfoCookie{Call: _m.Call.Return(_a0)}
+// SetTokenCookies is a helper method to define mock.On call
+//   - ctx context.Context
+//   - writer http.ResponseWriter
+//   - token *oauth2.Token
+func (_e *CookieHandler_Expecter) SetTokenCookies(ctx interface{}, writer interface{}, token interface{}) *CookieHandler_SetTokenCookies_Call {
+	return &CookieHandler_SetTokenCookies_Call{Call: _e.mock.On("SetTokenCookies", ctx, writer, token)}
 }
 
-func (_m *CookieHandler) OnSetUserInfoCookie(ctx context.Context, writer http.ResponseWriter, userInfo *service.UserInfoResponse) *CookieHandler_SetUserInfoCookie {
-	c_call := _m.On("SetUserInfoCookie", ctx, writer, userInfo)
-	return &CookieHandler_SetUserInfoCookie{Call: c_call}
+func (_c *CookieHandler_SetTokenCookies_Call) Run(run func(ctx context.Context, writer http.ResponseWriter, token *oauth2.Token)) *CookieHandler_SetTokenCookies_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(http.ResponseWriter), args[2].(*oauth2.Token))
+	})
+	return _c
 }
 
-func (_m *CookieHandler) OnSetUserInfoCookieMatch(matchers ...interface{}) *CookieHandler_SetUserInfoCookie {
-	c_call := _m.On("SetUserInfoCookie", matchers...)
-	return &CookieHandler_SetUserInfoCookie{Call: c_call}
+func (_c *CookieHandler_SetTokenCookies_Call) Return(_a0 error) *CookieHandler_SetTokenCookies_Call {
+	_c.Call.Return(_a0)
+	return _c
 }
 
 // SetUserInfoCookie provides a mock function with given fields: ctx, writer, userInfo
@@ -251,3 +297,43 @@ func (_m *CookieHandler) SetUserInfoCookie(ctx context.Context, writer http.Resp
 
 	return r0
 }
+
+// CookieHandler_SetUserInfoCookie_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetUserInfoCookie'
+type CookieHandler_SetUserInfoCookie_Call struct {
+	*mock.Call
+}
+
+// SetUserInfoCookie is a helper method to define mock.On call
+//   - ctx context.Context
+//   - writer http.ResponseWriter
+//   - userInfo *service.UserInfoResponse
+func (_e *CookieHandler_Expecter) SetUserInfoCookie(ctx interface{}, writer interface{}, userInfo interface{}) *CookieHandler_SetUserInfoCookie_Call {
+	return &CookieHandler_SetUserInfoCookie_Call{Call: _e.mock.On("SetUserInfoCookie", ctx, writer, userInfo)}
+}
+
+func (_c *CookieHandler_SetUserInfoCookie_Call) Run(run func(ctx context.Context, writer http.ResponseWriter, userInfo *service.UserInfoResponse)) *CookieHandler_SetUserInfoCookie_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(http.ResponseWriter), args[2].(*service.UserInfoResponse))
+	})
+	return _c
+}
+
+func (_c *CookieHandler_SetUserInfoCookie_Call) Return(_a0 error) *CookieHandler_SetUserInfoCookie_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+type mockConstructorTestingTNewCookieHandler interface {
+	mock.TestingT
+	Cleanup(func())
+}
+
+// NewCookieHandler creates a new instance of CookieHandler. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+func NewCookieHandler(t mockConstructorTestingTNewCookieHandler) *CookieHandler {
+	mock := &CookieHandler{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

From a254eda3a3bce972be37462f7fa0a7275f1fa0ed Mon Sep 17 00:00:00 2001
From: Iaroslav Ciupin <iaroslav@union.ai>
Date: Mon, 11 Sep 2023 22:25:11 +0300
Subject: [PATCH 4/4] fix

Signed-off-by: Iaroslav Ciupin <iaroslav@union.ai>
---
 auth/handlers_test.go                   |  20 ++-
 auth/interfaces/cookie.go               |   2 +-
 auth/interfaces/mocks/cookie_handler.go | 196 +++++++-----------------
 3 files changed, 65 insertions(+), 153 deletions(-)

diff --git a/auth/handlers_test.go b/auth/handlers_test.go
index 0ba0653bc..d6cfac743 100644
--- a/auth/handlers_test.go
+++ b/auth/handlers_test.go
@@ -52,8 +52,8 @@ func setupMockedAuthContextAtEndpoint(endpoint string) *mocks.AuthenticationCont
 		Timeout: IdpConnectionTimeout,
 	}
 	mockAuthCtx.OnCookieManagerMatch().Return(mockCookieHandler)
-	mockCookieHandler.EXPECT().SetTokenCookies(mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
-	mockCookieHandler.EXPECT().SetUserInfoCookie(mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
+	mockCookieHandler.OnSetTokenCookiesMatch(mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
+	mockCookieHandler.OnSetUserInfoCookieMatch(mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
 	mockAuthCtx.OnOAuth2ClientConfigMatch(mock.Anything).Return(&dummyOAuth2Config)
 	mockAuthCtx.OnGetHTTPClient().Return(dummyHTTPClient)
 	return mockAuthCtx
@@ -261,11 +261,10 @@ func TestGetLogoutHandler(t *testing.T) {
 	ctx := context.Background()
 
 	t.Run("no_hook_no_redirect", func(t *testing.T) {
-		cookieHandler := new(mocks.CookieHandler)
+		cookieHandler := &CookieManager{}
 		authCtx := mocks.AuthenticationContext{}
 		authCtx.OnCookieManager().Return(cookieHandler).Once()
 		w := httptest.NewRecorder()
-		cookieHandler.EXPECT().DeleteCookies(ctx, w).Once()
 		r := plugins.NewRegistry()
 		req, err := http.NewRequest(http.MethodGet, "/logout", nil)
 		require.NoError(t, err)
@@ -273,17 +272,16 @@ func TestGetLogoutHandler(t *testing.T) {
 		GetLogoutEndpointHandler(ctx, &authCtx, r)(w, req)
 
 		assert.Equal(t, http.StatusOK, w.Code)
+		require.Len(t, w.Result().Cookies(), 3)
 		authCtx.AssertExpectations(t)
-		cookieHandler.AssertExpectations(t)
 	})
 
 	t.Run("no_hook_with_redirect", func(t *testing.T) {
 		ctx := context.Background()
-		cookieHandler := new(mocks.CookieHandler)
+		cookieHandler := &CookieManager{}
 		authCtx := mocks.AuthenticationContext{}
 		authCtx.OnCookieManager().Return(cookieHandler).Once()
 		w := httptest.NewRecorder()
-		cookieHandler.EXPECT().DeleteCookies(ctx, w).Once()
 		r := plugins.NewRegistry()
 		req, err := http.NewRequest(http.MethodGet, "/logout?redirect_url=/foo", nil)
 		require.NoError(t, err)
@@ -292,16 +290,15 @@ func TestGetLogoutHandler(t *testing.T) {
 
 		assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
 		authCtx.AssertExpectations(t)
-		cookieHandler.AssertExpectations(t)
+		require.Len(t, w.Result().Cookies(), 3)
 	})
 
 	t.Run("with_hook_with_redirect", func(t *testing.T) {
 		ctx := context.Background()
-		cookieHandler := new(mocks.CookieHandler)
+		cookieHandler := &CookieManager{}
 		authCtx := mocks.AuthenticationContext{}
 		authCtx.OnCookieManager().Return(cookieHandler).Once()
 		w := httptest.NewRecorder()
-		cookieHandler.EXPECT().DeleteCookies(ctx, w).Once()
 		r := plugins.NewRegistry()
 		hook := new(mock.Mock)
 		err := r.Register(plugins.PluginIDLogoutHook, LogoutHookFunc(func(
@@ -319,8 +316,8 @@ func TestGetLogoutHandler(t *testing.T) {
 		GetLogoutEndpointHandler(ctx, &authCtx, r)(w, req)
 
 		assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
+		require.Len(t, w.Result().Cookies(), 3)
 		authCtx.AssertExpectations(t)
-		cookieHandler.AssertExpectations(t)
 		hook.AssertExpectations(t)
 	})
 
@@ -345,6 +342,7 @@ func TestGetLogoutHandler(t *testing.T) {
 		GetLogoutEndpointHandler(ctx, &authCtx, r)(w, req)
 
 		assert.Equal(t, http.StatusInternalServerError, w.Code)
+		assert.Empty(t, w.Result().Cookies())
 		authCtx.AssertExpectations(t)
 		hook.AssertExpectations(t)
 	})
diff --git a/auth/interfaces/cookie.go b/auth/interfaces/cookie.go
index 743f40293..988709a8f 100644
--- a/auth/interfaces/cookie.go
+++ b/auth/interfaces/cookie.go
@@ -9,7 +9,7 @@ import (
 	"golang.org/x/oauth2"
 )
 
-//go:generate mockery --name=CookieHandler --output=mocks/ --case=underscore --with-expecter
+//go:generate mockery -name=CookieHandler -output=mocks/ -case=underscore
 
 type CookieHandler interface {
 	SetTokenCookies(ctx context.Context, writer http.ResponseWriter, token *oauth2.Token) error
diff --git a/auth/interfaces/mocks/cookie_handler.go b/auth/interfaces/mocks/cookie_handler.go
index 942979a40..74371b0ce 100644
--- a/auth/interfaces/mocks/cookie_handler.go
+++ b/auth/interfaces/mocks/cookie_handler.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
 
@@ -18,41 +18,27 @@ type CookieHandler struct {
 	mock.Mock
 }
 
-type CookieHandler_Expecter struct {
-	mock *mock.Mock
-}
-
-func (_m *CookieHandler) EXPECT() *CookieHandler_Expecter {
-	return &CookieHandler_Expecter{mock: &_m.Mock}
-}
-
 // DeleteCookies provides a mock function with given fields: ctx, writer
 func (_m *CookieHandler) DeleteCookies(ctx context.Context, writer http.ResponseWriter) {
 	_m.Called(ctx, writer)
 }
 
-// CookieHandler_DeleteCookies_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DeleteCookies'
-type CookieHandler_DeleteCookies_Call struct {
+type CookieHandler_RetrieveAuthCodeRequest struct {
 	*mock.Call
 }
 
-// DeleteCookies is a helper method to define mock.On call
-//   - ctx context.Context
-//   - writer http.ResponseWriter
-func (_e *CookieHandler_Expecter) DeleteCookies(ctx interface{}, writer interface{}) *CookieHandler_DeleteCookies_Call {
-	return &CookieHandler_DeleteCookies_Call{Call: _e.mock.On("DeleteCookies", ctx, writer)}
+func (_m CookieHandler_RetrieveAuthCodeRequest) Return(authRequestURL string, err error) *CookieHandler_RetrieveAuthCodeRequest {
+	return &CookieHandler_RetrieveAuthCodeRequest{Call: _m.Call.Return(authRequestURL, err)}
 }
 
-func (_c *CookieHandler_DeleteCookies_Call) Run(run func(ctx context.Context, writer http.ResponseWriter)) *CookieHandler_DeleteCookies_Call {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(http.ResponseWriter))
-	})
-	return _c
+func (_m *CookieHandler) OnRetrieveAuthCodeRequest(ctx context.Context, request *http.Request) *CookieHandler_RetrieveAuthCodeRequest {
+	c_call := _m.On("RetrieveAuthCodeRequest", ctx, request)
+	return &CookieHandler_RetrieveAuthCodeRequest{Call: c_call}
 }
 
-func (_c *CookieHandler_DeleteCookies_Call) Return() *CookieHandler_DeleteCookies_Call {
-	_c.Call.Return()
-	return _c
+func (_m *CookieHandler) OnRetrieveAuthCodeRequestMatch(matchers ...interface{}) *CookieHandler_RetrieveAuthCodeRequest {
+	c_call := _m.On("RetrieveAuthCodeRequest", matchers...)
+	return &CookieHandler_RetrieveAuthCodeRequest{Call: c_call}
 }
 
 // RetrieveAuthCodeRequest provides a mock function with given fields: ctx, request
@@ -76,28 +62,22 @@ func (_m *CookieHandler) RetrieveAuthCodeRequest(ctx context.Context, request *h
 	return r0, r1
 }
 
-// CookieHandler_RetrieveAuthCodeRequest_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RetrieveAuthCodeRequest'
-type CookieHandler_RetrieveAuthCodeRequest_Call struct {
+type CookieHandler_RetrieveTokenValues struct {
 	*mock.Call
 }
 
-// RetrieveAuthCodeRequest is a helper method to define mock.On call
-//   - ctx context.Context
-//   - request *http.Request
-func (_e *CookieHandler_Expecter) RetrieveAuthCodeRequest(ctx interface{}, request interface{}) *CookieHandler_RetrieveAuthCodeRequest_Call {
-	return &CookieHandler_RetrieveAuthCodeRequest_Call{Call: _e.mock.On("RetrieveAuthCodeRequest", ctx, request)}
+func (_m CookieHandler_RetrieveTokenValues) Return(idToken string, accessToken string, refreshToken string, err error) *CookieHandler_RetrieveTokenValues {
+	return &CookieHandler_RetrieveTokenValues{Call: _m.Call.Return(idToken, accessToken, refreshToken, err)}
 }
 
-func (_c *CookieHandler_RetrieveAuthCodeRequest_Call) Run(run func(ctx context.Context, request *http.Request)) *CookieHandler_RetrieveAuthCodeRequest_Call {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(*http.Request))
-	})
-	return _c
+func (_m *CookieHandler) OnRetrieveTokenValues(ctx context.Context, request *http.Request) *CookieHandler_RetrieveTokenValues {
+	c_call := _m.On("RetrieveTokenValues", ctx, request)
+	return &CookieHandler_RetrieveTokenValues{Call: c_call}
 }
 
-func (_c *CookieHandler_RetrieveAuthCodeRequest_Call) Return(authRequestURL string, err error) *CookieHandler_RetrieveAuthCodeRequest_Call {
-	_c.Call.Return(authRequestURL, err)
-	return _c
+func (_m *CookieHandler) OnRetrieveTokenValuesMatch(matchers ...interface{}) *CookieHandler_RetrieveTokenValues {
+	c_call := _m.On("RetrieveTokenValues", matchers...)
+	return &CookieHandler_RetrieveTokenValues{Call: c_call}
 }
 
 // RetrieveTokenValues provides a mock function with given fields: ctx, request
@@ -135,28 +115,22 @@ func (_m *CookieHandler) RetrieveTokenValues(ctx context.Context, request *http.
 	return r0, r1, r2, r3
 }
 
-// CookieHandler_RetrieveTokenValues_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RetrieveTokenValues'
-type CookieHandler_RetrieveTokenValues_Call struct {
+type CookieHandler_RetrieveUserInfo struct {
 	*mock.Call
 }
 
-// RetrieveTokenValues is a helper method to define mock.On call
-//   - ctx context.Context
-//   - request *http.Request
-func (_e *CookieHandler_Expecter) RetrieveTokenValues(ctx interface{}, request interface{}) *CookieHandler_RetrieveTokenValues_Call {
-	return &CookieHandler_RetrieveTokenValues_Call{Call: _e.mock.On("RetrieveTokenValues", ctx, request)}
+func (_m CookieHandler_RetrieveUserInfo) Return(_a0 *service.UserInfoResponse, _a1 error) *CookieHandler_RetrieveUserInfo {
+	return &CookieHandler_RetrieveUserInfo{Call: _m.Call.Return(_a0, _a1)}
 }
 
-func (_c *CookieHandler_RetrieveTokenValues_Call) Run(run func(ctx context.Context, request *http.Request)) *CookieHandler_RetrieveTokenValues_Call {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(*http.Request))
-	})
-	return _c
+func (_m *CookieHandler) OnRetrieveUserInfo(ctx context.Context, request *http.Request) *CookieHandler_RetrieveUserInfo {
+	c_call := _m.On("RetrieveUserInfo", ctx, request)
+	return &CookieHandler_RetrieveUserInfo{Call: c_call}
 }
 
-func (_c *CookieHandler_RetrieveTokenValues_Call) Return(idToken string, accessToken string, refreshToken string, err error) *CookieHandler_RetrieveTokenValues_Call {
-	_c.Call.Return(idToken, accessToken, refreshToken, err)
-	return _c
+func (_m *CookieHandler) OnRetrieveUserInfoMatch(matchers ...interface{}) *CookieHandler_RetrieveUserInfo {
+	c_call := _m.On("RetrieveUserInfo", matchers...)
+	return &CookieHandler_RetrieveUserInfo{Call: c_call}
 }
 
 // RetrieveUserInfo provides a mock function with given fields: ctx, request
@@ -182,28 +156,22 @@ func (_m *CookieHandler) RetrieveUserInfo(ctx context.Context, request *http.Req
 	return r0, r1
 }
 
-// CookieHandler_RetrieveUserInfo_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RetrieveUserInfo'
-type CookieHandler_RetrieveUserInfo_Call struct {
+type CookieHandler_SetAuthCodeCookie struct {
 	*mock.Call
 }
 
-// RetrieveUserInfo is a helper method to define mock.On call
-//   - ctx context.Context
-//   - request *http.Request
-func (_e *CookieHandler_Expecter) RetrieveUserInfo(ctx interface{}, request interface{}) *CookieHandler_RetrieveUserInfo_Call {
-	return &CookieHandler_RetrieveUserInfo_Call{Call: _e.mock.On("RetrieveUserInfo", ctx, request)}
+func (_m CookieHandler_SetAuthCodeCookie) Return(_a0 error) *CookieHandler_SetAuthCodeCookie {
+	return &CookieHandler_SetAuthCodeCookie{Call: _m.Call.Return(_a0)}
 }
 
-func (_c *CookieHandler_RetrieveUserInfo_Call) Run(run func(ctx context.Context, request *http.Request)) *CookieHandler_RetrieveUserInfo_Call {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(*http.Request))
-	})
-	return _c
+func (_m *CookieHandler) OnSetAuthCodeCookie(ctx context.Context, writer http.ResponseWriter, authRequestURL string) *CookieHandler_SetAuthCodeCookie {
+	c_call := _m.On("SetAuthCodeCookie", ctx, writer, authRequestURL)
+	return &CookieHandler_SetAuthCodeCookie{Call: c_call}
 }
 
-func (_c *CookieHandler_RetrieveUserInfo_Call) Return(_a0 *service.UserInfoResponse, _a1 error) *CookieHandler_RetrieveUserInfo_Call {
-	_c.Call.Return(_a0, _a1)
-	return _c
+func (_m *CookieHandler) OnSetAuthCodeCookieMatch(matchers ...interface{}) *CookieHandler_SetAuthCodeCookie {
+	c_call := _m.On("SetAuthCodeCookie", matchers...)
+	return &CookieHandler_SetAuthCodeCookie{Call: c_call}
 }
 
 // SetAuthCodeCookie provides a mock function with given fields: ctx, writer, authRequestURL
@@ -220,29 +188,22 @@ func (_m *CookieHandler) SetAuthCodeCookie(ctx context.Context, writer http.Resp
 	return r0
 }
 
-// CookieHandler_SetAuthCodeCookie_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetAuthCodeCookie'
-type CookieHandler_SetAuthCodeCookie_Call struct {
+type CookieHandler_SetTokenCookies struct {
 	*mock.Call
 }
 
-// SetAuthCodeCookie is a helper method to define mock.On call
-//   - ctx context.Context
-//   - writer http.ResponseWriter
-//   - authRequestURL string
-func (_e *CookieHandler_Expecter) SetAuthCodeCookie(ctx interface{}, writer interface{}, authRequestURL interface{}) *CookieHandler_SetAuthCodeCookie_Call {
-	return &CookieHandler_SetAuthCodeCookie_Call{Call: _e.mock.On("SetAuthCodeCookie", ctx, writer, authRequestURL)}
+func (_m CookieHandler_SetTokenCookies) Return(_a0 error) *CookieHandler_SetTokenCookies {
+	return &CookieHandler_SetTokenCookies{Call: _m.Call.Return(_a0)}
 }
 
-func (_c *CookieHandler_SetAuthCodeCookie_Call) Run(run func(ctx context.Context, writer http.ResponseWriter, authRequestURL string)) *CookieHandler_SetAuthCodeCookie_Call {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(http.ResponseWriter), args[2].(string))
-	})
-	return _c
+func (_m *CookieHandler) OnSetTokenCookies(ctx context.Context, writer http.ResponseWriter, token *oauth2.Token) *CookieHandler_SetTokenCookies {
+	c_call := _m.On("SetTokenCookies", ctx, writer, token)
+	return &CookieHandler_SetTokenCookies{Call: c_call}
 }
 
-func (_c *CookieHandler_SetAuthCodeCookie_Call) Return(_a0 error) *CookieHandler_SetAuthCodeCookie_Call {
-	_c.Call.Return(_a0)
-	return _c
+func (_m *CookieHandler) OnSetTokenCookiesMatch(matchers ...interface{}) *CookieHandler_SetTokenCookies {
+	c_call := _m.On("SetTokenCookies", matchers...)
+	return &CookieHandler_SetTokenCookies{Call: c_call}
 }
 
 // SetTokenCookies provides a mock function with given fields: ctx, writer, token
@@ -259,29 +220,22 @@ func (_m *CookieHandler) SetTokenCookies(ctx context.Context, writer http.Respon
 	return r0
 }
 
-// CookieHandler_SetTokenCookies_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetTokenCookies'
-type CookieHandler_SetTokenCookies_Call struct {
+type CookieHandler_SetUserInfoCookie struct {
 	*mock.Call
 }
 
-// SetTokenCookies is a helper method to define mock.On call
-//   - ctx context.Context
-//   - writer http.ResponseWriter
-//   - token *oauth2.Token
-func (_e *CookieHandler_Expecter) SetTokenCookies(ctx interface{}, writer interface{}, token interface{}) *CookieHandler_SetTokenCookies_Call {
-	return &CookieHandler_SetTokenCookies_Call{Call: _e.mock.On("SetTokenCookies", ctx, writer, token)}
+func (_m CookieHandler_SetUserInfoCookie) Return(_a0 error) *CookieHandler_SetUserInfoCookie {
+	return &CookieHandler_SetUserInfoCookie{Call: _m.Call.Return(_a0)}
 }
 
-func (_c *CookieHandler_SetTokenCookies_Call) Run(run func(ctx context.Context, writer http.ResponseWriter, token *oauth2.Token)) *CookieHandler_SetTokenCookies_Call {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(http.ResponseWriter), args[2].(*oauth2.Token))
-	})
-	return _c
+func (_m *CookieHandler) OnSetUserInfoCookie(ctx context.Context, writer http.ResponseWriter, userInfo *service.UserInfoResponse) *CookieHandler_SetUserInfoCookie {
+	c_call := _m.On("SetUserInfoCookie", ctx, writer, userInfo)
+	return &CookieHandler_SetUserInfoCookie{Call: c_call}
 }
 
-func (_c *CookieHandler_SetTokenCookies_Call) Return(_a0 error) *CookieHandler_SetTokenCookies_Call {
-	_c.Call.Return(_a0)
-	return _c
+func (_m *CookieHandler) OnSetUserInfoCookieMatch(matchers ...interface{}) *CookieHandler_SetUserInfoCookie {
+	c_call := _m.On("SetUserInfoCookie", matchers...)
+	return &CookieHandler_SetUserInfoCookie{Call: c_call}
 }
 
 // SetUserInfoCookie provides a mock function with given fields: ctx, writer, userInfo
@@ -297,43 +251,3 @@ func (_m *CookieHandler) SetUserInfoCookie(ctx context.Context, writer http.Resp
 
 	return r0
 }
-
-// CookieHandler_SetUserInfoCookie_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetUserInfoCookie'
-type CookieHandler_SetUserInfoCookie_Call struct {
-	*mock.Call
-}
-
-// SetUserInfoCookie is a helper method to define mock.On call
-//   - ctx context.Context
-//   - writer http.ResponseWriter
-//   - userInfo *service.UserInfoResponse
-func (_e *CookieHandler_Expecter) SetUserInfoCookie(ctx interface{}, writer interface{}, userInfo interface{}) *CookieHandler_SetUserInfoCookie_Call {
-	return &CookieHandler_SetUserInfoCookie_Call{Call: _e.mock.On("SetUserInfoCookie", ctx, writer, userInfo)}
-}
-
-func (_c *CookieHandler_SetUserInfoCookie_Call) Run(run func(ctx context.Context, writer http.ResponseWriter, userInfo *service.UserInfoResponse)) *CookieHandler_SetUserInfoCookie_Call {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(http.ResponseWriter), args[2].(*service.UserInfoResponse))
-	})
-	return _c
-}
-
-func (_c *CookieHandler_SetUserInfoCookie_Call) Return(_a0 error) *CookieHandler_SetUserInfoCookie_Call {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-type mockConstructorTestingTNewCookieHandler interface {
-	mock.TestingT
-	Cleanup(func())
-}
-
-// NewCookieHandler creates a new instance of CookieHandler. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
-func NewCookieHandler(t mockConstructorTestingTNewCookieHandler) *CookieHandler {
-	mock := &CookieHandler{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}