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

## Problem
Currently we use a `sleep(8000)` function which causes us to wait an
arbitrary 8 seconds for Amazon Q to implement the fibonacci sequence.
This can cause flakiness and is generally not a good practice. We also
have not fully updated our inline test and our writeToTextEditor
function does not work well. We also have not implemented an inline test
which uses keybinds to start an inline chat suggestion.

## Solution
I have implemented a function which constantly checks for a "stable"
state in which we check for any new number of lines generated by Amazon
Q. If this stable state is repeated more than 2 times, meaning the
number of lines has stayed the same for 3 seconds (1 second wait in
between), we will confirm that AmazonQ has finished generating its
response. We have a maximum generate time of 15 seconds, as a fibonacci
sequence should not take more than that, or something is clearly wrong
with the model and we will error out. We have implemented the inline
tests and the writeToTextEditor function fix.


---

- Treat all work as PUBLIC. Private `feature/x` branches will not be
squash-merged at release time.
- Your code changes must meet the guidelines in
[CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines).
- License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

---------

Co-authored-by: laura-codess <[email protected]>

Author: surajrdy-aws

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

@@ -109,7 +109,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)
@@ -215,6 +215,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