codewhisperer: enable UTG(Unit test generator) for java/python (#3755)

Author: Will-ShaoHua

.changes/next-release/feature-581591a9-066f-4696-bc5c-1eb8cd01b642.json

@@ -0,0 +1,4 @@
+{
+  "type" : "feature",
+  "description" : "CodeWhisperer: Improve suggestion quality with enhanced file context fetching"
+}

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererJava.kt

@@ -23,6 +23,8 @@ class CodeWhispererJava private constructor() : CodeWhispererProgrammingLanguage
 
     override fun isSupplementalContextSupported() = true
 
+    override fun isUTGSupported() = true
+
     companion object {
         const val ID = "java"
 

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/language/languages/CodeWhispererPython.kt

@@ -21,6 +21,8 @@ class CodeWhispererPython private constructor() : CodeWhispererProgrammingLangua
 
     override fun isAllClassifier(): Boolean = true
 
+    override fun isUTGSupported() = true
+
     companion object {
         const val ID = "python"
 

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt

@@ -37,7 +37,6 @@ import software.amazon.awssdk.services.codewhispererruntime.model.Recommendation
 import software.amazon.awssdk.services.codewhispererruntime.model.SupplementalContext
 import software.amazon.awssdk.services.codewhispererruntime.model.ThrottlingException
 import software.aws.toolkits.core.utils.debug
-import software.aws.toolkits.core.utils.error
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
 import software.aws.toolkits.jetbrains.core.coroutines.disposableCoroutineScope
@@ -512,7 +511,7 @@ class CodeWhispererService {
         // 2. supplemental context
         val startFetchingTimestamp = System.currentTimeMillis()
         val isTstFile = FileContextProvider.getInstance(project).isTestFile(psiFile)
-        val supplementalContext = if (CodeWhispererUserGroupSettings.getInstance().getUserGroup() == CodeWhispererUserGroup.CrossFile && !isTstFile) {
+        val supplementalContext = if (CodeWhispererUserGroupSettings.getInstance().getUserGroup() == CodeWhispererUserGroup.CrossFile) {
             runBlocking {
                 try {
                     withTimeout(SUPPLEMENTAL_CONTEXT_TIMEOUT) {
@@ -530,13 +529,18 @@ class CodeWhispererService {
                             targetFileName = fileContext.filename
                         )
                     } else {
-                        LOG.error(e) { "Run into unexpected error when fetching supplemental context" }
+                        LOG.debug { "Run into unexpected error when fetching supplemental context, error: ${e.message}" }
                         null
                     }
                 }
             }
         } else {
-            null
+            SupplementalContextInfo(
+                isUtg = isTstFile,
+                contents = emptyList(),
+                latency = 0,
+                targetFileName = fileContext.filename
+            )
         }
 
         // 3. caret position

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

@@ -3,11 +3,14 @@
 
 package software.aws.toolkits.jetbrains.services.codewhisperer
 
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiFile
+import com.intellij.testFramework.DisposableRule
 import com.intellij.testFramework.fixtures.JavaCodeInsightTestFixture
+import com.intellij.testFramework.replaceService
 import com.intellij.testFramework.runInEdtAndGet
 import com.intellij.testFramework.runInEdtAndWait
 import kotlinx.coroutines.runBlocking
@@ -17,12 +20,16 @@ import org.junit.Ignore
 import org.junit.Rule
 import org.junit.Test
 import org.mockito.kotlin.any
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
 import org.mockito.kotlin.spy
 import org.mockito.kotlin.times
 import org.mockito.kotlin.verify
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererCsharp
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererKotlin
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroup
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupSettings
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.DefaultCodeWhispererFileContextProvider
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.FileContextProvider
 import software.aws.toolkits.jetbrains.utils.rules.HeavyJavaCodeInsightTestFixtureRule
@@ -35,6 +42,10 @@ class CodeWhispererFileContextProviderTest {
     @Rule
     val projectRule = HeavyJavaCodeInsightTestFixtureRule()
 
+    @JvmField
+    @Rule
+    val disposableRule = DisposableRule()
+
     lateinit var sut: DefaultCodeWhispererFileContextProvider
 
     lateinit var fixture: JavaCodeInsightTestFixture
@@ -213,10 +224,13 @@ class CodeWhispererFileContextProviderTest {
      *          - MainTest.java
      *
      */
-    // TODO: unignore this case once we configure isUtgSupport = true
-    @Ignore
     @Test
     fun `extractSupplementalFileContext from tst file should extract focal file`() {
+        ApplicationManager.getApplication().replaceService(
+            CodeWhispererUserGroupSettings::class.java,
+            mock { on { getUserGroup() } doReturn CodeWhispererUserGroup.CrossFile },
+            disposableRule.disposable
+        )
         val module = fixture.addModule("main")
         fixture.addClass(module, JAVA_MAIN)
 

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

@@ -86,7 +86,7 @@ class CodeWhispererServiceTest {
     }
 
     @Test
-    fun `getRequestContext - cross file context should be null for non-cross-file user group`() {
+    fun `getRequestContext - cross file context should be empty for non-cross-file user group`() {
         whenever(userGroupSetting.getUserGroup()).thenReturn(CodeWhispererUserGroup.Control)
         val file = projectRule.fixture.addFileToProject("main.java", "public class Main {}")
 
@@ -102,6 +102,8 @@ class CodeWhispererServiceTest {
             LatencyContext()
         )
 
-        assertThat(actual.supplementalContext).isNull()
+        assertThat(actual.supplementalContext).isNotNull
+        assertThat(actual.supplementalContext?.contents).isEmpty()
+        assertThat(actual.supplementalContext?.contentLength).isEqualTo(0)
     }
 }