tst

Author: rli

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

@@ -13,11 +13,12 @@ import org.assertj.core.api.Assertions.assertThat
 import org.junit.Before
 import org.junit.Test
 import org.mockito.kotlin.any
+import org.mockito.kotlin.doCallRealMethod
 import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.spy
-import org.mockito.kotlin.stub
 import org.mockito.kotlin.whenever
+import org.mockito.kotlin.wheneverBlocking
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionTriggerKind
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.util.LspEditorUtil.toUriString
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
@@ -55,10 +56,8 @@ class CodeWhispererServiceTest : CodeWhispererTestBase() {
         val fileContextProviderSpy = spy(fileContextProvider)
         projectRule.project.replaceService(FileContextProvider::class.java, fileContextProviderSpy, disposableRule.disposable)
 
-        codewhispererService.stub {
-            onBlocking {
-                getRequestContext(any(), any(), any(), any(), any())
-            }.thenCallRealMethod()
+        doCallRealMethod().wheneverBlocking(codewhispererService) {
+            getRequestContext(any(), any(), any(), any(), any())
         }
 
         val requestContext = codewhispererService.getRequestContext(
@@ -95,10 +94,8 @@ class CodeWhispererServiceTest : CodeWhispererTestBase() {
         }
 
         projectRule.project.replaceService(FileContextProvider::class.java, mockFileContextProvider, disposableRule.disposable)
-        codewhispererService.stub {
-            onBlocking {
-                getRequestContext(any(), any(), any(), any(), any())
-            }.thenCallRealMethod()
+        doCallRealMethod().wheneverBlocking(codewhispererService) {
+            getRequestContext(any(), any(), any(), any(), any())
         }
 
         val actual = codewhispererService.getRequestContext(

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

@@ -20,7 +20,6 @@ import kotlinx.coroutines.Deferred
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.first
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runTest
 import kotlinx.coroutines.yield
 import org.assertj.core.api.Assertions.assertThat
@@ -33,7 +32,6 @@ import org.mockito.kotlin.any
 import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.doAnswer
 import org.mockito.kotlin.doNothing
-import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.doSuspendableAnswer
 import org.mockito.kotlin.spy
 import org.mockito.kotlin.stub
@@ -314,27 +312,26 @@ open class CodeWhispererTestBase {
         val projectCaptor = argumentCaptor<Project>()
         val psiFileCaptor = argumentCaptor<PsiFile>()
         val latencyContextCaptor = argumentCaptor<LatencyContext>()
-        codewhispererService.stub {
-            onBlocking {
-                getRequestContext(
-                    triggerTypeCaptor.capture(),
-                    editorCaptor.capture(),
-                    projectCaptor.capture(),
-                    psiFileCaptor.capture(),
-                    latencyContextCaptor.capture()
-                )
-            }.doSuspendableAnswer {
-                val requestContext = codewhispererService.getRequestContext(
-                    triggerTypeCaptor.firstValue,
-                    editorCaptor.firstValue,
-                    projectCaptor.firstValue,
-                    psiFileCaptor.firstValue,
-                    latencyContextCaptor.firstValue
-                )
-                projectRule.fixture.type(userInput)
-                requestContext
-            }.thenCallRealMethod()
-        }
+
+        wheneverBlocking {
+            codewhispererService.getRequestContext(
+                triggerTypeCaptor.capture(),
+                editorCaptor.capture(),
+                projectCaptor.capture(),
+                psiFileCaptor.capture(),
+                latencyContextCaptor.capture()
+            )
+        }.doSuspendableAnswer {
+            val requestContext = codewhispererService.getRequestContext(
+                triggerTypeCaptor.firstValue,
+                editorCaptor.firstValue,
+                projectCaptor.firstValue,
+                psiFileCaptor.firstValue,
+                latencyContextCaptor.firstValue
+            )
+            projectRule.fixture.type(userInput)
+            requestContext
+        }.thenCallRealMethod()
     }
 
     fun mockLspInlineCompletionResponse(response: InlineCompletionListWithReferences) {

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt

@@ -37,7 +37,6 @@ import kotlinx.coroutines.channels.BufferOverflow
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.asSharedFlow
-import kotlinx.coroutines.flow.firstOrNull
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.future.asCompletableFuture
 import kotlinx.coroutines.isActive