Skip to content

Commit 3f0651b

Browse files
smoring2Xi Shen
smoring2
and
Xi Shen
authored
Add unit tests for Gumby client (#4045)
* add ut * add ut for upload * temp * add unit test for gumby client * use the jobid in base for the unit test --------- Co-authored-by: Xi Shen <[email protected]>
1 parent 9d7191f commit 3f0651b

File tree

7 files changed

+268
-1
lines changed

7 files changed

+268
-1
lines changed

‎jetbrains-core/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeModernizerSession.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -285,7 +285,7 @@ class CodeModernizerSession(
285285
/**
286286
* Adapted from [CodeWhispererCodeScanSession]
287287
*/
288-
private fun uploadPayload(payload: File): String {
288+
fun uploadPayload(payload: File): String {
289289
val sha256checksum: String = Base64.getEncoder().encodeToString(DigestUtils.sha256(FileInputStream(payload)))
290290
LOG.warn { "About to create an upload url" }
291291
if (isDisposed.get()) {

‎jetbrains-core/tst-resources/codemodernizer/expectedFile

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
expected

‎jetbrains-core/tst-resources/codemodernizer/overwrittenFile

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
overwritten

‎jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/codemodernizer/CodeWhispererCodeModernizerGumbyClientTest.kt

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
1+
// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.codewhisperer.codemodernizer
5+
6+
import com.intellij.testFramework.RuleChain
7+
import com.intellij.testFramework.replaceService
8+
import kotlinx.coroutines.launch
9+
import org.assertj.core.api.Assertions.assertThat
10+
import org.junit.After
11+
import org.junit.Before
12+
import org.junit.Rule
13+
import org.junit.Test
14+
import org.mockito.ArgumentCaptor
15+
import org.mockito.kotlin.any
16+
import org.mockito.kotlin.argumentCaptor
17+
import org.mockito.kotlin.doReturn
18+
import org.mockito.kotlin.mock
19+
import org.mockito.kotlin.stub
20+
import org.mockito.kotlin.verify
21+
import org.mockito.kotlin.verifyNoInteractions
22+
import org.mockito.kotlin.whenever
23+
import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
24+
import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlRequest
25+
import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlResponse
26+
import software.amazon.awssdk.services.codewhispererruntime.model.GetTransformationPlanRequest
27+
import software.amazon.awssdk.services.codewhispererruntime.model.GetTransformationPlanResponse
28+
import software.amazon.awssdk.services.codewhispererruntime.model.GetTransformationRequest
29+
import software.amazon.awssdk.services.codewhispererruntime.model.GetTransformationResponse
30+
import software.amazon.awssdk.services.codewhispererruntime.model.StartTransformationRequest
31+
import software.amazon.awssdk.services.codewhispererruntime.model.StartTransformationResponse
32+
import software.amazon.awssdk.services.codewhispererruntime.model.StopTransformationRequest
33+
import software.amazon.awssdk.services.codewhispererruntime.model.StopTransformationResponse
34+
import software.amazon.awssdk.services.codewhispererruntime.model.TransformationLanguage
35+
import software.amazon.awssdk.services.codewhispererstreaming.CodeWhispererStreamingAsyncClient
36+
import software.amazon.awssdk.services.codewhispererstreaming.model.ExportResultArchiveRequest
37+
import software.amazon.awssdk.services.codewhispererstreaming.model.ExportResultArchiveResponseHandler
38+
import software.amazon.awssdk.services.ssooidc.SsoOidcClient
39+
import software.aws.toolkits.core.TokenConnectionSettings
40+
import software.aws.toolkits.core.utils.test.aString
41+
import software.aws.toolkits.jetbrains.core.MockClientManagerRule
42+
import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
43+
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
44+
import software.aws.toolkits.jetbrains.core.credentials.BearerSsoConnection
45+
import software.aws.toolkits.jetbrains.core.credentials.ManagedSsoProfile
46+
import software.aws.toolkits.jetbrains.core.credentials.MockCredentialManagerRule
47+
import software.aws.toolkits.jetbrains.core.credentials.MockToolkitAuthManagerRule
48+
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
49+
import software.aws.toolkits.jetbrains.services.codemodernizer.client.GumbyClient
50+
import software.aws.toolkits.jetbrains.services.codemodernizer.model.JobId
51+
import software.aws.toolkits.jetbrains.settings.AwsSettings
52+
import java.util.concurrent.CompletableFuture
53+
54+
class CodeWhispererCodeModernizerGumbyClientTest : CodeWhispererCodeModernizerTestBase() {
55+
val mockClientManagerRule = MockClientManagerRule()
56+
val mockCredentialRule = MockCredentialManagerRule()
57+
val authManagerRule = MockToolkitAuthManagerRule()
58+
59+
@Rule
60+
@JvmField
61+
val ruleChain = RuleChain(projectRule, mockCredentialRule, mockClientManagerRule, disposableRule)
62+
63+
private lateinit var bearerClient: CodeWhispererRuntimeClient
64+
private lateinit var streamingBearerClient: CodeWhispererStreamingAsyncClient
65+
private lateinit var ssoClient: SsoOidcClient
66+
67+
private lateinit var gumbyClient: GumbyClient
68+
private lateinit var connectionManager: ToolkitConnectionManager
69+
private var isTelemetryEnabledDefault: Boolean = false
70+
71+
@Before
72+
override fun setup() {
73+
super.setup()
74+
gumbyClient = GumbyClient.getInstance(projectRule.project)
75+
ssoClient = mockClientManagerRule.create()
76+
77+
bearerClient = mockClientManagerRule.create<CodeWhispererRuntimeClient>().stub {
78+
on { createUploadUrl(any<CreateUploadUrlRequest>()) } doReturn exampleCreateUploadUrlResponse
79+
on { getTransformation(any<GetTransformationRequest>()) } doReturn exampleGetCodeMigrationResponse
80+
on { startTransformation(any<StartTransformationRequest>()) } doReturn exampleStartCodeMigrationResponse
81+
on { getTransformationPlan(any<GetTransformationPlanRequest>()) } doReturn exampleGetCodeMigrationPlanResponse
82+
on { stopTransformation(any<StopTransformationRequest>()) } doReturn exampleStopTransformationResponse
83+
}
84+
85+
streamingBearerClient = mockClientManagerRule.create<CodeWhispererStreamingAsyncClient>().stub {
86+
on { exportResultArchive(any<ExportResultArchiveRequest>(), any<ExportResultArchiveResponseHandler>()) } doReturn CompletableFuture()
87+
}
88+
89+
val mockConnection = mock<BearerSsoConnection>()
90+
whenever(mockConnection.getConnectionSettings()) doReturn mock<TokenConnectionSettings>()
91+
92+
connectionManager = mock {
93+
on {
94+
activeConnectionForFeature(any())
95+
} doReturn authManagerRule.createConnection(ManagedSsoProfile("us-east-1", aString(), emptyList())) as AwsBearerTokenConnection
96+
}
97+
projectRule.project.replaceService(ToolkitConnectionManager::class.java, connectionManager, disposableRule.disposable)
98+
99+
isTelemetryEnabledDefault = AwsSettings.getInstance().isTelemetryEnabled
100+
}
101+
102+
@After
103+
fun tearDown() {
104+
AwsSettings.getInstance().isTelemetryEnabled = isTelemetryEnabledDefault
105+
}
106+
107+
@Test
108+
fun `check createUploadUrl`() {
109+
val actual = gumbyClient.createGumbyUploadUrl("test")
110+
argumentCaptor<CreateUploadUrlRequest>().apply {
111+
verify(bearerClient).createUploadUrl(capture())
112+
verifyNoInteractions(streamingBearerClient)
113+
assertThat(actual).isInstanceOf(CreateUploadUrlResponse::class.java)
114+
assertThat(actual).usingRecursiveComparison().comparingOnlyFields("uploadUrl", "uploadId", "kmsKeyArn")
115+
.isEqualTo(exampleCreateUploadUrlResponse)
116+
}
117+
}
118+
119+
@Test
120+
fun `check getCodeModernizationJob`() {
121+
val actual = gumbyClient.getCodeModernizationJob("jobId")
122+
argumentCaptor<GetTransformationRequest>().apply {
123+
verify(bearerClient).getTransformation(capture())
124+
verifyNoInteractions(streamingBearerClient)
125+
assertThat(actual).isInstanceOf(GetTransformationResponse::class.java)
126+
assertThat(actual).usingRecursiveComparison().comparingOnlyFields("jobId", "status", "transformationType", "source")
127+
.isEqualTo(exampleGetCodeMigrationResponse)
128+
}
129+
}
130+
131+
@Test
132+
fun `check startCodeModernization`() {
133+
val actual = gumbyClient.startCodeModernization("jobId", TransformationLanguage.JAVA_8, TransformationLanguage.JAVA_17)
134+
argumentCaptor<StartTransformationRequest>().apply {
135+
verify(bearerClient).startTransformation(capture())
136+
verifyNoInteractions(streamingBearerClient)
137+
assertThat(actual).isInstanceOf(StartTransformationResponse::class.java)
138+
assertThat(actual).usingRecursiveComparison().comparingOnlyFields("transformationJobId").isEqualTo(exampleStartCodeMigrationResponse)
139+
}
140+
}
141+
142+
@Test
143+
fun `check getCodeModernizationPlan`() {
144+
val actual = gumbyClient.getCodeModernizationPlan(JobId("JobId"))
145+
argumentCaptor<GetTransformationPlanRequest>().apply {
146+
verify(bearerClient).getTransformationPlan(capture())
147+
verifyNoInteractions(streamingBearerClient)
148+
assertThat(actual).isInstanceOf(GetTransformationPlanResponse::class.java)
149+
assertThat(actual).usingRecursiveComparison().comparingOnlyFields("transformationSteps").isEqualTo(exampleGetCodeMigrationPlanResponse)
150+
}
151+
}
152+
153+
@Test
154+
fun `check stopTransformation`() {
155+
val actual = gumbyClient.stopTransformation("JobId")
156+
argumentCaptor<StopTransformationRequest>().apply {
157+
verify(bearerClient).stopTransformation(capture())
158+
verifyNoInteractions(streamingBearerClient)
159+
assertThat(actual).isInstanceOf(StopTransformationResponse::class.java)
160+
assertThat(actual).usingRecursiveComparison().comparingOnlyFields("transformationStatus").isEqualTo(exampleStopTransformationResponse)
161+
}
162+
}
163+
164+
@Test
165+
fun `check downloadExportResultArchive`() {
166+
val requestCaptor = ArgumentCaptor.forClass(ExportResultArchiveRequest::class.java)
167+
val handlerCaptor = ArgumentCaptor.forClass(ExportResultArchiveResponseHandler::class.java)
168+
projectCoroutineScope(project).launch {
169+
gumbyClient.downloadExportResultArchive(jobId)
170+
argumentCaptor<ExportResultArchiveRequest, ExportResultArchiveResponseHandler>().apply {
171+
verify(streamingBearerClient).exportResultArchive(requestCaptor.capture(), handlerCaptor.capture())
172+
verifyNoInteractions(bearerClient)
173+
}
174+
}
175+
}
176+
}

‎jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/codemodernizer/CodeWhispererCodeModernizerSessionTest.kt

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,32 +3,51 @@
33

44
package software.aws.toolkits.jetbrains.services.codewhisperer.codemodernizer
55

6+
import com.github.tomakehurst.wiremock.client.WireMock.aResponse
7+
import com.github.tomakehurst.wiremock.client.WireMock.put
8+
import com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo
9+
import com.github.tomakehurst.wiremock.core.WireMockConfiguration
10+
import com.github.tomakehurst.wiremock.junit.WireMockRule
611
import com.intellij.openapi.projectRoots.JavaSdkVersion
712
import com.intellij.openapi.roots.ModuleRootManager
813
import com.intellij.testFramework.runInEdtAndWait
914
import kotlinx.coroutines.runBlocking
15+
import org.apache.commons.codec.digest.DigestUtils
16+
import org.assertj.core.api.Assertions.assertThat
17+
import org.gradle.internal.impldep.com.amazonaws.ResponseMetadata
1018
import org.junit.Assert.assertEquals
1119
import org.junit.Assert.assertFalse
1220
import org.junit.Before
21+
import org.junit.Rule
1322
import org.junit.Test
1423
import org.mockito.Mockito.doReturn
1524
import org.mockito.Mockito.mock
1625
import org.mockito.Mockito.`when`
1726
import org.mockito.kotlin.any
1827
import org.mockito.kotlin.atLeastOnce
1928
import org.mockito.kotlin.doNothing
29+
import org.mockito.kotlin.eq
30+
import org.mockito.kotlin.inOrder
31+
import org.mockito.kotlin.stub
2032
import org.mockito.kotlin.times
2133
import org.mockito.kotlin.verify
2234
import org.mockito.kotlin.verifyNoMoreInteractions
2335
import org.mockito.kotlin.whenever
36+
import software.amazon.awssdk.awscore.DefaultAwsResponseMetadata
37+
import software.amazon.awssdk.http.SdkHttpResponse
38+
import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlResponse
2439
import software.amazon.awssdk.services.codewhispererruntime.model.TransformationStatus
2540
import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeModernizerJobCompletedResult
2641
import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeModernizerSessionContext
2742
import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeModernizerStartJobResult
2843
import software.aws.toolkits.jetbrains.services.codemodernizer.model.ZipCreationResult
44+
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil
45+
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
2946
import software.aws.toolkits.jetbrains.utils.rules.HeavyJavaCodeInsightTestFixtureRule
3047
import software.aws.toolkits.jetbrains.utils.rules.addFileToModule
3148
import java.io.File
49+
import java.io.FileInputStream
50+
import java.util.Base64
3251
import java.util.zip.ZipFile
3352
import kotlin.io.path.Path
3453
import kotlin.test.assertNotNull
@@ -39,9 +58,25 @@ class CodeWhispererCodeModernizerSessionTest : CodeWhispererCodeModernizerTestBa
3958
path.forEach { projectRule.fixture.addFileToModule(module, it, it) }
4059
}
4160

61+
@Rule
62+
@JvmField
63+
val wireMock = WireMockRule(WireMockConfiguration.wireMockConfig().dynamicPort())
64+
65+
lateinit var gumbyUploadUrlResponse: CreateUploadUrlResponse
66+
4267
@Before
4368
override fun setup() {
4469
super.setup()
70+
val s3endpoint = "http://127.0.0.1:${wireMock.port()}"
71+
gumbyUploadUrlResponse = CreateUploadUrlResponse.builder()
72+
.uploadUrl(s3endpoint)
73+
.uploadId("1234")
74+
.kmsKeyArn("0000000000000000000000000000000000:key/1234abcd")
75+
.responseMetadata(DefaultAwsResponseMetadata.create(mapOf(ResponseMetadata.AWS_REQUEST_ID to CodeWhispererTestUtil.testRequestId)))
76+
.sdkHttpResponse(
77+
SdkHttpResponse.builder().headers(mapOf(CodeWhispererService.KET_SESSION_ID to listOf(CodeWhispererTestUtil.testSessionId))).build()
78+
)
79+
.build() as CreateUploadUrlResponse
4580
}
4681

4782
@Test
@@ -279,4 +314,35 @@ class CodeWhispererCodeModernizerSessionTest : CodeWhispererCodeModernizerTestBa
279314
verify(clientAdaptorSpy, times(4)).getCodeModernizationJob(any())
280315
verify(clientAdaptorSpy, atLeastOnce()).getCodeModernizationPlan(any())
281316
}
317+
318+
@Test
319+
fun `overwritten files would have different checksum from expected files`() {
320+
val expectedSha256checksum: String = Base64.getEncoder().encodeToString(
321+
DigestUtils.sha256(FileInputStream(expectedFilePath.toAbsolutePath().toString()))
322+
)
323+
val fakeSha256checksum: String = Base64.getEncoder().encodeToString(
324+
DigestUtils.sha256(FileInputStream(overwrittenFilePath.toAbsolutePath().toString()))
325+
)
326+
assertThat(expectedSha256checksum).isNotEqualTo(fakeSha256checksum)
327+
}
328+
329+
@Test
330+
fun `test uploadPayload()`() {
331+
val expectedSha256checksum: String =
332+
Base64.getEncoder().encodeToString(DigestUtils.sha256(FileInputStream(expectedFilePath.toAbsolutePath().toString())))
333+
clientAdaptorSpy.stub {
334+
onGeneric { clientAdaptorSpy.createGumbyUploadUrl(any()) }
335+
.thenReturn(gumbyUploadUrlResponse)
336+
}
337+
wireMock.stubFor(put(urlEqualTo("/")).willReturn(aResponse().withStatus(200)))
338+
testSessionSpy.uploadPayload(expectedFilePath.toFile())
339+
340+
val inOrder = inOrder(testSessionSpy)
341+
inOrder.verify(testSessionSpy).uploadArtifactToS3(
342+
eq(gumbyUploadUrlResponse.uploadUrl()),
343+
eq(expectedFilePath.toFile()),
344+
eq(expectedSha256checksum),
345+
eq(gumbyUploadUrlResponse.kmsKeyArn())
346+
)
347+
}
282348
}

‎jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/codemodernizer/CodeWhispererCodeModernizerTest.kt

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,11 @@ import software.aws.toolkits.jetbrains.services.codemodernizer.filterOnlyParentF
3030
import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeModernizerArtifact
3131
import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeModernizerStartJobResult
3232
import software.aws.toolkits.jetbrains.services.codemodernizer.model.CustomerSelection
33+
import software.aws.toolkits.jetbrains.services.codemodernizer.model.InvalidTelemetryReason
34+
import software.aws.toolkits.jetbrains.services.codemodernizer.model.ValidationResult
3335
import software.aws.toolkits.jetbrains.services.codemodernizer.unzipFile
36+
import software.aws.toolkits.resources.message
37+
import software.aws.toolkits.telemetry.CodeTransformPreValidationError
3438
import kotlin.io.path.Path
3539
import kotlin.io.path.createTempDirectory
3640
import kotlin.io.path.exists
@@ -141,4 +145,17 @@ class CodeWhispererCodeModernizerTest : CodeWhispererCodeModernizerTestBase() {
141145
codeModernizerManagerSpy.userInitiatedStopCodeModernization()
142146
verify(codeModernizerManagerSpy, times(1)).notifyTransformationStartStopping()
143147
}
148+
149+
@Test
150+
fun `start transformation without IdC connection`() {
151+
val result = codeModernizerManagerSpy.validate(project)
152+
val expectedResult = ValidationResult(
153+
false,
154+
message("codemodernizer.notification.warn.invalid_project.description.reason.not_logged_in"),
155+
InvalidTelemetryReason(
156+
CodeTransformPreValidationError.NonSSOLogin
157+
)
158+
)
159+
assertEquals(expectedResult, result)
160+
}
144161
}

‎jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/codemodernizer/CodeWhispererCodeModernizerTestBase.kt

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUr
3131
import software.amazon.awssdk.services.codewhispererruntime.model.GetTransformationPlanResponse
3232
import software.amazon.awssdk.services.codewhispererruntime.model.GetTransformationResponse
3333
import software.amazon.awssdk.services.codewhispererruntime.model.StartTransformationResponse
34+
import software.amazon.awssdk.services.codewhispererruntime.model.StopTransformationResponse
3435
import software.amazon.awssdk.services.codewhispererruntime.model.TransformationJob
3536
import software.amazon.awssdk.services.codewhispererruntime.model.TransformationLanguage
3637
import software.amazon.awssdk.services.codewhispererruntime.model.TransformationPlan
@@ -92,6 +93,8 @@ open class CodeWhispererCodeModernizerTestBase(
9293
internal val migrationStep = MigrationStep("Test migration step")
9394
internal lateinit var testCodeModernizerArtifact: CodeModernizerArtifact
9495
internal val exampleZipPath = "simple.zip".toResourceFile().toPath()
96+
internal val expectedFilePath = "expectedFile".toResourceFile().toPath()
97+
internal val overwrittenFilePath = "overwrittenFile".toResourceFile().toPath()
9598
internal val validZipPatchDirPath = "patch/"
9699
internal val validZipArtifactsPath = "artifacts/"
97100
internal val validZipSummaryPath = "summary/"
@@ -155,6 +158,9 @@ open class CodeWhispererCodeModernizerTestBase(
155158
.sdkHttpResponse(SdkHttpResponse.builder().headers(mapOf(CodeWhispererService.KET_SESSION_ID to listOf(CodeWhispererTestUtil.testSessionId))).build())
156159
.build() as GetTransformationPlanResponse
157160

161+
internal val exampleStopTransformationResponse = StopTransformationResponse.builder()
162+
.transformationStatus(TransformationStatus.STOPPED).build() as StopTransformationResponse
163+
158164
internal val exampleGetCodeMigrationResponse = GetTransformationResponse.builder()
159165
.transformationJob(
160166
TransformationJob.builder()

0 commit comments

Comments
 (0)