Skip to content

Commit 9eaf004

Browse files
rlirli
committed
tst
1 parent f1bec1c commit 9eaf004

File tree

2 files changed

+132
-120
lines changed

2 files changed

+132
-120
lines changed

‎plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/workspace/context/ProjectContextProviderTest.kt‎

Lines changed: 123 additions & 107 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import com.intellij.testFramework.DisposableRule
1919
import com.intellij.testFramework.replaceService
2020
import io.mockk.every
2121
import io.mockk.spyk
22+
import kotlinx.coroutines.Dispatchers
2223
import kotlinx.coroutines.ExperimentalCoroutinesApi
2324
import kotlinx.coroutines.TimeoutCancellationException
2425
import kotlinx.coroutines.test.StandardTestDispatcher
@@ -225,43 +226,49 @@ class ProjectContextProviderTest {
225226

226227
@Test
227228
fun `query should send correct encrypted request to lsp`() = runTest {
228-
sut = ProjectContextProvider(project, encoderServer, this)
229-
val r = sut.query("foo", null)
230-
advanceUntilIdle()
229+
// use real time
230+
withContext(Dispatchers.Default.limitedParallelism(1)) {
231+
sut = ProjectContextProvider(project, encoderServer, this)
232+
val r = sut.query("foo", null)
233+
advanceUntilIdle()
231234

232-
val request = QueryChatRequest("foo")
233-
val requestJson = mapper.writeValueAsString(request)
235+
val request = QueryChatRequest("foo")
236+
val requestJson = mapper.writeValueAsString(request)
234237

235-
assertThat(mapper.readTree(requestJson)).isEqualTo(mapper.readTree("""{ "query": "foo" }"""))
238+
assertThat(mapper.readTree(requestJson)).isEqualTo(mapper.readTree("""{ "query": "foo" }"""))
236239

237-
val encryptedRequest = encoderServer.encrypt(requestJson)
240+
val encryptedRequest = encoderServer.encrypt(requestJson)
238241

239-
wireMock.verify(
240-
1,
241-
postRequestedFor(urlPathEqualTo("/query"))
242-
.withHeader("Content-Type", equalTo("text/plain"))
243-
.withRequestBody(equalTo(encryptedRequest))
244-
)
242+
wireMock.verify(
243+
1,
244+
postRequestedFor(urlPathEqualTo("/query"))
245+
.withHeader("Content-Type", equalTo("text/plain"))
246+
.withRequestBody(equalTo(encryptedRequest))
247+
)
248+
}
245249
}
246250

247251
@Test
248252
fun `queryInline should send correct encrypted request to lsp`() = runTest {
249-
sut = ProjectContextProvider(project, encoderServer, this)
250-
sut.queryInline("foo", "Foo.java", InlineContextTarget.CODEMAP)
251-
advanceUntilIdle()
253+
// use real time
254+
withContext(Dispatchers.Default.limitedParallelism(1)) {
255+
sut = ProjectContextProvider(project, encoderServer, this)
256+
sut.queryInline("foo", "Foo.java", InlineContextTarget.CODEMAP)
257+
advanceUntilIdle()
252258

253-
val request = QueryInlineCompletionRequest("foo", "Foo.java", "codemap")
254-
val requestJson = mapper.writeValueAsString(request)
259+
val request = QueryInlineCompletionRequest("foo", "Foo.java", "codemap")
260+
val requestJson = mapper.writeValueAsString(request)
255261

256-
assertThat(mapper.readTree(requestJson)).isEqualTo(mapper.readTree("""{ "query": "foo", "filePath": "Foo.java", "target": "codemap" }"""))
262+
assertThat(mapper.readTree(requestJson)).isEqualTo(mapper.readTree("""{ "query": "foo", "filePath": "Foo.java", "target": "codemap" }"""))
257263

258-
val encryptedRequest = encoderServer.encrypt(requestJson)
259-
wireMock.verify(
260-
1,
261-
postRequestedFor(urlPathEqualTo("/queryInlineProjectContext"))
262-
.withHeader("Content-Type", equalTo("text/plain"))
263-
.withRequestBody(equalTo(encryptedRequest))
264-
)
264+
val encryptedRequest = encoderServer.encrypt(requestJson)
265+
wireMock.verify(
266+
1,
267+
postRequestedFor(urlPathEqualTo("/queryInlineProjectContext"))
268+
.withHeader("Content-Type", equalTo("text/plain"))
269+
.withRequestBody(equalTo(encryptedRequest))
270+
)
271+
}
265272
}
266273

267274
@Test
@@ -287,6 +294,8 @@ class ProjectContextProviderTest {
287294

288295
@Test
289296
fun `query chat should return deserialized relevantDocument`() = runTest {
297+
// use real time
298+
withContext(Dispatchers.Default.limitedParallelism(1)) {
290299
sut = ProjectContextProvider(project, encoderServer, this)
291300
val r = sut.query("foo", null)
292301
advanceUntilIdle()
@@ -304,6 +313,7 @@ class ProjectContextProviderTest {
304313
)
305314
)
306315
}
316+
}
307317

308318
@Test
309319
fun `query inline should throw if resultset not deserializable`() {
@@ -334,31 +344,34 @@ class ProjectContextProviderTest {
334344

335345
@Test
336346
fun `query inline should return deserialized bm25 chunks`() = runTest {
337-
sut = ProjectContextProvider(project, encoderServer, this)
338-
advanceUntilIdle()
339-
val r = sut.queryInline("foo", "filepath", InlineContextTarget.CODEMAP)
340-
assertThat(r).hasSize(3)
341-
assertThat(r[0]).isEqualTo(
342-
InlineBm25Chunk(
343-
"content1",
344-
"file1",
345-
0.1
347+
// use real time
348+
withContext(Dispatchers.Default.limitedParallelism(1)) {
349+
sut = ProjectContextProvider(project, encoderServer, this)
350+
advanceUntilIdle()
351+
val r = sut.queryInline("foo", "filepath", InlineContextTarget.CODEMAP)
352+
assertThat(r).hasSize(3)
353+
assertThat(r[0]).isEqualTo(
354+
InlineBm25Chunk(
355+
"content1",
356+
"file1",
357+
0.1
358+
)
346359
)
347-
)
348-
assertThat(r[1]).isEqualTo(
349-
InlineBm25Chunk(
350-
"content2",
351-
"file2",
352-
0.2
360+
assertThat(r[1]).isEqualTo(
361+
InlineBm25Chunk(
362+
"content2",
363+
"file2",
364+
0.2
365+
)
353366
)
354-
)
355-
assertThat(r[2]).isEqualTo(
356-
InlineBm25Chunk(
357-
"content3",
358-
"file3",
359-
0.3
367+
assertThat(r[2]).isEqualTo(
368+
InlineBm25Chunk(
369+
"content3",
370+
"file3",
371+
0.3
372+
)
360373
)
361-
)
374+
}
362375
}
363376

364377
@Test
@@ -431,78 +444,81 @@ class ProjectContextProviderTest {
431444

432445
@Test
433446
fun `test query payload is encrypted`() = runTest {
434-
sut = ProjectContextProvider(project, encoderServer, this)
435-
sut.query("what does this project do", null)
436-
advanceUntilIdle()
437-
verify(encoderServer, times(1)).encrypt(any())
447+
// use real time
448+
withContext(Dispatchers.Default.limitedParallelism(1)) {
449+
sut = ProjectContextProvider(project, encoderServer, this)
450+
sut.query("what does this project do", null)
451+
advanceUntilIdle()
452+
verify(encoderServer, times(1)).encrypt(any())
453+
}
438454
}
439455

440456
private fun createMockServer() = WireMockRule(wireMockConfig().dynamicPort())
441457
}
442458

443459
// language=JSON
444460
val validQueryInlineResponse = """
445-
[
446-
{
447-
"content": "content1",
448-
"filePath": "file1",
449-
"score": 0.1
450-
},
451-
{
452-
"content": "content2",
453-
"filePath": "file2",
454-
"score": 0.2
455-
},
456-
{
457-
"content": "content3",
458-
"filePath": "file3",
459-
"score": 0.3
460-
}
461-
]
461+
[
462+
{
463+
"content": "content1",
464+
"filePath": "file1",
465+
"score": 0.1
466+
},
467+
{
468+
"content": "content2",
469+
"filePath": "file2",
470+
"score": 0.2
471+
},
472+
{
473+
"content": "content3",
474+
"filePath": "file3",
475+
"score": 0.3
476+
}
477+
]
462478
""".trimIndent()
463479

464480
// language=JSON
465481
val validQueryChatResponse = """
466-
[
467-
{
468-
"filePath": "file1",
469-
"content": "content1",
470-
"id": "id1",
471-
"index": "index1",
472-
"vec": [
473-
"vec_1-1",
474-
"vec_1-2",
475-
"vec_1-3"
476-
],
477-
"context": "context1",
478-
"prev": "prev1",
479-
"next": "next1",
480-
"relativePath": "relativeFilePath1",
481-
"programmingLanguage": "language1"
482-
},
483-
{
484-
"filePath": "file2",
485-
"content": "content2",
486-
"id": "id2",
487-
"index": "index2",
488-
"vec": [
489-
"vec_2-1",
490-
"vec_2-2",
491-
"vec_2-3"
492-
],
493-
"context": "context2",
494-
"prev": "prev2",
495-
"next": "next2",
496-
"relativePath": "relativeFilePath2",
497-
"programmingLanguage": "language2"
498-
}
499-
]
482+
[
483+
{
484+
"filePath": "file1",
485+
"content": "content1",
486+
"id": "id1",
487+
"index": "index1",
488+
"vec": [
489+
"vec_1-1",
490+
"vec_1-2",
491+
"vec_1-3"
492+
],
493+
"context": "context1",
494+
"prev": "prev1",
495+
"next": "next1",
496+
"relativePath": "relativeFilePath1",
497+
"programmingLanguage": "language1"
498+
},
499+
{
500+
"filePath": "file2",
501+
"content": "content2",
502+
"id": "id2",
503+
"index": "index2",
504+
"vec": [
505+
"vec_2-1",
506+
"vec_2-2",
507+
"vec_2-3"
508+
],
509+
"context": "context2",
510+
"prev": "prev2",
511+
"next": "next2",
512+
"relativePath": "relativeFilePath2",
513+
"programmingLanguage": "language2"
514+
}
515+
]
500516
""".trimIndent()
501517

502518
// language=JSON
503519
val validGetUsageResponse = """
504-
{
505-
"memoryUsage":123,
506-
"cpuUsage":456
507-
}
520+
{
521+
"memoryUsage":123,
522+
"cpuUsage":456
523+
}
508524
""".trimIndent()

‎plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextProvider.kt‎

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ import com.intellij.openapi.vfs.VirtualFileVisitor
1717
import com.intellij.openapi.vfs.isFile
1818
import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
1919
import kotlinx.coroutines.CoroutineScope
20-
import kotlinx.coroutines.async
2120
import kotlinx.coroutines.delay
2221
import kotlinx.coroutines.launch
2322
import kotlinx.coroutines.runBlocking
@@ -170,20 +169,17 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
170169

171170
// TODO: rename queryChat
172171
suspend fun query(prompt: String, timeout: Long?): List<RelevantDocument> = withTimeout(timeout ?: CHAT_EXPLICIT_PROJECT_CONTEXT_TIMEOUT) {
173-
cs.async {
174-
val encrypted = encryptRequest(QueryChatRequest(prompt))
175-
val response = sendMsgToLsp(LspMessage.QueryChat, encrypted) ?: return@async emptyList()
176-
val parsedResponse = mapper.readValue<List<Chunk>>(response.responseBody)
177-
queryResultToRelevantDocuments(parsedResponse)
178-
}.await()
172+
val encrypted = encryptRequest(QueryChatRequest(prompt))
173+
val response = sendMsgToLsp(LspMessage.QueryChat, encrypted) ?: return@withTimeout emptyList()
174+
val parsedResponse = mapper.readValue<List<Chunk>>(response.responseBody)
175+
176+
return@withTimeout queryResultToRelevantDocuments(parsedResponse)
179177
}
180178

181179
suspend fun queryInline(query: String, filePath: String, target: InlineContextTarget): List<InlineBm25Chunk> = withTimeout(SUPPLEMENTAL_CONTEXT_TIMEOUT) {
182-
cs.async {
183-
val encrypted = encryptRequest(QueryInlineCompletionRequest(query, filePath, target.toString()))
184-
val r = sendMsgToLsp(LspMessage.QueryInlineCompletion, encrypted) ?: return@async emptyList()
185-
return@async mapper.readValue<List<InlineBm25Chunk>>(r.responseBody)
186-
}.await()
180+
val encrypted = encryptRequest(QueryInlineCompletionRequest(query, filePath, target.toString()))
181+
val r = sendMsgToLsp(LspMessage.QueryInlineCompletion, encrypted) ?: return@withTimeout emptyList()
182+
return@withTimeout mapper.readValue<List<InlineBm25Chunk>>(r.responseBody)
187183
}
188184

189185
suspend fun getUsage(): Usage? {
@@ -200,7 +196,7 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En
200196
@RequiresBackgroundThread
201197
fun updateIndex(filePaths: List<String>, mode: IndexUpdateMode) {
202198
val encrypted = encryptRequest(UpdateIndexRequest(filePaths, mode.command))
203-
runBlocking(IO) { sendMsgToLsp(LspMessage.UpdateIndex, encrypted) }
199+
runBlocking { sendMsgToLsp(LspMessage.UpdateIndex, encrypted) }
204200
}
205201

206202
private fun recordIndexWorkspace(

0 commit comments

Comments
 (0)