fix unit test

leigaol · leigaol · commit 8f49fcd4d1fb · 2024-09-16T14:25:42.000-07:00
diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/chat/messenger/ChatPromptHandler.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/chat/messenger/ChatPromptHandler.kt
@@ -90,6 +90,7 @@ class ChatPromptHandler(private val telemetryHelper: TelemetryHelper) {
                     ChatMessage(tabId = tabId, triggerId = triggerId, messageId = requestId, messageType = ChatMessageType.Answer, followUps = followUps)
 
                 telemetryHelper.setResponseStreamTotalTime(tabId)
+                telemetryHelper.setResponseHasProjectContext(data, response)
                 telemetryHelper.recordAddMessage(data, response, responseText.length, statusCode, countTotalNumberOfCodeBlocks(responseText))
                 emit(response)
             }
diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/chat/telemetry/TelemetryHelper.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/chat/telemetry/TelemetryHelper.kt
@@ -96,6 +96,7 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
     // When Chat API responds to a user message (full response streamed)
     fun recordAddMessage(data: ChatRequestData, response: ChatMessage, responseLength: Int, statusCode: Int, numberOfCodeBlocks: Int) {
         val hasProjectContext = getIsProjectContextEnabled() && data.useRelevantDocuments && data.relevantTextDocuments.isNotEmpty()
+        responseHasProjectContext[response.messageId] = hasProjectContext
         AmazonqTelemetry.addMessage(
             cwsprChatConversationId = getConversationId(response.tabId).orEmpty(),
             cwsprChatMessageId = response.messageId,
@@ -120,7 +121,6 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
             codewhispererCustomizationArn = data.customization?.arn,
             cwsprChatHasProjectContext = hasProjectContext
         )
-        responseHasProjectContext[response.messageId] = hasProjectContext
         val programmingLanguage = data.activeFileContext.fileContext?.fileLanguage
         val validProgrammingLanguage = if (ChatSessionV1.validLanguages.contains(programmingLanguage)) programmingLanguage else null
 
@@ -137,7 +137,7 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
             data.message.length,
             responseLength,
             numberOfCodeBlocks,
-            getIsProjectContextEnabled() && data.useRelevantDocuments && data.relevantTextDocuments.isNotEmpty(),
+            getMessageHasProjectContext(response.messageId),
             data.customization
         ).also {
             logger.debug {
@@ -164,9 +164,6 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
 
     // When user interacts with a message (e.g. copy code, insert code, vote)
     suspend fun recordInteractWithMessage(message: IncomingCwcMessage) {
-        val hasProjectContext = { messageId: String ->
-            responseHasProjectContext.getOrDefault(messageId, false)
-        }
         val event: ChatInteractWithMessageEvent? = when (message) {
             is IncomingCwcMessage.ChatItemVoted -> {
                 AmazonqTelemetry.interactWithMessage(
@@ -178,7 +175,7 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
                         else -> CwsprChatInteractionType.Unknown
                     },
                     credentialStartUrl = getStartUrl(context.project),
-                    cwsprChatHasProjectContext = hasProjectContext(message.messageId)
+                    cwsprChatHasProjectContext = getMessageHasProjectContext(message.messageId)
                 )
                 ChatInteractWithMessageEvent.builder().apply {
                     conversationId(getConversationId(message.tabId).orEmpty())
@@ -190,7 +187,7 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
                             else -> ChatMessageInteractionType.UNKNOWN_TO_SDK_VERSION
                         }
                     )
-                    hasProjectLevelContext(hasProjectContext(message.messageId))
+                    hasProjectLevelContext(getMessageHasProjectContext(message.messageId))
                 }.build()
             }
 
@@ -200,13 +197,13 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
                     cwsprChatMessageId = message.messageId.orEmpty(),
                     cwsprChatInteractionType = CwsprChatInteractionType.ClickFollowUp,
                     credentialStartUrl = getStartUrl(context.project),
-                    cwsprChatHasProjectContext = hasProjectContext(message.messageId.orEmpty())
+                    cwsprChatHasProjectContext = getMessageHasProjectContext(message.messageId.orEmpty())
                 )
                 ChatInteractWithMessageEvent.builder().apply {
                     conversationId(getConversationId(message.tabId).orEmpty())
                     messageId(message.messageId.orEmpty())
                     interactionType(ChatMessageInteractionType.CLICK_FOLLOW_UP)
-                    hasProjectLevelContext(hasProjectContext(message.messageId.orEmpty()))
+                    hasProjectLevelContext(getMessageHasProjectContext(message.messageId.orEmpty()))
                 }.build()
             }
 
