From 6bb5f1c023b657b0e3bc363935a347b84b3eeccb Mon Sep 17 00:00:00 2001
From: hbc <bcxxxxxx@gmail.com>
Date: Tue, 19 Dec 2023 11:39:10 -0800
Subject: [PATCH] feat: implement library token provider

---
 pkg/token/options_ctor_test.go |  1 -
 pkg/token/provider.go          | 41 +++++++++++++++++++
 pkg/token/provider_test.go     | 74 ++++++++++++++++++++++++++++++++++
 3 files changed, 115 insertions(+), 1 deletion(-)
 create mode 100644 pkg/token/provider.go
 create mode 100644 pkg/token/provider_test.go

diff --git a/pkg/token/options_ctor_test.go b/pkg/token/options_ctor_test.go
index 9f922173..bf0032f9 100644
--- a/pkg/token/options_ctor_test.go
+++ b/pkg/token/options_ctor_test.go
@@ -122,7 +122,6 @@ func TestOptions_toInternalOptions(t *testing.T) {
 		}, o.toInternalOptions())
 	})
 
-
 	// this test uses reflection to ensure all fields in *Options
 	// are copied to *token.Options without modification.
 	t.Run("fields assignment", func(t *testing.T) {
diff --git a/pkg/token/provider.go b/pkg/token/provider.go
new file mode 100644
index 00000000..479d9ce1
--- /dev/null
+++ b/pkg/token/provider.go
@@ -0,0 +1,41 @@
+package token
+
+import (
+	"context"
+
+	"github.com/Azure/kubelogin/pkg/internal/token"
+)
+
+type tokenProviderShim struct {
+	impl token.TokenProvider
+}
+
+var _ TokenProvider = (*tokenProviderShim)(nil)
+
+func (tp *tokenProviderShim) GetAccessToken(ctx context.Context) (AccessToken, error) {
+	t, err := tp.impl.Token(ctx)
+	if err != nil {
+		return AccessToken{}, err
+	}
+
+	rv := AccessToken{
+		Token:     t.AccessToken,
+		ExpiresOn: t.Expires(),
+	}
+
+	return rv, nil
+}
+
+// GetTokenProvider returns a token provider based on the given options.
+func GetTokenProvider(options *Options) (TokenProvider, error) {
+	impl, err := token.NewTokenProvider(options.toInternalOptions())
+	if err != nil {
+		return nil, err
+	}
+
+	rv := &tokenProviderShim{
+		impl: impl,
+	}
+
+	return rv, nil
+}
diff --git a/pkg/token/provider_test.go b/pkg/token/provider_test.go
new file mode 100644
index 00000000..4bc41e00
--- /dev/null
+++ b/pkg/token/provider_test.go
@@ -0,0 +1,74 @@
+package token
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+
+	"github.com/Azure/go-autorest/autorest/adal"
+	"github.com/Azure/kubelogin/pkg/internal/token/mock_token"
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/mock/gomock"
+)
+
+func TestGetTokenProvider(t *testing.T) {
+	t.Run("invalid login method", func(t *testing.T) {
+		opts := &Options{
+			LoginMethod: "invalid-login-method",
+		}
+		tp, err := GetTokenProvider(opts)
+		assert.Error(t, err)
+		assert.Nil(t, tp)
+	})
+
+	t.Run("basic", func(t *testing.T) {
+		opts := &Options{
+			LoginMethod:        MSILogin,
+			ClientID:           "client-id",
+			IdentityResourceID: "identity-resource-id",
+			ServerID:           "server-id",
+		}
+		tp, err := GetTokenProvider(opts)
+		assert.NoError(t, err)
+		assert.NotNil(t, tp)
+	})
+}
+
+func TestTokenProviderShim_GetAccessToken(t *testing.T) {
+	t.Run("failure case", func(t *testing.T) {
+		mockCtrl := gomock.NewController(t)
+		defer mockCtrl.Finish()
+
+		mockTokenProvider := mock_token.NewMockTokenProvider(mockCtrl)
+		mockTokenProvider.EXPECT().Token(gomock.Any()).Return(adal.Token{}, assert.AnError)
+
+		tp := &tokenProviderShim{
+			impl: mockTokenProvider,
+		}
+
+		token, err := tp.GetAccessToken(context.Background())
+		assert.Equal(t, AccessToken{}, token)
+		assert.Equal(t, assert.AnError, err)
+	})
+
+	t.Run("success case", func(t *testing.T) {
+		mockCtrl := gomock.NewController(t)
+		defer mockCtrl.Finish()
+
+		adalToken := adal.Token{
+			AccessToken: "access-token",
+			ExpiresOn:   json.Number("1700000000"),
+		}
+		mockTokenProvider := mock_token.NewMockTokenProvider(mockCtrl)
+		mockTokenProvider.EXPECT().Token(gomock.Any()).Return(adalToken, nil)
+
+		tp := &tokenProviderShim{
+			impl: mockTokenProvider,
+		}
+
+		token, err := tp.GetAccessToken(context.Background())
+		assert.NoError(t, err)
+		assert.Equal(t, adalToken.AccessToken, token.Token)
+		assert.Equal(t, adalToken.Expires(), token.ExpiresOn)
+	})
+}