fix(amazonq): Addition of wait function for Amazon Q inline chat UI E2E Tests and Inline Keybind Shortcut Test

packages/amazonq/test/e2e_new/amazonq/tests/inline.test.ts

@@ -5,7 +5,7 @@
 import '../utils/setup'
 import { Workbench, EditorView, InputBox, TextEditor, WebviewView, Key } from 'vscode-extension-tester'
 import { testContext } from '../utils/testContext'
-import { createNewTextFile, writeToTextEditor, sleep } from '../utils/generalUtils'
+import { createNewTextFile, writeToTextEditor, waitForInlineGeneration, pressShortcut } from '../utils/generalUtils'
 import assert from 'assert'
 
 describe('Amazon Q Inline Completion / Chat Functionality', function () {
@@ -16,35 +16,43 @@ describe('Amazon Q Inline Completion / Chat Functionality', function () {
     let textEditor: TextEditor
     let webviewView: WebviewView
 
-    before(async function () {
+    beforeEach(async function () {
         webviewView = testContext.webviewView
         await webviewView.switchBack()
         workbench = testContext.workbench
         editorView = new EditorView()
         testContext.editorView = editorView
         textEditor = await createNewTextFile(workbench, editorView)
     })
-    after(async function () {
-        // Switch back to Webview Iframe when dealing with external webviews from Amazon Q.
+    afterEach(async function () {
+        // Switch back to iframe when dealing with external webviews from Amazon Q.
         await editorView.closeAllEditors()
         await webviewView.switchToFrame()
     })
-    it('Inline Test Shortcut', async () => {
-        await writeToTextEditor(textEditor, 'def factorial(n):')
-        const text = await textEditor.getText()
-        assert.equal(text, 'def factorial(n): ')
-        await textEditor.clearText()
-
+    it('Inline Test Generate', async () => {
         const textBefore = await textEditor.getText()
         await workbench.executeCommand('Amazon Q: Inline Chat')
         const input = new InputBox()
         await input.sendKeys('Generate the fibonacci sequence through iteration')
         await input.sendKeys(Key.ENTER)
-        // Must wait for response to be generated.
-        await sleep(8000)
+        // Wait for Amazon Q to finish generating code
+        await waitForInlineGeneration(textEditor)
 
         const textAfter = await textEditor.getText()
-        assert(textAfter.length > textBefore.length, 'Amazon Q should have generated code')
+        assert(textAfter.length > textBefore.length, 'Amazon Q generated code')
+        await textEditor.clearText()
+    })
+    it('Inline Keybind Shortcut', async () => {
+        await writeToTextEditor(textEditor, 'def fibonacci(n):')
+        await textEditor.selectText('def fibonacci(n):')
+
+        const driver = webviewView.getDriver()
+        await pressShortcut(driver, Key.COMMAND, 'i')
+        const input = new InputBox()
+        await input.sendKeys('Generate the fibonacci sequence through recursion')
+        await input.sendKeys(Key.ENTER)
+        await waitForInlineGeneration(textEditor)
+        // Clean Up Text
         await textEditor.clearText()
     })
 })

packages/amazonq/test/e2e_new/amazonq/utils/generalUtils.ts

@@ -105,7 +105,7 @@ export async function pressKey(driver: WebDriver, key: keyof typeof Key): Promis
  * Ctrl + C | await pressShortcut(driver, Key.CONTROL, 'c')
  * Ctrl + Shift + T | await pressShortcut(driver, Key.CONTROL, Key.SHIFT, 't')
  */
-export async function pressShortcut(driver: WebDriver, ...keys: (keyof typeof Key)[]): Promise<void> {
+export async function pressShortcut(driver: WebDriver, ...keys: (string | keyof typeof Key)[]): Promise<void> {
     const actions = driver.actions()
     for (const key of keys) {
         actions.keyDown(key)
@@ -211,6 +211,39 @@ export async function writeToTextEditor(textEditor: TextEditor, text: string): P
     const currentLines = await textEditor.getNumberOfLines()
     await textEditor.typeTextAt(currentLines, 1, text)
 }
+/**
+ * Waits for Inline Generation by Amazon Q by checking if line count stops changing.
+ * The function checks for a "stable state" by monitoring the number of lines in the editor.
+ * A stable state is achieved when the line count remains unchanged for 3 consecutive checks (3 seconds).
+ * Checks are performed every 1 second.
+ * @param editor The TextEditor instance
+ * @param timeout Maximum time to wait in milliseconds (default: 15000). Function will throw an error if generation takes longer than this timeout.
+ * @returns Promise<void>
+ * @throws Error if timeout is exceeded before a stable state is reached
+ */
+export async function waitForInlineGeneration(editor: TextEditor, timeout = 15000): Promise<void> {
+    const startTime = Date.now()
+    let previousLines = await editor.getNumberOfLines()
+    let stableCount = 0
+
+    while (Date.now() - startTime < timeout) {
+        await sleep(1000)
+        const currentLines = await editor.getNumberOfLines()
+
+        if (currentLines === previousLines) {
+            stableCount++
+            if (stableCount >= 3) {
+                return
+            }
+        } else {
+            stableCount = 0
+        }
+
+        previousLines = currentLines
+    }
+
+    throw new Error(`Editor stabilization timed out after ${timeout}ms`)
+}
 
 /**
  * Finds an item based on the text