@@ -221,15 +218,15 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
                     credentialStartUrl = getStartUrl(context.project),
                     cwsprChatCodeBlockIndex = message.codeBlockIndex,
                     cwsprChatTotalCodeBlocks = message.totalCodeBlocks,
-                    cwsprChatHasProjectContext = hasProjectContext(message.messageId)
+                    cwsprChatHasProjectContext = getMessageHasProjectContext(message.messageId)
                 )
                 ChatInteractWithMessageEvent.builder().apply {
                     conversationId(getConversationId(message.tabId).orEmpty())
                     messageId(message.messageId)
                     interactionType(ChatMessageInteractionType.COPY_SNIPPET)
                     interactionTarget(message.insertionTargetType)
                     acceptedCharacterCount(message.code.length)
-                    hasProjectLevelContext(hasProjectContext(message.messageId))
+                    hasProjectLevelContext(getMessageHasProjectContext(message.messageId))
                 }.build()
             }
 
@@ -245,7 +242,7 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
                     credentialStartUrl = getStartUrl(context.project),
                     cwsprChatCodeBlockIndex = message.codeBlockIndex,
                     cwsprChatTotalCodeBlocks = message.totalCodeBlocks,
-                    cwsprChatHasProjectContext = hasProjectContext(message.messageId)
+                    cwsprChatHasProjectContext = getMessageHasProjectContext(message.messageId)
                 )
                 ChatInteractWithMessageEvent.builder().apply {
                     conversationId(getConversationId(message.tabId).orEmpty())
@@ -254,7 +251,7 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
                     interactionTarget(message.insertionTargetType)
                     acceptedCharacterCount(message.code.length)
                     acceptedLineCount(message.code.lines().size)
-                    hasProjectLevelContext(hasProjectContext(message.messageId))
+                    hasProjectLevelContext(getMessageHasProjectContext(message.messageId))
                 }.build()
             }
 
@@ -274,7 +271,7 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
                     cwsprChatInteractionTarget = message.link,
                     cwsprChatHasReference = null,
                     credentialStartUrl = getStartUrl(context.project),
-                    cwsprChatHasProjectContext = hasProjectContext(message.messageId)
+                    cwsprChatHasProjectContext = getMessageHasProjectContext(message.messageId)
                 )
                 ChatInteractWithMessageEvent.builder().apply {
                     conversationId(getConversationId(message.tabId).orEmpty())
@@ -287,7 +284,7 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
                         }
                     )
                     interactionTarget(message.link)
-                    hasProjectLevelContext(hasProjectContext(message.messageId))
+                    hasProjectLevelContext(getMessageHasProjectContext(message.messageId))
                 }.build()
             }
 
@@ -374,6 +371,14 @@ class TelemetryHelper(private val context: AmazonQAppInitContext, private val se
         responseStreamTotalTime[tabId] = totalTime
     }
 
+    fun setResponseHasProjectContext(data: ChatRequestData, response: ChatMessage) {
+        responseHasProjectContext[response.messageId] = getIsProjectContextEnabled() && data.useRelevantDocuments && data.relevantTextDocuments.isNotEmpty()
+    }
+
+    fun getMessageHasProjectContext(messageId: String): Boolean {
+        return responseHasProjectContext.getOrDefault(messageId, false)
+    }
+
     @VisibleForTesting
     fun getResponseStreamTimeToFirstChunk(tabId: String): Double {
         val chunkTimes = responseStreamTimeForChunks[tabId] ?: return 0.0
diff --git a/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/TelemetryHelperTest.kt b/plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/TelemetryHelperTest.kt
@@ -331,6 +331,7 @@ class TelemetryHelperTest {
                     messageId(messageId)
                     interactionType(ChatMessageInteractionType.UPVOTE)
                     customizationArn(customizationArn)
+                    hasProjectLevelContext(false)
                 }.build()
             )
         )
@@ -373,6 +374,7 @@ class TelemetryHelperTest {
                     messageId(messageId)
                     interactionType(ChatMessageInteractionType.CLICK_FOLLOW_UP)
                     customizationArn(customizationArn)
+                    hasProjectLevelContext(false)
                 }.build()
             )
         )
@@ -429,6 +431,7 @@ class TelemetryHelperTest {
                     interactionTarget("insertionTargetType")
                     acceptedCharacterCount("println()".length)
                     customizationArn(customizationArn)
+                    hasProjectLevelContext(false)
                 }.build()
             )
         )
