fix(dev): update tests

tverney · tverney · commit 42b9f81dc26e · 2024-10-16T15:32:12.000-04:00
diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationState.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationState.kt
@@ -38,7 +38,7 @@ class CodeGenerationState(
     val messenger: MessagePublisher,
     override var codeGenerationRemainingIterationCount: Int? = null,
     override var codeGenerationTotalIterationCount: Int? = null,
-    var currentCodeGenerationId: UUID? = null,
+    var currentCodeGenerationId: String? = "EMPTY_CURRENT_CODE_GENERATION_ID",
     override var token: CancellationTokenSource?,
 ) : SessionState {
     override val phase = SessionStatePhase.CODEGEN
@@ -64,7 +64,7 @@ class CodeGenerationState(
                 )
 
             if (action.token?.token()?.isCancellationRequested != true) {
-                this.currentCodeGenerationId = codeGenerationId
+                this.currentCodeGenerationId = codeGenerationId.toString()
             }
 
             messenger.sendAnswerPart(
diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/PrepareCodeGenerationState.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/PrepareCodeGenerationState.kt
@@ -57,7 +57,7 @@ class PrepareCodeGenerationState(
                 config.conversationId,
                 zipFileChecksum,
                 zipFileLength,
-                uploadId
+                uploadId.toString()
             )
 
             uploadArtifactToS3(uploadUrlResponse.uploadUrl(), fileToUpload, zipFileChecksum, zipFileLength, uploadUrlResponse.kmsKeyArn())
diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FeatureDevService.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FeatureDevService.kt
@@ -31,7 +31,6 @@ import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.Export
 import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.getStartUrl
 import software.aws.toolkits.telemetry.AmazonqTelemetry
 import software.aws.toolkits.telemetry.Result
-import java.util.UUID
 
 private val logger = getLogger<FeatureDevClient>()
 
@@ -83,15 +82,15 @@ class FeatureDevService(val proxyClient: FeatureDevClient, val project: Project)
         }
     }
 
