feat(remoteoauth): Add WithToken option (#1898)

instead of `WithAccessToken`, which needed explicit default values and didn't support refresh tokens in non-cloud mode. --------- Co-authored-by: Kemal Hadimli <disq@users.noreply.github.com>
cloudquery · Sep 19, 2024 · ff7a485 · ff7a485
1 parent 4201fa6
commit ff7a485
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 3 deletions.
diff --git a/helpers/remoteoauth/token_test.go b/helpers/remoteoauth/token_test.go
@@ -9,11 +9,24 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/oauth2"
 )
 
 const testAPIKey = "test-key"
 
 func TestLocalTokenAccess(t *testing.T) {
+	r := require.New(t)
+	_, cloud := os.LookupEnv("CQ_CLOUD")
+	r.False(cloud, "CQ_CLOUD should not be set")
+	tok, err := NewTokenSource(WithToken(oauth2.Token{AccessToken: "token", TokenType: "bearer"}))
+	r.NoError(err)
+	tk, err := tok.Token()
+	r.NoError(err)
+	r.True(tk.Valid())
+	r.Equal("token", tk.AccessToken)
+}
+
+func TestLocalTokenAccessWithDeprecatedTokenOpt(t *testing.T) {
 	r := require.New(t)
 	_, cloud := os.LookupEnv("CQ_CLOUD")
 	r.False(cloud, "CQ_CLOUD should not be set")
@@ -41,7 +54,7 @@ func TestFirstLocalTokenAccess(t *testing.T) {
 		"_CQ_CONNECTOR_ID":   connID,
 	})
 	r := require.New(t)
-	tok, err := NewTokenSource(WithAccessToken("token", "bearer", time.Time{}))
+	tok, err := NewTokenSource(WithToken(oauth2.Token{AccessToken: "token"}))
 	r.NoError(err)
 	tk, err := tok.Token()
 	r.NoError(err)
@@ -63,7 +76,7 @@ func TestInvalidAPIKeyTokenAccess(t *testing.T) {
 		"_CQ_CONNECTOR_ID":   connID,
 	})
 	r := require.New(t)
-	tok, err := NewTokenSource(WithAccessToken("token", "bearer", time.Time{}))
+	tok, err := NewTokenSource(WithToken(oauth2.Token{AccessToken: "token"}))
 	r.NoError(err)
 	tk, err := tok.Token()
 	r.Nil(tk)
@@ -110,7 +123,7 @@ func TestTestConnectionTokenAccess(t *testing.T) {
 		"_CQ_CONNECTOR_ID":            connID,
 	})
 	r := require.New(t)
-	tok, err := NewTokenSource(WithAccessToken("token", "bearer", time.Time{}))
+	tok, err := NewTokenSource(WithToken(oauth2.Token{AccessToken: "token"}))
 	r.NoError(err)
 	tk, err := tok.Token()
 	r.NoError(err)

diff --git a/helpers/remoteoauth/tokenoptions.go b/helpers/remoteoauth/tokenoptions.go
@@ -9,6 +9,8 @@ import (
 
 type TokenSourceOption func(source *tokenSource)
 
+// WithAccessToken sets the access token, token type and expiry time for the token source.
+// Deprecated: Use WithToken instead.
 func WithAccessToken(token, tokenType string, expiry time.Time) TokenSourceOption {
 	return func(t *tokenSource) {
 		t.currentToken = oauth2.Token{
@@ -19,6 +21,14 @@ func WithAccessToken(token, tokenType string, expiry time.Time) TokenSourceOptio
 	}
 }
 
+// WithToken sets the default token for the token source.
+func WithToken(token oauth2.Token) TokenSourceOption {
+	return func(t *tokenSource) {
+		t.currentToken = token
+	}
+}
+
+// WithDefaultContext sets the default context for the token source, used when creating a new token request.
 func WithDefaultContext(ctx context.Context) TokenSourceOption {
 	return func(t *tokenSource) {
 		t.defaultContext = ctx