Skip to content

Commit 3213793

Browse files
parameja1parameja1
committed
Add unit tests to validate sign out and chat intro message changes
Update changelog to add SageMaker Unified Studio features Address review comments and linting
1 parent e4a1f1c commit 3213793

File tree

9 files changed

+180
-10
lines changed

9 files changed

+180
-10
lines changed

‎packages/amazonq/.changes/next-release/Feature-390d2096-ea1e-43b9-a889-2494d6e9e6d3.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type": "Feature",
3+
"description": "SageMaker Unified Studio: Disable Sign out"
4+
}

‎packages/amazonq/.changes/next-release/Feature-3c42f632-46b8-4e59-b5fa-75d638f3f08d.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type": "Feature",
3+
"description": "SageMaker Unified Studio: Update Q Chat Introduction message"
4+
}

‎packages/amazonq/package.json‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -351,7 +351,7 @@
351351
},
352352
{
353353
"command": "aws.amazonq.signout",
354-
"when": "(view == aws.amazonq.AmazonQChatView) && aws.codewhisperer.connected && !aws.amazonq.isSagemakerStudio",
354+
"when": "(view == aws.amazonq.AmazonQChatView) && aws.codewhisperer.connected && !aws.amazonq.isSagemakerUnifiedStudio",
355355
"group": "2_amazonQ@4"
356356
},
357357
{
@@ -627,7 +627,7 @@
627627
"title": "%AWS.command.codewhisperer.signout%",
628628
"category": "%AWS.amazonq.title%",
629629
"icon": "$(debug-disconnect)",
630-
"enablement": "aws.codewhisperer.connected && !aws.amazonq.isSagemakerStudio"
630+
"enablement": "aws.codewhisperer.connected && !aws.amazonq.isSagemakerUnifiedStudio"
631631
},
632632
{
633633
"command": "aws.amazonq.learnMore",

‎packages/amazonq/src/app/chat/activation.ts‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,7 @@ export async function activate(context: ExtensionContext) {
1616
const appInitContext = amazonq.DefaultAmazonQAppInitContext.instance
1717

1818
registerApps(appInitContext, context)
19-
const serviceName = process.env.SERVICE_NAME ?? 'SageMakerUnifiedStudio' // Fallback service name to a generic name AmazonQ
20-
const amazonQState = amazonq.AmazonQState.initialize(serviceName)
21-
getLogger().info(`In Extension activation - q state is initialized ${amazonQState.serviceName}`)
22-
// Create a context key for SageMaker Studio state and set in context
23-
await setContext('aws.amazonq.isSagemakerStudio', amazonQState.isSageMakerUnifiedStudio())
19+
await setupAmazonQState()
2420

2521
const provider = new amazonq.AmazonQChatViewProvider(
2622
context,
@@ -68,6 +64,14 @@ export async function activate(context: ExtensionContext) {
6864

6965
void setupLsp()
7066
void setupAuthNotification()
67+
68+
async function setupAmazonQState() {
69+
const serviceName = process.env.SERVICE_NAME ?? 'AmazonQ' // Fallback service name to a generic name AmazonQ
70+
const amazonQState = amazonq.AmazonQState.initialize(serviceName)
71+
getLogger().info(`In Extension activation - q state is initialized ${amazonQState.serviceName}`)
72+
// Create a context key for SageMaker Studio state and set in context
73+
await setContext('aws.amazonq.isSagemakerUnifiedStudio', amazonQState.isSageMakerUnifiedStudio())
74+
}
7175
}
7276

7377
function registerApps(appInitContext: amazonq.AmazonQAppInitContext, context: ExtensionContext) {

‎packages/core/src/amazonq/util/amazonQState.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,10 @@
1+
/**
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*
5+
* Manages the state of Amazon Q service in the extension
6+
* using Singleton pattern.
7+
*/
18
export class AmazonQState {
29
private static _instance: AmazonQState | undefined
310
private readonly _serviceName: string

‎packages/core/src/amazonq/webview/ui/tabs/generator.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ export class TabDataGenerator {
5353
return {}
5454
}
5555

56-
const isSMUSProTier = serviceName ? serviceName === 'SageMakerUnifiedStudio' : false
56+
const isSMUS = serviceName ? serviceName === 'SageMakerUnifiedStudio' : false
5757

5858
const tabData: MynahUIDataModel = {
5959
tabTitle: taskName ?? TabTypeDataMap[tabType].title,
@@ -66,7 +66,7 @@ export class TabDataGenerator {
6666
? [
6767
{
6868
type: ChatItemType.ANSWER,
69-
body: isSMUSProTier ? qChatIntroMessageForSMUS : TabTypeDataMap[tabType].welcome,
69+
body: isSMUS ? qChatIntroMessageForSMUS : TabTypeDataMap[tabType].welcome,
7070
},
7171
{
7272
type: ChatItemType.ANSWER,

‎packages/core/src/shared/vscode/setContext.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ export type contextKey =
3737
| 'gumby.wasQCodeTransformationUsed'
3838
| 'amazonq.inline.codelensShortcutEnabled'
3939
| 'aws.toolkit.lambda.walkthroughSelected'
40-
| 'aws.amazonq.isSagemakerStudio'
40+
| 'aws.amazonq.isSagemakerUnifiedStudio'
4141

4242
const contextMap: Partial<Record<contextKey, any>> = {}
4343

‎packages/core/src/test/amazonq/webview/ui/tabs/generator.test.ts‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
/*!
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
import assert from 'assert'
6+
import { ChatItemType } from '@aws/mynah-ui'
7+
import { TabDataGenerator } from '../../../../../amazonq/webview/ui/tabs/generator'
8+
import { TabType } from '../../../../../amazonq/webview/ui/storages/tabsStorage'
9+
import { qChatIntroMessageForSMUS } from '../../../../../amazonq/webview/ui/tabs/constants'
10+
11+
describe('TabDataGenerator', () => {
12+
let tabGenerator: TabDataGenerator
13+
14+
const defaultProps = {
15+
isFeatureDevEnabled: false,
16+
isGumbyEnabled: false,
17+
isScanEnabled: false,
18+
isTestEnabled: false,
19+
isDocEnabled: false,
20+
isChatEnabled: true,
21+
isAuthEnabled: true,
22+
}
23+
24+
beforeEach(() => {
25+
tabGenerator = new TabDataGenerator(defaultProps)
26+
})
27+
28+
describe('getTabData tests', () => {
29+
it('returns empty object for welcome tab', () => {
30+
const result = tabGenerator.getTabData('welcome', true)
31+
assert.deepStrictEqual(result, {})
32+
})
33+
34+
it('returns SMUS intro message when serviceName is SageMakerUnifiedStudio', () => {
35+
const result = tabGenerator.getTabData('cwc', true, 'TestTask', 'SageMakerUnifiedStudio')
36+
37+
assert.strictEqual(result.chatItems?.length, 2)
38+
assert.deepStrictEqual(result.chatItems?.[0], {
39+
type: ChatItemType.ANSWER,
40+
body: qChatIntroMessageForSMUS,
41+
})
42+
})
43+
44+
it('returns default intro message for non-SMUS service', () => {
45+
const result = tabGenerator.getTabData('cwc', true, 'TestTask', 'OtherService')
46+
47+
assert.strictEqual(result.chatItems?.length, 2)
48+
assert.ok(result.chatItems?.[0].body)
49+
assert.ok(!result.chatItems?.[0].body.includes(qChatIntroMessageForSMUS))
50+
})
51+
52+
it('returns no chat items when needWelcomeMessages is false', () => {
53+
const result = tabGenerator.getTabData('cwc', false, 'TestTask', 'SageMakerUnifiedStudio')
54+
55+
assert.strictEqual(result.chatItems?.length, 0)
56+
})
57+
58+
it('returns correct data structure for cwc tab type', () => {
59+
const tabType: TabType = 'cwc'
60+
const result = tabGenerator.getTabData(tabType, true)
61+
62+
assert.ok(typeof result.tabTitle === 'string')
63+
assert.ok(result.promptInputInfo !== undefined)
64+
if (result.promptInputInfo) {
65+
assert.ok(result.promptInputInfo.includes('Amazon Q'))
66+
}
67+
assert.ok(Array.isArray(result.quickActionCommands))
68+
assert.ok(typeof result.promptInputPlaceholder === 'string')
69+
assert.ok(Array.isArray(result.contextCommands))
70+
assert.ok(result.chatItems?.some((item) => item.type === ChatItemType.ANSWER))
71+
})
72+
})
73+
})

‎packages/core/src/test/codewhisperer/commands/basicCommands.test.ts‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ import { CodeWhispererSettings } from '../../../codewhisperer/util/codewhisperer
6868
import { confirm } from '../../../shared'
6969
import * as commentUtils from '../../../shared/utilities/commentUtils'
7070
import * as startCodeFixGeneration from '../../../codewhisperer/commands/startCodeFixGeneration'
71+
import { AmazonQState } from '../../../amazonq/util/amazonQState'
7172

7273
describe('CodeWhisperer-basicCommands', function () {
7374
let targetCommand: Command<any> & vscode.Disposable
@@ -520,6 +521,83 @@ describe('CodeWhisperer-basicCommands', function () {
520521
})
521522
await listCodeWhispererCommands.execute()
522523
})
524+
525+
it('excludes sign out when in SageMaker Unified Studio', async function () {
526+
sinon.stub(AuthUtil.instance, 'isConnected').returns(true)
527+
sinon.stub(AuthUtil.instance, 'isConnectionExpired').returns(false)
528+
sinon.stub(AmazonQState.instance, 'isSageMakerUnifiedStudio').get(() => true)
529+
530+
getTestWindow().onDidShowQuickPick((e) => {
531+
e.assertContainsItems(
532+
createAutoSuggestions(true),
533+
createOpenReferenceLog(),
534+
createGettingStarted(),
535+
switchToAmazonQNode(),
536+
...genericItems(),
537+
createSettingsNode()
538+
)
539+
e.dispose()
540+
})
541+
542+
await listCodeWhispererCommands.execute()
543+
})
544+
545+
it('includes sign out when authenticated and not in SageMaker Unified Studio', async function () {
546+
sinon.stub(AuthUtil.instance, 'isConnected').returns(true)
547+
sinon.stub(AuthUtil.instance, 'isConnectionExpired').returns(false)
548+
sinon.stub(AmazonQState.instance, 'isSageMakerUnifiedStudio').get(() => false)
549+
await CodeScansState.instance.setScansEnabled(false)
550+
551+
getTestWindow().onDidShowQuickPick((e) => {
552+
e.assertContainsItems(
553+
createAutoSuggestions(true),
554+
createOpenReferenceLog(),
555+
createGettingStarted(),
556+
switchToAmazonQNode(),
557+
createAutoScans(false),
558+
...genericItems(),
559+
createSettingsNode(),
560+
createSignout()
561+
)
562+
e.dispose()
563+
})
564+
565+
await listCodeWhispererCommands.execute()
566+
})
567+
568+
it('excludes sign out when using vended credentials in SageMaker Unified Studio', async function () {
569+
sinon.stub(AuthUtil.instance, 'isConnected').returns(true)
570+
sinon.stub(AuthUtil.instance, 'isConnectionExpired').returns(false)
571+
sinon.stub(AmazonQState.instance, 'isSageMakerUnifiedStudio').get(() => true)
572+
sinon.stub(AuthUtil.instance, 'isBuilderIdInUse').returns(true)
573+
574+
getTestWindow().onDidShowQuickPick((e) => {
575+
e.assertContainsItems(
576+
createAutoSuggestions(true),
577+
createOpenReferenceLog(),
578+
createGettingStarted(),
579+
switchToAmazonQNode(),
580+
...genericItems(),
581+
createSettingsNode()
582+
)
583+
e.dispose()
584+
})
585+
586+
await listCodeWhispererCommands.execute()
587+
})
588+
589+
it('shows expected items when connection is expired in SageMaker Unified Studio', async function () {
590+
sinon.stub(AuthUtil.instance, 'isConnected').returns(true)
591+
sinon.stub(AuthUtil.instance, 'isConnectionExpired').returns(true)
592+
sinon.stub(AmazonQState.instance, 'isSageMakerUnifiedStudio').get(() => true)
593+
594+
getTestWindow().onDidShowQuickPick((e) => {
595+
e.assertContainsItems(createReconnect(), createLearnMore(), ...genericItems())
596+
e.dispose()
597+
})
598+
599+
await listCodeWhispererCommands.execute()
600+
})
523601
})
524602

525603
describe('applySecurityFix', function () {

0 commit comments

Comments
 (0)