From fe8b376900a85437fe4519e3257a806240fdae2f Mon Sep 17 00:00:00 2001 From: Zhuowen Chen Date: Thu, 16 Jan 2025 11:35:13 -0800 Subject: [PATCH 1/2] test(amazonq): update tests for messages at different remaining iteration counts --- .../controller/FeatureDevControllerTest.kt | 115 ++++++++++++++++++ 1 file changed, 115 insertions(+) diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerTest.kt index 10e79dc58f0..5d8a1ba318a 100644 --- a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerTest.kt +++ b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerTest.kt @@ -433,6 +433,121 @@ class FeatureDevControllerTest : FeatureDevTestBase() { } } + @Test + fun `test handleChat onCodeGeneration sends correct add code messages`() = runTest { + val totalIterations = 10 + + for (remainingIterations in 0 until totalIterations) { + val message = if (remainingIterations > 2) { + message("amazonqFeatureDev.code_generation.iteration_counts_ask_to_add_code_or_feedback") + } else if (remainingIterations > 0) { + message( + "amazonqFeatureDev.code_generation.iteration_counts", + remainingIterations, + totalIterations, + ) + } else { + message( + "amazonqFeatureDev.code_generation.iteration_counts_ask_to_add_code", + remainingIterations, + totalIterations, + ) + } + val mockSession = mock() + val featureDevService = mockk() + val repoContext = mock() + val sessionStateConfig = SessionStateConfig(testConversationId, repoContext, featureDevService) + val mockInteraction = mock() + whenever(mockSession.send(userMessage)).thenReturn(mockInteraction) + whenever(mockSession.sessionState).thenReturn( + PrepareCodeGenerationState( + testTabId, + CancellationTokenSource(), + "test-command", + sessionStateConfig, + newFileContents, + deletedFiles, + testReferences, + testUploadId, + 1, + messenger, + remainingIterations, + totalIterations, + diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet()), + ), + ) + + controller.onCodeGeneration(mockSession, userMessage, testTabId) + + coVerify { + messenger.sendAnswer( + tabId = testTabId, + messageType = FeatureDevMessageType.Answer, + message = message + ) + } + } + } + + @Test + fun `test handleChat onCodeGeneration sends correct messages after cancellation`() = runTest { + val totalIterations = 10 + + for (remainingIterations in -1 until totalIterations) { + // remainingIterations < 0 to represent the null case + val message = if (remainingIterations > 2 || remainingIterations < 0) { + message("amazonqFeatureDev.code_generation.stopped_code_generation_no_iteration_count_display") + } else if (remainingIterations > 0) { + message( + "amazonqFeatureDev.code_generation.stopped_code_generation", + remainingIterations, + totalIterations, + ) + } else { + message( + "amazonqFeatureDev.code_generation.stopped_code_generation_no_iterations", + remainingIterations, + totalIterations, + ) + } + val mockSession = mock() + val featureDevService = mockk() + val repoContext = mock() + val sessionStateConfig = SessionStateConfig(testConversationId, repoContext, featureDevService) + val mockInteraction = mock() + val token = CancellationTokenSource() + token.cancel() + whenever(mockSession.send(userMessage)).thenReturn(mockInteraction) + whenever(mockSession.sessionState).thenReturn( + PrepareCodeGenerationState( + testTabId, + token, + "test-command", + sessionStateConfig, + newFileContents, + deletedFiles, + testReferences, + testUploadId, + 1, + messenger, + (if (remainingIterations < 0) null else remainingIterations), + totalIterations, + diffMetricsProcessed = DiffMetricsProcessed(HashSet(), HashSet()), + ), + ) + + controller.onCodeGeneration(mockSession, userMessage, testTabId) + + coVerify { + messenger.sendAnswer( + tabId = testTabId, + messageType = FeatureDevMessageType.Answer, + message = message + ) + } + } + } + @Test fun `test handleChat onCodeGeneration sends success metrics`() = runTest { val mockSession = mock() From 6874c6a71217a4a77f73e62df733a637ee21b99a Mon Sep 17 00:00:00 2001 From: Zhuowen Chen Date: Thu, 16 Jan 2025 14:12:52 -0800 Subject: [PATCH 2/2] test(amazonq): update tests for messages at different remaining iteration counts --- .../amazonqFeatureDev/controller/FeatureDevControllerTest.kt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerTest.kt index 5d8a1ba318a..dfba40646fa 100644 --- a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerTest.kt +++ b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerTest.kt @@ -494,7 +494,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() { val totalIterations = 10 for (remainingIterations in -1 until totalIterations) { - // remainingIterations < 0 to represent the null case + // remainingIterations < 0 is to represent the null case val message = if (remainingIterations > 2 || remainingIterations < 0) { message("amazonqFeatureDev.code_generation.stopped_code_generation_no_iteration_count_display") } else if (remainingIterations > 0) {