tst

Author: rli

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererRecommendationAction.kt

@@ -8,12 +8,15 @@ import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.CommonDataKeys
 import com.intellij.openapi.project.DumbAware
+import com.intellij.openapi.util.Key
+import kotlinx.coroutines.Job
 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.CodeWhispererAutomatedTriggerType
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
 import software.aws.toolkits.resources.message
 import software.aws.toolkits.telemetry.CodewhispererTriggerType
+import java.util.concurrent.atomic.AtomicReference
 
 class CodeWhispererRecommendationAction : AnAction(message("codewhisperer.trigger.service")), DumbAware {
     override fun getActionUpdateThread() = ActionUpdateThread.BGT
@@ -32,6 +35,12 @@ class CodeWhispererRecommendationAction : AnAction(message("codewhisperer.trigge
         }
 
         val triggerType = TriggerTypeInfo(CodewhispererTriggerType.OnDemand, CodeWhispererAutomatedTriggerType.Unknown())
-        CodeWhispererService.getInstance().showRecommendationsInPopup(editor, triggerType, latencyContext)
+        val job = CodeWhispererService.getInstance().showRecommendationsInPopup(editor, triggerType, latencyContext)
+
+        e.getData(CommonDataKeys.EDITOR)?.getUserData(ACTION_JOB_KEY)?.set(job)
+    }
+
+    companion object {
+        val ACTION_JOB_KEY = Key.create<AtomicReference<Job?>>("amazonq.codewhisperer.job")
     }
 }

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt

@@ -104,15 +104,17 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
         editor: Editor,
         triggerTypeInfo: TriggerTypeInfo,
         latencyContext: LatencyContext
-    ) {
+    ): Job? {
         if (job == null || job?.isCompleted == true) {
             job = cs.launch(getCoroutineBgContext()) {
                 doShowRecommendationsInPopup(editor, triggerTypeInfo, latencyContext)
             }
         }
+
+        // did some wrangling, but compiler didn't believe this can't be null
+        return job
     }
 
-    @RequiresEdt
     private suspend fun doShowRecommendationsInPopup(
         editor: Editor,
         triggerTypeInfo: TriggerTypeInfo,
@@ -123,6 +125,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
 
         latencyContext.credentialFetchingStart = System.nanoTime()
 
+        // try to refresh automatically if possible, otherwise ask user to login again
         if (isQExpired(project)) {
             // consider changing to only running once a ~minute since this is relatively expensive
             // say the connection is un-refreshable if refresh fails for 3 times

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

@@ -8,6 +8,10 @@ import com.intellij.testFramework.DisposableRule
 import com.intellij.testFramework.RuleChain
 import com.intellij.testFramework.replaceService
 import com.intellij.testFramework.runInEdtAndWait
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.yield
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.After
 import org.junit.Before
@@ -32,6 +36,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestU
 import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponse
 import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonTestLeftContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.testValidAccessToken
+import software.aws.toolkits.jetbrains.services.codewhisperer.actions.CodeWhispererRecommendationAction
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererLoginType
 import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorManager
@@ -51,6 +56,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhisp
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererColorUtil.POPUP_DIM_HEX
 import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
 import software.aws.toolkits.resources.message
+import java.util.concurrent.atomic.AtomicReference
 
 // TODO: restructure testbase, too bulky and hard to debug
 open class CodeWhispererTestBase {
@@ -180,12 +186,26 @@ open class CodeWhispererTestBase {
         runInEdtAndWait {}
     }
 
+    /**
+     * Block until manual action has either failed or completed
+     */
     fun invokeCodeWhispererService() {
+        val jobRef = AtomicReference<Job?>()
         runInEdtAndWait {
+            projectRule.fixture.editor.putUserData(CodeWhispererRecommendationAction.ACTION_JOB_KEY, jobRef)
+            // does not block, so we need to extract something to track the async task
             projectRule.fixture.performEditorAction(codeWhispererRecommendationActionId)
         }
-        while (CodeWhispererInvocationStatus.getInstance().hasExistingInvocation()) {
-            Thread.sleep(10)
+
+        runTest {
+            // wait for CodeWhispererService#showRecommendationsInPopup to complete, if started
+            jobRef.get()?.join()
+
+            // wait for subsequent background operations to be complete
+            while (CodeWhispererInvocationStatus.getInstance().hasExistingInvocation()) {
+                yield()
+                delay(10)
+            }
         }
     }