in plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatController.kt [169:319]
override suspend fun processStartTestGen(message: IncomingCodeTestMessage.StartTestGen) {
codeTestChatHelper.setActiveCodeTestTabId(message.tabId)
val session = codeTestChatHelper.getActiveSession()
if (session.isGeneratingTests) {
return
}
sessionCleanUp(session.tabId)
// check if IDE has active file open, yes return (fileName and filePath) else return null
val project = context.project
val fileInfo = checkActiveFileInIDE(project, message) ?: return
session.programmingLanguage = fileInfo.fileLanguage
session.startTimeOfTestGeneration = Instant.now().toEpochMilli().toDouble()
session.isGeneratingTests = true
var requestId: String = ""
var statusCode: Int = 0
var conversationId: String? = null
var testResponseMessageId: String? = null
var testResponseText: String = ""
val userMessage = when {
message.prompt != "" -> {
"/test ${message.prompt}"
}
else -> "/test Generate unit tests for `${fileInfo.fileName}`"
}
session.hasUserPromptSupplied = message.prompt.isNotEmpty()
// Send user prompt to chat
codeTestChatHelper.addNewMessage(
CodeTestChatMessageContent(message = userMessage, type = ChatMessageType.Prompt, canBeVoted = false),
message.tabId,
false
)
if (fileInfo.fileInWorkspace && isLanguageSupported(fileInfo.fileLanguage.languageId)) {
// Send Capability card to chat
codeTestChatHelper.addNewMessage(
CodeTestChatMessageContent(informationCard = true, message = null, type = ChatMessageType.Answer, canBeVoted = false),
message.tabId,
false
)
var selectionRange = getEditorSelectionRange(project)
session.isCodeBlockSelected = selectionRange !== null
// This removes /test if user accidentally added while doing Regenerate unit tests and truncates the user response to 4096 characters.
val userPrompt = message.prompt
.let { if (it.startsWith("/test")) it.substringAfter("/test ").trim() else it }
.take(4096)
CodeWhispererUTGChatManager.getInstance(project).generateTests(userPrompt, codeTestChatHelper, null, selectionRange)
} else {
// Not adding a progress bar to unsupported language cases
val responseHandler = GenerateAssistantResponseResponseHandler.builder()
.onResponse {
requestId = it.responseMetadata().requestId()
statusCode = it.sdkHttpResponse().statusCode()
conversationId = it.conversationId()
}
.subscriber { stream: ChatResponseStream ->
stream.accept(object : GenerateAssistantResponseResponseHandler.Visitor {
override fun visitAssistantResponseEvent(event: AssistantResponseEvent) {
testResponseText += event.content()
cs.launch {
codeTestChatHelper.updateAnswer(
CodeTestChatMessageContent(
message = testResponseText,
type = ChatMessageType.AnswerPart
),
messageIdOverride = testResponseMessageId
)
}
}
})
}
.build()
if (!fileInfo.fileInWorkspace) {
val messageContent =
"<span style=\"color: #EE9D28;\">⚠<b> I can't generate tests for ${fileInfo.fileName}" +
" because it's outside the project directory.</b><br></span> " +
"I can still provide examples, instructions and code suggestions."
codeTestChatHelper.addNewMessage(
CodeTestChatMessageContent(
message = messageContent,
type = ChatMessageType.Answer,
canBeVoted = false
),
message.tabId,
false
)
}
testResponseMessageId = codeTestChatHelper.addAnswer(
CodeTestChatMessageContent(
message = "",
type = ChatMessageType.AnswerStream
)
)
codeTestChatHelper.updateUI(
loadingChat = true,
promptInputDisabledState = true,
)
// Send Request to Sync UTG API
val contextExtractor = ActiveFileContextExtractor.create(fqnWebviewAdapter = null, project = project)
val activeFileContext = ActiveFileContext(
fileContext = FileContext(
fileLanguage = fileInfo.fileLanguage.languageId,
filePath = fileInfo.filePath,
matchPolicy = null
),
focusAreaContext = contextExtractor.extractContextForTrigger(ExtractionTriggerType.ChatMessage).focusAreaContext,
)
val requestData = ChatRequestData(
tabId = session.tabId,
message = "Generate unit tests for the following part of my code: ${message.prompt.ifBlank { fileInfo.fileName }}",
activeFileContext = activeFileContext,
userIntent = UserIntent.GENERATE_UNIT_TESTS,
triggerType = TriggerType.ContextMenu,
customization = CodeWhispererModelConfigurator.getInstance().activeCustomization(context.project),
relevantTextDocuments = emptyList(),
useRelevantDocuments = false,
)
val client = QRegionProfileManager.getInstance().getQClient<CodeWhispererStreamingAsyncClient>(project)
val request = requestData.toChatRequest()
client.generateAssistantResponse(request, responseHandler).await()
// TODO: Need to send isCodeBlockSelected field
requestId.let { id ->
LOG.debug { "$FEATURE_NAME: Unit test generation requestId: $id" }
AmazonqTelemetry.utgGenerateTests(
cwsprChatProgrammingLanguage = session.programmingLanguage.languageId,
hasUserPromptSupplied = session.hasUserPromptSupplied,
isFileInWorkspace = fileInfo.fileInWorkspace,
isSupportedLanguage = isLanguageSupported(fileInfo.fileLanguage.languageId),
credentialStartUrl = getStartUrl(project),
result = MetricResult.Succeeded,
perfClientLatency = (Instant.now().toEpochMilli() - session.startTimeOfTestGeneration),
requestId = id,
status = Status.ACCEPTED,
)
}
session.isGeneratingTests = false
codeTestChatHelper.updateUI(
loadingChat = false,
promptInputDisabledState = false
)
}
}