feedback

Author: samgst-amazon

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

@@ -19,5 +19,5 @@ interface AmazonQLanguageServer : LanguageServer {
     fun updateTokenCredentials(payload: UpdateCredentialsPayload): CompletableFuture<ResponseMessage>
 
     @JsonNotification("aws/credentials/token/delete")
-    fun deleteTokenCredentials()
+    fun deleteTokenCredentials(): CompletableFuture<Unit>
 }

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

@@ -154,28 +154,26 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS
         instance = start()
     }
 
-    suspend fun execute(runnable: suspend (AmazonQLanguageServer) -> Unit) {
+    suspend fun<T> execute(runnable: suspend (AmazonQLanguageServer) -> T): T {
         val lsp = withTimeout(10.seconds) {
             val holder = mutex.withLock { instance }.await()
             holder.initializer.join()
 
             holder.languageServer
         }
-
-        runnable(lsp)
+        return runnable(lsp)
     }
 
-    fun executeSync(runnable: suspend (AmazonQLanguageServer) -> Unit) {
+    fun<T> executeSync(runnable: suspend (AmazonQLanguageServer) -> T): T =
         runBlocking(cs.coroutineContext) {
             execute(runnable)
         }
-    }
 
     companion object {
         private val LOG = getLogger<AmazonQLspService>()
         fun getInstance(project: Project) = project.service<AmazonQLspService>()
 
-        fun executeIfRunning(project: Project, runnable: (AmazonQLanguageServer) -> Unit) =
+        fun <T> executeIfRunning(project: Project, runnable: (AmazonQLanguageServer) -> T): T? =
             project.serviceIfCreated<AmazonQLspService>()?.executeSync(runnable)
     }
 }

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

@@ -26,28 +26,18 @@ class DefaultAuthCredentialsService(
 
         val payload = createUpdateCredentialsPayload(token)
 
-        return CompletableFuture<ResponseMessage>().also { completableFuture ->
-            AmazonQLspService.executeIfRunning(project) { server ->
-                server.updateTokenCredentials(payload)
-                    .whenComplete { response, throwable ->
-                        if (throwable != null) {
-                            completableFuture.completeExceptionally(throwable)
-                        } else {
-                            completableFuture.complete(response)
-                        }
-                    }
-            } ?: completableFuture.completeExceptionally(IllegalStateException("LSP Server not running"))
-        }
+        return AmazonQLspService.executeIfRunning(project) { server ->
+            server.updateTokenCredentials(payload)
+        } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
     }
 
-    override fun deleteTokenCredentials(): CompletableFuture<Unit> {
-        return CompletableFuture<Unit>().also { completableFuture ->
+    override fun deleteTokenCredentials(): CompletableFuture<Unit> =
+        CompletableFuture<Unit>().also { completableFuture ->
             AmazonQLspService.executeIfRunning(project) { server ->
                 server.deleteTokenCredentials()
                 completableFuture.complete(null)
             } ?: completableFuture.completeExceptionally(IllegalStateException("LSP Server not running"))
         }
-    }
 
     private fun createUpdateCredentialsPayload(token: String): UpdateCredentialsPayload =
         UpdateCredentialsPayload(

plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsServiceTest.kt

@@ -5,11 +5,8 @@ package software.aws.toolkits.jetbrains.services.amazonq.lsp.auth
 
 import com.intellij.openapi.components.serviceIfCreated
 import com.intellij.openapi.project.Project
-import io.mockk.coEvery
 import io.mockk.every
-import io.mockk.just
 import io.mockk.mockk
-import io.mockk.runs
 import io.mockk.verify
 import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
 import org.junit.Before
@@ -38,10 +35,10 @@ class DefaultAuthCredentialsServiceTest {
         every { project.serviceIfCreated<AmazonQLspService>() } returns mockLspService
 
         // Mock the LSP service's executeSync method as a suspend function
-        coEvery {
-            mockLspService.executeSync(any())
+        every {
+            mockLspService.executeSync<CompletableFuture<ResponseMessage>>(any())
         } coAnswers {
-            val func = firstArg<suspend (AmazonQLanguageServer) -> Unit>()
+            val func = firstArg<suspend (AmazonQLanguageServer) -> CompletableFuture<ResponseMessage>>()
             func.invoke(mockLanguageServer)
         }
 
@@ -87,7 +84,7 @@ class DefaultAuthCredentialsServiceTest {
 
     @Test
     fun `test deleteTokenCredentials success`() {
-        every { mockLanguageServer.deleteTokenCredentials() } just runs
+        every { mockLanguageServer.deleteTokenCredentials() } returns CompletableFuture.completedFuture(Unit)
 
         sut.deleteTokenCredentials()