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.TestCodeAction"
+                    class="software.aws.toolkits.jetbrains.services.cwc.commands.TestCodeAction">
+                <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",
     ),
+    Test(
+        verb = "Generate unit tests for",
+        actionId = "aws.amazonq.testCode",
+    ),
     SendToPrompt(
         verb = "SendToPrompt",
         actionId = "aws.amazonq.sendToPrompt",

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

@@ -0,0 +1,25 @@
+// 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.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
+import java.security.MessageDigest
+import kotlin.text.Charsets.UTF_8
+
+class TestCodeAction : CustomAction(EditorContextCommand.Test) {
+    override fun update(e: AnActionEvent) {
+        val project = e.getData(CommonDataKeys.PROJECT) ?: return
+        val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) as? AwsBearerTokenConnection
+        if (connection == null) {
+            // Hide the action by default if no connection found.
+            e.presentation.isEnabledAndVisible = false
+            return
+        }
+        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 @@
         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 @@
             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,9 @@
         }
 
         // Create prompt
-        val prompt = "${message.command} the following part of my code for me: $codeSelection"
+        val prompt = if (EditorContextCommand.Test == 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

@@ -57,6 +57,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

@@ -3,7 +3,11 @@
 
 package software.aws.toolkits.jetbrains.services.cwc.controller.chat.userIntent
 
+import com.intellij.openapi.actionSystem.CommonDataKeys
 import software.amazon.awssdk.services.codewhispererstreaming.model.UserIntent
+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
 import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteraction
 import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteractionType
 import software.aws.toolkits.jetbrains.services.cwc.clients.chat.model.FollowUpType
@@ -16,14 +20,16 @@
         EditorContextCommand.Fix -> UserIntent.APPLY_COMMON_BEST_PRACTICES
         EditorContextCommand.Optimize -> UserIntent.IMPROVE_CODE
         EditorContextCommand.ExplainCodeScanIssue -> UserIntent.EXPLAIN_CODE_SELECTION
+        EditorContextCommand.Test -> 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 +49,8 @@
     fun getUserIntentFromOnboardingPageInteraction(interaction: OnboardingPageInteraction) = when (interaction.type) {
         OnboardingPageInteractionType.CwcButtonClick -> null
     }
+
+    private fun isInternalAmazonUser(startUrl: String?): Boolean {
+        return 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

@@ -25,6 +25,8 @@ action.aws.toolkit.jetbrains.core.services.cwc.commands.OptimizeCodeAction.descr
 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
 action.aws.toolkit.jetbrains.core.services.cwc.commands.RefactorCodeAction.text = Refactor Code
+action.aws.toolkit.jetbrains.core.services.cwc.commands.TestCodeAction.description = Generates unit tests for the selected code
+action.aws.toolkit.jetbrains.core.services.cwc.commands.TestCodeAction.text = Test Code (Beta)
 action.aws.toolkit.jetbrains.core.services.cwc.commands.SendToPromptAction.description = Sends selected code to chat
 action.aws.toolkit.jetbrains.core.services.cwc.commands.SendToPromptAction.text = Send to Prompt
 action.aws.toolkit.open.arn.browser.text=Open ARN in AWS Console

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":{