feat(amazonq): user can generate unit tests for selected code #5577

- Add command to Amazon Q to generate unit tests for selected code. The
  feature will be limited to internal Amazon users, initially.
- Add "Generate Tests" command.

Author: bweedop

packages/amazonq/package.json

@@ -349,9 +349,14 @@
                     "command": "aws.amazonq.optimizeCode",
                     "group": "cw_chat@4"
                 },
+                {
+                    "command": "aws.amazonq.generateUnitTests",
+                    "group": "cw_chat@5",
+                    "when": "aws.codewhisperer.connected && aws.isInternalUser"
+                },
                 {
                     "command": "aws.amazonq.sendToPrompt",
-                    "group": "cw_chat@5"
+                    "group": "cw_chat@6"
                 }
             ],
             "editor/context": [
@@ -423,6 +428,12 @@
                 "category": "%AWS.amazonq.title%",
                 "enablement": "aws.codewhisperer.connected"
             },
+            {
+                "command": "aws.amazonq.generateUnitTests",
+                "title": "%AWS.command.amazonq.generateUnitTests%",
+                "category": "%AWS.amazonq.title%",
+                "enablement": "aws.codewhisperer.connected && aws.isInternalUser"
+            },
             {
                 "command": "aws.amazonq.reconnect",
                 "title": "%AWS.command.codewhisperer.reconnect%",
@@ -599,6 +610,13 @@
                 "mac": "cmd+alt+q",
                 "linux": "meta+alt+q"
             },
+            {
+                "command": "aws.amazonq.generateUnitTests",
+                "key": "win+alt+t",
+                "mac": "cmd+alt+t",
+                "linux": "meta+alt+t",
+                "when": "aws.codewhisperer.connected && aws.isInternalUser"
+            },
             {
                 "command": "aws.amazonq.invokeInlineCompletion",
                 "key": "alt+c",

packages/core/package.nls.json

@@ -111,6 +111,7 @@
     "AWS.command.amazonq.fixCode": "Fix",
     "AWS.command.amazonq.optimizeCode": "Optimize",
     "AWS.command.amazonq.sendToPrompt": "Send to prompt",
+    "AWS.command.amazonq.generateUnitTests": "Generate Tests (Beta)",
     "AWS.command.amazonq.security.scan": "Run Project Scan",
     "AWS.command.deploySamApplication": "Deploy SAM Application",
     "AWS.command.aboutToolkit": "About",

packages/core/src/auth/auth.ts

@@ -59,13 +59,15 @@ import {
     AwsConnection,
     scopesCodeWhispererCore,
     ProfileNotFoundError,
+    isSsoConnection,
 } from './connection'
 import { isSageMaker, isCloud9, isAmazonQ } from '../shared/extensionUtilities'
 import { telemetry } from '../shared/telemetry/telemetry'
 import { randomUUID } from '../shared/crypto'
 import { asStringifiedStack } from '../shared/telemetry/spans'
 import { withTelemetryContext } from '../shared/telemetry/util'
 import { DiskCacheError } from '../shared/utilities/cacheUtils'
+import { setContext } from '../shared/vscode/setContext'
 
 interface AuthService {
     /**
@@ -166,6 +168,30 @@ export class Auth implements AuthService, ConnectionManager {
         return this.#ssoCacheWatcher
     }
 
+    public get startUrl(): string | undefined {
+        return isSsoConnection(this.activeConnection)
+            ? this.normalizeStartUrl(this.activeConnection.startUrl)
+            : undefined
+    }
+
+    public isConnected(): boolean {
+        return this.activeConnection !== undefined
+    }
+
+    /**
+     * Normalizes the provided URL
+     *
+     *  Any trailing '/' and `#` is removed from the URL
+     *  e.g. https://view.awsapps.com/start/# will become https://view.awsapps.com/start
+     */
+    public normalizeStartUrl(startUrl: string | undefined) {
+        return !startUrl ? undefined : startUrl.replace(/[\/#]+$/g, '')
+    }
+
+    public isInternalAmazonUser(): boolean {
+        return this.isConnected() && this.startUrl === 'https://amzn.awsapps.com/start'
+    }
+
     /**
      * Map startUrl -> declared connections
      */
@@ -223,6 +249,8 @@ export class Auth implements AuthService, ConnectionManager {
         this.#onDidChangeActiveConnection.fire(conn)
         await this.store.setCurrentProfileId(id)
 
+        await setContext('aws.isInternalUser', this.isInternalAmazonUser())
+
         return conn
     }
 
@@ -373,6 +401,7 @@ export class Auth implements AuthService, ConnectionManager {
             }
         }
         this.#onDidDeleteConnection.fire({ connId, storedProfile: profile })
+        await setContext('aws.isInternalUser', false)
     }
 
     @withTelemetryContext({ name: 'clearStaleLinkedIamConnections', class: authClassName })
@@ -405,6 +434,7 @@ export class Auth implements AuthService, ConnectionManager {
         await provider.invalidate('devModeManualExpiration')
         // updates the state of the connection
         await this.refreshConnectionState(conn)
+        await setContext('aws.isInternalUser', false)
     }
 
     public async getConnection(connection: Pick<Connection, 'id'>): Promise<Connection | undefined> {

packages/core/src/codewhispererChat/commands/registerCommands.ts

@@ -94,6 +94,14 @@ export function registerCommands(controllerPublishers: ChatControllerMessagePubl
             })
         })
     })
