Skip to content

Commit 9b1dd41

Browse files
Randall-JiangRandall-Jiang
committed
enable build&execute for internal amazon users
1 parent beabdf1 commit 9b1dd41

File tree

8 files changed

+163
-52
lines changed

8 files changed

+163
-52
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: 97 additions & 24 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 {
@@ -966,23 +972,90 @@ class CodeTestChatController(
966972
return
967973
}
968974

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)
975+
// TODO: only go to future iterations when buildExitCode or testExitCode > 0, right now iterate regardless
976+
if (taskContext.buildExitCode > 0) {
977+
// handle build failure case
978+
979+
taskContext.progressStatus = BuildAndExecuteProgressStatus.BUILD_FAILED
980+
updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration)
981+
taskContext.progressStatus = BuildAndExecuteProgressStatus.FIXING_TEST_CASES
982+
val buildAndExecuteMessageId = updateBuildAndExecuteProgressCard(
983+
taskContext.progressStatus,
984+
messageId,
985+
session.iteration + 1
986+
)
972987

973-
val previousUTGIterationContext = PreviousUTGIterationContext(
974-
buildLogFile = buildLogsFile,
975-
testLogFile = testLogsFile,
976-
selectedFile = session.selectedFile,
977-
buildAndExecuteMessageId = buildAndExecuteMessageId
978-
)
988+
val previousUTGIterationContext = PreviousUTGIterationContext(
989+
buildLogFile = buildLogsFile,
990+
testLogFile = null,
991+
selectedFile = session.selectedFile,
992+
buildAndExecuteMessageId = buildAndExecuteMessageId
993+
)
979994

980-
val job = CodeWhispererUTGChatManager.getInstance(context.project).generateTests("", codeTestChatHelper, previousUTGIterationContext, null)
981-
job?.join()
995+
session.isGeneratingTests = false
996+
session.conversationState = ConversationState.IDLE
982997

