From 483591b983d4c174c73fcefe249315a791143b18 Mon Sep 17 00:00:00 2001
From: Laxman Reddy <141967714+laileni-aws@users.noreply.github.com>
Date: Tue, 15 Oct 2024 10:50:42 -0700
Subject: [PATCH] Moving the extractCodeBlockLanguage function to
 shared/markdown.ts and added test cases for this function

---
 .../controllers/chat/messenger/messenger.ts   | 21 +---------
 packages/core/src/shared/markdown.ts          | 22 +++++++++++
 .../core/src/test/shared/markdown.test.ts     | 38 +++++++++++++++++++
 3 files changed, 62 insertions(+), 19 deletions(-)
 create mode 100644 packages/core/src/shared/markdown.ts
 create mode 100644 packages/core/src/test/shared/markdown.test.ts

diff --git a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts
index 706dad51954..41a31d30edd 100644
--- a/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts
+++ b/packages/core/src/codewhispererChat/controllers/chat/messenger/messenger.ts
@@ -33,6 +33,7 @@ import { CodeScanIssue } from '../../../../codewhisperer/models/model'
 import { marked } from 'marked'
 import { JSDOM } from 'jsdom'
 import { LspController } from '../../../../amazonq/lsp/lspController'
+import { extractCodeBlockLanguage } from '../../../../shared/markdown'
 
 export type StaticTextResponseType = 'quick-action-help' | 'onboarding-help' | 'transform' | 'help'
 
@@ -185,7 +186,7 @@ export class Messenger {
                     ) {
                         message += chatEvent.assistantResponseEvent.content
                         if (codeBlockLanguage === 'plaintext') {
-                            codeBlockLanguage = this.extractCodeBlockLanguage(message)
+                            codeBlockLanguage = extractCodeBlockLanguage(message)
                         }
                         this.dispatcher.sendChatMessage(
                             new ChatMessage(
@@ -342,24 +343,6 @@ export class Messenger {
             })
     }
 
-    private 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.
-        const codeBlockStart = message.indexOf('```')
-        if (codeBlockStart === -1) {
-            return 'plaintext'
-        }
-
-        const languageStart = codeBlockStart + 3
-        const languageEnd = message.indexOf('\n', languageStart)
-
-        if (languageEnd === -1) {
-            return 'plaintext'
-        }
-
-        const language = message.substring(languageStart, languageEnd).trim()
-        return language !== '' ? language : 'plaintext'
-    }
-
     public sendErrorMessage(errorMessage: string | undefined, tabID: string, requestID: string | undefined) {
         this.showChatExceptionMessage(
             {
diff --git a/packages/core/src/shared/markdown.ts b/packages/core/src/shared/markdown.ts
new file mode 100644
index 00000000000..215e6603906
--- /dev/null
+++ b/packages/core/src/shared/markdown.ts
@@ -0,0 +1,22 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export const extractCodeBlockLanguage = (message: string) => {
+    // This fulfills both the cases of unit test generation(java, python) and general use case(Non java and Non python) languages.
+    const codeBlockStart = message.indexOf('```')
+    if (codeBlockStart === -1) {
+        return 'plaintext'
+    }
+
+    const languageStart = codeBlockStart + 3
+    const languageEnd = message.indexOf('\n', languageStart)
+
+    if (languageEnd === -1) {
+        return 'plaintext'
+    }
+
+    const language = message.substring(languageStart, languageEnd).trim()
+    return language !== '' ? language : 'plaintext'
+}
diff --git a/packages/core/src/test/shared/markdown.test.ts b/packages/core/src/test/shared/markdown.test.ts
new file mode 100644
index 00000000000..306ff06ce48
--- /dev/null
+++ b/packages/core/src/test/shared/markdown.test.ts
@@ -0,0 +1,38 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import { extractCodeBlockLanguage } from '../../shared/markdown'
+
+describe('extractCodeBlockLanguage', () => {
+    it('should return "plaintext" when no code block is present', () => {
+        const message = 'This is a message without a code block'
+        assert.strictEqual(extractCodeBlockLanguage(message), 'plaintext')
+    })
+
+    it('should return the language when a code block with language is present', () => {
+        const message = 'Here is some code:\n```javascript\nconsole.log("Hello");\n```'
+        assert.strictEqual(extractCodeBlockLanguage(message), 'javascript')
+    })
+
+    it('should return "plaintext" when a code block is present but no language is specified', () => {
+        const message = 'Here is some code:\n```\nconsole.log("Hello");\n```'
+        assert.strictEqual(extractCodeBlockLanguage(message), 'plaintext')
+    })
+
+    it('should handle whitespace before the language specification', () => {
+        const message = 'Code:\n```   typescript\nconst x: number = 5;\n```'
+        assert.strictEqual(extractCodeBlockLanguage(message), 'typescript')
+    })
+
+    it('should return "plaintext" when the code block is not closed', () => {
+        const message = 'Incomplete code block:\n```javascript\nconsole.log("Hello");'
+        assert.strictEqual(extractCodeBlockLanguage(message), 'plaintext')
+    })
+
+    it('should handle empty messages', () => {
+        assert.strictEqual(extractCodeBlockLanguage(''), 'plaintext')
+    })
+})