Skip to content

Commit ebf63bc

Browse files
rlirli
committed
tst
1 parent f3fd620 commit ebf63bc

File tree

5 files changed

+49
-24
lines changed

5 files changed

+49
-24
lines changed

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,13 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
191191
val params = createInlineCompletionParams(requestContext.editor, requestContext.triggerTypeInfo, nextToken)
192192
server.inlineCompletionWithReferences(params)
193193
}
194-
val completion = result?.await() ?: break
194+
val completion = result?.await()
195+
if (completion == null) {
196+
// no result / not running
197+
CodeWhispererInvocationStatus.getInstance().finishInvocation()
198+
break
199+
}
200+
195201
nextToken = completion.partialResultToken
196202
val endTime = System.nanoTime()
197203
val latency = TimeUnit.NANOSECONDS.toMillis(endTime - startTime).toDouble()

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

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ package software.aws.toolkits.jetbrains.services.codewhisperer
66
import com.intellij.openapi.application.ApplicationManager
77
import com.intellij.openapi.editor.Editor
88
import com.intellij.openapi.project.Project
9+
import com.intellij.openapi.util.Disposer
910
import com.intellij.psi.PsiFile
1011
import com.intellij.testFramework.DisposableRule
1112
import com.intellij.testFramework.RuleChain
@@ -18,6 +19,7 @@ import kotlinx.coroutines.CoroutineScope
1819
import kotlinx.coroutines.Deferred
1920
import kotlinx.coroutines.Job
2021
import kotlinx.coroutines.delay
22+
import kotlinx.coroutines.flow.first
2123
import kotlinx.coroutines.test.TestScope
2224
import kotlinx.coroutines.test.runTest
2325
import kotlinx.coroutines.yield
@@ -31,12 +33,14 @@ import org.mockito.kotlin.any
3133
import org.mockito.kotlin.argumentCaptor
3234
import org.mockito.kotlin.doAnswer
3335
import org.mockito.kotlin.doNothing
36+
import org.mockito.kotlin.doReturn
3437
import org.mockito.kotlin.doSuspendableAnswer
3538
import org.mockito.kotlin.spy
3639
import org.mockito.kotlin.stub
3740
import org.mockito.kotlin.timeout
3841
import org.mockito.kotlin.verify
3942
import org.mockito.kotlin.whenever
43+
import org.mockito.kotlin.wheneverBlocking
4044
import software.amazon.awssdk.services.ssooidc.SsoOidcClient
4145
import software.aws.toolkits.jetbrains.core.MockClientManagerRule
4246
import software.aws.toolkits.jetbrains.core.credentials.ManagedSsoProfile
@@ -112,7 +116,7 @@ open class CodeWhispererTestBase {
112116
protected lateinit var codeScanManager: CodeWhispererCodeScanManager
113117

114118
@Before
115-
open fun setUp() {
119+
open fun setUp() = runTest {
116120
mockLanguageServer = mockk()
117121
val starter = object : AmazonQServerInstanceStarter {
118122
override fun start(
@@ -141,10 +145,13 @@ open class CodeWhispererTestBase {
141145
}
142146
}
143147

144-
mockLspService = spy(AmazonQLspService(starter, projectRule.project, TestScope()))
148+
mockLspService = spy(AmazonQLspService(starter, projectRule.project, this))
145149

146150
// Mock the service methods on Project
147151
projectRule.project.replaceService(AmazonQLspService::class.java, mockLspService, disposableRule.disposable)
152+
// wait for init to finish
153+
mockLspService.instanceFlow.first()
154+
148155
mockLspInlineCompletionResponse(pythonResponse)
149156

150157
mockClientManagerRule.create<SsoOidcClient>()
@@ -165,12 +172,10 @@ open class CodeWhispererTestBase {
165172
stateManager = spy(CodeWhispererExplorerActionManager.getInstance())
166173
recommendationManager = CodeWhispererRecommendationManager.getInstance()
167174
codewhispererService = spy(CodeWhispererService.getInstance())
168-
codewhispererService.stub {
169-
onBlocking {
170-
getWorkspaceIds(any())
171-
} doAnswer {
172-
CompletableFuture.completedFuture(LspServerConfigurations(listOf(WorkspaceInfo("file:///", "workspaceId"))))
173-
}
175+
doAnswer {
176+
CompletableFuture.completedFuture(LspServerConfigurations(listOf(WorkspaceInfo("file:///", "workspaceId"))))
177+
}.wheneverBlocking(codewhispererService) {
178+
getWorkspaceIds(any())
174179
}
175180
ApplicationManager.getApplication().replaceService(CodeWhispererService::class.java, codewhispererService, disposableRule.disposable)
176181
editorManager = CodeWhispererEditorManager.getInstance()
@@ -233,6 +238,8 @@ open class CodeWhispererTestBase {
233238
runInEdtAndWait {
234239
popupManagerSpy.closePopup()
235240
}
241+
242+
Disposer.dispose(mockLspService)
236243
}
237244

238245
fun withCodeWhispererServiceInvokedAndWait(runnable: (InvocationContext) -> Unit) {

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

Lines changed: 22 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import com.intellij.execution.process.ProcessListener
1414
import com.intellij.execution.process.ProcessOutputType
1515
import com.intellij.notification.NotificationAction
1616
import com.intellij.openapi.Disposable
17+
import com.intellij.openapi.application.ApplicationManager
1718
import com.intellij.openapi.components.Service
1819
import com.intellij.openapi.components.service
1920
import com.intellij.openapi.components.serviceIfCreated
@@ -175,20 +176,16 @@ class AmazonQLspService @VisibleForTesting constructor(
175176
var attempts = 0
176177
while (attempts < 3) {
177178
try {
178-
val result = withTimeout(30.seconds) {
179-
val instance = starter.start(project, cs).also {
180-
Disposer.register(this@AmazonQLspService, it)
181-
}
182-
// wait for handshake to complete
183-
instance.initializeResult.join()
184-
185-
instance.also {
186-
_flowInstance.emit(it)
187-
}
179+
// no timeout; start() can download which may take long time
180+
val instance = starter.start(project, cs).also {
181+
Disposer.register(this@AmazonQLspService, it)
188182
}
183+
// wait for handshake to complete
184+
instance.initializeResult.join()
189185

190-
// withTimeout can throw
191-
return@async result
186+
return@async instance.also {
187+
_flowInstance.emit(it)
188+
}
192189
} catch (e: Exception) {
193190
LOG.warn(e) { "Failed to start LSP server" }
194191
}
@@ -203,6 +200,9 @@ class AmazonQLspService @VisibleForTesting constructor(
203200

204201
// Initialize heartbeat job
205202
heartbeatJob = cs.launch {
203+
if (ApplicationManager.getApplication().isUnitTestMode) {
204+
return@launch
205+
}
206206
while (isActive) {
207207
delay(5.seconds) // Check every 5 seconds
208208
val shouldLoop = checkConnectionStatus()
@@ -300,11 +300,20 @@ class AmazonQLspService @VisibleForTesting constructor(
300300
}
301301

302302
suspend fun<T> executeIfRunning(runnable: suspend AmazonQLspService.(AmazonQLanguageServer) -> T): T? = withContext(dispatcher) {
303-
instanceFlow.firstOrNull()?.let { runnable(it) } ?: run {
303+
val lsp = try {
304+
withTimeout(5.seconds) {
305+
val holder = mutex.withLock { instance }.await()
306+
holder.initializeResult.join()
307+
308+
holder.languageServer
309+
}
310+
} catch (_: Exception) {
304311
LOG.debug { "LSP not running" }
305312

306313
null
307314
}
315+
316+
lsp?.let { runnable(it) }
308317
}
309318

310319
fun<T> syncExecuteIfRunning(runnable: suspend AmazonQLspService.(AmazonQLanguageServer) -> T): T? =

‎plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/textdocument/TextDocumentServiceHandler.kt‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -192,8 +192,10 @@ class TextDocumentServiceHandler(
192192
)
193193

194194
// Send notification to the language server
195-
AmazonQLspService.executeIfRunning(project) { _ ->
196-
rawEndpoint.notify(ACTIVE_EDITOR_CHANGED_NOTIFICATION, params)
195+
cs.launch {
196+
AmazonQLspService.executeAsyncIfRunning(project) { _ ->
197+
rawEndpoint.notify(ACTIVE_EDITOR_CHANGED_NOTIFICATION, params)
198+
}
197199
}
198200
}
199201

‎plugins/core/jetbrains-community/tstFixtures/software/aws/toolkits/jetbrains/utils/rules/CodeInsightTestFixtureRule.kt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ open class CodeInsightTestFixtureRule(protected val testDescription: LightProjec
6060
this.description = description
6161
// This timer is cancelled but it still continues running when the test is over since it cancels lazily. This is fine, so suppress the leak
6262
ThreadLeakTracker.longRunningThreadCreated(ApplicationManager.getApplication(), "Debugger Worker launch timer")
63+
ThreadLeakTracker.longRunningThreadCreated(ApplicationManager.getApplication(), "Test worker")
6364
}
6465

6566
override fun after() {

0 commit comments

Comments
 (0)