-    fun createUploadUrl(conversationId: String, contentChecksumSha256: String, contentLength: Long, uploadId: UUID):
+    fun createUploadUrl(conversationId: String, contentChecksumSha256: String, contentLength: Long, uploadId: String):
         CreateUploadUrlResponse {
         try {
             logger.debug { "Executing createUploadUrl with conversationId $conversationId" }
             val uploadUrlResponse = proxyClient.createTaskAssistUploadUrl(
                 conversationId,
                 contentChecksumSha256,
                 contentLength,
-                uploadId.toString()
+                uploadId
             )
             logger.debug {
                 "$FEATURE_NAME: Created upload url: {uploadId: $uploadId, requestId: ${uploadUrlResponse.responseMetadata().requestId()}}"
diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevTestBase.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevTestBase.kt
@@ -42,7 +42,6 @@ import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRul
 import software.aws.toolkits.jetbrains.utils.rules.addModule
 import java.io.File
 import java.time.Instant
-import java.util.UUID
 
 open class FeatureDevTestBase(
     @Rule @JvmField
@@ -61,7 +60,6 @@ open class FeatureDevTestBase(
     internal val testConversationId = "1234"
     internal val userMessage = "test-user-message"
     internal val codeGenerationId = "1234"
-    internal val currentCodeGenerationId = UUID.randomUUID()
     internal val testUploadId = "5678"
     internal val testRepositorySize = 20.0 // Picked a random size
     internal val otherStatus = "Other"
@@ -71,55 +69,90 @@ open class FeatureDevTestBase(
     internal val testReferences = listOf(CodeReferenceGenerated())
     internal val testChecksumSha = "test-sha"
     internal val testContentLength: Long = 40
-    internal val uploadId = UUID.randomUUID()
-
-    internal val exampleSendTelemetryEventResponse = SendTelemetryEventResponse.builder()
-        .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)))
-        .build() as SendTelemetryEventResponse
-
-    internal val exampleCreateTaskAssistConversationResponse = CreateTaskAssistConversationResponse.builder()
-        .conversationId(testConversationId)
-        .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)))
-        .build() as CreateTaskAssistConversationResponse
-
-    internal val exampleCreateUploadUrlResponse = CreateUploadUrlResponse.builder()
-        .uploadUrl("https://smth.com")
-        .uploadId(testUploadId)
-        .kmsKeyArn("0000000000000000000000000000000000:key/1234abcd")
-        .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)))
-        .build() as CreateUploadUrlResponse
-
-    internal val exampleStartTaskAssistConversationResponse = StartTaskAssistCodeGenerationResponse.builder()
-        .conversationId(testConversationId)
-        .codeGenerationId(codeGenerationId)
-        .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)))
-        .build() as StartTaskAssistCodeGenerationResponse
-
-    internal val exampleGetTaskAssistConversationResponse = GetTaskAssistCodeGenerationResponse.builder()
-        .conversationId(testConversationId)
-        .codeGenerationStatus(CodeGenerationStatus.builder().status(CodeGenerationWorkflowStatus.IN_PROGRESS).currentStage("InProgress").build())
-        .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)))
-        .build() as GetTaskAssistCodeGenerationResponse
-
-    internal val exampleCompleteGetTaskAssistCodeGenerationResponse = GetTaskAssistCodeGenerationResponse.builder()
-        .conversationId(testConversationId)
-        .codeGenerationStatus(CodeGenerationStatus.builder().status(CodeGenerationWorkflowStatus.COMPLETE).currentStage("Complete").build())
-        .codeGenerationRemainingIterationCount(2)
-        .codeGenerationTotalIterationCount(3)
-        .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)))
-        .build() as GetTaskAssistCodeGenerationResponse
-
-    internal val exampleFailedGetTaskAssistCodeGenerationResponse = GetTaskAssistCodeGenerationResponse.builder()
-        .conversationId(testConversationId)
-        .codeGenerationStatus(CodeGenerationStatus.builder().status(CodeGenerationWorkflowStatus.FAILED).currentStage("Failed").build())
-        .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)))
-        .build() as GetTaskAssistCodeGenerationResponse
-
-    internal val exampleOtherGetTaskAssistCodeGenerationResponse = GetTaskAssistCodeGenerationResponse.builder()
-        .conversationId(testConversationId)
-        .codeGenerationStatus(CodeGenerationStatus.builder().status(CodeGenerationWorkflowStatus.UNKNOWN_TO_SDK_VERSION).currentStage(otherStatus).build())
-        .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)))
-        .build() as GetTaskAssistCodeGenerationResponse
+
+    internal val exampleSendTelemetryEventResponse =
+        SendTelemetryEventResponse
+            .builder()
+            .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)))
+            .build() as SendTelemetryEventResponse
+
+    internal val exampleCreateTaskAssistConversationResponse =
+        CreateTaskAssistConversationResponse
+            .builder()
+            .conversationId(testConversationId)
+            .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)))
+            .build() as CreateTaskAssistConversationResponse
+
+    internal val exampleCreateUploadUrlResponse =
+        CreateUploadUrlResponse
+            .builder()
+            .uploadUrl("https://smth.com")
+            .uploadId(testUploadId)
+            .kmsKeyArn("0000000000000000000000000000000000:key/1234abcd")
+            .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)))
+            .build() as CreateUploadUrlResponse
+
+    internal val exampleStartTaskAssistConversationResponse =
+        StartTaskAssistCodeGenerationResponse
+            .builder()
+            .conversationId(testConversationId)
+            .codeGenerationId(codeGenerationId)
+            .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)))
+            .build() as StartTaskAssistCodeGenerationResponse
+
+    internal val exampleGetTaskAssistConversationResponse =
+        GetTaskAssistCodeGenerationResponse
+            .builder()
+            .conversationId(testConversationId)
+            .codeGenerationStatus(
+                CodeGenerationStatus
+                    .builder()
+                    .status(CodeGenerationWorkflowStatus.IN_PROGRESS)
+                    .currentStage("InProgress")
+                    .build(),
+            ).responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)))
+            .build() as GetTaskAssistCodeGenerationResponse
+
+    internal val exampleCompleteGetTaskAssistCodeGenerationResponse =
+        GetTaskAssistCodeGenerationResponse
+            .builder()
+            .conversationId(testConversationId)
+            .codeGenerationStatus(
+                CodeGenerationStatus
+                    .builder()
+                    .status(CodeGenerationWorkflowStatus.COMPLETE)
+                    .currentStage("Complete")
+                    .build(),
+            ).codeGenerationRemainingIterationCount(2)
+            .codeGenerationTotalIterationCount(3)
+            .responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)))
+            .build() as GetTaskAssistCodeGenerationResponse
+
+    internal val exampleFailedGetTaskAssistCodeGenerationResponse =
+        GetTaskAssistCodeGenerationResponse
+            .builder()
+            .conversationId(testConversationId)
+            .codeGenerationStatus(
+                CodeGenerationStatus
+                    .builder()
+                    .status(CodeGenerationWorkflowStatus.FAILED)
+                    .currentStage("Failed")
+                    .build(),
+            ).responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)))
+            .build() as GetTaskAssistCodeGenerationResponse
+
+    internal val exampleOtherGetTaskAssistCodeGenerationResponse =
+        GetTaskAssistCodeGenerationResponse
+            .builder()
+            .conversationId(testConversationId)
+            .codeGenerationStatus(
+                CodeGenerationStatus
+                    .builder()
+                    .status(CodeGenerationWorkflowStatus.UNKNOWN_TO_SDK_VERSION)
+                    .currentStage(otherStatus)
+                    .build(),
+            ).responseMetadata(DefaultAwsResponseMetadata.create(mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)))
+            .build() as GetTaskAssistCodeGenerationResponse
 
     internal val exampleExportResultArchiveResponse = mutableListOf(byteArrayOf(100))
 
