CodeWhisperer: Add telemetry optout status to SendTelemetryEventRequest (#3887)

Updated the API and put the field in SendTelemetryEvent API requests.

Author: andrewyuq

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt

@@ -21,6 +21,7 @@ import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUr
 import software.amazon.awssdk.services.codewhispererruntime.model.CreateUploadUrlResponse
 import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsRequest
 import software.amazon.awssdk.services.codewhispererruntime.model.GenerateCompletionsResponse
+import software.amazon.awssdk.services.codewhispererruntime.model.OptOutPreference
 import software.amazon.awssdk.services.codewhispererruntime.model.SendTelemetryEventResponse
 import software.amazon.awssdk.services.codewhispererruntime.model.SuggestionState
 import software.aws.toolkits.core.utils.getLogger
@@ -39,6 +40,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.service.RequestCon
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.ResponseContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.transform
+import software.aws.toolkits.jetbrains.settings.AwsSettings
 import software.aws.toolkits.telemetry.CodewhispererCompletionType
 import software.aws.toolkits.telemetry.CodewhispererSuggestionState
 import java.time.Instant
@@ -211,6 +213,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
                     it.generatedLine(lineCount)
                 }
             }
+            requestBuilder.optOutPreference(getTelemetryOptoutPreference())
         }
     }
 
@@ -227,6 +230,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
                 it.timestamp(Instant.now())
             }
         }
+        requestBuilder.optOutPreference(getTelemetryOptoutPreference())
     }
 
     override fun sendUserModificationTelemetry(
@@ -246,6 +250,7 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
                 it.timestamp(Instant.now())
             }
         }
+        requestBuilder.optOutPreference(getTelemetryOptoutPreference())
     }
 
     override fun sendCodeScanTelemetry(
@@ -261,8 +266,16 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
                 it.timestamp(Instant.now())
             }
         }
+        requestBuilder.optOutPreference(getTelemetryOptoutPreference())
     }
 
