Skip to content

Commit f1d44db

Browse files
Randall-JiangRandall-Jiang
committed
address comment
1 parent 63fb0ed commit f1d44db

File tree

4 files changed

+70
-302
lines changed

4 files changed

+70
-302
lines changed

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

+58-157
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,11 @@ import kotlinx.coroutines.delay
1717
import kotlinx.coroutines.launch
1818
import software.amazon.awssdk.core.exception.SdkServiceException
1919
import software.amazon.awssdk.services.codewhispererruntime.model.GetTestGenerationResponse
20+
import software.amazon.awssdk.services.codewhispererruntime.model.PackageInfo
2021
import software.amazon.awssdk.services.codewhispererruntime.model.Range
2122
import software.amazon.awssdk.services.codewhispererruntime.model.StartTestGenerationResponse
2223
import software.amazon.awssdk.services.codewhispererruntime.model.TargetCode
24+
import software.amazon.awssdk.services.codewhispererruntime.model.TargetFileInfo
2325
import software.amazon.awssdk.services.codewhispererruntime.model.TestGenerationJobStatus
2426
import software.amazon.awssdk.services.codewhispererstreaming.model.ExportContext
2527
import software.amazon.awssdk.services.codewhispererstreaming.model.ExportIntent
@@ -32,11 +34,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.clients.AmazonQStreaming
3234
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.controller.CodeTestChatHelper
3335
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.Button
3436
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.messages.CodeTestChatMessageContent
35-
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.PackageInfo
36-
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.PackageInfoList
3737
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.PreviousUTGIterationContext
38-
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.TargetFileInfo
39-
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.model.TargetFileInfoList
4038
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.BuildAndExecuteProgressStatus
4139
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.Session
4240
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.utils.combineBuildAndExecuteLogFiles
@@ -61,7 +59,6 @@ import java.io.ByteArrayInputStream
6159
import java.io.ByteArrayOutputStream
6260
import java.io.File
6361
import java.io.IOException
64-
import java.nio.file.Path
6562
import java.nio.file.Paths
6663
import java.time.Duration
6764
import java.time.Instant
@@ -91,7 +88,7 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
9188
session.isGeneratingTests = true
9289
session.iteration++
9390

