make codewhispererservice spy on testBase

Author: evanliu048

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

@@ -12,16 +12,13 @@ import com.intellij.openapi.actionSystem.IdeActions
 import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_MOVE_CARET_LEFT
 import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_MOVE_CARET_RIGHT
 import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_TAB
-import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.editor.actionSystem.EditorActionManager
-import com.intellij.testFramework.replaceService
 import com.intellij.testFramework.runInEdtAndWait
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.Before
 import org.junit.Test
 import org.mockito.kotlin.any
 import org.mockito.kotlin.doAnswer
-import org.mockito.kotlin.spy
 import org.mockito.kotlin.stub
 import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsRequest
 import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsResponse
@@ -32,20 +29,13 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestU
 import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.metadata
 import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.sdkHttpResponse
 import software.aws.toolkits.jetbrains.services.codewhisperer.actions.CodeWhispererActionPromoter
-import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
 import kotlin.test.fail
-import org.mockito.kotlin.doNothing
 
 class CodeWhispererAcceptTest : CodeWhispererTestBase() {
 
-    private lateinit var codewhispererServiceSpy: CodeWhispererService
-
     @Before
     override fun setUp() {
         super.setUp()
-        codewhispererServiceSpy = spy(codewhispererService)
-        ApplicationManager.getApplication().replaceService(CodeWhispererService::class.java, codewhispererServiceSpy, disposableRule.disposable)
-        doNothing().`when`(codewhispererServiceSpy).promoteNextInvocationIfAvailable()
 
         // Use java code to test curly braces behavior
         projectRule.fixture.configureByText(javaFileName, javaTestContext)

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

@@ -34,14 +34,14 @@ import kotlin.test.fail
 
 class CodeWhispererSettingsTest : CodeWhispererTestBase() {
 
-    private lateinit var codewhispererServiceSpy: CodeWhispererService
+//    private lateinit var codewhispererServiceSpy: CodeWhispererService
     private lateinit var toolWindowHeadlessManager: ToolWindowHeadlessManagerImpl
 
     @Before
     override fun setUp() {
         super.setUp()
-        codewhispererServiceSpy = spy(codewhispererService)
-        ApplicationManager.getApplication().replaceService(CodeWhispererService::class.java, codewhispererServiceSpy, disposableRule.disposable)
+//        codewhispererServiceSpy = spy(codewhispererService)
+//        ApplicationManager.getApplication().replaceService(CodeWhispererService::class.java, codewhispererServiceSpy, disposableRule.disposable)
 
         // Create a mock ToolWindowManager with working implementation of setAvailable() and isAvailable()
         toolWindowHeadlessManager = object : ToolWindowHeadlessManagerImpl(projectRule.project) {
@@ -83,7 +83,7 @@ class CodeWhispererSettingsTest : CodeWhispererTestBase() {
         whenever(stateManager.checkActiveCodeWhispererConnectionType(projectRule.project)).thenReturn(CodeWhispererLoginType.Logout)
         assertThat(isCodeWhispererEnabled(projectRule.project)).isFalse
         invokeCodeWhispererService()
-        verify(codewhispererServiceSpy, never()).showRecommendationsInPopup(any(), any(), any())
+        verify(codewhispererService, never()).showRecommendationsInPopup(any(), any(), any())
     }
 
     @Test
@@ -92,7 +92,7 @@ class CodeWhispererSettingsTest : CodeWhispererTestBase() {
         assertThat(stateManager.isAutoEnabled()).isFalse
         runInEdtAndWait {
             projectRule.fixture.type(':')
-            verify(codewhispererServiceSpy, never()).showRecommendationsInPopup(any(), any(), any())
+            verify(codewhispererService, never()).showRecommendationsInPopup(any(), any(), any())
         }
     }
 

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

@@ -105,11 +105,9 @@ class CodeWhispererTelemetryTest : CodeWhispererTestBase() {
 
     @Test
     fun `test pre-setup failure will send service invocation event with failed status`() {
-        val codewhispererServiceSpy = spy(codewhispererService) {
-            onGeneric { getRequestContext(any(), any(), any(), any(), any()) }
-                .doAnswer { throw Exception() }
-        }
-        ApplicationManager.getApplication().replaceService(CodeWhispererService::class.java, codewhispererServiceSpy, disposableRule.disposable)
+        doAnswer { throw Exception() }
+            .`when`(codewhispererService)
+            .getRequestContext(any(), any(), any(), any(), any())
 
         invokeCodeWhispererService()
 
@@ -164,7 +162,8 @@ class CodeWhispererTelemetryTest : CodeWhispererTestBase() {
             // 1 serviceInvocation + 1 userModification + 1 userDecision for accepted +
             // (count - 1) userDecisions for ignored
             verify(batcher, atLeast(2 + count)).enqueue(capture())
-            assertEventsContainsFieldsAndCount(allValues, serviceInvocation, 3)
+            // one service invocation for current recommendation and one for next recommendation
+            assertEventsContainsFieldsAndCount(allValues, serviceInvocation, 2)
             assertEventsContainsFieldsAndCount(
                 allValues,
                 userModification,
@@ -411,15 +410,13 @@ class CodeWhispererTelemetryTest : CodeWhispererTestBase() {
 
     @Test
     fun `test showing IntelliSense after triggering CodeWhisperer will send userDecision events of state Discard`() {
-        val codewhispererServiceSpy = spy(codewhispererService)
-        codewhispererServiceSpy.stub {
+        codewhispererService.stub {
             onGeneric {
                 canDoInvocation(any(), any())
             } doAnswer {
                 true
             }
         }
-        ApplicationManager.getApplication().replaceService(CodeWhispererService::class.java, codewhispererServiceSpy, disposableRule.disposable)
         popupManagerSpy.stub {
             onGeneric {
                 hasConflictingPopups(any())

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

@@ -114,7 +114,10 @@ open class CodeWhispererTestBase {
 
         stateManager = spy(CodeWhispererExplorerActionManager.getInstance())
         recommendationManager = CodeWhispererRecommendationManager.getInstance()
-        codewhispererService = CodeWhispererService.getInstance()
+        codewhispererService = spy(CodeWhispererService.getInstance())
+        ApplicationManager.getApplication().replaceService(CodeWhispererService::class.java, codewhispererService, disposableRule.disposable)
+        doNothing().`when`(codewhispererService).promoteNextInvocationIfAvailable()
+
         editorManager = CodeWhispererEditorManager.getInstance()
         settingsManager = CodeWhispererSettings.getInstance()
 

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

@@ -100,13 +100,12 @@ class CodeWhispererUserInputTest : CodeWhispererTestBase() {
     }
 
     private fun addUserInputAfterInvocation(userInput: String) {
-        val codewhispererServiceSpy = spy(codewhispererService)
         val triggerTypeCaptor = argumentCaptor<TriggerTypeInfo>()
         val editorCaptor = argumentCaptor<Editor>()
         val projectCaptor = argumentCaptor<Project>()
         val psiFileCaptor = argumentCaptor<PsiFile>()
         val latencyContextCaptor = argumentCaptor<LatencyContext>()
-        codewhispererServiceSpy.stub {
+        codewhispererService.stub {
             onGeneric {
                 getRequestContext(
                     triggerTypeCaptor.capture(),
@@ -116,7 +115,7 @@ class CodeWhispererUserInputTest : CodeWhispererTestBase() {
                     latencyContextCaptor.capture()
                 )
             }.doAnswer {
-                val requestContext = codewhispererServiceSpy.getRequestContext(
+                val requestContext = codewhispererService.getRequestContext(
                     triggerTypeCaptor.firstValue,
                     editorCaptor.firstValue,
                     projectCaptor.firstValue,
@@ -127,6 +126,5 @@ class CodeWhispererUserInputTest : CodeWhispererTestBase() {
                 requestContext
             }.thenCallRealMethod()
         }
-        ApplicationManager.getApplication().replaceService(CodeWhispererService::class.java, codewhispererServiceSpy, disposableRule.disposable)
     }
 }