fixed the PR based on the comments

Author: laura-codess

packages/amazonq/test/e2e_new/amazonq/chat.test.ts

@@ -7,6 +7,7 @@ import { WebviewView } from 'vscode-extension-tester'
 import { closeAllTabs } from './framework/cleanupHelper'
 import { testContext } from './utils/testContext'
 import { waitForChatResponse, writeToChat } from './framework/chatHelper'
+import assert from 'assert'
 
 describe('Amazon Q Chat Basic Functionality', function () {
     // this timeout is the general timeout for the entire test suite
@@ -18,7 +19,11 @@ describe('Amazon Q Chat Basic Functionality', function () {
     })
 
     afterEach(async () => {
-        await closeAllTabs(webviewView)
+        try {
+            await closeAllTabs(webviewView)
+        } catch (e) {
+            assert.fail(`Failed to clean up tabs: ${e}`)
+        }
     })
 
     it('Chat Prompt Test', async () => {

packages/amazonq/test/e2e_new/amazonq/framework/chatHelper.ts

@@ -5,13 +5,6 @@
 import { By, WebviewView } from 'vscode-extension-tester'
 import { waitForElement } from './generalHelper'
 
-/* Writes a prompt to the chat input and waits for a response
-
-Logic:
-Finds the chat input element using the .mynah-chat-prompt-input CSS selector,
-sends the provided prompt test, clicks the send button, and waits for a chat
-response. Returns true if successful, throws an error if the response times out */
-
 export async function writeToChat(prompt: string, webview: WebviewView): Promise<boolean> {
     const chatInput = await waitForElement(webview, By.css('.mynah-chat-prompt-input'))
     await chatInput.sendKeys(prompt)
@@ -20,14 +13,6 @@ export async function writeToChat(prompt: string, webview: WebviewView): Promise
     return true
 }
 
-/* Waits for a chat response and outputs whether the response is "correct"
-
-Logic: 
-The overall conversation container's css is .mynah-chat-items-conversation-container. 
-Within that container we can check how many elements exist. If there is 2 elements,
-we can assume that the chat response has been generated. However, we must grab the 
-latest conversation container, as there can be multiple conversations in the webview. */
-
 export async function waitForChatResponse(webview: WebviewView, timeout = 15000): Promise<boolean> {
     const startTime = Date.now()
 

packages/amazonq/test/e2e_new/amazonq/framework/cleanupHelper.ts

@@ -4,14 +4,7 @@
  */
 import { By, WebviewView } from 'vscode-extension-tester'
 
-/* Finds all the tabs by looking for the close buttons and then closes them one by one.
-
-Logic:
-There is a button with the css mynah-tabs-close-button, we need to click that button and 
-close all the tabs after the test is done to avoid memory from a previous test. To double
-check if all the tabs are closed, we can check if the mynah-tabs-container is empty. */
-
-export async function closeAllTabs(webview: WebviewView) {
+export async function closeAllTabs(webview: WebviewView): Promise<boolean> {
     try {
         const closeButtons = await webview.findWebElements(By.css('.mynah-tabs-close-button'))
 
@@ -20,15 +13,19 @@ export async function closeAllTabs(webview: WebviewView) {
             await new Promise((resolve) => setTimeout(resolve, 500))
         }
 
-        // double check that all tabs are closed by checking if the mynah-tabs-container is empty
         const tabsContainer = await webview.findWebElements(By.css('.mynah-tabs-container'))
-        if (
+        const allClosed =
             tabsContainer.length === 0 ||
             (await tabsContainer[0].findElements(By.css('.mynah-tab-item-label'))).length === 0
-        ) {
+
+        if (allClosed) {
             console.log('All chat tabs successfully closed')
+            return true
+        } else {
+            throw new Error('Failed to close all tabs')
         }
     } catch (error) {
-        console.log('Error closing tabs:', error)
+        console.error('Error closing tabs:', error)
+        throw error
     }
 }

packages/amazonq/test/e2e_new/amazonq/framework/generalHelper.ts

@@ -5,16 +5,7 @@
 import { By, WebviewView, WebElement } from 'vscode-extension-tester'
 import { until } from 'selenium-webdriver'
 
-/* Waits for an element (or multiple elements) to appear based on the parameters
-
-Logic:
-The function utilizes the Selenium wait driver. We can call that driver from our
-WebviewView but it can also be called on parts of the VSCode Editor that are not
-part of the WebviewView.
-
-(TO-DO: create a more general function that can be called on any part of the VSCode
-Editor. Will do when a use case appears for it.*/
-
+/* Note: If multiple is set to False, then it will return the first element it finds that matches the Locator*/
 export async function waitForElement(
     webview: WebviewView,
     locator: By,
@@ -38,13 +29,6 @@ export async function waitForElement(
     return multiple ? await webview.findWebElements(locator) : await webview.findWebElement(locator)
 }
 
-/* General function for finding WebElement by their text content
-
-Logic:
-It searches through an array of WebElements and looks for an element with 
-the ".tittle" CSS class within each item. Compares the text content and returns
-the first matching parent element, or throws an error if not found. */
-
 export async function findItemByText(items: WebElement[], text: string) {
     for (const item of items) {
         const titleDivs = await item.findElements(By.css('.title'))

packages/amazonq/test/e2e_new/amazonq/framework/loginHelper.ts

@@ -13,8 +13,7 @@ Currently, the function will
 3. Inputs the Start URL
 4. IMPORTANT: you must click manually open yourself when the popup window asks to open the browser and complete the authentication in the browser** 
 
-TO-DO: Currently this loginToAmazonQ is not fully autonomous as we ran into a blocker when the browser window pops up
-Documentation: https://quip-amazon.com/PoJOAyt4ja8H/Authentication-for-UI-Tests-Documentation */
+TO-DO: Currently this loginToAmazonQ is not fully autonomous as we ran into a blocker when the browser window pops up */
 
 export async function loginToAmazonQ(): Promise<void> {
     const workbench = new Workbench()

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

@@ -5,7 +5,7 @@
 import { loginToAmazonQ } from '../framework/loginHelper'
 
 before(async function () {
-    this.timeout(60000) // Increase timeout to 60 seconds
+    this.timeout(60000)
     console.log('\n\n*** MANUAL INTERVENTION REQUIRED ***')
     console.log('When prompted, you must manually click to open the browser and complete authentication')
     console.log('You have 60 seconds to complete this step\n\n')

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

@@ -5,9 +5,26 @@
 import { Workbench, WebviewView } from 'vscode-extension-tester'
 
 export interface TestContext {
-    workbench?: Workbench
-    webviewView?: WebviewView
+    workbench: Workbench
+    webviewView: WebviewView
 }
 
-// arr to store shared context
-export const testContext: TestContext = {}
+export const testContext = new Proxy<TestContext>({} as TestContext, {
+    get(target, prop) {
+        if (prop in target && target[prop as keyof TestContext] !== undefined) {
+            return target[prop as keyof TestContext]
+        }
+        throw new Error(
+            `TestContext.${String(prop)} is undefined. Make sure setup.ts has properly initialized the test context.`
+        )
+    },
+    set(target, prop, value) {
+        target[prop as keyof TestContext] = value
+        return true
+    },
+})
+
+export function initializeTestContext(workbench: Workbench, webviewView: WebviewView): void {
+    testContext.workbench = workbench
+    testContext.webviewView = webviewView
+}