Skip to content

Commit 337b30b

Browse files
yuxianrzyuxianrz
committed
rename login to login_sso and login_iam
1 parent 962aaee commit 337b30b

File tree

10 files changed

+83
-62
lines changed

10 files changed

+83
-62
lines changed

‎packages/amazonq/test/e2e/amazonq/utils/setup.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,5 +22,5 @@ export async function loginToIdC() {
2222
)
2323
}
2424

25-
await AuthUtil.instance.login(startUrl, region, 'sso')
25+
await AuthUtil.instance.login_sso(startUrl, region)
2626
}

‎packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,11 @@ describe('RegionProfileManager', async function () {
2626

2727
async function setupConnection(type: 'builderId' | 'idc') {
2828
if (type === 'builderId') {
29-
await AuthUtil.instance.login(constants.builderIdStartUrl, region, 'sso')
29+
await AuthUtil.instance.login_sso(constants.builderIdStartUrl, region)
3030
assert.ok(AuthUtil.instance.isSsoSession())
3131
assert.ok(AuthUtil.instance.isBuilderIdConnection())
3232
} else if (type === 'idc') {
33-
await AuthUtil.instance.login(enterpriseSsoStartUrl, region, 'sso')
33+
await AuthUtil.instance.login_sso(enterpriseSsoStartUrl, region)
3434
assert.ok(AuthUtil.instance.isSsoSession())
3535
assert.ok(AuthUtil.instance.isIdcConnection())
3636
}

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

Lines changed: 48 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -26,19 +26,19 @@ describe('AuthUtil', async function () {
2626

2727
describe('Auth state', function () {
2828
it('login with BuilderId', async function () {
29-
await auth.login(constants.builderIdStartUrl, constants.builderIdRegion)
29+
await auth.login_sso(constants.builderIdStartUrl, constants.builderIdRegion)
3030
assert.ok(auth.isConnected())
3131
assert.ok(auth.isBuilderIdConnection())
3232
})
3333

3434
it('login with IDC', async function () {
35-
await auth.login('https://example.awsapps.com/start', 'us-east-1')
35+
await auth.login_sso('https://example.awsapps.com/start', 'us-east-1')
3636
assert.ok(auth.isConnected())
3737
assert.ok(auth.isIdcConnection())
3838
})
3939

4040
it('identifies internal users', async function () {
41-
await auth.login(constants.internalStartUrl, 'us-east-1')
41+
await auth.login_sso(constants.internalStartUrl, 'us-east-1')
4242
assert.ok(auth.isInternalAmazonUser())
4343
})
4444

@@ -55,7 +55,7 @@ describe('AuthUtil', async function () {
5555

5656
describe('Token management', function () {
5757
it('can get token when connected with SSO', async function () {
58-
await auth.login(constants.builderIdStartUrl, constants.builderIdRegion)
58+
await auth.login_sso(constants.builderIdStartUrl, constants.builderIdRegion)
5959
const token = await auth.getToken()
6060
assert.ok(token)
6161
})
@@ -68,14 +68,14 @@ describe('AuthUtil', async function () {
6868

6969
describe('getTelemetryMetadata', function () {
7070
it('returns valid metadata for BuilderId connection', async function () {
71-
await auth.login(constants.builderIdStartUrl, constants.builderIdRegion)
71+
await auth.login_sso(constants.builderIdStartUrl, constants.builderIdRegion)
7272
const metadata = await auth.getTelemetryMetadata()
7373
assert.strictEqual(metadata.credentialSourceId, 'awsId')
7474
assert.strictEqual(metadata.credentialStartUrl, constants.builderIdStartUrl)
7575
})
7676

7777
it('returns valid metadata for IDC connection', async function () {
78-
await auth.login('https://example.awsapps.com/start', 'us-east-1')
78+
await auth.login_sso('https://example.awsapps.com/start', 'us-east-1')
7979
const metadata = await auth.getTelemetryMetadata()
8080
assert.strictEqual(metadata.credentialSourceId, 'iamIdentityCenter')
8181
assert.strictEqual(metadata.credentialStartUrl, 'https://example.awsapps.com/start')
@@ -96,37 +96,38 @@ describe('AuthUtil', async function () {
9696
})
9797

9898
it('returns BuilderId forms when using BuilderId', async function () {
99-
await auth.login(constants.builderIdStartUrl, constants.builderIdRegion)
99+
await auth.login_sso(constants.builderIdStartUrl, constants.builderIdRegion)
100100
const forms = await auth.getAuthFormIds()
101101
assert.deepStrictEqual(forms, ['builderIdCodeWhisperer'])
102102
})
103103

104104
it('returns IDC forms when using IDC without SSO account access', async function () {
105105
const session = (auth as any).session
106-
sinon.stub(session, 'getProfile').resolves({
106+
session && sinon.stub(session, 'getProfile').resolves({
107107
ssoSession: {
108108
settings: {
109109
sso_registration_scopes: ['codewhisperer:*'],
110110
},
111111
},
112112
})
113113

114-
await auth.login('https://example.awsapps.com/start', 'us-east-1')
114+
await auth.login_sso('https://example.awsapps.com/start', 'us-east-1')
115115
const forms = await auth.getAuthFormIds()
116116
assert.deepStrictEqual(forms, ['identityCenterCodeWhisperer'])
117117
})
118118

119119
it('returns IDC forms with explorer when using IDC with SSO account access', async function () {
120+
await auth.login_sso('https://example.awsapps.com/start', 'us-east-1')
120121
const session = (auth as any).session
121-
sinon.stub(session, 'getProfile').resolves({
122+
123+
session && sinon.stub(session, 'getProfile').resolves({
122124
ssoSession: {
123125
settings: {
124126
sso_registration_scopes: ['codewhisperer:*', 'sso:account:access'],
125127
},
126128
},
127129
})
128130

129-
await auth.login('https://example.awsapps.com/start', 'us-east-1')
130131
const forms = await auth.getAuthFormIds()
131132
assert.deepStrictEqual(forms.sort(), ['identityCenterCodeWhisperer', 'identityCenterExplorer'].sort())
132133
})
@@ -178,7 +179,7 @@ describe('AuthUtil', async function () {
178179
})
179180

180181
it('updates bearer token when state is refreshed', async function () {
181-
await auth.login(constants.builderIdStartUrl, 'us-east-1')
182+
await auth.login_sso(constants.builderIdStartUrl, 'us-east-1')
182183

183184
await (auth as any).stateChangeHandler({ state: 'refreshed' })
184185

@@ -187,7 +188,7 @@ describe('AuthUtil', async function () {
187188
})
188189

189190
it('cleans up when connection expires', async function () {
190-
await auth.login(constants.builderIdStartUrl, 'us-east-1')
191+
await auth.login_sso(constants.builderIdStartUrl, 'us-east-1')
191192

192193
await (auth as any).stateChangeHandler({ state: 'expired' })
193194

@@ -197,13 +198,15 @@ describe('AuthUtil', async function () {
197198
it('deletes bearer token when disconnected', async function () {
198199
await (auth as any).stateChangeHandler({ state: 'notConnected' })
199200

200-
assert.ok(mockLspAuth.deleteBearerToken.called)
201+
if (auth.isSsoSession(auth.session)){
202+
assert.ok(mockLspAuth.deleteBearerToken.called)
203+
}
201204
})
202205

203206
it('updates bearer token and restores profile on reconnection', async function () {
204207
const restoreProfileSelectionSpy = sinon.spy(regionProfileManager, 'restoreProfileSelection')
205208

206-
await auth.login('https://example.awsapps.com/start', 'us-east-1')
209+
await auth.login_sso('https://example.awsapps.com/start', 'us-east-1')
207210

208211
await (auth as any).stateChangeHandler({ state: 'connected' })
209212

@@ -215,7 +218,7 @@ describe('AuthUtil', async function () {
215218
const invalidateProfileSpy = sinon.spy(regionProfileManager, 'invalidateProfile')
216219
const clearCacheSpy = sinon.spy(regionProfileManager, 'clearCache')
217220

218-
await auth.login('https://example.awsapps.com/start', 'us-east-1')
221+
await auth.login_sso('https://example.awsapps.com/start', 'us-east-1')
219222

220223
await (auth as any).stateChangeHandler({ state: 'expired' })
221224

@@ -280,12 +283,16 @@ describe('AuthUtil', async function () {
280283
await auth.migrateSsoConnectionToLsp('test-client')
281284

282285
assert.ok(memento.update.calledWith('auth.profiles', undefined))
283-
assert.ok(!auth.session.updateProfile?.called)
286+
assert.ok(!auth.session?.updateProfile?.called)
284287
})
285288

286289
it('proceeds with migration if LSP token check throws', async function () {
287290
memento.get.returns({ profile1: validProfile })
288291
mockLspAuth.getSsoToken.rejects(new Error('Token check failed'))
292+
293+
if (!(auth as any).session){
294+
auth.session = new auth2.SsoLogin(auth.profileName, auth.lspAuth, auth.eventEmitter)
295+
}
289296
const updateProfileStub = sinon.stub((auth as any).session, 'updateProfile').resolves()
290297

291298
await auth.migrateSsoConnectionToLsp('test-client')
@@ -297,22 +304,24 @@ describe('AuthUtil', async function () {
297304
it('migrates valid SSO connection', async function () {
298305
memento.get.returns({ profile1: validProfile })
299306

300-
const updateProfileStub = sinon.stub((auth as any).session, 'updateProfile').resolves()
307+
if ((auth as any).session) {
308+
const updateProfileStub = sinon.stub((auth as any).session, 'updateProfile').resolves()
301309

302-
await auth.migrateSsoConnectionToLsp('test-client')
310+
await auth.migrateSsoConnectionToLsp('test-client')
303311

304-
assert.ok(updateProfileStub.calledOnce)
305-
assert.ok(memento.update.calledWith('auth.profiles', undefined))
312+
assert.ok(updateProfileStub.calledOnce)
313+
assert.ok(memento.update.calledWith('auth.profiles', undefined))
306314

307-
const files = await fs.readdir(cacheDir)
308-
assert.strictEqual(files.length, 2) // Should have both the token and registration file
309-
310-
// Verify file contents were preserved
311-
const newFiles = files.map((f) => path.join(cacheDir, f[0]))
312-
for (const file of newFiles) {
313-
const content = await fs.readFileText(file)
314-
const parsed = JSON.parse(content)
315-
assert.ok(parsed.test === 'registration' || parsed.test === 'token')
315+
const files = await fs.readdir(cacheDir)
316+
assert.strictEqual(files.length, 2) // Should have both the token and registration file
317+
318+
// Verify file contents were preserved
319+
const newFiles = files.map((f) => path.join(cacheDir, f[0]))
320+
for (const file of newFiles) {
321+
const content = await fs.readFileText(file)
322+
const parsed = JSON.parse(content)
323+
assert.ok(parsed.test === 'registration' || parsed.test === 'token')
324+
}
316325
}
317326
})
318327

@@ -351,6 +360,10 @@ describe('AuthUtil', async function () {
351360
}
352361
memento.get.returns(mockProfiles)
353362

363+
if (!(auth as any).session){
364+
auth.session = new auth2.SsoLogin(auth.profileName, auth.lspAuth, auth.eventEmitter)
365+
}
366+
354367
const updateProfileStub = sinon.stub((auth as any).session, 'updateProfile').resolves()
355368

356369
await auth.migrateSsoConnectionToLsp('test-client')
@@ -376,6 +389,10 @@ describe('AuthUtil', async function () {
376389
}
377390
memento.get.returns(mockProfiles)
378391

392+
if (!(auth as any).session){
393+
auth.session = new auth2.SsoLogin(auth.profileName, auth.lspAuth, auth.eventEmitter)
394+
}
395+
379396
const updateProfileStub = sinon.stub((auth as any).session, 'updateProfile').resolves()
380397

381398
await auth.migrateSsoConnectionToLsp('test-client')
@@ -389,4 +406,4 @@ describe('AuthUtil', async function () {
389406
)
390407
})
391408
})
392-
})
409+
})

‎packages/amazonq/test/unit/codewhisperer/util/showSsoPrompt.test.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ describe('showConnectionPrompt', function () {
2828
})
2929

3030
it('can select connect to AwsBuilderId', async function () {
31-
sinon.stub(AuthUtil.instance, 'login').resolves()
31+
sinon.stub(AuthUtil.instance, 'login_sso').resolves()
3232

3333
getTestWindow().onDidShowQuickPick(async (picker) => {
3434
await picker.untilReady()
@@ -44,7 +44,7 @@ describe('showConnectionPrompt', function () {
4444

4545
it('connectToAwsBuilderId calls AuthUtil login with builderIdStartUrl', async function () {
4646
sinon.stub(vscode.commands, 'executeCommand')
47-
const loginStub = sinon.stub(AuthUtil.instance, 'login').resolves()
47+
const loginStub = sinon.stub(AuthUtil.instance, 'login_sso').resolves()
4848

4949
await awsIdSignIn()
5050

‎packages/core/src/auth/auth2.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -283,6 +283,7 @@ export class LanguageClientAuth {
283283
* Abstract class for connection management
284284
*/
285285
export abstract class BaseLogin {
286+
protected loginType: LoginType | undefined
286287
protected connectionState: AuthState = 'notConnected'
287288
protected cancellationToken: CancellationTokenSource | undefined
288289
protected _data: { startUrl?: string; region?: string; accessKey?: string; secretKey?: string } | undefined
@@ -360,6 +361,7 @@ export abstract class BaseLogin {
360361
*/
361362
export class SsoLogin extends BaseLogin {
362363
// Cached information from the identity server for easy reference
364+
override readonly loginType = LoginTypes.SSO
363365
private ssoTokenId: string | undefined
364366

365367
constructor(profileName: string, lspAuth: LanguageClientAuth, eventEmitter: vscode.EventEmitter<AuthStateEvent>) {
@@ -501,6 +503,7 @@ export class SsoLogin extends BaseLogin {
501503
*/
502504
export class IamLogin extends BaseLogin {
503505
// Cached information from the identity server for easy reference
506+
override readonly loginType = LoginTypes.IAM
504507
// private iamCredentialId: string | undefined
505508

506509
constructor(profileName: string, lspAuth: LanguageClientAuth, eventEmitter: vscode.EventEmitter<AuthStateEvent>) {

‎packages/core/src/codewhisperer/ui/codeWhispererNodes.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -271,7 +271,7 @@ export function createSignIn(): DataQuickPickItem<'signIn'> {
271271
if (isWeb()) {
272272
// TODO: nkomonen, call a Command instead
273273
onClick = () => {
274-
void AuthUtil.instance.login(builderIdStartUrl, builderIdRegion, 'sso')
274+
void AuthUtil.instance.login_sso(builderIdStartUrl, builderIdRegion)
275275
}
276276
}
277277

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

Lines changed: 23 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ import { showAmazonQWalkthroughOnce } from '../../amazonq/onboardingPage/walkthr
3030
import { setContext } from '../../shared/vscode/setContext'
3131
import { openUrl } from '../../shared/utilities/vsCodeUtils'
3232
import { telemetry } from '../../shared/telemetry/telemetry'
33-
import { AuthStateEvent, cacheChangedEvent, LanguageClientAuth, Login, SsoLogin, IamLogin } from '../../auth/auth2'
33+
import { AuthStateEvent, cacheChangedEvent, LanguageClientAuth, Login, SsoLogin, IamLogin, LoginTypes } from '../../auth/auth2'
3434
import { builderIdStartUrl, internalStartUrl } from '../../auth/sso/constants'
3535
import { VSCODE_EXTENSION_ID } from '../../shared/extensions'
3636
import { RegionProfileManager } from '../region/regionProfileManager'
@@ -108,11 +108,11 @@ export class AuthUtil implements IAuthProvider {
108108
}
109109

110110
isSsoSession(): boolean {
111-
return this.session instanceof SsoLogin
111+
return this.session?.loginType === LoginTypes.SSO || this.session instanceof SsoLogin
112112
}
113113

114114
isIamSession(): boolean {
115-
return this.session instanceof IamLogin
115+
return this.session?.loginType === LoginTypes.IAM || this.session instanceof IamLogin
116116
}
117117

118118
/**
@@ -134,7 +134,7 @@ export class AuthUtil implements IAuthProvider {
134134
if (!this.isConnected()) {
135135
// Try to restore an IAM session
136136
this.session = new IamLogin(this.profileName, this.lspAuth, this.eventEmitter)
137-
// await this.session.restore()
137+
await this.session.restore()
138138
if (!this.isConnected()) {
139139
// If both fail, reset the session
140140
this.session = undefined
@@ -160,25 +160,26 @@ export class AuthUtil implements IAuthProvider {
160160
}
161161
}
162162

163-
// Log into the desired session type using the authentication parameters
164-
async login(accessKey: string, secretKey: string, loginType: 'iam'): Promise<GetIamCredentialResult | undefined>
165-
async login(startUrl: string, region: string, loginType: 'sso'): Promise<GetSsoTokenResult | undefined>
166-
async login(
167-
first: string,
168-
second: string,
169-
loginType: 'iam' | 'sso'
170-
): Promise<GetSsoTokenResult | GetIamCredentialResult | undefined> {
171-
let response: GetSsoTokenResult | GetIamCredentialResult | undefined
172-
173-
// Start session if the current session type does not match the desired type
174-
if (loginType === 'sso' && !this.isSsoSession()) {
163+
// Log in using SSO
164+
async login_sso(startUrl: string, region: string): Promise<GetSsoTokenResult | undefined> {
165+
let response: GetSsoTokenResult | undefined
166+
// Create SSO login session
167+
if (!this.isSsoSession()) {
175168
this.session = new SsoLogin(this.profileName, this.lspAuth, this.eventEmitter)
176-
response = await this.session.login({ startUrl: first, region: second, scopes: amazonQScopes })
177-
} else if (loginType === 'iam' && !this.isIamSession()) {
178-
this.session = new IamLogin(this.profileName, this.lspAuth, this.eventEmitter)
179-
response = await this.session.login({ accessKey: first, secretKey: second })
180169
}
170+
response = await (this.session as SsoLogin).login({ startUrl: startUrl, region: region, scopes: amazonQScopes })
171+
await showAmazonQWalkthroughOnce()
172+
return response
173+
}
181174

175+
// Log in using IAM or STS credentials
176+
async login_iam(accessKey: string, secretKey: string, sessionToken?: string): Promise<GetIamCredentialResult | undefined> {
177+
let response: GetIamCredentialResult | undefined
178+
// Create IAM login session
179+
if (!this.isIamSession()) {
180+
this.session = new IamLogin(this.profileName, this.lspAuth, this.eventEmitter)
181+
}
182+
response = await (this.session as IamLogin).login({ accessKey: accessKey, secretKey: secretKey })
182183
await showAmazonQWalkthroughOnce()
183184
return response
184185
}
@@ -197,8 +198,8 @@ export class AuthUtil implements IAuthProvider {
197198
}
198199

199200
async getToken() {
200-
if (this.isSsoSession()) {
201-
const token = (await this.session!.getCredential()).credential
201+
if (this.session) {
202+
const token = (await this.session.getCredential()).credential
202203
if (typeof token !== 'string') {
203204
throw new ToolkitError('Cannot get token with IAM session')
204205
}

‎packages/core/src/codewhisperer/util/getStartUrl.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ export const getStartUrl = async () => {
2929

3030
export async function connectToEnterpriseSso(startUrl: string, region: Region['id']) {
3131
try {
32-
await AuthUtil.instance.login(startUrl, region, 'sso')
32+
await AuthUtil.instance.login_sso(startUrl, region)
3333
} catch (e) {
3434
throw ToolkitError.chain(e, CodeWhispererConstants.failedToConnectIamIdentityCenter, {
3535
code: 'FailedToConnect',

‎packages/core/src/codewhisperer/util/showSsoPrompt.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ export const showCodeWhispererConnectionPrompt = async () => {
4747
export async function awsIdSignIn() {
4848
getLogger().info('selected AWS ID sign in')
4949
try {
50-
await AuthUtil.instance.login(builderIdStartUrl, builderIdRegion, 'sso')
50+
await AuthUtil.instance.login_sso(builderIdStartUrl, builderIdRegion)
5151
} catch (e) {
5252
throw ToolkitError.chain(e, failedToConnectAwsBuilderId, { code: 'FailedToConnect' })
5353
}

0 commit comments

Comments
 (0)