Merge 398ad59 into a7fc664

flyteorg · Oct 30, 2023 · 160a457 · 160a457
2 parents a7fc664 + 398ad59
commit 160a457
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 3 deletions.
diff --git a/flyteidl/clients/go/admin/deviceflow/payload.go b/flyteidl/clients/go/admin/deviceflow/payload.go
@@ -40,5 +40,6 @@ type DeviceAccessTokenRequest struct {
 
 type DeviceAccessTokenResponse struct {
 	oauth2.Token
-	Error string `json:"error"`
+	Error     string `json:"error"`
+	ExpiresIn int64  `json:"expires_in"` // relative seconds from now
 }
diff --git a/flyteidl/clients/go/admin/deviceflow/token_orchestrator.go b/flyteidl/clients/go/admin/deviceflow/token_orchestrator.go
@@ -125,12 +125,14 @@ func (t TokenOrchestrator) PollTokenEndpoint(ctx context.Context, tokReq DeviceA
 		// Unmarshalled response if it contains an error then check if we need to increase the polling interval
 		if len(tokResp.Error) > 0 {
 			if tokResp.Error == errSlowDown || tokResp.Error == errAuthPending {
-				pollInterval = pollInterval * 2
-
+				logger.Debugf(ctx, "going to poll again due to error %v", tokResp.Error)
 			} else {
 				return nil, fmt.Errorf("oauth error : %v", tokResp.Error)
 			}
 		} else {
+			if secs := tokResp.ExpiresIn; secs > 0 {
+				tokResp.Token.Expiry = time.Now().Add(time.Duration(secs) * time.Second)
+			}
 			// Got the auth token in the response and save it in the cache
 			err = t.TokenCache.SaveToken(&tokResp.Token)
 			// Saving into the cache is only considered to be a warning in this case.

diff --git a/flyteidl/clients/go/admin/deviceflow/token_orchestrator_test.go b/flyteidl/clients/go/admin/deviceflow/token_orchestrator_test.go
@@ -87,6 +87,7 @@ func TestFetchFromAuthFlow(t *testing.T) {
 					Token: oauth2.Token{
 						AccessToken: "access_token",
 					},
+					ExpiresIn: 300,
 				}
 				darBytes, err := json.Marshal(dar)
 				assert.Nil(t, err)
@@ -121,5 +122,6 @@ func TestFetchFromAuthFlow(t *testing.T) {
 		assert.Nil(t, err)
 		assert.NotNil(t, authToken)
 		assert.Equal(t, "access_token", authToken.AccessToken)
+		assert.True(t, authToken.Expiry.After(time.Now().Add(time.Second*200)))
 	})
 }