Skip to content

Commit 24a0832

Browse files
Randall-JiangRandall-Jiang
committed
add build-execute code
1 parent beabdf1 commit 24a0832

File tree

8 files changed

+174
-87
lines changed

8 files changed

+174
-87
lines changed

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt‎

Lines changed: 20 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,6 @@ import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.PreviousUT
3737
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.ShortAnswer
3838
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.BuildAndExecuteProgressStatus
3939
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.Session
40-
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.utils.combineBuildAndExecuteLogFiles
4140
import software.aws.toolkits.jetbrains.services.codemodernizer.utils.calculateTotalLatency
4241
import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.CodeTestException
4342
import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.fileTooLarge
@@ -59,10 +58,10 @@ import java.io.ByteArrayInputStream
5958
import java.io.ByteArrayOutputStream
6059
import java.io.File
6160
import java.io.IOException
62-
import java.nio.file.Path
6361
import java.nio.file.Paths
6462
import java.time.Duration
6563
import java.time.Instant
64+
import java.util.UUID
6665
import java.util.concurrent.atomic.AtomicBoolean
6766
import java.util.zip.ZipInputStream
6867

@@ -89,6 +88,10 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
8988
val session = codeTestChatHelper.getActiveSession()
9089
session.isGeneratingTests = true
9190
session.iteration++
91+
if (session.testGenerationJobGroupName.isNullOrEmpty()) {
92+
session.testGenerationJobGroupName = UUID.randomUUID().toString()
93+
}
94+
val final = session.testGenerationJobGroupName
9295

9396
// Set the Progress bar to "Generating unit tests..."
9497
codeTestChatHelper.updateUI(
@@ -137,7 +140,8 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
137140
)
138141
.build()
139142
),
140-
userInput = prompt
143+
userInput = prompt,
144+
testGenerationJobGroupName = final
141145
)
142146
delay(200)
143147
response?.testGenerationJob() != null
@@ -175,7 +179,7 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
175179
var shortAnswer = ShortAnswer()
176180
LOG.debug {
177181
"Q TestGen session: ${codeTestChatHelper.getActiveCodeTestTabId()}: " +
178-
"polling result for id: ${job.testGenerationJobId()}, group name: ${job.testGenerationJobGroupName()}, " +
182+
"polling result for id: ${job.testGenerationJobId()}, group name: ${session.testGenerationJobGroupName}, " +
179183
"request id: ${startTestGenerationResponse.responseMetadata().requestId()}"
180184
}
181185

@@ -198,6 +202,7 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
198202
// Setting default value to 0 if the value is null or invalid
199203
session.numberOfUnitTestCasesGenerated = shortAnswer.numberOfTestMethods
200204
session.testFileRelativePathToProjectRoot = getTestFilePathRelativeToRoot(shortAnswer)
205+
session.shortAnswer = shortAnswer
201206

