Add unit test

Author: opieter-aws

packages/amazonq/src/lsp/client.ts

@@ -178,6 +178,9 @@ export async function startLanguageServer(
     await client.onReady()
     AuthUtil.create(new auth2.LanguageClientAuth(client, clientId, encryptionKey))
 
+    // Migrate SSO connections from old Auth to the LSP identity server
+    // This function only migrates connections once
+    // This call can be removed once all/most users have updated to the latest AmazonQ version
     try {
         await AuthUtil.instance.migrateSsoConnectionToLsp(clientName)
     } catch (e) {

packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts

@@ -5,10 +5,12 @@
 
 import assert from 'assert'
 import * as sinon from 'sinon'
-import { AuthUtil } from 'aws-core-vscode/codewhisperer'
+import * as path from 'path'
+import { AuthUtil, amazonQScopes } from 'aws-core-vscode/codewhisperer'
 import { createTestAuthUtil } from 'aws-core-vscode/test'
-import { constants } from 'aws-core-vscode/auth'
+import { constants, cache } from 'aws-core-vscode/auth'
 import { auth2 } from 'aws-core-vscode/auth'
+import { mementoUtils, makeTemporaryToolkitFolder, tryRemoveFolder, fs } from 'aws-core-vscode/shared'
 
 describe('AuthUtil', async function () {
     let auth: any
@@ -193,4 +195,145 @@ describe('AuthUtil', async function () {
             assert.ok(clearCacheSpy.called)
         })
     })
+
+    describe('migrateSsoConnectionToLsp', function () {
+        let memento: any
+        let cacheDir: string
+        let fromRegistrationFile: string
+        let fromTokenFile: string
+
+        const validProfile = {
+            type: 'sso',
+            startUrl: 'https://test2.com',
+            ssoRegion: 'us-east-1',
+            scopes: amazonQScopes,
+            metadata: {
+                connectionState: 'valid',
+            },
+        }
+        beforeEach(async function () {
+            memento = {
+                get: sinon.stub(),
+                update: sinon.stub().resolves(),
+            }
+            cacheDir = await makeTemporaryToolkitFolder()
+
+            sinon.stub(mementoUtils, 'getEnvironmentSpecificMemento').returns(memento)
+            sinon.stub(cache, 'getCacheDir').returns(cacheDir)
+
+            fromTokenFile = cache.getTokenCacheFile(cacheDir, 'profile1')
+            const registrationKey = {
+                startUrl: validProfile.startUrl,
+                region: validProfile.ssoRegion,
+                scopes: amazonQScopes,
+            }
+            fromRegistrationFile = cache.getRegistrationCacheFile(cacheDir, registrationKey)
+
+            const registrationData = { test: 'registration' }
+            const tokenData = { test: 'token' }
+
+            await fs.writeFile(fromRegistrationFile, JSON.stringify(registrationData))
+            await fs.writeFile(fromTokenFile, JSON.stringify(tokenData))
+        })
+
+        afterEach(async function () {
+            await Promise.all([tryRemoveFolder(cacheDir)])
+            sinon.restore()
+        })
+
+        it('migrates valid SSO connection', async function () {
+            memento.get.returns({ profile1: validProfile })
+
+            const updateProfileStub = sinon.stub((auth as any).session, 'updateProfile').resolves()
+
+            await auth.migrateSsoConnectionToLsp('test-client')
+
+            assert.ok(updateProfileStub.calledOnce)
+            assert.ok(memento.update.calledWith('auth.profiles', undefined))
+
+            const files = await fs.readdir(cacheDir)
+            assert.strictEqual(files.length, 2) // Should have both the token and registration file
+
+            // Verify file contents were preserved
+            const newFiles = files.map((f) => path.join(cacheDir, f[0]))
+            for (const file of newFiles) {
+                const content = await fs.readFileText(file)
+                const parsed = JSON.parse(content)
+                assert.ok(parsed.test === 'registration' || parsed.test === 'token')
+            }
+        })
+
+        it('does not migrate if no matching SSO profile exists', async function () {
+            const mockProfiles = {
+                'test-profile': {
+                    type: 'iam',
+                    startUrl: 'https://test.com',
+                    ssoRegion: 'us-east-1',
+                },
+            }
+            memento.get.returns(mockProfiles)
+
+            await auth.migrateSsoConnectionToLsp('test-client')
+
+            // Assert that the file names have not updated
+            const files = await fs.readdir(cacheDir)
+            assert.ok(files.length === 2)
+            assert.ok(await fs.exists(fromRegistrationFile))
+            assert.ok(await fs.exists(fromTokenFile))
+            assert.ok(!memento.update.called)
+        })
+
+        it('migrates only profile with matching scopes', async function () {
+            const mockProfiles = {
+                profile1: validProfile,
+                profile2: {
+                    type: 'sso',
+                    startUrl: 'https://test.com',
+                    ssoRegion: 'us-east-1',
+                    scopes: ['different:scope'],
+                    metadata: {
+                        connectionState: 'valid',
+                    },
+                },
+            }
+            memento.get.returns(mockProfiles)
+
+            const updateProfileStub = sinon.stub((auth as any).session, 'updateProfile').resolves()
+
+            await auth.migrateSsoConnectionToLsp('test-client')
+
+            assert.ok(updateProfileStub.calledOnce)
+            assert.ok(memento.update.calledWith('auth.profiles', undefined))
+            assert.deepStrictEqual(updateProfileStub.firstCall.args[0], {
+                startUrl: validProfile.startUrl,
+                region: validProfile.ssoRegion,
+                scopes: validProfile.scopes,
+            })
+        })
+
+        it('uses valid connection state when multiple profiles exist', async function () {
+            const mockProfiles = {
+                profile2: {
+                    ...validProfile,
+                    metadata: {
+                        connectionState: 'invalid',
+                    },
+                },
+                profile1: validProfile,
+            }
+            memento.get.returns(mockProfiles)
+
+            const updateProfileStub = sinon.stub((auth as any).session, 'updateProfile').resolves()
+
+            await auth.migrateSsoConnectionToLsp('test-client')
+
+            assert.ok(
+                updateProfileStub.calledWith({
+                    startUrl: validProfile.startUrl,
+                    region: validProfile.ssoRegion,
+                    scopes: validProfile.scopes,
+                })
+            )
+        })
+    })
 })

