wip

Author: rli

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/BrowserConnector.kt

@@ -112,7 +112,7 @@ class BrowserConnector(
     private val themeBrowserAdapter: ThemeBrowserAdapter = ThemeBrowserAdapter(),
     private val project: Project,
 ) {
-    var uiReady = CompletableDeferred<Boolean>()
+    val uiReady = CompletableDeferred<Boolean>()
     private val chatCommunicationManager = ChatCommunicationManager.getInstance(project)
     private val chatAsyncResultManager = ChatAsyncResultManager.getInstance(project)
 

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

@@ -13,13 +13,19 @@ import com.intellij.testFramework.replaceService
 import com.intellij.testFramework.runInEdtAndWait
 import io.mockk.every
 import io.mockk.mockk
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.Deferred
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.delay
+import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.test.runTest
 import kotlinx.coroutines.yield
 import org.assertj.core.api.Assertions.assertThat
+import org.eclipse.lsp4j.jsonrpc.Launcher
+import org.eclipse.lsp4j.jsonrpc.RemoteEndpoint
 import org.junit.After
 import org.junit.Before
+import org.junit.BeforeClass
 import org.junit.Rule
 import org.mockito.kotlin.any
 import org.mockito.kotlin.argumentCaptor
@@ -38,6 +44,9 @@ import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.sono.Q_SCOPES
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLanguageServer
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQServerInstanceFacade
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AwsExtendedInitializeResult
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LspServerConfigurations
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.WorkspaceInfo
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionListWithReferences
@@ -69,6 +78,7 @@ import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings
 import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
 import software.aws.toolkits.resources.message
 import java.util.concurrent.CompletableFuture
+import java.util.concurrent.Future
 import java.util.concurrent.atomic.AtomicReference
 
 // TODO: restructure testbase, too bulky and hard to debug
@@ -100,8 +110,30 @@ open class CodeWhispererTestBase {
 
     @Before
     open fun setUp() {
-        mockLspService = spy(AmazonQLspService.getInstance(projectRule.project))
         mockLanguageServer = mockk()
+        mockLspService = spy(object : AmazonQLspService(projectRule.project, mockk()) {
+            override fun start() = CompletableDeferred(object : AmazonQServerInstanceFacade {
+                override val launcher: Launcher<AmazonQLanguageServer>
+                    get() = TODO("Not yet implemented")
+
+                override val launcherFuture: Future<Void>
+                    get() = TODO("Not yet implemented")
+
+                override val initializeResult: Deferred<AwsExtendedInitializeResult>
+                    get() = TODO("Not yet implemented")
+
+                override val encryptionManager: JwtEncryptionManager
+                    get() = TODO("Not yet implemented")
+
+                override val languageServer: AmazonQLanguageServer
+                    get() = mockLanguageServer
+
+                override val rawEndpoint: RemoteEndpoint
+                    get() = TODO("Not yet implemented")
+
+                override fun dispose() {}
+            })
+        })
 
         // Mock the service methods on Project
         projectRule.project.replaceService(AmazonQLspService::class.java, mockLspService, disposableRule.disposable)

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

@@ -61,6 +61,7 @@ import org.eclipse.lsp4j.jsonrpc.RemoteEndpoint
 import org.eclipse.lsp4j.jsonrpc.json.JsonRpcMethod
 import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
 import org.eclipse.lsp4j.launch.LSPLauncher
+import org.jetbrains.annotations.VisibleForTesting
 import org.slf4j.event.Level
 import software.aws.toolkits.core.utils.debug
 import software.aws.toolkits.core.utils.getLogger
@@ -134,11 +135,11 @@ internal class LSPProcessListener : ProcessListener {
 }
 
 @Service(Service.Level.PROJECT)
-class AmazonQLspService(private val project: Project, private val cs: CoroutineScope) : Disposable {
+open class AmazonQLspService(private val project: Project, private val cs: CoroutineScope) : Disposable {
     private val _flowInstance = MutableSharedFlow<AmazonQServerInstance>(replay = 1, onBufferOverflow = BufferOverflow.DROP_OLDEST)
     val instanceFlow = _flowInstance.asSharedFlow().map { it.languageServer }
 
-    private var instance: Deferred<AmazonQServerInstance>
+    private var instance: Deferred<AmazonQServerInstanceFacade>
     val capabilities
         get() = instance.getCompleted().initializeResult.getCompleted().capabilities
 
@@ -154,7 +155,8 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS
     // dont allow lsp commands if server is restarting
     private val mutex = Mutex(false)
 
-    private fun start() = cs.async {
+    @VisibleForTesting
+    protected open fun start(): Deferred<AmazonQServerInstanceFacade> = cs.async {
         // manage lifecycle RAII-like so we can restart at arbitrary time
         // and suppress IDE error if server fails to start
         var attempts = 0
@@ -331,21 +333,27 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS
     }
 }
 
-private class AmazonQServerInstance(private val project: Project, private val cs: CoroutineScope) : Disposable {
-    val encryptionManager = JwtEncryptionManager()
-
-    private val launcher: Launcher<AmazonQLanguageServer>
+interface AmazonQServerInstanceFacade : Disposable {
+    val launcher: Launcher<AmazonQLanguageServer>
+    val launcherFuture: Future<Void>
+    val initializeResult: Deferred<AwsExtendedInitializeResult>
+    val encryptionManager: JwtEncryptionManager
 
     val languageServer: AmazonQLanguageServer
         get() = launcher.remoteProxy
 
     val rawEndpoint: RemoteEndpoint
         get() = launcher.remoteEndpoint
+}
+
+private class AmazonQServerInstance(private val project: Project, private val cs: CoroutineScope) : Disposable, AmazonQServerInstanceFacade {
+    override val encryptionManager = JwtEncryptionManager()
+    override val launcher: Launcher<AmazonQLanguageServer>
 
     @Suppress("ForbiddenVoid")
-    val launcherFuture: Future<Void>
+    override val launcherFuture: Future<Void>
     private val launcherHandler: KillableProcessHandler
-    val initializeResult: Deferred<InitializeResult>
+    override val initializeResult: Deferred<AwsExtendedInitializeResult>
 
     private fun createClientCapabilities(): ClientCapabilities =
         ClientCapabilities().apply {
@@ -546,7 +554,7 @@ private class AmazonQServerInstance(private val project: Project, private val cs
             }
             languageServer.initialized(InitializedParams())
 
-            initializeResult
+            initializeResult as AwsExtendedInitializeResult
         }
 
         // invokeOnCompletion results in weird lock/timeout error