fix(amazonq): Fix Generate unit tests from right click context menu (#5705)

* initial changes

* fix gen unit tests

Author: manodnyab

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatController.kt

@@ -56,6 +56,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitConte
 import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthController
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionReference
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionReferencePosition
+import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage
 import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
 import software.aws.toolkits.jetbrains.services.amazonq.project.RelevantDocument
 import software.aws.toolkits.jetbrains.services.amazonqCodeTest.CodeWhispererUTGChatManager
@@ -110,6 +111,12 @@ import java.util.UUID
 import software.amazon.awssdk.services.codewhispererstreaming.model.Position as StreamingPosition
 import software.amazon.awssdk.services.codewhispererstreaming.model.Range as StreamingRange
 
+data class TestCommandMessage(
+    val sender: String = "codetest",
+    val command: String = "test",
+    val type: String = "addAnswer",
+) : AmazonQMessage
+
 class CodeTestChatController(
     private val context: AmazonQAppInitContext,
     private val chatSessionStorage: ChatSessionStorage,

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

@@ -1,8 +1,9 @@
 // Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-
+@file:Suppress("BannedImports")
 package software.aws.toolkits.jetbrains.services.cwc.commands
 
+import com.google.gson.Gson
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.Project
 import kotlinx.coroutines.flow.MutableSharedFlow
@@ -14,6 +15,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.Gener
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SendToPromptParams
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TriggerType
 import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage
+import software.aws.toolkits.jetbrains.services.amazonqCodeTest.controller.TestCommandMessage
 import software.aws.toolkits.jetbrains.services.cwc.editor.context.ActiveFileContextExtractor
 import software.aws.toolkits.jetbrains.services.cwc.editor.context.ExtractionTriggerType
 
@@ -25,8 +27,7 @@ class ActionRegistrar {
 
     fun reportMessageClick(command: EditorContextCommand, project: Project) {
         if (command == EditorContextCommand.GenerateUnitTests) {
-            // pre-existing old chat code path
-            _messages.tryEmit(ContextMenuActionMessage(command, project))
+            AsyncChatUiListener.notifyPartialMessageUpdate(Gson().toJson(TestCommandMessage()))
         } else {
             // new agentic chat route
             ApplicationManager.getApplication().executeOnPooledThread {

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

@@ -38,7 +38,6 @@ import software.aws.toolkits.jetbrains.core.coroutines.EDT
 import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext
 import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthController
 import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthNeededState
-import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage
 import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher
 import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteraction
 import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteractionType
@@ -52,7 +51,6 @@ import software.aws.toolkits.jetbrains.services.cwc.clients.chat.model.TriggerTy
 import software.aws.toolkits.jetbrains.services.cwc.clients.chat.v1.ChatSessionFactoryV1
 import software.aws.toolkits.jetbrains.services.cwc.commands.CodeScanIssueActionMessage
 import software.aws.toolkits.jetbrains.services.cwc.commands.ContextMenuActionMessage
-import software.aws.toolkits.jetbrains.services.cwc.commands.EditorContextCommand
 import software.aws.toolkits.jetbrains.services.cwc.controller.chat.StaticPrompt
 import software.aws.toolkits.jetbrains.services.cwc.controller.chat.StaticTextResponse
 import software.aws.toolkits.jetbrains.services.cwc.controller.chat.messenger.ChatPromptHandler
@@ -76,12 +74,6 @@ import software.aws.toolkits.jetbrains.services.cwc.storage.ChatSessionStorage
 import software.aws.toolkits.telemetry.CwsprChatCommandType
 import java.util.UUID
 
-data class TestCommandMessage(
-    val sender: String = "codetest",
-    val command: String = "test",
-    val type: String = "addAnswer",
-) : AmazonQMessage
-
 class ChatController private constructor(
     private val context: AmazonQAppInitContext,
     private val chatSessionStorage: ChatSessionStorage,
@@ -285,16 +277,7 @@ class ChatController private constructor(
 
     // JB specific (not in vscode)
     override suspend fun processContextMenuCommand(message: ContextMenuActionMessage) {
-        // Extract context
-        if (message.project != context.project) {
-            return
-        }
-
-        if (message.command == EditorContextCommand.GenerateUnitTests) {
-            // Publish an event to "codetest" tab with command as "test" and type as "addAnswer"
-            val messageToPublish = TestCommandMessage()
-            context.messagesFromAppToUi.publish(messageToPublish)
-        }
+        // No-op since context commands are handled elsewhere. This function will be deprecated once we remove this class
     }
 
     private suspend fun processPromptActions(

plugins/amazonq/mynah-ui/src/mynah-ui/ui/quickActions/handler.ts

@@ -331,7 +331,7 @@ private handleDocCommand(chatPrompt: ChatPrompt, tabID: string, taskName: string
         this.showScanInTab(affectedTabId)
     }
 
-    private handleCodeTestCommand(chatPrompt: ChatPrompt, tabID: string, eventId: string | undefined) {
+    private handleCodeTestCommand(chatPrompt: ChatPrompt, tabID: string | undefined, eventId: string | undefined) {
         if (!this.isCodeTestEnabled) {
             return
         }
@@ -343,10 +343,17 @@ private handleDocCommand(chatPrompt: ChatPrompt, tabID: string, taskName: string
             this.connector.startTestGen(testTabId, realPromptText)
             return
         }
+        /**
+         * right click -> generate test has no tab id
+         * we have to manually create one if a testgen tab
+         * wasn't previously created
+         */
+        if (!tabID) {
+            tabID = this.mynahUI?.updateStore('', {})
+        }
         const affectedTabId: string | undefined = this.addTab(tabID)
 
         // if there is no test tab, open a new one
-
         if (affectedTabId === undefined) {
             this.mynahUI?.notify({
                 content: uiComponentsTexts.noMoreTabsTooltip,