Skip to content

Commit 382c674

Browse files
andrewyuqrli
andrewyuq
and
rli
authored
CodeWhipserer completion integration tests (#3481)
* CodeWhipserer completion integration tests * Add CodeScan Integration tests * Add Reference Tracker integration tests --------- Co-authored-by: Richard Li <[email protected]>
1 parent 3d5ff59 commit 382c674

File tree

15 files changed

+505
-22
lines changed

15 files changed

+505
-22
lines changed

‎jetbrains-core/it/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererCodeScanIntegrationTest.kt‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
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
5+
6+
import com.intellij.testFramework.runInEdtAndWait
7+
import org.assertj.core.api.Assertions.assertThat
8+
import org.junit.Test
9+
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.cppFileName
10+
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.cppTestLeftContext
11+
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.sessionconfig.CodeScanSessionConfig
12+
import software.aws.toolkits.jetbrains.utils.rules.RunWithRealCredentials.RequiresRealCredentials
13+
import software.aws.toolkits.resources.message
14+
15+
@RequiresRealCredentials
16+
class CodeWhispererCodeScanIntegrationTest : CodeWhispererIntegrationTestBase() {
17+
private val filePromptWithSecurityIssues = """
18+
from flask import app
19+
20+
@app.route('/')
21+
def execute_input_noncompliant():
22+
from flask import request
23+
module_version = request.args.get("module_version")
24+
# Noncompliant: executes unsanitized inputs.
25+
exec("import urllib%s as urllib" % module_version)
26+
27+
@app.route('/')
28+
def execute_input_compliant():
29+
from flask import request
30+
module_version = request.args.get("module_version")
31+
# Compliant: executes sanitized inputs.
32+
exec("import urllib%d as urllib" % int(module_version))
33+
""".trimIndent()
34+
35+
@Test
36+
fun testCodeScanValidWithIssues() {
37+
projectRule.fixture.addFileToProject("test2.py", filePromptWithSecurityIssues)
38+
val response = runCodeScan()
39+
assertThat(response.issues.size).isEqualTo(3)
40+
assertThat(response.responseContext.codeScanTotalIssues).isEqualTo(3)
41+
assertThat(response.responseContext.codeScanJobId).isNotNull
42+
}
43+
44+
@Test
45+
fun testCodeScanValidWithNoIssues() {
46+
val response = runCodeScan()
47+
assertThat(response.issues.size).isEqualTo(0)
48+
assertThat(response.responseContext.codeScanTotalIssues).isEqualTo(0)
49+
assertThat(response.responseContext.codeScanJobId).isNotNull
50+
}
51+
52+
@Test
53+
fun testCodeScanFileTooLarge() {
54+
val largePrompt = "a".repeat(1024 * 300)
55+
val file = projectRule.fixture.addFileToProject("test2.py", largePrompt)
56+
runInEdtAndWait {
57+
projectRule.fixture.openFileInEditor(file.virtualFile)
58+
}
59+
testCodeScanWithErrorMessage(
60+
message(
61+
"codewhisperer.codescan.file_too_large",
62+
CodeScanSessionConfig.create(file.virtualFile, projectRule.project).getPresentablePayloadLimit()
63+
)
64+
)
65+
}
66+
67+
@Test
68+
fun testCodeScanUnsupportedLanguage() {
69+
val file = projectRule.fixture.addFileToProject(cppFileName, cppTestLeftContext)
70+
runInEdtAndWait {
71+
projectRule.fixture.openFileInEditor(file.virtualFile)
72+
}
73+
testCodeScanWithErrorMessage(
74+
message("codewhisperer.codescan.file_ext_not_supported", file.virtualFile.extension ?: "")
75+
)
76+
}
77+
}

‎jetbrains-core/it/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererCodeScanJavaIntegrationTest.kt‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
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
5+
6+
import com.intellij.testFramework.runInEdtAndWait
7+
import org.junit.Test
8+
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.javaTestContext
9+
import software.aws.toolkits.jetbrains.utils.rules.HeavyJavaCodeInsightTestFixtureRule
10+
import software.aws.toolkits.jetbrains.utils.rules.RunWithRealCredentials.RequiresRealCredentials
11+
import software.aws.toolkits.jetbrains.utils.rules.addClass
12+
import software.aws.toolkits.jetbrains.utils.rules.addModule
13+
import software.aws.toolkits.resources.message
14+
15+
@RequiresRealCredentials
16+
class CodeWhispererCodeScanJavaIntegrationTest : CodeWhispererIntegrationTestBase(HeavyJavaCodeInsightTestFixtureRule()) {
17+
@Test
18+
fun testCodeScanJavaProjectNoBuild() {
19+
projectRule as HeavyJavaCodeInsightTestFixtureRule
20+
val module = projectRule.fixture.addModule("main")
21+
22+
val psiClass = projectRule.fixture.addClass(module, javaTestContext)
23+
runInEdtAndWait {
24+
projectRule.fixture.openFileInEditor(psiClass.containingFile.virtualFile)
25+
}
26+
testCodeScanWithErrorMessage(message("codewhisperer.codescan.build_artifacts_not_found"))
27+
}
28+
}

‎jetbrains-core/it/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererCompletionIntegrationTest.kt‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
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
5+
6+
import org.assertj.core.api.Assertions.assertThat
7+
import org.junit.Test
8+
import org.junit.jupiter.api.assertDoesNotThrow
9+
import org.mockito.kotlin.any
10+
import org.mockito.kotlin.never
11+
import org.mockito.kotlin.verify
12+
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.cppFileName
13+
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
14+
import software.aws.toolkits.jetbrains.utils.rules.RunWithRealCredentials.RequiresRealCredentials
15+
import software.aws.toolkits.resources.message
16+
17+
@RequiresRealCredentials
18+
class CodeWhispererCompletionIntegrationTest : CodeWhispererIntegrationTestBase() {
19+
@Test
20+
fun testInvokeCompletionManualTrigger() {
21+
assertDoesNotThrow {
22+
withCodeWhispererServiceInvokedAndWait { response ->
23+
val requestId = response.responseMetadata().requestId()
24+
assertThat(requestId).isNotNull
25+
val sessionId = response.sdkHttpResponse().headers().getOrDefault(
26+
CodeWhispererService.KET_SESSION_ID,
27+
listOf(requestId)
28+
)[0]
29+
assertThat(sessionId).isNotNull
30+
}
31+
}
32+
}
33+
34+
@Test
35+
fun testInvokeCompletionAutoTrigger() {
36+
assertDoesNotThrow {
37+
stateManager.setAutoEnabled(true)
38+
withCodeWhispererServiceInvokedAndWait(false) { response ->
39+
val requestId = response.responseMetadata().requestId()
40+
assertThat(requestId).isNotNull
41+
val sessionId = response.sdkHttpResponse().headers().getOrDefault(
42+
CodeWhispererService.KET_SESSION_ID,
43+
listOf(requestId)
44+
)[0]
45+
assertThat(sessionId).isNotNull
46+
}
47+
}
48+
}
49+
50+
@Test
51+
fun testInvokeCompletionUnsupportedLanguage() {
52+
val file = setFileContext(cppFileName, CodeWhispererTestUtil.cppTestLeftContext, "")
53+
assertDoesNotThrow {
54+
invokeCodeWhispererService()
55+
verify(popupManager, never()).showPopup(any(), any(), any(), any(), any())
56+
verify(clientAdaptor, never()).listRecommendationsPaginator(any(), any())
57+
testMessageShown(message("codewhisperer.language.error", file.fileType.name))
58+
}
59+
}
60+
}

‎jetbrains-core/it/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererIntegrationTestBase.kt‎

Lines changed: 229 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,229 @@
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
5+
6+
import com.intellij.analysis.problemsView.toolWindow.ProblemsView
7+
import com.intellij.openapi.application.ApplicationManager
8+
import com.intellij.openapi.wm.RegisterToolWindowTask
9+
import com.intellij.openapi.wm.ToolWindowManager
10+
import com.intellij.psi.PsiFile
11+
import com.intellij.testFramework.DisposableRule
12+
import com.intellij.testFramework.RuleChain
13+
import com.intellij.testFramework.replaceService
14+
import com.intellij.testFramework.runInEdtAndWait
15+
import kotlinx.coroutines.runBlocking
16+
import org.assertj.core.api.Assertions.assertThat
17+
import org.junit.After
18+
import org.junit.Before
19+
import org.junit.Rule
20+
import org.mockito.kotlin.any
21+
import org.mockito.kotlin.anyOrNull
22+
import org.mockito.kotlin.argumentCaptor
23+
import org.mockito.kotlin.atLeastOnce
24+
import org.mockito.kotlin.doNothing
25+
import org.mockito.kotlin.spy
26+
import org.mockito.kotlin.timeout
27+
import org.mockito.kotlin.times
28+
import org.mockito.kotlin.verify
29+
import org.mockito.kotlin.whenever
30+
import software.amazon.awssdk.services.codewhisperer.model.ListRecommendationsResponse
31+
import software.aws.toolkits.core.TokenConnectionSettings
32+
import software.aws.toolkits.core.credentials.ToolkitBearerTokenProvider
33+
import software.aws.toolkits.jetbrains.core.MockClientManager
34+
import software.aws.toolkits.jetbrains.core.credentials.ToolkitAuthManager
35+
import software.aws.toolkits.jetbrains.core.credentials.loginSso
36+
import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_URL
37+
import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider
38+
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.codeWhispererRecommendationActionId
39+
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonFileName
40+
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonTestLeftContext
41+
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeScanResponse
42+
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanException
43+
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanIssue
44+
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanManager
45+
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
46+
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExploreActionState
47+
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExploreStateType
48+
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
49+
import software.aws.toolkits.jetbrains.services.codewhisperer.model.CodeScanTelemetryEvent
50+
import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext
51+
import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererPopupManager
52+
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererInvocationStatus
53+
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
54+
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererConfiguration
55+
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererConfigurationType
56+
import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererSettings
57+
import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererTelemetryService
58+
import software.aws.toolkits.jetbrains.utils.rules.CodeInsightTestFixtureRule
59+
import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
60+
import software.aws.toolkits.jetbrains.utils.rules.RunWithRealCredentials
61+
62+
open class CodeWhispererIntegrationTestBase(val projectRule: CodeInsightTestFixtureRule = PythonCodeInsightTestFixtureRule()) {
63+
64+
private val realCredentials = RunWithRealCredentials(projectRule)
65+
internal val disposableRule = DisposableRule()
66+
67+
@Rule
68+
@JvmField
69+
val ruleChain = RuleChain(projectRule, disposableRule, realCredentials)
70+
71+
protected lateinit var popupManager: CodeWhispererPopupManager
72+
protected lateinit var clientAdaptor: CodeWhispererClientAdaptor
73+
protected lateinit var stateManager: CodeWhispererExplorerActionManager
74+
protected lateinit var codewhispererService: CodeWhispererService
75+
protected lateinit var settingsManager: CodeWhispererSettings
76+
private lateinit var originalExplorerActionState: CodeWhispererExploreActionState
77+
private lateinit var originalSettings: CodeWhispererConfiguration
78+
internal lateinit var scanManager: CodeWhispererCodeScanManager
79+
protected lateinit var telemetryServiceSpy: CodeWhispererTelemetryService
80+
81+
@Suppress("UnreachableCode")
82+
@Before
83+
open fun setUp() {
84+
MockClientManager.useRealImplementations(disposableRule.disposable)
85+
86+
loginSso(projectRule.project, SONO_URL)
87+
val connectionId = ToolkitBearerTokenProvider.ssoIdentifier(SONO_URL)
88+
val connection = ToolkitAuthManager.getInstance().getConnection(connectionId) ?: return
89+
val tokenProvider = (connection.getConnectionSettings() as TokenConnectionSettings).tokenProvider.delegate as BearerTokenProvider
90+
tokenProvider.resolveToken()
91+
92+
ToolWindowManager.getInstance(projectRule.project).registerToolWindow(
93+
RegisterToolWindowTask(id = ProblemsView.ID, canWorkInDumbMode = true)
94+
)
95+
96+
scanManager = spy(CodeWhispererCodeScanManager.getInstance(projectRule.project))
97+
doNothing().whenever(scanManager).addCodeScanUI(any())
98+
projectRule.project.replaceService(CodeWhispererCodeScanManager::class.java, scanManager, disposableRule.disposable)
99+
100+
telemetryServiceSpy = spy(CodeWhispererTelemetryService.getInstance())
101+
ApplicationManager.getApplication().replaceService(CodeWhispererTelemetryService::class.java, telemetryServiceSpy, disposableRule.disposable)
102+
103+
stateManager = CodeWhispererExplorerActionManager.getInstance()
104+
stateManager.setHasAcceptedTermsOfService(true)
105+
stateManager.setAutoEnabled(false)
106+
107+
popupManager = spy(CodeWhispererPopupManager.getInstance())
108+
popupManager.reset()
109+
doNothing().whenever(popupManager).showPopup(any(), any(), any(), any(), any())
110+
ApplicationManager.getApplication().replaceService(CodeWhispererPopupManager::class.java, popupManager, disposableRule.disposable)
111+
112+
codewhispererService = spy(CodeWhispererService.getInstance())
113+
ApplicationManager.getApplication().replaceService(CodeWhispererService::class.java, codewhispererService, disposableRule.disposable)
114+
115+
settingsManager = CodeWhispererSettings.getInstance()
116+
117+
clientAdaptor = spy(CodeWhispererClientAdaptor.getInstance(projectRule.project))
118+
projectRule.project.replaceService(CodeWhispererClientAdaptor::class.java, clientAdaptor, disposableRule.disposable)
119+
120+
originalExplorerActionState = stateManager.state
121+
originalSettings = settingsManager.state
122+
stateManager.loadState(
123+
CodeWhispererExploreActionState().apply {
124+
CodeWhispererExploreStateType.values().forEach {
125+
value[it] = true
126+
}
127+
}
128+
)
129+
settingsManager.loadState(
130+
CodeWhispererConfiguration().apply {
131+
value[CodeWhispererConfigurationType.IsIncludeCodeWithReference] = true
132+
}
133+
)
134+
135+
setFileContext(pythonFileName, pythonTestLeftContext, "")
136+
}
137+
138+
@After
139+
open fun tearDown() {
140+
runInEdtAndWait {
141+
stateManager.loadState(originalExplorerActionState)
142+
settingsManager.loadState(originalSettings)
143+
popupManager.reset()
144+
}
145+
}
146+
147+
fun withCodeWhispererServiceInvokedAndWait(manual: Boolean = true, runnable: (ListRecommendationsResponse) -> Unit) {
148+
val responseCaptor = argumentCaptor<ListRecommendationsResponse>()
149+
val statesCaptor = argumentCaptor<InvocationContext>()
150+
invokeCodeWhispererService(manual)
151+
verify(codewhispererService, timeout(5000).atLeastOnce()).validateResponse(responseCaptor.capture())
152+
val response = responseCaptor.lastValue
153+
verify(popupManager, timeout(5000).atLeastOnce()).showPopup(statesCaptor.capture(), any(), any(), any(), any())
154+
val states = statesCaptor.lastValue
155+
156+
runInEdtAndWait {
157+
try {
158+
runnable(response)
159+
} finally {
160+
CodeWhispererPopupManager.getInstance().closePopup(states.popup)
161+
}
162+
}
163+
}
164+
165+
fun invokeCodeWhispererService(manual: Boolean = true) {
166+
if (manual) {
167+
runInEdtAndWait {
168+
projectRule.fixture.performEditorAction(codeWhispererRecommendationActionId)
169+
}
170+
} else {
171+
runInEdtAndWait {
172+
projectRule.fixture.type('(')
173+
}
174+
}
175+
while (CodeWhispererInvocationStatus.getInstance().hasExistingInvocation()) {
176+
Thread.sleep(10)
177+
}
178+
}
179+
180+
fun setFileContext(filename: String, leftContext: String, rightContext: String): PsiFile {
181+
val file = projectRule.fixture.configureByText(filename, leftContext + rightContext)
182+
runInEdtAndWait {
183+
projectRule.fixture.editor.caretModel.primaryCaret.moveToOffset(leftContext.length)
184+
}
185+
return file
186+
}
187+
188+
fun runCodeScan(success: Boolean = true): CodeScanResponse {
189+
runInEdtAndWait {
190+
projectRule.fixture.performEditorAction(CodeWhispererTestUtil.codeWhispererCodeScanActionId)
191+
}
192+
val issuesCaptor = argumentCaptor<List<CodeWhispererCodeScanIssue>>()
193+
val codeScanEventCaptor = argumentCaptor<CodeScanTelemetryEvent>()
194+
return runBlocking {
195+
var issues = emptyList<CodeWhispererCodeScanIssue>()
196+
if (success) {
197+
verify(scanManager, timeout(60000).atLeastOnce()).renderResponseOnUIThread(issuesCaptor.capture())
198+
issues = issuesCaptor.lastValue
199+
}
200+
verify(telemetryServiceSpy, timeout(60000).atLeastOnce()).sendSecurityScanEvent(codeScanEventCaptor.capture(), anyOrNull())
201+
val codeScanResponseContext = codeScanEventCaptor.lastValue.codeScanResponseContext
202+
CodeScanResponse.Success(issues, codeScanResponseContext)
203+
}
204+
}
205+
206+
fun testCodeScanWithErrorMessage(message: String) {
207+
val response = runCodeScan(success = false)
208+
assertThat(response.issues.size).isEqualTo(0)
209+
assertThat(response.responseContext.codeScanTotalIssues).isEqualTo(0)
210+
assertThat(response.responseContext.codeScanJobId).isNull()
211+
val exceptionCaptor = argumentCaptor<Exception>()
212+
verify(scanManager, atLeastOnce()).handleException(exceptionCaptor.capture())
213+
val e = exceptionCaptor.lastValue
214+
assertThat(e is CodeWhispererCodeScanException).isTrue
215+
assertThat(e.message).isEqualTo(message)
216+
}
217+
218+
fun testMessageShown(message: String, info: Boolean = true) {
219+
val messageCaptor = argumentCaptor<String>()
220+
if (info) {
221+
verify(codewhispererService, timeout(5000).times(1))
222+
.showCodeWhispererInfoHint(any(), messageCaptor.capture())
223+
} else {
224+
verify(codewhispererService, timeout(5000).times(1))
225+
.showCodeWhispererErrorHint(any(), messageCaptor.capture())
226+
}
227+
assertThat(messageCaptor.lastValue).isEqualTo(message)
228+
}
229+
}

0 commit comments

Comments
 (0)