+    Commands.register('aws.amazonq.generateUnitTests', async (data) => {
+        return focusAmazonQPanel.execute(placeholder, 'amazonq.generateUnitTests').then(() => {
+            controllerPublishers.processContextMenuCommand.publish({
+                type: 'aws.amazonq.generateUnitTests',
+                triggerType: getCommandTriggerType(data),
+            })
+        })
+    })
 }
 
 export type EditorContextBaseCommandType =
@@ -102,6 +110,7 @@ export type EditorContextBaseCommandType =
     | 'aws.amazonq.fixCode'
     | 'aws.amazonq.optimizeCode'
     | 'aws.amazonq.sendToPrompt'
+    | 'aws.amazonq.generateUnitTests'
 
 export type CodeScanIssueCommandType = 'aws.amazonq.explainIssue'
 

packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts

@@ -340,6 +340,7 @@ export class Messenger {
         ['aws.amazonq.fixCode', 'Fix'],
         ['aws.amazonq.optimizeCode', 'Optimize'],
         ['aws.amazonq.sendToPrompt', 'Send to prompt'],
+        ['aws.amazonq.generateUnitTests', 'Generate unit tests for'],
     ])
 
     public sendStaticTextResponse(type: StaticTextResponseType, triggerID: string, tabID: string) {

packages/core/src/codewhispererChat/controllers/chat/prompts/promptsGenerator.ts

@@ -14,6 +14,7 @@ export class PromptsGenerator {
         ['aws.amazonq.fixCode', 'Fix'],
         ['aws.amazonq.optimizeCode', 'Optimize'],
         ['aws.amazonq.sendToPrompt', 'Send to prompt'],
+        ['aws.amazonq.generateUnitTests', 'Generate unit tests for'],
     ])
 
     public generateForContextMenuCommand(command: EditorContextCommand): string {

packages/core/src/codewhispererChat/controllers/chat/telemetryHelper.ts

@@ -86,6 +86,8 @@ export class CWCTelemetryHelper {
                 return 'explainLineByLine'
             case UserIntent.SHOW_EXAMPLES:
                 return 'showExample'
+            case UserIntent.GENERATE_UNIT_TESTS:
+                return 'generateUnitTests'
             default:
                 return undefined
         }

packages/core/src/codewhispererChat/controllers/chat/userIntent/userIntentRecognizer.ts

@@ -6,6 +6,7 @@
 import { UserIntent } from '@amzn/codewhisperer-streaming'
 import { EditorContextCommand } from '../../../commands/registerCommands'
 import { PromptMessage } from '../model'
+import { Auth } from '../../../../auth'
 
 export class UserIntentRecognizer {
     public getFromContextMenuCommand(command: EditorContextCommand): UserIntent | undefined {
@@ -18,6 +19,8 @@ export class UserIntentRecognizer {
                 return UserIntent.APPLY_COMMON_BEST_PRACTICES
             case 'aws.amazonq.optimizeCode':
                 return UserIntent.IMPROVE_CODE
+            case 'aws.amazonq.generateUnitTests':
+                return UserIntent.GENERATE_UNIT_TESTS
             default:
                 return undefined
         }
@@ -36,6 +39,8 @@ export class UserIntentRecognizer {
             return UserIntent.APPLY_COMMON_BEST_PRACTICES
         } else if (prompt.message.startsWith('Optimize')) {
             return UserIntent.IMPROVE_CODE
+        } else if (prompt.message.startsWith('Generate unit tests') && Auth.instance.isInternalAmazonUser()) {
+            return UserIntent.GENERATE_UNIT_TESTS
         }
         return undefined
     }

packages/core/src/shared/telemetry/vscodeTelemetry.json

@@ -58,7 +58,8 @@
                 "showExample",
                 "citeSources",
                 "explainLineByLine",
-                "explainCodeSelection"
+                "explainCodeSelection",
+                "generateUnitTests"
             ],
             "description": "Explict user intent associated with a chat message"
         },

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

@@ -14,6 +14,7 @@ type contextKey =
     | 'aws.isDevMode'
     | 'aws.isSageMaker'
     | 'aws.isWebExtHost'
+    | 'aws.isInternalUser'
     | 'aws.amazonq.showLoginView'
     | 'aws.codecatalyst.connected'
     | 'aws.codewhisperer.connected'