@@ -462,7 +465,6 @@ class TelemetryHelperTest {
         mockClient.stub {
             on { this.sendChatInteractWithMessageTelemetry(any<ChatInteractWithMessageEvent>()) } doReturn mockSteResponse
         }
-
         val codeBlockIndex = 1
         val totalCodeBlocks = 10
         val inserTionTargetType = "insertionTargetType"
@@ -493,6 +495,7 @@ class TelemetryHelperTest {
                     acceptedCharacterCount(code.length)
                     acceptedLineCount(code.lines().size)
                     customizationArn(customizationArn)
+                    hasProjectLevelContext(false)
                 }.build()
             )
         )
@@ -552,6 +555,7 @@ class TelemetryHelperTest {
                     interactionType(ChatMessageInteractionType.CLICK_LINK)
                     interactionTarget(link)
                     customizationArn(customizationArn)
+                    hasProjectLevelContext(false)
                 }.build()
             )
         )

Original file line number	Diff line number	Diff line change
`@@ -90,6 +90,7 @@ class ChatPromptHandler(private val telemetryHelper: TelemetryHelper) {`
`90`	`90`	`ChatMessage(tabId = tabId, triggerId = triggerId, messageId = requestId, messageType = ChatMessageType.Answer, followUps = followUps)`
`91`	`91`
`92`	`92`	`telemetryHelper.setResponseStreamTotalTime(tabId)`
	`93`	`+ telemetryHelper.setResponseHasProjectContext(data, response)`
`93`	`94`	`telemetryHelper.recordAddMessage(data, response, responseText.length, statusCode, countTotalNumberOfCodeBlocks(responseText))`
`94`	`95`	`emit(response)`
`95`	`96`	`}`
Original file line number	Diff line number	Diff line change
`@@ -331,6 +331,7 @@ class TelemetryHelperTest {`
`331`	`331`	`messageId(messageId)`
`332`	`332`	`interactionType(ChatMessageInteractionType.UPVOTE)`
`333`	`333`	`customizationArn(customizationArn)`
	`334`	`+ hasProjectLevelContext(false)`
`334`	`335`	`}.build()`
`335`	`336`	`)`
`336`	`337`	`)`
`@@ -373,6 +374,7 @@ class TelemetryHelperTest {`
`373`	`374`	`messageId(messageId)`
`374`	`375`	`interactionType(ChatMessageInteractionType.CLICK_FOLLOW_UP)`
`375`	`376`	`customizationArn(customizationArn)`
	`377`	`+ hasProjectLevelContext(false)`
`376`	`378`	`}.build()`
`377`	`379`	`)`
`378`	`380`	`)`
`@@ -429,6 +431,7 @@ class TelemetryHelperTest {`
`429`	`431`	`interactionTarget("insertionTargetType")`
`430`	`432`	`acceptedCharacterCount("println()".length)`
`431`	`433`	`customizationArn(customizationArn)`
	`434`	`+ hasProjectLevelContext(false)`
`432`	`435`	`}.build()`
`433`	`436`	`)`
`434`	`437`	`)`
`@@ -462,7 +465,6 @@ class TelemetryHelperTest {`
`462`	`465`	`mockClient.stub {`
`463`	`466`	`on { this.sendChatInteractWithMessageTelemetry(any<ChatInteractWithMessageEvent>()) } doReturn mockSteResponse`
`464`	`467`	`}`
`465`		`-`
`466`	`468`	`val codeBlockIndex = 1`
`467`	`469`	`val totalCodeBlocks = 10`
`468`	`470`	`val inserTionTargetType = "insertionTargetType"`
`@@ -493,6 +495,7 @@ class TelemetryHelperTest {`
`493`	`495`	`acceptedCharacterCount(code.length)`
`494`	`496`	`acceptedLineCount(code.lines().size)`
`495`	`497`	`customizationArn(customizationArn)`
	`498`	`+ hasProjectLevelContext(false)`
`496`	`499`	`}.build()`
`497`	`500`	`)`
`498`	`501`	`)`
`@@ -552,6 +555,7 @@ class TelemetryHelperTest {`
`552`	`555`	`interactionType(ChatMessageInteractionType.CLICK_LINK)`
`553`	`556`	`interactionTarget(link)`
`554`	`557`	`customizationArn(customizationArn)`
	`558`	`+ hasProjectLevelContext(false)`
`555`	`559`	`}.build()`
`556`	`560`	`)`
`557`	`561`	`)`