202207
// update test summary card in success case
203208
if (previousIterationContext == null) {
@@ -250,10 +255,9 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
250255
if (shortAnswer.stopIteration == "true") {
251256
throw CodeTestException("TestGenFailedError: ${shortAnswer.planSummary}", "TestGenFailedError", shortAnswer.planSummary)
252257
}
253-
val fileName = shortAnswer.sourceFilePath?.let { Path.of(it).fileName.toString() } ?: path.fileName.toString()
254258
codeTestChatHelper.updateAnswer(
255259
CodeTestChatMessageContent(
256-
message = generateSummaryMessage(fileName) + shortAnswer.planSummary,
260+
message = generateSummaryMessage(path.fileName.toString()) + shortAnswer.planSummary,
257261
type = ChatMessageType.Answer
258262
),
259263
messageIdOverride = codeTestResponseContext.testSummaryMessageId
@@ -297,7 +301,7 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
297301
// TODO: Modify text according to FnF
298302
codeTestChatHelper.addAnswer(
299303
CodeTestChatMessageContent(
300-
message = message("testgen.error.generic_technical_error_message"),
304+
message = message("testgen.message.failed"),
301305
type = ChatMessageType.Answer,
302306
canBeVoted = true
303307
)
@@ -469,10 +473,8 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
469473
previousIterationContext.selectedFile
470474
}
471475

472-
val combinedBuildAndExecuteLogFile = combineBuildAndExecuteLogFiles(
473-
previousIterationContext?.buildLogFile,
474-
previousIterationContext?.testLogFile
475-
)
476+
val combinedBuildAndExecuteLogFile =
477+
previousIterationContext?.buildLogFile
476478
val codeTestSessionConfig = CodeTestSessionConfig(file, project, combinedBuildAndExecuteLogFile)
477479
codeTestChatHelper.getActiveSession().projectRoot = codeTestSessionConfig.projectRoot.path
478480

@@ -481,8 +483,13 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
481483
return codeWhispererCodeTestSession.run(codeTestChatHelper, previousIterationContext)
482484
}
483485

484-
private fun startTestGeneration(uploadId: String, targetCode: List<TargetCode>, userInput: String): StartTestGenerationResponse =
485-
CodeWhispererClientAdaptor.getInstance(project).startTestGeneration(uploadId, targetCode, userInput)
486+
private fun startTestGeneration(
487+
uploadId: String,
488+
targetCode: List<TargetCode>,
489+
userInput: String,
490+
testGenerationJobGroupName: String,
491+
): StartTestGenerationResponse =
492+
CodeWhispererClientAdaptor.getInstance(project).startTestGeneration(uploadId, targetCode, userInput, testGenerationJobGroupName)
486493

487494
private fun getTestGenerationStatus(jobId: String, jobGroupName: String): GetTestGenerationResponse =
488495
CodeWhispererClientAdaptor.getInstance(project).getTestGeneration(jobId, jobGroupName)

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

Lines changed: 100 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ import software.aws.toolkits.jetbrains.core.AwsClientManager
5353
import software.aws.toolkits.jetbrains.core.coroutines.EDT
5454
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
5555
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
56+
import software.aws.toolkits.jetbrains.core.credentials.sono.isInternalUser
5657
import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext
5758
import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthController
5859
import software.aws.toolkits.jetbrains.services.amazonq.project.RelevantDocument
@@ -67,6 +68,7 @@ import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.Incomin
6768
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.PreviousUTGIterationContext
6869
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.BuildAndExecuteProgressStatus
6970
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.Session
71+
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.UTG_CHAT_MAX_ITERATION
7072
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.storage.ChatSessionStorage
7173
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.utils.constructBuildAndExecutionSummaryText
7274
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.utils.runBuildOrTestCommand
@@ -617,6 +619,8 @@ class CodeTestChatController(
617619
"utg_accept" -> {
618620
// open the file at test path relative to the project root
619621
val testFileAbsolutePath = Paths.get(session.projectRoot, session.testFileRelativePathToProjectRoot)
622+
val startUrl = getStartUrl(context.project)
623+
val isInternalUser = isInternalUser(startUrl)
620624
openOrCreateTestFileAndApplyDiff(context.project, testFileAbsolutePath, session.generatedTestDiffs.values.first(), session.openedDiffFile)
621625
session.codeReferences?.let { references ->
622626
LOG.debug { "Accepted unit tests with references: $references" }
@@ -691,24 +695,27 @@ class CodeTestChatController(
691695
canBeVoted = false
692696
)
693697
)
694-
sessionCleanUp(session.tabId)
698+
if (!isInternalUser) {
699+
sessionCleanUp(message.tabId)
700+
return
701+
}
702+
695703
codeTestChatHelper.updateUI(
696704
promptInputDisabledState = false,
697705
promptInputPlaceholder = message("testgen.placeholder.enter_slash_quick_actions"),
698706
)
699-
/*
707+
700708
val taskContext = session.buildAndExecuteTaskContext
701709
if (session.iteration < 2) {
702-
taskContext.buildCommand = getBuildCommand(message.tabId)
703-
taskContext.executionCommand = getExecutionCommand(message.tabId)
710+
taskContext.buildCommand = codeTestChatHelper.getSession(message.tabId).shortAnswer.buildCommand.toString()
711+
704712
codeTestChatHelper.addAnswer(
705713
CodeTestChatMessageContent(
706714
message = """
707715
Would you like me to help build and execute the test? I'll run following commands
708716
709717
```sh
710718
${taskContext.buildCommand}
711-
${taskContext.executionCommand}
712719
```
713720
""".trimIndent(),
714721
type = ChatMessageType.Answer,
@@ -790,10 +797,9 @@ class CodeTestChatController(
790797
promptInputPlaceholder = message("testgen.placeholder.newtab")
791798
)
792799
}
793-
*/
794800
}
795-
/*
796-
//TODO: this is for unit test regeneration build iteration loop
801+
802+
// TODO: this is for unit test regeneration build iteration loop
797803
"utg_regenerate" -> {
798804
// close the existing open diff in the editor.
799805
ApplicationManager.getApplication().invokeLater {
@@ -810,6 +816,7 @@ class CodeTestChatController(
810816
session.testGenerationJob,
811817
session.testGenerationJobGroupName,
812818
session.programmingLanguage,
819+
IdeCategory.JETBRAINS,
813820
session.numberOfUnitTestCasesGenerated,
814821
0,
815822
session.linesOfCodeGenerated,
@@ -821,13 +828,12 @@ class CodeTestChatController(
821828
"Successfully sent test generation telemetry. RequestId: ${
822829
testGenerationEventResponse.responseMetadata().requestId()}"
823830
}
824-
sessionCleanUp(session.tabId)
831+
// sessionCleanUp(session.tabId)
825832
codeTestChatHelper.updateUI(
826833
promptInputDisabledState = false,
827834
promptInputPlaceholder = message("testgen.placeholder.waiting_on_your_inputs"),
828835
)
829836
}
830-
*/
831837

832838
"utg_reject" -> {
833839
ApplicationManager.getApplication().invokeLater {
@@ -917,43 +923,14 @@ class CodeTestChatController(
917923
"tmpFile for build logs:\n ${buildLogsFile.path}"
918924
}
919925

920-
runBuildOrTestCommand(taskContext.buildCommand, buildLogsFile, context.project, isBuildCommand = true, taskContext)
926+
runBuildOrTestCommand(taskContext.buildCommand, buildLogsFile, context.project, isBuildCommand = true, taskContext, session.projectRoot)
921927
while (taskContext.buildExitCode < 0) {
922928
// wait until build command finished
923929
delay(1000)
924930
}
925931

926-
// TODO: only go to future iterations when buildExitCode or testExitCode > 0, right now iterate regardless
927-
if (taskContext.buildExitCode > 0) {
928-
// TODO: handle build failure case
929-
// ...
930-
// return
931-
}
932-
taskContext.progressStatus = BuildAndExecuteProgressStatus.RUN_EXECUTION_TESTS
933-
updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration)
934-
935-
val testLogsFile = VirtualFileManager.getInstance().findFileByNioPath(
936-
withContext(currentCoroutineContext()) {
937-
Files.createTempFile(null, null)
938-
}
939-
)
940-
if (testLogsFile == null) {
941-
// TODO: handle no log file case
942-
return
943-
}
944-
LOG.debug {
945-
"Q TestGen session: ${codeTestChatHelper.getActiveCodeTestTabId()}: " +
946-
"tmpFile for test logs:\n ${buildLogsFile.path}"
947-
}
948-
delay(1000)
949-
runBuildOrTestCommand(taskContext.executionCommand, testLogsFile, context.project, isBuildCommand = false, taskContext)
950-
while (taskContext.testExitCode < 0) {
951-
// wait until test command finished
952-
delay(1000)
953-
}
954-
955-
if (taskContext.testExitCode == 0) {
956-
taskContext.progressStatus = BuildAndExecuteProgressStatus.TESTS_EXECUTED
932+
if (taskContext.buildExitCode == 0) {
933+
taskContext.progressStatus = BuildAndExecuteProgressStatus.PROCESS_TEST_RESULTS
957934
updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration)
958935
codeTestChatHelper.addAnswer(
959936
CodeTestChatMessageContent(
@@ -966,23 +943,90 @@ class CodeTestChatController(
966943
return
967944
}
968945

969-
// has test failure, we will zip the latest project and invoke backend again
970-
taskContext.progressStatus = BuildAndExecuteProgressStatus.FIXING_TEST_CASES
971-
val buildAndExecuteMessageId = updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration)
946+
// TODO: only go to future iterations when buildExitCode or testExitCode > 0, right now iterate regardless
947+
if (taskContext.buildExitCode > 0) {
948+
// handle build failure case
972949

973-
val previousUTGIterationContext = PreviousUTGIterationContext(
974-
buildLogFile = buildLogsFile,
975-
testLogFile = testLogsFile,
976-
selectedFile = session.selectedFile,
977-
buildAndExecuteMessageId = buildAndExecuteMessageId
978-
)
950+
taskContext.progressStatus = BuildAndExecuteProgressStatus.BUILD_FAILED
951+
updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration)
952+
taskContext.progressStatus = BuildAndExecuteProgressStatus.FIXING_TEST_CASES
953+
val buildAndExecuteMessageId = updateBuildAndExecuteProgressCard(
954+
taskContext.progressStatus,
955+
messageId,
956+
session.iteration + 1
957+
)
958+
959+
val previousUTGIterationContext = PreviousUTGIterationContext(
960+
buildLogFile = buildLogsFile,
961+
testLogFile = null,
962+
selectedFile = session.selectedFile,
963+
buildAndExecuteMessageId = buildAndExecuteMessageId
964+
)
965+
966+
session.isGeneratingTests = false
967+
session.conversationState = ConversationState.IDLE
979968

980-
val job = CodeWhispererUTGChatManager.getInstance(context.project).generateTests("", codeTestChatHelper, previousUTGIterationContext, null)
981-
job?.join()
969+
val job = CodeWhispererUTGChatManager.getInstance(
970+
context.project
971+
).generateTests(session.userPrompt, codeTestChatHelper, previousUTGIterationContext, null)
972+
job?.join()
982973

983-
taskContext.progressStatus = BuildAndExecuteProgressStatus.PROCESS_TEST_RESULTS
984-
// session.iteration already updated in generateTests
985-
updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration - 1)
974+
return
975+
}
976+
// taskContext.progressStatus = BuildAndExecuteProgressStatus.RUN_EXECUTION_TESTS
977+
// updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration)
978+
//
979+
// val testLogsFile = VirtualFileManager.getInstance().findFileByNioPath(
980+
// withContext(currentCoroutineContext()) {
981+
// Files.createTempFile(null, null)
982+
// }
983+
// )
984+
// if (testLogsFile == null) {
985+
// // TODO: handle no log file case
986+
// return
987+
// }
988+
// LOG.debug {
989+
// "Q TestGen session: ${codeTestChatHelper.getActiveCodeTestTabId()}: " +
990+
// "tmpFile for test logs:\n ${buildLogsFile.path}"
991+
// }
992+
// delay(1000)
993+
// runBuildOrTestCommand(taskContext.executionCommand, testLogsFile, context.project, isBuildCommand = false, taskContext)
994+
// while (taskContext.testExitCode < 0) {
995+
// // wait until test command finished
996+
// delay(1000)
997+
// }
998+
//
999+
// if (taskContext.testExitCode == 0) {
1000+
// taskContext.progressStatus = BuildAndExecuteProgressStatus.TESTS_EXECUTED
1001+
// updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration)
1002+
// codeTestChatHelper.addAnswer(
1003+
// CodeTestChatMessageContent(
1004+
// message = message("testgen.message.success"),
1005+
// type = ChatMessageType.Answer,
1006+
// canBeVoted = false
1007+
// )
1008+
// )
1009+
// // sessionCleanUp(message.tabId)
1010+
// return
1011+
// }
1012+
//
1013+
// // has test failure, we will zip the latest project and invoke backend again
1014+
// taskContext.progressStatus = BuildAndExecuteProgressStatus.FIXING_TEST_CASES
1015+
// val buildAndExecuteMessageId = updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration)
1016+
//
1017+
// val previousUTGIterationContext = PreviousUTGIterationContext(
1018+
// buildLogFile = buildLogsFile,
1019+
// testLogFile = testLogsFile,
1020+
// selectedFile = session.selectedFile,
1021+
// buildAndExecuteMessageId = buildAndExecuteMessageId
1022+
// )
1023+
//
1024+
// val job = CodeWhispererUTGChatManager.getInstance(context.project).generateTests("", codeTestChatHelper, previousUTGIterationContext, null)
1025+
// job?.join()
1026+
//
1027+
// taskContext.progressStatus = BuildAndExecuteProgressStatus.PROCESS_TEST_RESULTS
1028+
// // session.iteration already updated in generateTests
1029+
// updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration - 1)
9861030
}
9871031
"utg_modify_command" -> {
9881032
// TODO allow user input to modify the command

0 commit comments

Comments
 (0)