Skip to content

Commit c046fb9

Browse files
sannicmsannicm
authored
feat(amazonqFeatureDev): add code generation unit tests (#4229)
amazonqFeatureDev codebase is lacking unit tests. This PR is extending the coverage for the critical logic that needs to be verified. It includes testing for the codegen chat states, session functions, api utility functions and minor refactors to allow simpler testing.
1 parent c13e5e9 commit c046fb9

File tree

12 files changed

+1014
-183
lines changed

12 files changed

+1014
-183
lines changed

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

Lines changed: 15 additions & 132 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ import com.intellij.diff.contents.EmptyContent
99
import com.intellij.diff.requests.SimpleDiffRequest
1010
import com.intellij.diff.util.DiffUserDataKeys
1111
import com.intellij.ide.BrowserUtil
12-
import com.intellij.notification.NotificationAction
1312
import com.intellij.openapi.application.runInEdt
1413
import com.intellij.openapi.command.WriteCommandAction
1514
import com.intellij.openapi.editor.Caret
@@ -49,19 +48,19 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendA
4948
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAuthNeededException
5049
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAuthenticationInProgressMessage
5150
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendChatInputEnabledMessage
52-
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendCodeResult
5351
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendError
5452
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendSystemPrompt
5553
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendUpdatePlaceholder
54+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.updateFileComponent
5655
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.DeletedFileInfo
5756
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.NewFileZipInfo
5857
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.PrepareCodeGenerationState
5958
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Session
6059
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.SessionStatePhase
6160
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.storage.ChatSessionStorage
61+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getFollowUpOptions
6262
import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference
6363
import software.aws.toolkits.jetbrains.ui.feedback.FeatureDevFeedbackDialog
64-
import software.aws.toolkits.jetbrains.utils.notifyInfo
6564
import software.aws.toolkits.resources.message
6665
import software.aws.toolkits.telemetry.AmazonqTelemetry
6766
import java.util.UUID
@@ -72,8 +71,8 @@ class FeatureDevController(
7271
private val authController: AuthController = AuthController()
7372
) : InboundAppMessagesHandler {
7473

75-
private val messenger = context.messagesFromAppToUi
76-
private val toolWindow = ToolWindowManager.getInstance(context.project).getToolWindow(AmazonQToolWindowFactory.WINDOW_ID)
74+
val messenger = context.messagesFromAppToUi
75+
val toolWindow = ToolWindowManager.getInstance(context.project).getToolWindow(AmazonQToolWindowFactory.WINDOW_ID)
7776

7877
override suspend fun processPromptChatMessage(message: IncomingFeatureDevMessage.ChatPrompt) {
7978
handleChat(
@@ -225,12 +224,7 @@ class FeatureDevController(
225224
filePaths.find { it.zipFilePath == fileToUpdate }?.let { it.rejected = !it.rejected }
226225
deletedFiles.find { it.zipFilePath == fileToUpdate }?.let { it.rejected = !it.rejected }
227226

228-
session.updateFilesPaths(
229-
messenger = messenger,
230-
tabId = message.tabId,
231-
filePaths = filePaths,
232-
deletedFiles = deletedFiles
233-
)
227+
messenger.updateFileComponent(message.tabId, filePaths, deletedFiles)
234228
}
235229

236230
private suspend fun newTabOpened(tabId: String) {
@@ -292,13 +286,18 @@ class FeatureDevController(
292286
references = state.references
293287
}
294288
}
289+
295290
AmazonqTelemetry.isAcceptedCodeChanges(
296-
project = null,
297291
amazonqNumberOfFilesAccepted = (filePaths.filterNot { it.rejected }.size + deletedFiles.filterNot { it.rejected }.size) * 1.0,
298292
amazonqConversationId = session.conversationId,
299293
enabled = true
300294
)
301-
session.insertChanges(filePaths = filePaths, deletedFiles = deletedFiles, references = references)
295+
296+
session.insertChanges(
297+
filePaths = filePaths.filterNot { it.rejected },
298+
deletedFiles = deletedFiles.filterNot { it.rejected },
299+
references = references
300+
)
302301

303302
messenger.sendAnswer(
304303
tabId = tabId,
@@ -502,124 +501,6 @@ class FeatureDevController(
502501
messenger.sendAsyncEventProgress(tabId = tabId, inProgress = false)
503502
}
504503

505-
private suspend fun onCodeGeneration(session: Session, message: String, tabId: String) {
506-
messenger.sendAsyncEventProgress(
507-
tabId = tabId,
508-
inProgress = true,
509-
message = message("amazonqFeatureDev.chat_message.start_code_generation"),
510-
)
511-
512-
try {
513-
messenger.sendAnswer(
514-
tabId = tabId,
515-
message = message("amazonqFeatureDev.chat_message.requesting_changes"),
516-
messageType = FeatureDevMessageType.AnswerStream,
517-
)
518-
519-
messenger.sendUpdatePlaceholder(tabId = tabId, newPlaceholder = message("amazonqFeatureDev.placeholder.generating_code"))
520-
521-
session.send(message) // Trigger code generation
522-
523-
val state = session.sessionState
524-
525-
var filePaths: List<NewFileZipInfo> = emptyList()
526-
var deletedFiles: List<DeletedFileInfo> = emptyList()
527-
var references: List<CodeReference> = emptyList()
528-
var uploadId = ""
529-
530-
when (state) {
531-
is PrepareCodeGenerationState -> {
532-
filePaths = state.filePaths
533-
deletedFiles = state.deletedFiles
534-
references = state.references
535-
uploadId = state.uploadId
536-
}
537-
}
538-
539-
// Atm this is the only possible path as codegen is mocked to return empty.
540-
if (filePaths.size or deletedFiles.size == 0) {
541-
messenger.sendAnswer(
542-
tabId = tabId,
543-
messageType = FeatureDevMessageType.Answer,
544-
message = message("amazonqFeatureDev.code_generation.no_file_changes")
545-
)
546-
messenger.sendSystemPrompt(
547-
tabId = tabId,
548-
followUp = if (retriesRemaining(session) > 0) {
549-
listOf(
550-
FollowUp(
551-
pillText = message("amazonqFeatureDev.follow_up.retry"),
552-
type = FollowUpTypes.RETRY,
553-
status = FollowUpStatusType.Warning
554-
)
555-
)
556-
} else {
557-
emptyList()
558-
}
559-
)
560-
messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = false) // Lock chat input until retry is clicked.
561-
return
562-
}
563-
564-
messenger.sendCodeResult(tabId = tabId, uploadId = uploadId, filePaths = filePaths, deletedFiles = deletedFiles, references = references)
565-
566-
messenger.sendSystemPrompt(tabId = tabId, followUp = getFollowUpOptions(session.sessionState.phase, interactionSucceeded = true))
567-
568-
messenger.sendUpdatePlaceholder(tabId = tabId, newPlaceholder = message("amazonqFeatureDev.placeholder.after_code_generation"))
569-
} finally {
570-
messenger.sendAsyncEventProgress(tabId = tabId, inProgress = false) // Finish processing the event
571-
messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = false) // Lock chat input until a follow-up is clicked.
572-
573-
if (toolWindow != null && !toolWindow.isVisible) {
574-
notifyInfo(
575-
title = message("amazonqFeatureDev.code_generation.notification_title"),
576-
content = message("amazonqFeatureDev.code_generation.notification_message"),
577-
project = context.project,
578-
notificationActions = listOf(openChatNotificationAction())
579-
)
580-
}
581-
}
582-
}
583-
584-
private fun openChatNotificationAction() = NotificationAction.createSimple(message("amazonqFeatureDev.code_generation.notification_open_link")) {
585-
toolWindow?.show()
586-
}
587-
588-
private fun getFollowUpOptions(phase: SessionStatePhase?, interactionSucceeded: Boolean): List<FollowUp> {
589-
when (phase) {
590-
SessionStatePhase.APPROACH -> {
591-
return when (interactionSucceeded) {
592-
true -> listOf(
593-
FollowUp(
594-
pillText = message("amazonqFeatureDev.follow_up.generate_code"),
595-
type = FollowUpTypes.GENERATE_CODE,
596-
status = FollowUpStatusType.Info,
597-
)
598-
)
599-
600-
false -> emptyList()
601-
}
602-
}
603-
SessionStatePhase.CODEGEN -> {
604-
return listOf(
605-
FollowUp(
606-
pillText = message("amazonqFeatureDev.follow_up.insert_code"),
607-
type = FollowUpTypes.INSERT_CODE,
608-
icon = FollowUpIcons.Ok,
609-
status = FollowUpStatusType.Success
610-
),
611-
FollowUp(
612-
pillText = message("amazonqFeatureDev.follow_up.provide_feedback_and_regenerate"),
613-
type = FollowUpTypes.PROVIDE_FEEDBACK_AND_REGENERATE_CODE,
614-
icon = FollowUpIcons.Refresh,
615-
status = FollowUpStatusType.Info
616-
)
617-
)
618-
}
619-
else -> return emptyList()
620-
}
621-
}
622-
623504
private suspend fun retryRequests(tabId: String) {
624505
var session: Session? = null
625506
try {
@@ -728,9 +609,11 @@ class FeatureDevController(
728609
}
729610
}
730611

612+
fun getProject() = context.project
613+
731614
private fun getSessionInfo(tabId: String) = chatSessionStorage.getSession(tabId, context.project)
732615

733-
private fun retriesRemaining(session: Session?): Int = session?.retries ?: DEFAULT_RETRY_LIMIT
616+
fun retriesRemaining(session: Session?): Int = session?.retries ?: DEFAULT_RETRY_LIMIT
734617

735618
companion object {
736619
private val logger = getLogger<FeatureDevController>()

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

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
// Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.amazonqFeatureDev.controller
5+
6+
import com.intellij.notification.NotificationAction
7+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FeatureDevMessageType
8+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUp
9+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUpStatusType
10+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.FollowUpTypes
11+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAnswer
12+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAsyncEventProgress
13+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendChatInputEnabledMessage
14+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendCodeResult
15+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendSystemPrompt
16+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendUpdatePlaceholder
17+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.DeletedFileInfo
18+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.NewFileZipInfo
19+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.PrepareCodeGenerationState
20+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Session
21+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.getFollowUpOptions
22+
import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference
23+
import software.aws.toolkits.jetbrains.utils.notifyInfo
24+
import software.aws.toolkits.resources.message
25+
26+
suspend fun FeatureDevController.onCodeGeneration(session: Session, message: String, tabId: String) {
27+
messenger.sendAsyncEventProgress(
28+
tabId = tabId,
29+
inProgress = true,
30+
message = message("amazonqFeatureDev.chat_message.start_code_generation"),
31+
)
32+
33+
try {
34+
this.messenger.sendAnswer(
35+
tabId = tabId,
36+
message = message("amazonqFeatureDev.chat_message.requesting_changes"),
37+
messageType = FeatureDevMessageType.AnswerStream,
38+
)
39+
40+
messenger.sendUpdatePlaceholder(tabId = tabId, newPlaceholder = message("amazonqFeatureDev.placeholder.generating_code"))
41+
42+
session.send(message) // Trigger code generation
43+
44+
val state = session.sessionState
45+
46+
var filePaths: List<NewFileZipInfo> = emptyList()
47+
var deletedFiles: List<DeletedFileInfo> = emptyList()
48+
var references: List<CodeReference> = emptyList()
49+
var uploadId = ""
50+
51+
when (state) {
52+
is PrepareCodeGenerationState -> {
53+
filePaths = state.filePaths
54+
deletedFiles = state.deletedFiles
55+
references = state.references
56+
uploadId = state.uploadId
57+
}
58+
}
59+
60+
// Atm this is the only possible path as codegen is mocked to return empty.
61+
if (filePaths.size or deletedFiles.size == 0) {
62+
messenger.sendAnswer(
63+
tabId = tabId,
64+
messageType = FeatureDevMessageType.Answer,
65+
message = message("amazonqFeatureDev.code_generation.no_file_changes")
66+
)
67+
messenger.sendSystemPrompt(
68+
tabId = tabId,
69+
followUp = if (retriesRemaining(session) > 0) {
70+
listOf(
71+
FollowUp(
72+
pillText = message("amazonqFeatureDev.follow_up.retry"),
73+
type = FollowUpTypes.RETRY,
74+
status = FollowUpStatusType.Warning
75+
)
76+
)
77+
} else {
78+
emptyList()
79+
}
80+
)
81+
messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = false) // Lock chat input until retry is clicked.
82+
return
83+
}
84+
85+
messenger.sendCodeResult(tabId = tabId, uploadId = uploadId, filePaths = filePaths, deletedFiles = deletedFiles, references = references)
86+
87+
messenger.sendSystemPrompt(tabId = tabId, followUp = getFollowUpOptions(session.sessionState.phase, interactionSucceeded = true))
88+
89+
messenger.sendUpdatePlaceholder(tabId = tabId, newPlaceholder = message("amazonqFeatureDev.placeholder.after_code_generation"))
90+
} finally {
91+
messenger.sendAsyncEventProgress(tabId = tabId, inProgress = false) // Finish processing the event
92+
messenger.sendChatInputEnabledMessage(tabId = tabId, enabled = false) // Lock chat input until a follow-up is clicked.
93+
94+
if (toolWindow != null && !toolWindow.isVisible) {
95+
notifyInfo(
96+
title = message("amazonqFeatureDev.code_generation.notification_title"),
97+
content = message("amazonqFeatureDev.code_generation.notification_message"),
98+
project = getProject(),
99+
notificationActions = listOf(openChatNotificationAction())
100+
)
101+
}
102+
}
103+
}
104+
105+
private fun FeatureDevController.openChatNotificationAction() = NotificationAction.createSimple(
106+
message("amazonqFeatureDev.code_generation.notification_open_link")
107+
) {
108+
toolWindow?.show()
109+
}

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationState.kt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -64,8 +64,8 @@ class CodeGenerationState(
6464
messenger = messenger,
6565
)
6666

67-
// It is not needed to interact right away with the code generation state.
68-
// returns a SessionStateInteraction object to be handled by the controller.
67+
// It is not needed to interact right away with the PrepareCodeGeneration.
68+
// returns therefore a SessionStateInteraction object to be handled by the controller.
6969
return SessionStateInteraction(
7070
nextState = nextState,
7171
interaction = Interaction(content = "", interactionSucceeded = true)

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/Session.kt‎

Lines changed: 5 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -11,14 +11,12 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.CODE_GENERATIO
1111
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.clients.FeatureDevClient
1212
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.conversationIdNotFound
1313
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.sendAsyncEventProgress
14-
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.messages.updateFileComponent
1514
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.createConversation
15+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.resolveAndCreateOrUpdateFile
16+
import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.util.resolveAndDeleteFile
1617
import software.aws.toolkits.jetbrains.services.cwc.controller.ReferenceLogController
1718
import software.aws.toolkits.jetbrains.services.cwc.messages.CodeReference
1819
import software.aws.toolkits.telemetry.AmazonqTelemetry
19-
import kotlin.io.path.createDirectories
20-
import kotlin.io.path.deleteIfExists
21-
import kotlin.io.path.writeBytes
2220

2321
class Session(val tabID: String, val project: Project) {
2422
var context: FeatureDevSessionContext
@@ -91,35 +89,15 @@ class Session(val tabID: String, val project: Project) {
9189
AmazonqTelemetry.isApproachAccepted(amazonqConversationId = conversationId, enabled = true)
9290
}
9391

94-
suspend fun updateFilesPaths(
95-
messenger: MessagePublisher,
96-
tabId: String,
97-
filePaths: List<NewFileZipInfo>,
98-
deletedFiles: List<DeletedFileInfo>
99-
) {
100-
messenger.updateFileComponent(tabId, filePaths, deletedFiles)
101-
}
102-
10392
/**
10493
* Triggered by the Insert code follow-up button to apply code changes.
10594
*/
10695
fun insertChanges(filePaths: List<NewFileZipInfo>, deletedFiles: List<DeletedFileInfo>, references: List<CodeReference>) {
10796
val projectRootPath = context.projectRoot.toNioPath()
10897

109-
filePaths
110-
.filterNot { it.rejected }
111-
.forEach {
112-
val filePath = projectRootPath.resolve(it.zipFilePath)
113-
filePath.parent.createDirectories() // Create directories if needed
114-
filePath.writeBytes(it.fileContent.toByteArray(Charsets.UTF_8))
115-
}
98+
filePaths.forEach { resolveAndCreateOrUpdateFile(projectRootPath, it.zipFilePath, it.fileContent) }
11699

117-
deletedFiles
118-
.filterNot { it.rejected }
119-
.forEach {
120-
val deleteFilePath = projectRootPath.resolve(it.zipFilePath)
121-
deleteFilePath.deleteIfExists()
122-
}
100+
deletedFiles.forEach { resolveAndDeleteFile(projectRootPath, it.zipFilePath) }
123101

124102
ReferenceLogController.addReferenceLog(references, project)
125103

@@ -170,6 +148,7 @@ class Session(val tabID: String, val project: Project) {
170148
return _conversationId as String
171149
}
172150
}
151+
173152
val sessionState: SessionState
174153
get() {
175154
if (_state == null) {

0 commit comments

Comments
 (0)