From 2cf362a3e23e055edcf566b24dbc17ab3e5944e2 Mon Sep 17 00:00:00 2001
From: Bodie Weedop <weedop@amazon.com>
Date: Mon, 30 Sep 2024 11:08:45 -0700
Subject: [PATCH] Revert "revert(amazonq): Disable generate unit tests command
 and context menu for all users (#5674)"

This reverts commit c311da22f968103c986cc287b568b227a7a699e4.
---
 packages/amazonq/package.json                               | 6 +++---
 .../controllers/chat/userIntent/userIntentRecognizer.ts     | 3 +++
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/packages/amazonq/package.json b/packages/amazonq/package.json
index 0f60e5c36f3..5c98c6eaaf0 100644
--- a/packages/amazonq/package.json
+++ b/packages/amazonq/package.json
@@ -352,7 +352,7 @@
                 {
                     "command": "aws.amazonq.generateUnitTests",
                     "group": "cw_chat@5",
-                    "when": "false"
+                    "when": "aws.codewhisperer.connected && aws.isInternalUser"
                 },
                 {
                     "command": "aws.amazonq.sendToPrompt",
@@ -432,7 +432,7 @@
                 "command": "aws.amazonq.generateUnitTests",
                 "title": "%AWS.command.amazonq.generateUnitTests%",
                 "category": "%AWS.amazonq.title%",
-                "enablement": "false"
+                "enablement": "aws.codewhisperer.connected && aws.isInternalUser"
             },
             {
                 "command": "aws.amazonq.reconnect",
@@ -615,7 +615,7 @@
                 "key": "win+alt+t",
                 "mac": "cmd+alt+t",
                 "linux": "meta+alt+t",
-                "when": "false"
+                "when": "aws.codewhisperer.connected && aws.isInternalUser"
             },
             {
                 "command": "aws.amazonq.invokeInlineCompletion",
diff --git a/packages/core/src/codewhispererChat/controllers/chat/userIntent/userIntentRecognizer.ts b/packages/core/src/codewhispererChat/controllers/chat/userIntent/userIntentRecognizer.ts
index 058f5e8effa..32b6b274a51 100644
--- a/packages/core/src/codewhispererChat/controllers/chat/userIntent/userIntentRecognizer.ts
+++ b/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 {
@@ -38,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
     }