more tests, should cover above 80%

Author: ndyakov

identity/managed_identity_provider.go

@@ -6,9 +6,18 @@ import (
 	"fmt"
 
 	mi "github.com/AzureAD/microsoft-authentication-library-for-go/apps/managedidentity"
+	"github.com/AzureAD/microsoft-authentication-library-for-go/apps/public"
 	"github.com/redis-developer/go-redis-entraid/shared"
 )
 
+// ManagedIdentityClient is an interface that defines the methods for a managed identity client.
+// It is used to acquire a token using the managed identity.
+type ManagedIdentityClient interface {
+	// AcquireToken acquires a token using the managed identity.
+	// It returns the token and an error if any.
+	AcquireToken(ctx context.Context, resource string, opts ...mi.AcquireTokenOption) (public.AuthResult, error)
+}
+
 // ManagedIdentityProviderOptions represents the options for the managed identity provider.
 // It is used to configure the identity provider when requesting a manager.
 type ManagedIdentityProviderOptions struct {
@@ -38,14 +47,22 @@ type ManagedIdentityProvider struct {
 	scopes []string
 
 	// client is the managed identity client used to request a manager.
-	client *mi.Client
+	client ManagedIdentityClient
+}
+
+// realManagedIdentityClient is a wrapper around the real mi.Client that implements our interface
+type realManagedIdentityClient struct {
+	client mi.Client
+}
+
+func (c *realManagedIdentityClient) AcquireToken(ctx context.Context, resource string, opts ...mi.AcquireTokenOption) (public.AuthResult, error) {
+	return c.client.AcquireToken(ctx, resource, opts...)
 }
 
 // NewManagedIdentityProvider creates a new managed identity provider for Azure with managed identity.
 // It is used to configure the identity provider when requesting a manager.
 func NewManagedIdentityProvider(opts ManagedIdentityProviderOptions) (*ManagedIdentityProvider, error) {
-	var client mi.Client
-	var err error
+	var client ManagedIdentityClient
 
 	if opts.ManagedIdentityType != SystemAssignedIdentity && opts.ManagedIdentityType != UserAssignedIdentity {
 		return nil, errors.New("invalid managed identity type")
@@ -56,25 +73,29 @@ func NewManagedIdentityProvider(opts ManagedIdentityProviderOptions) (*ManagedId
 		// SystemAssignedIdentity is the type of identity that is automatically managed by Azure.
 		// This type of identity is automatically created and managed by Azure.
 		// It is used to authenticate the identity when requesting a manager.
-		client, err = mi.New(mi.SystemAssigned())
+		miClient, err := mi.New(mi.SystemAssigned())
+		if err != nil {
+			return nil, fmt.Errorf("couldn't create managed identity client: %w", err)
+		}
+		client = &realManagedIdentityClient{client: miClient}
 	case UserAssignedIdentity:
 		// UserAssignedIdentity is required to be specified when using a user assigned identity.
 		if opts.UserAssignedClientID == "" {
 			return nil, errors.New("user assigned client ID is required when using user assigned identity")
 		}
 		// UserAssignedIdentity is the type of identity that is managed by the user.
-		client, err = mi.New(mi.UserAssignedClientID(opts.UserAssignedClientID))
-	}
-
-	if err != nil {
-		return nil, fmt.Errorf("couldn't create managed identity client: %w", err)
+		miClient, err := mi.New(mi.UserAssignedClientID(opts.UserAssignedClientID))
+		if err != nil {
+			return nil, fmt.Errorf("couldn't create managed identity client: %w", err)
+		}
+		client = &realManagedIdentityClient{client: miClient}
 	}
 
 	return &ManagedIdentityProvider{
 		userAssignedClientID: opts.UserAssignedClientID,
 		managedIdentityType:  opts.ManagedIdentityType,
 		scopes:               opts.Scopes,
-		client:               &client,
+		client:               client,
 	}, nil
 }
 

identity/managed_identity_provider_test.go

@@ -0,0 +1,218 @@
+package identity
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	mi "github.com/AzureAD/microsoft-authentication-library-for-go/apps/managedidentity"
+	"github.com/AzureAD/microsoft-authentication-library-for-go/apps/public"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+)
+
+// MockManagedIdentityClient is a mock implementation of the managed identity client
+type MockManagedIdentityClient struct {
+	mock.Mock
+}
+
+func (m *MockManagedIdentityClient) AcquireToken(ctx context.Context, resource string, opts ...mi.AcquireTokenOption) (public.AuthResult, error) {
+	args := m.Called(ctx, resource)
+	return args.Get(0).(public.AuthResult), args.Error(1)
+}
+
+func TestNewManagedIdentityProvider(t *testing.T) {
+	tests := []struct {
+		name          string
+		opts          ManagedIdentityProviderOptions
+		expectedError string
+	}{
+		{
+			name: "System assigned identity",
+			opts: ManagedIdentityProviderOptions{
+				ManagedIdentityType: SystemAssignedIdentity,
+				Scopes:              []string{"https://redis.azure.com"},
+			},
+			expectedError: "",
+		},
+		{
+			name: "User assigned identity with client ID",
+			opts: ManagedIdentityProviderOptions{
+				ManagedIdentityType:  UserAssignedIdentity,
+				UserAssignedClientID: "test-client-id",
+				Scopes:               []string{"https://redis.azure.com"},
+			},
+			expectedError: "",
+		},
+		{
+			name: "User assigned identity without client ID",
+			opts: ManagedIdentityProviderOptions{
+				ManagedIdentityType: UserAssignedIdentity,
+				Scopes:              []string{"https://redis.azure.com"},
+			},
+			expectedError: "user assigned client ID is required when using user assigned identity",
+		},
+		{
+			name: "Invalid identity type",
+			opts: ManagedIdentityProviderOptions{
+				ManagedIdentityType: "invalid-type",
+				Scopes:              []string{"https://redis.azure.com"},
+			},
+			expectedError: "invalid managed identity type",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			provider, err := NewManagedIdentityProvider(tt.opts)
+
+			if tt.expectedError != "" {
+				assert.Error(t, err)
+				assert.Contains(t, err.Error(), tt.expectedError)
+				assert.Nil(t, provider)
+			} else {
+				assert.NoError(t, err)
+				assert.NotNil(t, provider)
+				assert.Equal(t, tt.opts.ManagedIdentityType, provider.managedIdentityType)
+				assert.Equal(t, tt.opts.UserAssignedClientID, provider.userAssignedClientID)
+				assert.Equal(t, tt.opts.Scopes, provider.scopes)
+				assert.NotNil(t, provider.client)
+			}
+		})
+	}
+}
+
+func TestRequestToken(t *testing.T) {
+	tests := []struct {
+		name          string
+		provider      *ManagedIdentityProvider
+		expectedError string
+	}{
+		{
+			name: "Success with default resource",
+			provider: &ManagedIdentityProvider{
+				scopes: []string{},
+				client: new(MockManagedIdentityClient),
+			},
+			expectedError: "",
+		},
+		{
+			name: "Success with custom resource",
+			provider: &ManagedIdentityProvider{
+				scopes: []string{"custom-resource"},
+				client: new(MockManagedIdentityClient),
+			},
+			expectedError: "",
+		},
+		{
+			name: "Error when client is nil",
+			provider: &ManagedIdentityProvider{
+				scopes: []string{},
+				client: nil,
+			},
+			expectedError: "managed identity client is not initialized",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Set up the mock expectations if we have a mock client
+			if tt.provider.client != nil {
+				mockClient := tt.provider.client.(*MockManagedIdentityClient)
+				expectedResource := RedisResource
+				if len(tt.provider.scopes) > 0 {
+					expectedResource = tt.provider.scopes[0]
+				}
+
+				if tt.expectedError == "" {
+					mockClient.On("AcquireToken", mock.Anything, expectedResource).
+						Return(public.AuthResult{
+							AccessToken: "test-token",
+							ExpiresOn:   time.Now().Add(time.Hour),
+						}, nil)
+				} else {
+					mockClient.On("AcquireToken", mock.Anything, expectedResource).
+						Return(public.AuthResult{}, errors.New(tt.expectedError))
+				}
+			}
+
+			response, err := tt.provider.RequestToken()
+
+			if tt.expectedError != "" {
+				assert.Error(t, err)
+				assert.Contains(t, err.Error(), tt.expectedError)
+				assert.Nil(t, response)
+			} else {
+				assert.NoError(t, err)
+				assert.NotNil(t, response)
+			}
+
+			// Verify mock expectations
+			if tt.provider.client != nil {
+				mockClient := tt.provider.client.(*MockManagedIdentityClient)
+				mockClient.AssertExpectations(t)
+			}
+		})
+	}
+}
+
+func TestRequestToken_ErrorCases(t *testing.T) {
+	tests := []struct {
+		name          string
+		provider      *ManagedIdentityProvider
+		setupMock     func(*MockManagedIdentityClient)
+		expectedError string
+	}{
+		{
+			name: "AcquireToken fails",
+			provider: &ManagedIdentityProvider{
+				scopes: []string{},
+				client: new(MockManagedIdentityClient),
+			},
+			setupMock: func(m *MockManagedIdentityClient) {
+				m.On("AcquireToken", mock.Anything, RedisResource).
+					Return(public.AuthResult{}, errors.New("failed to acquire token"))
+			},
+			expectedError: "coudn't acquire manager: failed to acquire token",
+		},
+		{
+			name: "AcquireToken fails with custom resource",
+			provider: &ManagedIdentityProvider{
+				scopes: []string{"custom-resource"},
+				client: new(MockManagedIdentityClient),
+			},
+			setupMock: func(m *MockManagedIdentityClient) {
+				m.On("AcquireToken", mock.Anything, "custom-resource").
+					Return(public.AuthResult{}, errors.New("failed to acquire token"))
+			},
+			expectedError: "coudn't acquire manager: failed to acquire token",
+		},
+		{
+			name: "AcquireToken fails with invalid resource",
+			provider: &ManagedIdentityProvider{
+				scopes: []string{"invalid-resource"},
+				client: new(MockManagedIdentityClient),
+			},
+			setupMock: func(m *MockManagedIdentityClient) {
+				m.On("AcquireToken", mock.Anything, "invalid-resource").
+					Return(public.AuthResult{}, errors.New("invalid resource"))
+			},
+			expectedError: "coudn't acquire manager: invalid resource",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mockClient := tt.provider.client.(*MockManagedIdentityClient)
+			tt.setupMock(mockClient)
+
+			response, err := tt.provider.RequestToken()
+
+			assert.Error(t, err)
+			assert.Contains(t, err.Error(), tt.expectedError)
+			assert.Nil(t, response)
+			mockClient.AssertExpectations(t)
+		})
+	}
+}