Skip to content

Commit c4acf81

Browse files
yuxianrzyuxianrz
committed
change login_sso and login_iam, fix auth2.ts
1 parent 1b28bf8 commit c4acf81

File tree

10 files changed

+86
-79
lines changed

10 files changed

+86
-79
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_sso(startUrl, region)
25+
await AuthUtil.instance.loginSso(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_sso(constants.builderIdStartUrl, region)
29+
await AuthUtil.instance.loginSso(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_sso(enterpriseSsoStartUrl, region)
33+
await AuthUtil.instance.loginSso(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: 74 additions & 63 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_sso(constants.builderIdStartUrl, constants.builderIdRegion)
29+
await auth.loginSso(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_sso('https://example.awsapps.com/start', 'us-east-1')
35+
await auth.loginSso('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_sso(constants.internalStartUrl, 'us-east-1')
41+
await auth.loginSso(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_sso(constants.builderIdStartUrl, constants.builderIdRegion)
58+
await auth.loginSso(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_sso(constants.builderIdStartUrl, constants.builderIdRegion)
71+
await auth.loginSso(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_sso('https://example.awsapps.com/start', 'us-east-1')
78+
await auth.loginSso('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,39 @@ describe('AuthUtil', async function () {
9696
})
9797

9898
it('returns BuilderId forms when using BuilderId', async function () {
99-
await auth.login_sso(constants.builderIdStartUrl, constants.builderIdRegion)
99+
await auth.loginSso(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-
session && sinon.stub(session, 'getProfile').resolves({
107-
ssoSession: {
108-
settings: {
109-
sso_registration_scopes: ['codewhisperer:*'],
106+
session &&
107+
sinon.stub(session, 'getProfile').resolves({
108+
ssoSession: {
109+
settings: {
110+
sso_registration_scopes: ['codewhisperer:*'],
111+
},
110112
},
111-
},
112-
})
113+
})
113114

114-
await auth.login_sso('https://example.awsapps.com/start', 'us-east-1')
115+
await auth.loginSso('https://example.awsapps.com/start', 'us-east-1')
115116
const forms = await auth.getAuthFormIds()
116117
assert.deepStrictEqual(forms, ['identityCenterCodeWhisperer'])
117118
})
118119

119120
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')
121+
await auth.loginSso('https://example.awsapps.com/start', 'us-east-1')
121122
const session = (auth as any).session
122123

123-
session && sinon.stub(session, 'getProfile').resolves({
124-
ssoSession: {
125-
settings: {
126-
sso_registration_scopes: ['codewhisperer:*', 'sso:account:access'],
124+
session &&
125+
sinon.stub(session, 'getProfile').resolves({
126+
ssoSession: {
127+
settings: {
128+
sso_registration_scopes: ['codewhisperer:*', 'sso:account:access'],
129+
},
127130
},
128-
},
129-
})
131+
})
130132

131133
const forms = await auth.getAuthFormIds()
132134
assert.deepStrictEqual(forms.sort(), ['identityCenterCodeWhisperer', 'identityCenterExplorer'].sort())
@@ -180,7 +182,7 @@ describe('AuthUtil', async function () {
180182
})
181183

182184
it('updates bearer token when state is refreshed', async function () {
183-
await auth.login_sso(constants.builderIdStartUrl, 'us-east-1')
185+
await auth.loginSso(constants.builderIdStartUrl, 'us-east-1')
184186

185187
await (auth as any).stateChangeHandler({ state: 'refreshed' })
186188

@@ -189,7 +191,7 @@ describe('AuthUtil', async function () {
189191
})
190192

191193
it('cleans up when connection expires', async function () {
192-
await auth.login_sso(constants.builderIdStartUrl, 'us-east-1')
194+
await auth.loginSso(constants.builderIdStartUrl, 'us-east-1')
193195

194196
await (auth as any).stateChangeHandler({ state: 'expired' })
195197

@@ -199,15 +201,15 @@ describe('AuthUtil', async function () {
199201
it('deletes bearer token when disconnected', async function () {
200202
await (auth as any).stateChangeHandler({ state: 'notConnected' })
201203

202-
if (auth.isSsoSession(auth.session)){
204+
if (auth.isSsoSession(auth.session)) {
203205
assert.ok(mockLspAuth.deleteBearerToken.called)
204206
}
205207
})
206208

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

210-
await auth.login_sso('https://example.awsapps.com/start', 'us-east-1')
212+
await auth.loginSso('https://example.awsapps.com/start', 'us-east-1')
211213

212214
await (auth as any).stateChangeHandler({ state: 'connected' })
213215

@@ -219,7 +221,7 @@ describe('AuthUtil', async function () {
219221
const invalidateProfileSpy = sinon.spy(regionProfileManager, 'invalidateProfile')
220222
const clearCacheSpy = sinon.spy(regionProfileManager, 'clearCache')
221223

222-
await auth.login_sso('https://example.awsapps.com/start', 'us-east-1')
224+
await auth.loginSso('https://example.awsapps.com/start', 'us-east-1')
223225

224226
await (auth as any).stateChangeHandler({ state: 'expired' })
225227

@@ -291,7 +293,7 @@ describe('AuthUtil', async function () {
291293
memento.get.returns({ profile1: validProfile })
292294
mockLspAuth.getSsoToken.rejects(new Error('Token check failed'))
293295

294-
if (!(auth as any).session){
296+
if (!(auth as any).session) {
295297
auth.session = new auth2.SsoLogin(auth.profileName, auth.lspAuth, auth.eventEmitter)
296298
}
297299
const updateProfileStub = sinon.stub((auth as any).session, 'updateProfile').resolves()
@@ -361,7 +363,7 @@ describe('AuthUtil', async function () {
361363
}
362364
memento.get.returns(mockProfiles)
363365

364-
if (!(auth as any).session){
366+
if (!(auth as any).session) {
365367
auth.session = new auth2.SsoLogin(auth.profileName, auth.lspAuth, auth.eventEmitter)
366368
}
367369

@@ -390,7 +392,7 @@ describe('AuthUtil', async function () {
390392
}
391393
memento.get.returns(mockProfiles)
392394

393-
if (!(auth as any).session){
395+
if (!(auth as any).session) {
394396
auth.session = new auth2.SsoLogin(auth.profileName, auth.lspAuth, auth.eventEmitter)
395397
}
396398

@@ -408,7 +410,7 @@ describe('AuthUtil', async function () {
408410
})
409411
})
410412

411-
describe('login_iam', function () {
413+
describe('loginIam', function () {
412414
it('creates IAM session and logs in', async function () {
413415
const mockResponse = {
414416
id: 'test-credential-id',
@@ -421,23 +423,25 @@ describe('AuthUtil', async function () {
421423
data: 'credential-data',
422424
},
423425
}
424-
426+
425427
const mockIamLogin = {
426428
login: sinon.stub().resolves(mockResponse),
427429
loginType: 'iam',
428430
}
429-
431+
430432
sinon.stub(auth2, 'IamLogin').returns(mockIamLogin as any)
431-
432-
const response = await auth.login_iam('accessKey', 'secretKey', 'sessionToken')
433-
433+
434+
const response = await auth.loginIam('accessKey', 'secretKey', 'sessionToken')
435+
434436
assert.ok(mockIamLogin.login.calledOnce)
435-
assert.ok(mockIamLogin.login.calledWith({
436-
accessKey: 'accessKey',
437-
secretKey: 'secretKey',
438-
sessionToken: 'sessionToken',
439-
roleArn: undefined,
440-
}))
437+
assert.ok(
438+
mockIamLogin.login.calledWith({
439+
accessKey: 'accessKey',
440+
secretKey: 'secretKey',
441+
sessionToken: 'sessionToken',
442+
roleArn: undefined,
443+
})
444+
)
441445
assert.strictEqual(response, mockResponse)
442446
})
443447

@@ -453,23 +457,30 @@ describe('AuthUtil', async function () {
453457
data: 'credential-data',
454458
},
455459
}
456-
460+
457461
const mockIamLogin = {
458462
login: sinon.stub().resolves(mockResponse),
459463
loginType: 'iam',
460464
}
461-
465+
462466
sinon.stub(auth2, 'IamLogin').returns(mockIamLogin as any)
463-
464-
const response = await auth.login_iam('accessKey', 'secretKey', 'sessionToken', 'arn:aws:iam::123456789012:role/TestRole')
465-
467+
468+
const response = await auth.loginIam(
469+
'accessKey',
470+
'secretKey',
471+
'sessionToken',
472+
'arn:aws:iam::123456789012:role/TestRole'
473+
)
474+
466475
assert.ok(mockIamLogin.login.calledOnce)
467-
assert.ok(mockIamLogin.login.calledWith({
468-
accessKey: 'accessKey',
469-
secretKey: 'secretKey',
470-
sessionToken: 'sessionToken',
471-
roleArn: 'arn:aws:iam::123456789012:role/TestRole',
472-
}))
476+
assert.ok(
477+
mockIamLogin.login.calledWith({
478+
accessKey: 'accessKey',
479+
secretKey: 'secretKey',
480+
sessionToken: 'sessionToken',
481+
roleArn: 'arn:aws:iam::123456789012:role/TestRole',
482+
})
483+
)
473484
assert.strictEqual(response, mockResponse)
474485
})
475486
})
@@ -481,19 +492,19 @@ describe('AuthUtil', async function () {
481492
secretAccessKey: 'test-secret-key',
482493
sessionToken: 'test-session-token',
483494
}
484-
495+
485496
const mockSession = {
486497
getCredential: sinon.stub().resolves({
487498
credential: mockCredentials,
488499
updateCredentialsParams: { data: 'test' },
489500
}),
490501
loginType: 'iam',
491502
}
492-
503+
493504
;(auth as any).session = mockSession
494-
505+
495506
const result = await auth.getIamCredential()
496-
507+
497508
assert.ok(mockSession.getCredential.calledOnce)
498509
assert.deepStrictEqual(result, mockCredentials)
499510
})
@@ -506,9 +517,9 @@ describe('AuthUtil', async function () {
506517
}),
507518
loginType: 'sso',
508519
}
509-
520+
510521
;(auth as any).session = mockSession
511-
522+
512523
try {
513524
await auth.getIamCredential()
514525
assert.fail('Should have thrown an error')
@@ -519,7 +530,7 @@ describe('AuthUtil', async function () {
519530

520531
it('throws error when not logged in', async function () {
521532
;(auth as any).session = undefined
522-
533+
523534
try {
524535
await auth.getIamCredential()
525536
assert.fail('Should have thrown an error')
@@ -533,20 +544,20 @@ describe('AuthUtil', async function () {
533544
it('returns true for IAM session', function () {
534545
const mockSession = { loginType: 'iam' }
535546
;(auth as any).session = mockSession
536-
547+
537548
assert.strictEqual(auth.isIamSession(), true)
538549
})
539550

540551
it('returns false for SSO session', function () {
541552
const mockSession = { loginType: 'sso' }
542553
;(auth as any).session = mockSession
543-
554+
544555
assert.strictEqual(auth.isIamSession(), false)
545556
})
546557

547558
it('returns false when no session', function () {
548559
;(auth as any).session = undefined
549-
560+
550561
assert.strictEqual(auth.isIamSession(), false)
551562
})
552563
})
@@ -592,4 +603,4 @@ describe('AuthUtil', async function () {
592603
assert.ok(mockLspAuth.deleteIamCredential.called)
593604
})
594605
})
595-
})
606+
})

‎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_sso').resolves()
31+
sinon.stub(AuthUtil.instance, 'loginSso').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_sso').resolves()
47+
const loginStub = sinon.stub(AuthUtil.instance, 'loginSso').resolves()
4848

4949
await awsIdSignIn()
5050

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

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@ export class LanguageClientAuth {
193193
): Promise<UpdateProfileResult> {
194194
// Add credentials and delete SSO settings from profile
195195
let profile: Profile
196-
if (roleArn) {
196+
if (roleArn && sourceProfile) {
197197
profile = {
198198
kinds: [ProfileKind.IamSourceProfileProfile],
199199
name: profileName,
@@ -234,10 +234,6 @@ export class LanguageClientAuth {
234234
}
235235
return this.client.sendRequest(updateProfileRequestType.method, {
236236
profile: profile,
237-
ssoSession: {
238-
name: profileName,
239-
settings: undefined,
240-
},
241237
} satisfies UpdateProfileParams)
242238
}
243239

0 commit comments

Comments
 (0)