Skip to content

Commit a21572d

Browse files
yuxianrzyuxianrz
committed
fix: sync with IAM PR changes
1 parent fc740fa commit a21572d

File tree

14 files changed

+88
-116
lines changed

14 files changed

+88
-116
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: 64 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -26,36 +26,36 @@ 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

45-
it('identifies SSO session', function () {
46-
;(auth as any).session = { loginType: auth2.LoginTypes.SSO }
45+
it('identifies SSO session', async function () {
46+
await auth.loginSso(constants.internalStartUrl, 'us-east-1')
4747
assert.strictEqual(auth.isSsoSession(), true)
4848
})
4949

50-
it('identifies non-SSO session', function () {
51-
;(auth as any).session = { loginType: auth2.LoginTypes.IAM }
50+
it('identifies non-SSO session', async function () {
51+
await auth.loginIam('accessKey', 'secretKey', 'sessionToken')
5252
assert.strictEqual(auth.isSsoSession(), false)
5353
})
5454
})
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,23 @@ 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+
})
442+
)
441443
assert.strictEqual(response, mockResponse)
442444
})
443445

@@ -461,7 +463,7 @@ describe('AuthUtil', async function () {
461463

462464
sinon.stub(auth2, 'IamLogin').returns(mockIamLogin as any)
463465

464-
const response = await auth.login_iam('accessKey', 'secretKey', 'sessionToken', 'arn:aws:iam::123456789012:role/TestRole')
466+
const response = await auth.loginIam('accessKey', 'secretKey', 'sessionToken', 'arn:aws:iam::123456789012:role/TestRole')
465467

466468
assert.ok(mockIamLogin.login.calledOnce)
467469
assert.ok(mockIamLogin.login.calledWith({
@@ -481,19 +483,19 @@ describe('AuthUtil', async function () {
481483
secretAccessKey: 'test-secret-key',
482484
sessionToken: 'test-session-token',
483485
}
484-
486+
485487
const mockSession = {
486488
getCredential: sinon.stub().resolves({
487489
credential: mockCredentials,
488490
updateCredentialsParams: { data: 'test' },
489491
}),
490492
loginType: 'iam',
491493
}
492-
494+
493495
;(auth as any).session = mockSession
494-
496+
495497
const result = await auth.getIamCredential()
496-
498+
497499
assert.ok(mockSession.getCredential.calledOnce)
498500
assert.deepStrictEqual(result, mockCredentials)
499501
})
@@ -506,9 +508,9 @@ describe('AuthUtil', async function () {
506508
}),
507509
loginType: 'sso',
508510
}
509-
511+
510512
;(auth as any).session = mockSession
511-
513+
512514
try {
513515
await auth.getIamCredential()
514516
assert.fail('Should have thrown an error')
@@ -519,7 +521,7 @@ describe('AuthUtil', async function () {
519521

520522
it('throws error when not logged in', async function () {
521523
;(auth as any).session = undefined
522-
524+
523525
try {
524526
await auth.getIamCredential()
525527
assert.fail('Should have thrown an error')
@@ -531,22 +533,22 @@ describe('AuthUtil', async function () {
531533

532534
describe('isIamSession', function () {
533535
it('returns true for IAM session', function () {
534-
const mockSession = { loginType: 'iam' }
536+
const mockSession = new auth2.IamLogin(auth.profileName, auth.lspAuth, auth.eventEmitter)
535537
;(auth as any).session = mockSession
536-
538+
537539
assert.strictEqual(auth.isIamSession(), true)
538540
})
539541

540542
it('returns false for SSO session', function () {
541543
const mockSession = { loginType: 'sso' }
542544
;(auth as any).session = mockSession
543-
545+
544546
assert.strictEqual(auth.isIamSession(), false)
545547
})
546548

547549
it('returns false when no session', function () {
548550
;(auth as any).session = undefined
549-
551+
550552
assert.strictEqual(auth.isIamSession(), false)
551553
})
552554
})
@@ -559,37 +561,17 @@ describe('AuthUtil', async function () {
559561
})
560562

561563
it('updates IAM credential when state is refreshed', async function () {
562-
const mockSession = {
563-
getCredential: sinon.stub().resolves({
564-
credential: { accessKeyId: 'key', secretAccessKey: 'secret' },
565-
updateCredentialsParams: { data: 'fake-data' },
566-
}),
567-
loginType: 'iam',
568-
}
564+
const mockSession = new auth2.IamLogin(auth.profileName, auth.lspAuth, auth.eventEmitter)
565+
sinon.stub(mockSession, 'getCredential').resolves({
566+
credential: { accessKeyId: 'key', secretAccessKey: 'secret' },
567+
updateCredentialsParams: { data: 'fake-data' },
568+
})
569569
;(auth as any).session = mockSession
570570

571571
await (auth as any).stateChangeHandler({ state: 'refreshed' })
572572

573573
assert.ok(mockLspAuth.updateIamCredential.called)
574574
assert.strictEqual(mockLspAuth.updateIamCredential.firstCall.args[0].data, 'fake-data')
575575
})
576-
577-
it('cleans up IAM credential when connection expires', async function () {
578-
const mockSession = { loginType: 'iam' }
579-
;(auth as any).session = mockSession
580-
581-
await (auth as any).stateChangeHandler({ state: 'expired' })
582-
583-
assert.ok(mockLspAuth.deleteIamCredential.called)
584-
})
585-
586-
it('deletes IAM credential when disconnected', async function () {
587-
const mockSession = { loginType: 'iam' }
588-
;(auth as any).session = mockSession
589-
590-
await (auth as any).stateChangeHandler({ state: 'notConnected' })
591-
592-
assert.ok(mockLspAuth.deleteIamCredential.called)
593-
})
594576
})
595577
})

‎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

0 commit comments

Comments
 (0)