Moved extractCodeBlockLanguage to TextUtils and added test cases

Author: laileni-aws

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/chat/messenger/ChatPromptHandler.kt

@@ -11,6 +11,7 @@ import kotlinx.coroutines.flow.onStart
 import software.amazon.awssdk.awscore.exception.AwsServiceException
 import software.amazon.awssdk.services.codewhispererstreaming.model.CodeWhispererStreamingException
 import software.aws.toolkits.core.utils.convertMarkdownToHTML
+import software.aws.toolkits.core.utils.extractCodeBlockLanguage
 import software.aws.toolkits.jetbrains.services.cwc.clients.chat.exceptions.ChatApiException
 import software.aws.toolkits.jetbrains.services.cwc.clients.chat.model.ChatRequestData
 import software.aws.toolkits.jetbrains.services.cwc.clients.chat.model.ChatResponseEvent
@@ -145,24 +146,6 @@ class ChatPromptHandler(private val telemetryHelper: TelemetryHelper) {
             }
     }
 
-    private fun extractCodeBlockLanguage(message: String): String {
-        // This fulfills both the cases of unit test generation(java, python) and general use case(Non java and Non python) languages.
-        val codeBlockStart = message.indexOf("```")
-        if (codeBlockStart == -1) {
-            return defaultTestGenResponseLanguage
-        }
-
-        val languageStart = codeBlockStart + 3
-        val languageEnd = message.indexOf('\n', languageStart)
-
-        if (languageEnd == -1) {
-            return defaultTestGenResponseLanguage
-        }
-
-        val language = message.substring(languageStart, languageEnd).trim()
-        return if (language.isNotEmpty()) language else defaultTestGenResponseLanguage
-    }
-
     private fun processChatEvent(
         tabId: String,
         triggerId: String,

plugins/core/core/src/software/aws/toolkits/core/utils/TextUtils.kt

@@ -17,6 +17,24 @@ fun convertMarkdownToHTML(markdown: String): String {
     return htmlRenderer.render(document)
 }
 
+fun extractCodeBlockLanguage(message: String): String {
+    // This fulfills both the cases of unit test generation(java, python) and general use case(Non java and Non python) languages.
+    val defaultTestGenResponseLanguage: String = "plaintext"
+    val codeBlockStart = message.indexOf("```")
+    if (codeBlockStart == -1) {
+        return defaultTestGenResponseLanguage
+    }
+
+    val languageStart = codeBlockStart + 3
+    val languageEnd = message.indexOf('\n', languageStart)
+
+    if (languageEnd == -1) {
+        return defaultTestGenResponseLanguage
+    }
+
+    return message.substring(languageStart, languageEnd).trim().ifEmpty { defaultTestGenResponseLanguage }
+}
+
 class CodeBlockRenderer(private val html: HtmlWriter) : NodeRenderer {
     override fun getNodeTypes(): Set<Class<out Node>> = setOf(FencedCodeBlock::class.java)
     override fun render(node: Node?) {

plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/utils/TextUtilsTest.kt

@@ -11,12 +11,15 @@ import org.intellij.lang.annotations.Language
 import org.junit.Rule
 import org.junit.Test
 import software.aws.toolkits.core.utils.convertMarkdownToHTML
+import software.aws.toolkits.core.utils.extractCodeBlockLanguage
 
 class TextUtilsTest {
     @Rule
     @JvmField
     val projectRule = ProjectRule()
 
+    private val defaultTestGenResponseLanguage = "plaintext"
+
     @Test
     fun textGetsFormatted() {
         @Language("JSON")
@@ -155,4 +158,46 @@ class TextUtilsTest {
         val actual = applyPatch(inputPatch, fileContent, inputFilePath)
         assertThat(actual).isEqualTo("dummy\ndummy\ndummy\ndummy\nfirst line\nthird line\nforth line")
     }
+
+    @Test
+    fun `extractCodeBlockLanguage returns default when no code block is present`() {
+        val message = "This is a message without a code block"
+        assertThat(defaultTestGenResponseLanguage).isEqualTo(extractCodeBlockLanguage(message))
+    }
+
+    @Test
+    fun `extractCodeBlockLanguage returns language when code block with language is present`() {
+        val message = "Here's a code block:\n```kotlin\nval x = 5\n```"
+        assertThat("kotlin").isEqualTo(extractCodeBlockLanguage(message))
+    }
+
+    @Test
+    fun `extractCodeBlockLanguage returns default when code block has no language specified`() {
+        val message = "Here's a code block:\n```\nval x = 5\n```"
+        assertThat(defaultTestGenResponseLanguage).isEqualTo(extractCodeBlockLanguage(message))
+    }
+
+    @Test
+    fun `extractCodeBlockLanguage returns language when multiple code blocks are present`() {
+        val message = "First block:\n```java\nint x = 5;\n```\nSecond block:\n```python\nx = 5\n```"
+        assertThat("java").isEqualTo(extractCodeBlockLanguage(message))
+    }
+
+    @Test
+    fun `extractCodeBlockLanguage returns default when code block is not closed`() {
+        val message = "Incomplete code block:\n```kotlin\nval x = 5"
+        assertThat("kotlin").isEqualTo(extractCodeBlockLanguage(message))
+    }
+
+    @Test
+    fun `extractCodeBlockLanguage trims whitespace from language`() {
+        val message = "Code block with spaces:\n```  kotlin  \nval x = 5\n```"
+        assertThat("kotlin").isEqualTo(extractCodeBlockLanguage(message))
+    }
+
+    @Test
+    fun `extractCodeBlockLanguage handles empty message`() {
+        val message = ""
+        assertThat(defaultTestGenResponseLanguage).isEqualTo(extractCodeBlockLanguage(message))
+    }
 }