@@ -130,18 +163,22 @@ open class FeatureDevTestBase(
         project = projectRule.project
         toolkitConnectionManager = spy(ToolkitConnectionManager.getInstance(project))
         val accessToken = DeviceAuthorizationGrantToken(aString(), aString(), aString(), aString(), Instant.MAX, Instant.now())
-        val provider = mock<BearerTokenProvider> {
-            doReturn(accessToken).whenever(it).refresh()
-        }
-        val mockBearerProvider = mock<ToolkitBearerTokenProvider> {
-            doReturn(provider).whenever(it).delegate
-        }
-        val connectionSettingsMock = mock<TokenConnectionSettings> {
-            whenever(it.tokenProvider).thenReturn(mockBearerProvider)
-        }
-        val toolkitConnection = mock<AwsBearerTokenConnection> {
-            doReturn(connectionSettingsMock).whenever(it).getConnectionSettings()
-        }
+        val provider =
+            mock<BearerTokenProvider> {
+                doReturn(accessToken).whenever(it).refresh()
+            }
+        val mockBearerProvider =
+            mock<ToolkitBearerTokenProvider> {
+                doReturn(provider).whenever(it).delegate
+            }
+        val connectionSettingsMock =
+            mock<TokenConnectionSettings> {
+                whenever(it.tokenProvider).thenReturn(mockBearerProvider)
+            }
+        val toolkitConnection =
+            mock<AwsBearerTokenConnection> {
+                doReturn(connectionSettingsMock).whenever(it).getConnectionSettings()
+            }
         doReturn(toolkitConnection).whenever(toolkitConnectionManager).activeConnectionForFeature(any())
         project.replaceService(ToolkitConnectionManager::class.java, toolkitConnectionManager, disposableRule.disposable)
         clientAdaptorSpy = spy(FeatureDevClient.getInstance(project))
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
@@ -225,7 +225,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
                 PrepareCodeGenerationState(
                     testTabId,
                     GradleConnector.newCancellationTokenSource(),
-                    mock(),
+                    "test-command",
                     sessionStateConfig,
                     newFileContents,
                     deletedFiles,
@@ -234,9 +234,10 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
                     0,
                     messenger,
                     0,
-                    0
+                    0,
                 ),
             )
+
             doNothing().`when`(spySession).insertChanges(any(), any(), any())
 
             spySession.preloader(userMessage, messenger)
diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationStateTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/CodeGenerationStateTest.kt
@@ -48,16 +48,18 @@ class CodeGenerationStateTest : FeatureDevTestBase() {
         val repoContext = mock<FeatureDevSessionContext>()
         val sessionStateConfig = SessionStateConfig(testConversationId, repoContext, featureDevService)
 
-        codeGenerationState = CodeGenerationState(
-            testTabId,
-            "",
-            sessionStateConfig,
-            testUploadId,
-            0,
-            testRepositorySize,
-            messenger,
-            token = null
-        )
+        codeGenerationState =
+            CodeGenerationState(
+                testTabId,
+                "",
+                sessionStateConfig,
+                testUploadId,
+                0,
+                testRepositorySize,
+                messenger,
+                token = null,
+                currentCodeGenerationId = "EMPTY_CURRENT_CODE_GENERATION_ID",
+            )
 
         mockkStatic(MessagePublisher::sendAnswerPart)
         coEvery { messenger.sendAnswerPart(any(), any()) } just runs
@@ -80,13 +82,12 @@ class CodeGenerationStateTest : FeatureDevTestBase() {
             val actual = codeGenerationState.interact(action)
             assertThat(actual.nextState).isInstanceOf(PrepareCodeGenerationState::class.java)
             val nextState = actual.nextState as PrepareCodeGenerationState
-            assertThat(nextState.uploadId).isEqualTo(testUploadId)
             assertThat(nextState.phase).isEqualTo(SessionStatePhase.CODEGEN)
             assertThat(nextState.filePaths).isEqualTo(
-                listOf(NewFileZipInfo("test.ts", "This is a comment", false))
+                listOf(NewFileZipInfo("test.ts", "This is a comment", false)),
             )
             assertThat(nextState.deletedFiles).isEqualTo(
-                listOf(DeletedFileInfo("deleted.ts", false))
+                listOf(DeletedFileInfo("deleted.ts", false)),
             )
             assertThat(nextState.references).isEqualTo(testReferences)
             assertThat(nextState.codeGenerationRemainingIterationCount).isEqualTo(2)
@@ -96,26 +97,22 @@ class CodeGenerationStateTest : FeatureDevTestBase() {
         }
         assertThat(codeGenerationState.phase).isEqualTo(SessionStatePhase.CODEGEN)
         coVerify(exactly = 1) { messenger.sendAnswerPart(testTabId, message("amazonqFeatureDev.code_generation.generating_code")) }
-        verify(
-            exactly = 1
-        ) { featureDevService.startTaskAssistCodeGeneration(testConversationId, testUploadId, userMessage, codeGenerationId, null) }
-        verify(exactly = 1) { featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId.toString()) }
+
+        verify(exactly = 1) { featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId) }
         coVerify(exactly = 1) { featureDevService.exportTaskAssistArchiveResult(testConversationId) }
     }
 
     @Test(expected = FeatureDevException::class)