packages/core/src/auth/index.ts

@@ -23,6 +23,7 @@ export { Auth } from './auth'
 export { CredentialsStore } from './credentials/store'
 export { LoginManager } from './deprecated/loginManager'
 export * as constants from './sso/constants'
+export * as cache from './sso/cache'
 export * as authUtils from './utils'
 export * as auth2 from './auth2'
 export * as SsoAccessTokenProvider from './sso/ssoAccessTokenProvider'

packages/core/src/codewhisperer/util/authUtil.ts

@@ -6,7 +6,7 @@
 import * as vscode from 'vscode'
 import * as localizedText from '../../shared/localizedText'
 import * as nls from 'vscode-nls'
-import * as fs from 'fs'
+import { fs } from '../../shared/fs/fs'
 import * as path from 'path'
 import * as crypto from 'crypto'
 import { ToolkitError } from '../../shared/errors'
@@ -403,8 +403,8 @@ export class AuthUtil implements IAuthProvider {
                 const toTokenFile = filePath(this.profileName)
 
                 try {
-                    fs.renameSync(fromRegistrationFile, toRegistrationFile)
-                    fs.renameSync(fromTokenFile, toTokenFile)
+                    await fs.rename(fromRegistrationFile, toRegistrationFile)
+                    await fs.rename(fromTokenFile, toTokenFile)
                     getLogger().debug('Successfully renamed registration and token files')
                 } catch (err) {
                     getLogger().error(`Failed to rename files during migration: ${err}`)

packages/core/src/shared/index.ts

@@ -74,3 +74,4 @@ export * as BaseLspInstaller from './lsp/baseLspInstaller'
 export * as collectionUtil from './utilities/collectionUtils'
 export * from './datetime'
 export * from './performance/marks'
+export * as mementoUtils from './utilities/mementos'