+    private fun getTelemetryOptoutPreference() =
+        if (AwsSettings.getInstance().isTelemetryEnabled) {
+            OptOutPreference.OPTIN
+        } else {
+            OptOutPreference.OPTOUT
+        }
+
     override fun dispose() {
         if (this::mySigv4Client.isLazyInitialized) {
             mySigv4Client.close()

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

@@ -8,10 +8,12 @@ import com.intellij.testFramework.ProjectRule
 import com.intellij.testFramework.RuleChain
 import com.intellij.testFramework.replaceService
 import org.assertj.core.api.Assertions.assertThat
+import org.junit.After
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
 import org.mockito.kotlin.any
+import org.mockito.kotlin.argThat
 import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.doReturnConsecutively
@@ -42,6 +44,9 @@ import software.amazon.awssdk.services.codewhispererruntime.model.GetCodeAnalysi
 import software.amazon.awssdk.services.codewhispererruntime.model.GetCodeAnalysisResponse
 import software.amazon.awssdk.services.codewhispererruntime.model.ListCodeAnalysisFindingsRequest
 import software.amazon.awssdk.services.codewhispererruntime.model.ListCodeAnalysisFindingsResponse
+import software.amazon.awssdk.services.codewhispererruntime.model.OptOutPreference
+import software.amazon.awssdk.services.codewhispererruntime.model.SendTelemetryEventRequest
+import software.amazon.awssdk.services.codewhispererruntime.model.SendTelemetryEventResponse
 import software.amazon.awssdk.services.codewhispererruntime.model.StartCodeAnalysisRequest
 import software.amazon.awssdk.services.codewhispererruntime.model.StartCodeAnalysisResponse
 import software.amazon.awssdk.services.codewhispererruntime.paginators.GenerateCompletionsIterable
@@ -62,6 +67,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestU
 import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.sdkHttpResponse
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptorImpl
+import software.aws.toolkits.jetbrains.settings.AwsSettings
 
 class CodeWhispererClientAdaptorTest {
     val projectRule = ProjectRule()
@@ -80,6 +86,7 @@ class CodeWhispererClientAdaptorTest {
 
     private lateinit var sut: CodeWhispererClientAdaptor
     private lateinit var connectionManager: ToolkitConnectionManager
+    private var isTelemetryEnabledDefault: Boolean = false
 
     @Before
     fun setup() {
@@ -99,6 +106,7 @@ class CodeWhispererClientAdaptorTest {
             on { startCodeAnalysis(any<StartCodeAnalysisRequest>()) } doReturn startCodeAnalysisResponse
             on { getCodeAnalysis(any<GetCodeAnalysisRequest>()) } doReturn getCodeAnalysisResponse
             on { listCodeAnalysisFindings(any<ListCodeAnalysisFindingsRequest>()) } doReturn listCodeAnalysisFindingsResponse
+            on { sendTelemetryEvent(any<SendTelemetryEventRequest>()) } doReturn sendtelemetryEventResponse
         }
 
         val mockConnection = mock<BearerSsoConnection>()
@@ -110,6 +118,13 @@ class CodeWhispererClientAdaptorTest {
             } doReturn authManagerRule.createConnection(ManagedSsoProfile("us-east-1", aString(), emptyList())) as AwsBearerTokenConnection
         }
         projectRule.project.replaceService(ToolkitConnectionManager::class.java, connectionManager, disposableRule.disposable)
+
+        isTelemetryEnabledDefault = AwsSettings.getInstance().isTelemetryEnabled
+    }
+
+    @After
+    fun tearDown() {
+        AwsSettings.getInstance().isTelemetryEnabled = isTelemetryEnabledDefault
     }
 
     @Test
@@ -239,6 +254,55 @@ class CodeWhispererClientAdaptorTest {
         }
     }
 
+    @Test
+    fun `sendTelemetryEvent for userTriggerDecision respects telemetry optin status`() {
+        sendTelemetryEventOptOutCheckHelper {
+            sut.putUserTriggerDecisionTelemetry(
+                aRequestContext(projectRule.project),
+                aResponseContext(),
+                aCompletionType(),
+                aSuggestionState(),
+                0,
+                1
+            )
+        }
+    }
+
+    @Test
+    fun `sendTelemetryEvent for codePercentage respects telemetry optin status`() {
+        sendTelemetryEventOptOutCheckHelper {
+            sut.putCodePercentageTelemetry(aProgrammingLanguage(), 0, 1)
+        }
+    }
+
+    @Test
+    fun `sendTelemetryEvent for codeScan respects telemetry optin status`() {
+        sendTelemetryEventOptOutCheckHelper {
+            sut.sendCodeScanTelemetry(aProgrammingLanguage(), null)
+        }
+    }
+
+    @Test
+    fun `sendTelemetryEvent for userModification respects telemetry optin status`() {
+        sendTelemetryEventOptOutCheckHelper {
+            sut.sendUserModificationTelemetry(aString(), aString(), aProgrammingLanguage(), 0.0)
+        }
+    }
+
+    private fun sendTelemetryEventOptOutCheckHelper(mockApiCall: () -> Unit) {
+        AwsSettings.getInstance().isTelemetryEnabled = true
+        mockApiCall()
+        verify(bearerClient).sendTelemetryEvent(
+            argThat<SendTelemetryEventRequest> { optOutPreference() == OptOutPreference.OPTIN }
+        )
+
+        AwsSettings.getInstance().isTelemetryEnabled = false
+        mockApiCall()
+        verify(bearerClient).sendTelemetryEvent(
+            argThat<SendTelemetryEventRequest> { optOutPreference() == OptOutPreference.OPTOUT }
+        )
+    }
+
     private companion object {
         val createCodeScanRequest = CreateCodeScanRequest.builder()
             .artifacts(mapOf(ArtifactType.SOURCE_CODE to "foo"))
@@ -294,6 +358,8 @@ class CodeWhispererClientAdaptorTest {
             .sdkHttpResponse(sdkHttpResponse)
             .build() as ListCodeAnalysisFindingsResponse
 
+        val sendtelemetryEventResponse = SendTelemetryEventResponse.builder().build()
+
         private val generateCompletionsPaginatorResponse: GenerateCompletionsIterable = mock()
 
         private val createCodeScanUploadUrlResponse: CreateCodeScanUploadUrlResponse = mock()