diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt index 0115491f9a..9d638fc60a 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt @@ -76,6 +76,10 @@ interface CodeWhispererClientAdaptor : Disposable { firstRequest: GenerateCompletionsRequest, ): Sequence + fun generateCompletions( + firstRequest: GenerateCompletionsRequest, + ): GenerateCompletionsResponse + fun createUploadUrl( request: CreateUploadUrlRequest, ): CreateUploadUrlResponse @@ -322,6 +326,8 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW yield(response) } while (!nextToken.isNullOrEmpty()) } + override fun generateCompletions(firstRequest: GenerateCompletionsRequest): GenerateCompletionsResponse = + bearerClient().generateCompletions(firstRequest) override fun createUploadUrl(request: CreateUploadUrlRequest): CreateUploadUrlResponse = bearerClient().createUploadUrl(request) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupListener.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupListener.kt index 145e3cf624..60e7448034 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupListener.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupListener.kt @@ -7,6 +7,7 @@ import com.intellij.openapi.ui.popup.JBPopupListener import com.intellij.openapi.ui.popup.LightweightWindowEvent import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererInvocationStatus +import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererServiceNew import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererTelemetryService import java.time.Duration @@ -27,7 +28,8 @@ class CodeWhispererPopupListener(private val states: InvocationContext) : JBPopu recommendationContext, CodeWhispererPopupManager.getInstance().sessionContext, event.isOk, - CodeWhispererInvocationStatus.getInstance().popupStartTimestamp?.let { Duration.between(it, Instant.now()) } + CodeWhispererInvocationStatus.getInstance().popupStartTimestamp?.let { Duration.between(it, Instant.now()) }, + CodeWhispererService.getInstance().getNextInvocationContext() ) CodeWhispererInvocationStatus.getInstance().setPopupActive(false) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupManager.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupManager.kt index 29ffb766a8..3c32413617 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupManager.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupManager.kt @@ -68,6 +68,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.popup.listeners.Co import software.aws.toolkits.jetbrains.services.codewhisperer.popup.listeners.CodeWhispererPrevButtonActionListener import software.aws.toolkits.jetbrains.services.codewhisperer.popup.listeners.CodeWhispererScrollListener import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererInvocationStatus +import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererTelemetryService import software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow.CodeWhispererCodeReferenceManager import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererColorUtil.POPUP_DIM_HEX @@ -455,6 +456,9 @@ class CodeWhispererPopupManager { CodeWhispererEditorManager.getInstance().updateEditorWithRecommendation(states, sessionContext) } closePopup(states.popup) + if (sessionContext.selectedIndex == 0) { + CodeWhispererService.getInstance().promoteNextInvocationIfAvailable() + } } } ) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt index 54451e949c..7a92e36e6b 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt @@ -16,6 +16,7 @@ import com.intellij.openapi.editor.Editor import com.intellij.openapi.editor.VisualPosition import com.intellij.openapi.project.Project import com.intellij.openapi.ui.popup.JBPopup +import com.intellij.openapi.ui.popup.JBPopupFactory import com.intellij.openapi.util.Disposer import com.intellij.psi.PsiDocumentManager import com.intellij.psi.PsiFile @@ -24,6 +25,7 @@ import com.intellij.util.messages.Topic import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Deferred import kotlinx.coroutines.Job +import kotlinx.coroutines.SupervisorJob import kotlinx.coroutines.async import kotlinx.coroutines.delay import kotlinx.coroutines.isActive @@ -98,6 +100,7 @@ import java.util.concurrent.TimeUnit class CodeWhispererService(private val cs: CoroutineScope) : Disposable { private val codeInsightSettingsFacade = CodeInsightsSettingsFacade() private var refreshFailure: Int = 0 + private var nextInvocationContext: InvocationContext? = null init { Disposer.register(this, codeInsightSettingsFacade) @@ -209,7 +212,110 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { invokeCodeWhispererInBackground(requestContext) } - internal suspend fun invokeCodeWhispererInBackground(requestContext: RequestContext): Job { + internal suspend fun invokeCodeWhispererInBackground( + requestContext: RequestContext, + currStates: InvocationContext? = null, + ): Job { + // current states != null means that it's prefetch + if (currStates != null) { + val firstValidRecommendation = currStates.recommendationContext.details + .firstOrNull { + !it.isDiscarded && it.recommendation.content().isNotEmpty() + } ?: return SupervisorJob().apply { complete() } + val job = cs.launch(getCoroutineBgContext()) { + val latencyContext = LatencyContext().apply { + codewhispererPreprocessingStart = System.nanoTime() + codewhispererEndToEndStart = System.nanoTime() + } + + val nextCaretPosition = calculateNextCaretPosition(requestContext, firstValidRecommendation) + val nextFileContextInfo = createNextFileContextInfo(requestContext, firstValidRecommendation) + + val nextRequestContext = requestContext.copy( + caretPosition = nextCaretPosition, + fileContextInfo = nextFileContextInfo, + latencyContext = latencyContext + ) + val newVisualPosition = withContext(EDT) { + runReadAction { + nextRequestContext.editor.offsetToVisualPosition(nextRequestContext.caretPosition.offset) + } + } + try { + val nextResponse = CodeWhispererClientAdaptor + .getInstance(nextRequestContext.project) + .generateCompletions( + buildCodeWhispererRequest( + nextRequestContext.fileContextInfo, + nextRequestContext.awaitSupplementalContext(), + nextRequestContext.customizationArn + ) + ) + val startTime = System.nanoTime() + nextRequestContext.latencyContext.codewhispererPreprocessingEnd = System.nanoTime() + nextRequestContext.latencyContext.paginationAllCompletionsStart = System.nanoTime() + CodeWhispererInvocationStatus.getInstance().setInvocationStart() + nextResponse.let { + val endTime = System.nanoTime() + val latency = TimeUnit.NANOSECONDS.toMillis(endTime - startTime).toDouble() + val requestId = nextResponse.responseMetadata().requestId() + val sessionId = nextResponse.sdkHttpResponse().headers().getOrDefault(KET_SESSION_ID, listOf(requestId))[0] + + nextRequestContext.latencyContext.apply { + codewhispererPostprocessingStart = System.nanoTime() + paginationFirstCompletionTime = (endTime - codewhispererEndToEndStart).toDouble() + firstRequestId = requestId + } + + CodeWhispererInvocationStatus.getInstance().setInvocationSessionId(sessionId) + + val nextResponseContext = ResponseContext(sessionId) + CodeWhispererTelemetryService.getInstance().sendServiceInvocationEvent( + nextResponse.responseMetadata().requestId(), + nextRequestContext, + nextResponseContext, + nextResponse.completions().size, + true, + latency, + null + ) + val validatedResponse = validateResponse(it) + val detailContexts = withContext(EDT) { + runReadAction { + CodeWhispererRecommendationManager.getInstance().buildDetailContext( + nextRequestContext, + "", + validatedResponse.completions(), + validatedResponse.responseMetadata().requestId() + ) + } + } + val nextRecommendationContext = RecommendationContext(detailContexts, "", "", newVisualPosition) + val newPopup = withContext(EDT) { + JBPopupFactory.getInstance().createMessage("Dummy popup") + } + + // send userDecision and trigger decision when next recommendation haven't been seen + if (currStates.popup.isDisposed) { + CodeWhispererTelemetryService.getInstance().sendUserDecisionEventForAll( + nextRequestContext, + nextResponseContext, + nextRecommendationContext, + SessionContext(), + false + ) + } else { + nextInvocationContext = InvocationContext(nextRequestContext, nextResponseContext, nextRecommendationContext, newPopup) + } + LOG.debug { "Prefetched next invocation stored in nextInvocationContext" } + } + } catch (ex: Exception) { + LOG.warn { "Failed to prefetch next codewhisperer invocation: ${ex.message}" } + } + } + return job + } + val popup = withContext(EDT) { CodeWhispererPopupManager.getInstance().initPopup().also { Disposer.register(it) { CodeWhispererInvocationStatus.getInstance().finishInvocation() } @@ -491,6 +597,9 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { CodeWhispererPopupManager.getInstance().cancelPopup(popup) return null } + cs.launch(getCoroutineBgContext()) { + invokeCodeWhispererInBackground(requestContext, nextStates) + } } else { // subsequent responses nextStates = updateStates(currStates, response) @@ -621,6 +730,67 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { CodeWhispererPopupManager.getInstance().changeStates(states, 0, "", true, recommendationAdded) } + fun promoteNextInvocationIfAvailable() { + val nextStates = nextInvocationContext ?: run { + LOG.debug { "No nextInvocationContext found, nothing to promote." } + return + } + nextInvocationContext?.popup?.let { Disposer.dispose(it) } + nextInvocationContext = null + + runInEdt { + val newPopup = CodeWhispererPopupManager.getInstance().initPopup() + val updatedNextStates = nextStates.copy(popup = newPopup).also { + addPopupChildDisposables(it.popup) + Disposer.register(newPopup, it) + } + CodeWhispererPopupManager.getInstance().initPopupListener(updatedNextStates) + CodeWhispererPopupManager.getInstance().changeStates( + updatedNextStates, + 0, + "", + typeaheadAdded = true, + recommendationAdded = false + ) + cs.launch(getCoroutineBgContext()) { + invokeCodeWhispererInBackground(updatedNextStates.requestContext, updatedNextStates) + } + } + + LOG.debug { "Promoted nextInvocationContext to current session and displayed next recommendation." } + } + + private fun calculateNextCaretPosition( + currentRequestContext: RequestContext, + firstValidRecommendation: DetailContext, + ): CaretPosition { + val indent = currentRequestContext.fileContextInfo.caretContext.leftContextOnCurrentLine + .takeWhile { it.isWhitespace() } + val recommendedText = buildString { + append(indent) + append(firstValidRecommendation.recommendation.content()) + if (!endsWith("\n")) { + append("\n") + } + } + val lineCount = recommendedText.count { it == '\n' } + + return CaretPosition( + line = currentRequestContext.caretPosition.line + lineCount, + offset = currentRequestContext.caretPosition.offset + recommendedText.length + ) + } + + private fun createNextFileContextInfo( + requestContext: RequestContext, + firstValidRecommendation: DetailContext, + ): FileContextInfo { + val updatedCaretContext = requestContext.fileContextInfo.caretContext.copy( + leftFileContext = requestContext.fileContextInfo.caretContext.leftFileContext + firstValidRecommendation.recommendation.content() + ) + return requestContext.fileContextInfo.copy(caretContext = updatedCaretContext) + } + private fun sendDiscardedUserDecisionEventForAll( requestContext: RequestContext, responseContext: ResponseContext, @@ -781,6 +951,8 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable { override fun dispose() {} + fun getNextInvocationContext(): InvocationContext? = nextInvocationContext + companion object { private val LOG = getLogger() private const val MAX_REFRESH_ATTEMPT = 3 diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererTelemetryService.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererTelemetryService.kt index 59d82c3b44..e8593c8c07 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererTelemetryService.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererTelemetryService.kt @@ -453,6 +453,7 @@ class CodeWhispererTelemetryService { sessionContext: SessionContext, hasUserAccepted: Boolean, popupShownTime: Duration? = null, + nextInvocationContext: InvocationContext? = null, ) { val detailContexts = recommendationContext.details val decisions = mutableListOf() @@ -500,6 +501,19 @@ class CodeWhispererTelemetryService { previousUserTriggerDecisions.add(this) // we need this as well because AutotriggerService will reset the queue periodically CodeWhispererAutoTriggerService.getInstance().addPreviousDecision(this) + // send possible next session event if current action is reject and next popup haven't shown up + if (CodewhispererSuggestionState.from(this.toString()) == CodewhispererSuggestionState.Reject) { + nextInvocationContext?.let { + sendUserDecisionEventForAll( + it.requestContext, + it.responseContext, + it.recommendationContext, + SessionContext(), + false, + nextInvocationContext = null + ) + } + } } } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererSettingsTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererSettingsTest.kt index 3a640c15bc..e9e7a30e18 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererSettingsTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererSettingsTest.kt @@ -4,7 +4,6 @@ package software.aws.toolkits.jetbrains.services.codewhisperer import com.intellij.analysis.problemsView.toolWindow.ProblemsView -import com.intellij.openapi.application.ApplicationManager import com.intellij.openapi.components.service import com.intellij.openapi.wm.RegisterToolWindowTask import com.intellij.openapi.wm.ToolWindow @@ -18,14 +17,12 @@ import org.junit.Ignore import org.junit.Test import org.mockito.kotlin.any import org.mockito.kotlin.never -import org.mockito.kotlin.spy import org.mockito.kotlin.verify import org.mockito.kotlin.whenever import software.aws.toolkits.jetbrains.core.ToolWindowHeadlessManagerImpl import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererLoginType import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExploreActionState import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererEnabled -import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService import software.aws.toolkits.jetbrains.services.codewhisperer.status.CodeWhispererStatusBarWidgetFactory import software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow.CodeWhispererCodeReferenceToolWindowFactory import software.aws.toolkits.jetbrains.settings.CodeWhispererConfiguration @@ -34,14 +31,11 @@ import kotlin.test.fail class CodeWhispererSettingsTest : CodeWhispererTestBase() { - 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) // Create a mock ToolWindowManager with working implementation of setAvailable() and isAvailable() toolWindowHeadlessManager = object : ToolWindowHeadlessManagerImpl(projectRule.project) { @@ -83,7 +77,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 +86,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()) } } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTelemetryTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTelemetryTest.kt index 6bd9d00629..b2a5b49a4c 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTelemetryTest.kt +++ b/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, 1) + // one service invocation for current recommendation and one for next recommendation + assertEventsContainsFieldsAndCount(allValues, serviceInvocation, 2) assertEventsContainsFieldsAndCount( allValues, userModification, @@ -194,10 +193,12 @@ class CodeWhispererTelemetryTest : CodeWhispererTestBase() { val count = pythonResponse.completions().size argumentCaptor().apply { verify(batcher, atLeast(1 + count)).enqueue(capture()) + // send 2 invocation events for current and next recommendation session, + // one reject and at lease (count -1) Unseen events assertEventsContainsFieldsAndCount( allValues, serviceInvocation, - 1, + 2, "result" to Result.Succeeded.toString(), ) assertEventsContainsFieldsAndCount( @@ -211,6 +212,7 @@ class CodeWhispererTelemetryTest : CodeWhispererTestBase() { userDecision, count - 1, codewhispererSuggestionState to CodewhispererSuggestionState.Unseen.toString(), + atLeast = true ) } } @@ -218,13 +220,14 @@ class CodeWhispererTelemetryTest : CodeWhispererTestBase() { @Test fun `test invoking CodeWhisperer will send service invocation event with succeeded status`() { + // serviceInvocation event will be sent for current session and preloaded next session withCodeWhispererServiceInvokedAndWait { argumentCaptor().apply { verify(batcher, atLeastOnce()).enqueue(capture()) assertEventsContainsFieldsAndCount( allValues, serviceInvocation, - 1, + 2, "result" to Result.Succeeded.toString(), ) } @@ -287,12 +290,13 @@ class CodeWhispererTelemetryTest : CodeWhispererTestBase() { val prefixNotMatchCount = pythonResponse.completions().filter { !it.content().startsWith(typeahead) }.size + // serviceInvocation event will be sent for current session and preloaded next session argumentCaptor().apply { verify(batcher, atLeast(1 + prefixNotMatchCount)).enqueue(capture()) assertEventsContainsFieldsAndCount( allValues, serviceInvocation, - 1, + 2, "result" to Result.Succeeded.toString(), ) assertEventsContainsFieldsAndCount( @@ -373,13 +377,14 @@ class CodeWhispererTelemetryTest : CodeWhispererTestBase() { @Test fun `test invoking CodeWhisperer will send service invocation event with sessionId and requestId from response`() { + // serviceInvocation event will be sent for current session and preloaded next session withCodeWhispererServiceInvokedAndWait { states -> val metricCaptor = argumentCaptor() verify(batcher, atLeastOnce()).enqueue(metricCaptor.capture()) assertEventsContainsFieldsAndCount( metricCaptor.allValues, serviceInvocation, - 1, + 2, "codewhispererSessionId" to states.responseContext.sessionId, "codewhispererRequestId" to states.recommendationContext.details[0].requestId, ) @@ -405,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()) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestBase.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestBase.kt index c6687d501b..8cdca24258 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestBase.kt +++ b/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() diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererUserInputTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererUserInputTest.kt index 8c4c562481..a1948f0534 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererUserInputTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererUserInputTest.kt @@ -3,21 +3,17 @@ package software.aws.toolkits.jetbrains.services.codewhisperer -import com.intellij.openapi.application.ApplicationManager import com.intellij.openapi.editor.Editor import com.intellij.openapi.project.Project import com.intellij.psi.PsiFile -import com.intellij.testFramework.replaceService import org.assertj.core.api.Assertions.assertThat import org.junit.Test import org.mockito.kotlin.argumentCaptor import org.mockito.kotlin.doAnswer -import org.mockito.kotlin.spy import org.mockito.kotlin.stub import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponse import software.aws.toolkits.jetbrains.services.codewhisperer.model.LatencyContext import software.aws.toolkits.jetbrains.services.codewhisperer.model.TriggerTypeInfo -import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService class CodeWhispererUserInputTest : CodeWhispererTestBase() { @@ -100,13 +96,12 @@ class CodeWhispererUserInputTest : CodeWhispererTestBase() { } private fun addUserInputAfterInvocation(userInput: String) { - val codewhispererServiceSpy = spy(codewhispererService) val triggerTypeCaptor = argumentCaptor() val editorCaptor = argumentCaptor() val projectCaptor = argumentCaptor() val psiFileCaptor = argumentCaptor() val latencyContextCaptor = argumentCaptor() - codewhispererServiceSpy.stub { + codewhispererService.stub { onGeneric { getRequestContext( triggerTypeCaptor.capture(), @@ -116,7 +111,7 @@ class CodeWhispererUserInputTest : CodeWhispererTestBase() { latencyContextCaptor.capture() ) }.doAnswer { - val requestContext = codewhispererServiceSpy.getRequestContext( + val requestContext = codewhispererService.getRequestContext( triggerTypeCaptor.firstValue, editorCaptor.firstValue, projectCaptor.firstValue, @@ -127,6 +122,5 @@ class CodeWhispererUserInputTest : CodeWhispererTestBase() { requestContext }.thenCallRealMethod() } - ApplicationManager.getApplication().replaceService(CodeWhispererService::class.java, codewhispererServiceSpy, disposableRule.disposable) } }