lint and fix unit test

flyteorg · Mar 26, 2024 · cb9c6b8 · cb9c6b8
1 parent 4ea162c
commit cb9c6b8
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 7 deletions.
diff --git a/flyteadmin/auth/cookie.go b/flyteadmin/auth/cookie.go
@@ -19,7 +19,7 @@ import (
 const (
 	// #nosec
 	accessTokenCookieName = "flyte_at"
-	// nosec
+	// #nosec
 	accessTokenCookieNameSplit = "flyte_at_1"
 	// #nosec
 	idTokenCookieName = "flyte_idt"

diff --git a/flyteadmin/auth/cookie_manager.go b/flyteadmin/auth/cookie_manager.go
@@ -172,7 +172,11 @@ func (c CookieManager) SetTokenCookies(ctx context.Context, writer http.Response
 		return errors.Errorf(ErrTokenNil, "Attempting to set cookies with nil token")
 	}
 
-	c.StoreAccessToken(ctx, token.AccessToken, writer)
+	err = c.StoreAccessToken(ctx, token.AccessToken, writer)
+
+	if err != nil {
+		return logger.Errorf(ctx, "Error storing access token %s", err)
+	}
 
 	if idTokenRaw, converted := token.Extra(idTokenExtra).(string); converted {
 		idCookie, err := NewSecureCookie(idTokenCookieName, idTokenRaw, c.hashKey, c.blockKey, c.domain, c.getHTTPSameSitePolicy())

diff --git a/flyteadmin/auth/handlers_test.go b/flyteadmin/auth/handlers_test.go
@@ -305,7 +305,7 @@ 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)
+		require.Len(t, w.Result().Cookies(), 4)
 		authCtx.AssertExpectations(t)
 	})
 
@@ -323,7 +323,7 @@ func TestGetLogoutHandler(t *testing.T) {
 
 		assert.Equal(t, http.StatusTemporaryRedirect, w.Code)
 		authCtx.AssertExpectations(t)
-		require.Len(t, w.Result().Cookies(), 3)
+		require.Len(t, w.Result().Cookies(), 4)
 	})
 
 	t.Run("with_hook_with_redirect", func(t *testing.T) {
@@ -349,7 +349,7 @@ 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)
+		require.Len(t, w.Result().Cookies(), 4)
 		authCtx.AssertExpectations(t)
 		hook.AssertExpectations(t)
 	})
@@ -403,11 +403,15 @@ func TestGetHTTPRequestCookieToMetadataHandler(t *testing.T) {
 	assert.NoError(t, err)
 	req.AddCookie(&accessTokenCookie)
 
-	idCookie, err := NewSecureCookie(idTokenCookieName, "a.b.c", cookieManager.hashKey, cookieManager.blockKey, "localhost", http.SameSiteDefaultMode)
+	accessTokenCookieSplit, err = NewSecureCookie(accessTokenCookieNameSplit, ".d.e.f", cookieManager.hashKey, cookieManager.blockKey, "localhost", http.SameSiteDefaultMode)
+	assert.NoError(t, err)
+	req.AddCookie(&accessTokenCookieSplit)
+
+	idCookie, err := NewSecureCookie(idTokenCookieName, "a.b.c.d.e.f", cookieManager.hashKey, cookieManager.blockKey, "localhost", http.SameSiteDefaultMode)
 	assert.NoError(t, err)
 	req.AddCookie(&idCookie)
 
-	assert.Equal(t, "IDToken a.b.c", handler(ctx, req)["authorization"][0])
+	assert.Equal(t, "IDToken a.b.c.d.e.f", handler(ctx, req)["authorization"][0])
 }
 
 func TestGetHTTPMetadataTaggingHandler(t *testing.T) {