feat(amazonq): Add command to generate unit tests for selected code

plugins/amazonq/chat/jetbrains-community/resources/META-INF/plugin-chat.xml

@@ -7,7 +7,7 @@
 
     <projectListeners>
         <listener class="software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AmazonQToolWindowListener"
-                  topic="com.intellij.openapi.wm.ex.ToolWindowManagerListener"/>   
+                  topic="com.intellij.openapi.wm.ex.ToolWindowManagerListener"/>
         <listener class="software.aws.toolkits.jetbrains.services.cwc.editor.context.project.ProjectContextEditorListener"
                   topic="com.intellij.openapi.fileEditor.FileEditorManagerListener"/>
     </projectListeners>
@@ -70,6 +70,11 @@
                 <keyboard-shortcut keymap="$default" first-keystroke="meta alt A" />
             </action>
 
+            <action id="aws.toolkit.jetbrains.core.services.cwc.commands.GenerateUnitTestsAction"
+                    class="software.aws.toolkits.jetbrains.services.cwc.commands.GenerateUnitTestsAction">
+                <keyboard-shortcut keymap="$default" first-keystroke="meta alt T" />
+            </action>
+
             <action id="aws.toolkit.jetbrains.core.services.cwc.commands.SendToPromptAction"
                     class="software.aws.toolkits.jetbrains.services.cwc.commands.SendToPromptAction">
                 <keyboard-shortcut keymap="$default" first-keystroke="meta alt S" />

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/clients/chat/v1/ChatSessionV1.kt

@@ -297,6 +297,7 @@ class ChatSessionV1(
         UserIntent.EXPLAIN_LINE_BY_LINE -> FollowUpType.LineByLine
         UserIntent.EXPLAIN_CODE_SELECTION -> FollowUpType.ExplainInDetail
         UserIntent.UNKNOWN_TO_SDK_VERSION -> FollowUpType.Generated
+        UserIntent.GENERATE_UNIT_TESTS -> FollowUpType.Generated
         null -> FollowUpType.Generated
     }
 

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/commands/EditorContextCommand.kt

@@ -26,6 +26,10 @@ enum class EditorContextCommand(
         verb = "Optimize",
         actionId = "aws.amazonq.optimizeCode",
     ),
+    GenerateUnitTests(
+        verb = "Generate unit tests for",
+        actionId = "aws.amazonq.generateUnitTests",
+    ),
     SendToPrompt(
         verb = "SendToPrompt",
         actionId = "aws.amazonq.sendToPrompt",

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/commands/GenerateUnitTestsAction.kt

@@ -0,0 +1,21 @@
+// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.cwc.commands
+
+import com.intellij.openapi.actionSystem.ActionUpdateThread
+import com.intellij.openapi.actionSystem.AnActionEvent
+import com.intellij.openapi.actionSystem.CommonDataKeys
+import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
+import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
+import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
+
+class GenerateUnitTestsAction : CustomAction(EditorContextCommand.GenerateUnitTests) {
+    override fun getActionUpdateThread() = ActionUpdateThread.BGT
+
+    override fun update(e: AnActionEvent) {
+        val project = e.getData(CommonDataKeys.PROJECT) ?: return
+        val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) as? AwsBearerTokenConnection
+        e.presentation.isEnabledAndVisible = connection?.startUrl == "https://amzn.awsapps.com/start"
+    }
+}

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ChatController.kt