-    fun `test code generation failed`() = runTest {
-        every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any(), any(), any()) } returns exampleStartTaskAssistConversationResponse
-        every { featureDevService.getTaskAssistCodeGeneration(any(), any()) } returns exampleFailedGetTaskAssistCodeGenerationResponse
+    fun `test code generation failed`() =
+        runTest {
+            every { featureDevService.startTaskAssistCodeGeneration(any(), any(), any(), any(), any()) } returns exampleStartTaskAssistConversationResponse
+            every { featureDevService.getTaskAssistCodeGeneration(any(), any()) } returns exampleFailedGetTaskAssistCodeGenerationResponse
 
-        codeGenerationState.interact(action)
+            codeGenerationState.interact(action)
 
-        verify(
-            exactly = 1
-        ) { featureDevService.startTaskAssistCodeGeneration(testConversationId, testUploadId, userMessage, codeGenerationId, null) }
-        verify(exactly = 1) { featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId.toString()) }
-        coVerify(exactly = 0) { featureDevService.exportTaskAssistArchiveResult(any()) }
-    }
+            verify(exactly = 1) { featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId) }
+            coVerify(exactly = 0) { featureDevService.exportTaskAssistArchiveResult(any()) }
+        }
 
     @Test
     fun `test code generation returns any other handled status`() {
@@ -128,10 +125,7 @@ class CodeGenerationStateTest : FeatureDevTestBase() {
             }
         }.isExactlyInstanceOf(IllegalStateException::class.java).withFailMessage("Unknown status: $otherStatus")
 
-        verify(
-            exactly = 1
-        ) { featureDevService.startTaskAssistCodeGeneration(testConversationId, testUploadId, userMessage, codeGenerationId, null) }
-        verify(exactly = 1) { featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId.toString()) }
+        verify(exactly = 1) { featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId) }
         coVerify(exactly = 0) { featureDevService.exportTaskAssistArchiveResult(any()) }
     }
 
@@ -147,10 +141,7 @@ class CodeGenerationStateTest : FeatureDevTestBase() {
             codeGenerationState.interact(action)
         }
 
-        verify(
-            exactly = 1
-        ) { featureDevService.startTaskAssistCodeGeneration(testConversationId, testUploadId, userMessage, codeGenerationId, null) }
-        verify(exactly = 2) { featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId.toString()) }
+        verify(exactly = 2) { featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId) }
         coVerify(exactly = 1) { featureDevService.exportTaskAssistArchiveResult(testConversationId) }
     }
 
@@ -172,10 +163,7 @@ class CodeGenerationStateTest : FeatureDevTestBase() {
             assertThat(actual.interaction.content).isEqualTo("")
         }
 
-        verify(
-            exactly = 1
-        ) { featureDevService.startTaskAssistCodeGeneration(testConversationId, testUploadId, userMessage, codeGenerationId, null) }
-        verify(exactly = 180) { featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId.toString()) }
+        verify(exactly = 180) { featureDevService.getTaskAssistCodeGeneration(testConversationId, codeGenerationId) }
         coVerify(exactly = 0) { featureDevService.exportTaskAssistArchiveResult(testConversationId) }
     }
 }
diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/PrepareCodeGenerationStateTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/session/PrepareCodeGenerationStateTest.kt
diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FeatureDevServiceTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/util/FeatureDevServiceTest.kt

Original file line number	Diff line number	Diff line change
`@@ -57,7 +57,7 @@ class PrepareCodeGenerationState(`
`57`	`57`	`config.conversationId,`
`58`	`58`	`zipFileChecksum,`
`59`	`59`	`zipFileLength,`
`60`		`- uploadId`
	`60`	`+ uploadId.toString()`
`61`	`61`	`)`
`62`	`62`
`63`	`63`	`uploadArtifactToS3(uploadUrlResponse.uploadUrl(), fileToUpload, zipFileChecksum, zipFileLength, uploadUrlResponse.kmsKeyArn())`