33
44package software.aws.toolkits.jetbrains.services.amazonq.workspace.context
55
6+ import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
67import com.github.tomakehurst.wiremock.client.WireMock.aResponse
78import com.github.tomakehurst.wiremock.client.WireMock.any
9+ import com.github.tomakehurst.wiremock.client.WireMock.equalTo
10+ import com.github.tomakehurst.wiremock.client.WireMock.postRequestedFor
811import com.github.tomakehurst.wiremock.client.WireMock.stubFor
912import com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo
1013import com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig
@@ -25,9 +28,12 @@ import org.mockito.kotlin.times
2528import org.mockito.kotlin.verify
2629import org.mockito.kotlin.whenever
2730import software.aws.toolkits.jetbrains.services.amazonq.project.EncoderServer
31+ import software.aws.toolkits.jetbrains.services.amazonq.project.IndexRequest
2832import software.aws.toolkits.jetbrains.services.amazonq.project.LspMessage
2933import software.aws.toolkits.jetbrains.services.amazonq.project.ProjectContextProvider
34+ import software.aws.toolkits.jetbrains.services.amazonq.project.QueryChatRequest
3035import software.aws.toolkits.jetbrains.services.amazonq.project.RelevantDocument
36+ import software.aws.toolkits.jetbrains.services.amazonq.project.UpdateIndexRequest
3137import software.aws.toolkits.jetbrains.utils.rules.CodeInsightTestFixtureRule
3238import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule
3339import java.net.ConnectException
@@ -48,6 +54,8 @@ class ProjectContextProviderTest {
4854 private lateinit var encoderServer: EncoderServer
4955 private lateinit var sut: ProjectContextProvider
5056
57+ private val mapper = jacksonObjectMapper()
58+
5159 @Before
5260 fun setup () {
5361 encoderServer = spy(EncoderServer (project))
@@ -91,6 +99,85 @@ class ProjectContextProviderTest {
9199 assertThat(LspMessage .GetUsageMetrics .endpoint).isEqualTo(" getUsage" )
92100 }
93101
102+ @Test
103+ fun `index should send files within the project to lsp` () {
104+ projectRule.fixture.addFileToProject(" Foo.java" , " foo" )
105+ projectRule.fixture.addFileToProject(" Bar.java" , " bar" )
106+ projectRule.fixture.addFileToProject(" Baz.java" , " baz" )
107+
108+ sut.index()
109+
110+ val request = IndexRequest (listOf (" /src/Foo.java" , " /src/Bar.java" , " /src/Baz.java" ), " /src" , false )
111+ assertThat(request.filePaths).hasSize(3 )
112+ assertThat(request.filePaths).satisfies({
113+ it.contains(" /src/Foo.java" ) &&
114+ it.contains(" /src/Baz.java" ) &&
115+ it.contains(" /src/Bar.java" )
116+
117+ })
118+
119+ wireMock.verify(
120+ 1 ,
121+ postRequestedFor(urlPathEqualTo(" /indexFiles" ))
122+ .withHeader(" Content-Type" , equalTo(" text/plain" ))
123+ // comment it out because order matters and will cause json string different
124+ // .withRequestBody(equalTo(encryptedRequest))
125+ )
126+ }
127+
128+ @Test
129+ fun `updateIndex will not send message to lsp if index is not complete` () {
130+ sut.isIndexComplete.set(false )
131+
132+ sut.updateIndex(" foo.java" )
133+
134+ assertThat(wireMock.allServeEvents).isEmpty()
135+ wireMock.verify(
136+ 0 ,
137+ postRequestedFor(urlPathEqualTo(" /updateIndex" ))
138+ .withHeader(" Content-Type" , equalTo(" text/plain" ))
139+ )
140+ }
141+
142+ @Test
143+ fun `updateIndex should send correct encrypted request to lsp` () {
144+ sut.isIndexComplete.set(true )
145+
146+ sut.updateIndex(" foo.java" )
147+ val request = UpdateIndexRequest (" foo.java" )
148+ val requestJson = mapper.writeValueAsString(request)
149+
150+ assertThat(mapper.readTree(requestJson)).isEqualTo(mapper.readTree(""" { "filePath": "foo.java" }""" ))
151+
152+ val encryptedRequest = encoderServer.encrypt(requestJson)
153+
154+ wireMock.verify(
155+ 1 ,
156+ postRequestedFor(urlPathEqualTo(" /updateIndex" ))
157+ .withHeader(" Content-Type" , equalTo(" text/plain" ))
158+ .withRequestBody(equalTo(encryptedRequest))
159+ )
160+ }
161+
162+ @Test
163+ fun `query should send correct encrypted request to lsp` () {
164+ sut.query(" foo" )
165+
166+ val request = QueryChatRequest (" foo" )
167+ val requestJson = mapper.writeValueAsString(request)
168+
169+ assertThat(mapper.readTree(requestJson)).isEqualTo(mapper.readTree(""" { "query": "foo" }""" ))
170+
171+ val encryptedRequest = encoderServer.encrypt(requestJson)
172+
173+ wireMock.verify(
174+ 1 ,
175+ postRequestedFor(urlPathEqualTo(" /query" ))
176+ .withHeader(" Content-Type" , equalTo(" text/plain" ))
177+ .withRequestBody(equalTo(encryptedRequest))
178+ )
179+ }
180+
94181 @Test
95182 fun `query chat should return empty if result set non deserializable` () = runTest {
96183 stubFor(
0 commit comments