[CodeWhisperer] Fix the logic to detect invalid references (#3243)

* Fix the logic to detect invalid references

Author: andrewyuq

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt

@@ -451,9 +451,12 @@ class CodeWhispererService {
         // If contentSpans in reference are not consistent with content(recommendations),
         // remove the incorrect references.
         val validatedRecommendations = response.recommendations().map {
-            val validReference = it.hasReferences() && it.references().isNotEmpty() &&
-                it.content().length == it.references().last().recommendationContentSpan().end()
-            if (validReference) {
+            val validReferences = it.hasReferences() && it.references().isNotEmpty() &&
+                it.references().none { reference ->
+                    val span = reference.recommendationContentSpan()
+                    span.start() > span.end() || span.start() < 0 || span.end() > it.content().length
+                }
+            if (validReferences) {
                 it
             } else {
                 it.toBuilder().references(DefaultSdkAutoConstructList.getInstance()).build()

jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererReferencesTest.kt

@@ -0,0 +1,67 @@
+// Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.codewhisperer
+
+import com.intellij.testFramework.runInEdtAndWait
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.Test
+import org.mockito.kotlin.any
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.doAnswer
+import org.mockito.kotlin.stub
+import org.mockito.kotlin.timeout
+import org.mockito.kotlin.verify
+import software.amazon.awssdk.services.codewhisperer.model.ListRecommendationsRequest
+import software.amazon.awssdk.services.codewhisperer.model.ListRecommendationsResponse
+import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.generateMockRecommendationDetail
+import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.getReferenceInfo
+import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.metadata
+import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.sdkHttpResponse
+import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext
+
+class CodeWhispererReferencesTest : CodeWhispererTestBase() {
+
+    @Test
+    fun `test references with invalid content span will not remove the references from the recommendation`() {
+        testReferencesRangeValidity("test", -1, 4, invalid = true, 1)
+        testReferencesRangeValidity("test", 0, 5, invalid = true, 2)
+        testReferencesRangeValidity("test", 2, 1, invalid = true, 3)
+        testReferencesRangeValidity("test", 2, 2, invalid = false, 4)
+        testReferencesRangeValidity("test", 0, 4, invalid = false, 5)
+        testReferencesRangeValidity("test", 0, 2, invalid = false, 6)
+        testReferencesRangeValidity("test", 3, 4, invalid = false, 7)
+    }
+
+    private fun testReferencesRangeValidity(content: String, start: Int, end: Int, invalid: Boolean, times: Int) {
+        val referenceInfo = getReferenceInfo()
+        val invalidReferencesResponse = ListRecommendationsResponse.builder()
+            .recommendations(
+                generateMockRecommendationDetail(content, referenceInfo.first, referenceInfo.second, start, end),
+            )
+            .nextToken("")
+            .responseMetadata(metadata)
+            .sdkHttpResponse(sdkHttpResponse)
+            .build() as ListRecommendationsResponse
+
+        mockClient.stub {
+            on {
+                mockClient.listRecommendations(any<ListRecommendationsRequest>())
+            } doAnswer {
+                invalidReferencesResponse
+            }
+        }
+
+        val statesCaptor = argumentCaptor<InvocationContext>()
+        withCodeWhispererServiceInvokedAndWait {
+            verify(popupManagerSpy, timeout(5000).times(times))
+                .render(statesCaptor.capture(), any(), any())
+            statesCaptor.lastValue.recommendationContext.details.forEach {
+                assertThat(it.recommendation.references().isEmpty()).isEqualTo(invalid)
+            }
+            runInEdtAndWait {
+                popupManagerSpy.popupComponents.acceptButton.doClick()
+            }
+        }
+    }
+}

jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestUtil.kt

@@ -26,10 +26,10 @@ object CodeWhispererTestUtil {
         Pair("Apache-2.0", "testRepo2"),
         Pair("BSD-4-Clause", "testRepo3")
     )
-    private val metadata: DefaultAwsResponseMetadata = DefaultAwsResponseMetadata.create(
+    val metadata: DefaultAwsResponseMetadata = DefaultAwsResponseMetadata.create(
         mapOf(AWS_REQUEST_ID to testRequestId)
     )
-    private val sdkHttpResponse = SdkHttpResponse.builder().headers(
+    val sdkHttpResponse = SdkHttpResponse.builder().headers(
         mapOf(CodeWhispererService.KET_SESSION_ID to listOf(testSessionId))
     ).build()
     private val errorDetail = AwsErrorDetails.builder()
@@ -100,21 +100,37 @@ object CodeWhispererTestUtil {
     const val javaTestContext = "public class Test {\n    public static void main\n}"
 
     internal fun generateMockRecommendationDetail(content: String): Recommendation {
-        val pair = testReferenceInfoPair[Random.nextInt(testReferenceInfoPair.size)]
-        return Recommendation.builder()
-            .content(content)
+        val referenceInfo = getReferenceInfo()
+        return Recommendation.builder().content(content)
             .references(
-                Reference.builder()
-                    .licenseName(pair.first)
-                    .repository(pair.second)
-                    .recommendationContentSpan(
-                        Span.builder()
-                            .start(0)
-                            .end(content.length)
-                            .build()
-                    )
-                    .build()
+                generateMockReferences(referenceInfo.first, referenceInfo.second, 0, content.length)
             )
             .build()
     }
+
+    internal fun getReferenceInfo() = testReferenceInfoPair[Random.nextInt(testReferenceInfoPair.size)]
+
+    internal fun generateMockRecommendationDetail(
+        content: String,
+        licenseName: String,
+        repository: String,
+        start: Int,
+        end: Int
+    ): Recommendation =
+        Recommendation.builder()
+            .content(content)
+            .references(generateMockReferences(licenseName, repository, start, end))
+            .build()
+
+    private fun generateMockReferences(licenseName: String, repository: String, start: Int, end: Int) =
+        Reference.builder()
+            .licenseName(licenseName)
+            .repository(repository)
+            .recommendationContentSpan(
+                Span.builder()
+                    .start(start)
+                    .end(end)
+                    .build()
+            )
+            .build()
 }