@@ -19,6 +19,7 @@ import org.junit.Test
1919import org.mockito.kotlin.any
2020import org.mockito.kotlin.argumentCaptor
2121import org.mockito.kotlin.doReturn
22+ import org.mockito.kotlin.eq
2223import org.mockito.kotlin.mock
2324import org.mockito.kotlin.spy
2425import org.mockito.kotlin.times
@@ -35,6 +36,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.customization.Code
3536import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
3637import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
3738import software.aws.toolkits.jetbrains.services.codewhisperer.language.languages.CodeWhispererJava
39+ import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretContext
3840import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretPosition
3941import software.aws.toolkits.jetbrains.services.codewhisperer.model.Chunk
4042import software.aws.toolkits.jetbrains.services.codewhisperer.model.FileContextInfo
@@ -93,6 +95,31 @@ class CodeWhispererServiceTest {
9395 projectRule.project.replaceService(AwsConnectionManager ::class .java, mock(), disposableRule.disposable)
9496 }
9597
98+ @Test
99+ fun `getRequestContext should use correct fileContext and timeout to fetch supplementalContext` () = runTest {
100+ val fileContextProvider = FileContextProvider .getInstance(projectRule.project)
101+ val fileContextProviderSpy = spy(fileContextProvider)
102+ projectRule.project.replaceService(FileContextProvider ::class .java, fileContextProviderSpy, disposableRule.disposable)
103+
104+ val requestContext = sut.getRequestContext(
105+ TriggerTypeInfo (CodewhispererTriggerType .AutoTrigger , CodeWhispererAutomatedTriggerType .Enter ()),
106+ editor = projectRule.fixture.editor,
107+ project = projectRule.project,
108+ file,
109+ LatencyContext ()
110+ )
111+
112+ requestContext.awaitSupplementalContext()
113+ val fileContextCaptor = argumentCaptor<FileContextInfo >()
114+ verify(fileContextProviderSpy, times(1 )).extractSupplementalFileContext(eq(file), fileContextCaptor.capture(), eq(100 ))
115+ assertThat(fileContextCaptor.firstValue).isEqualTo(FileContextInfo (
116+ CaretContext (leftFileContext = " " , rightFileContext = " public class Main {}" , leftContextOnCurrentLine = " " ),
117+ " main.java" ,
118+ CodeWhispererJava .INSTANCE ,
119+ " main.java"
120+ ))
121+ }
122+
96123 @Test
97124 fun `getRequestContext should have supplementalContext and customizatioArn if they're present` () {
98125 whenever(customizationConfig.activeCustomization(projectRule.project)).thenReturn(
0 commit comments