983-
taskContext.progressStatus = BuildAndExecuteProgressStatus.PROCESS_TEST_RESULTS
984-
// session.iteration already updated in generateTests
985-
updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration - 1)
998+
val job = CodeWhispererUTGChatManager.getInstance(
999+
context.project
1000+
).generateTests(session.userPrompt, codeTestChatHelper, previousUTGIterationContext, null)
1001+
job?.join()
1002+
1003+
return
1004+
}
1005+
// taskContext.progressStatus = BuildAndExecuteProgressStatus.RUN_EXECUTION_TESTS
1006+
// updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration)
1007+
//
1008+
// val testLogsFile = VirtualFileManager.getInstance().findFileByNioPath(
1009+
// withContext(currentCoroutineContext()) {
1010+
// Files.createTempFile(null, null)
1011+
// }
1012+
// )
1013+
// if (testLogsFile == null) {
1014+
// // TODO: handle no log file case
1015+
// return
1016+
// }
1017+
// LOG.debug {
1018+
// "Q TestGen session: ${codeTestChatHelper.getActiveCodeTestTabId()}: " +
1019+
// "tmpFile for test logs:\n ${buildLogsFile.path}"
1020+
// }
1021+
// delay(1000)
1022+
// runBuildOrTestCommand(taskContext.executionCommand, testLogsFile, context.project, isBuildCommand = false, taskContext)
1023+
// while (taskContext.testExitCode < 0) {
1024+
// // wait until test command finished
1025+
// delay(1000)
1026+
// }
1027+
//
1028+
// if (taskContext.testExitCode == 0) {
1029+
// taskContext.progressStatus = BuildAndExecuteProgressStatus.TESTS_EXECUTED
1030+
// updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration)
1031+
// codeTestChatHelper.addAnswer(
1032+
// CodeTestChatMessageContent(
1033+
// message = message("testgen.message.success"),
1034+
// type = ChatMessageType.Answer,
1035+
// canBeVoted = false
1036+
// )
1037+
// )
1038+
// // sessionCleanUp(message.tabId)
1039+
// return
1040+
// }
1041+
//
1042+
// // has test failure, we will zip the latest project and invoke backend again
1043+
// taskContext.progressStatus = BuildAndExecuteProgressStatus.FIXING_TEST_CASES
1044+
// val buildAndExecuteMessageId = updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration)
1045+
//
1046+
// val previousUTGIterationContext = PreviousUTGIterationContext(
1047+
// buildLogFile = buildLogsFile,
1048+
// testLogFile = testLogsFile,
1049+
// selectedFile = session.selectedFile,
1050+
// buildAndExecuteMessageId = buildAndExecuteMessageId
1051+
// )
1052+
//
1053+
// val job = CodeWhispererUTGChatManager.getInstance(context.project).generateTests("", codeTestChatHelper, previousUTGIterationContext, null)
1054+
// job?.join()
1055+
//
1056+
// taskContext.progressStatus = BuildAndExecuteProgressStatus.PROCESS_TEST_RESULTS
1057+
// // session.iteration already updated in generateTests
1058+
// updateBuildAndExecuteProgressCard(taskContext.progressStatus, messageId, session.iteration - 1)
9861059
}
9871060
"utg_modify_command" -> {
9881061
// TODO allow user input to modify the command

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

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,16 @@ enum class BuildAndExecuteStatusIcon(val icon: String) {
99
WAIT("<span>&#9744;</span>"),
1010
CURRENT("<span>&#9744;</span>"),
1111
DONE("<span style=\"color: green;\">&#10004;</span>"),
12+
FAILED("<span style=\"color: red;\">&#10060;</span>"),
1213
}
1314

1415
fun getBuildIcon(progressStatus: BuildAndExecuteProgressStatus) =
1516
if (progressStatus < BuildAndExecuteProgressStatus.RUN_BUILD) {
1617
BuildAndExecuteStatusIcon.WAIT.icon
1718
} else if (progressStatus == BuildAndExecuteProgressStatus.RUN_BUILD) {
1819
BuildAndExecuteStatusIcon.CURRENT.icon
20+
} else if (progressStatus == BuildAndExecuteProgressStatus.BUILD_FAILED || progressStatus == BuildAndExecuteProgressStatus.FIXING_TEST_CASES) {
21+
BuildAndExecuteStatusIcon.FAILED.icon
1922
} else {
2023
BuildAndExecuteStatusIcon.DONE.icon
2124
}
@@ -25,15 +28,19 @@ fun getExecutionIcon(progressStatus: BuildAndExecuteProgressStatus) =
2528
BuildAndExecuteStatusIcon.WAIT.icon
2629
} else if (progressStatus == BuildAndExecuteProgressStatus.RUN_EXECUTION_TESTS) {
2730
BuildAndExecuteStatusIcon.CURRENT.icon
31+
} else if (progressStatus == BuildAndExecuteProgressStatus.BUILD_FAILED || progressStatus == BuildAndExecuteProgressStatus.FIXING_TEST_CASES) {
32+
BuildAndExecuteStatusIcon.FAILED.icon
2833
} else {
2934
BuildAndExecuteStatusIcon.DONE.icon
3035
}
3136

3237
fun getFixingTestCasesIcon(progressStatus: BuildAndExecuteProgressStatus) =
33-
if (progressStatus < BuildAndExecuteProgressStatus.FIXING_TEST_CASES) {
34-
BuildAndExecuteStatusIcon.WAIT.icon
38+
if (progressStatus == BuildAndExecuteProgressStatus.BUILD_FAILED) {
39+
BuildAndExecuteStatusIcon.FAILED.icon
3540
} else if (progressStatus == BuildAndExecuteProgressStatus.FIXING_TEST_CASES) {
3641
BuildAndExecuteStatusIcon.CURRENT.icon
37-
} else {
42+
} else if (progressStatus >= BuildAndExecuteProgressStatus.PROCESS_TEST_RESULTS) {
3843
BuildAndExecuteStatusIcon.DONE.icon
44+
} else {
45+
BuildAndExecuteStatusIcon.WAIT.icon
3946
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import com.intellij.openapi.vfs.VirtualFile
77

88
data class PreviousUTGIterationContext(
99
val buildLogFile: VirtualFile,
10-
val testLogFile: VirtualFile,
10+
val testLogFile: VirtualFile?,
1111
val selectedFile: VirtualFile?,
1212
val buildAndExecuteMessageId: String?,
1313
)

0 commit comments

Comments
 (0)