@@ -130,6 +130,8 @@ class ChatController private constructor(
         var shouldAddIndexInProgressMessage: Boolean = false
         var shouldUseWorkspaceContext: Boolean = false
         val isDataCollectionGroup = CodeWhispererFeatureConfigService.getInstance().getIsDataCollectionEnabled()
+        val startUrl = getStartUrl(context.project)
+
         if (prompt.contains("@workspace")) {
             if (CodeWhispererSettings.getInstance().isProjectContextEnabled()) {
                 shouldUseWorkspaceContext = true
@@ -151,7 +153,7 @@ class ChatController private constructor(
             triggerId = triggerId,
             message = prompt,
             activeFileContext = contextExtractor.extractContextForTrigger(ExtractionTriggerType.ChatMessage),
-            userIntent = intentRecognizer.getUserIntentFromPromptChatMessage(message.chatMessage),
+            userIntent = intentRecognizer.getUserIntentFromPromptChatMessage(message.chatMessage, startUrl),
             TriggerType.Click,
             projectContextQueryResult = queryResult,
             shouldAddIndexInProgressMessage = shouldAddIndexInProgressMessage,
@@ -332,7 +334,11 @@ class ChatController private constructor(
         }
 
         // Create prompt
-        val prompt = "${message.command} the following part of my code for me: $codeSelection"
+        val prompt = if (EditorContextCommand.GenerateUnitTests == message.command) {
+            "${message.command.verb} the following part of my code for me: $codeSelection"
+        } else {
+            "${message.command} the following part of my code for me: $codeSelection"
+        }
 
         processPromptActions(prompt, message, triggerId, fileContext)
     }

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/chat/telemetry/TelemetryHelper.kt

@@ -58,6 +58,7 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
         UserIntent.CITE_SOURCES -> CwsprChatUserIntent.CiteSources
         UserIntent.EXPLAIN_LINE_BY_LINE -> CwsprChatUserIntent.ExplainLineByLine
         UserIntent.EXPLAIN_CODE_SELECTION -> CwsprChatUserIntent.ExplainCodeSelection
+        UserIntent.GENERATE_UNIT_TESTS -> CwsprChatUserIntent.GenerateUnitTests
         UserIntent.UNKNOWN_TO_SDK_VERSION -> CwsprChatUserIntent.Unknown
     }
 

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/chat/userIntent/UserIntentRecognizer.kt

@@ -16,14 +16,16 @@ class UserIntentRecognizer {
         EditorContextCommand.Fix -> UserIntent.APPLY_COMMON_BEST_PRACTICES
         EditorContextCommand.Optimize -> UserIntent.IMPROVE_CODE
         EditorContextCommand.ExplainCodeScanIssue -> UserIntent.EXPLAIN_CODE_SELECTION
+        EditorContextCommand.GenerateUnitTests -> UserIntent.GENERATE_UNIT_TESTS
         EditorContextCommand.SendToPrompt -> null
     }
 
-    fun getUserIntentFromPromptChatMessage(prompt: String) = when {
+    fun getUserIntentFromPromptChatMessage(prompt: String, startUrl: String?) = when {
         prompt.startsWith("Explain") -> UserIntent.EXPLAIN_CODE_SELECTION
         prompt.startsWith("Refactor") -> UserIntent.SUGGEST_ALTERNATE_IMPLEMENTATION
         prompt.startsWith("Fix") -> UserIntent.APPLY_COMMON_BEST_PRACTICES
         prompt.startsWith("Optimize") -> UserIntent.IMPROVE_CODE
+        prompt.startsWith("Generate unit tests") && isInternalAmazonUser(startUrl) -> UserIntent.GENERATE_UNIT_TESTS
         else -> null
     }
 
@@ -43,4 +45,6 @@ class UserIntentRecognizer {
     fun getUserIntentFromOnboardingPageInteraction(interaction: OnboardingPageInteraction) = when (interaction.type) {
         OnboardingPageInteractionType.CwcButtonClick -> null
     }
+
+    private fun isInternalAmazonUser(startUrl: String?): Boolean = startUrl == "https://amzn.awsapps.com/start"
 }

plugins/core/jetbrains-community/resources/telemetryOverride.json

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

plugins/core/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties

@@ -21,6 +21,8 @@ action.aws.toolkit.jetbrains.core.services.cwc.commands.ExplainCodeAction.descri
 action.aws.toolkit.jetbrains.core.services.cwc.commands.ExplainCodeAction.text = Explain Code
 action.aws.toolkit.jetbrains.core.services.cwc.commands.FixCodeAction.description = Fixes the selected code
 action.aws.toolkit.jetbrains.core.services.cwc.commands.FixCodeAction.text = Fix Code
+action.aws.toolkit.jetbrains.core.services.cwc.commands.GenerateUnitTestsAction.description = Generates unit tests for the selected code
+action.aws.toolkit.jetbrains.core.services.cwc.commands.GenerateUnitTestsAction.text = Generate Tests (Beta)
 action.aws.toolkit.jetbrains.core.services.cwc.commands.OptimizeCodeAction.description = Optimizes the selected code
 action.aws.toolkit.jetbrains.core.services.cwc.commands.OptimizeCodeAction.text = Optimize Code
 action.aws.toolkit.jetbrains.core.services.cwc.commands.RefactorCodeAction.description = Refactors the selected code

plugins/core/sdk-codegen/codegen-resources/codewhispererruntime/service-2.json

@@ -1952,7 +1952,8 @@
                 "SHOW_EXAMPLES",
                 "CITE_SOURCES",
                 "EXPLAIN_LINE_BY_LINE",
-                "EXPLAIN_CODE_SELECTION"
+                "EXPLAIN_CODE_SELECTION",
+                "GENERATE_UNIT_TESTS"
             ]
         },
         "UserModificationEvent":{

plugins/core/sdk-codegen/codegen-resources/codewhispererstreaming/service-2.json

@@ -741,7 +741,8 @@
                 "SHOW_EXAMPLES",
                 "CITE_SOURCES",
                 "EXPLAIN_LINE_BY_LINE",
-                "EXPLAIN_CODE_SELECTION"
+                "EXPLAIN_CODE_SELECTION",
+                "GENERATE_UNIT_TESTS"
             ]
         },
         "ValidationException":{