94-
// Set the Progress bar to "Generating unit tests..."
91+
// Set progress bar to "Generating unit tests..."
9592
codeTestChatHelper.updateUI(
9693
promptInputDisabledState = true,
9794
promptInputProgress = testGenProgressField(0),
@@ -102,6 +99,7 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
10299
session.srcZipFileSize = codeTestResponseContext.payloadContext.srcZipFileSize
103100
session.artifactUploadDuration = codeTestResponseContext.serviceInvocationContext.artifactsUploadDuration
104101
val path = codeTestResponseContext.currentFileRelativePath
102+
105103
if (codeTestResponseContext.payloadContext.payloadLimitCrossed == true) {
106104
fileTooLarge()
107105
}
@@ -110,54 +108,39 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
110108
throwIfCancelled(session)
111109

112110
LOG.debug {
113-
"Q TestGen StartTestGenerationRequest: TabId= ${codeTestChatHelper.getActiveCodeTestTabId()}: " +
114-
"uploadId: ${createUploadUrlResponse.uploadId()}, relativeTargetPath: ${codeTestResponseContext.currentFileRelativePath}, " +
115-
"selectionRange: $selectionRange, "
111+
"Q TestGen StartTestGenerationRequest: TabId=${codeTestChatHelper.getActiveCodeTestTabId()}, " +
112+
"uploadId=${createUploadUrlResponse.uploadId()}, relativeTargetPath=${codeTestResponseContext.currentFileRelativePath}, " +
113+
"selectionRange=$selectionRange"
116114
}
117115

118116
// 2nd API call: StartTestGeneration
119117
val startTestGenerationResponse = try {
120118
var response: StartTestGenerationResponse? = null
121-
122119
waitUntil(
123120
succeedOn = { response?.sdkHttpResponse()?.statusCode() == 200 },
124-
maxDuration = Duration.ofSeconds(1), // 1 second timeout
121+
maxDuration = Duration.ofSeconds(1), // 1-second timeout
125122
) {
126-
try {
127-
response = startTestGeneration(
128-
uploadId = createUploadUrlResponse.uploadId(),
129-
targetCode = listOf(
130-
TargetCode.builder()
131-
.relativeTargetPath(codeTestResponseContext.currentFileRelativePath.toString())
132-
.targetLineRangeList(
133-
if (selectionRange != null) {
134-
listOf(selectionRange)
135-
} else {
136-
emptyList()
137-
}
138-
)
139-
.build()
140-
),
141-
userInput = prompt
142-
)
143-
delay(200)
144-
response?.testGenerationJob() != null
145-
} catch (e: Exception) {
146-
throw e
147-
}
123+
response = startTestGeneration(
124+
uploadId = createUploadUrlResponse.uploadId(),
125+
targetCode = listOf(
126+
TargetCode.builder()
127+
.relativeTargetPath(codeTestResponseContext.currentFileRelativePath.toString())
128+
.targetLineRangeList(selectionRange?.let { listOf(it) } ?: emptyList())
129+
.build()
130+
),
131+
userInput = prompt
132+
)
133+
delay(200)
134+
response?.testGenerationJob() != null
148135
}
149-
150136
response ?: throw RuntimeException("Failed to start test generation")
151137
} catch (e: Exception) {
152138
LOG.error(e) { "Unexpected error while creating test generation job" }
153-
val errorMessage = getTelemetryErrorMessage(e, CodeWhispererConstants.FeatureName.TEST_GENERATION)
154-
155-
// Sending requestId to telemetry if there is Validation Exception
156139
if (e is SdkServiceException) {
157140
session.startTestGenerationRequestId = e.requestId()
158141
}
159142
throw CodeTestException(
160-
"CreateTestJobError: $errorMessage",
143+
"CreateTestJobError: ${getTelemetryErrorMessage(e, CodeWhispererConstants.FeatureName.TEST_GENERATION)}",
161144
"CreateTestJobError",
162145
message("testgen.error.generic_technical_error_message")
163146
)
@@ -169,57 +152,29 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
169152
session.testGenerationJob = job.testGenerationJobId()
170153
throwIfCancelled(session)
171154

172-
// 3rd API call: Step 3: Polling mechanism on test job status with getTestGenStatus getTestGeneration
155+
// 3rd API call: Polling mechanism on test job status
173156
var finished = false
174157
var testGenerationResponse: GetTestGenerationResponse? = null
175-
176-
var packageInfoList: List<software.amazon.awssdk.services.codewhispererruntime.model.PackageInfo> = emptyList()
177-
LOG.debug {
178-
"Q TestGen session: ${codeTestChatHelper.getActiveCodeTestTabId()}: " +
179-
"polling result for id: ${job.testGenerationJobId()}, group name: ${job.testGenerationJobGroupName()}, " +
180-
"request id: ${startTestGenerationResponse.responseMetadata().requestId()}"
181-
}
158+
var packageInfoList = emptyList<software.amazon.awssdk.services.codewhispererruntime.model.PackageInfo>()
159+
var packageInfo: software.amazon.awssdk.services.codewhispererruntime.model.PackageInfo? = null
160+
var targetFileInfo: software.amazon.awssdk.services.codewhispererruntime.model.TargetFileInfo? = null
182161

183162
while (!finished) {
184163
throwIfCancelled(session)
185164
testGenerationResponse = getTestGenerationStatus(job.testGenerationJobId(), job.testGenerationJobGroupName())
186-
187165
val status = testGenerationResponse.testGenerationJob().status()
188-
if (status == TestGenerationJobStatus.COMPLETED) {
189-
LOG.debug {
190-
"Q TestGen session: ${codeTestChatHelper.getActiveCodeTestTabId()}: " +
191-
"Test generation completed, package info: ${testGenerationResponse.testGenerationJob().packageInfoList()}"
192-
}
193-
finished = true
194-
packageInfoList = testGenerationResponse.testGenerationJob().packageInfoList()
195-
val packageInfo = packageInfoList.firstOrNull()
196-
val targetFileInfo = packageInfo?.targetFileInfoList()?.firstOrNull()
197-
if (packageInfo != null && targetFileInfo != null) {
198-
try {
199-
session.packageInfoList.member = PackageInfo(
200-
executionCommand = packageInfo.executionCommand(),
201-
buildCommand = packageInfo.buildCommand(),
202-
buildOrder = packageInfo.buildOrder(),
203-
testFramework = packageInfo.testFramework(),
204-
packageSummary = packageInfo.packageSummary(),
205-
packagePlan = packageInfo.packagePlan(),
206-
targetFileInfoList = TargetFileInfoList(
207-
member = listOf(
208-
TargetFileInfo(
209-
filePath = targetFileInfo.filePath(),
210-
testFilePath = targetFileInfo.testFilePath(),
211-
testCoverage = targetFileInfo.testCoverage(),
212-
fileSummary = targetFileInfo.fileSummary(),
213-
filePlan = targetFileInfo.filePlan(),
214-
numberOfTestMethods = targetFileInfo.numberOfTestMethods()
215-
216-
)
217-
)
218-
)
219-
)
220-
221-
val testFileName = targetFileInfo.testFilePath()?.let { File(it).name }.orEmpty()
222-
session.testFileName = testFileName
166+
packageInfoList = testGenerationResponse.testGenerationJob().packageInfoList()
167+
packageInfo = packageInfoList.firstOrNull()
168+
targetFileInfo = packageInfo?.targetFileInfoList()?.firstOrNull()
169+
170+
when (status) {
171+
TestGenerationJobStatus.COMPLETED -> {
172+
LOG.debug { "Test generation completed, package info: $packageInfoList" }
173+
finished = true
174+
175+
if (packageInfo != null && targetFileInfo != null) {
176+
session.packageInfoList = packageInfoList
177+
session.testFileName = targetFileInfo.testFilePath()?.let { File(it).name }.orEmpty()
223178
session.numberOfUnitTestCasesGenerated = targetFileInfo.numberOfTestMethods() ?: 0
224179
session.testFileRelativePathToProjectRoot = getTestFilePathRelativeToRoot(targetFileInfo)
225180

@@ -232,78 +187,24 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
232187
messageIdOverride = codeTestResponseContext.testSummaryMessageId
233188
)
234189
}
235-
packageInfoList = session.packageInfoList
236-
} catch (e: Exception) {
237-
throw CodeTestException(
238-
"TestGenFailedError: " + message("testgen.message.failed"),
239-
"TestGenFailedError",
240-
message("testgen.error.generic_technical_error_message")
241-
)
242190
}
243191
}
244-
} else if (status == TestGenerationJobStatus.FAILED) {
245-
LOG.debug {
246-
"Q TestGen session: ${codeTestChatHelper.getActiveCodeTestTabId()}: " +
247-
"Test generation failed, package info: ${testGenerationResponse.testGenerationJob().packageInfoList()}"
248-
}
249-
250-
packageInfoList = testGenerationResponse.testGenerationJob().packageInfoList()
251-
val packageInfo = packageInfoList.firstOrNull()
252-
val targetFileInfo = packageInfo?.targetFileInfoList()?.firstOrNull()
253-
254-
if (packageInfo != null) {
255-
session.packageInfoList.member = PackageInfo(
256-
executionCommand = packageInfo.executionCommand(),
257-
buildCommand = packageInfo.buildCommand(),
258-
buildOrder = packageInfo.buildOrder(),
259-
testFramework = packageInfo.testFramework(),
260-
packageSummary = packageInfo.packageSummary(),
261-
packagePlan = packageInfo.packagePlan(),
262-
targetFileInfoList = TargetFileInfoList(member = emptyList())
192+
TestGenerationJobStatus.FAILED -> {
193+
LOG.debug { "Test generation failed, package info: $packageInfoList" }
194+
throw CodeTestException(
195+
"TestGenFailedError: " + message("testgen.message.failed"),
196+
"TestGenFailedError",
197+
message("testgen.error.generic_technical_error_message")
263198
)
264199
}
265-
266-
throw CodeTestException(
267-
"TestGenFailedError: " + message("testgen.message.failed"),
268-
"TestGenFailedError",
269-
message("testgen.error.generic_technical_error_message")
270-
)
271-
} else {
272-
// In progress
273-
LOG.debug {
274-
"Q TestGen session: ${codeTestChatHelper.getActiveCodeTestTabId()}: " +
275-
"Test generation in progress, progress rate ${testGenerationResponse.testGenerationJob().progressRate()}}"
276-
}
277-
val progressRate = testGenerationResponse.testGenerationJob().progressRate() ?: 0
278-
279-
if (previousIterationContext == null) {
280-
val packageInfoList = testGenerationResponse.testGenerationJob().packageInfoList()
281-
val packageInfo = packageInfoList.firstOrNull()
282-
if (packageInfo != null) {
283-
try {
284-
val targetFileInfo = packageInfo?.targetFileInfoList()?.firstOrNull()
285-
if (targetFileInfo != null) {
286-
val fileName = targetFileInfo?.filePath()?.let { Path.of(it).fileName.toString() } ?: path.fileName.toString()
287-
codeTestChatHelper.updateAnswer(
288-
CodeTestChatMessageContent(
289-
message = generateSummaryMessage(fileName) + (targetFileInfo.filePlan() ?: ""),
290-
type = ChatMessageType.Answer
291-
),
292-
messageIdOverride = codeTestResponseContext.testSummaryMessageId
293-
)
294-
}
295-
} catch (e: Exception) {
296-
LOG.debug("failed to process package info: ${e.message}")
297-
}
298-
}
200+
else -> {
201+
LOG.debug { "Test generation in progress, progress rate: ${testGenerationResponse.testGenerationJob().progressRate()}" }
202+
codeTestChatHelper.updateUI(
203+
promptInputDisabledState = true,
204+
promptInputProgress = testGenProgressField(testGenerationResponse.testGenerationJob().progressRate() ?: 0),
205+
)
299206
}
300-
codeTestChatHelper.updateUI(
301-
promptInputDisabledState = true,
302-
promptInputProgress = testGenProgressField(progressRate),
303-
)
304207
}
305-
306-
// polling every 2 seconds to reduce # of API calls
307208
delay(2000)
308209
}
309210

@@ -343,15 +244,15 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
343244
return
344245
}
345246

346-
val targetFileInfo = session.packageInfoList.member?.targetFileInfoList?.member?.firstOrNull()
347-
val codeReference = targetFileInfo?.codeReferences?.map { ref ->
247+
targetFileInfo = session.packageInfoList.firstOrNull()?.targetFileInfoList()?.firstOrNull()
248+
val codeReference = targetFileInfo?.codeReferences()?.map { ref ->
348249
CodeReference(
349-
licenseName = ref.licenseName,
350-
url = ref.url,
351-
information = "${ref.licenseName} - <a href=\"${ref.url}\">${ref.repository}</a>"
250+
licenseName = ref.licenseName(),
251+
url = ref.url(),
252+
information = "${ref.licenseName()} - <a href=\"${ref.url()}\">${ref.repository()}</a>"
352253
)
353254
}
354-
targetFileInfo?.codeReferences?.let { session.codeReferences = it }
255+
targetFileInfo?.codeReferences()?.let { session.codeReferences = it }
355256
val isReferenceAllowed = CodeWhispererSettings.getInstance().isIncludeCodeWithReference()
356257
if (!isReferenceAllowed && codeReference?.isNotEmpty() == true) {
357258
codeTestChatHelper.addAnswer(
@@ -364,7 +265,7 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
364265
)
365266
} else {
366267
if (previousIterationContext == null) {
367-
// show another card as the answer
268+
// show another card as the answer, remove this once we remove 3 backticks from backend
368269
val jobSummary = testGenerationResponse?.testGenerationJob()?.jobSummary()?.trim() ?: ""
369270

370271
val cleanedPlanSummary = jobSummary
@@ -430,7 +331,7 @@ Please see the unit tests generated below. Click 'View Diff' to review the chang
430331
// both needs to be handled the same way which is remove the first sub-directory
431332
private fun getTestFilePathRelativeToRoot(targetFileInfo: Any?): String {
432333
val pathString = when (targetFileInfo) {
433-
is TargetFileInfo -> targetFileInfo.testFilePath
334+
is TargetFileInfo -> targetFileInfo.testFilePath()
434335
is software.amazon.awssdk.services.codewhispererruntime.model.TargetFileInfo -> targetFileInfo.testFilePath()
435336
else -> generatedTestDiffs.keys.firstOrNull()
436337
} ?: throw RuntimeException("No test file path found")
@@ -534,7 +435,7 @@ Please see the unit tests generated below. Click 'View Diff' to review the chang
534435
private fun resetTestGenFlowSession(session: Session) {
535436
// session.selectedFile doesn't need to be reset since it will remain unchanged
536437
session.conversationState = ConversationState.IN_PROGRESS
537-
session.packageInfoList = PackageInfoList()
438+
session.packageInfoList = emptyList()
538439
session.openedDiffFile = null
539440
session.testFileRelativePathToProjectRoot = ""
540441
session.testFileName = ""

0 commit comments

Comments
 (0)