fix(test): add unit test for auth activation initialize method

Author: yueny2020

packages/core/src/auth/activation.ts

@@ -12,7 +12,7 @@ import { isAmazonQ, isSageMaker } from '../shared/extensionUtilities'
 import { getLogger } from '../shared/logger/logger'
 import { getErrorMsg } from '../shared/errors'
 
-interface SagemakerCookie {
+export interface SagemakerCookie {
     authMode?: 'Sso' | 'Iam'
 }
 

packages/core/src/test/auth/activation.test.ts

@@ -0,0 +1,292 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as vscode from 'vscode'
+import * as sinon from 'sinon'
+import assert from 'assert'
+import { initialize, SagemakerCookie } from '../../auth/activation'
+import { LoginManager } from '../../auth/deprecated/loginManager'
+import * as extensionUtilities from '../../shared/extensionUtilities'
+import * as authUtils from '../../auth/utils'
+import * as errors from '../../shared/errors'
+import * as credentials from '../../auth/providers/credentials'
+
+describe('auth/activation', function () {
+    let sandbox: sinon.SinonSandbox
+    let mockLoginManager: LoginManager
+    let executeCommandStub: sinon.SinonStub
+    let isAmazonQStub: sinon.SinonStub
+    let isSageMakerStub: sinon.SinonStub
+    let initializeCredentialsProviderManagerStub: sinon.SinonStub
+    let getErrorMsgStub: sinon.SinonStub
+    let fromStringStub: sinon.SinonStub
+    let mockLogger: any
+    let onDidChangeActiveConnectionEvent: sinon.SinonStub
+
+    beforeEach(function () {
+        sandbox = sinon.createSandbox()
+
+        // Create mocks
+        mockLoginManager = {
+            login: sandbox.stub(),
+            logout: sandbox.stub(),
+        } as any
+
+        mockLogger = {
+            warn: sandbox.stub(),
+            info: sandbox.stub(),
+            error: sandbox.stub(),
+            debug: sandbox.stub(),
+        }
+
+        onDidChangeActiveConnectionEvent = sandbox.stub()
+
+        // Stub external dependencies
+        executeCommandStub = sandbox.stub(vscode.commands, 'executeCommand')
+        isAmazonQStub = sandbox.stub(extensionUtilities, 'isAmazonQ')
+        isSageMakerStub = sandbox.stub(extensionUtilities, 'isSageMaker')
+        initializeCredentialsProviderManagerStub = sandbox.stub(authUtils, 'initializeCredentialsProviderManager')
+        getErrorMsgStub = sandbox.stub(errors, 'getErrorMsg')
+        fromStringStub = sandbox.stub(credentials, 'fromString')
+    })
+
+    afterEach(function () {
+        sandbox.restore()
+    })
+
+    describe('initialize', function () {
+        it('should not execute sagemaker.parseCookies when not in AmazonQ and SageMaker environment', async function () {
+            isAmazonQStub.returns(false)
+            isSageMakerStub.returns(false)
+
+            await initialize(mockLoginManager)
+
+            assert.ok(!executeCommandStub.called)
+            assert.ok(!initializeCredentialsProviderManagerStub.called)
+        })
+
+        it('should not execute sagemaker.parseCookies when only in AmazonQ environment', async function () {
+            isAmazonQStub.returns(true)
+            isSageMakerStub.returns(false)
+
+            await initialize(mockLoginManager)
+
+            assert.ok(!executeCommandStub.called)
+            assert.ok(!initializeCredentialsProviderManagerStub.called)
+        })
+
+        it('should not execute sagemaker.parseCookies when only in SageMaker environment', async function () {
+            isAmazonQStub.returns(false)
+            isSageMakerStub.returns(true)
+
+            await initialize(mockLoginManager)
+
+            assert.ok(!executeCommandStub.called)
+            assert.ok(!initializeCredentialsProviderManagerStub.called)
+        })
+
+        it('should execute sagemaker.parseCookies when in both AmazonQ and SageMaker environment', async function () {
+            isAmazonQStub.returns(true)
+            isSageMakerStub.returns(true)
+            executeCommandStub.withArgs('sagemaker.parseCookies').resolves({ authMode: 'Sso' } as SagemakerCookie)
+
+            await initialize(mockLoginManager)
+
+            assert.ok(executeCommandStub.calledOnceWith('sagemaker.parseCookies'))
+            assert.ok(!initializeCredentialsProviderManagerStub.called)
+        })
+
+        it('should initialize credentials provider manager when authMode is not Sso', async function () {
+            isAmazonQStub.returns(true)
+            isSageMakerStub.returns(true)
+            executeCommandStub.withArgs('sagemaker.parseCookies').resolves({ authMode: 'Iam' } as SagemakerCookie)
+
+            await initialize(mockLoginManager)
+
+            assert.ok(executeCommandStub.calledOnceWith('sagemaker.parseCookies'))
+            assert.ok(initializeCredentialsProviderManagerStub.calledOnce)
+        })
+
+        it('should initialize credentials provider manager when authMode is undefined', async function () {
+            isAmazonQStub.returns(true)
+            isSageMakerStub.returns(true)
+            executeCommandStub.withArgs('sagemaker.parseCookies').resolves({} as SagemakerCookie)
+
+            await initialize(mockLoginManager)
+
+            assert.ok(executeCommandStub.calledOnceWith('sagemaker.parseCookies'))
+            assert.ok(initializeCredentialsProviderManagerStub.calledOnce)
+        })
+
+        it('should warn and not throw when sagemaker.parseCookies command is not found', async function () {
+            isAmazonQStub.returns(true)
+            isSageMakerStub.returns(true)
+            const error = new Error("command 'sagemaker.parseCookies' not found")
+            executeCommandStub.withArgs('sagemaker.parseCookies').rejects(error)
+            getErrorMsgStub.returns("command 'sagemaker.parseCookies' not found")
+
+            await initialize(mockLoginManager)
+
+            assert.ok(executeCommandStub.calledOnceWith('sagemaker.parseCookies'))
+            assert.ok(getErrorMsgStub.calledOnceWith(error))
+            assert.ok(
+                mockLogger.warn.calledOnceWith(
+                    'yueny testing Failed to execute command "sagemaker.parseCookies": Error: command \'sagemaker.parseCookies\' not found'
+                )
+            )
+            assert.ok(!initializeCredentialsProviderManagerStub.called)
+        })
+
+        it('should throw when sagemaker.parseCookies fails with non-command-not-found error', async function () {
+            isAmazonQStub.returns(true)
+            isSageMakerStub.returns(true)
+            const error = new Error('Some other error')
+            executeCommandStub.withArgs('sagemaker.parseCookies').rejects(error)
+            getErrorMsgStub.returns('Some other error')
+
+            await assert.rejects(initialize(mockLoginManager), /Some other error/)
+
+            assert.ok(executeCommandStub.calledOnceWith('sagemaker.parseCookies'))
+            assert.ok(getErrorMsgStub.calledOnceWith(error))
+            assert.ok(!mockLogger.warn.called)
+            assert.ok(!initializeCredentialsProviderManagerStub.called)
+        })
+
+        it('should set up connection change listener correctly', async function () {
+            isAmazonQStub.returns(false)
+            isSageMakerStub.returns(false)
+            let eventCallback: (conn: any) => void = () => {}
+
+            onDidChangeActiveConnectionEvent.callsFake((callback: any) => {
+                eventCallback = callback
+                return { dispose: sandbox.stub() }
+            })
+
+            await initialize(mockLoginManager)
+
+            assert.ok(onDidChangeActiveConnectionEvent.calledOnce)
+
+            // Test the event callback is set up
+            assert.strictEqual(typeof eventCallback, 'function')
+        })
+
+        describe('connection change event handler', function () {
+            let eventCallback: (conn: any) => Promise<void>
+
+            beforeEach(async function () {
+                isAmazonQStub.returns(false)
+                isSageMakerStub.returns(false)
+
+                onDidChangeActiveConnectionEvent.callsFake((callback: any) => {
+                    eventCallback = callback
+                    return { dispose: sandbox.stub() }
+                })
+
+                await initialize(mockLoginManager)
+            })
+
+            it('should call login when connection is valid IAM', async function () {
+                const mockConnection = {
+                    type: 'iam',
+                    state: 'valid',
+                    id: 'test-connection-id',
+                }
+                const mockProviderId = { credentialSource: 'test', credentialTypeId: 'test' }
+                fromStringStub.withArgs('test-connection-id').returns(mockProviderId)
+                ;(mockLoginManager.login as sinon.SinonStub).resolves(true)
+
+                await eventCallback(mockConnection)
+
+                assert.ok(fromStringStub.calledOnceWith('test-connection-id'))
+                assert.ok(
+                    (mockLoginManager.login as sinon.SinonStub).calledOnceWith({
+                        passive: true,
+                        providerId: mockProviderId,
+                    })
+                )
+                assert.ok(!(mockLoginManager.logout as sinon.SinonStub).called)
+            })
+
+            it('should call logout when connection is not IAM', async function () {
+                const mockConnection = {
+                    type: 'sso',
+                    state: 'valid',
+                    id: 'test-connection-id',
+                }
+                ;(mockLoginManager.logout as sinon.SinonStub).resolves()
+
+                await eventCallback(mockConnection)
+
+                assert.ok((mockLoginManager.logout as sinon.SinonStub).calledOnce)
+                assert.ok(!(mockLoginManager.login as sinon.SinonStub).called)
+                assert.ok(!fromStringStub.called)
+            })
+
+            it('should call logout when connection state is not valid', async function () {
+                const mockConnection = {
+                    type: 'iam',
+                    state: 'invalid',
+                    id: 'test-connection-id',
+                }
+                ;(mockLoginManager.logout as sinon.SinonStub).resolves()
+
+                await eventCallback(mockConnection)
+
+                assert.ok((mockLoginManager.logout as sinon.SinonStub).calledOnce)
+                assert.ok(!(mockLoginManager.login as sinon.SinonStub).called)
+                assert.ok(!fromStringStub.called)
+            })
+
+            it('should call logout when connection is undefined', async function () {
+                ;(mockLoginManager.logout as sinon.SinonStub).resolves()
+
+                await eventCallback(undefined)
+
+                assert.ok((mockLoginManager.logout as sinon.SinonStub).calledOnce)
+                assert.ok(!(mockLoginManager.login as sinon.SinonStub).called)
+                assert.ok(!fromStringStub.called)
+            })
+        })
+
+        it('should handle the complete flow for AmazonQ + SageMaker with Iam authMode', async function () {
+            isAmazonQStub.returns(true)
+            isSageMakerStub.returns(true)
+            executeCommandStub.withArgs('sagemaker.parseCookies').resolves({ authMode: 'Iam' } as SagemakerCookie)
+
+            let eventCallback: (conn: any) => Promise<void> = async () => {}
+            onDidChangeActiveConnectionEvent.callsFake((callback: any) => {
+                eventCallback = callback
+                return { dispose: sandbox.stub() }
+            })
+
+            const mockConnection = {
+                type: 'iam',
+                state: 'valid',
+                id: 'test-connection-id',
+            }
+            const mockProviderId = { credentialSource: 'test', credentialTypeId: 'test' }
+            fromStringStub.withArgs('test-connection-id').returns(mockProviderId)
+            ;(mockLoginManager.login as sinon.SinonStub).resolves(true)
+
+            await initialize(mockLoginManager)
+
+            // Verify SageMaker initialization
+            assert.ok(executeCommandStub.calledOnceWith('sagemaker.parseCookies'))
+            assert.ok(initializeCredentialsProviderManagerStub.calledOnce)
+
+            // Verify event listener is set up
+            assert.ok(onDidChangeActiveConnectionEvent.calledOnce)
+
+            // Test event handler works correctly
+            await eventCallback(mockConnection)
+            assert.ok(
+                (mockLoginManager.login as sinon.SinonStub).calledOnceWith({
+                    passive: true,
+                    providerId: mockProviderId,
+                })
+            )
+        })
+    })
+})