c

Author: rli

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererServiceTest.kt

@@ -56,7 +56,7 @@ class CodeWhispererServiceTest : CodeWhispererTestBase() {
         projectRule.project.replaceService(FileContextProvider::class.java, fileContextProviderSpy, disposableRule.disposable)
 
         codewhispererService.stub {
-            onGeneric {
+            onBlocking {
                 getRequestContext(any(), any(), any(), any(), any())
             }.thenCallRealMethod()
         }
@@ -81,7 +81,7 @@ class CodeWhispererServiceTest : CodeWhispererTestBase() {
     }
 
     @Test
-    fun `getRequestContext should have customizationArn if it's present`() {
+    fun `getRequestContext should have customizationArn if it's present`() = runTest {
         whenever(customizationConfig.activeCustomization(projectRule.project)).thenReturn(
             CodeWhispererCustomization(
                 "fake-arn",
@@ -96,7 +96,7 @@ class CodeWhispererServiceTest : CodeWhispererTestBase() {
 
         projectRule.project.replaceService(FileContextProvider::class.java, mockFileContextProvider, disposableRule.disposable)
         codewhispererService.stub {
-            onGeneric {
+            onBlocking {
                 getRequestContext(any(), any(), any(), any(), any())
             }.thenCallRealMethod()
         }

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestBase.kt

@@ -31,6 +31,7 @@ import org.mockito.kotlin.any
 import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.doAnswer
 import org.mockito.kotlin.doNothing
+import org.mockito.kotlin.doSuspendableAnswer
 import org.mockito.kotlin.spy
 import org.mockito.kotlin.stub
 import org.mockito.kotlin.timeout
@@ -147,7 +148,7 @@ open class CodeWhispererTestBase {
         mockLspInlineCompletionResponse(pythonResponse)
 
         mockClientManagerRule.create<SsoOidcClient>()
-        every { mockLanguageServer.logInlineCompletionSessionResults(any()) } returns CompletableFuture.completedFuture(Unit)
+        every { mockLanguageServer.logInlineCompletionSessionResults(any()) } returns Unit
 
         popupManagerSpy = spy(CodeWhispererPopupManager.getInstance())
         popupManagerSpy.reset()
@@ -165,7 +166,7 @@ open class CodeWhispererTestBase {
         recommendationManager = CodeWhispererRecommendationManager.getInstance()
         codewhispererService = spy(CodeWhispererService.getInstance())
         codewhispererService.stub {
-            onGeneric {
+            onBlocking {
                 getWorkspaceIds(any())
             } doAnswer {
                 CompletableFuture.completedFuture(LspServerConfigurations(listOf(WorkspaceInfo("file:///", "workspaceId"))))
@@ -307,15 +308,15 @@ open class CodeWhispererTestBase {
         val psiFileCaptor = argumentCaptor<PsiFile>()
         val latencyContextCaptor = argumentCaptor<LatencyContext>()
         codewhispererService.stub {
-            onGeneric {
+            onBlocking {
                 getRequestContext(
                     triggerTypeCaptor.capture(),
                     editorCaptor.capture(),
                     projectCaptor.capture(),
                     psiFileCaptor.capture(),
                     latencyContextCaptor.capture()
                 )
-            }.doAnswer {
+            }.doSuspendableAnswer {
                 val requestContext = codewhispererService.getRequestContext(
                     triggerTypeCaptor.firstValue,
                     editorCaptor.firstValue,