Skip to content

Commit 7eda113

Browse files
Will-ShaoHuaWill-ShaoHua
authored
(feat codewhisperer): dial up java utg support (#3799)
1 parent 0e756b7 commit 7eda113

File tree

3 files changed

+153
-22
lines changed

3 files changed

+153
-22
lines changed

‎.changes/next-release/feature-8d1feae5-8af2-4ca2-91e2-75e0f75159da.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type" : "feature",
3+
"description" : "CodeWhisperer: Improve file context fetching for Java test files"
4+
}

‎jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt‎

Lines changed: 40 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -119,31 +119,27 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
119119
override suspend fun extractSupplementalFileContext(psiFile: PsiFile, targetContext: FileContextInfo): SupplementalContextInfo? {
120120
val startFetchingTimestamp = System.currentTimeMillis()
121121
val isTst = isTestFile(psiFile)
122-
val userGroup = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
123122
val language = targetContext.programmingLanguage
124-
125-
val chunks = if (isTst && targetContext.programmingLanguage.isUTGSupported()) {
126-
if (userGroup == CodeWhispererUserGroup.CrossFile) {
127-
extractSupplementalFileContextForTst(psiFile, targetContext)
128-
} else {
129-
emptyList()
130-
}
131-
} else if (!isTst && targetContext.programmingLanguage.isSupplementalContextSupported()) {
132-
when (language) {
133-
is CodeWhispererJava -> extractSupplementalFileContextForSrc(psiFile, targetContext)
134-
135-
is CodeWhispererPython, is CodeWhispererJavaScript, is CodeWhispererTypeScript, is CodeWhispererJsx, is CodeWhispererTsx ->
136-
if (userGroup == CodeWhispererUserGroup.CrossFile) {
137-
extractSupplementalFileContextForSrc(psiFile, targetContext)
138-
} else {
139-
emptyList()
140-
}
141-
142-
else -> emptyList()
123+
val group = CodeWhispererUserGroupSettings.getInstance().getUserGroup()
124+
125+
val chunks = if (isTst) {
126+
when (shouldFetchUtgContext(language, group)) {
127+
true -> extractSupplementalFileContextForTst(psiFile, targetContext)
128+
false -> emptyList()
129+
null -> {
130+
LOG.debug { "UTG is not supporting ${targetContext.programmingLanguage.languageId}" }
131+
null
132+
}
143133
}
144134
} else {
145-
LOG.debug { "${if (isTst) "UTG" else "CrossFile"} not supported for ${targetContext.programmingLanguage.languageId}" }
146-
null
135+
when (shouldFetchCrossfileContext(language, group)) {
136+
true -> extractSupplementalFileContextForSrc(psiFile, targetContext)
137+
false -> emptyList()
138+
null -> {
139+
LOG.debug { "Crossfile is not supporting ${targetContext.programmingLanguage.languageId}" }
140+
null
141+
}
142+
}
147143
}
148144

149145
return chunks?.let {
@@ -282,5 +278,27 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
282278

283279
companion object {
284280
private val LOG = getLogger<DefaultCodeWhispererFileContextProvider>()
281+
282+
fun shouldFetchUtgContext(language: CodeWhispererProgrammingLanguage, userGroup: CodeWhispererUserGroup): Boolean? {
283+
if (!language.isUTGSupported()) {
284+
return null
285+
}
286+
287+
return when (language) {
288+
is CodeWhispererJava -> true
289+
else -> userGroup == CodeWhispererUserGroup.CrossFile
290+
}
291+
}
292+
293+
fun shouldFetchCrossfileContext(language: CodeWhispererProgrammingLanguage, userGroup: CodeWhispererUserGroup): Boolean? {
294+
if (!language.isSupplementalContextSupported()) {
295+
return null
296+
}
297+
298+
return when (language) {
299+
is CodeWhispererJava -> true
300+
else -> userGroup == CodeWhispererUserGroup.CrossFile
301+
}
302+
}
285303
}
286304
}

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

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,17 @@ import org.mockito.kotlin.mock
2525
import org.mockito.kotlin.spy
2626
import org.mockito.kotlin.times
2727
import org.mockito.kotlin.verify
28+
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererCpp
2829
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererCsharp
30+
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererGo
2931
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava
32+
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJavaScript
33+
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJsx
3034
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererKotlin
35+
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererPython
36+
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererRuby
37+
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererTsx
38+
import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererTypeScript
3139
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroup
3240
import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupSettings
3341
import software.aws.toolkits.jetbrains.services.codewhisperer.util.DefaultCodeWhispererFileContextProvider
@@ -59,6 +67,107 @@ class CodeWhispererFileContextProviderTest {
5967
sut = FileContextProvider.getInstance(project) as DefaultCodeWhispererFileContextProvider
6068
}
6169

70+
@Test
71+
fun `shouldFetchUtgContext - fully support`() {
72+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererJava.INSTANCE, CodeWhispererUserGroup.CrossFile)).isTrue
73+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererJava.INSTANCE, CodeWhispererUserGroup.Control)).isTrue
74+
}
75+
76+
@Test
77+
fun `shouldFetchUtgContext - partially support`() {
78+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererPython.INSTANCE, CodeWhispererUserGroup.CrossFile)).isTrue
79+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererPython.INSTANCE, CodeWhispererUserGroup.Control)).isFalse
80+
}
81+
82+
@Test
83+
fun `shouldFetchUtgContext - no support`() {
84+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererJavaScript.INSTANCE, CodeWhispererUserGroup.Control)).isNull()
85+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererJavaScript.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull()
86+
87+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererJsx.INSTANCE, CodeWhispererUserGroup.Control)).isNull()
88+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererJsx.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull()
89+
90+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererTypeScript.INSTANCE, CodeWhispererUserGroup.Control)).isNull()
91+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererTypeScript.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull()
92+
93+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererTsx.INSTANCE, CodeWhispererUserGroup.Control)).isNull()
94+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererTsx.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull()
95+
96+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererCsharp.INSTANCE, CodeWhispererUserGroup.Control)).isNull()
97+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererCsharp.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull()
98+
99+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererKotlin.INSTANCE, CodeWhispererUserGroup.Control)).isNull()
100+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererKotlin.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull()
101+
102+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererGo.INSTANCE, CodeWhispererUserGroup.Control)).isNull()
103+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererGo.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull()
104+
105+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererTsx.INSTANCE, CodeWhispererUserGroup.Control)).isNull()
106+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchUtgContext(CodeWhispererTsx.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull()
107+
}
108+
109+
@Test
110+
fun `shouldFetchCrossfileContext - fully support`() {
111+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererJava.INSTANCE, CodeWhispererUserGroup.CrossFile)).isTrue
112+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererJava.INSTANCE, CodeWhispererUserGroup.Control)).isTrue
113+
}
114+
115+
@Test
116+
fun `shouldFetchCrossfileContext - partially support`() {
117+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererPython.INSTANCE, CodeWhispererUserGroup.Control)).isFalse
118+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererPython.INSTANCE, CodeWhispererUserGroup.CrossFile)).isTrue
119+
120+
assertThat(
121+
DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(
122+
CodeWhispererJavaScript.INSTANCE,
123+
CodeWhispererUserGroup.Control
124+
)
125+
).isFalse
126+
assertThat(
127+
DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(
128+
CodeWhispererJavaScript.INSTANCE,
129+
CodeWhispererUserGroup.CrossFile
130+
)
131+
).isTrue
132+
133+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererJsx.INSTANCE, CodeWhispererUserGroup.Control)).isFalse
134+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererJsx.INSTANCE, CodeWhispererUserGroup.CrossFile)).isTrue
135+
136+
assertThat(
137+
DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(
138+
CodeWhispererTypeScript.INSTANCE,
139+
CodeWhispererUserGroup.Control
140+
)
141+
).isFalse
142+
assertThat(
143+
DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(
144+
CodeWhispererTypeScript.INSTANCE,
145+
CodeWhispererUserGroup.CrossFile
146+
)
147+
).isTrue
148+
149+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererTsx.INSTANCE, CodeWhispererUserGroup.Control)).isFalse
150+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererTsx.INSTANCE, CodeWhispererUserGroup.CrossFile)).isTrue
151+
}
152+
153+
@Test
154+
fun `shouldFetchCrossfileContext - no support`() {
155+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererCsharp.INSTANCE, CodeWhispererUserGroup.Control)).isNull()
156+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererCsharp.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull()
157+
158+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererKotlin.INSTANCE, CodeWhispererUserGroup.Control)).isNull()
159+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererKotlin.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull()
160+
161+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererGo.INSTANCE, CodeWhispererUserGroup.Control)).isNull()
162+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererGo.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull()
163+
164+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererCpp.INSTANCE, CodeWhispererUserGroup.Control)).isNull()
165+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererCpp.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull()
166+
167+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererRuby.INSTANCE, CodeWhispererUserGroup.Control)).isNull()
168+
assertThat(DefaultCodeWhispererFileContextProvider.shouldFetchCrossfileContext(CodeWhispererRuby.INSTANCE, CodeWhispererUserGroup.CrossFile)).isNull()
169+
}
170+
62171
@Test
63172
fun `languages not supporting supplemental context will return empty`() {
64173
val psiFiles = setupFixture(fixture)

0 commit comments

Comments
 (0)